mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
Compare commits
24
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6bc26cf591 | ||
|
|
911e057ccd | ||
|
|
815416cd62 | ||
|
|
f0bd4f5b88 | ||
|
|
c3343aa015 | ||
|
|
a0f61b39d2 | ||
|
|
41ed0a0f98 | ||
|
|
6daa2f155f | ||
|
|
706c3a1b47 | ||
|
|
656151783c | ||
|
|
82822d7b25 | ||
|
|
dcb2d89e3a | ||
|
|
f3da30a454 | ||
|
|
5901451603 | ||
|
|
d3d517fd67 | ||
|
|
01678be61d | ||
|
|
fb3548956b | ||
|
|
5847479fd8 | ||
|
|
3fdb187796 | ||
|
|
b441ebec02 | ||
|
|
fa51b9f52d | ||
|
|
dea789cca9 | ||
|
|
485d1568c4 | ||
|
|
acb590214a |
@@ -54,7 +54,7 @@ jobs:
|
|||||||
- group: features
|
- group: features
|
||||||
files: "tests/e2e/scenarios/test_skills.py tests/e2e/scenarios/test_tool_approval.py tests/e2e/scenarios/test_webhook.py"
|
files: "tests/e2e/scenarios/test_skills.py tests/e2e/scenarios/test_tool_approval.py tests/e2e/scenarios/test_webhook.py"
|
||||||
- group: extensions
|
- group: extensions
|
||||||
files: "tests/e2e/scenarios/test_extensions.py tests/e2e/scenarios/test_extension_oauth.py tests/e2e/scenarios/test_telegram_token_validation.py tests/e2e/scenarios/test_telegram_hot_activation.py tests/e2e/scenarios/test_wasm_lifecycle.py tests/e2e/scenarios/test_tool_execution.py tests/e2e/scenarios/test_pairing.py tests/e2e/scenarios/test_mcp_auth_flow.py tests/e2e/scenarios/test_oauth_credential_fallback.py tests/e2e/scenarios/test_routine_oauth_credential_injection.py"
|
files: "tests/e2e/scenarios/test_extensions.py tests/e2e/scenarios/test_extension_oauth.py tests/e2e/scenarios/test_oauth_url_parameters.py tests/e2e/scenarios/test_telegram_token_validation.py tests/e2e/scenarios/test_telegram_hot_activation.py tests/e2e/scenarios/test_wasm_lifecycle.py tests/e2e/scenarios/test_tool_execution.py tests/e2e/scenarios/test_pairing.py tests/e2e/scenarios/test_mcp_auth_flow.py tests/e2e/scenarios/test_oauth_credential_fallback.py tests/e2e/scenarios/test_routine_oauth_credential_injection.py"
|
||||||
- group: routines
|
- group: routines
|
||||||
files: "tests/e2e/scenarios/test_owner_scope.py tests/e2e/scenarios/test_routine_event_batch.py"
|
files: "tests/e2e/scenarios/test_owner_scope.py tests/e2e/scenarios/test_routine_event_batch.py"
|
||||||
steps:
|
steps:
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ jobs:
|
|||||||
tests:
|
tests:
|
||||||
name: Tests (${{ matrix.name }})
|
name: Tests (${{ matrix.name }})
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
timeout-minutes: 45
|
||||||
strategy:
|
strategy:
|
||||||
fail-fast: false
|
fail-fast: false
|
||||||
matrix:
|
matrix:
|
||||||
@@ -40,11 +41,14 @@ jobs:
|
|||||||
- name: Build WASM channels (for integration tests)
|
- name: Build WASM channels (for integration tests)
|
||||||
run: ./scripts/build-wasm-extensions.sh --channels
|
run: ./scripts/build-wasm-extensions.sh --channels
|
||||||
- name: Run Tests
|
- name: Run Tests
|
||||||
run: cargo test ${{ matrix.flags }} -- --nocapture
|
run: |
|
||||||
|
timeout --signal=INT --kill-after=30s 40m \
|
||||||
|
cargo test ${{ matrix.flags }} -- --nocapture
|
||||||
|
|
||||||
heavy-integration-tests:
|
heavy-integration-tests:
|
||||||
name: Heavy Integration Tests
|
name: Heavy Integration Tests
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
timeout-minutes: 20
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@v6
|
uses: actions/checkout@v6
|
||||||
@@ -58,9 +62,13 @@ jobs:
|
|||||||
- name: Build Telegram WASM channel
|
- name: Build Telegram WASM channel
|
||||||
run: cargo build --manifest-path channels-src/telegram/Cargo.toml --target wasm32-wasip2 --release
|
run: cargo build --manifest-path channels-src/telegram/Cargo.toml --target wasm32-wasip2 --release
|
||||||
- name: Run thread scheduling integration tests
|
- name: Run thread scheduling integration tests
|
||||||
run: cargo test --no-default-features --features libsql,integration --test e2e_thread_scheduling -- --nocapture
|
run: |
|
||||||
|
timeout --signal=INT --kill-after=30s 15m \
|
||||||
|
cargo test --no-default-features --features libsql,integration --test e2e_thread_scheduling -- --nocapture
|
||||||
- name: Run Telegram thread-scope regression test
|
- name: Run Telegram thread-scope regression test
|
||||||
run: cargo test --features integration --test telegram_auth_integration test_private_messages_use_chat_id_as_thread_scope -- --exact
|
run: |
|
||||||
|
timeout --signal=INT --kill-after=30s 10m \
|
||||||
|
cargo test --features integration --test telegram_auth_integration test_private_messages_use_chat_id_as_thread_scope -- --exact
|
||||||
|
|
||||||
telegram-tests:
|
telegram-tests:
|
||||||
name: Telegram Channel Tests
|
name: Telegram Channel Tests
|
||||||
@@ -68,6 +76,7 @@ jobs:
|
|||||||
github.event_name != 'pull_request' ||
|
github.event_name != 'pull_request' ||
|
||||||
github.base_ref != 'staging'
|
github.base_ref != 'staging'
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
timeout-minutes: 15
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@v6
|
uses: actions/checkout@v6
|
||||||
@@ -75,7 +84,9 @@ jobs:
|
|||||||
uses: dtolnay/rust-toolchain@stable
|
uses: dtolnay/rust-toolchain@stable
|
||||||
- uses: Swatinem/rust-cache@v2
|
- uses: Swatinem/rust-cache@v2
|
||||||
- name: Run Telegram Channel Tests
|
- name: Run Telegram Channel Tests
|
||||||
run: cargo test --manifest-path channels-src/telegram/Cargo.toml -- --nocapture
|
run: |
|
||||||
|
timeout --signal=INT --kill-after=30s 10m \
|
||||||
|
cargo test --manifest-path channels-src/telegram/Cargo.toml -- --nocapture
|
||||||
|
|
||||||
windows-build:
|
windows-build:
|
||||||
name: Windows Build (${{ matrix.name }})
|
name: Windows Build (${{ matrix.name }})
|
||||||
@@ -110,6 +121,7 @@ jobs:
|
|||||||
github.event_name != 'pull_request' ||
|
github.event_name != 'pull_request' ||
|
||||||
github.base_ref != 'staging'
|
github.base_ref != 'staging'
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
timeout-minutes: 30
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@v6
|
uses: actions/checkout@v6
|
||||||
@@ -125,7 +137,9 @@ jobs:
|
|||||||
- name: Build all WASM extensions against current WIT
|
- name: Build all WASM extensions against current WIT
|
||||||
run: ./scripts/build-wasm-extensions.sh
|
run: ./scripts/build-wasm-extensions.sh
|
||||||
- name: Instantiation test (host linker compatibility)
|
- name: Instantiation test (host linker compatibility)
|
||||||
run: cargo test --all-features wit_compat -- --nocapture
|
run: |
|
||||||
|
timeout --signal=INT --kill-after=30s 20m \
|
||||||
|
cargo test --all-features wit_compat -- --nocapture
|
||||||
|
|
||||||
bench-compile:
|
bench-compile:
|
||||||
name: Benchmark Compilation
|
name: Benchmark Compilation
|
||||||
|
|||||||
Generated
+21
-12
@@ -157,7 +157,7 @@ version = "1.1.5"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
|
checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.60.2",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -168,7 +168,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"anstyle",
|
"anstyle",
|
||||||
"once_cell_polyfill",
|
"once_cell_polyfill",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.60.2",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -2136,7 +2136,7 @@ dependencies = [
|
|||||||
"libc",
|
"libc",
|
||||||
"option-ext",
|
"option-ext",
|
||||||
"redox_users 0.5.2",
|
"redox_users 0.5.2",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.59.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -2323,7 +2323,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"libc",
|
"libc",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.52.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -3428,6 +3428,7 @@ dependencies = [
|
|||||||
"hyper-util",
|
"hyper-util",
|
||||||
"iana-time-zone",
|
"iana-time-zone",
|
||||||
"insta",
|
"insta",
|
||||||
|
"ironclaw_common",
|
||||||
"ironclaw_safety",
|
"ironclaw_safety",
|
||||||
"json5",
|
"json5",
|
||||||
"libsql",
|
"libsql",
|
||||||
@@ -3485,6 +3486,14 @@ dependencies = [
|
|||||||
"zip",
|
"zip",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ironclaw_common"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ironclaw_safety"
|
name = "ironclaw_safety"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
@@ -4134,7 +4143,7 @@ version = "0.50.3"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.59.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -5472,7 +5481,7 @@ dependencies = [
|
|||||||
"errno",
|
"errno",
|
||||||
"libc",
|
"libc",
|
||||||
"linux-raw-sys 0.12.1",
|
"linux-raw-sys 0.12.1",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.52.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -6154,7 +6163,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
|
checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"libc",
|
"libc",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.60.2",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -6354,9 +6363,9 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tar"
|
name = "tar"
|
||||||
version = "0.4.44"
|
version = "0.4.45"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "1d863878d212c87a19c1a610eb53bb01fe12951c0501cf5a0d65f724914a667a"
|
checksum = "22692a6476a21fa75fdfc11d452fda482af402c008cdbaf3476414e122040973"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"filetime",
|
"filetime",
|
||||||
"libc",
|
"libc",
|
||||||
@@ -6379,7 +6388,7 @@ dependencies = [
|
|||||||
"getrandom 0.4.2",
|
"getrandom 0.4.2",
|
||||||
"once_cell",
|
"once_cell",
|
||||||
"rustix 1.1.4",
|
"rustix 1.1.4",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.52.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -7179,7 +7188,7 @@ checksum = "f2f6fb2847f6742cd76af783a2a2c49e9375d0a111c7bef6f71cd9e738c72d6e"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"memoffset",
|
"memoffset",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.60.2",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -8029,7 +8038,7 @@ version = "0.1.11"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.48.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|||||||
+4
-1
@@ -1,5 +1,5 @@
|
|||||||
[workspace]
|
[workspace]
|
||||||
members = [".", "crates/ironclaw_safety"]
|
members = [".", "crates/ironclaw_common", "crates/ironclaw_safety"]
|
||||||
exclude = [
|
exclude = [
|
||||||
"channels-src/discord",
|
"channels-src/discord",
|
||||||
"channels-src/telegram",
|
"channels-src/telegram",
|
||||||
@@ -100,6 +100,9 @@ tower-http = { version = "0.6", features = ["trace", "cors", "set-header"] }
|
|||||||
# Cron scheduling for routines
|
# Cron scheduling for routines
|
||||||
cron = "0.13"
|
cron = "0.13"
|
||||||
|
|
||||||
|
# Shared types
|
||||||
|
ironclaw_common = { path = "crates/ironclaw_common", version = "0.1.0" }
|
||||||
|
|
||||||
# Safety/sanitization
|
# Safety/sanitization
|
||||||
ironclaw_safety = { path = "crates/ironclaw_safety", version = "0.1.0" }
|
ironclaw_safety = { path = "crates/ironclaw_safety", version = "0.1.0" }
|
||||||
regex = "1"
|
regex = "1"
|
||||||
|
|||||||
+1
-1
@@ -161,7 +161,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
| `config` | ✅ | ✅ | - | Read/write config plus validate/path helpers |
|
| `config` | ✅ | ✅ | - | Read/write config plus validate/path helpers |
|
||||||
| `backup` | ✅ | ❌ | P3 | Create/verify local backup archives |
|
| `backup` | ✅ | ❌ | P3 | Create/verify local backup archives |
|
||||||
| `channels` | ✅ | 🚧 | P2 | `list` implemented; `enable`/`disable`/`status` deferred pending config source unification |
|
| `channels` | ✅ | 🚧 | P2 | `list` implemented; `enable`/`disable`/`status` deferred pending config source unification |
|
||||||
| `models` | ✅ | 🚧 | - | Model selector in TUI |
|
| `models` | ✅ | 🚧 | P1 | `models list [<provider>]` (`--verbose`, `--json`; fetches live model list when provider specified), `models status` (`--json`), `models set <model>`, `models set-provider <provider> [--model model]` (alias normalization, config.toml + .env persistence). Remaining: `set` doesn't validate model against live list. |
|
||||||
| `status` | ✅ | ✅ | - | System status (enriched session details) |
|
| `status` | ✅ | ✅ | - | System status (enriched session details) |
|
||||||
| `agents` | ✅ | ❌ | P3 | Multi-agent management |
|
| `agents` | ✅ | ❌ | P3 | Multi-agent management |
|
||||||
| `sessions` | ✅ | ❌ | P3 | Session listing (shows subagent models) |
|
| `sessions` | ✅ | ❌ | P3 | Session listing (shows subagent models) |
|
||||||
|
|||||||
@@ -21,7 +21,7 @@
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"name": "feishu_app_secret",
|
"name": "feishu_app_secret",
|
||||||
"prompt": "Enter your Feishu/Lark App Secret",
|
"prompt": "Enter your Feishu/Lark App Secret (from your app settings at open.feishu.cn)",
|
||||||
"optional": false
|
"optional": false
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -0,0 +1,18 @@
|
|||||||
|
[package]
|
||||||
|
name = "ironclaw_common"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
rust-version = "1.92"
|
||||||
|
description = "Shared types and utilities for the IronClaw workspace"
|
||||||
|
authors = ["NEAR AI <[email protected]>"]
|
||||||
|
license = "MIT OR Apache-2.0"
|
||||||
|
homepage = "https://github.com/nearai/ironclaw"
|
||||||
|
repository = "https://github.com/nearai/ironclaw"
|
||||||
|
publish = false
|
||||||
|
|
||||||
|
[package.metadata.dist]
|
||||||
|
dist = false
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
@@ -0,0 +1,393 @@
|
|||||||
|
//! Application-wide event types.
|
||||||
|
//!
|
||||||
|
//! `AppEvent` is the real-time event protocol used across the entire
|
||||||
|
//! application. The web gateway serialises these to SSE / WebSocket
|
||||||
|
//! frames, but other subsystems (agent loop, orchestrator, extensions)
|
||||||
|
//! produce and consume them too.
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// A single tool decision in a reasoning update (SSE DTO).
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolDecisionDto {
|
||||||
|
pub tool_name: String,
|
||||||
|
pub rationale: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolDecisionDto {
|
||||||
|
/// Parse a list of tool decisions from a JSON array value.
|
||||||
|
pub fn from_json_array(value: &serde_json::Value) -> Vec<Self> {
|
||||||
|
value
|
||||||
|
.as_array()
|
||||||
|
.map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|d| {
|
||||||
|
Some(Self {
|
||||||
|
tool_name: d.get("tool_name")?.as_str()?.to_string(),
|
||||||
|
rationale: d.get("rationale")?.as_str()?.to_string(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(tag = "type")]
|
||||||
|
pub enum AppEvent {
|
||||||
|
#[serde(rename = "response")]
|
||||||
|
Response { content: String, thread_id: String },
|
||||||
|
#[serde(rename = "thinking")]
|
||||||
|
Thinking {
|
||||||
|
message: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "tool_started")]
|
||||||
|
ToolStarted {
|
||||||
|
name: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "tool_completed")]
|
||||||
|
ToolCompleted {
|
||||||
|
name: String,
|
||||||
|
success: bool,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
error: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
parameters: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "tool_result")]
|
||||||
|
ToolResult {
|
||||||
|
name: String,
|
||||||
|
preview: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "stream_chunk")]
|
||||||
|
StreamChunk {
|
||||||
|
content: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "status")]
|
||||||
|
Status {
|
||||||
|
message: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "job_started")]
|
||||||
|
JobStarted {
|
||||||
|
job_id: String,
|
||||||
|
title: String,
|
||||||
|
browse_url: String,
|
||||||
|
},
|
||||||
|
#[serde(rename = "approval_needed")]
|
||||||
|
ApprovalNeeded {
|
||||||
|
request_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
description: String,
|
||||||
|
parameters: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
/// Whether the "always" auto-approve option should be shown.
|
||||||
|
allow_always: bool,
|
||||||
|
},
|
||||||
|
#[serde(rename = "auth_required")]
|
||||||
|
AuthRequired {
|
||||||
|
extension_name: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
instructions: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
auth_url: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
setup_url: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "auth_completed")]
|
||||||
|
AuthCompleted {
|
||||||
|
extension_name: String,
|
||||||
|
success: bool,
|
||||||
|
message: String,
|
||||||
|
},
|
||||||
|
#[serde(rename = "error")]
|
||||||
|
Error {
|
||||||
|
message: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
#[serde(rename = "heartbeat")]
|
||||||
|
Heartbeat,
|
||||||
|
|
||||||
|
// Sandbox job streaming events (worker + Claude Code bridge)
|
||||||
|
#[serde(rename = "job_message")]
|
||||||
|
JobMessage {
|
||||||
|
job_id: String,
|
||||||
|
role: String,
|
||||||
|
content: String,
|
||||||
|
},
|
||||||
|
#[serde(rename = "job_tool_use")]
|
||||||
|
JobToolUse {
|
||||||
|
job_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
input: serde_json::Value,
|
||||||
|
},
|
||||||
|
#[serde(rename = "job_tool_result")]
|
||||||
|
JobToolResult {
|
||||||
|
job_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
output: String,
|
||||||
|
},
|
||||||
|
#[serde(rename = "job_status")]
|
||||||
|
JobStatus { job_id: String, message: String },
|
||||||
|
#[serde(rename = "job_result")]
|
||||||
|
JobResult {
|
||||||
|
job_id: String,
|
||||||
|
status: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
session_id: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
fallback_deliverable: Option<serde_json::Value>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// An image was generated by a tool.
|
||||||
|
#[serde(rename = "image_generated")]
|
||||||
|
ImageGenerated {
|
||||||
|
data_url: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
path: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Suggested follow-up messages for the user.
|
||||||
|
#[serde(rename = "suggestions")]
|
||||||
|
Suggestions {
|
||||||
|
suggestions: Vec<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Per-turn token usage and cost summary.
|
||||||
|
#[serde(rename = "turn_cost")]
|
||||||
|
TurnCost {
|
||||||
|
input_tokens: u64,
|
||||||
|
output_tokens: u64,
|
||||||
|
cost_usd: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Extension activation status change (WASM channels).
|
||||||
|
#[serde(rename = "extension_status")]
|
||||||
|
ExtensionStatus {
|
||||||
|
extension_name: String,
|
||||||
|
status: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
message: Option<String>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Agent reasoning update (why it chose specific tools).
|
||||||
|
#[serde(rename = "reasoning_update")]
|
||||||
|
ReasoningUpdate {
|
||||||
|
narrative: String,
|
||||||
|
decisions: Vec<ToolDecisionDto>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Reasoning update for a sandbox job.
|
||||||
|
#[serde(rename = "job_reasoning")]
|
||||||
|
JobReasoning {
|
||||||
|
job_id: String,
|
||||||
|
narrative: String,
|
||||||
|
decisions: Vec<ToolDecisionDto>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AppEvent {
|
||||||
|
/// The wire-format event type string (matches the `#[serde(rename)]` value).
|
||||||
|
pub fn event_type(&self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Response { .. } => "response",
|
||||||
|
Self::Thinking { .. } => "thinking",
|
||||||
|
Self::ToolStarted { .. } => "tool_started",
|
||||||
|
Self::ToolCompleted { .. } => "tool_completed",
|
||||||
|
Self::ToolResult { .. } => "tool_result",
|
||||||
|
Self::StreamChunk { .. } => "stream_chunk",
|
||||||
|
Self::Status { .. } => "status",
|
||||||
|
Self::JobStarted { .. } => "job_started",
|
||||||
|
Self::ApprovalNeeded { .. } => "approval_needed",
|
||||||
|
Self::AuthRequired { .. } => "auth_required",
|
||||||
|
Self::AuthCompleted { .. } => "auth_completed",
|
||||||
|
Self::Error { .. } => "error",
|
||||||
|
Self::Heartbeat => "heartbeat",
|
||||||
|
Self::JobMessage { .. } => "job_message",
|
||||||
|
Self::JobToolUse { .. } => "job_tool_use",
|
||||||
|
Self::JobToolResult { .. } => "job_tool_result",
|
||||||
|
Self::JobStatus { .. } => "job_status",
|
||||||
|
Self::JobResult { .. } => "job_result",
|
||||||
|
Self::ImageGenerated { .. } => "image_generated",
|
||||||
|
Self::Suggestions { .. } => "suggestions",
|
||||||
|
Self::TurnCost { .. } => "turn_cost",
|
||||||
|
Self::ExtensionStatus { .. } => "extension_status",
|
||||||
|
Self::ReasoningUpdate { .. } => "reasoning_update",
|
||||||
|
Self::JobReasoning { .. } => "job_reasoning",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// Verify that `event_type()` returns the same string as the serde
|
||||||
|
/// `"type"` field for every variant. This catches drift between the
|
||||||
|
/// `#[serde(rename)]` attributes and the manual match arms.
|
||||||
|
#[test]
|
||||||
|
fn event_type_matches_serde_type_field() {
|
||||||
|
let variants: Vec<AppEvent> = vec![
|
||||||
|
AppEvent::Response {
|
||||||
|
content: String::new(),
|
||||||
|
thread_id: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::Thinking {
|
||||||
|
message: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::ToolStarted {
|
||||||
|
name: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::ToolCompleted {
|
||||||
|
name: String::new(),
|
||||||
|
success: true,
|
||||||
|
error: None,
|
||||||
|
parameters: None,
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::ToolResult {
|
||||||
|
name: String::new(),
|
||||||
|
preview: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::StreamChunk {
|
||||||
|
content: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::Status {
|
||||||
|
message: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::JobStarted {
|
||||||
|
job_id: String::new(),
|
||||||
|
title: String::new(),
|
||||||
|
browse_url: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::ApprovalNeeded {
|
||||||
|
request_id: String::new(),
|
||||||
|
tool_name: String::new(),
|
||||||
|
description: String::new(),
|
||||||
|
parameters: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
allow_always: false,
|
||||||
|
},
|
||||||
|
AppEvent::AuthRequired {
|
||||||
|
extension_name: String::new(),
|
||||||
|
instructions: None,
|
||||||
|
auth_url: None,
|
||||||
|
setup_url: None,
|
||||||
|
},
|
||||||
|
AppEvent::AuthCompleted {
|
||||||
|
extension_name: String::new(),
|
||||||
|
success: true,
|
||||||
|
message: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::Error {
|
||||||
|
message: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::Heartbeat,
|
||||||
|
AppEvent::JobMessage {
|
||||||
|
job_id: String::new(),
|
||||||
|
role: String::new(),
|
||||||
|
content: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::JobToolUse {
|
||||||
|
job_id: String::new(),
|
||||||
|
tool_name: String::new(),
|
||||||
|
input: serde_json::Value::Null,
|
||||||
|
},
|
||||||
|
AppEvent::JobToolResult {
|
||||||
|
job_id: String::new(),
|
||||||
|
tool_name: String::new(),
|
||||||
|
output: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::JobStatus {
|
||||||
|
job_id: String::new(),
|
||||||
|
message: String::new(),
|
||||||
|
},
|
||||||
|
AppEvent::JobResult {
|
||||||
|
job_id: String::new(),
|
||||||
|
status: String::new(),
|
||||||
|
session_id: None,
|
||||||
|
fallback_deliverable: None,
|
||||||
|
},
|
||||||
|
AppEvent::ImageGenerated {
|
||||||
|
data_url: String::new(),
|
||||||
|
path: None,
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::Suggestions {
|
||||||
|
suggestions: vec![],
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::TurnCost {
|
||||||
|
input_tokens: 0,
|
||||||
|
output_tokens: 0,
|
||||||
|
cost_usd: String::new(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::ExtensionStatus {
|
||||||
|
extension_name: String::new(),
|
||||||
|
status: String::new(),
|
||||||
|
message: None,
|
||||||
|
},
|
||||||
|
AppEvent::ReasoningUpdate {
|
||||||
|
narrative: String::new(),
|
||||||
|
decisions: vec![],
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
AppEvent::JobReasoning {
|
||||||
|
job_id: String::new(),
|
||||||
|
narrative: String::new(),
|
||||||
|
decisions: vec![],
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
for variant in &variants {
|
||||||
|
let json: serde_json::Value = serde_json::to_value(variant).unwrap();
|
||||||
|
let serde_type = json["type"].as_str().unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
variant.event_type(),
|
||||||
|
serde_type,
|
||||||
|
"event_type() mismatch for variant: {:?}",
|
||||||
|
variant
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn round_trip_deserialize() {
|
||||||
|
let original = AppEvent::Response {
|
||||||
|
content: "hello".to_string(),
|
||||||
|
thread_id: "t1".to_string(),
|
||||||
|
};
|
||||||
|
let json = serde_json::to_string(&original).unwrap();
|
||||||
|
let deserialized: AppEvent = serde_json::from_str(&json).unwrap();
|
||||||
|
assert_eq!(deserialized.event_type(), "response");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
//! Shared types and utilities for the IronClaw workspace.
|
||||||
|
|
||||||
|
mod event;
|
||||||
|
mod util;
|
||||||
|
|
||||||
|
pub use event::{AppEvent, ToolDecisionDto};
|
||||||
|
pub use util::truncate_preview;
|
||||||
@@ -0,0 +1,100 @@
|
|||||||
|
//! Shared utility functions.
|
||||||
|
|
||||||
|
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
|
||||||
|
///
|
||||||
|
/// If the input is wrapped in `<tool_output ...>...</tool_output>` and truncation
|
||||||
|
/// removes the closing tag, the tag is re-appended so downstream XML parsers
|
||||||
|
/// never see an unclosed element.
|
||||||
|
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
|
||||||
|
if s.len() <= max_bytes {
|
||||||
|
return s.to_string();
|
||||||
|
}
|
||||||
|
// Walk backwards from max_bytes to find a valid char boundary
|
||||||
|
let mut end = max_bytes;
|
||||||
|
while end > 0 && !s.is_char_boundary(end) {
|
||||||
|
end -= 1;
|
||||||
|
}
|
||||||
|
let mut result = format!("{}...", &s[..end]);
|
||||||
|
|
||||||
|
// Re-close <tool_output> if truncation cut through the closing tag.
|
||||||
|
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
|
||||||
|
result.push_str("\n</tool_output>");
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_short_string() {
|
||||||
|
assert_eq!(truncate_preview("hello", 10), "hello");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_exact_boundary() {
|
||||||
|
assert_eq!(truncate_preview("hello", 5), "hello");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_truncates_ascii() {
|
||||||
|
assert_eq!(truncate_preview("hello world", 5), "hello...");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_empty_string() {
|
||||||
|
assert_eq!(truncate_preview("", 10), "");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_multibyte_char_boundary() {
|
||||||
|
let s = "a\u{20AC}b";
|
||||||
|
let result = truncate_preview(s, 3);
|
||||||
|
assert_eq!(result, "a...");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_emoji() {
|
||||||
|
let s = "hi\u{1F980}";
|
||||||
|
let result = truncate_preview(s, 4);
|
||||||
|
assert_eq!(result, "hi...");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_cjk() {
|
||||||
|
let s = "\u{4F60}\u{597D}\u{4E16}\u{754C}";
|
||||||
|
let result = truncate_preview(s, 7);
|
||||||
|
assert_eq!(result, "\u{4F60}\u{597D}...");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_zero_max_bytes() {
|
||||||
|
assert_eq!(truncate_preview("hello", 0), "...");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_closes_tool_output_tag() {
|
||||||
|
let s = "<tool_output name=\"search\">\nSome very long content here\n</tool_output>";
|
||||||
|
let result = truncate_preview(s, 60);
|
||||||
|
assert!(result.ends_with("</tool_output>"));
|
||||||
|
assert!(result.contains("..."));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_no_extra_close_when_intact() {
|
||||||
|
let s = "<tool_output name=\"echo\">\nshort\n</tool_output>";
|
||||||
|
let result = truncate_preview(s, 500);
|
||||||
|
assert_eq!(result, s);
|
||||||
|
assert_eq!(result.matches("</tool_output>").count(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_preview_non_xml_unaffected() {
|
||||||
|
let s = "Just a plain long string that gets truncated";
|
||||||
|
let result = truncate_preview(s, 10);
|
||||||
|
assert_eq!(result, "Just a pla...");
|
||||||
|
assert!(!result.contains("</tool_output>"));
|
||||||
|
}
|
||||||
|
}
|
||||||
+90
-11
@@ -16,6 +16,7 @@ use crate::agent::context_monitor::ContextMonitor;
|
|||||||
use crate::agent::heartbeat::spawn_heartbeat;
|
use crate::agent::heartbeat::spawn_heartbeat;
|
||||||
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
|
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
|
||||||
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
|
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
|
||||||
|
use crate::agent::session::ThreadState;
|
||||||
use crate::agent::session_manager::SessionManager;
|
use crate::agent::session_manager::SessionManager;
|
||||||
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
|
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
|
||||||
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler, SchedulerDeps};
|
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler, SchedulerDeps};
|
||||||
@@ -84,6 +85,15 @@ fn resolve_owner_scope_notification_user(
|
|||||||
trimmed_option(explicit_user).or_else(|| trimmed_option(owner_fallback))
|
trimmed_option(explicit_user).or_else(|| trimmed_option(owner_fallback))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_single_message_repl(message: &IncomingMessage) -> bool {
|
||||||
|
message.channel == "repl"
|
||||||
|
&& message
|
||||||
|
.metadata
|
||||||
|
.get("single_message_mode")
|
||||||
|
.and_then(|value| value.as_bool())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
async fn resolve_channel_notification_user(
|
async fn resolve_channel_notification_user(
|
||||||
extension_manager: Option<&Arc<ExtensionManager>>,
|
extension_manager: Option<&Arc<ExtensionManager>>,
|
||||||
channel: Option<&str>,
|
channel: Option<&str>,
|
||||||
@@ -157,8 +167,8 @@ pub struct AgentDeps {
|
|||||||
pub hooks: Arc<HookRegistry>,
|
pub hooks: Arc<HookRegistry>,
|
||||||
/// Cost enforcement guardrails (daily budget, hourly rate limits).
|
/// Cost enforcement guardrails (daily budget, hourly rate limits).
|
||||||
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
|
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
|
||||||
/// SSE broadcast sender for live job event streaming to the web gateway.
|
/// SSE manager for live job event streaming to the web gateway.
|
||||||
pub sse_tx: Option<tokio::sync::broadcast::Sender<crate::channels::web::types::SseEvent>>,
|
pub sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||||
/// HTTP interceptor for trace recording/replay.
|
/// HTTP interceptor for trace recording/replay.
|
||||||
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
|
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
|
||||||
/// Audio transcription middleware for voice messages.
|
/// Audio transcription middleware for voice messages.
|
||||||
@@ -169,6 +179,9 @@ pub struct AgentDeps {
|
|||||||
pub sandbox_readiness: crate::agent::routine_engine::SandboxReadiness,
|
pub sandbox_readiness: crate::agent::routine_engine::SandboxReadiness,
|
||||||
/// Software builder for self-repair tool rebuilding.
|
/// Software builder for self-repair tool rebuilding.
|
||||||
pub builder: Option<Arc<dyn crate::tools::SoftwareBuilder>>,
|
pub builder: Option<Arc<dyn crate::tools::SoftwareBuilder>>,
|
||||||
|
/// Resolved LLM backend identifier (e.g., "nearai", "openai", "groq").
|
||||||
|
/// Used by `/model` persistence to determine which env var to update.
|
||||||
|
pub llm_backend: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The main agent that coordinates all components.
|
/// The main agent that coordinates all components.
|
||||||
@@ -235,8 +248,8 @@ impl Agent {
|
|||||||
hooks: deps.hooks.clone(),
|
hooks: deps.hooks.clone(),
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
if let Some(ref tx) = deps.sse_tx {
|
if let Some(ref sse) = deps.sse_tx {
|
||||||
scheduler.set_sse_sender(tx.clone());
|
scheduler.set_sse_sender(Arc::clone(sse));
|
||||||
}
|
}
|
||||||
if let Some(ref interceptor) = deps.http_interceptor {
|
if let Some(ref interceptor) = deps.http_interceptor {
|
||||||
scheduler.set_http_interceptor(Arc::clone(interceptor));
|
scheduler.set_http_interceptor(Arc::clone(interceptor));
|
||||||
@@ -1052,10 +1065,11 @@ impl Agent {
|
|||||||
} else {
|
} else {
|
||||||
drop(sess);
|
drop(sess);
|
||||||
self.session_manager
|
self.session_manager
|
||||||
.resolve_thread(
|
.resolve_thread_with_parsed_uuid(
|
||||||
&message.user_id,
|
&message.user_id,
|
||||||
&message.channel,
|
&message.channel,
|
||||||
message.conversation_scope(),
|
message.conversation_scope(),
|
||||||
|
approval_thread_uuid,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -1136,9 +1150,14 @@ impl Agent {
|
|||||||
&& let Submission::UserInput { ref content } = submission
|
&& let Submission::UserInput { ref content } = submission
|
||||||
&& let Some(engine) = self.routine_engine().await
|
&& let Some(engine) = self.routine_engine().await
|
||||||
{
|
{
|
||||||
let fired = engine
|
let single_message_repl = is_single_message_repl(message);
|
||||||
.check_event_triggers(&message.user_id, &message.channel, content)
|
// Use post-hook content so that BeforeInbound hooks that rewrite
|
||||||
.await;
|
// input are respected by event trigger matching.
|
||||||
|
let fired = if single_message_repl {
|
||||||
|
engine.check_event_triggers_and_wait(message, content).await
|
||||||
|
} else {
|
||||||
|
engine.check_event_triggers(message, content).await
|
||||||
|
};
|
||||||
if fired > 0 {
|
if fired > 0 {
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
channel = %message.channel,
|
channel = %message.channel,
|
||||||
@@ -1146,10 +1165,16 @@ impl Agent {
|
|||||||
fired,
|
fired,
|
||||||
"Consumed inbound user message with matching event-triggered routine(s)"
|
"Consumed inbound user message with matching event-triggered routine(s)"
|
||||||
);
|
);
|
||||||
return Ok(Some(String::new()));
|
return if single_message_repl {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
Ok(Some(String::new()))
|
||||||
|
};
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let session_for_empty_exit = Arc::clone(&session);
|
||||||
|
|
||||||
// Process based on submission type
|
// Process based on submission type
|
||||||
let result = match submission {
|
let result = match submission {
|
||||||
Submission::UserInput { content } => {
|
Submission::UserInput { content } => {
|
||||||
@@ -1246,6 +1271,28 @@ impl Agent {
|
|||||||
command,
|
command,
|
||||||
message.channel
|
message.channel
|
||||||
);
|
);
|
||||||
|
// /reasoning is special-cased here (not in handle_system_command)
|
||||||
|
// because it needs the session + thread_id to read turn reasoning
|
||||||
|
// data, which handle_system_command's signature doesn't provide.
|
||||||
|
if command == "reasoning" {
|
||||||
|
let result = self
|
||||||
|
.handle_reasoning_command(&args, &session, thread_id)
|
||||||
|
.await;
|
||||||
|
return match result {
|
||||||
|
SubmissionResult::Response { content } => Ok(Some(content)),
|
||||||
|
SubmissionResult::Ok { message } => Ok(message),
|
||||||
|
SubmissionResult::Error { message } => {
|
||||||
|
Ok(Some(format!("Error: {}", message)))
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
if is_single_message_repl(message) {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
Ok(Some(String::new()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
// Authorization checks (including restart channel check) are enforced in handle_system_command
|
// Authorization checks (including restart channel check) are enforced in handle_system_command
|
||||||
self.handle_system_command(&command, &args, &message.channel)
|
self.handle_system_command(&command, &args, &message.channel)
|
||||||
.await
|
.await
|
||||||
@@ -1305,7 +1352,26 @@ impl Agent {
|
|||||||
Ok(Some(content))
|
Ok(Some(content))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
SubmissionResult::Ok { message } => Ok(message),
|
SubmissionResult::Ok {
|
||||||
|
message: output_message,
|
||||||
|
} => {
|
||||||
|
let should_exit =
|
||||||
|
if output_message.as_deref() == Some("") && is_single_message_repl(message) {
|
||||||
|
let sess = session_for_empty_exit.lock().await;
|
||||||
|
sess.threads
|
||||||
|
.get(&thread_id)
|
||||||
|
.map(|thread| thread.state != ThreadState::AwaitingApproval)
|
||||||
|
.unwrap_or(true)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
};
|
||||||
|
|
||||||
|
if should_exit {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
Ok(output_message)
|
||||||
|
}
|
||||||
|
}
|
||||||
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
||||||
SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())),
|
SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())),
|
||||||
SubmissionResult::NeedApproval { .. } => {
|
SubmissionResult::NeedApproval { .. } => {
|
||||||
@@ -1321,7 +1387,7 @@ impl Agent {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
chat_tool_execution_metadata, resolve_routine_notification_user,
|
chat_tool_execution_metadata, is_single_message_repl, resolve_routine_notification_user,
|
||||||
should_fallback_routine_notification, truncate_for_preview,
|
should_fallback_routine_notification, truncate_for_preview,
|
||||||
};
|
};
|
||||||
use crate::channels::IncomingMessage;
|
use crate::channels::IncomingMessage;
|
||||||
@@ -1483,4 +1549,17 @@ mod tests {
|
|||||||
|
|
||||||
assert!(should_fallback_routine_notification(&error)); // safety: test-only assertion
|
assert!(should_fallback_routine_notification(&error)); // safety: test-only assertion
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn single_message_repl_detection_requires_repl_channel_and_metadata_flag() {
|
||||||
|
let repl = IncomingMessage::new("repl", "owner-scope", "hello")
|
||||||
|
.with_metadata(serde_json::json!({ "single_message_mode": true }));
|
||||||
|
let gateway = IncomingMessage::new("gateway", "owner-scope", "hello")
|
||||||
|
.with_metadata(serde_json::json!({ "single_message_mode": true }));
|
||||||
|
let plain_repl = IncomingMessage::new("repl", "owner-scope", "hello");
|
||||||
|
|
||||||
|
assert!(is_single_message_repl(&repl)); // safety: test-only assertion
|
||||||
|
assert!(!is_single_message_repl(&gateway)); // safety: test-only assertion
|
||||||
|
assert!(!is_single_message_repl(&plain_repl)); // safety: test-only assertion
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -414,6 +414,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let delegate = MockDelegate::new(vec![
|
let delegate = MockDelegate::new(vec![
|
||||||
tool_calls_output(vec![tool_call]),
|
tool_calls_output(vec![tool_call]),
|
||||||
|
|||||||
+138
-3
@@ -465,6 +465,94 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Handle `/reasoning [N|all]` — show reasoning history for the active thread.
|
||||||
|
pub(super) async fn handle_reasoning_command(
|
||||||
|
&self,
|
||||||
|
args: &[String],
|
||||||
|
session: &Arc<Mutex<Session>>,
|
||||||
|
thread_id: Uuid,
|
||||||
|
) -> SubmissionResult {
|
||||||
|
// Clone the turn data we need, then drop the session lock.
|
||||||
|
let turns_snapshot: Vec<(
|
||||||
|
usize,
|
||||||
|
Option<String>,
|
||||||
|
Vec<crate::agent::session::TurnToolCall>,
|
||||||
|
)>;
|
||||||
|
{
|
||||||
|
let sess = session.lock().await;
|
||||||
|
let thread = match sess.threads.get(&thread_id) {
|
||||||
|
Some(t) => t,
|
||||||
|
None => return SubmissionResult::error("No active thread."),
|
||||||
|
};
|
||||||
|
|
||||||
|
if thread.turns.is_empty() {
|
||||||
|
return SubmissionResult::ok_with_message("No turns yet.");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse argument: default=last turn, "all"=all turns, N=specific turn (1-based).
|
||||||
|
let selected: Vec<&crate::agent::session::Turn> = match args.first().map(|s| s.as_str())
|
||||||
|
{
|
||||||
|
Some("all") => thread.turns.iter().collect(),
|
||||||
|
Some(n) => match n.parse::<usize>() {
|
||||||
|
Ok(0) => return SubmissionResult::error("Turn numbers start at 1."),
|
||||||
|
Ok(num) if num > thread.turns.len() => {
|
||||||
|
return SubmissionResult::error(format!(
|
||||||
|
"Turn {} does not exist (max: {}).",
|
||||||
|
num,
|
||||||
|
thread.turns.len()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(num) => vec![&thread.turns[num - 1]],
|
||||||
|
Err(_) => return SubmissionResult::error("Usage: /reasoning [N|all]"),
|
||||||
|
},
|
||||||
|
None => {
|
||||||
|
// Default: last turn that has tool calls
|
||||||
|
match thread.turns.iter().rev().find(|t| !t.tool_calls.is_empty()) {
|
||||||
|
Some(t) => vec![t],
|
||||||
|
None => {
|
||||||
|
return SubmissionResult::ok_with_message("No turns with tool calls.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
turns_snapshot = selected
|
||||||
|
.into_iter()
|
||||||
|
.map(|t| (t.turn_number, t.narrative.clone(), t.tool_calls.clone()))
|
||||||
|
.collect();
|
||||||
|
}
|
||||||
|
// Session lock is now dropped — format output without holding it.
|
||||||
|
|
||||||
|
let mut output = String::new();
|
||||||
|
for (turn_number, narrative, tool_calls) in &turns_snapshot {
|
||||||
|
output.push_str(&format!("--- Turn {} ---\n", turn_number + 1));
|
||||||
|
if let Some(narrative) = narrative {
|
||||||
|
output.push_str(&format!("Reasoning: {}\n", narrative));
|
||||||
|
}
|
||||||
|
if tool_calls.is_empty() {
|
||||||
|
output.push_str(" (no tool calls)\n");
|
||||||
|
} else {
|
||||||
|
for tc in tool_calls {
|
||||||
|
let status = if tc.error.is_some() {
|
||||||
|
"error"
|
||||||
|
} else if tc.result.is_some() {
|
||||||
|
"ok"
|
||||||
|
} else {
|
||||||
|
"pending"
|
||||||
|
};
|
||||||
|
output.push_str(&format!(" {} [{}]", tc.name, status));
|
||||||
|
if let Some(ref rationale) = tc.rationale {
|
||||||
|
output.push_str(&format!(" — {}", rationale));
|
||||||
|
}
|
||||||
|
output.push('\n');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
output.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
SubmissionResult::response(output.trim_end())
|
||||||
|
}
|
||||||
|
|
||||||
/// Handle system commands that bypass thread-state checks entirely.
|
/// Handle system commands that bypass thread-state checks entirely.
|
||||||
pub(super) async fn handle_system_command(
|
pub(super) async fn handle_system_command(
|
||||||
&self,
|
&self,
|
||||||
@@ -480,6 +568,7 @@ impl Agent {
|
|||||||
" /version Show version info\n",
|
" /version Show version info\n",
|
||||||
" /tools List available tools\n",
|
" /tools List available tools\n",
|
||||||
" /debug Toggle debug mode\n",
|
" /debug Toggle debug mode\n",
|
||||||
|
" /reasoning [N|all] Show agent reasoning for turns\n",
|
||||||
" /ping Connectivity check\n",
|
" /ping Connectivity check\n",
|
||||||
"\n",
|
"\n",
|
||||||
"Jobs:\n",
|
"Jobs:\n",
|
||||||
@@ -841,12 +930,50 @@ impl Agent {
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
tracing::warn!("Failed to persist model to DB: {}", e);
|
tracing::warn!("Failed to persist model to DB: {}", e);
|
||||||
|
} else {
|
||||||
|
tracing::debug!("Persisted selected_model to DB: {}", model);
|
||||||
}
|
}
|
||||||
|
} else {
|
||||||
|
tracing::warn!("No database store available — model choice will not persist to DB");
|
||||||
}
|
}
|
||||||
|
|
||||||
// 2. Update TOML config file if it exists (sync I/O in spawn_blocking).
|
// 2. Update .env and TOML config file (sync I/O in spawn_blocking).
|
||||||
let model_owned = model.to_string();
|
let model_owned = model.to_string();
|
||||||
|
let backend = self.deps.llm_backend.clone();
|
||||||
if let Err(e) = tokio::task::spawn_blocking(move || {
|
if let Err(e) = tokio::task::spawn_blocking(move || {
|
||||||
|
// 2a. Update the backend-specific model env var in ~/.ironclaw/.env.
|
||||||
|
//
|
||||||
|
// Env vars have the HIGHEST priority in LlmConfig::resolve_model()
|
||||||
|
// (env var > TOML > DB > default). If the .env file has e.g.
|
||||||
|
// NEARAI_MODEL=old-model, it shadows everything else. We must
|
||||||
|
// update this var or the /model change is invisible on restart.
|
||||||
|
let registry = crate::llm::ProviderRegistry::load();
|
||||||
|
let model_env = registry.model_env_var(&backend);
|
||||||
|
let env_var_prefix = format!("{}=", model_env);
|
||||||
|
|
||||||
|
// Only update the .env file if the var is actually set there
|
||||||
|
// (avoid injecting new vars the user never configured).
|
||||||
|
let env_path = crate::bootstrap::ironclaw_env_path();
|
||||||
|
let env_has_var = std::fs::read_to_string(&env_path)
|
||||||
|
.ok()
|
||||||
|
.is_some_and(|content| {
|
||||||
|
content.lines().any(|line| {
|
||||||
|
let trimmed = line.trim_start();
|
||||||
|
!trimmed.starts_with('#') && trimmed.starts_with(&env_var_prefix)
|
||||||
|
})
|
||||||
|
});
|
||||||
|
if env_has_var {
|
||||||
|
if let Err(e) = crate::bootstrap::upsert_bootstrap_var(model_env, &model_owned) {
|
||||||
|
tracing::warn!("Failed to update {} in .env: {}", model_env, e);
|
||||||
|
} else {
|
||||||
|
tracing::debug!("Updated {} in .env to {}", model_env, model_owned);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2b. Update (or create) the TOML config file.
|
||||||
|
//
|
||||||
|
// The TOML overlay has higher priority than DB settings on
|
||||||
|
// startup, so it MUST stay in sync with the DB.
|
||||||
let toml_path = crate::settings::Settings::default_toml_path();
|
let toml_path = crate::settings::Settings::default_toml_path();
|
||||||
match crate::settings::Settings::load_toml(&toml_path) {
|
match crate::settings::Settings::load_toml(&toml_path) {
|
||||||
Ok(Some(mut settings)) => {
|
Ok(Some(mut settings)) => {
|
||||||
@@ -856,7 +983,15 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(None) => {
|
Ok(None) => {
|
||||||
// No config file on disk; nothing to update.
|
// No config file yet — create one so the model choice
|
||||||
|
// survives restarts even when the DB is unavailable.
|
||||||
|
let settings = crate::settings::Settings {
|
||||||
|
selected_model: Some(model_owned),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
if let Err(e) = settings.save_toml(&toml_path) {
|
||||||
|
tracing::warn!("Failed to create config.toml for model persistence: {}", e);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!("Failed to load config.toml for model persistence: {}", e);
|
tracing::warn!("Failed to load config.toml for model persistence: {}", e);
|
||||||
@@ -865,7 +1000,7 @@ impl Agent {
|
|||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
tracing::warn!("Model TOML persistence task failed: {}", e);
|
tracing::warn!("Model persistence task failed: {}", e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+104
-9
@@ -420,6 +420,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
content: Option<String>,
|
content: Option<String>,
|
||||||
reason_ctx: &mut ReasoningContext,
|
reason_ctx: &mut ReasoningContext,
|
||||||
) -> Result<Option<LoopOutcome>, Error> {
|
) -> Result<Option<LoopOutcome>, Error> {
|
||||||
|
// Extract and sanitize the narrative before consuming `content`.
|
||||||
|
let narrative = content
|
||||||
|
.as_deref()
|
||||||
|
.filter(|c| !c.trim().is_empty())
|
||||||
|
.map(|c| {
|
||||||
|
let sanitized = self
|
||||||
|
.agent
|
||||||
|
.safety()
|
||||||
|
.sanitize_tool_output("agent_narrative", c);
|
||||||
|
sanitized.content
|
||||||
|
})
|
||||||
|
.filter(|c| !c.trim().is_empty());
|
||||||
|
|
||||||
// Add the assistant message with tool_calls to context.
|
// Add the assistant message with tool_calls to context.
|
||||||
// OpenAI protocol requires this before tool-result messages.
|
// OpenAI protocol requires this before tool-result messages.
|
||||||
reason_ctx
|
reason_ctx
|
||||||
@@ -440,6 +453,41 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
// Build per-tool decisions for the reasoning update.
|
||||||
|
// Sanitize each rationale through SafetyLayer (parity with JobDelegate).
|
||||||
|
let decisions: Vec<crate::channels::ToolDecision> = tool_calls
|
||||||
|
.iter()
|
||||||
|
.filter_map(|tc| {
|
||||||
|
tc.reasoning.as_ref().map(|r| {
|
||||||
|
let sanitized = self
|
||||||
|
.agent
|
||||||
|
.safety()
|
||||||
|
.sanitize_tool_output("tool_rationale", r)
|
||||||
|
.content;
|
||||||
|
crate::channels::ToolDecision {
|
||||||
|
tool_name: tc.name.clone(),
|
||||||
|
rationale: sanitized,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Emit reasoning update to channels.
|
||||||
|
if narrative.is_some() || !decisions.is_empty() {
|
||||||
|
let _ = self
|
||||||
|
.agent
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&self.message.channel,
|
||||||
|
StatusUpdate::ReasoningUpdate {
|
||||||
|
narrative: narrative.clone().unwrap_or_default(),
|
||||||
|
decisions: decisions.clone(),
|
||||||
|
},
|
||||||
|
&self.message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
// Record tool calls in the thread with sensitive params redacted.
|
// Record tool calls in the thread with sensitive params redacted.
|
||||||
{
|
{
|
||||||
let mut redacted_args: Vec<serde_json::Value> = Vec::with_capacity(tool_calls.len());
|
let mut redacted_args: Vec<serde_json::Value> = Vec::with_capacity(tool_calls.len());
|
||||||
@@ -455,8 +503,23 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
||||||
&& let Some(turn) = thread.last_turn_mut()
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
{
|
{
|
||||||
|
// Set turn-level narrative.
|
||||||
|
if turn.narrative.is_none() {
|
||||||
|
turn.narrative = narrative;
|
||||||
|
}
|
||||||
for (tc, safe_args) in tool_calls.iter().zip(redacted_args) {
|
for (tc, safe_args) in tool_calls.iter().zip(redacted_args) {
|
||||||
turn.record_tool_call(&tc.name, safe_args);
|
let sanitized_rationale = tc.reasoning.as_ref().map(|r| {
|
||||||
|
self.agent
|
||||||
|
.safety()
|
||||||
|
.sanitize_tool_output("tool_rationale", r)
|
||||||
|
.content
|
||||||
|
});
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
&tc.name,
|
||||||
|
safe_args,
|
||||||
|
sanitized_rationale,
|
||||||
|
Some(tc.id.clone()),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -726,7 +789,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
||||||
&& let Some(turn) = thread.last_turn_mut()
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
{
|
{
|
||||||
turn.record_tool_error(error_msg.clone());
|
turn.record_tool_error_for(&tc.id, error_msg.clone());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
reason_ctx
|
reason_ctx
|
||||||
@@ -852,16 +915,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
|
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Record sanitized result in thread
|
// Record sanitized result in thread (identity-based matching).
|
||||||
{
|
{
|
||||||
let mut sess = self.session.lock().await;
|
let mut sess = self.session.lock().await;
|
||||||
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
|
||||||
&& let Some(turn) = thread.last_turn_mut()
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
{
|
{
|
||||||
if is_tool_error {
|
if is_tool_error {
|
||||||
turn.record_tool_error(result_content.clone());
|
turn.record_tool_error_for(&tc.id, result_content.clone());
|
||||||
} else {
|
} else {
|
||||||
turn.record_tool_result(serde_json::json!(result_content));
|
turn.record_tool_result_for(
|
||||||
|
&tc.id,
|
||||||
|
serde_json::json!(result_content),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1098,15 +1164,23 @@ pub(crate) fn extract_suggestions(text: &str) -> (String, Vec<String>) {
|
|||||||
Regex::new(r"(?s)<suggestions>\s*(.*?)\s*</suggestions>").expect("valid regex") // safety: constant pattern
|
Regex::new(r"(?s)<suggestions>\s*(.*?)\s*</suggestions>").expect("valid regex") // safety: constant pattern
|
||||||
});
|
});
|
||||||
|
|
||||||
// Find the position of the last closing code fence to avoid matching inside code blocks
|
// Build a sorted list of code fence positions to determine open/close pairing.
|
||||||
let last_code_fence = text.rfind("```").unwrap_or(0);
|
// A position is "inside" a fenced block when it falls between an odd-numbered
|
||||||
|
// fence (opening) and the next even-numbered fence (closing).
|
||||||
|
let fence_positions: Vec<usize> = text.match_indices("```").map(|(pos, _)| pos).collect();
|
||||||
|
|
||||||
// Find all matches, take the last one that's after the last code fence
|
let is_inside_fence = |pos: usize| -> bool {
|
||||||
|
// Count how many fences appear before `pos`. If odd, we're inside a fence.
|
||||||
|
let count = fence_positions.iter().take_while(|&&fp| fp <= pos).count();
|
||||||
|
count % 2 == 1
|
||||||
|
};
|
||||||
|
|
||||||
|
// Find all matches, take the last one that's outside any code fence
|
||||||
let mut best_match: Option<regex::Match<'_>> = None;
|
let mut best_match: Option<regex::Match<'_>> = None;
|
||||||
let mut best_capture: Option<String> = None;
|
let mut best_capture: Option<String> = None;
|
||||||
for caps in RE.captures_iter(text) {
|
for caps in RE.captures_iter(text) {
|
||||||
if let (Some(full), Some(inner)) = (caps.get(0), caps.get(1))
|
if let (Some(full), Some(inner)) = (caps.get(0), caps.get(1))
|
||||||
&& full.start() >= last_code_fence
|
&& !is_inside_fence(full.start())
|
||||||
{
|
{
|
||||||
best_match = Some(full);
|
best_match = Some(full);
|
||||||
best_capture = Some(inner.as_str().to_string());
|
best_capture = Some(inner.as_str().to_string());
|
||||||
@@ -1225,6 +1299,7 @@ mod tests {
|
|||||||
document_extraction: None,
|
document_extraction: None,
|
||||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||||
builder: None,
|
builder: None,
|
||||||
|
llm_backend: "nearai".to_string(),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -1453,11 +1528,13 @@ mod tests {
|
|||||||
id: "call_2".to_string(),
|
id: "call_2".to_string(),
|
||||||
name: "http".to_string(),
|
name: "http".to_string(),
|
||||||
arguments: serde_json::json!({"url": "https://example.com"}),
|
arguments: serde_json::json!({"url": "https://example.com"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
ToolCall {
|
ToolCall {
|
||||||
id: "call_3".to_string(),
|
id: "call_3".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({"message": "done"}),
|
arguments: serde_json::json!({"message": "done"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
],
|
],
|
||||||
user_timezone: None,
|
user_timezone: None,
|
||||||
@@ -1643,6 +1720,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({"message": "hi"}),
|
arguments: serde_json::json!({"message": "hi"}),
|
||||||
|
reasoning: None,
|
||||||
}],
|
}],
|
||||||
),
|
),
|
||||||
ChatMessage::tool_result("call_1", "echo", "hi"),
|
ChatMessage::tool_result("call_1", "echo", "hi"),
|
||||||
@@ -1735,11 +1813,13 @@ mod tests {
|
|||||||
id: "c1".to_string(),
|
id: "c1".to_string(),
|
||||||
name: "http".to_string(),
|
name: "http".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
ToolCall {
|
ToolCall {
|
||||||
id: "c2".to_string(),
|
id: "c2".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
],
|
],
|
||||||
),
|
),
|
||||||
@@ -1773,6 +1853,7 @@ mod tests {
|
|||||||
id: "c1".to_string(),
|
id: "c1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
}],
|
}],
|
||||||
),
|
),
|
||||||
ChatMessage::tool_result("c1", "echo", "done"),
|
ChatMessage::tool_result("c1", "echo", "done"),
|
||||||
@@ -1903,6 +1984,7 @@ mod tests {
|
|||||||
id: crate::llm::generate_tool_call_id(0, 0),
|
id: crate::llm::generate_tool_call_id(0, 0),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({"message": "looping"}),
|
arguments: serde_json::json!({"message": "looping"}),
|
||||||
|
reasoning: None,
|
||||||
}],
|
}],
|
||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
@@ -2056,6 +2138,7 @@ mod tests {
|
|||||||
id: crate::llm::generate_tool_call_id(0, 0),
|
id: crate::llm::generate_tool_call_id(0, 0),
|
||||||
name: "nonexistent_tool".to_string(),
|
name: "nonexistent_tool".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
}],
|
}],
|
||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
@@ -2092,6 +2175,7 @@ mod tests {
|
|||||||
document_extraction: None,
|
document_extraction: None,
|
||||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||||
builder: None,
|
builder: None,
|
||||||
|
llm_backend: "nearai".to_string(),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -2212,6 +2296,7 @@ mod tests {
|
|||||||
document_extraction: None,
|
document_extraction: None,
|
||||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||||
builder: None,
|
builder: None,
|
||||||
|
llm_backend: "nearai".to_string(),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -2345,6 +2430,16 @@ mod tests {
|
|||||||
assert!(suggestions.is_empty()); // safety: test
|
assert!(suggestions.is_empty()); // safety: test
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_suggestions_inside_unclosed_code_fence() {
|
||||||
|
// Regression: odd number of fences (unclosed fence) must still be
|
||||||
|
// treated as "inside a code block".
|
||||||
|
let input = "```\ncode\n<suggestions>[\"bar\"]</suggestions>";
|
||||||
|
let (text, suggestions) = super::extract_suggestions(input);
|
||||||
|
assert_eq!(text, input); // safety: test
|
||||||
|
assert!(suggestions.is_empty()); // safety: test
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_extract_suggestions_after_code_fence() {
|
fn test_extract_suggestions_after_code_fence() {
|
||||||
let input = "```\ncode\n```\nAnswer.\n<suggestions>[\"foo\"]</suggestions>";
|
let input = "```\ncode\n```\nAnswer.\n<suggestions>[\"foo\"]</suggestions>";
|
||||||
|
|||||||
+35
-25
@@ -21,8 +21,8 @@ use tokio::task::JoinHandle;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::IncomingMessage;
|
use crate::channels::IncomingMessage;
|
||||||
use crate::channels::web::types::SseEvent;
|
|
||||||
use crate::context::{ContextManager, JobState};
|
use crate::context::{ContextManager, JobState};
|
||||||
|
use ironclaw_common::AppEvent;
|
||||||
|
|
||||||
/// Route context for forwarding job monitor events back to the user's channel.
|
/// Route context for forwarding job monitor events back to the user's channel.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -36,15 +36,15 @@ pub struct JobMonitorRoute {
|
|||||||
/// injects assistant messages into the agent loop.
|
/// injects assistant messages into the agent loop.
|
||||||
///
|
///
|
||||||
/// The monitor forwards:
|
/// The monitor forwards:
|
||||||
/// - `SseEvent::JobMessage` (assistant role): injected as incoming messages so
|
/// - `AppEvent::JobMessage` (assistant role): injected as incoming messages so
|
||||||
/// the main agent can read and relay to the user.
|
/// the main agent can read and relay to the user.
|
||||||
/// - `SseEvent::JobResult`: injected as a completion notice, then the task exits.
|
/// - `AppEvent::JobResult`: injected as a completion notice, then the task exits.
|
||||||
///
|
///
|
||||||
/// Tool use/result and status events are intentionally skipped (too noisy for
|
/// Tool use/result and status events are intentionally skipped (too noisy for
|
||||||
/// the main agent's context window).
|
/// the main agent's context window).
|
||||||
pub fn spawn_job_monitor(
|
pub fn spawn_job_monitor(
|
||||||
job_id: Uuid,
|
job_id: Uuid,
|
||||||
event_rx: broadcast::Receiver<(Uuid, SseEvent)>,
|
event_rx: broadcast::Receiver<(Uuid, String, AppEvent)>,
|
||||||
inject_tx: mpsc::Sender<IncomingMessage>,
|
inject_tx: mpsc::Sender<IncomingMessage>,
|
||||||
route: JobMonitorRoute,
|
route: JobMonitorRoute,
|
||||||
) -> JoinHandle<()> {
|
) -> JoinHandle<()> {
|
||||||
@@ -56,7 +56,7 @@ pub fn spawn_job_monitor(
|
|||||||
/// jobs don't stay `InProgress` forever in the `ContextManager`.
|
/// jobs don't stay `InProgress` forever in the `ContextManager`.
|
||||||
pub fn spawn_job_monitor_with_context(
|
pub fn spawn_job_monitor_with_context(
|
||||||
job_id: Uuid,
|
job_id: Uuid,
|
||||||
mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>,
|
mut event_rx: broadcast::Receiver<(Uuid, String, AppEvent)>,
|
||||||
inject_tx: mpsc::Sender<IncomingMessage>,
|
inject_tx: mpsc::Sender<IncomingMessage>,
|
||||||
route: JobMonitorRoute,
|
route: JobMonitorRoute,
|
||||||
context_manager: Option<Arc<ContextManager>>,
|
context_manager: Option<Arc<ContextManager>>,
|
||||||
@@ -68,13 +68,13 @@ pub fn spawn_job_monitor_with_context(
|
|||||||
|
|
||||||
loop {
|
loop {
|
||||||
match event_rx.recv().await {
|
match event_rx.recv().await {
|
||||||
Ok((ev_job_id, event)) => {
|
Ok((ev_job_id, _user_id, event)) => {
|
||||||
if ev_job_id != job_id {
|
if ev_job_id != job_id {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
match event {
|
match event {
|
||||||
SseEvent::JobMessage { role, content, .. } if role == "assistant" => {
|
AppEvent::JobMessage { role, content, .. } if role == "assistant" => {
|
||||||
let mut msg = IncomingMessage::new(
|
let mut msg = IncomingMessage::new(
|
||||||
route.channel.clone(),
|
route.channel.clone(),
|
||||||
route.user_id.clone(),
|
route.user_id.clone(),
|
||||||
@@ -92,7 +92,7 @@ pub fn spawn_job_monitor_with_context(
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
SseEvent::JobResult { status, .. } => {
|
AppEvent::JobResult { status, .. } => {
|
||||||
// Transition in-memory state so the job frees its
|
// Transition in-memory state so the job frees its
|
||||||
// max_jobs slot and query tools show the final state.
|
// max_jobs slot and query tools show the final state.
|
||||||
if let Some(ref cm) = context_manager {
|
if let Some(ref cm) = context_manager {
|
||||||
@@ -162,7 +162,7 @@ pub fn spawn_job_monitor_with_context(
|
|||||||
/// inject messages into) but we still need to free the `max_jobs` slot.
|
/// inject messages into) but we still need to free the `max_jobs` slot.
|
||||||
pub fn spawn_completion_watcher(
|
pub fn spawn_completion_watcher(
|
||||||
job_id: Uuid,
|
job_id: Uuid,
|
||||||
mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>,
|
mut event_rx: broadcast::Receiver<(Uuid, String, AppEvent)>,
|
||||||
context_manager: Arc<ContextManager>,
|
context_manager: Arc<ContextManager>,
|
||||||
) -> JoinHandle<()> {
|
) -> JoinHandle<()> {
|
||||||
let short_id = job_id.to_string()[..8].to_string();
|
let short_id = job_id.to_string()[..8].to_string();
|
||||||
@@ -170,7 +170,9 @@ pub fn spawn_completion_watcher(
|
|||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
loop {
|
loop {
|
||||||
match event_rx.recv().await {
|
match event_rx.recv().await {
|
||||||
Ok((ev_job_id, SseEvent::JobResult { status, .. })) if ev_job_id == job_id => {
|
Ok((ev_job_id, _user_id, AppEvent::JobResult { status, .. }))
|
||||||
|
if ev_job_id == job_id =>
|
||||||
|
{
|
||||||
let target = if status == "completed" {
|
let target = if status == "completed" {
|
||||||
JobState::Completed
|
JobState::Completed
|
||||||
} else {
|
} else {
|
||||||
@@ -227,7 +229,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_monitor_forwards_assistant_messages() {
|
async fn test_monitor_forwards_assistant_messages() {
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -237,7 +239,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobMessage {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobMessage {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
role: "assistant".to_string(),
|
role: "assistant".to_string(),
|
||||||
content: "I found a bug".to_string(),
|
content: "I found a bug".to_string(),
|
||||||
@@ -259,7 +262,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_monitor_ignores_other_jobs() {
|
async fn test_monitor_ignores_other_jobs() {
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -270,7 +273,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
other_job_id,
|
other_job_id,
|
||||||
SseEvent::JobMessage {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobMessage {
|
||||||
job_id: other_job_id.to_string(),
|
job_id: other_job_id.to_string(),
|
||||||
role: "assistant".to_string(),
|
role: "assistant".to_string(),
|
||||||
content: "wrong job".to_string(),
|
content: "wrong job".to_string(),
|
||||||
@@ -289,7 +293,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_monitor_exits_on_job_result() {
|
async fn test_monitor_exits_on_job_result() {
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -299,7 +303,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobResult {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobResult {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
status: "completed".to_string(),
|
status: "completed".to_string(),
|
||||||
session_id: None,
|
session_id: None,
|
||||||
@@ -324,7 +329,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_monitor_skips_tool_events() {
|
async fn test_monitor_skips_tool_events() {
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -334,7 +339,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobToolUse {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobToolUse {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
tool_name: "shell".to_string(),
|
tool_name: "shell".to_string(),
|
||||||
input: serde_json::json!({"command": "ls"}),
|
input: serde_json::json!({"command": "ls"}),
|
||||||
@@ -346,7 +352,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobMessage {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobMessage {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
role: "user".to_string(),
|
role: "user".to_string(),
|
||||||
content: "user prompt".to_string(),
|
content: "user prompt".to_string(),
|
||||||
@@ -402,7 +409,7 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let handle = spawn_job_monitor_with_context(
|
let handle = spawn_job_monitor_with_context(
|
||||||
@@ -417,7 +424,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobResult {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobResult {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
status: "completed".to_string(),
|
status: "completed".to_string(),
|
||||||
session_id: None,
|
session_id: None,
|
||||||
@@ -450,7 +458,7 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||||
|
|
||||||
let handle = spawn_job_monitor_with_context(
|
let handle = spawn_job_monitor_with_context(
|
||||||
@@ -465,7 +473,8 @@ mod tests {
|
|||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobResult {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobResult {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
status: "failed".to_string(),
|
status: "failed".to_string(),
|
||||||
session_id: None,
|
session_id: None,
|
||||||
@@ -498,13 +507,14 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
let (event_tx, _) = broadcast::channel::<(Uuid, String, AppEvent)>(16);
|
||||||
let handle = spawn_completion_watcher(job_id, event_tx.subscribe(), Arc::clone(&cm));
|
let handle = spawn_completion_watcher(job_id, event_tx.subscribe(), Arc::clone(&cm));
|
||||||
|
|
||||||
event_tx
|
event_tx
|
||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
SseEvent::JobResult {
|
"test-user".to_string(),
|
||||||
|
AppEvent::JobResult {
|
||||||
job_id: job_id.to_string(),
|
job_id: job_id.to_string(),
|
||||||
status: "completed".to_string(),
|
status: "completed".to_string(),
|
||||||
session_id: None,
|
session_id: None,
|
||||||
|
|||||||
+267
-31
@@ -18,13 +18,14 @@ use std::time::Duration;
|
|||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use regex::Regex;
|
use regex::Regex;
|
||||||
use tokio::sync::{RwLock, mpsc};
|
use tokio::sync::{RwLock, mpsc};
|
||||||
|
use tokio::task::JoinHandle;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::Scheduler;
|
use crate::agent::Scheduler;
|
||||||
use crate::agent::routine::{
|
use crate::agent::routine::{
|
||||||
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
|
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
|
||||||
};
|
};
|
||||||
use crate::channels::OutgoingResponse;
|
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||||
use crate::config::RoutineConfig;
|
use crate::config::RoutineConfig;
|
||||||
use crate::context::{JobContext, JobState};
|
use crate::context::{JobContext, JobState};
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
@@ -45,6 +46,11 @@ enum EventMatcher {
|
|||||||
System { routine: Routine },
|
System { routine: Routine },
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct TriggeredRoutine {
|
||||||
|
routine: Routine,
|
||||||
|
detail: String,
|
||||||
|
}
|
||||||
|
|
||||||
/// Distinguishes why sandbox is unavailable so error messages are accurate.
|
/// Distinguishes why sandbox is unavailable so error messages are accurate.
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub enum SandboxReadiness {
|
pub enum SandboxReadiness {
|
||||||
@@ -56,6 +62,40 @@ pub enum SandboxReadiness {
|
|||||||
DockerUnavailable,
|
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.
|
/// The routine execution engine.
|
||||||
pub struct RoutineEngine {
|
pub struct RoutineEngine {
|
||||||
config: RoutineConfig,
|
config: RoutineConfig,
|
||||||
@@ -167,10 +207,45 @@ impl RoutineEngine {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Check incoming message against event triggers. Returns number of routines fired.
|
/// Check incoming message against event triggers. Returns number of routines fired.
|
||||||
|
pub async fn check_event_triggers(&self, message: &IncomingMessage, content: &str) -> usize {
|
||||||
|
let triggered = self.matching_event_triggers(message, content).await;
|
||||||
|
let fired = triggered.len();
|
||||||
|
for triggered in triggered {
|
||||||
|
std::mem::drop(self.spawn_fire(triggered.routine, "event", Some(triggered.detail)));
|
||||||
|
}
|
||||||
|
fired
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fire matching event-triggered routines and wait for them to complete.
|
||||||
///
|
///
|
||||||
/// Accepts only the three fields needed for matching (user scope, channel,
|
/// Used by single-message REPL mode so the process does not exit before
|
||||||
/// message content) so callers never need to clone a full `IncomingMessage`.
|
/// background event-triggered routines finish.
|
||||||
pub async fn check_event_triggers(&self, user_id: &str, channel: &str, content: &str) -> usize {
|
pub async fn check_event_triggers_and_wait(
|
||||||
|
&self,
|
||||||
|
message: &IncomingMessage,
|
||||||
|
content: &str,
|
||||||
|
) -> usize {
|
||||||
|
let triggered = self.matching_event_triggers(message, content).await;
|
||||||
|
let fired = triggered.len();
|
||||||
|
let handles: Vec<JoinHandle<()>> = triggered
|
||||||
|
.into_iter()
|
||||||
|
.map(|triggered| self.spawn_fire(triggered.routine, "event", Some(triggered.detail)))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
for handle in handles {
|
||||||
|
if let Err(e) = handle.await {
|
||||||
|
tracing::warn!(error = %e, "Event-triggered routine task failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fired
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn matching_event_triggers(
|
||||||
|
&self,
|
||||||
|
message: &IncomingMessage,
|
||||||
|
content: &str,
|
||||||
|
) -> Vec<TriggeredRoutine> {
|
||||||
let cache = self.event_cache.read().await;
|
let cache = self.event_cache.read().await;
|
||||||
|
|
||||||
// Early return if there are no message matchers at all.
|
// Early return if there are no message matchers at all.
|
||||||
@@ -178,10 +253,9 @@ impl RoutineEngine {
|
|||||||
.iter()
|
.iter()
|
||||||
.any(|m| matches!(m, EventMatcher::Message { .. }))
|
.any(|m| matches!(m, EventMatcher::Message { .. }))
|
||||||
{
|
{
|
||||||
return 0;
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
let mut triggered = Vec::new();
|
||||||
let mut fired = 0;
|
|
||||||
|
|
||||||
// Collect routine IDs for batch query
|
// Collect routine IDs for batch query
|
||||||
let routine_ids: Vec<Uuid> = cache
|
let routine_ids: Vec<Uuid> = cache
|
||||||
@@ -193,13 +267,13 @@ impl RoutineEngine {
|
|||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if routine_ids.is_empty() {
|
if routine_ids.is_empty() {
|
||||||
return 0;
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Single batch query instead of N queries
|
// Single batch query instead of N queries
|
||||||
let concurrent_counts = match self.batch_concurrent_counts(&routine_ids).await {
|
let concurrent_counts = match self.batch_concurrent_counts(&routine_ids).await {
|
||||||
Some(counts) => counts,
|
Some(counts) => counts,
|
||||||
None => return 0,
|
None => return Vec::new(),
|
||||||
};
|
};
|
||||||
|
|
||||||
for matcher in cache.iter() {
|
for matcher in cache.iter() {
|
||||||
@@ -208,16 +282,24 @@ impl RoutineEngine {
|
|||||||
EventMatcher::System { .. } => continue,
|
EventMatcher::System { .. } => continue,
|
||||||
};
|
};
|
||||||
|
|
||||||
if routine.user_id != user_id {
|
// User ownership + channel filter (extracted for testability).
|
||||||
continue;
|
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.
|
||||||
// Channel filter
|
if routine.user_id != message.user_id {
|
||||||
if let Trigger::Event {
|
tracing::trace!(
|
||||||
channel: Some(ch), ..
|
routine = %routine.name,
|
||||||
} = &routine.trigger
|
routine_user = %routine.user_id,
|
||||||
&& ch != channel
|
message_user = %message.user_id,
|
||||||
{
|
"Skipped: user scope mismatch"
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
tracing::debug!(
|
||||||
|
routine = %routine.name,
|
||||||
|
channel = %message.channel,
|
||||||
|
"Skipped: channel mismatch"
|
||||||
|
);
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -228,14 +310,14 @@ impl RoutineEngine {
|
|||||||
|
|
||||||
// Cooldown check
|
// Cooldown check
|
||||||
if !self.check_cooldown(routine) {
|
if !self.check_cooldown(routine) {
|
||||||
tracing::trace!(routine = %routine.name, "Skipped: cooldown active");
|
tracing::debug!(routine = %routine.name, "Skipped: cooldown active");
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Concurrent run check (using batch-loaded counts)
|
// Concurrent run check (using batch-loaded counts)
|
||||||
let running_count = concurrent_counts.get(&routine.id).copied().unwrap_or(0);
|
let running_count = concurrent_counts.get(&routine.id).copied().unwrap_or(0);
|
||||||
if running_count >= routine.guardrails.max_concurrent as i64 {
|
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;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -246,11 +328,13 @@ impl RoutineEngine {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let detail = truncate(content, 200);
|
let detail = truncate(content, 200);
|
||||||
self.spawn_fire(routine.clone(), "event", Some(detail));
|
triggered.push(TriggeredRoutine {
|
||||||
fired += 1;
|
routine: routine.clone(),
|
||||||
|
detail,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
fired
|
triggered
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Emit a structured event to system-event routines.
|
/// Emit a structured event to system-event routines.
|
||||||
@@ -806,7 +890,12 @@ impl RoutineEngine {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Spawn a fire in a background task.
|
/// Spawn a fire in a background task.
|
||||||
fn spawn_fire(&self, routine: Routine, trigger_type: &str, trigger_detail: Option<String>) {
|
fn spawn_fire(
|
||||||
|
&self,
|
||||||
|
routine: Routine,
|
||||||
|
trigger_type: &str,
|
||||||
|
trigger_detail: Option<String>,
|
||||||
|
) -> JoinHandle<()> {
|
||||||
let run = RoutineRun {
|
let run = RoutineRun {
|
||||||
id: Uuid::new_v4(),
|
id: Uuid::new_v4(),
|
||||||
routine_id: routine.id,
|
routine_id: routine.id,
|
||||||
@@ -843,7 +932,7 @@ impl RoutineEngine {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
execute_routine(engine, routine, run).await;
|
execute_routine(engine, routine, run).await;
|
||||||
});
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn check_cooldown(&self, routine: &Routine) -> bool {
|
fn check_cooldown(&self, routine: &Routine) -> bool {
|
||||||
@@ -1305,6 +1394,19 @@ async fn execute_lightweight(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Sanitize a user-controlled string before interpolation into an LLM prompt.
|
||||||
|
/// Strips newlines (which could break prompt structure) and truncates to a
|
||||||
|
/// reasonable length to limit abuse surface.
|
||||||
|
fn sanitize_prompt_field(value: &str) -> String {
|
||||||
|
const MAX_LEN: usize = 128;
|
||||||
|
value
|
||||||
|
.chars()
|
||||||
|
.filter(|&c| c != '\n' && c != '\r')
|
||||||
|
.take(MAX_LEN)
|
||||||
|
.map(|c| if c == '`' { '\'' } else { c })
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
fn build_lightweight_prompt(
|
fn build_lightweight_prompt(
|
||||||
prompt: &str,
|
prompt: &str,
|
||||||
context_parts: &[String],
|
context_parts: &[String],
|
||||||
@@ -1323,14 +1425,16 @@ fn build_lightweight_prompt(
|
|||||||
);
|
);
|
||||||
|
|
||||||
if let Some(channel) = notify.channel.as_deref() {
|
if let Some(channel) = notify.channel.as_deref() {
|
||||||
|
let sanitized = sanitize_prompt_field(channel);
|
||||||
full_prompt.push_str(&format!(
|
full_prompt.push_str(&format!(
|
||||||
"The configured delivery channel for this routine is `{channel}`.\n"
|
"The configured delivery channel for this routine is `{sanitized}`.\n"
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(user) = notify.user.as_deref() {
|
if let Some(user) = notify.user.as_deref() {
|
||||||
|
let sanitized = sanitize_prompt_field(user);
|
||||||
full_prompt.push_str(&format!(
|
full_prompt.push_str(&format!(
|
||||||
"The configured delivery target for this routine is `{user}`.\n"
|
"The configured delivery target for this routine is `{sanitized}`.\n"
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1440,6 +1544,7 @@ fn handle_text_response(
|
|||||||
/// This is a simplified version of the full dispatcher loop:
|
/// This is a simplified version of the full dispatcher loop:
|
||||||
/// - Max 3-5 iterations (configurable)
|
/// - Max 3-5 iterations (configurable)
|
||||||
/// - Sequential tool execution (not parallel)
|
/// - Sequential tool execution (not parallel)
|
||||||
|
/// - Uses the owner's live autonomous tool scope when lightweight tools are enabled
|
||||||
/// - Auto-approval of non-Always tools
|
/// - Auto-approval of non-Always tools
|
||||||
/// - No hooks or approval dialogs
|
/// - No hooks or approval dialogs
|
||||||
async fn execute_lightweight_with_tools(
|
async fn execute_lightweight_with_tools(
|
||||||
@@ -1486,7 +1591,10 @@ async fn execute_lightweight_with_tools(
|
|||||||
let force_text = iteration >= max_iterations;
|
let force_text = iteration >= max_iterations;
|
||||||
|
|
||||||
if force_text {
|
if force_text {
|
||||||
// Final iteration: no tools, just get text response
|
// Final iteration: no tools, just get text response.
|
||||||
|
// Claude 4.6 rejects assistant prefill; NEAR AI rejects any non-user-ending
|
||||||
|
// conversation. Ensure the last message is user-role.
|
||||||
|
crate::util::ensure_ends_with_user_message(&mut messages);
|
||||||
let request = CompletionRequest::new(messages)
|
let request = CompletionRequest::new(messages)
|
||||||
.with_max_tokens(effective_max_tokens)
|
.with_max_tokens(effective_max_tokens)
|
||||||
.with_temperature(0.3);
|
.with_temperature(0.3);
|
||||||
@@ -1765,6 +1873,13 @@ pub fn spawn_cron_ticker(
|
|||||||
engine.check_cron_triggers().await;
|
engine.check_cron_triggers().await;
|
||||||
|
|
||||||
let mut ticker = tokio::time::interval(interval);
|
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 {
|
loop {
|
||||||
ticker.tick().await;
|
ticker.tick().await;
|
||||||
@@ -1772,7 +1887,11 @@ pub fn spawn_cron_ticker(
|
|||||||
// never races with FullJobWatcher instances from this process.
|
// never races with FullJobWatcher instances from this process.
|
||||||
engine.sync_dispatched_runs().await;
|
engine.sync_dispatched_runs().await;
|
||||||
engine.check_cron_triggers().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();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -1838,7 +1957,13 @@ fn strip_html_tags(s: &str) -> String {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
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;
|
use crate::config::RoutineConfig;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -2036,6 +2161,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]
|
#[test]
|
||||||
fn test_routine_tool_denylist_blocks_self_management_tools() {
|
fn test_routine_tool_denylist_blocks_self_management_tools() {
|
||||||
let denylisted = vec![
|
let denylisted = vec![
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ use tokio::task::JoinHandle;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::task::{Task, TaskContext, TaskOutput};
|
use crate::agent::task::{Task, TaskContext, TaskOutput};
|
||||||
use crate::channels::web::types::SseEvent;
|
|
||||||
use crate::config::AgentConfig;
|
use crate::config::AgentConfig;
|
||||||
use crate::context::{ContextManager, JobContext, JobState};
|
use crate::context::{ContextManager, JobContext, JobState};
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
@@ -67,8 +66,8 @@ pub struct Scheduler {
|
|||||||
extension_manager: Option<Arc<ExtensionManager>>,
|
extension_manager: Option<Arc<ExtensionManager>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<Arc<dyn Database>>,
|
||||||
hooks: Arc<HookRegistry>,
|
hooks: Arc<HookRegistry>,
|
||||||
/// SSE broadcast sender for live job event streaming.
|
/// SSE manager for live job event streaming.
|
||||||
sse_tx: Option<tokio::sync::broadcast::Sender<SseEvent>>,
|
sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||||
/// HTTP interceptor for trace recording/replay (propagated to workers).
|
/// HTTP interceptor for trace recording/replay (propagated to workers).
|
||||||
http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
|
http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
|
||||||
/// Running jobs (main LLM-driven jobs).
|
/// Running jobs (main LLM-driven jobs).
|
||||||
@@ -102,9 +101,9 @@ impl Scheduler {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the SSE broadcast sender for live job event streaming.
|
/// Set the SSE manager for live job event streaming.
|
||||||
pub fn set_sse_sender(&mut self, tx: tokio::sync::broadcast::Sender<SseEvent>) {
|
pub fn set_sse_sender(&mut self, sse: Arc<crate::channels::web::sse::SseManager>) {
|
||||||
self.sse_tx = Some(tx);
|
self.sse_tx = Some(sse);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the HTTP interceptor for trace recording/replay.
|
/// Set the HTTP interceptor for trace recording/replay.
|
||||||
|
|||||||
+193
-2
@@ -16,8 +16,8 @@ use chrono::{DateTime, TimeDelta, Utc};
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::web::util::truncate_preview;
|
|
||||||
use crate::llm::{ChatMessage, ToolCall, generate_tool_call_id};
|
use crate::llm::{ChatMessage, ToolCall, generate_tool_call_id};
|
||||||
|
use ironclaw_common::truncate_preview;
|
||||||
|
|
||||||
/// A session containing one or more threads.
|
/// A session containing one or more threads.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -449,6 +449,7 @@ impl Thread {
|
|||||||
id: call_id.clone(),
|
id: call_id.clone(),
|
||||||
name: tc.name.clone(),
|
name: tc.name.clone(),
|
||||||
arguments: tc.parameters.clone(),
|
arguments: tc.parameters.clone(),
|
||||||
|
reasoning: None,
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -522,7 +523,12 @@ impl Thread {
|
|||||||
&& let Some(ref tcs) = assistant_msg.tool_calls
|
&& let Some(ref tcs) = assistant_msg.tool_calls
|
||||||
{
|
{
|
||||||
for tc in tcs {
|
for tc in tcs {
|
||||||
turn.record_tool_call(&tc.name, tc.arguments.clone());
|
turn.record_tool_call_with_reasoning(
|
||||||
|
&tc.name,
|
||||||
|
tc.arguments.clone(),
|
||||||
|
tc.reasoning.clone(),
|
||||||
|
Some(tc.id.clone()),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -602,6 +608,10 @@ pub struct Turn {
|
|||||||
pub completed_at: Option<DateTime<Utc>>,
|
pub completed_at: Option<DateTime<Utc>>,
|
||||||
/// Error message (if failed).
|
/// Error message (if failed).
|
||||||
pub error: Option<String>,
|
pub error: Option<String>,
|
||||||
|
/// Agent's reasoning narrative for this turn.
|
||||||
|
/// Cleaned via `clean_response` and sanitized through `SafetyLayer` before storage.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub narrative: Option<String>,
|
||||||
/// Transient image content parts for multimodal LLM input.
|
/// Transient image content parts for multimodal LLM input.
|
||||||
/// Not serialized — images are only needed for the current LLM call.
|
/// Not serialized — images are only needed for the current LLM call.
|
||||||
/// The text description in `user_input` persists for compaction/context.
|
/// The text description in `user_input` persists for compaction/context.
|
||||||
@@ -621,6 +631,7 @@ impl Turn {
|
|||||||
started_at: Utc::now(),
|
started_at: Utc::now(),
|
||||||
completed_at: None,
|
completed_at: None,
|
||||||
error: None,
|
error: None,
|
||||||
|
narrative: None,
|
||||||
image_content_parts: Vec::new(),
|
image_content_parts: Vec::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -656,6 +667,26 @@ impl Turn {
|
|||||||
parameters: params,
|
parameters: params,
|
||||||
result: None,
|
result: None,
|
||||||
error: None,
|
error: None,
|
||||||
|
rationale: None,
|
||||||
|
tool_call_id: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a tool call with reasoning context.
|
||||||
|
pub fn record_tool_call_with_reasoning(
|
||||||
|
&mut self,
|
||||||
|
name: impl Into<String>,
|
||||||
|
params: serde_json::Value,
|
||||||
|
rationale: Option<String>,
|
||||||
|
tool_call_id: Option<String>,
|
||||||
|
) {
|
||||||
|
self.tool_calls.push(TurnToolCall {
|
||||||
|
name: name.into(),
|
||||||
|
parameters: params,
|
||||||
|
result: None,
|
||||||
|
error: None,
|
||||||
|
rationale,
|
||||||
|
tool_call_id,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -672,6 +703,60 @@ impl Turn {
|
|||||||
call.error = Some(error.into());
|
call.error = Some(error.into());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Record a tool result by tool_call_id, with fallback to first pending call.
|
||||||
|
pub fn record_tool_result_for(&mut self, tool_call_id: &str, result: serde_json::Value) {
|
||||||
|
if let Some(call) = self
|
||||||
|
.tool_calls
|
||||||
|
.iter_mut()
|
||||||
|
.find(|c| c.tool_call_id.as_deref() == Some(tool_call_id))
|
||||||
|
{
|
||||||
|
call.result = Some(result);
|
||||||
|
} else if let Some(call) = self
|
||||||
|
.tool_calls
|
||||||
|
.iter_mut()
|
||||||
|
.find(|c| c.result.is_none() && c.error.is_none())
|
||||||
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
tool_call_id = %tool_call_id,
|
||||||
|
fallback_tool = %call.name,
|
||||||
|
"tool_call_id not found, falling back to first pending call"
|
||||||
|
);
|
||||||
|
call.result = Some(result);
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
tool_call_id = %tool_call_id,
|
||||||
|
"Tool result dropped: no matching or pending tool call"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a tool error by tool_call_id, with fallback to first pending call.
|
||||||
|
pub fn record_tool_error_for(&mut self, tool_call_id: &str, error: impl Into<String>) {
|
||||||
|
if let Some(call) = self
|
||||||
|
.tool_calls
|
||||||
|
.iter_mut()
|
||||||
|
.find(|c| c.tool_call_id.as_deref() == Some(tool_call_id))
|
||||||
|
{
|
||||||
|
call.error = Some(error.into());
|
||||||
|
} else if let Some(call) = self
|
||||||
|
.tool_calls
|
||||||
|
.iter_mut()
|
||||||
|
.find(|c| c.result.is_none() && c.error.is_none())
|
||||||
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
tool_call_id = %tool_call_id,
|
||||||
|
fallback_tool = %call.name,
|
||||||
|
"tool_call_id not found, falling back to first pending call"
|
||||||
|
);
|
||||||
|
call.error = Some(error.into());
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
tool_call_id = %tool_call_id,
|
||||||
|
"Tool error dropped: no matching or pending tool call"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Record of a tool call made during a turn.
|
/// Record of a tool call made during a turn.
|
||||||
@@ -685,6 +770,12 @@ pub struct TurnToolCall {
|
|||||||
pub result: Option<serde_json::Value>,
|
pub result: Option<serde_json::Value>,
|
||||||
/// Error from the tool (if failed).
|
/// Error from the tool (if failed).
|
||||||
pub error: Option<String>,
|
pub error: Option<String>,
|
||||||
|
/// Agent's reasoning for choosing this tool.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub rationale: Option<String>,
|
||||||
|
/// The tool_call_id from the LLM, for identity-based result matching.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tool_call_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -1309,6 +1400,7 @@ mod tests {
|
|||||||
id: "call_0".to_string(),
|
id: "call_0".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "test"}),
|
arguments: serde_json::json!({"q": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
ChatMessage::user("Find test"),
|
ChatMessage::user("Find test"),
|
||||||
@@ -1339,6 +1431,7 @@ mod tests {
|
|||||||
id: "call_0".to_string(),
|
id: "call_0".to_string(),
|
||||||
name: "http".to_string(),
|
name: "http".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
ChatMessage::user("Fetch URL"),
|
ChatMessage::user("Fetch URL"),
|
||||||
@@ -1404,11 +1497,13 @@ mod tests {
|
|||||||
id: "call_a".to_string(),
|
id: "call_a".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "data"}),
|
arguments: serde_json::json!({"q": "data"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let tc2 = ToolCall {
|
let tc2 = ToolCall {
|
||||||
id: "call_b".to_string(),
|
id: "call_b".to_string(),
|
||||||
name: "write".to_string(),
|
name: "write".to_string(),
|
||||||
arguments: serde_json::json!({"path": "out.txt"}),
|
arguments: serde_json::json!({"path": "out.txt"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
ChatMessage::user("Find and save"),
|
ChatMessage::user("Find and save"),
|
||||||
@@ -1620,4 +1715,100 @@ mod tests {
|
|||||||
let merged = thread.drain_pending_messages().unwrap();
|
let merged = thread.drain_pending_messages().unwrap();
|
||||||
assert_eq!(merged, "failed batch\nnew msg");
|
assert_eq!(merged, "failed batch\nnew msg");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_record_tool_result_for_by_id() {
|
||||||
|
let mut turn = Turn::new(0, "test");
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_a",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_a".into()),
|
||||||
|
);
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_b",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_b".into()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Record result for second tool by ID
|
||||||
|
turn.record_tool_result_for("id_b", serde_json::json!("result_b"));
|
||||||
|
assert!(turn.tool_calls[0].result.is_none());
|
||||||
|
assert_eq!(
|
||||||
|
turn.tool_calls[1].result.as_ref().unwrap(),
|
||||||
|
&serde_json::json!("result_b")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_record_tool_error_for_by_id() {
|
||||||
|
let mut turn = Turn::new(0, "test");
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_a",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_a".into()),
|
||||||
|
);
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_b",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_b".into()),
|
||||||
|
);
|
||||||
|
|
||||||
|
turn.record_tool_error_for("id_a", "failed");
|
||||||
|
assert_eq!(turn.tool_calls[0].error.as_deref(), Some("failed"));
|
||||||
|
assert!(turn.tool_calls[1].error.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_record_tool_result_for_fallback_to_pending() {
|
||||||
|
let mut turn = Turn::new(0, "test");
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_a",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_a".into()),
|
||||||
|
);
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_b",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_b".into()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// First tool already has a result
|
||||||
|
turn.tool_calls[0].result = Some(serde_json::json!("done"));
|
||||||
|
|
||||||
|
// Unknown ID should fall back to first pending (tool_b)
|
||||||
|
turn.record_tool_result_for("unknown_id", serde_json::json!("fallback"));
|
||||||
|
assert_eq!(
|
||||||
|
turn.tool_calls[0].result.as_ref().unwrap(),
|
||||||
|
&serde_json::json!("done")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
turn.tool_calls[1].result.as_ref().unwrap(),
|
||||||
|
&serde_json::json!("fallback")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_record_tool_result_for_no_pending_is_noop() {
|
||||||
|
let mut turn = Turn::new(0, "test");
|
||||||
|
turn.record_tool_call_with_reasoning(
|
||||||
|
"tool_a",
|
||||||
|
serde_json::json!({}),
|
||||||
|
None,
|
||||||
|
Some("id_a".into()),
|
||||||
|
);
|
||||||
|
turn.tool_calls[0].result = Some(serde_json::json!("done"));
|
||||||
|
|
||||||
|
// No pending calls, unknown ID — should be a no-op
|
||||||
|
turn.record_tool_result_for("unknown_id", serde_json::json!("lost"));
|
||||||
|
assert_eq!(
|
||||||
|
turn.tool_calls[0].result.as_ref().unwrap(),
|
||||||
|
&serde_json::json!("done")
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+189
-34
@@ -102,11 +102,30 @@ impl SessionManager {
|
|||||||
/// Resolve an external thread ID to an internal thread.
|
/// Resolve an external thread ID to an internal thread.
|
||||||
///
|
///
|
||||||
/// Returns the session and thread ID. Creates both if they don't exist.
|
/// Returns the session and thread ID. Creates both if they don't exist.
|
||||||
|
/// Delegates to [`resolve_thread_with_parsed_uuid`](Self::resolve_thread_with_parsed_uuid)
|
||||||
|
/// with `parsed_uuid: None`.
|
||||||
pub async fn resolve_thread(
|
pub async fn resolve_thread(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
user_id: &str,
|
||||||
channel: &str,
|
channel: &str,
|
||||||
external_thread_id: Option<&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) {
|
) -> (Arc<Mutex<Session>>, Uuid) {
|
||||||
let session = self.get_or_create_session(user_id).await;
|
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),
|
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;
|
let thread_map = self.thread_map.read().await;
|
||||||
|
|
||||||
|
// Fast path: exact key match
|
||||||
if let Some(&thread_id) = thread_map.get(&key) {
|
if let Some(&thread_id) = thread_map.get(&key) {
|
||||||
// Verify thread still exists in session
|
|
||||||
let sess = session.lock().await;
|
let sess = session.lock().await;
|
||||||
if sess.threads.contains_key(&thread_id) {
|
if sess.threads.contains_key(&thread_id) {
|
||||||
return (Arc::clone(&session), thread_id);
|
return (Arc::clone(&session), thread_id);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
// Check if external_thread_id is itself a known thread UUID that
|
// UUID adoption check (still under the same read lock).
|
||||||
// exists in the session but was never registered in the thread_map
|
// If external_thread_id is a valid UUID not mapped elsewhere,
|
||||||
// (e.g. created by chat_new_thread_handler or hydrated from DB).
|
// it may be a thread created by chat_new_thread_handler or
|
||||||
// We only adopt it if no thread_map entry maps to this UUID —
|
// hydrated from DB that we can adopt.
|
||||||
// otherwise it belongs to a different channel scope.
|
// Only attempt adoption when external_thread_id is Some, preserving
|
||||||
if let Some(ext_tid) = external_thread_id
|
// the invariant that None external_thread_id never triggers adoption.
|
||||||
&& let Ok(ext_uuid) = Uuid::parse_str(ext_tid)
|
if external_thread_id.is_some() {
|
||||||
{
|
ext_uuid.filter(|&uuid| !thread_map.values().any(|&v| v == uuid))
|
||||||
let thread_map = self.thread_map.read().await;
|
} else {
|
||||||
let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid);
|
None
|
||||||
drop(thread_map);
|
}
|
||||||
|
}; // Single read lock dropped here
|
||||||
|
|
||||||
if !mapped_elsewhere {
|
// If we found an adoptable UUID, verify it exists in session and acquire write lock
|
||||||
let sess = session.lock().await;
|
if let Some(ext_uuid) = adoptable_uuid {
|
||||||
if sess.threads.contains_key(&ext_uuid) {
|
let sess = session.lock().await;
|
||||||
drop(sess);
|
if sess.threads.contains_key(&ext_uuid) {
|
||||||
|
drop(sess);
|
||||||
|
|
||||||
let mut thread_map = self.thread_map.write().await;
|
let mut thread_map = self.thread_map.write().await;
|
||||||
// Re-check after acquiring write lock to prevent race condition
|
// Re-check after acquiring write lock to prevent race condition
|
||||||
// where another task mapped this UUID between our read and write.
|
// where another task mapped this UUID between our read and write.
|
||||||
if !thread_map.values().any(|&v| v == ext_uuid) {
|
if !thread_map.values().any(|&v| v == ext_uuid) {
|
||||||
thread_map.insert(key, ext_uuid);
|
thread_map.insert(key, ext_uuid);
|
||||||
drop(thread_map);
|
drop(thread_map);
|
||||||
// Ensure undo manager exists
|
// Ensure undo manager exists
|
||||||
let mut undo_managers = self.undo_managers.write().await;
|
let mut undo_managers = self.undo_managers.write().await;
|
||||||
undo_managers
|
undo_managers
|
||||||
.entry(ext_uuid)
|
.entry(ext_uuid)
|
||||||
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
|
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
|
||||||
return (session, ext_uuid);
|
return (session, ext_uuid);
|
||||||
}
|
|
||||||
// If it was mapped elsewhere while we were unlocked, fall through
|
|
||||||
// to create a new thread, preserving channel isolation.
|
|
||||||
}
|
}
|
||||||
|
// 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]
|
#[tokio::test]
|
||||||
async fn test_resolve_thread_finds_existing_session_thread_by_uuid() {
|
async fn test_resolve_thread_finds_existing_session_thread_by_uuid() {
|
||||||
use crate::agent::session::{Session, Thread};
|
use crate::agent::session::{Session, Thread};
|
||||||
@@ -947,4 +1018,88 @@ mod tests {
|
|||||||
"should have exactly 1 thread, not a duplicate"
|
"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"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -92,6 +92,17 @@ impl SubmissionParser {
|
|||||||
args: vec![],
|
args: vec![],
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
if lower == "/reasoning" || lower.starts_with("/reasoning ") {
|
||||||
|
let args: Vec<String> = trimmed
|
||||||
|
.split_whitespace()
|
||||||
|
.skip(1)
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.collect();
|
||||||
|
return Submission::SystemCommand {
|
||||||
|
command: "reasoning".to_string(),
|
||||||
|
args,
|
||||||
|
};
|
||||||
|
}
|
||||||
if lower == "/restart" {
|
if lower == "/restart" {
|
||||||
tracing::debug!("[SubmissionParser::parse] Recognized /restart command");
|
tracing::debug!("[SubmissionParser::parse] Recognized /restart command");
|
||||||
return Submission::SystemCommand {
|
return Submission::SystemCommand {
|
||||||
|
|||||||
+60
-13
@@ -16,12 +16,12 @@ use crate::agent::dispatcher::{
|
|||||||
};
|
};
|
||||||
use crate::agent::session::{MAX_PENDING_MESSAGES, PendingApproval, Session, ThreadState};
|
use crate::agent::session::{MAX_PENDING_MESSAGES, PendingApproval, Session, ThreadState};
|
||||||
use crate::agent::submission::SubmissionResult;
|
use crate::agent::submission::SubmissionResult;
|
||||||
use crate::channels::web::util::truncate_preview;
|
|
||||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
use crate::llm::{ChatMessage, ToolCall};
|
use crate::llm::{ChatMessage, ToolCall};
|
||||||
use crate::tools::redact_params;
|
use crate::tools::redact_params;
|
||||||
|
use ironclaw_common::truncate_preview;
|
||||||
|
|
||||||
const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID.";
|
const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID.";
|
||||||
|
|
||||||
@@ -513,10 +513,10 @@ impl Agent {
|
|||||||
};
|
};
|
||||||
|
|
||||||
thread.complete_turn(&response);
|
thread.complete_turn(&response);
|
||||||
let (turn_number, tool_calls) = thread
|
let (turn_number, tool_calls, narrative) = thread
|
||||||
.turns
|
.turns
|
||||||
.last()
|
.last()
|
||||||
.map(|t| (t.turn_number, t.tool_calls.clone()))
|
.map(|t| (t.turn_number, t.tool_calls.clone(), t.narrative.clone()))
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let _ = self
|
let _ = self
|
||||||
.channels
|
.channels
|
||||||
@@ -534,6 +534,7 @@ impl Agent {
|
|||||||
&message.user_id,
|
&message.user_id,
|
||||||
turn_number,
|
turn_number,
|
||||||
&tool_calls,
|
&tool_calls,
|
||||||
|
narrative.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
self.persist_assistant_response(
|
self.persist_assistant_response(
|
||||||
@@ -725,7 +726,9 @@ impl Agent {
|
|||||||
///
|
///
|
||||||
/// Stored between the user and assistant messages so that
|
/// Stored between the user and assistant messages so that
|
||||||
/// `build_turns_from_db_messages` can reconstruct the tool call history.
|
/// `build_turns_from_db_messages` can reconstruct the tool call history.
|
||||||
/// Content is a JSON array of tool call summaries.
|
/// Content is a JSON object: `{ "calls": [...], "narrative": "..." }`.
|
||||||
|
/// The `calls` array contains tool call summaries with optional `rationale`
|
||||||
|
/// and `tool_call_id` fields. Legacy rows may be plain JSON arrays.
|
||||||
pub(super) async fn persist_tool_calls(
|
pub(super) async fn persist_tool_calls(
|
||||||
&self,
|
&self,
|
||||||
thread_id: Uuid,
|
thread_id: Uuid,
|
||||||
@@ -733,6 +736,7 @@ impl Agent {
|
|||||||
user_id: &str,
|
user_id: &str,
|
||||||
turn_number: usize,
|
turn_number: usize,
|
||||||
tool_calls: &[crate::agent::session::TurnToolCall],
|
tool_calls: &[crate::agent::session::TurnToolCall],
|
||||||
|
narrative: Option<&str>,
|
||||||
) {
|
) {
|
||||||
if tool_calls.is_empty() {
|
if tool_calls.is_empty() {
|
||||||
return;
|
return;
|
||||||
@@ -767,11 +771,30 @@ impl Agent {
|
|||||||
if let Some(ref error) = tc.error {
|
if let Some(ref error) = tc.error {
|
||||||
obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
|
obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
|
||||||
}
|
}
|
||||||
|
if let Some(ref rationale) = tc.rationale {
|
||||||
|
obj["rationale"] = serde_json::Value::String(truncate_preview(rationale, 500));
|
||||||
|
}
|
||||||
|
if let Some(ref tool_call_id) = tc.tool_call_id {
|
||||||
|
obj["tool_call_id"] =
|
||||||
|
serde_json::Value::String(truncate_preview(tool_call_id, 128));
|
||||||
|
}
|
||||||
obj
|
obj
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
let content = match serde_json::to_string(&summaries) {
|
// Wrap in an object with optional narrative so it can be reconstructed.
|
||||||
|
// safety: no byte-index slicing here; comment describes JSON shape
|
||||||
|
let wrapper = if let Some(n) = narrative {
|
||||||
|
serde_json::json!({
|
||||||
|
"narrative": truncate_preview(n, 1000),
|
||||||
|
"calls": summaries,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
serde_json::json!({
|
||||||
|
"calls": summaries,
|
||||||
|
})
|
||||||
|
};
|
||||||
|
let content = match serde_json::to_string(&wrapper) {
|
||||||
Ok(c) => c,
|
Ok(c) => c,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!("Failed to serialize tool calls: {}", e);
|
tracing::warn!("Failed to serialize tool calls: {}", e);
|
||||||
@@ -1104,9 +1127,12 @@ impl Agent {
|
|||||||
&& let Some(turn) = thread.last_turn_mut()
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
{
|
{
|
||||||
if is_tool_error {
|
if is_tool_error {
|
||||||
turn.record_tool_error(result_content.clone());
|
turn.record_tool_error_for(&pending.tool_call_id, result_content.clone());
|
||||||
} else {
|
} else {
|
||||||
turn.record_tool_result(serde_json::json!(result_content));
|
turn.record_tool_result_for(
|
||||||
|
&pending.tool_call_id,
|
||||||
|
serde_json::json!(result_content),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1358,9 +1384,12 @@ impl Agent {
|
|||||||
&& let Some(turn) = thread.last_turn_mut()
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
{
|
{
|
||||||
if is_deferred_error {
|
if is_deferred_error {
|
||||||
turn.record_tool_error(deferred_content.clone());
|
turn.record_tool_error_for(&tc.id, deferred_content.clone());
|
||||||
} else {
|
} else {
|
||||||
turn.record_tool_result(serde_json::json!(deferred_content));
|
turn.record_tool_result_for(
|
||||||
|
&tc.id,
|
||||||
|
serde_json::json!(deferred_content),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1459,10 +1488,10 @@ impl Agent {
|
|||||||
let (response, suggestions) =
|
let (response, suggestions) =
|
||||||
crate::agent::dispatcher::extract_suggestions(&response);
|
crate::agent::dispatcher::extract_suggestions(&response);
|
||||||
thread.complete_turn(&response);
|
thread.complete_turn(&response);
|
||||||
let (turn_number, tool_calls) = thread
|
let (turn_number, tool_calls, narrative) = thread
|
||||||
.turns
|
.turns
|
||||||
.last()
|
.last()
|
||||||
.map(|t| (t.turn_number, t.tool_calls.clone()))
|
.map(|t| (t.turn_number, t.tool_calls.clone(), t.narrative.clone()))
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
// User message already persisted at turn start; save tool calls then assistant response
|
// User message already persisted at turn start; save tool calls then assistant response
|
||||||
self.persist_tool_calls(
|
self.persist_tool_calls(
|
||||||
@@ -1471,6 +1500,7 @@ impl Agent {
|
|||||||
&message.user_id,
|
&message.user_id,
|
||||||
turn_number,
|
turn_number,
|
||||||
&tool_calls,
|
&tool_calls,
|
||||||
|
narrative.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
self.persist_assistant_response(
|
self.persist_assistant_response(
|
||||||
@@ -1646,7 +1676,7 @@ impl Agent {
|
|||||||
};
|
};
|
||||||
|
|
||||||
match ext_mgr
|
match ext_mgr
|
||||||
.configure_token(&pending.extension_name, token)
|
.configure_token(&pending.extension_name, token, &message.user_id)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(result) if result.activated => {
|
Ok(result) if result.activated => {
|
||||||
@@ -1816,7 +1846,20 @@ fn rebuild_chat_messages_from_db(
|
|||||||
"assistant" => result.push(ChatMessage::assistant(&msg.content)),
|
"assistant" => result.push(ChatMessage::assistant(&msg.content)),
|
||||||
"tool_calls" => {
|
"tool_calls" => {
|
||||||
// Try to parse the enriched JSON and rebuild tool messages.
|
// Try to parse the enriched JSON and rebuild tool messages.
|
||||||
if let Ok(calls) = serde_json::from_str::<Vec<serde_json::Value>>(&msg.content) {
|
// Supports two formats:
|
||||||
|
// - Old: plain JSON array of tool call summaries
|
||||||
|
// - New: wrapped object { "calls": [...], "narrative": "..." }
|
||||||
|
let calls: Vec<serde_json::Value> =
|
||||||
|
match serde_json::from_str::<serde_json::Value>(&msg.content) {
|
||||||
|
Ok(serde_json::Value::Array(arr)) => arr,
|
||||||
|
Ok(serde_json::Value::Object(obj)) => obj
|
||||||
|
.get("calls")
|
||||||
|
.and_then(|v| v.as_array())
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_default(),
|
||||||
|
_ => Vec::new(),
|
||||||
|
};
|
||||||
|
{
|
||||||
if calls.is_empty() {
|
if calls.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -1839,6 +1882,10 @@ fn rebuild_chat_messages_from_db(
|
|||||||
.get("parameters")
|
.get("parameters")
|
||||||
.cloned()
|
.cloned()
|
||||||
.unwrap_or(serde_json::json!({})),
|
.unwrap_or(serde_json::json!({})),
|
||||||
|
reasoning: c
|
||||||
|
.get("rationale")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(String::from),
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
|
|||||||
+31
-9
@@ -312,13 +312,7 @@ impl AppBuilder {
|
|||||||
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
|
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
|
||||||
|
|
||||||
// Register memory tools if database is available
|
// Register memory tools if database is available
|
||||||
let workspace_user_id = self
|
let workspace_user_id = self.config.owner_id.as_str();
|
||||||
.config
|
|
||||||
.channels
|
|
||||||
.gateway
|
|
||||||
.as_ref()
|
|
||||||
.map(|gw| gw.user_id.as_str())
|
|
||||||
.unwrap_or("default");
|
|
||||||
let workspace = if let Some(ref db) = self.db {
|
let workspace = if let Some(ref db) = self.db {
|
||||||
let emb_cache_config = EmbeddingCacheConfig {
|
let emb_cache_config = EmbeddingCacheConfig {
|
||||||
max_entries: self.config.embeddings.cache_size,
|
max_entries: self.config.embeddings.cache_size,
|
||||||
@@ -327,7 +321,7 @@ impl AppBuilder {
|
|||||||
.with_search_config(&self.config.search);
|
.with_search_config(&self.config.search);
|
||||||
|
|
||||||
if let Some(ref emb) = embeddings {
|
if let Some(ref emb) = embeddings {
|
||||||
ws = ws.with_embeddings_cached(emb.clone(), emb_cache_config);
|
ws = ws.with_embeddings_cached(emb.clone(), emb_cache_config.clone());
|
||||||
}
|
}
|
||||||
|
|
||||||
// Wire workspace-level settings (read scopes, memory layers)
|
// Wire workspace-level settings (read scopes, memory layers)
|
||||||
@@ -341,7 +335,35 @@ impl AppBuilder {
|
|||||||
}
|
}
|
||||||
ws = ws.with_memory_layers(self.config.workspace.memory_layers.clone());
|
ws = ws.with_memory_layers(self.config.workspace.memory_layers.clone());
|
||||||
let ws = Arc::new(ws);
|
let ws = Arc::new(ws);
|
||||||
tools.register_memory_tools(Arc::clone(&ws));
|
|
||||||
|
// Detect multi-tenant mode: when GATEWAY_USER_TOKENS is configured,
|
||||||
|
// each authenticated user needs their own workspace scope. Use
|
||||||
|
// WorkspacePool (which implements WorkspaceResolver) to create
|
||||||
|
// per-user workspaces on demand instead of sharing the startup
|
||||||
|
// workspace across all users.
|
||||||
|
let is_multi_tenant = self
|
||||||
|
.config
|
||||||
|
.channels
|
||||||
|
.gateway
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|gw| gw.user_tokens.is_some());
|
||||||
|
|
||||||
|
if is_multi_tenant {
|
||||||
|
let pool = Arc::new(crate::channels::web::server::WorkspacePool::new(
|
||||||
|
Arc::clone(db),
|
||||||
|
embeddings.clone(),
|
||||||
|
emb_cache_config,
|
||||||
|
self.config.search.clone(),
|
||||||
|
self.config.workspace.clone(),
|
||||||
|
));
|
||||||
|
tools.register_memory_tools_with_resolver(pool);
|
||||||
|
tracing::info!(
|
||||||
|
"Memory tools configured with per-user workspace resolver (multi-tenant mode)"
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
tools.register_memory_tools(Arc::clone(&ws));
|
||||||
|
}
|
||||||
|
|
||||||
Some(ws)
|
Some(ws)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
|
|||||||
@@ -265,6 +265,15 @@ impl OutgoingResponse {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A single tool decision within a reasoning update.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ToolDecision {
|
||||||
|
/// Tool name.
|
||||||
|
pub tool_name: String,
|
||||||
|
/// Agent's reasoning for choosing this tool.
|
||||||
|
pub rationale: String,
|
||||||
|
}
|
||||||
|
|
||||||
/// Status update types for showing agent activity.
|
/// Status update types for showing agent activity.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum StatusUpdate {
|
pub enum StatusUpdate {
|
||||||
@@ -333,6 +342,13 @@ pub enum StatusUpdate {
|
|||||||
},
|
},
|
||||||
/// Suggested follow-up messages for the user.
|
/// Suggested follow-up messages for the user.
|
||||||
Suggestions { suggestions: Vec<String> },
|
Suggestions { suggestions: Vec<String> },
|
||||||
|
/// Agent reasoning update (why it chose specific tools).
|
||||||
|
ReasoningUpdate {
|
||||||
|
/// Human-readable summary of the agent's decision.
|
||||||
|
narrative: String,
|
||||||
|
/// Per-tool decisions.
|
||||||
|
decisions: Vec<ToolDecision>,
|
||||||
|
},
|
||||||
/// Per-turn token usage and cost summary (shown as subtle metadata).
|
/// Per-turn token usage and cost summary (shown as subtle metadata).
|
||||||
TurnCost {
|
TurnCost {
|
||||||
input_tokens: u64,
|
input_tokens: u64,
|
||||||
|
|||||||
+1
-1
@@ -39,7 +39,7 @@ mod webhook_server;
|
|||||||
|
|
||||||
pub use channel::{
|
pub use channel::{
|
||||||
AttachmentKind, Channel, ChannelSecretUpdater, IncomingAttachment, IncomingMessage,
|
AttachmentKind, Channel, ChannelSecretUpdater, IncomingAttachment, IncomingMessage,
|
||||||
MessageStream, OutgoingResponse, StatusUpdate, routing_target_from_metadata,
|
MessageStream, OutgoingResponse, StatusUpdate, ToolDecision, routing_target_from_metadata,
|
||||||
};
|
};
|
||||||
pub use http::{HttpChannel, HttpChannelState};
|
pub use http::{HttpChannel, HttpChannelState};
|
||||||
pub use manager::ChannelManager;
|
pub use manager::ChannelManager;
|
||||||
|
|||||||
+65
-9
@@ -75,6 +75,7 @@ const SLASH_COMMANDS: &[&str] = &[
|
|||||||
"/suggest",
|
"/suggest",
|
||||||
"/thread",
|
"/thread",
|
||||||
"/resume",
|
"/resume",
|
||||||
|
"/reasoning",
|
||||||
];
|
];
|
||||||
|
|
||||||
/// Rustyline helper for slash-command tab completion.
|
/// Rustyline helper for slash-command tab completion.
|
||||||
@@ -430,6 +431,18 @@ impl ReplChannel {
|
|||||||
let _ = execute!(stderr, terminal::Clear(terminal::ClearType::FromCursorDown));
|
let _ = execute!(stderr, terminal::Clear(terminal::ClearType::FromCursorDown));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn finish_single_message_turn(&self) {
|
||||||
|
if self.single_message.is_none() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let tx = self.msg_tx.lock().ok().and_then(|mut guard| guard.take());
|
||||||
|
if let Some(tx) = tx {
|
||||||
|
let msg = IncomingMessage::new("repl", &self.user_id, "/quit");
|
||||||
|
let _ = tx.send(msg).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for ReplChannel {
|
impl Default for ReplChannel {
|
||||||
@@ -479,7 +492,9 @@ impl Channel for ReplChannel {
|
|||||||
|
|
||||||
async fn start(&self) -> Result<MessageStream, ChannelError> {
|
async fn start(&self) -> Result<MessageStream, ChannelError> {
|
||||||
let (tx, rx) = mpsc::channel(32);
|
let (tx, rx) = mpsc::channel(32);
|
||||||
// Store tx so send_status can inject approval responses directly
|
// Approval prompts inject responses back through this sender.
|
||||||
|
// In single-message mode we keep it until the turn finishes, then
|
||||||
|
// drop it after enqueuing /quit so the receiver stream can close.
|
||||||
if let Ok(mut guard) = self.msg_tx.lock() {
|
if let Ok(mut guard) = self.msg_tx.lock() {
|
||||||
*guard = Some(tx.clone());
|
*guard = Some(tx.clone());
|
||||||
}
|
}
|
||||||
@@ -495,11 +510,10 @@ impl Channel for ReplChannel {
|
|||||||
|
|
||||||
// Single message mode: send it and return
|
// Single message mode: send it and return
|
||||||
if let Some(msg) = single_message {
|
if let Some(msg) = single_message {
|
||||||
let incoming = IncomingMessage::new("repl", &user_id, &msg).with_timezone(&sys_tz);
|
let incoming = IncomingMessage::new("repl", &user_id, &msg)
|
||||||
|
.with_metadata(serde_json::json!({ "single_message_mode": true }))
|
||||||
|
.with_timezone(&sys_tz);
|
||||||
let _ = tx.blocking_send(incoming);
|
let _ = tx.blocking_send(incoming);
|
||||||
// Ensure the agent exits after handling exactly one turn in -m mode,
|
|
||||||
// even when other channels (gateway/http) are enabled.
|
|
||||||
let _ = tx.blocking_send(IncomingMessage::new("repl", &user_id, "/quit"));
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -662,6 +676,7 @@ impl Channel for ReplChannel {
|
|||||||
println!();
|
println!();
|
||||||
println!();
|
println!();
|
||||||
self.stdin_locked.store(false, Ordering::Relaxed);
|
self.stdin_locked.store(false, Ordering::Relaxed);
|
||||||
|
self.finish_single_message_turn().await;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -680,6 +695,7 @@ impl Channel for ReplChannel {
|
|||||||
println!();
|
println!();
|
||||||
// Unlock stdin so readline can resume
|
// Unlock stdin so readline can resume
|
||||||
self.stdin_locked.store(false, Ordering::Relaxed);
|
self.stdin_locked.store(false, Ordering::Relaxed);
|
||||||
|
self.finish_single_message_turn().await;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -779,6 +795,7 @@ impl Channel for ReplChannel {
|
|||||||
let msg_tx = Arc::clone(&self.msg_tx);
|
let msg_tx = Arc::clone(&self.msg_tx);
|
||||||
let user_id = self.user_id.clone();
|
let user_id = self.user_id.clone();
|
||||||
let lock_flag = Arc::clone(&self.stdin_locked);
|
let lock_flag = Arc::clone(&self.stdin_locked);
|
||||||
|
let single_message_mode = self.single_message.is_some();
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
let action = run_approval_selector(allow_always).unwrap_or("n");
|
let action = run_approval_selector(allow_always).unwrap_or("n");
|
||||||
// Unlock stdin so readline can resume after approval
|
// Unlock stdin so readline can resume after approval
|
||||||
@@ -787,7 +804,12 @@ impl Channel for ReplChannel {
|
|||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
if let Some(tx) = guard.as_ref() {
|
if let Some(tx) = guard.as_ref() {
|
||||||
let msg = IncomingMessage::new("repl", &user_id, action);
|
let msg = if single_message_mode {
|
||||||
|
IncomingMessage::new("repl", &user_id, action)
|
||||||
|
.with_metadata(serde_json::json!({ "single_message_mode": true }))
|
||||||
|
} else {
|
||||||
|
IncomingMessage::new("repl", &user_id, action)
|
||||||
|
};
|
||||||
let _ = tx.blocking_send(msg);
|
let _ = tx.blocking_send(msg);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
@@ -841,6 +863,19 @@ impl Channel for ReplChannel {
|
|||||||
StatusUpdate::Suggestions { .. } => {
|
StatusUpdate::Suggestions { .. } => {
|
||||||
// Suggestions are only rendered by the web gateway
|
// Suggestions are only rendered by the web gateway
|
||||||
}
|
}
|
||||||
|
StatusUpdate::ReasoningUpdate {
|
||||||
|
narrative,
|
||||||
|
decisions,
|
||||||
|
} => {
|
||||||
|
if !narrative.is_empty() {
|
||||||
|
let display = truncate_for_preview(&narrative, CLI_STATUS_MAX);
|
||||||
|
eprintln!(" \x1b[94m\u{25B6} {display}\x1b[0m");
|
||||||
|
}
|
||||||
|
for d in &decisions {
|
||||||
|
let display = truncate_for_preview(&d.rationale, CLI_STATUS_MAX);
|
||||||
|
eprintln!(" \x1b[90m\u{2192} {}: {display}\x1b[0m", d.tool_name);
|
||||||
|
}
|
||||||
|
}
|
||||||
StatusUpdate::TurnCost { .. } => {
|
StatusUpdate::TurnCost { .. } => {
|
||||||
// Cost display is handled by the TUI channel
|
// Cost display is handled by the TUI channel
|
||||||
}
|
}
|
||||||
@@ -875,6 +910,7 @@ impl Channel for ReplChannel {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
|
use tokio::time::{Duration, timeout};
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
@@ -883,16 +919,36 @@ mod tests {
|
|||||||
let repl = ReplChannel::with_message("hi".to_string());
|
let repl = ReplChannel::with_message("hi".to_string());
|
||||||
let mut stream = repl.start().await.expect("repl start should succeed");
|
let mut stream = repl.start().await.expect("repl start should succeed");
|
||||||
|
|
||||||
let first = stream.next().await.expect("first message missing");
|
let first = timeout(Duration::from_secs(1), stream.next())
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for first message")
|
||||||
|
.expect("first message missing");
|
||||||
assert_eq!(first.channel, "repl");
|
assert_eq!(first.channel, "repl");
|
||||||
assert_eq!(first.content, "hi");
|
assert_eq!(first.content, "hi");
|
||||||
|
|
||||||
let second = stream.next().await.expect("quit message missing");
|
assert!(
|
||||||
|
timeout(Duration::from_millis(100), stream.next())
|
||||||
|
.await
|
||||||
|
.is_err(),
|
||||||
|
"single-message mode should wait for the turn to finish before quitting"
|
||||||
|
);
|
||||||
|
|
||||||
|
repl.respond(&first, OutgoingResponse::text("done"))
|
||||||
|
.await
|
||||||
|
.expect("respond should succeed");
|
||||||
|
|
||||||
|
let second = timeout(Duration::from_secs(1), stream.next())
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for quit message")
|
||||||
|
.expect("quit message missing");
|
||||||
assert_eq!(second.channel, "repl");
|
assert_eq!(second.channel, "repl");
|
||||||
assert_eq!(second.content, "/quit");
|
assert_eq!(second.content, "/quit");
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
stream.next().await.is_none(),
|
timeout(Duration::from_secs(1), stream.next())
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for stream to close")
|
||||||
|
.is_none(),
|
||||||
"stream should end after /quit"
|
"stream should end after /quit"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -333,6 +333,9 @@ async fn webhook_handler(
|
|||||||
|
|
||||||
let channel_name = channel.channel_name();
|
let channel_name = channel.channel_name();
|
||||||
|
|
||||||
|
// Track whether any authentication was performed and passed.
|
||||||
|
let mut did_authenticate = false;
|
||||||
|
|
||||||
// Check if secret is required
|
// Check if secret is required
|
||||||
if state.router.requires_secret(channel_name).await {
|
if state.router.requires_secret(channel_name).await {
|
||||||
// Get the secret header name for this channel (from capabilities or default)
|
// Get the secret header name for this channel (from capabilities or default)
|
||||||
@@ -382,6 +385,7 @@ async fn webhook_handler(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
tracing::debug!(channel = %channel_name, "Webhook secret validated");
|
tracing::debug!(channel = %channel_name, "Webhook secret validated");
|
||||||
|
did_authenticate = true;
|
||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -433,6 +437,7 @@ async fn webhook_handler(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
tracing::debug!(channel = %channel_name, "Ed25519 signature verified");
|
tracing::debug!(channel = %channel_name, "Ed25519 signature verified");
|
||||||
|
did_authenticate = true;
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -484,6 +489,7 @@ async fn webhook_handler(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
tracing::debug!(channel = %channel_name, "HMAC-SHA256 signature verified");
|
tracing::debug!(channel = %channel_name, "HMAC-SHA256 signature verified");
|
||||||
|
did_authenticate = true;
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -510,8 +516,9 @@ async fn webhook_handler(
|
|||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
// Call the WASM channel
|
// Call the WASM channel. `did_authenticate` was set above by whichever
|
||||||
let secret_validated = state.router.requires_secret(channel_name).await;
|
// auth guard (secret / Ed25519 / HMAC) successfully validated the request.
|
||||||
|
let secret_validated = did_authenticate;
|
||||||
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
channel = %channel_name,
|
channel = %channel_name,
|
||||||
|
|||||||
@@ -3061,6 +3061,20 @@ fn status_to_wit(
|
|||||||
},
|
},
|
||||||
// Suggestions and turn cost are web-gateway-only; skip for WASM channels
|
// Suggestions and turn cost are web-gateway-only; skip for WASM channels
|
||||||
StatusUpdate::Suggestions { .. } | StatusUpdate::TurnCost { .. } => return None,
|
StatusUpdate::Suggestions { .. } | StatusUpdate::TurnCost { .. } => return None,
|
||||||
|
StatusUpdate::ReasoningUpdate {
|
||||||
|
narrative,
|
||||||
|
decisions,
|
||||||
|
} => {
|
||||||
|
let mut msg = narrative.clone();
|
||||||
|
for d in decisions {
|
||||||
|
msg.push_str(&format!("\n → {}: {}", d.tool_name, d.rationale));
|
||||||
|
}
|
||||||
|
wit_channel::StatusUpdate {
|
||||||
|
status: wit_channel::StatusType::Status,
|
||||||
|
message: msg,
|
||||||
|
metadata_json,
|
||||||
|
}
|
||||||
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+383
-22
@@ -1,17 +1,133 @@
|
|||||||
//! Bearer token authentication middleware for the web gateway.
|
//! Bearer token authentication middleware for the web gateway.
|
||||||
|
//!
|
||||||
|
//! Supports multi-user mode: each token maps to a `UserIdentity` that carries
|
||||||
|
//! the user_id. The identity is inserted into request extensions so downstream
|
||||||
|
//! handlers can extract it via `AuthenticatedUser`.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{Request, State},
|
extract::{FromRequestParts, Request, State},
|
||||||
http::{HeaderMap, Method, StatusCode},
|
http::{HeaderMap, Method, StatusCode, request::Parts},
|
||||||
middleware::Next,
|
middleware::Next,
|
||||||
response::{IntoResponse, Response},
|
response::{IntoResponse, Response},
|
||||||
};
|
};
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
use subtle::ConstantTimeEq;
|
use subtle::ConstantTimeEq;
|
||||||
|
|
||||||
/// Shared auth state injected via axum middleware state.
|
/// Identity resolved from a bearer token.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct UserIdentity {
|
||||||
|
pub user_id: String,
|
||||||
|
/// Additional user scopes this identity can read from.
|
||||||
|
pub workspace_read_scopes: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hash a token with SHA-256 for constant-size, timing-safe storage.
|
||||||
|
fn hash_token(token: &str) -> [u8; 32] {
|
||||||
|
let mut hasher = Sha256::new();
|
||||||
|
hasher.update(token.as_bytes());
|
||||||
|
hasher.finalize().into()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Multi-user auth state: maps token hashes to user identities.
|
||||||
|
///
|
||||||
|
/// Tokens are SHA-256 hashed on construction so they are never stored in
|
||||||
|
/// plaintext. Authentication compares fixed-size (32-byte) digests using
|
||||||
|
/// constant-time comparison, eliminating both length-oracle timing leaks
|
||||||
|
/// and accidental token exposure in memory dumps.
|
||||||
|
///
|
||||||
|
/// In single-user mode (the default), contains exactly one entry.
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct AuthState {
|
pub struct MultiAuthState {
|
||||||
pub token: String,
|
/// Maps SHA-256(token) → identity. Tokens are never stored in cleartext.
|
||||||
|
hashed_tokens: Vec<([u8; 32], UserIdentity)>,
|
||||||
|
/// Original first token kept only for single-user startup printing.
|
||||||
|
/// Not used for authentication.
|
||||||
|
display_token: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MultiAuthState {
|
||||||
|
/// Create a single-user auth state (backwards compatible).
|
||||||
|
pub fn single(token: String, user_id: String) -> Self {
|
||||||
|
let hash = hash_token(&token);
|
||||||
|
Self {
|
||||||
|
hashed_tokens: vec![(
|
||||||
|
hash,
|
||||||
|
UserIdentity {
|
||||||
|
user_id,
|
||||||
|
workspace_read_scopes: Vec::new(),
|
||||||
|
},
|
||||||
|
)],
|
||||||
|
display_token: Some(token),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a multi-user auth state from a map of tokens to identities.
|
||||||
|
pub fn multi(tokens: HashMap<String, UserIdentity>) -> Self {
|
||||||
|
let hashed_tokens: Vec<([u8; 32], UserIdentity)> = tokens
|
||||||
|
.into_iter()
|
||||||
|
.map(|(tok, identity)| (hash_token(&tok), identity))
|
||||||
|
.collect();
|
||||||
|
Self {
|
||||||
|
hashed_tokens,
|
||||||
|
display_token: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Authenticate a token, returning the associated identity if valid.
|
||||||
|
///
|
||||||
|
/// Uses SHA-256 hashing + constant-time comparison (`subtle::ConstantTimeEq`)
|
||||||
|
/// to prevent timing side-channels. Both the candidate and stored tokens are
|
||||||
|
/// hashed to 32-byte digests, eliminating length-oracle leaks. Iterates all
|
||||||
|
/// entries regardless of match to avoid early-exit timing differences.
|
||||||
|
/// O(n) in the number of configured users — negligible for typical
|
||||||
|
/// deployments (< 10 users).
|
||||||
|
pub fn authenticate(&self, candidate: &str) -> Option<&UserIdentity> {
|
||||||
|
let candidate_hash = hash_token(candidate);
|
||||||
|
let mut matched: Option<&UserIdentity> = None;
|
||||||
|
for (stored_hash, identity) in &self.hashed_tokens {
|
||||||
|
if bool::from(candidate_hash.ct_eq(stored_hash)) {
|
||||||
|
matched = Some(identity);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
matched
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the first token for backwards-compatible printing at startup.
|
||||||
|
///
|
||||||
|
/// Only available in single-user mode; returns `None` in multi-user mode
|
||||||
|
/// to avoid exposing tokens.
|
||||||
|
pub fn first_token(&self) -> Option<&str> {
|
||||||
|
self.display_token.as_deref()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the first user identity (for single-user fallback).
|
||||||
|
pub fn first_identity(&self) -> Option<&UserIdentity> {
|
||||||
|
self.hashed_tokens.first().map(|(_, id)| id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Axum extractor that provides the authenticated user identity.
|
||||||
|
///
|
||||||
|
/// Only available on routes behind `auth_middleware`. Extracts the
|
||||||
|
/// `UserIdentity` that the middleware inserted into request extensions.
|
||||||
|
pub struct AuthenticatedUser(pub UserIdentity);
|
||||||
|
|
||||||
|
impl<S> FromRequestParts<S> for AuthenticatedUser
|
||||||
|
where
|
||||||
|
S: Send + Sync,
|
||||||
|
{
|
||||||
|
type Rejection = (StatusCode, &'static str);
|
||||||
|
|
||||||
|
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
|
||||||
|
parts
|
||||||
|
.extensions
|
||||||
|
.get::<UserIdentity>()
|
||||||
|
.cloned()
|
||||||
|
.map(AuthenticatedUser)
|
||||||
|
.ok_or((StatusCode::UNAUTHORIZED, "Not authenticated"))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Whether query-string token auth is allowed for this request.
|
/// Whether query-string token auth is allowed for this request.
|
||||||
@@ -51,29 +167,34 @@ fn query_token(request: &Request) -> Option<String> {
|
|||||||
/// Auth middleware that validates bearer token from header or query param.
|
/// Auth middleware that validates bearer token from header or query param.
|
||||||
///
|
///
|
||||||
/// SSE connections can't set headers from `EventSource`, so we also accept
|
/// SSE connections can't set headers from `EventSource`, so we also accept
|
||||||
/// `?token=xxx` as a query parameter, but only on SSE endpoints.
|
/// `?token=xxx` as a query parameter, but only on SSE/WS endpoints.
|
||||||
|
///
|
||||||
|
/// On successful authentication, inserts the matching `UserIdentity` into
|
||||||
|
/// request extensions for downstream extraction via `AuthenticatedUser`.
|
||||||
pub async fn auth_middleware(
|
pub async fn auth_middleware(
|
||||||
State(auth): State<AuthState>,
|
State(auth): State<MultiAuthState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
request: Request,
|
mut request: Request,
|
||||||
next: Next,
|
next: Next,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
// Try Authorization header first (constant-time comparison).
|
// Try Authorization header first.
|
||||||
// RFC 6750 Section 2.1: auth-scheme comparison is case-insensitive.
|
// RFC 6750 Section 2.1: auth-scheme comparison is case-insensitive.
|
||||||
if let Some(auth_header) = headers.get("authorization")
|
if let Some(auth_header) = headers.get("authorization")
|
||||||
&& let Ok(value) = auth_header.to_str()
|
&& let Ok(value) = auth_header.to_str()
|
||||||
&& value.len() > 7
|
&& value.len() > 7
|
||||||
&& value[..7].eq_ignore_ascii_case("Bearer ")
|
&& value[..7].eq_ignore_ascii_case("Bearer ")
|
||||||
&& bool::from(value.as_bytes()[7..].ct_eq(auth.token.as_bytes()))
|
&& let Some(identity) = auth.authenticate(&value[7..])
|
||||||
{
|
{
|
||||||
|
request.extensions_mut().insert(identity.clone());
|
||||||
return next.run(request).await;
|
return next.run(request).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fall back to query parameter, but only for SSE endpoints (constant-time comparison).
|
// Fall back to query parameter, but only for SSE/WS endpoints.
|
||||||
if allows_query_token_auth(&request)
|
if allows_query_token_auth(&request)
|
||||||
&& let Some(token) = query_token(&request)
|
&& let Some(token) = query_token(&request)
|
||||||
&& bool::from(token.as_bytes().ct_eq(auth.token.as_bytes()))
|
&& let Some(identity) = auth.authenticate(&token)
|
||||||
{
|
{
|
||||||
|
request.extensions_mut().insert(identity.clone());
|
||||||
return next.run(request).await;
|
return next.run(request).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -83,15 +204,61 @@ pub async fn auth_middleware(
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::testing::credentials::{TEST_AUTH_SECRET_TOKEN, TEST_BEARER_TOKEN};
|
use crate::testing::credentials::TEST_AUTH_SECRET_TOKEN;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_auth_state_clone() {
|
fn test_multi_auth_state_single() {
|
||||||
let state = AuthState {
|
let state = MultiAuthState::single("tok-123".to_string(), "alice".to_string());
|
||||||
token: TEST_BEARER_TOKEN.to_string(),
|
let identity = state.authenticate("tok-123");
|
||||||
};
|
assert!(identity.is_some());
|
||||||
let cloned = state.clone();
|
assert_eq!(identity.unwrap().user_id, "alice");
|
||||||
assert_eq!(cloned.token, TEST_BEARER_TOKEN);
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multi_auth_state_reject_wrong_token() {
|
||||||
|
let state = MultiAuthState::single("tok-123".to_string(), "alice".to_string());
|
||||||
|
assert!(state.authenticate("wrong-token").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multi_auth_state_multi_users() {
|
||||||
|
let mut tokens = HashMap::new();
|
||||||
|
tokens.insert(
|
||||||
|
"tok-alice".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: Vec::new(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
tokens.insert(
|
||||||
|
"tok-bob".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "bob".to_string(),
|
||||||
|
workspace_read_scopes: Vec::new(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
let state = MultiAuthState::multi(tokens);
|
||||||
|
|
||||||
|
let alice = state.authenticate("tok-alice").unwrap();
|
||||||
|
assert_eq!(alice.user_id, "alice");
|
||||||
|
|
||||||
|
let bob = state.authenticate("tok-bob").unwrap();
|
||||||
|
assert_eq!(bob.user_id, "bob");
|
||||||
|
|
||||||
|
assert!(state.authenticate("tok-charlie").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multi_auth_state_first_token() {
|
||||||
|
let state = MultiAuthState::single("my-token".to_string(), "user1".to_string());
|
||||||
|
assert_eq!(state.first_token(), Some("my-token"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multi_auth_state_first_identity() {
|
||||||
|
let state = MultiAuthState::single("my-token".to_string(), "user1".to_string());
|
||||||
|
let identity = state.first_identity().unwrap();
|
||||||
|
assert_eq!(identity.user_id, "user1");
|
||||||
}
|
}
|
||||||
|
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
@@ -107,9 +274,7 @@ mod tests {
|
|||||||
/// Router with streaming endpoints (query auth allowed) and regular
|
/// Router with streaming endpoints (query auth allowed) and regular
|
||||||
/// endpoints (query auth rejected).
|
/// endpoints (query auth rejected).
|
||||||
fn test_app(token: &str) -> Router {
|
fn test_app(token: &str) -> Router {
|
||||||
let state = AuthState {
|
let state = MultiAuthState::single(token.to_string(), "test-user".to_string());
|
||||||
token: token.to_string(),
|
|
||||||
};
|
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/chat/events", get(dummy_handler))
|
.route("/api/chat/events", get(dummy_handler))
|
||||||
.route("/api/logs/events", get(dummy_handler))
|
.route("/api/logs/events", get(dummy_handler))
|
||||||
@@ -306,4 +471,200 @@ mod tests {
|
|||||||
let resp = app.oneshot(req).await.unwrap();
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- Multi-tenant auth integration tests ---
|
||||||
|
|
||||||
|
/// Handler that extracts `AuthenticatedUser` and returns the resolved user_id.
|
||||||
|
async fn identity_handler(AuthenticatedUser(identity): AuthenticatedUser) -> String {
|
||||||
|
identity.user_id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handler that extracts `AuthenticatedUser` and returns workspace_read_scopes as JSON.
|
||||||
|
async fn scopes_handler(AuthenticatedUser(identity): AuthenticatedUser) -> String {
|
||||||
|
serde_json::to_string(&identity.workspace_read_scopes).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a multi-user router where each token maps to a distinct identity.
|
||||||
|
fn multi_user_app(tokens: HashMap<String, UserIdentity>) -> Router {
|
||||||
|
let state = MultiAuthState::multi(tokens);
|
||||||
|
Router::new()
|
||||||
|
.route("/api/chat/events", get(identity_handler))
|
||||||
|
.route("/api/chat/send", post(identity_handler))
|
||||||
|
.route("/api/scopes", get(scopes_handler))
|
||||||
|
.layer(middleware::from_fn_with_state(state, auth_middleware))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn two_user_tokens() -> HashMap<String, UserIdentity> {
|
||||||
|
let mut tokens = HashMap::new();
|
||||||
|
tokens.insert(
|
||||||
|
"tok-alice".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec!["shared".to_string()],
|
||||||
|
},
|
||||||
|
);
|
||||||
|
tokens.insert(
|
||||||
|
"tok-bob".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "bob".to_string(),
|
||||||
|
workspace_read_scopes: vec!["shared".to_string(), "alice".to_string()],
|
||||||
|
},
|
||||||
|
);
|
||||||
|
tokens
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_alice_token_resolves_to_alice() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "alice");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_bob_token_resolves_to_bob() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events")
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "bob");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_sequential_tokens_resolve_independently() {
|
||||||
|
// Send both alice and bob tokens sequentially and verify each gets
|
||||||
|
// the correct identity — guards against token map corruption.
|
||||||
|
let tokens = two_user_tokens();
|
||||||
|
|
||||||
|
let app1 = multi_user_app(tokens.clone());
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app1.oneshot(req).await.unwrap();
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "alice");
|
||||||
|
|
||||||
|
let app2 = multi_user_app(tokens);
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events")
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app2.oneshot(req).await.unwrap();
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "bob");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_unknown_token_rejected() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events")
|
||||||
|
.header("Authorization", "Bearer tok-charlie")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_workspace_read_scopes_propagated() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
|
||||||
|
// Alice has ["shared"]
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/scopes")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
let scopes: Vec<String> = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(scopes, vec!["shared"]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_bob_has_two_scopes() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
|
||||||
|
// Bob has ["shared", "alice"]
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/scopes")
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
let scopes: Vec<String> = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(scopes, vec!["shared", "alice"]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_query_param_resolves_correct_identity() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/chat/events?token=tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "bob");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_post_with_bearer_resolves_identity() {
|
||||||
|
let app = multi_user_app(two_user_tokens());
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri("/api/chat/send")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
assert_eq!(body, "alice");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multi_user_empty_scopes_for_single_user() {
|
||||||
|
// Single-user mode creates identity with empty workspace_read_scopes.
|
||||||
|
let state = MultiAuthState::single("tok-only".to_string(), "solo".to_string());
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/api/scopes", get(scopes_handler))
|
||||||
|
.layer(middleware::from_fn_with_state(state, auth_middleware));
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/scopes")
|
||||||
|
.header("Authorization", "Bearer tok-only")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
|
||||||
|
let scopes: Vec<String> = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert!(scopes.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_prefix_and_extension_tokens_rejected() {
|
||||||
|
// Verifies that prefix/suffix variants of valid tokens are rejected.
|
||||||
|
// Note: the constant-time property is enforced structurally by use of
|
||||||
|
// subtle::ConstantTimeEq and cannot be verified via outcome testing.
|
||||||
|
let state = MultiAuthState::single("long-secret-token".to_string(), "user".to_string());
|
||||||
|
assert!(state.authenticate("long-secret").is_none());
|
||||||
|
assert!(state.authenticate("long-secret-token-extra").is_none());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -12,22 +12,24 @@ use serde::Deserialize;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::IncomingMessage;
|
use crate::channels::IncomingMessage;
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
use crate::channels::web::util::{build_turns_from_db_messages, truncate_preview};
|
use crate::channels::web::util::{build_turns_from_db_messages, truncate_preview};
|
||||||
|
|
||||||
pub async fn chat_send_handler(
|
pub async fn chat_send_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
Json(req): Json<SendMessageRequest>,
|
Json(req): Json<SendMessageRequest>,
|
||||||
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
||||||
if !state.chat_rate_limiter.check() {
|
if !state.chat_rate_limiter.check(&identity.user_id) {
|
||||||
return Err((
|
return Err((
|
||||||
StatusCode::TOO_MANY_REQUESTS,
|
StatusCode::TOO_MANY_REQUESTS,
|
||||||
"Rate limit exceeded. Try again shortly.".to_string(),
|
"Rate limit exceeded. Try again shortly.".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content);
|
let mut msg = IncomingMessage::new("gateway", &identity.user_id, &req.content);
|
||||||
|
|
||||||
if let Some(ref thread_id) = req.thread_id {
|
if let Some(ref thread_id) = req.thread_id {
|
||||||
msg = msg.with_thread(thread_id);
|
msg = msg.with_thread(thread_id);
|
||||||
@@ -74,6 +76,7 @@ pub async fn chat_send_handler(
|
|||||||
|
|
||||||
pub async fn chat_approval_handler(
|
pub async fn chat_approval_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
Json(req): Json<ApprovalRequest>,
|
Json(req): Json<ApprovalRequest>,
|
||||||
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
||||||
let (approved, always) = match req.action.as_str() {
|
let (approved, always) = match req.action.as_str() {
|
||||||
@@ -109,7 +112,7 @@ pub async fn chat_approval_handler(
|
|||||||
)
|
)
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let mut msg = IncomingMessage::new("gateway", &state.user_id, content);
|
let mut msg = IncomingMessage::new("gateway", &identity.user_id, content);
|
||||||
|
|
||||||
if let Some(ref thread_id) = req.thread_id {
|
if let Some(ref thread_id) = req.thread_id {
|
||||||
msg = msg.with_thread(thread_id);
|
msg = msg.with_thread(thread_id);
|
||||||
@@ -150,6 +153,7 @@ pub async fn chat_approval_handler(
|
|||||||
/// The token never touches the LLM, chat history, or SSE stream.
|
/// The token never touches the LLM, chat history, or SSE stream.
|
||||||
pub async fn chat_auth_token_handler(
|
pub async fn chat_auth_token_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Json(req): Json<AuthTokenRequest>,
|
Json(req): Json<AuthTokenRequest>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
||||||
@@ -158,7 +162,7 @@ pub async fn chat_auth_token_handler(
|
|||||||
))?;
|
))?;
|
||||||
|
|
||||||
match ext_mgr
|
match ext_mgr
|
||||||
.configure_token(&req.extension_name, &req.token)
|
.configure_token(&req.extension_name, &req.token, &user.user_id)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
@@ -169,20 +173,26 @@ pub async fn chat_auth_token_handler(
|
|||||||
resp.instructions = result.verification.as_ref().map(|v| v.instructions.clone());
|
resp.instructions = result.verification.as_ref().map(|v| v.instructions.clone());
|
||||||
|
|
||||||
if result.verification.is_some() {
|
if result.verification.is_some() {
|
||||||
state.sse.broadcast(SseEvent::AuthRequired {
|
state.sse.broadcast_for_user(
|
||||||
extension_name: req.extension_name.clone(),
|
&user.user_id,
|
||||||
instructions: Some(result.message),
|
AppEvent::AuthRequired {
|
||||||
auth_url: None,
|
extension_name: req.extension_name.clone(),
|
||||||
setup_url: None,
|
instructions: Some(result.message),
|
||||||
});
|
auth_url: None,
|
||||||
|
setup_url: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
} else {
|
} else {
|
||||||
clear_auth_mode(&state).await;
|
clear_auth_mode(&state, &user.user_id).await;
|
||||||
|
|
||||||
state.sse.broadcast(SseEvent::AuthCompleted {
|
state.sse.broadcast_for_user(
|
||||||
extension_name: req.extension_name.clone(),
|
&user.user_id,
|
||||||
success: true,
|
AppEvent::AuthCompleted {
|
||||||
message: result.message,
|
extension_name: req.extension_name.clone(),
|
||||||
});
|
success: true,
|
||||||
|
message: result.message,
|
||||||
|
},
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(Json(resp))
|
Ok(Json(resp))
|
||||||
@@ -190,12 +200,15 @@ pub async fn chat_auth_token_handler(
|
|||||||
Err(e) => {
|
Err(e) => {
|
||||||
let msg = e.to_string();
|
let msg = e.to_string();
|
||||||
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
||||||
state.sse.broadcast(SseEvent::AuthRequired {
|
state.sse.broadcast_for_user(
|
||||||
extension_name: req.extension_name.clone(),
|
&user.user_id,
|
||||||
instructions: Some(msg.clone()),
|
AppEvent::AuthRequired {
|
||||||
auth_url: None,
|
extension_name: req.extension_name.clone(),
|
||||||
setup_url: None,
|
instructions: Some(msg.clone()),
|
||||||
});
|
auth_url: None,
|
||||||
|
setup_url: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
}
|
}
|
||||||
Ok(Json(ActionResponse::fail(msg)))
|
Ok(Json(ActionResponse::fail(msg)))
|
||||||
}
|
}
|
||||||
@@ -205,16 +218,17 @@ pub async fn chat_auth_token_handler(
|
|||||||
/// Cancel an in-progress auth flow.
|
/// Cancel an in-progress auth flow.
|
||||||
pub async fn chat_auth_cancel_handler(
|
pub async fn chat_auth_cancel_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
Json(_req): Json<AuthCancelRequest>,
|
Json(_req): Json<AuthCancelRequest>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
clear_auth_mode(&state).await;
|
clear_auth_mode(&state, &identity.user_id).await;
|
||||||
Ok(Json(ActionResponse::ok("Auth cancelled")))
|
Ok(Json(ActionResponse::ok("Auth cancelled")))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Clear pending auth mode on the active thread.
|
/// Clear pending auth mode on the active thread.
|
||||||
pub async fn clear_auth_mode(state: &GatewayState) {
|
pub async fn clear_auth_mode(state: &GatewayState, user_id: &str) {
|
||||||
if let Some(ref sm) = state.session_manager {
|
if let Some(ref sm) = state.session_manager {
|
||||||
let session = sm.get_or_create_session(&state.user_id).await;
|
let session = sm.get_or_create_session(user_id).await;
|
||||||
let mut sess = session.lock().await;
|
let mut sess = session.lock().await;
|
||||||
if let Some(thread_id) = sess.active_thread
|
if let Some(thread_id) = sess.active_thread
|
||||||
&& let Some(thread) = sess.threads.get_mut(&thread_id)
|
&& let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||||
@@ -226,8 +240,9 @@ pub async fn clear_auth_mode(state: &GatewayState) {
|
|||||||
|
|
||||||
pub async fn chat_events_handler(
|
pub async fn chat_events_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||||
state.sse.subscribe().ok_or((
|
state.sse.subscribe(Some(user.user_id)).ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
"Too many connections".to_string(),
|
"Too many connections".to_string(),
|
||||||
))
|
))
|
||||||
@@ -237,6 +252,7 @@ pub async fn chat_ws_handler(
|
|||||||
headers: axum::http::HeaderMap,
|
headers: axum::http::HeaderMap,
|
||||||
ws: WebSocketUpgrade,
|
ws: WebSocketUpgrade,
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||||
// Validate Origin header to prevent cross-site WebSocket hijacking.
|
// Validate Origin header to prevent cross-site WebSocket hijacking.
|
||||||
let origin = headers
|
let origin = headers
|
||||||
@@ -262,7 +278,9 @@ pub async fn chat_ws_handler(
|
|||||||
"WebSocket origin not allowed".to_string(),
|
"WebSocket origin not allowed".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
Ok(ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state)))
|
Ok(ws.on_upgrade(move |socket| {
|
||||||
|
crate::channels::web::ws::handle_ws_connection(socket, state, identity)
|
||||||
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
@@ -274,6 +292,7 @@ pub struct HistoryQuery {
|
|||||||
|
|
||||||
pub async fn chat_history_handler(
|
pub async fn chat_history_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
Query(query): Query<HistoryQuery>,
|
Query(query): Query<HistoryQuery>,
|
||||||
) -> Result<Json<HistoryResponse>, (StatusCode, String)> {
|
) -> Result<Json<HistoryResponse>, (StatusCode, String)> {
|
||||||
let session_manager = state.session_manager.as_ref().ok_or((
|
let session_manager = state.session_manager.as_ref().ok_or((
|
||||||
@@ -281,7 +300,9 @@ pub async fn chat_history_handler(
|
|||||||
"Session manager not available".to_string(),
|
"Session manager not available".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
let session = session_manager.get_or_create_session(&state.user_id).await;
|
let session = session_manager
|
||||||
|
.get_or_create_session(&identity.user_id)
|
||||||
|
.await;
|
||||||
|
|
||||||
let limit = query.limit.unwrap_or(50);
|
let limit = query.limit.unwrap_or(50);
|
||||||
let before_cursor = query
|
let before_cursor = query
|
||||||
@@ -314,7 +335,7 @@ pub async fn chat_history_handler(
|
|||||||
&& let Some(ref store) = state.store
|
&& let Some(ref store) = state.store
|
||||||
{
|
{
|
||||||
let owned = store
|
let owned = store
|
||||||
.conversation_belongs_to_user(thread_id, &state.user_id)
|
.conversation_belongs_to_user(thread_id, &identity.user_id)
|
||||||
.await
|
.await
|
||||||
.unwrap_or(false);
|
.unwrap_or(false);
|
||||||
if !owned {
|
if !owned {
|
||||||
@@ -377,8 +398,10 @@ pub async fn chat_history_handler(
|
|||||||
truncate_preview(&s, 500)
|
truncate_preview(&s, 500)
|
||||||
}),
|
}),
|
||||||
error: tc.error.clone(),
|
error: tc.error.clone(),
|
||||||
|
rationale: tc.rationale.clone(),
|
||||||
})
|
})
|
||||||
.collect(),
|
.collect(),
|
||||||
|
narrative: t.narrative.clone(),
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -434,24 +457,27 @@ pub async fn chat_history_handler(
|
|||||||
|
|
||||||
pub async fn chat_threads_handler(
|
pub async fn chat_threads_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
) -> Result<Json<ThreadListResponse>, (StatusCode, String)> {
|
) -> Result<Json<ThreadListResponse>, (StatusCode, String)> {
|
||||||
let session_manager = state.session_manager.as_ref().ok_or((
|
let session_manager = state.session_manager.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
"Session manager not available".to_string(),
|
"Session manager not available".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
let session = session_manager.get_or_create_session(&state.user_id).await;
|
let session = session_manager
|
||||||
|
.get_or_create_session(&identity.user_id)
|
||||||
|
.await;
|
||||||
|
|
||||||
// Try DB first for persistent thread list
|
// Try DB first for persistent thread list
|
||||||
if let Some(ref store) = state.store {
|
if let Some(ref store) = state.store {
|
||||||
// Auto-create assistant thread if it doesn't exist
|
// Auto-create assistant thread if it doesn't exist
|
||||||
let assistant_id = store
|
let assistant_id = store
|
||||||
.get_or_create_assistant_conversation(&state.user_id, "gateway")
|
.get_or_create_assistant_conversation(&identity.user_id, "gateway")
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
if let Ok(summaries) = store
|
if let Ok(summaries) = store
|
||||||
.list_conversations_all_channels(&state.user_id, 50)
|
.list_conversations_all_channels(&identity.user_id, 50)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
let mut assistant_thread = None;
|
let mut assistant_thread = None;
|
||||||
@@ -534,13 +560,16 @@ pub async fn chat_threads_handler(
|
|||||||
|
|
||||||
pub async fn chat_new_thread_handler(
|
pub async fn chat_new_thread_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(identity): AuthenticatedUser,
|
||||||
) -> Result<Json<ThreadInfo>, (StatusCode, String)> {
|
) -> Result<Json<ThreadInfo>, (StatusCode, String)> {
|
||||||
let session_manager = state.session_manager.as_ref().ok_or((
|
let session_manager = state.session_manager.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
"Session manager not available".to_string(),
|
"Session manager not available".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
let session = session_manager.get_or_create_session(&state.user_id).await;
|
let session = session_manager
|
||||||
|
.get_or_create_session(&identity.user_id)
|
||||||
|
.await;
|
||||||
let (thread_id, info) = {
|
let (thread_id, info) = {
|
||||||
let mut sess = session.lock().await;
|
let mut sess = session.lock().await;
|
||||||
let thread = sess.create_thread();
|
let thread = sess.create_thread();
|
||||||
@@ -562,12 +591,12 @@ pub async fn chat_new_thread_handler(
|
|||||||
// so that the subsequent loadThreads() call from the frontend sees it.
|
// so that the subsequent loadThreads() call from the frontend sees it.
|
||||||
if let Some(ref store) = state.store {
|
if let Some(ref store) = state.store {
|
||||||
match store
|
match store
|
||||||
.ensure_conversation(thread_id, "gateway", &state.user_id, None)
|
.ensure_conversation(thread_id, "gateway", &identity.user_id, None)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(true) => {}
|
Ok(true) => {}
|
||||||
Ok(false) => tracing::warn!(
|
Ok(false) => tracing::warn!(
|
||||||
user = %state.user_id,
|
user = %identity.user_id,
|
||||||
thread_id = %thread_id,
|
thread_id = %thread_id,
|
||||||
"Skipped persisting new thread due to ownership/channel conflict"
|
"Skipped persisting new thread due to ownership/channel conflict"
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -8,11 +8,13 @@ use axum::{
|
|||||||
http::StatusCode,
|
http::StatusCode,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
|
||||||
pub async fn extensions_list_handler(
|
pub async fn extensions_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<ExtensionListResponse>, (StatusCode, String)> {
|
) -> Result<Json<ExtensionListResponse>, (StatusCode, String)> {
|
||||||
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
||||||
StatusCode::NOT_IMPLEMENTED,
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
@@ -20,7 +22,7 @@ pub async fn extensions_list_handler(
|
|||||||
))?;
|
))?;
|
||||||
|
|
||||||
let installed = ext_mgr
|
let installed = ext_mgr
|
||||||
.list(None, false)
|
.list(None, false, &user.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
@@ -80,6 +82,7 @@ pub async fn extensions_list_handler(
|
|||||||
|
|
||||||
pub async fn extensions_tools_handler(
|
pub async fn extensions_tools_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(_user): AuthenticatedUser,
|
||||||
) -> Result<Json<ToolListResponse>, (StatusCode, String)> {
|
) -> Result<Json<ToolListResponse>, (StatusCode, String)> {
|
||||||
let registry = state.tool_registry.as_ref().ok_or((
|
let registry = state.tool_registry.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
@@ -100,6 +103,7 @@ pub async fn extensions_tools_handler(
|
|||||||
|
|
||||||
pub async fn extensions_install_handler(
|
pub async fn extensions_install_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Json(req): Json<InstallExtensionRequest>,
|
Json(req): Json<InstallExtensionRequest>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
||||||
@@ -116,7 +120,7 @@ pub async fn extensions_install_handler(
|
|||||||
});
|
});
|
||||||
|
|
||||||
match ext_mgr
|
match ext_mgr
|
||||||
.install(&req.name, req.url.as_deref(), kind_hint)
|
.install(&req.name, req.url.as_deref(), kind_hint, &user.user_id)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(result) => Ok(Json(ActionResponse::ok(result.message))),
|
Ok(result) => Ok(Json(ActionResponse::ok(result.message))),
|
||||||
@@ -126,6 +130,7 @@ pub async fn extensions_install_handler(
|
|||||||
|
|
||||||
pub async fn extensions_remove_handler(
|
pub async fn extensions_remove_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(name): Path<String>,
|
Path(name): Path<String>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
let ext_mgr = state.extension_manager.as_ref().ok_or((
|
||||||
@@ -133,7 +138,7 @@ pub async fn extensions_remove_handler(
|
|||||||
"Extension manager not available (secrets store required)".to_string(),
|
"Extension manager not available (secrets store required)".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
match ext_mgr.remove(&name).await {
|
match ext_mgr.remove(&name, &user.user_id).await {
|
||||||
Ok(message) => Ok(Json(ActionResponse::ok(message))),
|
Ok(message) => Ok(Json(ActionResponse::ok(message))),
|
||||||
Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))),
|
Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))),
|
||||||
}
|
}
|
||||||
|
|||||||
+400
-277
@@ -11,11 +11,13 @@ use axum::{
|
|||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
|
||||||
pub async fn jobs_list_handler(
|
pub async fn jobs_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<JobListResponse>, (StatusCode, String)> {
|
) -> Result<Json<JobListResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
@@ -25,8 +27,8 @@ pub async fn jobs_list_handler(
|
|||||||
let mut jobs: Vec<JobInfo> = Vec::new();
|
let mut jobs: Vec<JobInfo> = Vec::new();
|
||||||
let mut seen_ids: HashSet<Uuid> = HashSet::new();
|
let mut seen_ids: HashSet<Uuid> = HashSet::new();
|
||||||
|
|
||||||
// Fetch sandbox jobs from database.
|
// Fetch sandbox jobs scoped to this user.
|
||||||
match store.list_sandbox_jobs().await {
|
match store.list_sandbox_jobs_for_user(&user.user_id).await {
|
||||||
Ok(sandbox_jobs) => {
|
Ok(sandbox_jobs) => {
|
||||||
for j in &sandbox_jobs {
|
for j in &sandbox_jobs {
|
||||||
let ui_state = match j.status.as_str() {
|
let ui_state = match j.status.as_str() {
|
||||||
@@ -50,8 +52,8 @@ pub async fn jobs_list_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fetch agent (non-sandbox) jobs from database, deduplicating by ID.
|
// Fetch agent (non-sandbox) jobs scoped to this user, deduplicating by ID.
|
||||||
match store.list_agent_jobs().await {
|
match store.list_agent_jobs_for_user(&user.user_id).await {
|
||||||
Ok(agent_jobs) => {
|
Ok(agent_jobs) => {
|
||||||
for j in &agent_jobs {
|
for j in &agent_jobs {
|
||||||
if seen_ids.contains(&j.id) {
|
if seen_ids.contains(&j.id) {
|
||||||
@@ -80,6 +82,7 @@ pub async fn jobs_list_handler(
|
|||||||
|
|
||||||
pub async fn jobs_summary_handler(
|
pub async fn jobs_summary_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<JobSummaryResponse>, (StatusCode, String)> {
|
) -> Result<Json<JobSummaryResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
@@ -93,8 +96,8 @@ pub async fn jobs_summary_handler(
|
|||||||
let mut failed = 0;
|
let mut failed = 0;
|
||||||
let mut stuck = 0;
|
let mut stuck = 0;
|
||||||
|
|
||||||
// Sandbox job counts.
|
// Sandbox job counts scoped to this user.
|
||||||
match store.sandbox_job_summary().await {
|
match store.sandbox_job_summary_for_user(&user.user_id).await {
|
||||||
Ok(s) => {
|
Ok(s) => {
|
||||||
total += s.total;
|
total += s.total;
|
||||||
pending += s.creating;
|
pending += s.creating;
|
||||||
@@ -107,8 +110,8 @@ pub async fn jobs_summary_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Agent job counts.
|
// Agent job counts scoped to this user.
|
||||||
match store.agent_job_summary().await {
|
match store.agent_job_summary_for_user(&user.user_id).await {
|
||||||
Ok(s) => {
|
Ok(s) => {
|
||||||
total += s.total;
|
total += s.total;
|
||||||
pending += s.pending;
|
pending += s.pending;
|
||||||
@@ -134,6 +137,7 @@ pub async fn jobs_summary_handler(
|
|||||||
|
|
||||||
pub async fn jobs_detail_handler(
|
pub async fn jobs_detail_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<JobDetailResponse>, (StatusCode, String)> {
|
) -> Result<Json<JobDetailResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -145,169 +149,213 @@ pub async fn jobs_detail_handler(
|
|||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||||
|
|
||||||
// Try sandbox job from DB first.
|
// Try sandbox job from DB first.
|
||||||
if let Ok(Some(job)) = store.get_sandbox_job(job_id).await {
|
match store.get_sandbox_job(job_id).await {
|
||||||
let browse_id = std::path::Path::new(&job.project_dir)
|
Ok(Some(job)) => {
|
||||||
.file_name()
|
if job.user_id != user.user_id {
|
||||||
.map(|n| n.to_string_lossy().to_string())
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
.unwrap_or_else(|| job.id.to_string());
|
}
|
||||||
|
let browse_id = std::path::Path::new(&job.project_dir)
|
||||||
|
.file_name()
|
||||||
|
.map(|n| n.to_string_lossy().to_string())
|
||||||
|
.unwrap_or_else(|| job.id.to_string());
|
||||||
|
|
||||||
let ui_state = match job.status.as_str() {
|
let ui_state = match job.status.as_str() {
|
||||||
"creating" => "pending",
|
"creating" => "pending",
|
||||||
"running" => "in_progress",
|
"running" => "in_progress",
|
||||||
s => s,
|
s => s,
|
||||||
};
|
};
|
||||||
|
|
||||||
let elapsed_secs = job.started_at.map(|start| {
|
let elapsed_secs = job.started_at.map(|start| {
|
||||||
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
|
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
|
||||||
(end - start).num_seconds().max(0) as u64
|
(end - start).num_seconds().max(0) as u64
|
||||||
});
|
|
||||||
|
|
||||||
// Synthesize transitions from timestamps.
|
|
||||||
let mut transitions = Vec::new();
|
|
||||||
if let Some(started) = job.started_at {
|
|
||||||
transitions.push(TransitionInfo {
|
|
||||||
from: "creating".to_string(),
|
|
||||||
to: "running".to_string(),
|
|
||||||
timestamp: started.to_rfc3339(),
|
|
||||||
reason: None,
|
|
||||||
});
|
});
|
||||||
}
|
|
||||||
if let Some(completed) = job.completed_at {
|
|
||||||
transitions.push(TransitionInfo {
|
|
||||||
from: "running".to_string(),
|
|
||||||
to: job.status.clone(),
|
|
||||||
timestamp: completed.to_rfc3339(),
|
|
||||||
reason: job.failure_reason.clone(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
|
// Synthesize transitions from timestamps.
|
||||||
let is_claude_code = mode.as_deref() == Some("claude_code");
|
let mut transitions = Vec::new();
|
||||||
|
if let Some(started) = job.started_at {
|
||||||
|
transitions.push(TransitionInfo {
|
||||||
|
from: "creating".to_string(),
|
||||||
|
to: "running".to_string(),
|
||||||
|
timestamp: started.to_rfc3339(),
|
||||||
|
reason: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if let Some(completed) = job.completed_at {
|
||||||
|
transitions.push(TransitionInfo {
|
||||||
|
from: "running".to_string(),
|
||||||
|
to: job.status.clone(),
|
||||||
|
timestamp: completed.to_rfc3339(),
|
||||||
|
reason: job.failure_reason.clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
return Ok(Json(JobDetailResponse {
|
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
|
||||||
id: job.id,
|
let is_claude_code = mode.as_deref() == Some("claude_code");
|
||||||
title: job.task.clone(),
|
|
||||||
description: String::new(),
|
return Ok(Json(JobDetailResponse {
|
||||||
state: ui_state.to_string(),
|
id: job.id,
|
||||||
user_id: job.user_id.clone(),
|
title: job.task.clone(),
|
||||||
created_at: job.created_at.to_rfc3339(),
|
description: String::new(),
|
||||||
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
|
state: ui_state.to_string(),
|
||||||
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
|
user_id: job.user_id.clone(),
|
||||||
elapsed_secs,
|
created_at: job.created_at.to_rfc3339(),
|
||||||
project_dir: Some(job.project_dir.clone()),
|
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
|
||||||
browse_url: Some(format!("/projects/{}/", browse_id)),
|
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
|
||||||
job_mode: mode.filter(|m| m != "worker"),
|
elapsed_secs,
|
||||||
transitions,
|
project_dir: Some(job.project_dir.clone()),
|
||||||
can_restart: state.job_manager.is_some(),
|
browse_url: Some(format!("/projects/{}/", browse_id)),
|
||||||
can_prompt: is_claude_code && state.prompt_queue.is_some(),
|
job_mode: mode.filter(|m| m != "worker"),
|
||||||
job_kind: Some("sandbox".to_string()),
|
transitions,
|
||||||
}));
|
can_restart: state.job_manager.is_some(),
|
||||||
|
can_prompt: is_claude_code && state.prompt_queue.is_some(),
|
||||||
|
job_kind: Some("sandbox".to_string()),
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
Ok(None) => {}
|
||||||
|
Err(e) => {
|
||||||
|
return Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fall back to agent job from DB.
|
// Fall back to agent job from DB.
|
||||||
if let Ok(Some(ctx)) = store.get_job(job_id).await {
|
match store.get_job(job_id).await {
|
||||||
let elapsed_secs = ctx.started_at.map(|start| {
|
Ok(Some(ctx)) => {
|
||||||
let end = ctx.completed_at.unwrap_or_else(chrono::Utc::now);
|
if ctx.user_id != user.user_id {
|
||||||
(end - start).num_seconds().max(0) as u64
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
});
|
}
|
||||||
|
let elapsed_secs = ctx.started_at.map(|start| {
|
||||||
|
let end = ctx.completed_at.unwrap_or_else(chrono::Utc::now);
|
||||||
|
(end - start).num_seconds().max(0) as u64
|
||||||
|
});
|
||||||
|
|
||||||
// Only show prompt bar for jobs that have a running worker (Pending/InProgress).
|
// Only show prompt bar for jobs that have a running worker (Pending/InProgress).
|
||||||
// Stuck jobs have no active worker loop, so messages would be silently dropped.
|
// Stuck jobs have no active worker loop, so messages would be silently dropped.
|
||||||
let is_promptable = matches!(
|
let is_promptable = matches!(
|
||||||
ctx.state,
|
ctx.state,
|
||||||
crate::context::JobState::Pending | crate::context::JobState::InProgress
|
crate::context::JobState::Pending | crate::context::JobState::InProgress
|
||||||
);
|
);
|
||||||
return Ok(Json(JobDetailResponse {
|
Ok(Json(JobDetailResponse {
|
||||||
id: ctx.job_id,
|
id: ctx.job_id,
|
||||||
title: ctx.title.clone(),
|
title: ctx.title.clone(),
|
||||||
description: ctx.description.clone(),
|
description: ctx.description.clone(),
|
||||||
state: ctx.state.to_string(),
|
state: ctx.state.to_string(),
|
||||||
user_id: ctx.user_id.clone(),
|
user_id: ctx.user_id.clone(),
|
||||||
created_at: ctx.created_at.to_rfc3339(),
|
created_at: ctx.created_at.to_rfc3339(),
|
||||||
started_at: ctx.started_at.map(|dt| dt.to_rfc3339()),
|
started_at: ctx.started_at.map(|dt| dt.to_rfc3339()),
|
||||||
completed_at: ctx.completed_at.map(|dt| dt.to_rfc3339()),
|
completed_at: ctx.completed_at.map(|dt| dt.to_rfc3339()),
|
||||||
elapsed_secs,
|
elapsed_secs,
|
||||||
project_dir: None,
|
project_dir: None,
|
||||||
browse_url: None,
|
browse_url: None,
|
||||||
job_mode: None,
|
job_mode: None,
|
||||||
transitions: Vec::new(),
|
transitions: Vec::new(),
|
||||||
can_restart: state.scheduler.is_some(),
|
can_restart: state.scheduler.is_some(),
|
||||||
can_prompt: is_promptable && state.scheduler.is_some(),
|
can_prompt: is_promptable && state.scheduler.is_some(),
|
||||||
job_kind: Some("agent".to_string()),
|
job_kind: Some("agent".to_string()),
|
||||||
}));
|
}))
|
||||||
|
}
|
||||||
|
Ok(None) => Err((StatusCode::NOT_FOUND, "Job not found".to_string())),
|
||||||
|
Err(e) => Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
)),
|
||||||
}
|
}
|
||||||
|
|
||||||
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn jobs_cancel_handler(
|
pub async fn jobs_cancel_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let job_id = Uuid::parse_str(&id)
|
let job_id = Uuid::parse_str(&id)
|
||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||||
|
|
||||||
// Try sandbox job cancellation.
|
// Try sandbox job cancellation.
|
||||||
if let Some(ref store) = state.store
|
if let Some(ref store) = state.store {
|
||||||
&& let Ok(Some(job)) = store.get_sandbox_job(job_id).await
|
match store.get_sandbox_job(job_id).await {
|
||||||
{
|
Ok(Some(job)) => {
|
||||||
if job.status == "running" || job.status == "creating" {
|
if job.user_id != user.user_id {
|
||||||
// Stop the container if we have a job manager.
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
if let Some(ref jm) = state.job_manager
|
}
|
||||||
&& let Err(e) = jm.stop_job(job_id).await
|
if job.status == "running" || job.status == "creating" {
|
||||||
{
|
if let Some(ref jm) = state.job_manager
|
||||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container during cancellation");
|
&& let Err(e) = jm.stop_job(job_id).await
|
||||||
|
{
|
||||||
|
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container during cancellation");
|
||||||
|
}
|
||||||
|
store
|
||||||
|
.update_sandbox_job_status(
|
||||||
|
job_id,
|
||||||
|
"failed",
|
||||||
|
Some(false),
|
||||||
|
Some("Cancelled by user"),
|
||||||
|
None,
|
||||||
|
Some(chrono::Utc::now()),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
}
|
||||||
|
return Ok(Json(serde_json::json!({
|
||||||
|
"status": "cancelled",
|
||||||
|
"job_id": job_id,
|
||||||
|
})));
|
||||||
|
}
|
||||||
|
Ok(None) => {}
|
||||||
|
Err(e) => {
|
||||||
|
return Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
store
|
|
||||||
.update_sandbox_job_status(
|
|
||||||
job_id,
|
|
||||||
"failed",
|
|
||||||
Some(false),
|
|
||||||
Some("Cancelled by user"),
|
|
||||||
None,
|
|
||||||
Some(chrono::Utc::now()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
|
||||||
}
|
}
|
||||||
return Ok(Json(serde_json::json!({
|
|
||||||
"status": "cancelled",
|
|
||||||
"job_id": job_id,
|
|
||||||
})));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fall back to agent job cancellation: stop the worker via the scheduler
|
// Fall back to agent job cancellation: stop the worker via the scheduler
|
||||||
// (which updates the in-memory ContextManager AND aborts the task handle),
|
// (which updates the in-memory ContextManager AND aborts the task handle),
|
||||||
// then persist the status to the DB as a fallback.
|
// then persist the status to the DB as a fallback.
|
||||||
if let Some(ref store) = state.store
|
if let Some(ref store) = state.store {
|
||||||
&& let Ok(Some(job)) = store.get_job(job_id).await
|
match store.get_job(job_id).await {
|
||||||
{
|
Ok(Some(job)) => {
|
||||||
if job.state.is_active() {
|
if job.user_id != user.user_id {
|
||||||
// Try to stop via scheduler (aborts the worker task + updates
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
// in-memory ContextManager). This is best-effort — the job may
|
}
|
||||||
// not be in the scheduler map if it already finished.
|
if job.state.is_active() {
|
||||||
if let Some(ref slot) = state.scheduler
|
// Try to stop via scheduler (aborts the worker task + updates
|
||||||
&& let Some(ref scheduler) = *slot.read().await
|
// in-memory ContextManager). This is best-effort — the job may
|
||||||
{
|
// not be in the scheduler map if it already finished.
|
||||||
let _ = scheduler.stop(job_id).await;
|
if let Some(ref slot) = state.scheduler
|
||||||
}
|
&& let Some(ref scheduler) = *slot.read().await
|
||||||
|
{
|
||||||
|
let _ = scheduler.stop(job_id).await;
|
||||||
|
}
|
||||||
|
|
||||||
// Always persist cancellation to the DB so the state is
|
// Always persist cancellation to the DB so the state is
|
||||||
// consistent even if the scheduler wasn't available or the
|
// consistent even if the scheduler wasn't available or the
|
||||||
// job wasn't in its in-memory map.
|
// job wasn't in its in-memory map.
|
||||||
store
|
store
|
||||||
.update_job_status(
|
.update_job_status(
|
||||||
job_id,
|
job_id,
|
||||||
crate::context::JobState::Cancelled,
|
crate::context::JobState::Cancelled,
|
||||||
Some("Cancelled by user"),
|
Some("Cancelled by user"),
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
}
|
||||||
|
return Ok(Json(serde_json::json!({
|
||||||
|
"status": "cancelled",
|
||||||
|
"job_id": job_id,
|
||||||
|
})));
|
||||||
|
}
|
||||||
|
Ok(None) => {}
|
||||||
|
Err(e) => {
|
||||||
|
return Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return Ok(Json(serde_json::json!({
|
|
||||||
"status": "cancelled",
|
|
||||||
"job_id": job_id,
|
|
||||||
})));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
||||||
@@ -315,6 +363,7 @@ pub async fn jobs_cancel_handler(
|
|||||||
|
|
||||||
pub async fn jobs_restart_handler(
|
pub async fn jobs_restart_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -326,146 +375,166 @@ pub async fn jobs_restart_handler(
|
|||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||||
|
|
||||||
// Try sandbox job restart first.
|
// Try sandbox job restart first.
|
||||||
if let Ok(Some(old_job)) = store.get_sandbox_job(old_job_id).await {
|
match store.get_sandbox_job(old_job_id).await {
|
||||||
if old_job.status != "interrupted" && old_job.status != "failed" {
|
Ok(Some(old_job)) => {
|
||||||
|
if old_job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
if old_job.status != "interrupted" && old_job.status != "failed" {
|
||||||
|
return Err((
|
||||||
|
StatusCode::CONFLICT,
|
||||||
|
format!("Cannot restart job in state '{}'", old_job.status),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let jm = state.job_manager.as_ref().ok_or((
|
||||||
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
|
"Sandbox not enabled".to_string(),
|
||||||
|
))?;
|
||||||
|
|
||||||
|
// Enrich the task with failure context.
|
||||||
|
let task = if let Some(ref reason) = old_job.failure_reason {
|
||||||
|
format!(
|
||||||
|
"Previous attempt failed: {}. Retry: {}",
|
||||||
|
reason, old_job.task
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
old_job.task.clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
let new_job_id = Uuid::new_v4();
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
|
||||||
|
let record = crate::history::SandboxJobRecord {
|
||||||
|
id: new_job_id,
|
||||||
|
task: task.clone(),
|
||||||
|
status: "creating".to_string(),
|
||||||
|
user_id: old_job.user_id.clone(),
|
||||||
|
project_dir: old_job.project_dir.clone(),
|
||||||
|
success: None,
|
||||||
|
failure_reason: None,
|
||||||
|
created_at: now,
|
||||||
|
started_at: None,
|
||||||
|
completed_at: None,
|
||||||
|
credential_grants_json: old_job.credential_grants_json.clone(),
|
||||||
|
};
|
||||||
|
store
|
||||||
|
.save_sandbox_job(&record)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
let mode = match store.get_sandbox_job_mode(old_job_id).await {
|
||||||
|
Ok(Some(m)) if m == "claude_code" => {
|
||||||
|
crate::orchestrator::job_manager::JobMode::ClaudeCode
|
||||||
|
}
|
||||||
|
_ => crate::orchestrator::job_manager::JobMode::Worker,
|
||||||
|
};
|
||||||
|
|
||||||
|
let credential_grants: Vec<crate::orchestrator::auth::CredentialGrant> =
|
||||||
|
serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
job_id = %old_job.id,
|
||||||
|
"Failed to deserialize credential grants from stored job: {}. \
|
||||||
|
Restarted job will have no credentials.",
|
||||||
|
e
|
||||||
|
);
|
||||||
|
vec![]
|
||||||
|
});
|
||||||
|
|
||||||
|
let project_dir = std::path::PathBuf::from(&old_job.project_dir);
|
||||||
|
let _token = jm
|
||||||
|
.create_job(
|
||||||
|
new_job_id,
|
||||||
|
&task,
|
||||||
|
Some(project_dir),
|
||||||
|
mode,
|
||||||
|
credential_grants,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| {
|
||||||
|
(
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Failed to create container: {}", e),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
store
|
||||||
|
.update_sandbox_job_status(new_job_id, "running", None, None, Some(now), None)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
return Ok(Json(serde_json::json!({
|
||||||
|
"status": "restarted",
|
||||||
|
"old_job_id": old_job_id,
|
||||||
|
"new_job_id": new_job_id,
|
||||||
|
})));
|
||||||
|
}
|
||||||
|
Ok(None) => {}
|
||||||
|
Err(e) => {
|
||||||
return Err((
|
return Err((
|
||||||
StatusCode::CONFLICT,
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
format!("Cannot restart job in state '{}'", old_job.status),
|
format!("Database error: {}", e),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let jm = state.job_manager.as_ref().ok_or((
|
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Sandbox not enabled".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
// Enrich the task with failure context.
|
|
||||||
let task = if let Some(ref reason) = old_job.failure_reason {
|
|
||||||
format!(
|
|
||||||
"Previous attempt failed: {}. Retry: {}",
|
|
||||||
reason, old_job.task
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
old_job.task.clone()
|
|
||||||
};
|
|
||||||
|
|
||||||
let new_job_id = Uuid::new_v4();
|
|
||||||
let now = chrono::Utc::now();
|
|
||||||
|
|
||||||
let record = crate::history::SandboxJobRecord {
|
|
||||||
id: new_job_id,
|
|
||||||
task: task.clone(),
|
|
||||||
status: "creating".to_string(),
|
|
||||||
user_id: old_job.user_id.clone(),
|
|
||||||
project_dir: old_job.project_dir.clone(),
|
|
||||||
success: None,
|
|
||||||
failure_reason: None,
|
|
||||||
created_at: now,
|
|
||||||
started_at: None,
|
|
||||||
completed_at: None,
|
|
||||||
credential_grants_json: old_job.credential_grants_json.clone(),
|
|
||||||
};
|
|
||||||
store
|
|
||||||
.save_sandbox_job(&record)
|
|
||||||
.await
|
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
|
||||||
|
|
||||||
let mode = match store.get_sandbox_job_mode(old_job_id).await {
|
|
||||||
Ok(Some(m)) if m == "claude_code" => {
|
|
||||||
crate::orchestrator::job_manager::JobMode::ClaudeCode
|
|
||||||
}
|
|
||||||
_ => crate::orchestrator::job_manager::JobMode::Worker,
|
|
||||||
};
|
|
||||||
|
|
||||||
let credential_grants: Vec<crate::orchestrator::auth::CredentialGrant> =
|
|
||||||
serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| {
|
|
||||||
tracing::warn!(
|
|
||||||
job_id = %old_job.id,
|
|
||||||
"Failed to deserialize credential grants from stored job: {}. \
|
|
||||||
Restarted job will have no credentials.",
|
|
||||||
e
|
|
||||||
);
|
|
||||||
vec![]
|
|
||||||
});
|
|
||||||
|
|
||||||
let project_dir = std::path::PathBuf::from(&old_job.project_dir);
|
|
||||||
let _token = jm
|
|
||||||
.create_job(
|
|
||||||
new_job_id,
|
|
||||||
&task,
|
|
||||||
Some(project_dir),
|
|
||||||
mode,
|
|
||||||
credential_grants,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.map_err(|e| {
|
|
||||||
(
|
|
||||||
StatusCode::INTERNAL_SERVER_ERROR,
|
|
||||||
format!("Failed to create container: {}", e),
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
store
|
|
||||||
.update_sandbox_job_status(new_job_id, "running", None, None, Some(now), None)
|
|
||||||
.await
|
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
|
||||||
|
|
||||||
return Ok(Json(serde_json::json!({
|
|
||||||
"status": "restarted",
|
|
||||||
"old_job_id": old_job_id,
|
|
||||||
"new_job_id": new_job_id,
|
|
||||||
})));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Try agent job restart: dispatch a new job via the scheduler.
|
// Try agent job restart: dispatch a new job via the scheduler.
|
||||||
if let Ok(Some(old_job)) = store.get_job(old_job_id).await {
|
match store.get_job(old_job_id).await {
|
||||||
if old_job.state.is_active() {
|
Ok(Some(old_job)) => {
|
||||||
return Err((
|
if old_job.user_id != user.user_id {
|
||||||
StatusCode::CONFLICT,
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
format!("Cannot restart job in state '{}'", old_job.state),
|
}
|
||||||
));
|
if old_job.state.is_active() {
|
||||||
|
return Err((
|
||||||
|
StatusCode::CONFLICT,
|
||||||
|
format!("Cannot restart job in state '{}'", old_job.state),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let slot = state.scheduler.as_ref().ok_or((
|
||||||
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
|
"Scheduler not available".to_string(),
|
||||||
|
))?;
|
||||||
|
let scheduler_guard = slot.read().await;
|
||||||
|
let scheduler = scheduler_guard.as_ref().ok_or((
|
||||||
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
|
"Agent not started yet".to_string(),
|
||||||
|
))?;
|
||||||
|
|
||||||
|
// Look up failure reason (O(1) point lookup).
|
||||||
|
let failure_reason = store
|
||||||
|
.get_agent_job_failure_reason(old_job_id)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let title = if !failure_reason.is_empty() {
|
||||||
|
format!(
|
||||||
|
"Previous attempt failed: {}. Retry: {}",
|
||||||
|
failure_reason, old_job.title
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
old_job.title.clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
let new_job_id = scheduler
|
||||||
|
.dispatch_job(&old_job.user_id, &title, &old_job.description, None)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(Json(serde_json::json!({
|
||||||
|
"status": "restarted",
|
||||||
|
"old_job_id": old_job_id,
|
||||||
|
"new_job_id": new_job_id,
|
||||||
|
})))
|
||||||
}
|
}
|
||||||
|
Ok(None) => Err((StatusCode::NOT_FOUND, "Job not found".to_string())),
|
||||||
let slot = state.scheduler.as_ref().ok_or((
|
Err(e) => Err((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
"Scheduler not available".to_string(),
|
format!("Database error: {}", e),
|
||||||
))?;
|
)),
|
||||||
let scheduler_guard = slot.read().await;
|
|
||||||
let scheduler = scheduler_guard.as_ref().ok_or((
|
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Agent not started yet".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
// Look up failure reason (O(1) point lookup).
|
|
||||||
let failure_reason = store
|
|
||||||
.get_agent_job_failure_reason(old_job_id)
|
|
||||||
.await
|
|
||||||
.ok()
|
|
||||||
.flatten()
|
|
||||||
.unwrap_or_default();
|
|
||||||
|
|
||||||
let title = if !failure_reason.is_empty() {
|
|
||||||
format!(
|
|
||||||
"Previous attempt failed: {}. Retry: {}",
|
|
||||||
failure_reason, old_job.title
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
old_job.title.clone()
|
|
||||||
};
|
|
||||||
|
|
||||||
let new_job_id = scheduler
|
|
||||||
.dispatch_job(&old_job.user_id, &title, &old_job.description, None)
|
|
||||||
.await
|
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
|
||||||
|
|
||||||
return Ok(Json(serde_json::json!({
|
|
||||||
"status": "restarted",
|
|
||||||
"old_job_id": old_job_id,
|
|
||||||
"new_job_id": new_job_id,
|
|
||||||
})));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Submit a follow-up prompt to a running job.
|
/// Submit a follow-up prompt to a running job.
|
||||||
@@ -476,6 +545,7 @@ pub async fn jobs_restart_handler(
|
|||||||
/// - Worker-mode sandbox jobs → not supported (no mechanism to inject)
|
/// - Worker-mode sandbox jobs → not supported (no mechanism to inject)
|
||||||
pub async fn jobs_prompt_handler(
|
pub async fn jobs_prompt_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
Json(body): Json<serde_json::Value>,
|
Json(body): Json<serde_json::Value>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
@@ -494,10 +564,15 @@ pub async fn jobs_prompt_handler(
|
|||||||
|
|
||||||
let done = body.get("done").and_then(|v| v.as_bool()).unwrap_or(false);
|
let done = body.get("done").and_then(|v| v.as_bool()).unwrap_or(false);
|
||||||
|
|
||||||
// Try sandbox job path: check if we have a sandbox record for this ID.
|
// Try sandbox job path first: verify ownership, then route to Claude Code or reject.
|
||||||
if let Some(ref s) = state.store
|
if let Some(ref s) = state.store
|
||||||
&& let Ok(Some(_)) = s.get_sandbox_job(job_id).await
|
&& let Ok(Some(sandbox_job)) = s.get_sandbox_job(job_id).await
|
||||||
{
|
{
|
||||||
|
// Verify ownership.
|
||||||
|
if sandbox_job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
// It's a sandbox job. Check if Claude Code mode.
|
// It's a sandbox job. Check if Claude Code mode.
|
||||||
let mode = s.get_sandbox_job_mode(job_id).await.ok().flatten();
|
let mode = s.get_sandbox_job_mode(job_id).await.ok().flatten();
|
||||||
if mode.as_deref() == Some("claude_code") {
|
if mode.as_deref() == Some("claude_code") {
|
||||||
@@ -522,7 +597,26 @@ pub async fn jobs_prompt_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Try agent job path: send via scheduler.
|
// Try agent job path: verify ownership, then send via scheduler.
|
||||||
|
if let Some(ref store) = state.store {
|
||||||
|
match store.get_job(job_id).await {
|
||||||
|
Ok(Some(agent_job)) => {
|
||||||
|
if agent_job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
return Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let slot = state.scheduler.as_ref().ok_or((
|
let slot = state.scheduler.as_ref().ok_or((
|
||||||
StatusCode::NOT_IMPLEMENTED,
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
"Agent job prompts require the scheduler to be configured".to_string(),
|
"Agent job prompts require the scheduler to be configured".to_string(),
|
||||||
@@ -550,6 +644,7 @@ pub async fn jobs_prompt_handler(
|
|||||||
/// Load persisted job events for a job (for history replay on page open).
|
/// Load persisted job events for a job (for history replay on page open).
|
||||||
pub async fn jobs_events_handler(
|
pub async fn jobs_events_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -561,6 +656,24 @@ pub async fn jobs_events_handler(
|
|||||||
.parse()
|
.parse()
|
||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||||
|
|
||||||
|
// Verify ownership before returning events.
|
||||||
|
match store.get_sandbox_job(job_id).await {
|
||||||
|
Ok(Some(job)) => {
|
||||||
|
if job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
return Err((
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("Database error: {}", e),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let events = store
|
let events = store
|
||||||
.list_job_events(job_id, None)
|
.list_job_events(job_id, None)
|
||||||
.await
|
.await
|
||||||
@@ -593,6 +706,7 @@ pub struct FilePathQuery {
|
|||||||
|
|
||||||
pub async fn job_files_list_handler(
|
pub async fn job_files_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
Query(query): Query<FilePathQuery>,
|
Query(query): Query<FilePathQuery>,
|
||||||
) -> Result<Json<ProjectFilesResponse>, (StatusCode, String)> {
|
) -> Result<Json<ProjectFilesResponse>, (StatusCode, String)> {
|
||||||
@@ -610,6 +724,10 @@ pub async fn job_files_list_handler(
|
|||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
||||||
|
|
||||||
|
if job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let base = std::path::PathBuf::from(&job.project_dir);
|
let base = std::path::PathBuf::from(&job.project_dir);
|
||||||
let rel_path = query.path.as_deref().unwrap_or("");
|
let rel_path = query.path.as_deref().unwrap_or("");
|
||||||
let target = base.join(rel_path);
|
let target = base.join(rel_path);
|
||||||
@@ -656,6 +774,7 @@ pub async fn job_files_list_handler(
|
|||||||
|
|
||||||
pub async fn job_files_read_handler(
|
pub async fn job_files_read_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
Query(query): Query<FilePathQuery>,
|
Query(query): Query<FilePathQuery>,
|
||||||
) -> Result<Json<ProjectFileReadResponse>, (StatusCode, String)> {
|
) -> Result<Json<ProjectFileReadResponse>, (StatusCode, String)> {
|
||||||
@@ -673,6 +792,10 @@ pub async fn job_files_read_handler(
|
|||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
||||||
|
|
||||||
|
if job.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let path = query.path.as_deref().ok_or((
|
let path = query.path.as_deref().ok_or((
|
||||||
StatusCode::BAD_REQUEST,
|
StatusCode::BAD_REQUEST,
|
||||||
"path parameter required".to_string(),
|
"path parameter required".to_string(),
|
||||||
|
|||||||
@@ -9,8 +9,27 @@ use axum::{
|
|||||||
};
|
};
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use crate::channels::web::auth::{AuthenticatedUser, UserIdentity};
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
use crate::workspace::Workspace;
|
||||||
|
|
||||||
|
/// Resolve the workspace for the authenticated user.
|
||||||
|
///
|
||||||
|
/// Prefers `workspace_pool` (multi-user mode) when available, falling back
|
||||||
|
/// to the single-user `state.workspace`.
|
||||||
|
pub(crate) async fn resolve_workspace(
|
||||||
|
state: &GatewayState,
|
||||||
|
user: &UserIdentity,
|
||||||
|
) -> Result<Arc<Workspace>, (StatusCode, String)> {
|
||||||
|
if let Some(ref pool) = state.workspace_pool {
|
||||||
|
return Ok(pool.get_or_create(user).await);
|
||||||
|
}
|
||||||
|
state.workspace.as_ref().cloned().ok_or((
|
||||||
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
|
"Workspace not available".to_string(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
pub struct TreeQuery {
|
pub struct TreeQuery {
|
||||||
@@ -20,12 +39,10 @@ pub struct TreeQuery {
|
|||||||
|
|
||||||
pub async fn memory_tree_handler(
|
pub async fn memory_tree_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Query(_query): Query<TreeQuery>,
|
Query(_query): Query<TreeQuery>,
|
||||||
) -> Result<Json<MemoryTreeResponse>, (StatusCode, String)> {
|
) -> Result<Json<MemoryTreeResponse>, (StatusCode, String)> {
|
||||||
let workspace = state.workspace.as_ref().ok_or((
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Workspace not available".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
// Build tree from list_all (flat list of all paths)
|
// Build tree from list_all (flat list of all paths)
|
||||||
let all_paths = workspace
|
let all_paths = workspace
|
||||||
@@ -68,12 +85,10 @@ pub struct ListQuery {
|
|||||||
|
|
||||||
pub async fn memory_list_handler(
|
pub async fn memory_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Query(query): Query<ListQuery>,
|
Query(query): Query<ListQuery>,
|
||||||
) -> Result<Json<MemoryListResponse>, (StatusCode, String)> {
|
) -> Result<Json<MemoryListResponse>, (StatusCode, String)> {
|
||||||
let workspace = state.workspace.as_ref().ok_or((
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Workspace not available".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
let path = query.path.as_deref().unwrap_or("");
|
let path = query.path.as_deref().unwrap_or("");
|
||||||
let entries = workspace
|
let entries = workspace
|
||||||
@@ -104,12 +119,10 @@ pub struct ReadQuery {
|
|||||||
|
|
||||||
pub async fn memory_read_handler(
|
pub async fn memory_read_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Query(query): Query<ReadQuery>,
|
Query(query): Query<ReadQuery>,
|
||||||
) -> Result<Json<MemoryReadResponse>, (StatusCode, String)> {
|
) -> Result<Json<MemoryReadResponse>, (StatusCode, String)> {
|
||||||
let workspace = state.workspace.as_ref().ok_or((
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Workspace not available".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
let doc = workspace
|
let doc = workspace
|
||||||
.read(&query.path)
|
.read(&query.path)
|
||||||
@@ -123,17 +136,75 @@ pub async fn memory_read_handler(
|
|||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
// memory_write_handler lives in server.rs (layer-aware version with append,
|
pub async fn memory_write_handler(
|
||||||
// privacy redirect, and proper error status codes).
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
|
Json(req): Json<MemoryWriteRequest>,
|
||||||
|
) -> Result<Json<MemoryWriteResponse>, (StatusCode, String)> {
|
||||||
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
|
|
||||||
|
// Route through layer-aware methods when a layer is specified.
|
||||||
|
//
|
||||||
|
// Note: unlike MemoryWriteTool, this endpoint does NOT block writes to
|
||||||
|
// identity files (IDENTITY.md, SOUL.md, etc.). The HTTP API is an
|
||||||
|
// authenticated admin interface; the supervisor uses it to seed identity
|
||||||
|
// files at startup. Identity-file protection is enforced at the tool
|
||||||
|
// layer (LLM-facing) where the write originates from an untrusted agent.
|
||||||
|
if let Some(ref layer_name) = req.layer {
|
||||||
|
let result = if req.append {
|
||||||
|
workspace
|
||||||
|
.append_to_layer(layer_name, &req.path, &req.content, req.force)
|
||||||
|
.await
|
||||||
|
} else {
|
||||||
|
workspace
|
||||||
|
.write_to_layer(layer_name, &req.path, &req.content, req.force)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
.map_err(|e| {
|
||||||
|
use crate::error::WorkspaceError;
|
||||||
|
let status = match &e {
|
||||||
|
WorkspaceError::LayerNotFound { .. } => StatusCode::BAD_REQUEST,
|
||||||
|
WorkspaceError::LayerReadOnly { .. } => StatusCode::FORBIDDEN,
|
||||||
|
WorkspaceError::PrivacyRedirectFailed => StatusCode::UNPROCESSABLE_ENTITY,
|
||||||
|
_ => StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
};
|
||||||
|
(status, e.to_string())
|
||||||
|
})?;
|
||||||
|
return Ok(Json(MemoryWriteResponse {
|
||||||
|
path: req.path,
|
||||||
|
status: "written",
|
||||||
|
redirected: Some(result.redirected),
|
||||||
|
actual_layer: Some(result.actual_layer),
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Non-layer path: honor the append field
|
||||||
|
if req.append {
|
||||||
|
workspace
|
||||||
|
.append(&req.path, &req.content)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
} else {
|
||||||
|
workspace
|
||||||
|
.write(&req.path, &req.content)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Json(MemoryWriteResponse {
|
||||||
|
path: req.path,
|
||||||
|
status: "written",
|
||||||
|
redirected: None,
|
||||||
|
actual_layer: None,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn memory_search_handler(
|
pub async fn memory_search_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Json(req): Json<MemorySearchRequest>,
|
Json(req): Json<MemorySearchRequest>,
|
||||||
) -> Result<Json<MemorySearchResponse>, (StatusCode, String)> {
|
) -> Result<Json<MemorySearchResponse>, (StatusCode, String)> {
|
||||||
let workspace = state.workspace.as_ref().ok_or((
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
|
||||||
"Workspace not available".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
let limit = req.limit.unwrap_or(10);
|
let limit = req.limit.unwrap_or(10);
|
||||||
let results = workspace
|
let results = workspace
|
||||||
@@ -142,10 +213,10 @@ pub async fn memory_search_handler(
|
|||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
let hits: Vec<SearchHit> = results
|
let hits: Vec<SearchHit> = results
|
||||||
.into_iter()
|
.iter()
|
||||||
.map(|r| SearchHit {
|
.map(|r| SearchHit {
|
||||||
path: r.document_path,
|
path: r.document_id.to_string(),
|
||||||
content: r.content,
|
content: r.content.clone(),
|
||||||
score: r.score as f64,
|
score: r.score as f64,
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|||||||
@@ -1,13 +1,10 @@
|
|||||||
//! Handler modules for the web gateway API.
|
//! Handler modules for the web gateway API.
|
||||||
//!
|
//!
|
||||||
//! Each module groups related endpoint handlers by domain.
|
//! Each module groups related endpoint handlers by domain.
|
||||||
//!
|
|
||||||
//! # Migration status
|
|
||||||
//!
|
|
||||||
//! `skills` is the canonical implementation used by `server.rs`.
|
|
||||||
//! The remaining modules are in-progress migrations from inline server.rs
|
|
||||||
//! handlers; their functions are not yet wired up, hence the `dead_code` allow.
|
|
||||||
|
|
||||||
|
pub mod jobs;
|
||||||
|
pub mod memory;
|
||||||
|
pub mod routines;
|
||||||
pub mod skills;
|
pub mod skills;
|
||||||
|
|
||||||
// Modules not yet wired into server.rs router -- suppress dead_code until
|
// Modules not yet wired into server.rs router -- suppress dead_code until
|
||||||
@@ -17,12 +14,6 @@ pub mod chat;
|
|||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub mod extensions;
|
pub mod extensions;
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub mod jobs;
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub mod memory;
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub mod routines;
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub mod settings;
|
pub mod settings;
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub mod static_files;
|
pub mod static_files;
|
||||||
|
|||||||
@@ -11,12 +11,14 @@ use serde::Deserialize;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::routine::{Trigger, next_cron_fire};
|
use crate::agent::routine::{Trigger, next_cron_fire};
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
use crate::error::RoutineError;
|
use crate::error::RoutineError;
|
||||||
|
|
||||||
pub async fn routines_list_handler(
|
pub async fn routines_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<RoutineListResponse>, (StatusCode, String)> {
|
) -> Result<Json<RoutineListResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
@@ -24,7 +26,7 @@ pub async fn routines_list_handler(
|
|||||||
))?;
|
))?;
|
||||||
|
|
||||||
let routines = store
|
let routines = store
|
||||||
.list_all_routines()
|
.list_routines(&user.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
@@ -35,6 +37,7 @@ pub async fn routines_list_handler(
|
|||||||
|
|
||||||
pub async fn routines_summary_handler(
|
pub async fn routines_summary_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<RoutineSummaryResponse>, (StatusCode, String)> {
|
) -> Result<Json<RoutineSummaryResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
StatusCode::SERVICE_UNAVAILABLE,
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
@@ -42,7 +45,7 @@ pub async fn routines_summary_handler(
|
|||||||
))?;
|
))?;
|
||||||
|
|
||||||
let routines = store
|
let routines = store
|
||||||
.list_all_routines()
|
.list_routines(&user.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
@@ -78,6 +81,7 @@ pub async fn routines_summary_handler(
|
|||||||
|
|
||||||
pub async fn routines_detail_handler(
|
pub async fn routines_detail_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<RoutineDetailResponse>, (StatusCode, String)> {
|
) -> Result<Json<RoutineDetailResponse>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -94,6 +98,10 @@ pub async fn routines_detail_handler(
|
|||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
||||||
|
|
||||||
|
if routine.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Routine not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let runs = store
|
let runs = store
|
||||||
.list_routine_runs(routine_id, 20)
|
.list_routine_runs(routine_id, 20)
|
||||||
.await
|
.await
|
||||||
@@ -106,7 +114,7 @@ pub async fn routines_detail_handler(
|
|||||||
trigger_type: run.trigger_type.clone(),
|
trigger_type: run.trigger_type.clone(),
|
||||||
started_at: run.started_at.to_rfc3339(),
|
started_at: run.started_at.to_rfc3339(),
|
||||||
completed_at: run.completed_at.map(|dt| dt.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(),
|
result_summary: run.result_summary.clone(),
|
||||||
tokens_used: run.tokens_used,
|
tokens_used: run.tokens_used,
|
||||||
job_id: run.job_id,
|
job_id: run.job_id,
|
||||||
@@ -137,6 +145,7 @@ pub async fn routines_detail_handler(
|
|||||||
|
|
||||||
pub async fn routines_trigger_handler(
|
pub async fn routines_trigger_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
// Clone the Arc out of the lock to avoid holding the RwLock across .await.
|
// Clone the Arc out of the lock to avoid holding the RwLock across .await.
|
||||||
@@ -152,7 +161,7 @@ pub async fn routines_trigger_handler(
|
|||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
||||||
|
|
||||||
let run_id = engine
|
let run_id = engine
|
||||||
.fire_manual(routine_id, Some(&state.user_id))
|
.fire_manual(routine_id, Some(&user.user_id))
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (routine_error_status(&e), e.to_string()))?;
|
.map_err(|e| (routine_error_status(&e), e.to_string()))?;
|
||||||
|
|
||||||
@@ -170,6 +179,7 @@ pub struct ToggleRequest {
|
|||||||
|
|
||||||
pub async fn routines_toggle_handler(
|
pub async fn routines_toggle_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
body: Option<Json<ToggleRequest>>,
|
body: Option<Json<ToggleRequest>>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
@@ -187,6 +197,10 @@ pub async fn routines_toggle_handler(
|
|||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
||||||
|
|
||||||
|
if routine.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Routine not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let was_enabled = routine.enabled;
|
let was_enabled = routine.enabled;
|
||||||
// If a specific value was provided, use it; otherwise toggle.
|
// If a specific value was provided, use it; otherwise toggle.
|
||||||
routine.enabled = match body {
|
routine.enabled = match body {
|
||||||
@@ -230,6 +244,7 @@ pub async fn routines_toggle_handler(
|
|||||||
|
|
||||||
pub async fn routines_delete_handler(
|
pub async fn routines_delete_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -240,6 +255,17 @@ pub async fn routines_delete_handler(
|
|||||||
let routine_id = Uuid::parse_str(&id)
|
let routine_id = Uuid::parse_str(&id)
|
||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
||||||
|
|
||||||
|
// Verify ownership before deleting.
|
||||||
|
let routine = store
|
||||||
|
.get_routine(routine_id)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
|
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
||||||
|
|
||||||
|
if routine.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Routine not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let deleted = store
|
let deleted = store
|
||||||
.delete_routine(routine_id)
|
.delete_routine(routine_id)
|
||||||
.await
|
.await
|
||||||
@@ -261,8 +287,10 @@ pub async fn routines_delete_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)] // Used by server.rs inline version; kept in sync here for future migration.
|
||||||
pub async fn routines_runs_handler(
|
pub async fn routines_runs_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(id): Path<String>,
|
Path(id): Path<String>,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let store = state.store.as_ref().ok_or((
|
let store = state.store.as_ref().ok_or((
|
||||||
@@ -273,6 +301,17 @@ pub async fn routines_runs_handler(
|
|||||||
let routine_id = Uuid::parse_str(&id)
|
let routine_id = Uuid::parse_str(&id)
|
||||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?;
|
||||||
|
|
||||||
|
// Verify ownership before listing runs.
|
||||||
|
let routine = store
|
||||||
|
.get_routine(routine_id)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
|
.ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?;
|
||||||
|
|
||||||
|
if routine.user_id != user.user_id {
|
||||||
|
return Err((StatusCode::NOT_FOUND, "Routine not found".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let runs = store
|
let runs = store
|
||||||
.list_routine_runs(routine_id, 50)
|
.list_routine_runs(routine_id, 50)
|
||||||
.await
|
.await
|
||||||
@@ -285,7 +324,7 @@ pub async fn routines_runs_handler(
|
|||||||
trigger_type: run.trigger_type.clone(),
|
trigger_type: run.trigger_type.clone(),
|
||||||
started_at: run.started_at.to_rfc3339(),
|
started_at: run.started_at.to_rfc3339(),
|
||||||
completed_at: run.completed_at.map(|dt| dt.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(),
|
result_summary: run.result_summary.clone(),
|
||||||
tokens_used: run.tokens_used,
|
tokens_used: run.tokens_used,
|
||||||
job_id: run.job_id,
|
job_id: run.job_id,
|
||||||
|
|||||||
@@ -8,17 +8,19 @@ use axum::{
|
|||||||
http::StatusCode,
|
http::StatusCode,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
|
||||||
pub async fn settings_list_handler(
|
pub async fn settings_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<SettingsListResponse>, StatusCode> {
|
) -> Result<Json<SettingsListResponse>, StatusCode> {
|
||||||
let store = state
|
let store = state
|
||||||
.store
|
.store
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
let rows = store.list_settings(&state.user_id).await.map_err(|e| {
|
let rows = store.list_settings(&user.user_id).await.map_err(|e| {
|
||||||
tracing::error!("Failed to list settings: {}", e);
|
tracing::error!("Failed to list settings: {}", e);
|
||||||
StatusCode::INTERNAL_SERVER_ERROR
|
StatusCode::INTERNAL_SERVER_ERROR
|
||||||
})?;
|
})?;
|
||||||
@@ -37,6 +39,7 @@ pub async fn settings_list_handler(
|
|||||||
|
|
||||||
pub async fn settings_get_handler(
|
pub async fn settings_get_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(key): Path<String>,
|
Path(key): Path<String>,
|
||||||
) -> Result<Json<SettingResponse>, StatusCode> {
|
) -> Result<Json<SettingResponse>, StatusCode> {
|
||||||
let store = state
|
let store = state
|
||||||
@@ -44,7 +47,7 @@ pub async fn settings_get_handler(
|
|||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
let row = store
|
let row = store
|
||||||
.get_setting_full(&state.user_id, &key)
|
.get_setting_full(&user.user_id, &key)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
tracing::error!("Failed to get setting '{}': {}", key, e);
|
tracing::error!("Failed to get setting '{}': {}", key, e);
|
||||||
@@ -61,6 +64,7 @@ pub async fn settings_get_handler(
|
|||||||
|
|
||||||
pub async fn settings_set_handler(
|
pub async fn settings_set_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(key): Path<String>,
|
Path(key): Path<String>,
|
||||||
Json(body): Json<SettingWriteRequest>,
|
Json(body): Json<SettingWriteRequest>,
|
||||||
) -> Result<StatusCode, StatusCode> {
|
) -> Result<StatusCode, StatusCode> {
|
||||||
@@ -69,7 +73,7 @@ pub async fn settings_set_handler(
|
|||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
store
|
store
|
||||||
.set_setting(&state.user_id, &key, &body.value)
|
.set_setting(&user.user_id, &key, &body.value)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
tracing::error!("Failed to set setting '{}': {}", key, e);
|
tracing::error!("Failed to set setting '{}': {}", key, e);
|
||||||
@@ -81,6 +85,7 @@ pub async fn settings_set_handler(
|
|||||||
|
|
||||||
pub async fn settings_delete_handler(
|
pub async fn settings_delete_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Path(key): Path<String>,
|
Path(key): Path<String>,
|
||||||
) -> Result<StatusCode, StatusCode> {
|
) -> Result<StatusCode, StatusCode> {
|
||||||
let store = state
|
let store = state
|
||||||
@@ -88,7 +93,7 @@ pub async fn settings_delete_handler(
|
|||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
store
|
store
|
||||||
.delete_setting(&state.user_id, &key)
|
.delete_setting(&user.user_id, &key)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
tracing::error!("Failed to delete setting '{}': {}", key, e);
|
tracing::error!("Failed to delete setting '{}': {}", key, e);
|
||||||
@@ -100,12 +105,13 @@ pub async fn settings_delete_handler(
|
|||||||
|
|
||||||
pub async fn settings_export_handler(
|
pub async fn settings_export_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
) -> Result<Json<SettingsExportResponse>, StatusCode> {
|
) -> Result<Json<SettingsExportResponse>, StatusCode> {
|
||||||
let store = state
|
let store = state
|
||||||
.store
|
.store
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
let settings = store.get_all_settings(&state.user_id).await.map_err(|e| {
|
let settings = store.get_all_settings(&user.user_id).await.map_err(|e| {
|
||||||
tracing::error!("Failed to export settings: {}", e);
|
tracing::error!("Failed to export settings: {}", e);
|
||||||
StatusCode::INTERNAL_SERVER_ERROR
|
StatusCode::INTERNAL_SERVER_ERROR
|
||||||
})?;
|
})?;
|
||||||
@@ -115,6 +121,7 @@ pub async fn settings_export_handler(
|
|||||||
|
|
||||||
pub async fn settings_import_handler(
|
pub async fn settings_import_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
Json(body): Json<SettingsImportRequest>,
|
Json(body): Json<SettingsImportRequest>,
|
||||||
) -> Result<StatusCode, StatusCode> {
|
) -> Result<StatusCode, StatusCode> {
|
||||||
let store = state
|
let store = state
|
||||||
@@ -122,7 +129,7 @@ pub async fn settings_import_handler(
|
|||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
store
|
store
|
||||||
.set_all_settings(&state.user_id, &body.settings)
|
.set_all_settings(&user.user_id, &body.settings)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
tracing::error!("Failed to import settings: {}", e);
|
tracing::error!("Failed to import settings: {}", e);
|
||||||
|
|||||||
@@ -8,11 +8,13 @@ use axum::{
|
|||||||
http::StatusCode,
|
http::StatusCode,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::server::GatewayState;
|
use crate::channels::web::server::GatewayState;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
|
||||||
pub async fn skills_list_handler(
|
pub async fn skills_list_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(_user): AuthenticatedUser,
|
||||||
) -> Result<Json<SkillListResponse>, (StatusCode, String)> {
|
) -> Result<Json<SkillListResponse>, (StatusCode, String)> {
|
||||||
let registry = state.skill_registry.as_ref().ok_or((
|
let registry = state.skill_registry.as_ref().ok_or((
|
||||||
StatusCode::NOT_IMPLEMENTED,
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
@@ -45,6 +47,7 @@ pub async fn skills_list_handler(
|
|||||||
|
|
||||||
pub async fn skills_search_handler(
|
pub async fn skills_search_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(_user): AuthenticatedUser,
|
||||||
Json(req): Json<SkillSearchRequest>,
|
Json(req): Json<SkillSearchRequest>,
|
||||||
) -> Result<Json<SkillSearchResponse>, (StatusCode, String)> {
|
) -> Result<Json<SkillSearchResponse>, (StatusCode, String)> {
|
||||||
let registry = state.skill_registry.as_ref().ok_or((
|
let registry = state.skill_registry.as_ref().ok_or((
|
||||||
@@ -119,6 +122,7 @@ pub async fn skills_search_handler(
|
|||||||
|
|
||||||
pub async fn skills_install_handler(
|
pub async fn skills_install_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
headers: axum::http::HeaderMap,
|
headers: axum::http::HeaderMap,
|
||||||
Json(req): Json<SkillInstallRequest>,
|
Json(req): Json<SkillInstallRequest>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
@@ -135,6 +139,8 @@ pub async fn skills_install_handler(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::info!(user_id = %user.user_id, skill = %req.name, "skill install requested");
|
||||||
|
|
||||||
let registry = state.skill_registry.as_ref().ok_or((
|
let registry = state.skill_registry.as_ref().ok_or((
|
||||||
StatusCode::NOT_IMPLEMENTED,
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
"Skills system not enabled".to_string(),
|
"Skills system not enabled".to_string(),
|
||||||
@@ -219,6 +225,7 @@ pub async fn skills_install_handler(
|
|||||||
|
|
||||||
pub async fn skills_remove_handler(
|
pub async fn skills_remove_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
headers: axum::http::HeaderMap,
|
headers: axum::http::HeaderMap,
|
||||||
Path(name): Path<String>,
|
Path(name): Path<String>,
|
||||||
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
) -> Result<Json<ActionResponse>, (StatusCode, String)> {
|
||||||
@@ -234,6 +241,8 @@ pub async fn skills_remove_handler(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::info!(user_id = %user.user_id, skill = %name, "skill remove requested");
|
||||||
|
|
||||||
let registry = state.skill_registry.as_ref().ok_or((
|
let registry = state.skill_registry.as_ref().ok_or((
|
||||||
StatusCode::NOT_IMPLEMENTED,
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
"Skills system not enabled".to_string(),
|
"Skills system not enabled".to_string(),
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ use axum::{
|
|||||||
};
|
};
|
||||||
|
|
||||||
use crate::bootstrap::ironclaw_base_dir;
|
use crate::bootstrap::ironclaw_base_dir;
|
||||||
|
use crate::channels::web::auth::AuthenticatedUser;
|
||||||
use crate::channels::web::types::*;
|
use crate::channels::web::types::*;
|
||||||
|
|
||||||
// --- Static file handlers ---
|
// --- Static file handlers ---
|
||||||
@@ -113,6 +114,7 @@ use crate::channels::web::server::GatewayState;
|
|||||||
|
|
||||||
pub async fn logs_events_handler(
|
pub async fn logs_events_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(_user): AuthenticatedUser,
|
||||||
) -> Result<
|
) -> Result<
|
||||||
Sse<impl futures::Stream<Item = Result<Event, Infallible>> + Send + 'static>,
|
Sse<impl futures::Stream<Item = Result<Event, Infallible>> + Send + 'static>,
|
||||||
(StatusCode, String),
|
(StatusCode, String),
|
||||||
@@ -152,6 +154,7 @@ pub async fn logs_events_handler(
|
|||||||
|
|
||||||
pub async fn gateway_status_handler(
|
pub async fn gateway_status_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(_user): AuthenticatedUser,
|
||||||
) -> Json<GatewayStatusResponse> {
|
) -> Json<GatewayStatusResponse> {
|
||||||
let sse_connections = state.sse.connection_count();
|
let sse_connections = state.sse.connection_count();
|
||||||
let ws_connections = state
|
let ws_connections = state
|
||||||
|
|||||||
+137
-36
@@ -31,6 +31,9 @@ pub mod ws;
|
|||||||
/// [`TestGatewayBuilder`](test_helpers::TestGatewayBuilder).
|
/// [`TestGatewayBuilder`](test_helpers::TestGatewayBuilder).
|
||||||
pub mod test_helpers;
|
pub mod test_helpers;
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests;
|
||||||
|
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
@@ -52,22 +55,24 @@ use crate::workspace::Workspace;
|
|||||||
|
|
||||||
use self::log_layer::{LogBroadcaster, LogLevelHandle};
|
use self::log_layer::{LogBroadcaster, LogLevelHandle};
|
||||||
|
|
||||||
|
use self::auth::MultiAuthState;
|
||||||
use self::server::GatewayState;
|
use self::server::GatewayState;
|
||||||
use self::sse::SseManager;
|
use self::sse::SseManager;
|
||||||
use self::types::SseEvent;
|
use self::types::AppEvent;
|
||||||
|
|
||||||
/// Web gateway channel implementing the Channel trait.
|
/// Web gateway channel implementing the Channel trait.
|
||||||
pub struct GatewayChannel {
|
pub struct GatewayChannel {
|
||||||
config: GatewayConfig,
|
config: GatewayConfig,
|
||||||
state: Arc<GatewayState>,
|
state: Arc<GatewayState>,
|
||||||
/// The actual auth token in use (generated or from config).
|
/// Multi-user auth state (replaces bare auth_token).
|
||||||
auth_token: String,
|
auth: MultiAuthState,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl GatewayChannel {
|
impl GatewayChannel {
|
||||||
/// Create a new gateway channel.
|
/// Create a new gateway channel.
|
||||||
///
|
///
|
||||||
/// If no auth token is configured, generates a random one and prints it.
|
/// If no auth token is configured, generates a random one and prints it.
|
||||||
|
/// Builds a single-user `MultiAuthState` from the config.
|
||||||
pub fn new(config: GatewayConfig) -> Self {
|
pub fn new(config: GatewayConfig) -> Self {
|
||||||
let auth_token = config.auth_token.clone().unwrap_or_else(|| {
|
let auth_token = config.auth_token.clone().unwrap_or_else(|| {
|
||||||
use rand::RngCore;
|
use rand::RngCore;
|
||||||
@@ -77,10 +82,13 @@ impl GatewayChannel {
|
|||||||
bytes.iter().map(|b| format!("{b:02x}")).collect()
|
bytes.iter().map(|b| format!("{b:02x}")).collect()
|
||||||
});
|
});
|
||||||
|
|
||||||
|
let auth = MultiAuthState::single(auth_token, config.user_id.clone());
|
||||||
|
|
||||||
let state = Arc::new(GatewayState {
|
let state = Arc::new(GatewayState {
|
||||||
msg_tx: tokio::sync::RwLock::new(None),
|
msg_tx: tokio::sync::RwLock::new(None),
|
||||||
sse: SseManager::new(),
|
sse: Arc::new(SseManager::new()),
|
||||||
workspace: None,
|
workspace: None,
|
||||||
|
workspace_pool: None,
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
log_broadcaster: None,
|
log_broadcaster: None,
|
||||||
log_level_handle: None,
|
log_level_handle: None,
|
||||||
@@ -90,13 +98,14 @@ impl GatewayChannel {
|
|||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
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),
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||||
llm_provider: None,
|
llm_provider: None,
|
||||||
skill_registry: None,
|
skill_registry: None,
|
||||||
skill_catalog: None,
|
skill_catalog: None,
|
||||||
chat_rate_limiter: server::RateLimiter::new(30, 60),
|
chat_rate_limiter: server::PerUserRateLimiter::new(30, 60),
|
||||||
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
||||||
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||||
registry_entries: Vec::new(),
|
registry_entries: Vec::new(),
|
||||||
@@ -109,7 +118,63 @@ impl GatewayChannel {
|
|||||||
Self {
|
Self {
|
||||||
config,
|
config,
|
||||||
state,
|
state,
|
||||||
auth_token,
|
auth,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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 {
|
||||||
|
msg_tx: tokio::sync::RwLock::new(None),
|
||||||
|
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: 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,
|
||||||
|
skill_registry: None,
|
||||||
|
skill_catalog: None,
|
||||||
|
chat_rate_limiter: server::PerUserRateLimiter::new(30, 60),
|
||||||
|
oauth_rate_limiter: server::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(),
|
||||||
|
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||||
|
active_config: server::ActiveConfigSnapshot::default(),
|
||||||
|
});
|
||||||
|
|
||||||
|
Self {
|
||||||
|
config,
|
||||||
|
state,
|
||||||
|
auth,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -118,8 +183,9 @@ impl GatewayChannel {
|
|||||||
let mut new_state = GatewayState {
|
let mut new_state = GatewayState {
|
||||||
msg_tx: tokio::sync::RwLock::new(None),
|
msg_tx: tokio::sync::RwLock::new(None),
|
||||||
// Preserve the existing broadcast channel so sender handles remain valid.
|
// Preserve the existing broadcast channel so sender handles remain valid.
|
||||||
sse: SseManager::from_sender(self.state.sse.sender()),
|
sse: Arc::new(SseManager::from_sender(self.state.sse.sender())),
|
||||||
workspace: self.state.workspace.clone(),
|
workspace: self.state.workspace.clone(),
|
||||||
|
workspace_pool: self.state.workspace_pool.clone(),
|
||||||
session_manager: self.state.session_manager.clone(),
|
session_manager: self.state.session_manager.clone(),
|
||||||
log_broadcaster: self.state.log_broadcaster.clone(),
|
log_broadcaster: self.state.log_broadcaster.clone(),
|
||||||
log_level_handle: self.state.log_level_handle.clone(),
|
log_level_handle: self.state.log_level_handle.clone(),
|
||||||
@@ -129,13 +195,14 @@ impl GatewayChannel {
|
|||||||
job_manager: self.state.job_manager.clone(),
|
job_manager: self.state.job_manager.clone(),
|
||||||
prompt_queue: self.state.prompt_queue.clone(),
|
prompt_queue: self.state.prompt_queue.clone(),
|
||||||
scheduler: self.state.scheduler.clone(),
|
scheduler: self.state.scheduler.clone(),
|
||||||
user_id: self.state.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),
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
ws_tracker: self.state.ws_tracker.clone(),
|
ws_tracker: self.state.ws_tracker.clone(),
|
||||||
llm_provider: self.state.llm_provider.clone(),
|
llm_provider: self.state.llm_provider.clone(),
|
||||||
skill_registry: self.state.skill_registry.clone(),
|
skill_registry: self.state.skill_registry.clone(),
|
||||||
skill_catalog: self.state.skill_catalog.clone(),
|
skill_catalog: self.state.skill_catalog.clone(),
|
||||||
chat_rate_limiter: server::RateLimiter::new(30, 60),
|
chat_rate_limiter: server::PerUserRateLimiter::new(30, 60),
|
||||||
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
||||||
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||||
registry_entries: self.state.registry_entries.clone(),
|
registry_entries: self.state.registry_entries.clone(),
|
||||||
@@ -260,9 +327,15 @@ impl GatewayChannel {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get the auth token (for printing to console on startup).
|
/// Inject the per-user workspace pool for multi-user mode.
|
||||||
|
pub fn with_workspace_pool(mut self, pool: Arc<server::WorkspacePool>) -> Self {
|
||||||
|
self.rebuild_state(|s| s.workspace_pool = Some(pool));
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the first auth token (for printing to console on startup).
|
||||||
pub fn auth_token(&self) -> &str {
|
pub fn auth_token(&self) -> &str {
|
||||||
&self.auth_token
|
self.auth.first_token().unwrap_or("")
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get a reference to the shared gateway state (for the agent to push SSE events).
|
/// Get a reference to the shared gateway state (for the agent to push SSE events).
|
||||||
@@ -291,7 +364,7 @@ impl Channel for GatewayChannel {
|
|||||||
),
|
),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
|
server::start_server(addr, self.state.clone(), self.auth.clone()).await?;
|
||||||
|
|
||||||
Ok(Box::pin(ReceiverStream::new(rx)))
|
Ok(Box::pin(ReceiverStream::new(rx)))
|
||||||
}
|
}
|
||||||
@@ -311,10 +384,13 @@ impl Channel for GatewayChannel {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
self.state.sse.broadcast(SseEvent::Response {
|
self.state.sse.broadcast_for_user(
|
||||||
content: response.content,
|
&msg.user_id,
|
||||||
thread_id,
|
AppEvent::Response {
|
||||||
});
|
content: response.content,
|
||||||
|
thread_id,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -329,11 +405,11 @@ impl Channel for GatewayChannel {
|
|||||||
.and_then(|v| v.as_str())
|
.and_then(|v| v.as_str())
|
||||||
.map(String::from);
|
.map(String::from);
|
||||||
let event = match status {
|
let event = match status {
|
||||||
StatusUpdate::Thinking(msg) => SseEvent::Thinking {
|
StatusUpdate::Thinking(msg) => AppEvent::Thinking {
|
||||||
message: msg,
|
message: msg,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
StatusUpdate::ToolStarted { name } => SseEvent::ToolStarted {
|
StatusUpdate::ToolStarted { name } => AppEvent::ToolStarted {
|
||||||
name,
|
name,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
@@ -342,23 +418,23 @@ impl Channel for GatewayChannel {
|
|||||||
success,
|
success,
|
||||||
error,
|
error,
|
||||||
parameters,
|
parameters,
|
||||||
} => SseEvent::ToolCompleted {
|
} => AppEvent::ToolCompleted {
|
||||||
name,
|
name,
|
||||||
success,
|
success,
|
||||||
error,
|
error,
|
||||||
parameters,
|
parameters,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
StatusUpdate::ToolResult { name, preview } => SseEvent::ToolResult {
|
StatusUpdate::ToolResult { name, preview } => AppEvent::ToolResult {
|
||||||
name,
|
name,
|
||||||
preview,
|
preview,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
StatusUpdate::StreamChunk(content) => SseEvent::StreamChunk {
|
StatusUpdate::StreamChunk(content) => AppEvent::StreamChunk {
|
||||||
content,
|
content,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
StatusUpdate::Status(msg) => SseEvent::Status {
|
StatusUpdate::Status(msg) => AppEvent::Status {
|
||||||
message: msg,
|
message: msg,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
@@ -366,7 +442,7 @@ impl Channel for GatewayChannel {
|
|||||||
job_id,
|
job_id,
|
||||||
title,
|
title,
|
||||||
browse_url,
|
browse_url,
|
||||||
} => SseEvent::JobStarted {
|
} => AppEvent::JobStarted {
|
||||||
job_id,
|
job_id,
|
||||||
title,
|
title,
|
||||||
browse_url,
|
browse_url,
|
||||||
@@ -377,7 +453,7 @@ impl Channel for GatewayChannel {
|
|||||||
description,
|
description,
|
||||||
parameters,
|
parameters,
|
||||||
allow_always,
|
allow_always,
|
||||||
} => SseEvent::ApprovalNeeded {
|
} => AppEvent::ApprovalNeeded {
|
||||||
request_id,
|
request_id,
|
||||||
tool_name,
|
tool_name,
|
||||||
description,
|
description,
|
||||||
@@ -391,7 +467,7 @@ impl Channel for GatewayChannel {
|
|||||||
instructions,
|
instructions,
|
||||||
auth_url,
|
auth_url,
|
||||||
setup_url,
|
setup_url,
|
||||||
} => SseEvent::AuthRequired {
|
} => AppEvent::AuthRequired {
|
||||||
extension_name,
|
extension_name,
|
||||||
instructions,
|
instructions,
|
||||||
auth_url,
|
auth_url,
|
||||||
@@ -401,25 +477,39 @@ impl Channel for GatewayChannel {
|
|||||||
extension_name,
|
extension_name,
|
||||||
success,
|
success,
|
||||||
message,
|
message,
|
||||||
} => SseEvent::AuthCompleted {
|
} => AppEvent::AuthCompleted {
|
||||||
extension_name,
|
extension_name,
|
||||||
success,
|
success,
|
||||||
message,
|
message,
|
||||||
},
|
},
|
||||||
StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated {
|
StatusUpdate::ImageGenerated { data_url, path } => AppEvent::ImageGenerated {
|
||||||
data_url,
|
data_url,
|
||||||
path,
|
path,
|
||||||
thread_id: thread_id.clone(),
|
thread_id: thread_id.clone(),
|
||||||
},
|
},
|
||||||
StatusUpdate::Suggestions { suggestions } => SseEvent::Suggestions {
|
StatusUpdate::Suggestions { suggestions } => AppEvent::Suggestions {
|
||||||
suggestions,
|
suggestions,
|
||||||
|
thread_id: thread_id.clone(),
|
||||||
|
},
|
||||||
|
StatusUpdate::ReasoningUpdate {
|
||||||
|
narrative,
|
||||||
|
decisions,
|
||||||
|
} => AppEvent::ReasoningUpdate {
|
||||||
|
narrative,
|
||||||
|
decisions: decisions
|
||||||
|
.into_iter()
|
||||||
|
.map(|d| crate::channels::web::types::ToolDecisionDto {
|
||||||
|
tool_name: d.tool_name,
|
||||||
|
rationale: d.rationale,
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
thread_id,
|
thread_id,
|
||||||
},
|
},
|
||||||
StatusUpdate::TurnCost {
|
StatusUpdate::TurnCost {
|
||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
cost_usd,
|
cost_usd,
|
||||||
} => SseEvent::TurnCost {
|
} => AppEvent::TurnCost {
|
||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
cost_usd,
|
cost_usd,
|
||||||
@@ -427,13 +517,21 @@ impl Channel for GatewayChannel {
|
|||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
self.state.sse.broadcast(event);
|
// Scope events to the user when user_id is available in metadata.
|
||||||
|
// When user_id is missing (heartbeat, routines), events go to all
|
||||||
|
// subscribers. In multi-tenant mode this leaks status across users.
|
||||||
|
if let Some(uid) = metadata.get("user_id").and_then(|v| v.as_str()) {
|
||||||
|
self.state.sse.broadcast_for_user(uid, event);
|
||||||
|
} else {
|
||||||
|
tracing::debug!("Status event missing user_id in metadata; broadcasting globally");
|
||||||
|
self.state.sse.broadcast(event);
|
||||||
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn broadcast(
|
async fn broadcast(
|
||||||
&self,
|
&self,
|
||||||
_user_id: &str,
|
user_id: &str,
|
||||||
response: OutgoingResponse,
|
response: OutgoingResponse,
|
||||||
) -> Result<(), ChannelError> {
|
) -> Result<(), ChannelError> {
|
||||||
let thread_id = match response.thread_id {
|
let thread_id = match response.thread_id {
|
||||||
@@ -445,10 +543,13 @@ impl Channel for GatewayChannel {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
self.state.sse.broadcast(SseEvent::Response {
|
self.state.sse.broadcast_for_user(
|
||||||
content: response.content,
|
user_id,
|
||||||
thread_id,
|
AppEvent::Response {
|
||||||
});
|
content: response.content,
|
||||||
|
thread_id,
|
||||||
|
},
|
||||||
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -231,6 +231,7 @@ pub fn convert_messages(messages: &[OpenAiMessage]) -> Result<Vec<ChatMessage>,
|
|||||||
name: tc.function.name.clone(),
|
name: tc.function.name.clone(),
|
||||||
arguments: serde_json::from_str(&tc.function.arguments)
|
arguments: serde_json::from_str(&tc.function.arguments)
|
||||||
.unwrap_or(serde_json::Value::Object(Default::default())),
|
.unwrap_or(serde_json::Value::Object(Default::default())),
|
||||||
|
reasoning: None,
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
Ok(ChatMessage::assistant_with_tool_calls(
|
Ok(ChatMessage::assistant_with_tool_calls(
|
||||||
@@ -463,9 +464,10 @@ fn build_tool_request(
|
|||||||
|
|
||||||
pub async fn chat_completions_handler(
|
pub async fn chat_completions_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
super::auth::AuthenticatedUser(user): super::auth::AuthenticatedUser,
|
||||||
Json(req): Json<OpenAiChatRequest>,
|
Json(req): Json<OpenAiChatRequest>,
|
||||||
) -> Result<impl IntoResponse, (StatusCode, Json<OpenAiErrorResponse>)> {
|
) -> Result<impl IntoResponse, (StatusCode, Json<OpenAiErrorResponse>)> {
|
||||||
if !state.chat_rate_limiter.check() {
|
if !state.chat_rate_limiter.check(&user.user_id) {
|
||||||
return Err(openai_error(
|
return Err(openai_error(
|
||||||
StatusCode::TOO_MANY_REQUESTS,
|
StatusCode::TOO_MANY_REQUESTS,
|
||||||
"Rate limit exceeded. Please try again later.",
|
"Rate limit exceeded. Please try again later.",
|
||||||
@@ -953,6 +955,7 @@ mod tests {
|
|||||||
id: "call_abc".to_string(),
|
id: "call_abc".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "rust"}),
|
arguments: serde_json::json!({"query": "rust"}),
|
||||||
|
reasoning: None,
|
||||||
}];
|
}];
|
||||||
|
|
||||||
let converted = convert_tool_calls_to_openai(&calls);
|
let converted = convert_tool_calls_to_openai(&calls);
|
||||||
|
|||||||
+528
-377
File diff suppressed because it is too large
Load Diff
+138
-61
@@ -11,15 +11,31 @@ use tokio::sync::broadcast;
|
|||||||
use tokio_stream::StreamExt;
|
use tokio_stream::StreamExt;
|
||||||
use tokio_stream::wrappers::BroadcastStream;
|
use tokio_stream::wrappers::BroadcastStream;
|
||||||
|
|
||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::AppEvent;
|
||||||
|
|
||||||
/// Maximum number of concurrent SSE/WebSocket connections.
|
/// Maximum number of concurrent SSE/WebSocket connections.
|
||||||
/// Prevents resource exhaustion from connection flooding.
|
/// Prevents resource exhaustion from connection flooding.
|
||||||
const MAX_CONNECTIONS: u64 = 100;
|
const MAX_CONNECTIONS: u64 = 100;
|
||||||
|
|
||||||
|
/// Envelope for broadcast events: carries an optional user scope.
|
||||||
|
///
|
||||||
|
/// `user_id = None` means the event is global (e.g. Heartbeat) and delivered
|
||||||
|
/// to all subscribers. `user_id = Some(id)` means the event is only delivered
|
||||||
|
/// to subscribers that match that user_id.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub(crate) struct ScopedEvent {
|
||||||
|
pub(crate) user_id: Option<String>,
|
||||||
|
pub(crate) event: AppEvent,
|
||||||
|
}
|
||||||
|
|
||||||
/// Manages SSE broadcast to all connected browser tabs.
|
/// Manages SSE broadcast to all connected browser tabs.
|
||||||
|
///
|
||||||
|
/// In multi-user mode, events are scoped by user_id so that each subscriber
|
||||||
|
/// only receives events intended for their user (plus global events like
|
||||||
|
/// Heartbeat). In single-user mode, all events are delivered to all subscribers
|
||||||
|
/// (backwards compatible).
|
||||||
pub struct SseManager {
|
pub struct SseManager {
|
||||||
tx: broadcast::Sender<SseEvent>,
|
tx: broadcast::Sender<ScopedEvent>,
|
||||||
connection_count: Arc<AtomicU64>,
|
connection_count: Arc<AtomicU64>,
|
||||||
max_connections: u64,
|
max_connections: u64,
|
||||||
}
|
}
|
||||||
@@ -45,7 +61,7 @@ impl SseManager {
|
|||||||
/// only be called before the server starts accepting connections (i.e.,
|
/// only be called before the server starts accepting connections (i.e.,
|
||||||
/// during startup wiring). Calling it after connections are established
|
/// during startup wiring). Calling it after connections are established
|
||||||
/// will break connection tracking and allow exceeding `MAX_CONNECTIONS`.
|
/// will break connection tracking and allow exceeding `MAX_CONNECTIONS`.
|
||||||
pub fn from_sender(tx: broadcast::Sender<SseEvent>) -> Self {
|
pub(crate) fn from_sender(tx: broadcast::Sender<ScopedEvent>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
tx,
|
tx,
|
||||||
connection_count: Arc::new(AtomicU64::new(0)),
|
connection_count: Arc::new(AtomicU64::new(0)),
|
||||||
@@ -53,15 +69,28 @@ impl SseManager {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Broadcast an event to all connected clients.
|
/// Get a clone of the broadcast sender for use by other components.
|
||||||
pub fn broadcast(&self, event: SseEvent) {
|
pub(crate) fn sender(&self) -> broadcast::Sender<ScopedEvent> {
|
||||||
// Ignore send errors (no receivers is fine)
|
self.tx.clone()
|
||||||
let _ = self.tx.send(event);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get a clone of the broadcast sender for use by other components.
|
/// Broadcast an event to all connected clients (global/unscoped).
|
||||||
pub fn sender(&self) -> broadcast::Sender<SseEvent> {
|
pub fn broadcast(&self, event: AppEvent) {
|
||||||
self.tx.clone()
|
let _ = self.tx.send(ScopedEvent {
|
||||||
|
user_id: None,
|
||||||
|
event,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Broadcast an event scoped to a specific user.
|
||||||
|
///
|
||||||
|
/// Only subscribers for this user_id (or unscoped subscribers) will
|
||||||
|
/// receive the event.
|
||||||
|
pub fn broadcast_for_user(&self, user_id: &str, event: AppEvent) {
|
||||||
|
let _ = self.tx.send(ScopedEvent {
|
||||||
|
user_id: Some(user_id.to_string()),
|
||||||
|
event,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get current number of active connections.
|
/// Get current number of active connections.
|
||||||
@@ -71,11 +100,15 @@ impl SseManager {
|
|||||||
|
|
||||||
/// Create a raw broadcast subscription for non-SSE consumers (e.g. WebSocket).
|
/// Create a raw broadcast subscription for non-SSE consumers (e.g. WebSocket).
|
||||||
///
|
///
|
||||||
/// Returns a stream of `SseEvent` values and increments/decrements the
|
/// When `user_id` is `Some`, only events scoped to that user (or global
|
||||||
/// connection counter on creation/drop, just like `subscribe()` does for SSE.
|
/// events) are delivered. When `None`, all events are delivered (single-user
|
||||||
|
/// backwards compatibility).
|
||||||
///
|
///
|
||||||
/// Returns `None` if the maximum connection limit has been reached.
|
/// Returns `None` if the maximum connection limit has been reached.
|
||||||
pub fn subscribe_raw(&self) -> Option<impl Stream<Item = SseEvent> + Send + 'static + use<>> {
|
pub fn subscribe_raw(
|
||||||
|
&self,
|
||||||
|
user_id: Option<String>,
|
||||||
|
) -> Option<impl Stream<Item = AppEvent> + Send + 'static + use<>> {
|
||||||
// Atomically increment only if below the limit. This prevents
|
// Atomically increment only if below the limit. This prevents
|
||||||
// concurrent callers from overshooting max_connections.
|
// concurrent callers from overshooting max_connections.
|
||||||
let counter = Arc::clone(&self.connection_count);
|
let counter = Arc::clone(&self.connection_count);
|
||||||
@@ -91,7 +124,19 @@ impl SseManager {
|
|||||||
.ok()?;
|
.ok()?;
|
||||||
let rx = self.tx.subscribe();
|
let rx = self.tx.subscribe();
|
||||||
|
|
||||||
let stream = BroadcastStream::new(rx).filter_map(|result| result.ok());
|
let stream = BroadcastStream::new(rx).filter_map(move |result| match result {
|
||||||
|
Ok(scoped) => {
|
||||||
|
// Global events (user_id=None) always pass through.
|
||||||
|
// Scoped events only pass if the subscriber matches (or subscriber is unscoped).
|
||||||
|
match (&user_id, &scoped.user_id) {
|
||||||
|
(_, None) => Some(scoped.event), // global -> all
|
||||||
|
(None, _) => Some(scoped.event), // unscoped subscriber -> all
|
||||||
|
(Some(sub), Some(ev)) if sub == ev => Some(scoped.event), // match
|
||||||
|
_ => None, // different user -> skip
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(_) => None,
|
||||||
|
});
|
||||||
|
|
||||||
Some(CountedStream {
|
Some(CountedStream {
|
||||||
inner: stream,
|
inner: stream,
|
||||||
@@ -101,9 +146,13 @@ impl SseManager {
|
|||||||
|
|
||||||
/// Create a new SSE stream for a client connection.
|
/// Create a new SSE stream for a client connection.
|
||||||
///
|
///
|
||||||
|
/// When `user_id` is `Some`, only events for that user (or global events)
|
||||||
|
/// are delivered. When `None`, all events are delivered.
|
||||||
|
///
|
||||||
/// Returns `None` if the maximum connection limit has been reached.
|
/// Returns `None` if the maximum connection limit has been reached.
|
||||||
pub fn subscribe(
|
pub fn subscribe(
|
||||||
&self,
|
&self,
|
||||||
|
user_id: Option<String>,
|
||||||
) -> Option<Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>>> {
|
) -> Option<Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>>> {
|
||||||
// Atomically increment only if below the limit.
|
// Atomically increment only if below the limit.
|
||||||
let counter = Arc::clone(&self.connection_count);
|
let counter = Arc::clone(&self.connection_count);
|
||||||
@@ -120,34 +169,25 @@ impl SseManager {
|
|||||||
let rx = self.tx.subscribe();
|
let rx = self.tx.subscribe();
|
||||||
|
|
||||||
let stream = BroadcastStream::new(rx)
|
let stream = BroadcastStream::new(rx)
|
||||||
.filter_map(|result| result.ok())
|
.filter_map(move |result| match result {
|
||||||
.map(|event| {
|
Ok(scoped) => match (&user_id, &scoped.user_id) {
|
||||||
let data = serde_json::to_string(&event).unwrap_or_default();
|
(_, None) => Some(scoped.event),
|
||||||
let event_type = match &event {
|
(None, _) => Some(scoped.event),
|
||||||
SseEvent::Response { .. } => "response",
|
(Some(sub), Some(ev)) if sub == ev => Some(scoped.event),
|
||||||
SseEvent::Thinking { .. } => "thinking",
|
_ => None,
|
||||||
SseEvent::ToolStarted { .. } => "tool_started",
|
},
|
||||||
SseEvent::ToolCompleted { .. } => "tool_completed",
|
Err(_) => None,
|
||||||
SseEvent::ToolResult { .. } => "tool_result",
|
})
|
||||||
SseEvent::StreamChunk { .. } => "stream_chunk",
|
.filter_map(|event| {
|
||||||
SseEvent::Status { .. } => "status",
|
let data = match serde_json::to_string(&event) {
|
||||||
SseEvent::ApprovalNeeded { .. } => "approval_needed",
|
Ok(s) => s,
|
||||||
SseEvent::AuthRequired { .. } => "auth_required",
|
Err(e) => {
|
||||||
SseEvent::AuthCompleted { .. } => "auth_completed",
|
tracing::warn!("Failed to serialize SSE event: {}", e);
|
||||||
SseEvent::Error { .. } => "error",
|
return None;
|
||||||
SseEvent::JobStarted { .. } => "job_started",
|
}
|
||||||
SseEvent::JobMessage { .. } => "job_message",
|
|
||||||
SseEvent::JobToolUse { .. } => "job_tool_use",
|
|
||||||
SseEvent::JobToolResult { .. } => "job_tool_result",
|
|
||||||
SseEvent::JobStatus { .. } => "job_status",
|
|
||||||
SseEvent::JobResult { .. } => "job_result",
|
|
||||||
SseEvent::Heartbeat => "heartbeat",
|
|
||||||
SseEvent::ImageGenerated { .. } => "image_generated",
|
|
||||||
SseEvent::Suggestions { .. } => "suggestions",
|
|
||||||
SseEvent::TurnCost { .. } => "turn_cost",
|
|
||||||
SseEvent::ExtensionStatus { .. } => "extension_status",
|
|
||||||
};
|
};
|
||||||
Ok(Event::default().event(event_type).data(data))
|
let event_type = event.event_type();
|
||||||
|
Some(Ok(Event::default().event(event_type).data(data)))
|
||||||
});
|
});
|
||||||
|
|
||||||
// Wrap in a stream that decrements on drop
|
// Wrap in a stream that decrements on drop
|
||||||
@@ -209,24 +249,22 @@ mod tests {
|
|||||||
fn test_broadcast_without_receivers() {
|
fn test_broadcast_without_receivers() {
|
||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
// Should not panic even with no receivers
|
// Should not panic even with no receivers
|
||||||
manager.broadcast(SseEvent::Heartbeat);
|
manager.broadcast(AppEvent::Heartbeat);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_broadcast_to_receiver() {
|
async fn test_broadcast_to_receiver() {
|
||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
let mut rx = BroadcastStream::new(manager.tx.subscribe());
|
let mut stream = Box::pin(manager.subscribe_raw(None).expect("should subscribe"));
|
||||||
|
|
||||||
manager.broadcast(SseEvent::Status {
|
manager.broadcast(AppEvent::Status {
|
||||||
message: "test".to_string(),
|
message: "test".to_string(),
|
||||||
thread_id: None,
|
thread_id: None,
|
||||||
});
|
});
|
||||||
|
|
||||||
let event = rx.next().await;
|
let event = stream.next().await.unwrap();
|
||||||
assert!(event.is_some());
|
|
||||||
let event = event.unwrap().unwrap();
|
|
||||||
match event {
|
match event {
|
||||||
SseEvent::Status { message, .. } => assert_eq!(message, "test"),
|
AppEvent::Status { message, .. } => assert_eq!(message, "test"),
|
||||||
_ => panic!("unexpected event type"),
|
_ => panic!("unexpected event type"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -234,18 +272,18 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_subscribe_raw_receives_events() {
|
async fn test_subscribe_raw_receives_events() {
|
||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
let mut stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
let mut stream = Box::pin(manager.subscribe_raw(None).expect("should subscribe"));
|
||||||
|
|
||||||
assert_eq!(manager.connection_count(), 1);
|
assert_eq!(manager.connection_count(), 1);
|
||||||
|
|
||||||
manager.broadcast(SseEvent::Thinking {
|
manager.broadcast(AppEvent::Thinking {
|
||||||
message: "working".to_string(),
|
message: "working".to_string(),
|
||||||
thread_id: None,
|
thread_id: None,
|
||||||
});
|
});
|
||||||
|
|
||||||
let event = stream.next().await.unwrap();
|
let event = stream.next().await.unwrap();
|
||||||
match event {
|
match event {
|
||||||
SseEvent::Thinking { message, .. } => assert_eq!(message, "working"),
|
AppEvent::Thinking { message, .. } => assert_eq!(message, "working"),
|
||||||
_ => panic!("Expected Thinking event"),
|
_ => panic!("Expected Thinking event"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -254,7 +292,7 @@ mod tests {
|
|||||||
async fn test_subscribe_raw_decrements_on_drop() {
|
async fn test_subscribe_raw_decrements_on_drop() {
|
||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
{
|
{
|
||||||
let _stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
let _stream = Box::pin(manager.subscribe_raw(None).expect("should subscribe"));
|
||||||
assert_eq!(manager.connection_count(), 1);
|
assert_eq!(manager.connection_count(), 1);
|
||||||
}
|
}
|
||||||
// Stream dropped, counter should decrement
|
// Stream dropped, counter should decrement
|
||||||
@@ -264,16 +302,16 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_subscribe_raw_multiple_subscribers() {
|
async fn test_subscribe_raw_multiple_subscribers() {
|
||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
let mut s1 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
let mut s1 = Box::pin(manager.subscribe_raw(None).expect("should subscribe"));
|
||||||
let mut s2 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
let mut s2 = Box::pin(manager.subscribe_raw(None).expect("should subscribe"));
|
||||||
assert_eq!(manager.connection_count(), 2);
|
assert_eq!(manager.connection_count(), 2);
|
||||||
|
|
||||||
manager.broadcast(SseEvent::Heartbeat);
|
manager.broadcast(AppEvent::Heartbeat);
|
||||||
|
|
||||||
let e1 = s1.next().await.unwrap();
|
let e1 = s1.next().await.unwrap();
|
||||||
let e2 = s2.next().await.unwrap();
|
let e2 = s2.next().await.unwrap();
|
||||||
assert!(matches!(e1, SseEvent::Heartbeat));
|
assert!(matches!(e1, AppEvent::Heartbeat));
|
||||||
assert!(matches!(e2, SseEvent::Heartbeat));
|
assert!(matches!(e2, AppEvent::Heartbeat));
|
||||||
|
|
||||||
drop(s1);
|
drop(s1);
|
||||||
assert_eq!(manager.connection_count(), 1);
|
assert_eq!(manager.connection_count(), 1);
|
||||||
@@ -286,12 +324,51 @@ mod tests {
|
|||||||
let mut manager = SseManager::new();
|
let mut manager = SseManager::new();
|
||||||
manager.max_connections = 2; // Low limit for testing
|
manager.max_connections = 2; // Low limit for testing
|
||||||
|
|
||||||
let _s1 = Box::pin(manager.subscribe_raw().expect("first should succeed"));
|
let _s1 = Box::pin(manager.subscribe_raw(None).expect("first should succeed"));
|
||||||
let _s2 = Box::pin(manager.subscribe_raw().expect("second should succeed"));
|
let _s2 = Box::pin(manager.subscribe_raw(None).expect("second should succeed"));
|
||||||
assert_eq!(manager.connection_count(), 2);
|
assert_eq!(manager.connection_count(), 2);
|
||||||
|
|
||||||
// Third should be rejected
|
// Third should be rejected
|
||||||
assert!(manager.subscribe_raw().is_none());
|
assert!(manager.subscribe_raw(None).is_none());
|
||||||
assert!(manager.subscribe().is_none());
|
assert!(manager.subscribe(None).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_scoped_events_filtered_by_user() {
|
||||||
|
let manager = SseManager::new();
|
||||||
|
let mut alice = Box::pin(
|
||||||
|
manager
|
||||||
|
.subscribe_raw(Some("alice".to_string()))
|
||||||
|
.expect("subscribe"),
|
||||||
|
);
|
||||||
|
let mut bob = Box::pin(
|
||||||
|
manager
|
||||||
|
.subscribe_raw(Some("bob".to_string()))
|
||||||
|
.expect("subscribe"),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Send event scoped to alice
|
||||||
|
manager.broadcast_for_user(
|
||||||
|
"alice",
|
||||||
|
AppEvent::Status {
|
||||||
|
message: "alice only".to_string(),
|
||||||
|
thread_id: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
// Send global event
|
||||||
|
manager.broadcast(AppEvent::Heartbeat);
|
||||||
|
|
||||||
|
// Alice gets her scoped event
|
||||||
|
let e = alice.next().await.unwrap();
|
||||||
|
assert!(matches!(e, AppEvent::Status { .. }));
|
||||||
|
|
||||||
|
// Alice also gets the global heartbeat
|
||||||
|
let e = alice.next().await.unwrap();
|
||||||
|
assert!(matches!(e, AppEvent::Heartbeat));
|
||||||
|
|
||||||
|
// Bob only gets the global heartbeat (alice's event was filtered)
|
||||||
|
let e = bob.next().await.unwrap(); // safety: test-only
|
||||||
|
assert!(matches!(e, AppEvent::Heartbeat)); // safety: test assertion
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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>'
|
+ '<th>Trigger</th><th>Started</th><th>Completed</th><th>Status</th><th>Summary</th><th>Tokens</th>'
|
||||||
+ '</tr></thead><tbody>';
|
+ '</tr></thead><tbody>';
|
||||||
for (const run of routine.recent_runs) {
|
for (const run of routine.recent_runs) {
|
||||||
const runStatusClass = run.status === 'Ok' ? 'completed'
|
const runStatusClass = run.status === 'ok' ? 'completed'
|
||||||
: run.status === 'Failed' ? 'failed'
|
: run.status === 'failed' ? 'failed'
|
||||||
: run.status === 'Attention' ? 'stuck'
|
: run.status === 'attention' ? 'stuck'
|
||||||
: 'in_progress';
|
: 'in_progress';
|
||||||
html += '<tr>'
|
html += '<tr>'
|
||||||
+ '<td>' + escapeHtml(run.trigger_type) + '</td>'
|
+ '<td>' + escapeHtml(run.trigger_type) + '</td>'
|
||||||
|
|||||||
@@ -10,7 +10,8 @@ use std::sync::Arc;
|
|||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
use crate::channels::IncomingMessage;
|
use crate::channels::IncomingMessage;
|
||||||
use crate::channels::web::server::{GatewayState, RateLimiter, start_server};
|
use crate::channels::web::auth::MultiAuthState;
|
||||||
|
use crate::channels::web::server::{GatewayState, PerUserRateLimiter, RateLimiter, start_server};
|
||||||
use crate::channels::web::sse::SseManager;
|
use crate::channels::web::sse::SseManager;
|
||||||
use crate::channels::web::ws::WsConnectionTracker;
|
use crate::channels::web::ws::WsConnectionTracker;
|
||||||
|
|
||||||
@@ -64,8 +65,9 @@ impl TestGatewayBuilder {
|
|||||||
pub fn build(self) -> Arc<GatewayState> {
|
pub fn build(self) -> Arc<GatewayState> {
|
||||||
Arc::new(GatewayState {
|
Arc::new(GatewayState {
|
||||||
msg_tx: tokio::sync::RwLock::new(self.msg_tx),
|
msg_tx: tokio::sync::RwLock::new(self.msg_tx),
|
||||||
sse: SseManager::new(),
|
sse: Arc::new(SseManager::new()),
|
||||||
workspace: None,
|
workspace: None,
|
||||||
|
workspace_pool: None,
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
log_broadcaster: None,
|
log_broadcaster: None,
|
||||||
log_level_handle: None,
|
log_level_handle: None,
|
||||||
@@ -74,14 +76,15 @@ impl TestGatewayBuilder {
|
|||||||
store: None,
|
store: None,
|
||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
user_id: self.user_id,
|
owner_id: self.user_id.clone(),
|
||||||
|
default_sender_id: self.user_id,
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||||
llm_provider: self.llm_provider,
|
llm_provider: self.llm_provider,
|
||||||
skill_registry: None,
|
skill_registry: None,
|
||||||
skill_catalog: None,
|
skill_catalog: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
chat_rate_limiter: RateLimiter::new(30, 60),
|
chat_rate_limiter: PerUserRateLimiter::new(30, 60),
|
||||||
oauth_rate_limiter: RateLimiter::new(10, 60),
|
oauth_rate_limiter: RateLimiter::new(10, 60),
|
||||||
webhook_rate_limiter: RateLimiter::new(10, 60),
|
webhook_rate_limiter: RateLimiter::new(10, 60),
|
||||||
registry_entries: Vec::new(),
|
registry_entries: Vec::new(),
|
||||||
@@ -98,11 +101,26 @@ impl TestGatewayBuilder {
|
|||||||
self,
|
self,
|
||||||
auth_token: &str,
|
auth_token: &str,
|
||||||
) -> Result<(SocketAddr, Arc<GatewayState>), crate::error::ChannelError> {
|
) -> Result<(SocketAddr, Arc<GatewayState>), crate::error::ChannelError> {
|
||||||
|
let auth = MultiAuthState::single(auth_token.to_string(), "test-user".to_string());
|
||||||
let state = self.build();
|
let state = self.build();
|
||||||
let addr: SocketAddr = "127.0.0.1:0"
|
let addr: SocketAddr = "127.0.0.1:0"
|
||||||
.parse()
|
.parse()
|
||||||
.expect("hard-coded address must parse");
|
.expect("hard-coded address must parse"); // safety: constant literal
|
||||||
let bound = start_server(addr, state.clone(), auth_token.to_string()).await?;
|
let bound = start_server(addr, state.clone(), auth).await?;
|
||||||
|
Ok((bound, state))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build the state and start a gateway server with multi-user auth.
|
||||||
|
/// Returns the bound address and the shared state.
|
||||||
|
pub async fn start_multi(
|
||||||
|
self,
|
||||||
|
auth: MultiAuthState,
|
||||||
|
) -> Result<(SocketAddr, Arc<GatewayState>), crate::error::ChannelError> {
|
||||||
|
let state = self.build();
|
||||||
|
let addr: SocketAddr = "127.0.0.1:0"
|
||||||
|
.parse()
|
||||||
|
.expect("hard-coded address must parse"); // safety: constant literal
|
||||||
|
let bound = start_server(addr, state.clone(), auth).await?;
|
||||||
Ok((bound, state))
|
Ok((bound, state))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,3 @@
|
|||||||
|
//! Integration tests for the web gateway module.
|
||||||
|
|
||||||
|
mod multi_tenant;
|
||||||
@@ -0,0 +1,833 @@
|
|||||||
|
//! Multi-tenant isolation tests for the web gateway.
|
||||||
|
//!
|
||||||
|
//! Tests cover workspace pool scoping, job handler isolation, and auth
|
||||||
|
//! enforcement on protected endpoints. Uses `LibSqlBackend::new_local()`
|
||||||
|
//! with a temporary directory for a real (but ephemeral) database.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use axum::Router;
|
||||||
|
use axum::body::Body;
|
||||||
|
use axum::http::{Method, Request, StatusCode};
|
||||||
|
use axum::middleware;
|
||||||
|
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,
|
||||||
|
};
|
||||||
|
use crate::channels::web::server::{
|
||||||
|
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
||||||
|
};
|
||||||
|
use crate::channels::web::sse::SseManager;
|
||||||
|
use crate::config::GatewayConfig;
|
||||||
|
|
||||||
|
// ── Helpers ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Create a two-user `MultiAuthState` for alice and bob.
|
||||||
|
fn two_user_auth() -> MultiAuthState {
|
||||||
|
let mut tokens = HashMap::new();
|
||||||
|
tokens.insert(
|
||||||
|
"tok-alice".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec!["shared".to_string()],
|
||||||
|
},
|
||||||
|
);
|
||||||
|
tokens.insert(
|
||||||
|
"tok-bob".to_string(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: "bob".to_string(),
|
||||||
|
workspace_read_scopes: vec!["shared".to_string(), "alice".to_string()],
|
||||||
|
},
|
||||||
|
);
|
||||||
|
MultiAuthState::multi(tokens)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a `GatewayState` with configurable store and prompt queue.
|
||||||
|
fn build_state(
|
||||||
|
store: Option<Arc<dyn crate::db::Database>>,
|
||||||
|
prompt_queue: Option<PromptQueue>,
|
||||||
|
) -> Arc<GatewayState> {
|
||||||
|
Arc::new(GatewayState {
|
||||||
|
msg_tx: tokio::sync::RwLock::new(None),
|
||||||
|
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,
|
||||||
|
job_manager: None,
|
||||||
|
prompt_queue,
|
||||||
|
owner_id: "test".to_string(),
|
||||||
|
default_sender_id: "test".to_string(),
|
||||||
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
|
ws_tracker: None,
|
||||||
|
llm_provider: None,
|
||||||
|
skill_registry: None,
|
||||||
|
skill_catalog: None,
|
||||||
|
scheduler: 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: ActiveConfigSnapshot::default(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
/// deleted when the guard is dropped.
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
async fn test_db() -> (Arc<dyn crate::db::Database>, tempfile::TempDir) {
|
||||||
|
use crate::db::Database;
|
||||||
|
let dir = tempfile::tempdir().expect("failed to create temp dir"); // safety: test-only
|
||||||
|
let path = dir.path().join("test.db");
|
||||||
|
let backend = crate::db::libsql::LibSqlBackend::new_local(&path)
|
||||||
|
.await
|
||||||
|
.expect("failed to create test LibSqlBackend"); // safety: test-only
|
||||||
|
backend
|
||||||
|
.run_migrations()
|
||||||
|
.await
|
||||||
|
.expect("failed to run migrations"); // safety: test-only
|
||||||
|
(Arc::new(backend) as Arc<dyn crate::db::Database>, dir)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a minimal Routine for testing.
|
||||||
|
fn make_routine(user_id: &str, name: &str) -> crate::agent::routine::Routine {
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
crate::agent::routine::Routine {
|
||||||
|
id: Uuid::new_v4(),
|
||||||
|
name: name.to_string(),
|
||||||
|
description: format!("Test routine: {name}"),
|
||||||
|
user_id: user_id.to_string(),
|
||||||
|
enabled: true,
|
||||||
|
trigger: crate::agent::routine::Trigger::Cron {
|
||||||
|
schedule: "0 9 * * *".to_string(),
|
||||||
|
timezone: None,
|
||||||
|
},
|
||||||
|
action: crate::agent::routine::RoutineAction::Lightweight {
|
||||||
|
prompt: "hello".to_string(),
|
||||||
|
context_paths: vec![],
|
||||||
|
max_tokens: 1024,
|
||||||
|
use_tools: false,
|
||||||
|
max_tool_rounds: 3,
|
||||||
|
},
|
||||||
|
guardrails: crate::agent::routine::RoutineGuardrails {
|
||||||
|
cooldown: Duration::from_secs(60),
|
||||||
|
max_concurrent: 1,
|
||||||
|
dedup_window: None,
|
||||||
|
},
|
||||||
|
notify: crate::agent::routine::NotifyConfig {
|
||||||
|
channel: None,
|
||||||
|
user: None,
|
||||||
|
on_success: false,
|
||||||
|
on_failure: true,
|
||||||
|
on_attention: true,
|
||||||
|
},
|
||||||
|
last_run_at: None,
|
||||||
|
next_fire_at: None,
|
||||||
|
run_count: 0,
|
||||||
|
consecutive_failures: 0,
|
||||||
|
state: serde_json::json!({}),
|
||||||
|
created_at: now,
|
||||||
|
updated_at: now,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a minimal SandboxJobRecord for testing.
|
||||||
|
fn make_sandbox_job(user_id: &str, task: &str) -> crate::history::SandboxJobRecord {
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
crate::history::SandboxJobRecord {
|
||||||
|
id: Uuid::new_v4(),
|
||||||
|
task: task.to_string(),
|
||||||
|
status: "completed".to_string(),
|
||||||
|
user_id: user_id.to_string(),
|
||||||
|
project_dir: format!("/tmp/test-{}", Uuid::new_v4()),
|
||||||
|
success: Some(true),
|
||||||
|
failure_reason: None,
|
||||||
|
created_at: now,
|
||||||
|
started_at: Some(now),
|
||||||
|
completed_at: Some(now),
|
||||||
|
credential_grants_json: "[]".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
// WorkspacePool Tests
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
mod workspace_pool {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::{WorkspaceConfig, WorkspaceSearchConfig};
|
||||||
|
use crate::workspace::EmbeddingCacheConfig;
|
||||||
|
use crate::workspace::layer::MemoryLayer;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_applies_search_config() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
let search_config = WorkspaceSearchConfig {
|
||||||
|
rrf_k: 42,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let pool = WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
EmbeddingCacheConfig::default(),
|
||||||
|
search_config,
|
||||||
|
WorkspaceConfig::default(),
|
||||||
|
);
|
||||||
|
let identity = UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
};
|
||||||
|
let ws = pool.get_or_create(&identity).await;
|
||||||
|
assert_eq!(ws.user_id(), "alice");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_applies_memory_layers() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
let layers = vec![MemoryLayer {
|
||||||
|
name: "shared-layer".to_string(),
|
||||||
|
scope: "shared".to_string(),
|
||||||
|
writable: false,
|
||||||
|
sensitivity: Default::default(),
|
||||||
|
}];
|
||||||
|
let ws_config = WorkspaceConfig {
|
||||||
|
memory_layers: layers,
|
||||||
|
read_scopes: vec![],
|
||||||
|
};
|
||||||
|
let pool = WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
EmbeddingCacheConfig::default(),
|
||||||
|
WorkspaceSearchConfig::default(),
|
||||||
|
ws_config,
|
||||||
|
);
|
||||||
|
let identity = UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
};
|
||||||
|
let ws = pool.get_or_create(&identity).await;
|
||||||
|
// Memory layer scope "shared" should appear in read_user_ids.
|
||||||
|
assert!(
|
||||||
|
ws.read_user_ids().contains(&"shared".to_string()),
|
||||||
|
"expected 'shared' in read_user_ids, got {:?}",
|
||||||
|
ws.read_user_ids()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_applies_identity_read_scopes() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
let pool = WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
EmbeddingCacheConfig::default(),
|
||||||
|
WorkspaceSearchConfig::default(),
|
||||||
|
WorkspaceConfig::default(),
|
||||||
|
);
|
||||||
|
let identity = UserIdentity {
|
||||||
|
user_id: "bob".to_string(),
|
||||||
|
workspace_read_scopes: vec!["alice".to_string(), "shared".to_string()],
|
||||||
|
};
|
||||||
|
let ws = pool.get_or_create(&identity).await;
|
||||||
|
assert_eq!(ws.user_id(), "bob");
|
||||||
|
assert!(
|
||||||
|
ws.read_user_ids().contains(&"alice".to_string()),
|
||||||
|
"expected 'alice' in read_user_ids from identity scopes"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
ws.read_user_ids().contains(&"shared".to_string()),
|
||||||
|
"expected 'shared' in read_user_ids from identity scopes"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_caches_per_user() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
let pool = WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
EmbeddingCacheConfig::default(),
|
||||||
|
WorkspaceSearchConfig::default(),
|
||||||
|
WorkspaceConfig::default(),
|
||||||
|
);
|
||||||
|
let alice_id = UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
};
|
||||||
|
let bob_id = UserIdentity {
|
||||||
|
user_id: "bob".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let alice_ws1 = pool.get_or_create(&alice_id).await;
|
||||||
|
let alice_ws2 = pool.get_or_create(&alice_id).await;
|
||||||
|
let bob_ws = pool.get_or_create(&bob_id).await;
|
||||||
|
|
||||||
|
// Same user gets the same Arc.
|
||||||
|
assert!(Arc::ptr_eq(&alice_ws1, &alice_ws2));
|
||||||
|
// Different users get different instances.
|
||||||
|
assert!(!Arc::ptr_eq(&alice_ws1, &bob_ws));
|
||||||
|
assert_eq!(alice_ws1.user_id(), "alice");
|
||||||
|
assert_eq!(bob_ws.user_id(), "bob");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_combines_global_and_identity_scopes() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
let ws_config = WorkspaceConfig {
|
||||||
|
memory_layers: vec![],
|
||||||
|
read_scopes: vec!["global-shared".to_string()],
|
||||||
|
};
|
||||||
|
let pool = WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
EmbeddingCacheConfig::default(),
|
||||||
|
WorkspaceSearchConfig::default(),
|
||||||
|
ws_config,
|
||||||
|
);
|
||||||
|
let identity = UserIdentity {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
workspace_read_scopes: vec!["token-scope".to_string()],
|
||||||
|
};
|
||||||
|
let ws = pool.get_or_create(&identity).await;
|
||||||
|
let scopes = ws.read_user_ids();
|
||||||
|
// Primary scope
|
||||||
|
assert!(scopes.contains(&"alice".to_string()));
|
||||||
|
// Global config scope
|
||||||
|
assert!(
|
||||||
|
scopes.contains(&"global-shared".to_string()),
|
||||||
|
"expected global scope 'global-shared', got {:?}",
|
||||||
|
scopes
|
||||||
|
);
|
||||||
|
// Token identity scope
|
||||||
|
assert!(
|
||||||
|
scopes.contains(&"token-scope".to_string()),
|
||||||
|
"expected token scope 'token-scope', got {:?}",
|
||||||
|
scopes
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
// Jobs Handler Isolation Tests
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
mod jobs_isolation {
|
||||||
|
use super::*;
|
||||||
|
use crate::channels::web::handlers::jobs::{
|
||||||
|
jobs_cancel_handler, jobs_prompt_handler, jobs_restart_handler, jobs_summary_handler,
|
||||||
|
};
|
||||||
|
// SandboxStore methods are accessed through the Database supertrait.
|
||||||
|
|
||||||
|
/// Build a router with job endpoints behind multi-user auth.
|
||||||
|
fn jobs_router(state: Arc<GatewayState>, auth: MultiAuthState) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/api/jobs/summary", get(jobs_summary_handler))
|
||||||
|
.route("/api/jobs/{id}/cancel", post(jobs_cancel_handler))
|
||||||
|
.route("/api/jobs/{id}/restart", post(jobs_restart_handler))
|
||||||
|
.route("/api/jobs/{id}/prompt", post(jobs_prompt_handler))
|
||||||
|
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_jobs_summary_scoped_to_user() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
// Insert sandbox jobs for alice and bob.
|
||||||
|
let alice_job = make_sandbox_job("alice", "alice task");
|
||||||
|
let bob_job = make_sandbox_job("bob", "bob task");
|
||||||
|
db.save_sandbox_job(&alice_job).await.unwrap();
|
||||||
|
db.save_sandbox_job(&bob_job).await.unwrap();
|
||||||
|
|
||||||
|
let state = build_state(Some(db), None);
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = jobs_router(state, auth);
|
||||||
|
|
||||||
|
// Alice should see 1 job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/jobs/summary")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body: serde_json::Value =
|
||||||
|
serde_json::from_slice(&axum::body::to_bytes(resp.into_body(), 4096).await.unwrap())
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(body["total"], 1, "alice should see only her own jobs");
|
||||||
|
|
||||||
|
// Bob should see 1 job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/jobs/summary")
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body: serde_json::Value =
|
||||||
|
serde_json::from_slice(&axum::body::to_bytes(resp.into_body(), 4096).await.unwrap())
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(body["total"], 1, "bob should see only his own jobs");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_jobs_restart_rejects_other_user() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
// Insert a failed sandbox job owned by alice.
|
||||||
|
let mut alice_job = make_sandbox_job("alice", "alice task");
|
||||||
|
alice_job.status = "failed".to_string();
|
||||||
|
alice_job.success = Some(false);
|
||||||
|
db.save_sandbox_job(&alice_job).await.unwrap();
|
||||||
|
|
||||||
|
let state = build_state(Some(db), None);
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = jobs_router(state, auth);
|
||||||
|
|
||||||
|
// Bob tries to restart alice's job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri(format!("/api/jobs/{}/restart", alice_job.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not be able to restart alice's job"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_jobs_prompt_works_for_agent_jobs() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
// Insert a running sandbox job owned by alice in claude_code mode.
|
||||||
|
let mut alice_job = make_sandbox_job("alice", "prompt test");
|
||||||
|
alice_job.status = "running".to_string();
|
||||||
|
alice_job.success = None;
|
||||||
|
alice_job.completed_at = None;
|
||||||
|
db.save_sandbox_job(&alice_job).await.unwrap();
|
||||||
|
db.update_sandbox_job_mode(alice_job.id, "claude_code")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let prompt_queue: PromptQueue =
|
||||||
|
Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new()));
|
||||||
|
let state = build_state(Some(db), Some(prompt_queue.clone()));
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = jobs_router(state, auth);
|
||||||
|
|
||||||
|
// Alice prompts her own job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri(format!("/api/jobs/{}/prompt", alice_job.id))
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(
|
||||||
|
serde_json::to_string(&serde_json::json!({"content": "hello"})).unwrap(),
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::OK,
|
||||||
|
"alice should be able to prompt her own job"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Verify prompt was enqueued.
|
||||||
|
let queue = prompt_queue.lock().await;
|
||||||
|
assert!(
|
||||||
|
queue.contains_key(&alice_job.id),
|
||||||
|
"prompt queue should contain alice's job"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_jobs_prompt_rejects_other_user() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
let mut alice_job = make_sandbox_job("alice", "alice task");
|
||||||
|
alice_job.status = "running".to_string();
|
||||||
|
alice_job.success = None;
|
||||||
|
alice_job.completed_at = None;
|
||||||
|
db.save_sandbox_job(&alice_job).await.unwrap();
|
||||||
|
db.update_sandbox_job_mode(alice_job.id, "claude_code")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let prompt_queue: PromptQueue =
|
||||||
|
Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new()));
|
||||||
|
let state = build_state(Some(db), Some(prompt_queue));
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = jobs_router(state, auth);
|
||||||
|
|
||||||
|
// Bob tries to prompt alice's job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri(format!("/api/jobs/{}/prompt", alice_job.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(
|
||||||
|
serde_json::to_string(&serde_json::json!({"content": "sneaky"})).unwrap(),
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not be able to prompt alice's job"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_jobs_cancel_rejects_other_user() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
let mut alice_job = make_sandbox_job("alice", "alice running");
|
||||||
|
alice_job.status = "running".to_string();
|
||||||
|
alice_job.success = None;
|
||||||
|
alice_job.completed_at = None;
|
||||||
|
db.save_sandbox_job(&alice_job).await.unwrap();
|
||||||
|
|
||||||
|
let state = build_state(Some(db), None);
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = jobs_router(state, auth);
|
||||||
|
|
||||||
|
// Bob tries to cancel alice's job.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri(format!("/api/jobs/{}/cancel", alice_job.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not be able to cancel alice's job"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
// Routines Isolation Tests
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
mod routines_isolation {
|
||||||
|
use super::*;
|
||||||
|
use crate::channels::web::handlers::routines::{
|
||||||
|
routines_delete_handler, routines_detail_handler, routines_list_handler,
|
||||||
|
routines_summary_handler, routines_toggle_handler,
|
||||||
|
};
|
||||||
|
// RoutineStore methods are accessed through the Database supertrait.
|
||||||
|
|
||||||
|
fn routines_router(state: Arc<GatewayState>, auth: MultiAuthState) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/api/routines", get(routines_list_handler))
|
||||||
|
.route("/api/routines/summary", get(routines_summary_handler))
|
||||||
|
.route("/api/routines/{id}", get(routines_detail_handler))
|
||||||
|
.route("/api/routines/{id}/toggle", post(routines_toggle_handler))
|
||||||
|
.route("/api/routines/{id}", delete(routines_delete_handler))
|
||||||
|
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_routines_isolation() {
|
||||||
|
let (db, _dir) = test_db().await;
|
||||||
|
|
||||||
|
// Create routines for alice and bob.
|
||||||
|
let alice_routine = make_routine("alice", "alice-daily");
|
||||||
|
let bob_routine = make_routine("bob", "bob-daily");
|
||||||
|
db.create_routine(&alice_routine).await.unwrap();
|
||||||
|
db.create_routine(&bob_routine).await.unwrap();
|
||||||
|
|
||||||
|
let state = build_state(Some(db), None);
|
||||||
|
let auth = two_user_auth();
|
||||||
|
let app = routines_router(state, auth);
|
||||||
|
|
||||||
|
// Alice sees only her routine in the list.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/routines")
|
||||||
|
.header("Authorization", "Bearer tok-alice")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body: serde_json::Value =
|
||||||
|
serde_json::from_slice(&axum::body::to_bytes(resp.into_body(), 8192).await.unwrap())
|
||||||
|
.unwrap();
|
||||||
|
let routines = body["routines"].as_array().unwrap();
|
||||||
|
assert_eq!(routines.len(), 1, "alice should see only her routines");
|
||||||
|
assert_eq!(routines[0]["name"], "alice-daily");
|
||||||
|
|
||||||
|
// Bob sees only his routine.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/routines")
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
let body: serde_json::Value =
|
||||||
|
serde_json::from_slice(&axum::body::to_bytes(resp.into_body(), 8192).await.unwrap())
|
||||||
|
.unwrap();
|
||||||
|
let routines = body["routines"].as_array().unwrap();
|
||||||
|
assert_eq!(routines.len(), 1, "bob should see only his routines");
|
||||||
|
assert_eq!(routines[0]["name"], "bob-daily");
|
||||||
|
|
||||||
|
// Bob cannot view alice's routine detail.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri(format!("/api/routines/{}", alice_routine.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not see alice's routine detail"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Bob cannot toggle alice's routine.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::POST)
|
||||||
|
.uri(format!("/api/routines/{}/toggle", alice_routine.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not toggle alice's routine"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Bob cannot delete alice's routine.
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(Method::DELETE)
|
||||||
|
.uri(format!("/api/routines/{}", alice_routine.id))
|
||||||
|
.header("Authorization", "Bearer tok-bob")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"bob should not delete alice's routine"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
// Handler Auth Enforcement Tests
|
||||||
|
// ═══════════════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
|
mod auth_enforcement {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// Dummy handler that extracts `AuthenticatedUser` — if the auth middleware
|
||||||
|
/// rejects the request, this handler is never reached.
|
||||||
|
async fn authed_handler(AuthenticatedUser(_user): AuthenticatedUser) -> &'static str {
|
||||||
|
"ok"
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a router with the real auth middleware and dummy handlers at all
|
||||||
|
/// the paths we want to verify require authentication.
|
||||||
|
fn auth_test_router(auth: MultiAuthState) -> Router {
|
||||||
|
let state = build_state(None, None);
|
||||||
|
Router::new()
|
||||||
|
// Routines
|
||||||
|
.route("/api/routines", get(authed_handler))
|
||||||
|
.route("/api/routines/summary", get(authed_handler))
|
||||||
|
.route("/api/routines/{id}", get(authed_handler))
|
||||||
|
.route("/api/routines/{id}/toggle", post(authed_handler))
|
||||||
|
.route("/api/routines/{id}", delete(authed_handler))
|
||||||
|
// Skills
|
||||||
|
.route("/api/skills", get(authed_handler))
|
||||||
|
.route("/api/skills/search", post(authed_handler))
|
||||||
|
.route("/api/skills/install", post(authed_handler))
|
||||||
|
.route("/api/skills/{name}", delete(authed_handler))
|
||||||
|
// Logs
|
||||||
|
.route("/api/logs/events", get(authed_handler))
|
||||||
|
.route("/api/logs/level", get(authed_handler).put(authed_handler))
|
||||||
|
// Gateway status
|
||||||
|
.route("/api/gateway/status", get(authed_handler))
|
||||||
|
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Send a request without auth and assert it returns UNAUTHORIZED.
|
||||||
|
async fn assert_requires_auth(app: &Router, method: Method, uri: &str) {
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(method.clone())
|
||||||
|
.uri(uri)
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::UNAUTHORIZED,
|
||||||
|
"{} {} should require auth",
|
||||||
|
method,
|
||||||
|
uri
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Send a request with a valid token and assert it succeeds.
|
||||||
|
async fn assert_passes_with_token(app: &Router, method: Method, uri: &str, token: &str) {
|
||||||
|
let req = Request::builder()
|
||||||
|
.method(method.clone())
|
||||||
|
.uri(uri)
|
||||||
|
.header("Authorization", format!("Bearer {token}"))
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resp.status(),
|
||||||
|
StatusCode::OK,
|
||||||
|
"{} {} should pass with valid token",
|
||||||
|
method,
|
||||||
|
uri
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_routines_handlers_require_auth() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
let id = Uuid::new_v4();
|
||||||
|
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/routines").await;
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/routines/summary").await;
|
||||||
|
assert_requires_auth(&app, Method::GET, &format!("/api/routines/{id}")).await;
|
||||||
|
assert_requires_auth(&app, Method::POST, &format!("/api/routines/{id}/toggle")).await;
|
||||||
|
assert_requires_auth(&app, Method::DELETE, &format!("/api/routines/{id}")).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_skills_handlers_require_auth() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/skills").await;
|
||||||
|
assert_requires_auth(&app, Method::POST, "/api/skills/search").await;
|
||||||
|
assert_requires_auth(&app, Method::POST, "/api/skills/install").await;
|
||||||
|
assert_requires_auth(&app, Method::DELETE, "/api/skills/test-skill").await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_logs_handlers_require_auth() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/logs/events").await;
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/logs/level").await;
|
||||||
|
assert_requires_auth(&app, Method::PUT, "/api/logs/level").await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_gateway_status_requires_auth() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
|
||||||
|
assert_requires_auth(&app, Method::GET, "/api/gateway/status").await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_valid_token_passes_all_endpoints() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
let id = Uuid::new_v4();
|
||||||
|
|
||||||
|
assert_passes_with_token(&app, Method::GET, "/api/routines", "secret-tok").await;
|
||||||
|
assert_passes_with_token(&app, Method::GET, "/api/skills", "secret-tok").await;
|
||||||
|
assert_passes_with_token(&app, Method::GET, "/api/logs/events", "secret-tok").await;
|
||||||
|
assert_passes_with_token(&app, Method::GET, "/api/gateway/status", "secret-tok").await;
|
||||||
|
assert_passes_with_token(
|
||||||
|
&app,
|
||||||
|
Method::GET,
|
||||||
|
&format!("/api/routines/{id}"),
|
||||||
|
"secret-tok",
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_wrong_token_rejected_on_all_endpoints() {
|
||||||
|
let auth = MultiAuthState::single("secret-tok".to_string(), "user".to_string());
|
||||||
|
let app = auth_test_router(auth);
|
||||||
|
|
||||||
|
// Wrong token should be rejected.
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/routines")
|
||||||
|
.header("Authorization", "Bearer wrong-tok")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.clone().oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.uri("/api/gateway/status")
|
||||||
|
.header("Authorization", "Bearer wrong-tok")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap();
|
||||||
|
let resp = app.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||||
|
}
|
||||||
|
}
|
||||||
+33
-206
@@ -63,6 +63,9 @@ pub struct TurnInfo {
|
|||||||
pub started_at: String,
|
pub started_at: String,
|
||||||
pub completed_at: Option<String>,
|
pub completed_at: Option<String>,
|
||||||
pub tool_calls: Vec<ToolCallInfo>,
|
pub tool_calls: Vec<ToolCallInfo>,
|
||||||
|
/// Agent's reasoning narrative for this turn.
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub narrative: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -74,6 +77,9 @@ pub struct ToolCallInfo {
|
|||||||
pub result_preview: Option<String>,
|
pub result_preview: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub error: Option<String>,
|
pub error: Option<String>,
|
||||||
|
/// Agent's reasoning for choosing this tool.
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub rationale: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -114,165 +120,9 @@ pub struct ApprovalRequest {
|
|||||||
pub thread_id: Option<String>,
|
pub thread_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- SSE Event Types ---
|
// --- App Event (re-exported from ironclaw_common) ---
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize)]
|
pub use ironclaw_common::{AppEvent, ToolDecisionDto};
|
||||||
#[serde(tag = "type")]
|
|
||||||
pub enum SseEvent {
|
|
||||||
#[serde(rename = "response")]
|
|
||||||
Response { content: String, thread_id: String },
|
|
||||||
#[serde(rename = "thinking")]
|
|
||||||
Thinking {
|
|
||||||
message: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "tool_started")]
|
|
||||||
ToolStarted {
|
|
||||||
name: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "tool_completed")]
|
|
||||||
ToolCompleted {
|
|
||||||
name: String,
|
|
||||||
success: bool,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
error: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
parameters: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "tool_result")]
|
|
||||||
ToolResult {
|
|
||||||
name: String,
|
|
||||||
preview: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "stream_chunk")]
|
|
||||||
StreamChunk {
|
|
||||||
content: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "status")]
|
|
||||||
Status {
|
|
||||||
message: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "job_started")]
|
|
||||||
JobStarted {
|
|
||||||
job_id: String,
|
|
||||||
title: String,
|
|
||||||
browse_url: String,
|
|
||||||
},
|
|
||||||
#[serde(rename = "approval_needed")]
|
|
||||||
ApprovalNeeded {
|
|
||||||
request_id: String,
|
|
||||||
tool_name: String,
|
|
||||||
description: String,
|
|
||||||
parameters: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
/// Whether the "always" auto-approve option should be shown.
|
|
||||||
allow_always: bool,
|
|
||||||
},
|
|
||||||
#[serde(rename = "auth_required")]
|
|
||||||
AuthRequired {
|
|
||||||
extension_name: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
instructions: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
auth_url: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
setup_url: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "auth_completed")]
|
|
||||||
AuthCompleted {
|
|
||||||
extension_name: String,
|
|
||||||
success: bool,
|
|
||||||
message: String,
|
|
||||||
},
|
|
||||||
#[serde(rename = "error")]
|
|
||||||
Error {
|
|
||||||
message: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "heartbeat")]
|
|
||||||
Heartbeat,
|
|
||||||
|
|
||||||
// Sandbox job streaming events (worker + Claude Code bridge)
|
|
||||||
#[serde(rename = "job_message")]
|
|
||||||
JobMessage {
|
|
||||||
job_id: String,
|
|
||||||
role: String,
|
|
||||||
content: String,
|
|
||||||
},
|
|
||||||
#[serde(rename = "job_tool_use")]
|
|
||||||
JobToolUse {
|
|
||||||
job_id: String,
|
|
||||||
tool_name: String,
|
|
||||||
input: serde_json::Value,
|
|
||||||
},
|
|
||||||
#[serde(rename = "job_tool_result")]
|
|
||||||
JobToolResult {
|
|
||||||
job_id: String,
|
|
||||||
tool_name: String,
|
|
||||||
output: String,
|
|
||||||
},
|
|
||||||
#[serde(rename = "job_status")]
|
|
||||||
JobStatus { job_id: String, message: String },
|
|
||||||
#[serde(rename = "job_result")]
|
|
||||||
JobResult {
|
|
||||||
job_id: String,
|
|
||||||
status: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
session_id: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
fallback_deliverable: Option<serde_json::Value>,
|
|
||||||
},
|
|
||||||
|
|
||||||
/// An image was generated by a tool.
|
|
||||||
#[serde(rename = "image_generated")]
|
|
||||||
ImageGenerated {
|
|
||||||
data_url: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
path: Option<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
|
|
||||||
/// Suggested follow-up messages for the user.
|
|
||||||
#[serde(rename = "suggestions")]
|
|
||||||
Suggestions {
|
|
||||||
suggestions: Vec<String>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
|
|
||||||
/// Per-turn token usage and cost summary.
|
|
||||||
#[serde(rename = "turn_cost")]
|
|
||||||
TurnCost {
|
|
||||||
input_tokens: u64,
|
|
||||||
output_tokens: u64,
|
|
||||||
cost_usd: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
thread_id: Option<String>,
|
|
||||||
},
|
|
||||||
|
|
||||||
/// Extension activation status change (WASM channels).
|
|
||||||
#[serde(rename = "extension_status")]
|
|
||||||
ExtensionStatus {
|
|
||||||
extension_name: String,
|
|
||||||
status: String,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
message: Option<String>,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
// --- Memory ---
|
// --- Memory ---
|
||||||
|
|
||||||
@@ -784,32 +634,9 @@ pub enum WsServerMessage {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl WsServerMessage {
|
impl WsServerMessage {
|
||||||
/// Create a WsServerMessage from an SseEvent.
|
/// Create a WsServerMessage from an AppEvent.
|
||||||
pub fn from_sse_event(event: &SseEvent) -> Self {
|
pub fn from_app_event(event: &AppEvent) -> Self {
|
||||||
let event_type = match event {
|
let event_type = event.event_type();
|
||||||
SseEvent::Response { .. } => "response",
|
|
||||||
SseEvent::Thinking { .. } => "thinking",
|
|
||||||
SseEvent::ToolStarted { .. } => "tool_started",
|
|
||||||
SseEvent::ToolCompleted { .. } => "tool_completed",
|
|
||||||
SseEvent::ToolResult { .. } => "tool_result",
|
|
||||||
SseEvent::StreamChunk { .. } => "stream_chunk",
|
|
||||||
SseEvent::Status { .. } => "status",
|
|
||||||
SseEvent::JobStarted { .. } => "job_started",
|
|
||||||
SseEvent::ApprovalNeeded { .. } => "approval_needed",
|
|
||||||
SseEvent::AuthRequired { .. } => "auth_required",
|
|
||||||
SseEvent::AuthCompleted { .. } => "auth_completed",
|
|
||||||
SseEvent::Error { .. } => "error",
|
|
||||||
SseEvent::Heartbeat => "heartbeat",
|
|
||||||
SseEvent::JobMessage { .. } => "job_message",
|
|
||||||
SseEvent::JobToolUse { .. } => "job_tool_use",
|
|
||||||
SseEvent::JobToolResult { .. } => "job_tool_result",
|
|
||||||
SseEvent::JobStatus { .. } => "job_status",
|
|
||||||
SseEvent::JobResult { .. } => "job_result",
|
|
||||||
SseEvent::ImageGenerated { .. } => "image_generated",
|
|
||||||
SseEvent::Suggestions { .. } => "suggestions",
|
|
||||||
SseEvent::TurnCost { .. } => "turn_cost",
|
|
||||||
SseEvent::ExtensionStatus { .. } => "extension_status",
|
|
||||||
};
|
|
||||||
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
|
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
|
||||||
WsServerMessage::Event {
|
WsServerMessage::Event {
|
||||||
event_type: event_type.to_string(),
|
event_type: event_type.to_string(),
|
||||||
@@ -1101,12 +928,12 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_response() {
|
fn test_ws_server_from_app_event_response() {
|
||||||
let sse = SseEvent::Response {
|
let event = AppEvent::Response {
|
||||||
content: "hello".to_string(),
|
content: "hello".to_string(),
|
||||||
thread_id: "t1".to_string(),
|
thread_id: "t1".to_string(),
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "response");
|
assert_eq!(event_type, "response");
|
||||||
@@ -1118,12 +945,12 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_thinking() {
|
fn test_ws_server_from_app_event_thinking() {
|
||||||
let sse = SseEvent::Thinking {
|
let event = AppEvent::Thinking {
|
||||||
message: "reasoning...".to_string(),
|
message: "reasoning...".to_string(),
|
||||||
thread_id: None,
|
thread_id: None,
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "thinking");
|
assert_eq!(event_type, "thinking");
|
||||||
@@ -1134,8 +961,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_approval_needed() {
|
fn test_ws_server_from_app_event_approval_needed() {
|
||||||
let sse = SseEvent::ApprovalNeeded {
|
let event = AppEvent::ApprovalNeeded {
|
||||||
request_id: "r1".to_string(),
|
request_id: "r1".to_string(),
|
||||||
tool_name: "shell".to_string(),
|
tool_name: "shell".to_string(),
|
||||||
description: "Run ls".to_string(),
|
description: "Run ls".to_string(),
|
||||||
@@ -1143,7 +970,7 @@ mod tests {
|
|||||||
thread_id: Some("t1".to_string()),
|
thread_id: Some("t1".to_string()),
|
||||||
allow_always: true,
|
allow_always: true,
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "approval_needed");
|
assert_eq!(event_type, "approval_needed");
|
||||||
@@ -1155,9 +982,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_heartbeat() {
|
fn test_ws_server_from_app_event_heartbeat() {
|
||||||
let sse = SseEvent::Heartbeat;
|
let event = AppEvent::Heartbeat;
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, .. } => {
|
WsServerMessage::Event { event_type, .. } => {
|
||||||
assert_eq!(event_type, "heartbeat");
|
assert_eq!(event_type, "heartbeat");
|
||||||
@@ -1197,8 +1024,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_sse_auth_required_serialize() {
|
fn test_app_event_auth_required_serialize() {
|
||||||
let event = SseEvent::AuthRequired {
|
let event = AppEvent::AuthRequired {
|
||||||
extension_name: "notion".to_string(),
|
extension_name: "notion".to_string(),
|
||||||
instructions: Some("Get your token from...".to_string()),
|
instructions: Some("Get your token from...".to_string()),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
@@ -1214,8 +1041,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_sse_auth_completed_serialize() {
|
fn test_app_event_auth_completed_serialize() {
|
||||||
let event = SseEvent::AuthCompleted {
|
let event = AppEvent::AuthCompleted {
|
||||||
extension_name: "notion".to_string(),
|
extension_name: "notion".to_string(),
|
||||||
success: true,
|
success: true,
|
||||||
message: "notion authenticated (3 tools loaded)".to_string(),
|
message: "notion authenticated (3 tools loaded)".to_string(),
|
||||||
@@ -1228,14 +1055,14 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_auth_required() {
|
fn test_ws_server_from_app_event_auth_required() {
|
||||||
let sse = SseEvent::AuthRequired {
|
let event = AppEvent::AuthRequired {
|
||||||
extension_name: "openai".to_string(),
|
extension_name: "openai".to_string(),
|
||||||
instructions: Some("Enter API key".to_string()),
|
instructions: Some("Enter API key".to_string()),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
setup_url: None,
|
setup_url: None,
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "auth_required");
|
assert_eq!(event_type, "auth_required");
|
||||||
@@ -1246,13 +1073,13 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_ws_server_from_sse_auth_completed() {
|
fn test_ws_server_from_app_event_auth_completed() {
|
||||||
let sse = SseEvent::AuthCompleted {
|
let event = AppEvent::AuthCompleted {
|
||||||
extension_name: "slack".to_string(),
|
extension_name: "slack".to_string(),
|
||||||
success: false,
|
success: false,
|
||||||
message: "Invalid token".to_string(),
|
message: "Invalid token".to_string(),
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_app_event(&event);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "auth_completed");
|
assert_eq!(event_type, "auth_completed");
|
||||||
|
|||||||
+86
-115
@@ -2,28 +2,21 @@
|
|||||||
|
|
||||||
use crate::channels::web::types::{ToolCallInfo, TurnInfo};
|
use crate::channels::web::types::{ToolCallInfo, TurnInfo};
|
||||||
|
|
||||||
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
|
pub use ironclaw_common::truncate_preview;
|
||||||
///
|
|
||||||
/// If the input is wrapped in `<tool_output …>…</tool_output>` and truncation
|
|
||||||
/// removes the closing tag, the tag is re-appended so downstream XML parsers
|
|
||||||
/// never see an unclosed element.
|
|
||||||
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
|
|
||||||
if s.len() <= max_bytes {
|
|
||||||
return s.to_string();
|
|
||||||
}
|
|
||||||
// Walk backwards from max_bytes to find a valid char boundary
|
|
||||||
let mut end = max_bytes;
|
|
||||||
while end > 0 && !s.is_char_boundary(end) {
|
|
||||||
end -= 1;
|
|
||||||
}
|
|
||||||
let mut result = format!("{}...", &s[..end]);
|
|
||||||
|
|
||||||
// Re-close <tool_output> if truncation cut through the closing tag.
|
/// Parse tool call summary JSON objects into `ToolCallInfo` structs.
|
||||||
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
|
fn parse_tool_call_infos(calls: &[serde_json::Value]) -> Vec<ToolCallInfo> {
|
||||||
result.push_str("\n</tool_output>");
|
calls
|
||||||
}
|
.iter()
|
||||||
|
.map(|c| ToolCallInfo {
|
||||||
result
|
name: c["name"].as_str().unwrap_or("unknown").to_string(),
|
||||||
|
has_result: c.get("result_preview").is_some_and(|v| !v.is_null()),
|
||||||
|
has_error: c.get("error").is_some_and(|v| !v.is_null()),
|
||||||
|
result_preview: c["result_preview"].as_str().map(String::from),
|
||||||
|
error: c["error"].as_str().map(String::from),
|
||||||
|
rationale: c["rationale"].as_str().map(String::from),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
|
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
|
||||||
@@ -49,6 +42,7 @@ pub fn build_turns_from_db_messages(
|
|||||||
started_at: msg.created_at.to_rfc3339(),
|
started_at: msg.created_at.to_rfc3339(),
|
||||||
completed_at: None,
|
completed_at: None,
|
||||||
tool_calls: Vec::new(),
|
tool_calls: Vec::new(),
|
||||||
|
narrative: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Check if next message is a tool_calls record
|
// Check if next message is a tool_calls record
|
||||||
@@ -56,18 +50,28 @@ pub fn build_turns_from_db_messages(
|
|||||||
&& next.role == "tool_calls"
|
&& next.role == "tool_calls"
|
||||||
{
|
{
|
||||||
let tc_msg = iter.next().expect("peeked");
|
let tc_msg = iter.next().expect("peeked");
|
||||||
match serde_json::from_str::<Vec<serde_json::Value>>(&tc_msg.content) {
|
// Parse tool_calls JSON — supports two formats:
|
||||||
Ok(calls) => {
|
// safety: no byte-index slicing; comment describes JSON shape
|
||||||
turn.tool_calls = calls
|
match serde_json::from_str::<serde_json::Value>(&tc_msg.content) {
|
||||||
.iter()
|
Ok(serde_json::Value::Array(calls)) => {
|
||||||
.map(|c| ToolCallInfo {
|
// Old format: plain array
|
||||||
name: c["name"].as_str().unwrap_or("unknown").to_string(),
|
turn.tool_calls = parse_tool_call_infos(&calls);
|
||||||
has_result: c.get("result_preview").is_some(),
|
}
|
||||||
has_error: c.get("error").is_some(),
|
Ok(serde_json::Value::Object(obj)) => {
|
||||||
result_preview: c["result_preview"].as_str().map(String::from),
|
// New wrapped format with narrative
|
||||||
error: c["error"].as_str().map(String::from),
|
turn.narrative = obj
|
||||||
})
|
.get("narrative")
|
||||||
.collect();
|
.and_then(|v| v.as_str())
|
||||||
|
.map(String::from);
|
||||||
|
if let Some(serde_json::Value::Array(calls)) = obj.get("calls") {
|
||||||
|
turn.tool_calls = parse_tool_call_infos(calls);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(_) => {
|
||||||
|
tracing::warn!(
|
||||||
|
message_id = %tc_msg.id,
|
||||||
|
"Unexpected tool_calls JSON shape in DB, skipping"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -105,6 +109,7 @@ pub fn build_turns_from_db_messages(
|
|||||||
started_at: msg.created_at.to_rfc3339(),
|
started_at: msg.created_at.to_rfc3339(),
|
||||||
completed_at: Some(msg.created_at.to_rfc3339()),
|
completed_at: Some(msg.created_at.to_rfc3339()),
|
||||||
tool_calls: Vec::new(),
|
tool_calls: Vec::new(),
|
||||||
|
narrative: None,
|
||||||
});
|
});
|
||||||
turn_number += 1;
|
turn_number += 1;
|
||||||
}
|
}
|
||||||
@@ -118,88 +123,6 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
// ---- truncate_preview tests ----
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_short_string() {
|
|
||||||
assert_eq!(truncate_preview("hello", 10), "hello");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_exact_boundary() {
|
|
||||||
assert_eq!(truncate_preview("hello", 5), "hello");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_truncates_ascii() {
|
|
||||||
assert_eq!(truncate_preview("hello world", 5), "hello...");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_empty_string() {
|
|
||||||
assert_eq!(truncate_preview("", 10), "");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_multibyte_char_boundary() {
|
|
||||||
// '€' is 3 bytes (E2 82 AC). "a€b" = [61, E2, 82, AC, 62] = 5 bytes
|
|
||||||
// Truncating at max_bytes=3 should not split the euro sign.
|
|
||||||
let s = "a€b";
|
|
||||||
let result = truncate_preview(s, 3);
|
|
||||||
// max_bytes=3 lands mid-€, so it walks back to byte 1 ("a")
|
|
||||||
assert_eq!(result, "a...");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_emoji() {
|
|
||||||
// '🦀' is 4 bytes. "hi🦀" = 6 bytes
|
|
||||||
let s = "hi🦀";
|
|
||||||
let result = truncate_preview(s, 4);
|
|
||||||
// max_bytes=4 lands mid-🦀, walks back to byte 2 ("hi")
|
|
||||||
assert_eq!(result, "hi...");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_cjk() {
|
|
||||||
// CJK characters are 3 bytes each. "你好世界" = 12 bytes
|
|
||||||
let s = "你好世界";
|
|
||||||
let result = truncate_preview(s, 7);
|
|
||||||
// max_bytes=7 lands mid-character (byte 7 is inside 世), walks back to 6 ("你好")
|
|
||||||
assert_eq!(result, "你好...");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_zero_max_bytes() {
|
|
||||||
assert_eq!(truncate_preview("hello", 0), "...");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_closes_tool_output_tag() {
|
|
||||||
let s = "<tool_output name=\"search\">\nSome very long content here\n</tool_output>";
|
|
||||||
// Truncate so it cuts before the closing tag
|
|
||||||
let result = truncate_preview(s, 60);
|
|
||||||
assert!(result.ends_with("</tool_output>"));
|
|
||||||
assert!(result.contains("..."));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_no_extra_close_when_intact() {
|
|
||||||
let s = "<tool_output name=\"echo\">\nshort\n</tool_output>";
|
|
||||||
// The string is short enough not to be truncated
|
|
||||||
let result = truncate_preview(s, 500);
|
|
||||||
assert_eq!(result, s);
|
|
||||||
// Should not have a duplicate closing tag
|
|
||||||
assert_eq!(result.matches("</tool_output>").count(), 1);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_preview_non_xml_unaffected() {
|
|
||||||
let s = "Just a plain long string that gets truncated";
|
|
||||||
let result = truncate_preview(s, 10);
|
|
||||||
assert_eq!(result, "Just a pla...");
|
|
||||||
assert!(!result.contains("</tool_output>"));
|
|
||||||
}
|
|
||||||
|
|
||||||
// ---- build_turns_from_db_messages tests ----
|
// ---- build_turns_from_db_messages tests ----
|
||||||
|
|
||||||
fn make_msg(role: &str, content: &str, offset_ms: i64) -> crate::history::ConversationMessage {
|
fn make_msg(role: &str, content: &str, offset_ms: i64) -> crate::history::ConversationMessage {
|
||||||
@@ -305,4 +228,52 @@ mod tests {
|
|||||||
assert!(turns[0].tool_calls.is_empty());
|
assert!(turns[0].tool_calls.is_empty());
|
||||||
assert_eq!(turns[0].state, "Completed");
|
assert_eq!(turns[0].state, "Completed");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_turns_with_wrapped_tool_calls_format() {
|
||||||
|
let tc_json = serde_json::json!({
|
||||||
|
"narrative": "Searching memory for context before proceeding.",
|
||||||
|
"calls": [
|
||||||
|
{"name": "memory_search", "result_preview": "found 3 items", "rationale": "consult prior context"},
|
||||||
|
{"name": "shell", "error": "permission denied"}
|
||||||
|
]
|
||||||
|
});
|
||||||
|
let messages = vec![
|
||||||
|
make_msg("user", "Find info", 0),
|
||||||
|
make_msg("tool_calls", &tc_json.to_string(), 500),
|
||||||
|
make_msg("assistant", "Here's what I found", 1000),
|
||||||
|
];
|
||||||
|
let turns = build_turns_from_db_messages(&messages);
|
||||||
|
assert_eq!(turns.len(), 1);
|
||||||
|
assert_eq!(
|
||||||
|
turns[0].narrative.as_deref(),
|
||||||
|
Some("Searching memory for context before proceeding.")
|
||||||
|
);
|
||||||
|
assert_eq!(turns[0].tool_calls.len(), 2);
|
||||||
|
assert_eq!(turns[0].tool_calls[0].name, "memory_search");
|
||||||
|
assert_eq!(
|
||||||
|
turns[0].tool_calls[0].rationale.as_deref(),
|
||||||
|
Some("consult prior context")
|
||||||
|
);
|
||||||
|
assert!(turns[0].tool_calls[0].has_result);
|
||||||
|
assert_eq!(turns[0].tool_calls[1].name, "shell");
|
||||||
|
assert!(turns[0].tool_calls[1].has_error);
|
||||||
|
assert_eq!(turns[0].response.as_deref(), Some("Here's what I found"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_turns_wrapped_format_without_narrative() {
|
||||||
|
let tc_json = serde_json::json!({
|
||||||
|
"calls": [{"name": "echo", "result_preview": "hello"}]
|
||||||
|
});
|
||||||
|
let messages = vec![
|
||||||
|
make_msg("user", "Say hi", 0),
|
||||||
|
make_msg("tool_calls", &tc_json.to_string(), 500),
|
||||||
|
make_msg("assistant", "Done", 1000),
|
||||||
|
];
|
||||||
|
let turns = build_turns_from_db_messages(&messages);
|
||||||
|
assert_eq!(turns.len(), 1);
|
||||||
|
assert!(turns[0].narrative.is_none());
|
||||||
|
assert_eq!(turns[0].tool_calls.len(), 1);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+29
-17
@@ -62,7 +62,11 @@ impl Default for WsConnectionTracker {
|
|||||||
///
|
///
|
||||||
/// When either task ends (client disconnect or broadcast closed), both are
|
/// When either task ends (client disconnect or broadcast closed), both are
|
||||||
/// cleaned up.
|
/// cleaned up.
|
||||||
pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
|
pub async fn handle_ws_connection(
|
||||||
|
socket: WebSocket,
|
||||||
|
state: Arc<GatewayState>,
|
||||||
|
user: crate::channels::web::auth::UserIdentity,
|
||||||
|
) {
|
||||||
let (mut ws_sink, mut ws_stream) = socket.split();
|
let (mut ws_sink, mut ws_stream) = socket.split();
|
||||||
|
|
||||||
// Track connection
|
// Track connection
|
||||||
@@ -71,9 +75,9 @@ pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
|
|||||||
}
|
}
|
||||||
let tracker_for_drop = state.ws_tracker.clone();
|
let tracker_for_drop = state.ws_tracker.clone();
|
||||||
|
|
||||||
// Subscribe to broadcast events (same source as SSE).
|
// Subscribe to broadcast events (same source as SSE), scoped to this user.
|
||||||
// Reject if we've hit the connection limit.
|
// Reject if we've hit the connection limit.
|
||||||
let Some(raw_stream) = state.sse.subscribe_raw() else {
|
let Some(raw_stream) = state.sse.subscribe_raw(Some(user.user_id.clone())) else {
|
||||||
tracing::warn!("WebSocket rejected: too many connections");
|
tracing::warn!("WebSocket rejected: too many connections");
|
||||||
// Decrement the WS tracker we already incremented above.
|
// Decrement the WS tracker we already incremented above.
|
||||||
if let Some(ref tracker) = tracker_for_drop {
|
if let Some(ref tracker) = tracker_for_drop {
|
||||||
@@ -93,7 +97,7 @@ pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
|
|||||||
let msg = tokio::select! {
|
let msg = tokio::select! {
|
||||||
event = event_stream.next() => {
|
event = event_stream.next() => {
|
||||||
match event {
|
match event {
|
||||||
Some(sse_event) => WsServerMessage::from_sse_event(&sse_event),
|
Some(app_event) => WsServerMessage::from_app_event(&app_event),
|
||||||
None => break, // Broadcast channel closed
|
None => break, // Broadcast channel closed
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -117,7 +121,7 @@ pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
|
|||||||
});
|
});
|
||||||
|
|
||||||
// Receiver task: read client frames and route to agent
|
// Receiver task: read client frames and route to agent
|
||||||
let user_id = state.user_id.clone();
|
let user_id = user.user_id;
|
||||||
while let Some(Ok(frame)) = ws_stream.next().await {
|
while let Some(Ok(frame)) = ws_stream.next().await {
|
||||||
match frame {
|
match frame {
|
||||||
Message::Text(text) => {
|
Message::Text(text) => {
|
||||||
@@ -263,11 +267,15 @@ async fn handle_client_message(
|
|||||||
token,
|
token,
|
||||||
} => {
|
} => {
|
||||||
if let Some(ref ext_mgr) = state.extension_manager {
|
if let Some(ref ext_mgr) = state.extension_manager {
|
||||||
match ext_mgr.configure_token(&extension_name, &token).await {
|
match ext_mgr
|
||||||
|
.configure_token(&extension_name, &token, user_id)
|
||||||
|
.await
|
||||||
|
{
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
if result.verification.is_some() {
|
if result.verification.is_some() {
|
||||||
state.sse.broadcast(
|
state.sse.broadcast_for_user(
|
||||||
crate::channels::web::types::SseEvent::AuthRequired {
|
user_id,
|
||||||
|
crate::channels::web::types::AppEvent::AuthRequired {
|
||||||
extension_name: extension_name.clone(),
|
extension_name: extension_name.clone(),
|
||||||
instructions: Some(result.message),
|
instructions: Some(result.message),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
@@ -275,9 +283,10 @@ async fn handle_client_message(
|
|||||||
},
|
},
|
||||||
);
|
);
|
||||||
} else {
|
} else {
|
||||||
crate::channels::web::server::clear_auth_mode(state).await;
|
crate::channels::web::server::clear_auth_mode(state, user_id).await;
|
||||||
state.sse.broadcast(
|
state.sse.broadcast_for_user(
|
||||||
crate::channels::web::types::SseEvent::AuthCompleted {
|
user_id,
|
||||||
|
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||||
extension_name,
|
extension_name,
|
||||||
success: true,
|
success: true,
|
||||||
message: result.message,
|
message: result.message,
|
||||||
@@ -288,8 +297,9 @@ async fn handle_client_message(
|
|||||||
Err(e) => {
|
Err(e) => {
|
||||||
let msg = format!("Auth failed: {}", e);
|
let msg = format!("Auth failed: {}", e);
|
||||||
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
||||||
state.sse.broadcast(
|
state.sse.broadcast_for_user(
|
||||||
crate::channels::web::types::SseEvent::AuthRequired {
|
user_id,
|
||||||
|
crate::channels::web::types::AppEvent::AuthRequired {
|
||||||
extension_name: extension_name.clone(),
|
extension_name: extension_name.clone(),
|
||||||
instructions: Some(msg.clone()),
|
instructions: Some(msg.clone()),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
@@ -311,7 +321,7 @@ async fn handle_client_message(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
WsClientMessage::AuthCancel { .. } => {
|
WsClientMessage::AuthCancel { .. } => {
|
||||||
crate::channels::web::server::clear_auth_mode(state).await;
|
crate::channels::web::server::clear_auth_mode(state, user_id).await;
|
||||||
}
|
}
|
||||||
WsClientMessage::Ping => {
|
WsClientMessage::Ping => {
|
||||||
let _ = direct_tx.send(WsServerMessage::Pong).await;
|
let _ = direct_tx.send(WsServerMessage::Pong).await;
|
||||||
@@ -498,8 +508,9 @@ mod tests {
|
|||||||
|
|
||||||
GatewayState {
|
GatewayState {
|
||||||
msg_tx: tokio::sync::RwLock::new(msg_tx),
|
msg_tx: tokio::sync::RwLock::new(msg_tx),
|
||||||
sse: SseManager::new(),
|
sse: Arc::new(SseManager::new()),
|
||||||
workspace: None,
|
workspace: None,
|
||||||
|
workspace_pool: None,
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
log_broadcaster: None,
|
log_broadcaster: None,
|
||||||
log_level_handle: None,
|
log_level_handle: None,
|
||||||
@@ -509,13 +520,14 @@ mod tests {
|
|||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
user_id: "test".to_string(),
|
owner_id: "test".to_string(),
|
||||||
|
default_sender_id: "test".to_string(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||||
llm_provider: None,
|
llm_provider: None,
|
||||||
skill_registry: None,
|
skill_registry: None,
|
||||||
skill_catalog: None,
|
skill_catalog: None,
|
||||||
chat_rate_limiter: crate::channels::web::server::RateLimiter::new(30, 60),
|
chat_rate_limiter: crate::channels::web::server::PerUserRateLimiter::new(30, 60),
|
||||||
oauth_rate_limiter: crate::channels::web::server::RateLimiter::new(10, 60),
|
oauth_rate_limiter: crate::channels::web::server::RateLimiter::new(10, 60),
|
||||||
webhook_rate_limiter: crate::channels::web::server::RateLimiter::new(10, 60),
|
webhook_rate_limiter: crate::channels::web::server::RateLimiter::new(10, 60),
|
||||||
registry_entries: Vec::new(),
|
registry_entries: Vec::new(),
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ pub mod import;
|
|||||||
mod logs;
|
mod logs;
|
||||||
mod mcp;
|
mod mcp;
|
||||||
pub mod memory;
|
pub mod memory;
|
||||||
|
mod models;
|
||||||
pub mod oauth_defaults;
|
pub mod oauth_defaults;
|
||||||
mod pairing;
|
mod pairing;
|
||||||
mod registry;
|
mod registry;
|
||||||
@@ -45,6 +46,7 @@ pub use logs::{LogsCommand, run_logs_command};
|
|||||||
pub use mcp::{McpCommand, run_mcp_command};
|
pub use mcp::{McpCommand, run_mcp_command};
|
||||||
pub use memory::MemoryCommand;
|
pub use memory::MemoryCommand;
|
||||||
pub use memory::run_memory_command_with_db;
|
pub use memory::run_memory_command_with_db;
|
||||||
|
pub use models::{ModelsCommand, run_models_command};
|
||||||
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
|
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
|
||||||
pub use registry::{RegistryCommand, run_registry_command};
|
pub use registry::{RegistryCommand, run_registry_command};
|
||||||
pub use routines::{RoutinesCommand, run_routines_command};
|
pub use routines::{RoutinesCommand, run_routines_command};
|
||||||
@@ -217,6 +219,14 @@ pub enum Command {
|
|||||||
)]
|
)]
|
||||||
Hooks(HooksCommand),
|
Hooks(HooksCommand),
|
||||||
|
|
||||||
|
/// Manage LLM providers and models
|
||||||
|
#[command(
|
||||||
|
subcommand,
|
||||||
|
about = "Manage LLM providers and models",
|
||||||
|
long_about = "List providers, view current configuration, and set active provider/model.\nExamples:\n ironclaw models list\n ironclaw models list openai --verbose\n ironclaw models status\n ironclaw models set gpt-4o\n ironclaw models set-provider anthropic --model claude-sonnet-4-6-20250514"
|
||||||
|
)]
|
||||||
|
Models(ModelsCommand),
|
||||||
|
|
||||||
/// Probe external dependencies and validate configuration
|
/// Probe external dependencies and validate configuration
|
||||||
#[command(
|
#[command(
|
||||||
about = "Run diagnostics",
|
about = "Run diagnostics",
|
||||||
|
|||||||
@@ -0,0 +1,864 @@
|
|||||||
|
//! Models management CLI commands.
|
||||||
|
//!
|
||||||
|
//! Provides subcommands for listing providers, viewing current model
|
||||||
|
//! configuration, and setting the active provider/model. Settings are
|
||||||
|
//! persisted to both `config.toml` and `~/.ironclaw/.env` so changes
|
||||||
|
//! take effect immediately (no DB connection required).
|
||||||
|
|
||||||
|
use clap::Subcommand;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
use crate::llm::registry::ProviderRegistry;
|
||||||
|
use crate::settings::Settings;
|
||||||
|
|
||||||
|
#[derive(Subcommand, Debug, Clone)]
|
||||||
|
pub enum ModelsCommand {
|
||||||
|
/// List providers (or available models for a specific provider)
|
||||||
|
List {
|
||||||
|
/// Show only a specific provider (by ID or alias)
|
||||||
|
provider: Option<String>,
|
||||||
|
|
||||||
|
/// Show detailed information (env vars, base URL, protocol)
|
||||||
|
#[arg(short, long)]
|
||||||
|
verbose: bool,
|
||||||
|
|
||||||
|
/// Output as JSON
|
||||||
|
#[arg(long)]
|
||||||
|
json: bool,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Show current model configuration
|
||||||
|
Status {
|
||||||
|
/// Output as JSON
|
||||||
|
#[arg(long)]
|
||||||
|
json: bool,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Set the default model
|
||||||
|
Set {
|
||||||
|
/// Model name (e.g., "gpt-5-mini", "claude-sonnet-4-6-20250514")
|
||||||
|
model: String,
|
||||||
|
},
|
||||||
|
|
||||||
|
/// Set the LLM provider
|
||||||
|
SetProvider {
|
||||||
|
/// Provider ID or alias (e.g., "openai", "anthropic", "ollama")
|
||||||
|
provider: String,
|
||||||
|
|
||||||
|
/// Also set the model (defaults to provider's default model)
|
||||||
|
#[arg(long)]
|
||||||
|
model: Option<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run the models CLI subcommand.
|
||||||
|
pub async fn run_models_command(
|
||||||
|
cmd: ModelsCommand,
|
||||||
|
config_path: Option<&Path>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
match cmd {
|
||||||
|
ModelsCommand::List {
|
||||||
|
provider,
|
||||||
|
verbose,
|
||||||
|
json,
|
||||||
|
} => {
|
||||||
|
if let Some(ref id) = provider {
|
||||||
|
cmd_show_provider(id, verbose, json, config_path).await
|
||||||
|
} else {
|
||||||
|
cmd_list_providers(verbose, json, config_path).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ModelsCommand::Status { json } => cmd_status(json, config_path),
|
||||||
|
ModelsCommand::Set { model } => cmd_set_model(&model, config_path),
|
||||||
|
ModelsCommand::SetProvider { provider, model } => {
|
||||||
|
cmd_set_provider(&provider, model.as_deref(), config_path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Shared helpers ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Resolve the currently active backend and model from env + settings.
|
||||||
|
fn resolve_active(config_path: Option<&Path>) -> (String, String) {
|
||||||
|
let settings = load_settings(config_path);
|
||||||
|
resolve_active_from_settings(&settings)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve active backend + model from a pre-loaded Settings.
|
||||||
|
fn resolve_active_from_settings(settings: &Settings) -> (String, String) {
|
||||||
|
let backend = std::env::var("LLM_BACKEND")
|
||||||
|
.ok()
|
||||||
|
.or_else(|| settings.llm_backend.clone())
|
||||||
|
.unwrap_or_else(|| "nearai".to_string());
|
||||||
|
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
|
||||||
|
let canonical_backend = registry
|
||||||
|
.find(&backend)
|
||||||
|
.map(|d| d.id.clone())
|
||||||
|
.unwrap_or_else(|| backend.clone());
|
||||||
|
|
||||||
|
let model = if canonical_backend == "nearai" {
|
||||||
|
std::env::var("NEARAI_MODEL")
|
||||||
|
.ok()
|
||||||
|
.or_else(|| settings.selected_model.clone())
|
||||||
|
.unwrap_or_else(|| "qwen2.5-72b-instruct:free".to_string())
|
||||||
|
} else if let Some(def) = registry.find(&canonical_backend) {
|
||||||
|
std::env::var(&def.model_env)
|
||||||
|
.ok()
|
||||||
|
.or_else(|| settings.selected_model.clone())
|
||||||
|
.unwrap_or_else(|| def.default_model.clone())
|
||||||
|
} else {
|
||||||
|
settings
|
||||||
|
.selected_model
|
||||||
|
.clone()
|
||||||
|
.unwrap_or_else(|| "unknown".to_string())
|
||||||
|
};
|
||||||
|
|
||||||
|
(canonical_backend, model)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_settings(config_path: Option<&Path>) -> Settings {
|
||||||
|
if let Some(path) = config_path {
|
||||||
|
Settings::load_toml(path).ok().flatten().unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
let toml_path = config_toml_path();
|
||||||
|
if toml_path.exists() {
|
||||||
|
Settings::load_toml(&toml_path)
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
Settings::load()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn save_settings(settings: &Settings, config_path: Option<&Path>) -> anyhow::Result<()> {
|
||||||
|
let path = config_path
|
||||||
|
.map(|p| p.to_path_buf())
|
||||||
|
.unwrap_or_else(config_toml_path);
|
||||||
|
|
||||||
|
settings
|
||||||
|
.save_toml(&path)
|
||||||
|
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn config_toml_path() -> std::path::PathBuf {
|
||||||
|
crate::bootstrap::ironclaw_base_dir().join("config.toml")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to fetch the live model list from a provider.
|
||||||
|
///
|
||||||
|
/// Best-effort: returns `None` if config loading, provider creation, or the
|
||||||
|
/// `list_models()` call fails (missing API key, network error, etc.).
|
||||||
|
async fn try_fetch_models(provider_id: &str, config_path: Option<&Path>) -> Option<Vec<String>> {
|
||||||
|
let config = crate::config::Config::from_env_with_toml(config_path)
|
||||||
|
.await
|
||||||
|
.ok()?;
|
||||||
|
|
||||||
|
// Override backend to the requested provider so create_llm_provider
|
||||||
|
// constructs the right one.
|
||||||
|
let mut llm_config = config.llm.clone();
|
||||||
|
llm_config.backend = provider_id.to_string();
|
||||||
|
|
||||||
|
// For registry providers, resolve the RegistryProviderConfig if not
|
||||||
|
// already set for this backend.
|
||||||
|
if provider_id != "nearai" && provider_id != "bedrock" {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
if let Some(def) = registry.find(provider_id)
|
||||||
|
&& llm_config
|
||||||
|
.provider
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|p| p.provider_id != def.id)
|
||||||
|
{
|
||||||
|
// Build a minimal RegistryProviderConfig from env + registry
|
||||||
|
let api_key = def
|
||||||
|
.api_key_env
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|env| std::env::var(env).ok());
|
||||||
|
if def.api_key_required && api_key.is_none() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let base_url = def.default_base_url.clone().unwrap_or_default();
|
||||||
|
llm_config.provider = Some(crate::llm::RegistryProviderConfig {
|
||||||
|
protocol: def.protocol,
|
||||||
|
provider_id: def.id.clone(),
|
||||||
|
model: def.default_model.clone(),
|
||||||
|
api_key: api_key.map(secrecy::SecretString::from),
|
||||||
|
base_url,
|
||||||
|
extra_headers: Vec::new(),
|
||||||
|
oauth_token: None,
|
||||||
|
is_codex_chatgpt: false,
|
||||||
|
refresh_token: None,
|
||||||
|
auth_path: None,
|
||||||
|
cache_retention: Default::default(),
|
||||||
|
unsupported_params: def.unsupported_params.clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let session = crate::llm::create_session_manager(config.llm.session.clone()).await;
|
||||||
|
let provider = crate::llm::create_llm_provider(&llm_config, session)
|
||||||
|
.await
|
||||||
|
.ok()?;
|
||||||
|
provider.list_models().await.ok().filter(|m| !m.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Print available models section (text output).
|
||||||
|
fn print_model_list(models: &Option<Vec<String>>, active_model: Option<&String>) {
|
||||||
|
match models {
|
||||||
|
Some(models) => {
|
||||||
|
println!("\n Available models ({}):", models.len());
|
||||||
|
for m in models {
|
||||||
|
let marker = active_model
|
||||||
|
.filter(|a| a.as_str() == m)
|
||||||
|
.map(|_| " (active)")
|
||||||
|
.unwrap_or("");
|
||||||
|
println!(" {}{}", m, marker);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
println!(
|
||||||
|
"\n Could not fetch model list (missing credentials or provider unavailable)."
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Also update `~/.ironclaw/.env` so changes take effect immediately.
|
||||||
|
///
|
||||||
|
/// Skipped when `config_path` is `Some` (custom `--config`), because the user
|
||||||
|
/// is explicitly targeting a different config file and we must not pollute the
|
||||||
|
/// default profile's `.env`.
|
||||||
|
fn sync_to_dotenv(config_path: Option<&Path>, vars: &[(&str, &str)]) {
|
||||||
|
if config_path.is_some() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if let Err(e) = crate::bootstrap::upsert_bootstrap_vars(vars) {
|
||||||
|
eprintln!("Warning: failed to update .env: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── status ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn cmd_status(json: bool, config_path: Option<&Path>) -> anyhow::Result<()> {
|
||||||
|
let settings = load_settings(config_path);
|
||||||
|
let (backend, model) = resolve_active_from_settings(&settings);
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
|
||||||
|
let fallback = std::env::var("NEARAI_FALLBACK_MODEL").ok();
|
||||||
|
let cheap = std::env::var("NEARAI_CHEAP_MODEL").ok();
|
||||||
|
|
||||||
|
let description = if backend == "nearai" {
|
||||||
|
"NEAR AI inference (default)".to_string()
|
||||||
|
} else {
|
||||||
|
registry
|
||||||
|
.find(&backend)
|
||||||
|
.map(|d| d.description.clone())
|
||||||
|
.unwrap_or_default()
|
||||||
|
};
|
||||||
|
|
||||||
|
if json {
|
||||||
|
let v = serde_json::json!({
|
||||||
|
"provider": backend,
|
||||||
|
"model": model,
|
||||||
|
"description": description,
|
||||||
|
"fallback_model": fallback,
|
||||||
|
"cheap_model": cheap,
|
||||||
|
});
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".to_string())
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("Provider: {} ({})", backend, description);
|
||||||
|
println!("Model: {}", model);
|
||||||
|
if let Some(ref fb) = fallback {
|
||||||
|
println!("Fallback: {}", fb);
|
||||||
|
}
|
||||||
|
if let Some(ref ch) = cheap {
|
||||||
|
println!("Cheap: {}", ch);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── set ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn cmd_set_model(model: &str, config_path: Option<&Path>) -> anyhow::Result<()> {
|
||||||
|
let trimmed = model.trim();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
anyhow::bail!("Model name cannot be empty");
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut settings = load_settings(config_path);
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
|
||||||
|
// Warn if model name doesn't match any known provider's default model
|
||||||
|
let known_model = registry.all().iter().any(|d| d.default_model == trimmed)
|
||||||
|
|| trimmed.contains("qwen") // nearai models
|
||||||
|
|| trimmed.contains("llama")
|
||||||
|
|| trimmed.contains("gpt")
|
||||||
|
|| trimmed.contains("claude")
|
||||||
|
|| trimmed.contains("gemini")
|
||||||
|
|| trimmed.contains("mistral");
|
||||||
|
if !known_model {
|
||||||
|
eprintln!(
|
||||||
|
"Warning: '{}' is not a recognized model name. Proceeding anyway.",
|
||||||
|
trimmed
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
settings.selected_model = Some(trimmed.to_string());
|
||||||
|
save_settings(&settings, config_path)?;
|
||||||
|
|
||||||
|
let backend = std::env::var("LLM_BACKEND")
|
||||||
|
.ok()
|
||||||
|
.or_else(|| settings.llm_backend.clone())
|
||||||
|
.unwrap_or_else(|| "nearai".to_string());
|
||||||
|
|
||||||
|
// Also write to .env so the change takes effect immediately
|
||||||
|
let model_env = if backend == "nearai" {
|
||||||
|
"NEARAI_MODEL".to_string()
|
||||||
|
} else {
|
||||||
|
registry
|
||||||
|
.find(&backend)
|
||||||
|
.map(|d| d.model_env.clone())
|
||||||
|
.unwrap_or_default()
|
||||||
|
};
|
||||||
|
if !model_env.is_empty() {
|
||||||
|
sync_to_dotenv(config_path, &[(&model_env, trimmed)]);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("Model set to '{}' (provider: {})", trimmed, backend);
|
||||||
|
println!(
|
||||||
|
"Saved to {}",
|
||||||
|
config_path
|
||||||
|
.map(|p| p.display().to_string())
|
||||||
|
.unwrap_or_else(|| config_toml_path().display().to_string())
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── set-provider ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn cmd_set_provider(
|
||||||
|
provider: &str,
|
||||||
|
model: Option<&str>,
|
||||||
|
config_path: Option<&Path>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
|
||||||
|
// Validate and normalize provider
|
||||||
|
let canonical_id = if provider == "nearai" || provider == "near_ai" || provider == "near" {
|
||||||
|
"nearai".to_string()
|
||||||
|
} else {
|
||||||
|
let def = registry.find(provider).ok_or_else(|| {
|
||||||
|
let known: Vec<&str> = std::iter::once("nearai")
|
||||||
|
.chain(registry.all().iter().map(|d| d.id.as_str()))
|
||||||
|
.collect();
|
||||||
|
anyhow::anyhow!(
|
||||||
|
"Unknown provider '{}'. Known providers: {}",
|
||||||
|
provider,
|
||||||
|
known.join(", ")
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
def.id.clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
// Resolve model: explicit > provider default
|
||||||
|
let resolved_model = if let Some(m) = model {
|
||||||
|
m.to_string()
|
||||||
|
} else if canonical_id == "nearai" {
|
||||||
|
"qwen2.5-72b-instruct:free".to_string()
|
||||||
|
} else if let Some(def) = registry.find(&canonical_id) {
|
||||||
|
def.default_model.clone()
|
||||||
|
} else {
|
||||||
|
"default".to_string()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut settings = load_settings(config_path);
|
||||||
|
settings.llm_backend = Some(canonical_id.clone());
|
||||||
|
settings.selected_model = Some(resolved_model.clone());
|
||||||
|
save_settings(&settings, config_path)?;
|
||||||
|
|
||||||
|
// Also write to .env so the change takes effect immediately
|
||||||
|
let model_env = if canonical_id == "nearai" {
|
||||||
|
"NEARAI_MODEL".to_string()
|
||||||
|
} else {
|
||||||
|
registry
|
||||||
|
.find(&canonical_id)
|
||||||
|
.map(|d| d.model_env.clone())
|
||||||
|
.unwrap_or_default()
|
||||||
|
};
|
||||||
|
let mut vars: Vec<(&str, &str)> = vec![("LLM_BACKEND", &canonical_id)];
|
||||||
|
if !model_env.is_empty() {
|
||||||
|
vars.push((&model_env, &resolved_model));
|
||||||
|
}
|
||||||
|
sync_to_dotenv(config_path, &vars);
|
||||||
|
|
||||||
|
println!(
|
||||||
|
"Provider set to '{}', model set to '{}'",
|
||||||
|
canonical_id, resolved_model
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
"Saved to {}",
|
||||||
|
config_path
|
||||||
|
.map(|p| p.display().to_string())
|
||||||
|
.unwrap_or_else(|| config_toml_path().display().to_string())
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── list ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// List all providers with their default models.
|
||||||
|
async fn cmd_list_providers(
|
||||||
|
verbose: bool,
|
||||||
|
json: bool,
|
||||||
|
config_path: Option<&Path>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
let (active_backend, active_model) = resolve_active(config_path);
|
||||||
|
|
||||||
|
if json {
|
||||||
|
let mut entries: Vec<serde_json::Value> = Vec::new();
|
||||||
|
|
||||||
|
// NEAR AI (not in registry)
|
||||||
|
let nearai_active = active_backend == "nearai";
|
||||||
|
entries.push(serde_json::json!({
|
||||||
|
"id": "nearai",
|
||||||
|
"description": "NEAR AI inference (default)",
|
||||||
|
"default_model": "qwen2.5-72b-instruct:free",
|
||||||
|
"active": nearai_active,
|
||||||
|
"active_model": if nearai_active { Some(&active_model) } else { None },
|
||||||
|
}));
|
||||||
|
|
||||||
|
for def in registry.all() {
|
||||||
|
let is_active = active_backend == def.id;
|
||||||
|
let mut v = serde_json::json!({
|
||||||
|
"id": def.id,
|
||||||
|
"description": def.description,
|
||||||
|
"default_model": def.default_model,
|
||||||
|
"protocol": format!("{:?}", def.protocol),
|
||||||
|
"active": is_active,
|
||||||
|
});
|
||||||
|
if is_active {
|
||||||
|
v["active_model"] = serde_json::json!(active_model);
|
||||||
|
}
|
||||||
|
if verbose {
|
||||||
|
v["aliases"] = serde_json::json!(def.aliases);
|
||||||
|
v["model_env"] = serde_json::json!(def.model_env);
|
||||||
|
v["api_key_env"] = serde_json::json!(def.api_key_env);
|
||||||
|
v["api_key_required"] = serde_json::json!(def.api_key_required);
|
||||||
|
if let Some(ref url) = def.default_base_url {
|
||||||
|
v["base_url"] = serde_json::json!(url);
|
||||||
|
}
|
||||||
|
if let Some(ref setup) = def.setup {
|
||||||
|
v["can_list_models"] = serde_json::json!(setup.can_list_models());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
entries.push(v);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
serde_json::to_string_pretty(&entries).unwrap_or_else(|_| "[]".to_string())
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let providers = registry.all();
|
||||||
|
|
||||||
|
println!("Active: {} (model: {})\n", active_backend, active_model);
|
||||||
|
println!(
|
||||||
|
"{} provider(s) available:\n",
|
||||||
|
providers.len() + 1 // +1 for NEAR AI
|
||||||
|
);
|
||||||
|
|
||||||
|
// NEAR AI (not in registry)
|
||||||
|
let nearai_marker = if active_backend == "nearai" { " *" } else { "" };
|
||||||
|
if verbose {
|
||||||
|
println!(" nearai{}", nearai_marker);
|
||||||
|
println!(" Description: NEAR AI inference (default)");
|
||||||
|
println!(" Default model: qwen2.5-72b-instruct:free");
|
||||||
|
println!(" Model env: NEARAI_MODEL");
|
||||||
|
if active_backend == "nearai" {
|
||||||
|
println!(" Active model: {}", active_model);
|
||||||
|
}
|
||||||
|
println!();
|
||||||
|
} else {
|
||||||
|
println!(
|
||||||
|
" {:<22} {:<40} NEAR AI inference (default)",
|
||||||
|
format!("nearai{nearai_marker}"),
|
||||||
|
"qwen2.5-72b-instruct:free"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
for def in providers {
|
||||||
|
let is_active = active_backend == def.id;
|
||||||
|
let marker = if is_active { " *" } else { "" };
|
||||||
|
|
||||||
|
if verbose {
|
||||||
|
println!(" {}{}", def.id, marker);
|
||||||
|
println!(" Description: {}", def.description);
|
||||||
|
println!(" Default model: {}", def.default_model);
|
||||||
|
println!(" Protocol: {:?}", def.protocol);
|
||||||
|
println!(" Model env: {}", def.model_env);
|
||||||
|
if let Some(ref env) = def.api_key_env {
|
||||||
|
println!(
|
||||||
|
" API key env: {} ({})",
|
||||||
|
env,
|
||||||
|
if def.api_key_required {
|
||||||
|
"required"
|
||||||
|
} else {
|
||||||
|
"optional"
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if let Some(ref url) = def.default_base_url {
|
||||||
|
println!(" Base URL: {}", url);
|
||||||
|
}
|
||||||
|
if !def.aliases.is_empty() {
|
||||||
|
println!(" Aliases: {}", def.aliases.join(", "));
|
||||||
|
}
|
||||||
|
if is_active {
|
||||||
|
println!(" Active model: {}", active_model);
|
||||||
|
}
|
||||||
|
println!();
|
||||||
|
} else {
|
||||||
|
let model_display = if is_active {
|
||||||
|
active_model.clone()
|
||||||
|
} else {
|
||||||
|
def.default_model.clone()
|
||||||
|
};
|
||||||
|
println!(
|
||||||
|
" {:<22} {:<40} {}",
|
||||||
|
format!("{}{marker}", def.id),
|
||||||
|
model_display,
|
||||||
|
def.description,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !verbose {
|
||||||
|
println!();
|
||||||
|
println!("* = active provider. Use --verbose for details.");
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Show details for a specific provider.
|
||||||
|
async fn cmd_show_provider(
|
||||||
|
id: &str,
|
||||||
|
verbose: bool,
|
||||||
|
json: bool,
|
||||||
|
config_path: Option<&Path>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
let (active_backend, active_model) = resolve_active(config_path);
|
||||||
|
|
||||||
|
// Resolve canonical ID for model fetching
|
||||||
|
let canonical_id = if id == "nearai" || id == "near_ai" || id == "near" {
|
||||||
|
"nearai".to_string()
|
||||||
|
} else {
|
||||||
|
registry
|
||||||
|
.find(id)
|
||||||
|
.map(|d| d.id.clone())
|
||||||
|
.unwrap_or_else(|| id.to_string())
|
||||||
|
};
|
||||||
|
|
||||||
|
// Try to fetch live model list from the provider
|
||||||
|
let live_models = try_fetch_models(&canonical_id, config_path).await;
|
||||||
|
|
||||||
|
// Check NEAR AI first (not in registry)
|
||||||
|
if id == "nearai" || id == "near_ai" || id == "near" {
|
||||||
|
let is_active = active_backend == "nearai";
|
||||||
|
if json {
|
||||||
|
let mut v = serde_json::json!({
|
||||||
|
"id": "nearai",
|
||||||
|
"description": "NEAR AI inference (default)",
|
||||||
|
"default_model": "qwen2.5-72b-instruct:free",
|
||||||
|
"model_env": "NEARAI_MODEL",
|
||||||
|
"active": is_active,
|
||||||
|
});
|
||||||
|
if is_active {
|
||||||
|
v["active_model"] = serde_json::json!(active_model);
|
||||||
|
}
|
||||||
|
if let Some(ref models) = live_models {
|
||||||
|
v["available_models"] = serde_json::json!(models);
|
||||||
|
}
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".to_string())
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
println!("Provider: nearai");
|
||||||
|
println!(" Description: NEAR AI inference (default)");
|
||||||
|
println!(" Default model: qwen2.5-72b-instruct:free");
|
||||||
|
println!(" Model env: NEARAI_MODEL");
|
||||||
|
println!(" Active: {}", if is_active { "yes" } else { "no" });
|
||||||
|
if is_active {
|
||||||
|
println!(" Active model: {}", active_model);
|
||||||
|
}
|
||||||
|
print_model_list(&live_models, is_active.then_some(&active_model));
|
||||||
|
}
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let def = registry.find(id).ok_or_else(|| {
|
||||||
|
let known: Vec<&str> = std::iter::once("nearai")
|
||||||
|
.chain(registry.all().iter().map(|d| d.id.as_str()))
|
||||||
|
.collect();
|
||||||
|
anyhow::anyhow!(
|
||||||
|
"Unknown provider '{}'. Known providers: {}",
|
||||||
|
id,
|
||||||
|
known.join(", ")
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let is_active = active_backend == def.id;
|
||||||
|
|
||||||
|
if json {
|
||||||
|
let mut v = serde_json::json!({
|
||||||
|
"id": def.id,
|
||||||
|
"description": def.description,
|
||||||
|
"protocol": format!("{:?}", def.protocol),
|
||||||
|
"default_model": def.default_model,
|
||||||
|
"model_env": def.model_env,
|
||||||
|
"api_key_env": def.api_key_env,
|
||||||
|
"api_key_required": def.api_key_required,
|
||||||
|
"aliases": def.aliases,
|
||||||
|
"active": is_active,
|
||||||
|
});
|
||||||
|
if let Some(ref url) = def.default_base_url {
|
||||||
|
v["base_url"] = serde_json::json!(url);
|
||||||
|
}
|
||||||
|
if let Some(ref setup) = def.setup {
|
||||||
|
v["can_list_models"] = serde_json::json!(setup.can_list_models());
|
||||||
|
v["display_name"] = serde_json::json!(setup.display_name());
|
||||||
|
}
|
||||||
|
if is_active {
|
||||||
|
v["active_model"] = serde_json::json!(active_model);
|
||||||
|
}
|
||||||
|
if verbose && !def.unsupported_params.is_empty() {
|
||||||
|
v["unsupported_params"] = serde_json::json!(def.unsupported_params);
|
||||||
|
}
|
||||||
|
if let Some(ref models) = live_models {
|
||||||
|
v["available_models"] = serde_json::json!(models);
|
||||||
|
}
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".to_string())
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("Provider: {}", def.id);
|
||||||
|
println!(" Description: {}", def.description);
|
||||||
|
println!(" Protocol: {:?}", def.protocol);
|
||||||
|
println!(" Default model: {}", def.default_model);
|
||||||
|
println!(" Model env: {}", def.model_env);
|
||||||
|
if let Some(ref env) = def.api_key_env {
|
||||||
|
println!(
|
||||||
|
" API key env: {} ({})",
|
||||||
|
env,
|
||||||
|
if def.api_key_required {
|
||||||
|
"required"
|
||||||
|
} else {
|
||||||
|
"optional"
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if let Some(ref url) = def.default_base_url {
|
||||||
|
println!(" Base URL: {}", url);
|
||||||
|
}
|
||||||
|
if !def.aliases.is_empty() {
|
||||||
|
println!(" Aliases: {}", def.aliases.join(", "));
|
||||||
|
}
|
||||||
|
if let Some(ref setup) = def.setup {
|
||||||
|
println!(
|
||||||
|
" List models: {}",
|
||||||
|
if setup.can_list_models() {
|
||||||
|
"supported"
|
||||||
|
} else {
|
||||||
|
"not supported"
|
||||||
|
}
|
||||||
|
);
|
||||||
|
println!(" Display name: {}", setup.display_name());
|
||||||
|
}
|
||||||
|
if !def.unsupported_params.is_empty() {
|
||||||
|
println!(" Unsupported: {}", def.unsupported_params.join(", "));
|
||||||
|
}
|
||||||
|
println!(" Active: {}", if is_active { "yes" } else { "no" });
|
||||||
|
if is_active {
|
||||||
|
println!(" Active model: {}", active_model);
|
||||||
|
}
|
||||||
|
print_model_list(&live_models, is_active.then_some(&active_model));
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn resolve_active_defaults_to_nearai() {
|
||||||
|
let settings = Settings::default();
|
||||||
|
assert!(settings.llm_backend.is_none());
|
||||||
|
assert!(settings.selected_model.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn registry_loads_all_providers() {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
let all = registry.all();
|
||||||
|
assert!(
|
||||||
|
all.len() >= 10,
|
||||||
|
"should have at least 10 built-in providers, got {}",
|
||||||
|
all.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn registry_find_by_alias() {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
let def = registry
|
||||||
|
.find("claude")
|
||||||
|
.expect("claude alias should resolve");
|
||||||
|
assert_eq!(def.id, "anthropic");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn all_providers_have_description() {
|
||||||
|
let registry = ProviderRegistry::load();
|
||||||
|
for def in registry.all() {
|
||||||
|
assert!(
|
||||||
|
!def.description.is_empty(),
|
||||||
|
"provider {} should have a description",
|
||||||
|
def.id
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_model_persists_to_toml() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
cmd_set_model("gpt-5-mini", Some(&toml_path)).expect("set model");
|
||||||
|
|
||||||
|
let settings = Settings::load_toml(&toml_path)
|
||||||
|
.expect("read toml")
|
||||||
|
.expect("should have settings");
|
||||||
|
assert_eq!(settings.selected_model.as_deref(), Some("gpt-5-mini"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_provider_validates_unknown() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
let result = cmd_set_provider("nonexistent_provider", None, Some(&toml_path));
|
||||||
|
assert!(result.is_err());
|
||||||
|
let err = result.unwrap_err().to_string();
|
||||||
|
assert!(
|
||||||
|
err.contains("Unknown provider"),
|
||||||
|
"should mention unknown provider: {}",
|
||||||
|
err
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_provider_persists_to_toml() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
cmd_set_provider("groq", None, Some(&toml_path)).expect("set provider");
|
||||||
|
|
||||||
|
let settings = Settings::load_toml(&toml_path)
|
||||||
|
.expect("read toml")
|
||||||
|
.expect("should have settings");
|
||||||
|
assert_eq!(settings.llm_backend.as_deref(), Some("groq"));
|
||||||
|
assert_eq!(
|
||||||
|
settings.selected_model.as_deref(),
|
||||||
|
Some("llama-3.3-70b-versatile")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_provider_with_custom_model() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
cmd_set_provider("anthropic", Some("claude-opus-4-6"), Some(&toml_path))
|
||||||
|
.expect("set provider with model");
|
||||||
|
|
||||||
|
let settings = Settings::load_toml(&toml_path)
|
||||||
|
.expect("read toml")
|
||||||
|
.expect("should have settings");
|
||||||
|
assert_eq!(settings.llm_backend.as_deref(), Some("anthropic"));
|
||||||
|
assert_eq!(settings.selected_model.as_deref(), Some("claude-opus-4-6"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn custom_config_does_not_pollute_default_dotenv() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
// With a custom config path, sync_to_dotenv should be a no-op
|
||||||
|
// (it returns early when config_path is Some).
|
||||||
|
// We verify by checking that cmd_set_provider succeeds without
|
||||||
|
// trying to write to the default ~/.ironclaw/.env.
|
||||||
|
cmd_set_provider("groq", None, Some(&toml_path)).expect("set provider with custom config");
|
||||||
|
|
||||||
|
let settings = Settings::load_toml(&toml_path)
|
||||||
|
.expect("read toml")
|
||||||
|
.expect("should have settings");
|
||||||
|
assert_eq!(settings.llm_backend.as_deref(), Some("groq"));
|
||||||
|
// The key assertion is that no error was thrown trying to write
|
||||||
|
// to the default .env — sync_to_dotenv skipped it.
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_model_rejects_empty_name() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
let result = cmd_set_model("", Some(&toml_path));
|
||||||
|
assert!(result.is_err());
|
||||||
|
assert!(
|
||||||
|
result.unwrap_err().to_string().contains("cannot be empty"),
|
||||||
|
"should reject empty model name"
|
||||||
|
);
|
||||||
|
|
||||||
|
let result2 = cmd_set_model(" ", Some(&toml_path));
|
||||||
|
assert!(result2.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_provider_normalizes_alias() {
|
||||||
|
let dir = tempfile::tempdir().expect("create temp dir");
|
||||||
|
let toml_path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
cmd_set_provider("claude", None, Some(&toml_path)).expect("set via alias");
|
||||||
|
|
||||||
|
let settings = Settings::load_toml(&toml_path)
|
||||||
|
.expect("read toml")
|
||||||
|
.expect("should have settings");
|
||||||
|
assert_eq!(
|
||||||
|
settings.llm_backend.as_deref(),
|
||||||
|
Some("anthropic"),
|
||||||
|
"alias should be normalized to canonical ID"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
+402
-26
@@ -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 ──────────────────────────────────────────────
|
// ── Shared callback server ──────────────────────────────────────────────
|
||||||
|
|
||||||
// Core OAuth callback infrastructure is defined in `crate::llm::oauth_helpers`
|
// Core OAuth callback infrastructure is defined in `crate::llm::oauth_helpers`
|
||||||
@@ -447,8 +471,8 @@ pub struct PendingOAuthFlow {
|
|||||||
pub user_id: String,
|
pub user_id: String,
|
||||||
/// Secrets store reference for token persistence.
|
/// Secrets store reference for token persistence.
|
||||||
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||||
/// SSE broadcast sender for notifying the web UI.
|
/// SSE broadcast manager for notifying the web UI.
|
||||||
pub sse_sender: Option<tokio::sync::broadcast::Sender<crate::channels::web::types::SseEvent>>,
|
pub sse_manager: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||||
pub gateway_token: Option<String>,
|
pub gateway_token: Option<String>,
|
||||||
/// Additional form params for the token exchange request.
|
/// Additional form params for the token exchange request.
|
||||||
@@ -661,6 +685,48 @@ pub struct ProxyTokenExchangeRequest<'a> {
|
|||||||
pub extra_token_params: &'a HashMap<String, String>,
|
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.
|
/// Exchange an OAuth authorization code via the platform's token exchange proxy.
|
||||||
///
|
///
|
||||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
/// 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()
|
let client = reqwest::Client::builder()
|
||||||
.timeout(Duration::from_secs(60))
|
.timeout(Duration::from_secs(60))
|
||||||
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
.build()
|
.build()
|
||||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
|
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
|
||||||
let mut params = vec![
|
let mut params = vec![
|
||||||
@@ -724,41 +791,350 @@ pub async fn exchange_via_proxy(
|
|||||||
.json()
|
.json()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?;
|
.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
|
/// Refresh an OAuth access token via the platform's token refresh proxy.
|
||||||
.get(request.access_token_field)
|
///
|
||||||
.and_then(|v| v.as_str())
|
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||||
.ok_or_else(|| {
|
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||||
let fields: Vec<&str> = token_data
|
/// the provider requires it.
|
||||||
.as_object()
|
pub async fn refresh_token_via_proxy(
|
||||||
.map(|o| o.keys().map(|k| k.as_str()).collect())
|
request: ProxyRefreshTokenRequest<'_>,
|
||||||
.unwrap_or_default();
|
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||||
OAuthCallbackError::Io(format!(
|
if request.gateway_token.is_empty() {
|
||||||
"No '{}' field in proxy response (fields present: {:?})",
|
return Err(OAuthCallbackError::Io(
|
||||||
request.access_token_field, fields
|
"Gateway auth token is required for proxy token refresh".to_string(),
|
||||||
))
|
));
|
||||||
})?
|
}
|
||||||
.to_string();
|
|
||||||
|
|
||||||
let refresh_token = token_data
|
let refresh_url = format!("{}/oauth/refresh", request.proxy_url.trim_end_matches('/'));
|
||||||
.get("refresh_token")
|
let client = reqwest::Client::builder()
|
||||||
.and_then(|v| v.as_str())
|
.timeout(Duration::from_secs(15))
|
||||||
.map(String::from);
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
|
.build()
|
||||||
|
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
|
||||||
|
|
||||||
Ok(OAuthTokenResponse {
|
let mut params = vec![
|
||||||
access_token,
|
("refresh_token", request.refresh_token.to_string()),
|
||||||
refresh_token,
|
("token_url", request.token_url.to_string()),
|
||||||
expires_in,
|
("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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
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::{
|
use crate::cli::oauth_defaults::{
|
||||||
builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html,
|
builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html,
|
||||||
};
|
};
|
||||||
use crate::config::helpers::lock_env;
|
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]
|
#[test]
|
||||||
fn test_is_loopback_host() {
|
fn test_is_loopback_host() {
|
||||||
|
|||||||
+66
-10
@@ -10,7 +10,7 @@ use clap::Subcommand;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::routine::{
|
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;
|
use crate::db::Database;
|
||||||
|
|
||||||
@@ -251,15 +251,26 @@ async fn list(
|
|||||||
);
|
);
|
||||||
println!("{}", "-".repeat(130));
|
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 {
|
for r in &filtered {
|
||||||
let status = if r.enabled {
|
let last_run_status = last_run_results.get(&r.id).copied();
|
||||||
if r.consecutive_failures > 0 {
|
|
||||||
format!("err({})", r.consecutive_failures)
|
let status = if !r.enabled {
|
||||||
} else {
|
|
||||||
"active".to_string()
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
"disabled".to_string()
|
"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
|
let next_fire = r
|
||||||
@@ -340,8 +351,8 @@ async fn create(
|
|||||||
prompt: prompt.to_string(),
|
prompt: prompt.to_string(),
|
||||||
context_paths: Vec::new(),
|
context_paths: Vec::new(),
|
||||||
max_tokens: 4096,
|
max_tokens: 4096,
|
||||||
use_tools: false,
|
use_tools: true,
|
||||||
max_tool_rounds: 0,
|
max_tool_rounds: 3,
|
||||||
},
|
},
|
||||||
guardrails: RoutineGuardrails {
|
guardrails: RoutineGuardrails {
|
||||||
cooldown: std::time::Duration::from_secs(cooldown_secs),
|
cooldown: std::time::Duration::from_secs(cooldown_secs),
|
||||||
@@ -685,6 +696,7 @@ fn truncate(s: &str, max_chars: usize) -> String {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::agent::routine::RoutineAction;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn format_relative_future() {
|
fn format_relative_future() {
|
||||||
@@ -743,4 +755,48 @@ mod tests {
|
|||||||
assert!(notify.on_failure); // safety: test-only assertion
|
assert!(notify.on_failure); // safety: test-only assertion
|
||||||
assert!(!notify.on_success); // safety: test-only assertion
|
assert!(!notify.on_success); // safety: test-only assertion
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cli_create_defaults_lightweight_routines_to_tools_enabled() {
|
||||||
|
let harness = crate::testing::TestHarnessBuilder::new().build().await;
|
||||||
|
let db = harness.db.clone();
|
||||||
|
|
||||||
|
run_routines_command(
|
||||||
|
RoutinesCommand::Create {
|
||||||
|
name: "cli-digest".to_string(),
|
||||||
|
schedule: "0 0 9 * * *".to_string(),
|
||||||
|
prompt: "Prepare the morning digest.".to_string(),
|
||||||
|
description: "CLI created routine".to_string(),
|
||||||
|
timezone: Some("UTC".to_string()),
|
||||||
|
cooldown: 300,
|
||||||
|
notify_channel: None,
|
||||||
|
},
|
||||||
|
db.clone(),
|
||||||
|
"user1",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("create routine");
|
||||||
|
|
||||||
|
let routine = db
|
||||||
|
.get_routine_by_name("user1", "cli-digest")
|
||||||
|
.await
|
||||||
|
.expect("get routine by name")
|
||||||
|
.expect("cli-digest should exist");
|
||||||
|
|
||||||
|
match routine.action {
|
||||||
|
RoutineAction::Lightweight {
|
||||||
|
use_tools,
|
||||||
|
max_tool_rounds,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
assert!(
|
||||||
|
use_tools,
|
||||||
|
"CLI-created lightweight routines should default to tools"
|
||||||
|
);
|
||||||
|
assert_eq!(max_tool_rounds, 3);
|
||||||
|
}
|
||||||
|
other => panic!("expected lightweight action, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ Commands:
|
|||||||
service Manage OS service
|
service Manage OS service
|
||||||
skills Manage skills
|
skills Manage skills
|
||||||
hooks Manage lifecycle hooks
|
hooks Manage lifecycle hooks
|
||||||
|
models Manage LLM providers and models
|
||||||
doctor Run diagnostics
|
doctor Run diagnostics
|
||||||
logs View and manage gateway logs
|
logs View and manage gateway logs
|
||||||
status Show system status
|
status Show system status
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ Commands:
|
|||||||
service Manage OS service
|
service Manage OS service
|
||||||
skills Manage skills
|
skills Manage skills
|
||||||
hooks Manage lifecycle hooks
|
hooks Manage lifecycle hooks
|
||||||
|
models Manage LLM providers and models
|
||||||
doctor Run diagnostics
|
doctor Run diagnostics
|
||||||
logs View and manage gateway logs
|
logs View and manage gateway logs
|
||||||
status Show system status
|
status Show system status
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ Commands:
|
|||||||
service Manage OS service
|
service Manage OS service
|
||||||
skills Manage skills
|
skills Manage skills
|
||||||
hooks Manage lifecycle hooks
|
hooks Manage lifecycle hooks
|
||||||
|
models Manage LLM providers and models
|
||||||
doctor Run diagnostics
|
doctor Run diagnostics
|
||||||
logs View and manage gateway logs
|
logs View and manage gateway logs
|
||||||
status Show system status
|
status Show system status
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ Commands:
|
|||||||
service Manage OS service
|
service Manage OS service
|
||||||
skills Manage skills
|
skills Manage skills
|
||||||
hooks Manage lifecycle hooks
|
hooks Manage lifecycle hooks
|
||||||
|
models Manage LLM providers and models
|
||||||
doctor Run diagnostics
|
doctor Run diagnostics
|
||||||
logs View and manage gateway logs
|
logs View and manage gateway logs
|
||||||
status Show system status
|
status Show system status
|
||||||
|
|||||||
+270
-16
@@ -2,6 +2,7 @@
|
|||||||
//!
|
//!
|
||||||
//! Commands for installing, listing, removing, and authenticating WASM tools.
|
//! Commands for installing, listing, removing, and authenticating WASM tools.
|
||||||
|
|
||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -79,6 +80,10 @@ pub enum ToolCommand {
|
|||||||
/// Directory to look for tool (default: ~/.ironclaw/tools/)
|
/// Directory to look for tool (default: ~/.ironclaw/tools/)
|
||||||
#[arg(short, long)]
|
#[arg(short, long)]
|
||||||
dir: Option<PathBuf>,
|
dir: Option<PathBuf>,
|
||||||
|
|
||||||
|
/// User ID for checking credential status (default: "default")
|
||||||
|
#[arg(short, long, default_value = "default")]
|
||||||
|
user: String,
|
||||||
},
|
},
|
||||||
|
|
||||||
/// Configure authentication for a tool
|
/// Configure authentication for a tool
|
||||||
@@ -124,7 +129,11 @@ pub async fn run_tool_command(cmd: ToolCommand) -> anyhow::Result<()> {
|
|||||||
} => install_tool(path, name, capabilities, target, release, skip_build, force).await,
|
} => install_tool(path, name, capabilities, target, release, skip_build, force).await,
|
||||||
ToolCommand::List { dir, verbose } => list_tools(dir, verbose).await,
|
ToolCommand::List { dir, verbose } => list_tools(dir, verbose).await,
|
||||||
ToolCommand::Remove { name, dir } => remove_tool(name, dir).await,
|
ToolCommand::Remove { name, dir } => remove_tool(name, dir).await,
|
||||||
ToolCommand::Info { name_or_path, dir } => show_tool_info(name_or_path, dir).await,
|
ToolCommand::Info {
|
||||||
|
name_or_path,
|
||||||
|
dir,
|
||||||
|
user,
|
||||||
|
} => show_tool_info(name_or_path, dir, user).await,
|
||||||
ToolCommand::Auth { name, dir, user } => auth_tool(name, dir, user).await,
|
ToolCommand::Auth { name, dir, user } => auth_tool(name, dir, user).await,
|
||||||
ToolCommand::Setup { name, dir, user } => setup_tool(name, dir, user).await,
|
ToolCommand::Setup { name, dir, user } => setup_tool(name, dir, user).await,
|
||||||
}
|
}
|
||||||
@@ -388,7 +397,11 @@ async fn remove_tool(name: String, dir: Option<PathBuf>) -> anyhow::Result<()> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Show information about a tool.
|
/// Show information about a tool.
|
||||||
async fn show_tool_info(name_or_path: String, dir: Option<PathBuf>) -> anyhow::Result<()> {
|
async fn show_tool_info(
|
||||||
|
name_or_path: String,
|
||||||
|
dir: Option<PathBuf>,
|
||||||
|
user_id: String,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
let wasm_path = if name_or_path.ends_with(".wasm") {
|
let wasm_path = if name_or_path.ends_with(".wasm") {
|
||||||
PathBuf::from(&name_or_path)
|
PathBuf::from(&name_or_path)
|
||||||
} else {
|
} else {
|
||||||
@@ -423,7 +436,37 @@ async fn show_tool_info(name_or_path: String, dir: Option<PathBuf>) -> anyhow::R
|
|||||||
println!("\nCapabilities ({}):", caps_path.display());
|
println!("\nCapabilities ({}):", caps_path.display());
|
||||||
let content = fs::read_to_string(&caps_path).await?;
|
let content = fs::read_to_string(&caps_path).await?;
|
||||||
match CapabilitiesFile::from_json(&content) {
|
match CapabilitiesFile::from_json(&content) {
|
||||||
Ok(caps) => print_capabilities_detail(&caps),
|
Ok(caps) => {
|
||||||
|
// Lazily init secrets store only when auth secrets need checking.
|
||||||
|
let has_auth = caps.auth.is_some()
|
||||||
|
|| caps
|
||||||
|
.setup
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|s| !s.required_secrets.is_empty())
|
||||||
|
|| caps
|
||||||
|
.http
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|h| !h.credentials.is_empty());
|
||||||
|
let secrets_store = if has_auth {
|
||||||
|
match init_secrets_store().await {
|
||||||
|
Ok(store) => Some(store),
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!(" Warning: could not init secrets store: {}", e);
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
print_capabilities_detail(
|
||||||
|
&caps,
|
||||||
|
secrets_store
|
||||||
|
.as_ref()
|
||||||
|
.map(|s| s.as_ref() as &(dyn SecretsStore + Send + Sync)),
|
||||||
|
&user_id,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
Err(e) => println!(" Error parsing: {}", e),
|
Err(e) => println!(" Error parsing: {}", e),
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
@@ -476,8 +519,89 @@ fn print_capabilities_summary(caps: &CapabilitiesFile) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Per-secret info collected from all auth-related capability sections.
|
||||||
|
struct AuthSecretInfo {
|
||||||
|
secret_name: String,
|
||||||
|
/// Human-readable label (from auth.display_name or setup prompt).
|
||||||
|
description: Option<String>,
|
||||||
|
/// Injection location (from http.credentials).
|
||||||
|
location: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collected auth secrets and the set of secret names they cover.
|
||||||
|
struct CollectedAuthSecrets {
|
||||||
|
secrets: Vec<AuthSecretInfo>,
|
||||||
|
/// Secret names present in `secrets`, for filtering the Secrets capability section.
|
||||||
|
seen_names: HashSet<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collect and deduplicate auth secrets from all auth-related capability sections.
|
||||||
|
///
|
||||||
|
/// Priority for the description label: auth.display_name > setup.required_secrets.prompt.
|
||||||
|
/// Injection location is merged from http.credentials.
|
||||||
|
fn collect_auth_secrets(caps: &CapabilitiesFile) -> CollectedAuthSecrets {
|
||||||
|
let mut secrets: Vec<AuthSecretInfo> = Vec::new();
|
||||||
|
let mut seen: HashMap<String, usize> = HashMap::new();
|
||||||
|
|
||||||
|
// auth.display_name is the best label — seed first.
|
||||||
|
if let Some(ref auth) = caps.auth {
|
||||||
|
let index = secrets.len();
|
||||||
|
seen.insert(auth.secret_name.clone(), index);
|
||||||
|
secrets.push(AuthSecretInfo {
|
||||||
|
secret_name: auth.secret_name.clone(),
|
||||||
|
description: auth.display_name.clone(),
|
||||||
|
location: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// setup.required_secrets.prompt is second-best label.
|
||||||
|
if let Some(ref setup) = caps.setup {
|
||||||
|
for secret in &setup.required_secrets {
|
||||||
|
if !seen.contains_key(&secret.name) {
|
||||||
|
let index = secrets.len();
|
||||||
|
seen.insert(secret.name.clone(), index);
|
||||||
|
secrets.push(AuthSecretInfo {
|
||||||
|
secret_name: secret.name.clone(),
|
||||||
|
description: Some(secret.prompt.clone()),
|
||||||
|
location: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Merge injection location from http.credentials.
|
||||||
|
if let Some(ref http) = caps.http {
|
||||||
|
for cred in http.credentials.values() {
|
||||||
|
let loc = format!("{:?}", cred.location);
|
||||||
|
if let Some(&index) = seen.get(&cred.secret_name) {
|
||||||
|
secrets[index].location = Some(loc);
|
||||||
|
} else {
|
||||||
|
let index = secrets.len();
|
||||||
|
seen.insert(cred.secret_name.clone(), index);
|
||||||
|
secrets.push(AuthSecretInfo {
|
||||||
|
secret_name: cred.secret_name.clone(),
|
||||||
|
description: None,
|
||||||
|
location: Some(loc),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let seen_names = seen.into_keys().collect();
|
||||||
|
CollectedAuthSecrets {
|
||||||
|
secrets,
|
||||||
|
seen_names,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Print detailed capabilities.
|
/// Print detailed capabilities.
|
||||||
fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
async fn print_capabilities_detail(
|
||||||
|
caps: &CapabilitiesFile,
|
||||||
|
secrets_store: Option<&(dyn SecretsStore + Send + Sync)>,
|
||||||
|
user_id: &str,
|
||||||
|
) {
|
||||||
|
let mut collected = collect_auth_secrets(caps);
|
||||||
|
|
||||||
if let Some(ref http) = caps.http {
|
if let Some(ref http) = caps.http {
|
||||||
println!(" HTTP:");
|
println!(" HTTP:");
|
||||||
for endpoint in &http.allowlist {
|
for endpoint in &http.allowlist {
|
||||||
@@ -490,13 +614,6 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
|||||||
println!(" {} {} {}", methods, endpoint.host, path);
|
println!(" {} {} {}", methods, endpoint.host, path);
|
||||||
}
|
}
|
||||||
|
|
||||||
if !http.credentials.is_empty() {
|
|
||||||
println!(" Credentials:");
|
|
||||||
for (key, cred) in &http.credentials {
|
|
||||||
println!(" {}: {} -> {:?}", key, cred.secret_name, cred.location);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(ref rate) = http.rate_limit {
|
if let Some(ref rate) = http.rate_limit {
|
||||||
println!(
|
println!(
|
||||||
" Rate limit: {}/min, {}/hour",
|
" Rate limit: {}/min, {}/hour",
|
||||||
@@ -505,12 +622,24 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Filter secrets already covered by the auth section (always rendered when non-empty).
|
||||||
if let Some(ref secrets) = caps.secrets
|
if let Some(ref secrets) = caps.secrets
|
||||||
&& !secrets.allowed_names.is_empty()
|
&& !secrets.allowed_names.is_empty()
|
||||||
{
|
{
|
||||||
println!(" Secrets (existence check only):");
|
let extra: Vec<_> = if collected.secrets.is_empty() {
|
||||||
for name in &secrets.allowed_names {
|
secrets.allowed_names.iter().collect()
|
||||||
println!(" {}", name);
|
} else {
|
||||||
|
secrets
|
||||||
|
.allowed_names
|
||||||
|
.iter()
|
||||||
|
.filter(|name| !collected.seen_names.contains(name.as_str()))
|
||||||
|
.collect()
|
||||||
|
};
|
||||||
|
if !extra.is_empty() {
|
||||||
|
println!(" Secrets (existence check only):");
|
||||||
|
for name in extra {
|
||||||
|
println!(" {}", name);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -531,6 +660,38 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
|||||||
println!(" {}", prefix);
|
println!(" {}", prefix);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Consolidated auth status — sorted by secret name for deterministic output.
|
||||||
|
if !collected.secrets.is_empty() {
|
||||||
|
collected
|
||||||
|
.secrets
|
||||||
|
.sort_by(|a, b| a.secret_name.cmp(&b.secret_name));
|
||||||
|
println!(" Auth:");
|
||||||
|
for info in &collected.secrets {
|
||||||
|
let (icon, label) = match secrets_store {
|
||||||
|
Some(store) => match store.exists(user_id, &info.secret_name).await {
|
||||||
|
Ok(true) => ("\u{2713}", "configured"),
|
||||||
|
Ok(false) => ("\u{2717}", "missing"),
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!(
|
||||||
|
" Warning: failed to check secret `{}`: {}",
|
||||||
|
info.secret_name, e
|
||||||
|
);
|
||||||
|
("?", "unknown")
|
||||||
|
}
|
||||||
|
},
|
||||||
|
None => ("?", "unknown"),
|
||||||
|
};
|
||||||
|
let mut parts = info.secret_name.clone();
|
||||||
|
if let Some(ref desc) = info.description {
|
||||||
|
parts = format!("{} ({})", parts, desc);
|
||||||
|
}
|
||||||
|
if let Some(ref loc) = info.location {
|
||||||
|
parts = format!("{} -> {}", parts, loc);
|
||||||
|
}
|
||||||
|
println!(" {} {} {}", parts, icon, label);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Validate a tool name to prevent path traversal.
|
/// Validate a tool name to prevent path traversal.
|
||||||
@@ -677,8 +838,7 @@ async fn combine_provider_scopes(
|
|||||||
secret_name: &str,
|
secret_name: &str,
|
||||||
base_oauth: &crate::tools::wasm::OAuthConfigSchema,
|
base_oauth: &crate::tools::wasm::OAuthConfigSchema,
|
||||||
) -> crate::tools::wasm::OAuthConfigSchema {
|
) -> crate::tools::wasm::OAuthConfigSchema {
|
||||||
let mut all_scopes: std::collections::HashSet<String> =
|
let mut all_scopes: HashSet<String> = base_oauth.scopes.iter().cloned().collect();
|
||||||
base_oauth.scopes.iter().cloned().collect();
|
|
||||||
|
|
||||||
if let Ok(mut entries) = tokio::fs::read_dir(tools_dir).await {
|
if let Ok(mut entries) = tokio::fs::read_dir(tools_dir).await {
|
||||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||||
@@ -1127,6 +1287,8 @@ async fn setup_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyh
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::secrets::{CreateSecretParams, SecretsStore};
|
||||||
|
use crate::testing::credentials::test_secrets_store;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_format_size() {
|
fn test_format_size() {
|
||||||
@@ -1143,4 +1305,96 @@ mod tests {
|
|||||||
assert!(dir.to_string_lossy().contains(".ironclaw"));
|
assert!(dir.to_string_lossy().contains(".ironclaw"));
|
||||||
assert!(dir.to_string_lossy().contains("tools"));
|
assert!(dir.to_string_lossy().contains("tools"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Verify that auth secrets are deduplicated across auth, setup, and http.credentials,
|
||||||
|
/// and that credential status is checked against the secrets store.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_auth_secret_dedup_and_status() {
|
||||||
|
let caps = CapabilitiesFile::from_json(
|
||||||
|
r#"{
|
||||||
|
"auth": {
|
||||||
|
"secret_name": "gh_token",
|
||||||
|
"display_name": "GitHub"
|
||||||
|
},
|
||||||
|
"setup": {
|
||||||
|
"required_secrets": [
|
||||||
|
{ "name": "gh_token", "prompt": "GitHub PAT" },
|
||||||
|
{ "name": "extra_key", "prompt": "Extra API Key" }
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"http": {
|
||||||
|
"allowlist": [{ "host": "api.github.com" }],
|
||||||
|
"credentials": {
|
||||||
|
"github": {
|
||||||
|
"secret_name": "gh_token",
|
||||||
|
"location": { "type": "bearer" },
|
||||||
|
"host_patterns": ["api.github.com"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"secrets": {
|
||||||
|
"allowed_names": ["gh_token", "gh_*"]
|
||||||
|
}
|
||||||
|
}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let collected = collect_auth_secrets(&caps);
|
||||||
|
|
||||||
|
// gh_token should appear once (from auth), with location merged from credentials.
|
||||||
|
// extra_key should appear once (from setup).
|
||||||
|
assert_eq!(collected.secrets.len(), 2);
|
||||||
|
let gh = collected
|
||||||
|
.secrets
|
||||||
|
.iter()
|
||||||
|
.find(|s| s.secret_name == "gh_token")
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(gh.description.as_deref(), Some("GitHub"));
|
||||||
|
assert!(
|
||||||
|
gh.location.is_some(),
|
||||||
|
"location should be merged from http.credentials"
|
||||||
|
);
|
||||||
|
|
||||||
|
let extra = collected
|
||||||
|
.secrets
|
||||||
|
.iter()
|
||||||
|
.find(|s| s.secret_name == "extra_key")
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(extra.description.as_deref(), Some("Extra API Key"));
|
||||||
|
assert!(extra.location.is_none());
|
||||||
|
|
||||||
|
// Secrets section should filter gh_token (in seen_names) but keep gh_* (wildcard).
|
||||||
|
let secrets = caps.secrets.as_ref().unwrap();
|
||||||
|
let extra_secrets: Vec<_> = secrets
|
||||||
|
.allowed_names
|
||||||
|
.iter()
|
||||||
|
.filter(|name| !collected.seen_names.contains(name.as_str()))
|
||||||
|
.collect();
|
||||||
|
assert_eq!(extra_secrets, vec!["gh_*"]);
|
||||||
|
|
||||||
|
// Verify store check: missing secret -> exists returns false.
|
||||||
|
let store = test_secrets_store();
|
||||||
|
assert!(!store.exists("default", "gh_token").await.unwrap());
|
||||||
|
|
||||||
|
// Store gh_token and verify it's found.
|
||||||
|
store
|
||||||
|
.create(
|
||||||
|
"default",
|
||||||
|
CreateSecretParams::new("gh_token", "ghp_test123"),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert!(store.exists("default", "gh_token").await.unwrap());
|
||||||
|
// extra_key still missing.
|
||||||
|
assert!(!store.exists("default", "extra_key").await.unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// No auth sections → collect_auth_secrets returns empty.
|
||||||
|
#[test]
|
||||||
|
fn test_collect_auth_secrets_empty_caps() {
|
||||||
|
let caps = CapabilitiesFile::default();
|
||||||
|
let collected = collect_auth_secrets(&caps);
|
||||||
|
assert!(collected.secrets.is_empty());
|
||||||
|
assert!(collected.seen_names.is_empty());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use std::collections::HashMap;
|
|||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
use secrecy::SecretString;
|
use secrecy::SecretString;
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
use crate::bootstrap::ironclaw_base_dir;
|
use crate::bootstrap::ironclaw_base_dir;
|
||||||
use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env};
|
use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env};
|
||||||
@@ -45,6 +46,26 @@ pub struct GatewayConfig {
|
|||||||
/// Bearer token for authentication. Random hex generated at startup if unset.
|
/// Bearer token for authentication. Random hex generated at startup if unset.
|
||||||
pub auth_token: Option<String>,
|
pub auth_token: Option<String>,
|
||||||
pub user_id: String,
|
pub user_id: String,
|
||||||
|
/// Additional user scopes for workspace reads.
|
||||||
|
///
|
||||||
|
/// When set, the workspace will be able to read (search, read, list) from
|
||||||
|
/// these additional user scopes while writes remain isolated to `user_id`.
|
||||||
|
/// Parsed from `WORKSPACE_READ_SCOPES` (comma-separated).
|
||||||
|
pub workspace_read_scopes: Vec<String>,
|
||||||
|
/// Memory layer definitions (JSON in env var, or from external config).
|
||||||
|
pub memory_layers: Vec<crate::workspace::layer::MemoryLayer>,
|
||||||
|
/// Multi-user token map. When set, each token maps to a user identity.
|
||||||
|
/// Parsed from `GATEWAY_USER_TOKENS` (JSON string). When absent, falls back
|
||||||
|
/// to single-user mode via `auth_token` + `user_id`.
|
||||||
|
pub user_tokens: Option<HashMap<String, UserTokenConfig>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-user token configuration for multi-user mode.
|
||||||
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
|
pub struct UserTokenConfig {
|
||||||
|
pub user_id: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub workspace_read_scopes: Vec<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Signal channel configuration (signal-cli daemon HTTP/JSON-RPC).
|
/// Signal channel configuration (signal-cli daemon HTTP/JSON-RPC).
|
||||||
@@ -115,6 +136,118 @@ impl ChannelsConfig {
|
|||||||
.or_else(|| cs.gateway_user_id.clone())
|
.or_else(|| cs.gateway_user_id.clone())
|
||||||
.unwrap_or_else(|| owner_id.to_string());
|
.unwrap_or_else(|| owner_id.to_string());
|
||||||
|
|
||||||
|
let memory_layers: Vec<crate::workspace::layer::MemoryLayer> =
|
||||||
|
match optional_env("MEMORY_LAYERS")? {
|
||||||
|
Some(json_str) => {
|
||||||
|
serde_json::from_str(&json_str).map_err(|e| ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: format!("must be valid JSON array of layer objects: {e}"),
|
||||||
|
})?
|
||||||
|
}
|
||||||
|
None => crate::workspace::layer::MemoryLayer::default_for_user(&user_id),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Validate layer names and scopes
|
||||||
|
for layer in &memory_layers {
|
||||||
|
if layer.name.trim().is_empty() {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: "layer name must not be empty".to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if layer.name.len() > 64 {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: format!("layer name '{}' exceeds 64 characters", layer.name),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if !layer
|
||||||
|
.name
|
||||||
|
.chars()
|
||||||
|
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
||||||
|
{
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: format!(
|
||||||
|
"layer name '{}' contains invalid characters \
|
||||||
|
(allowed: a-z, A-Z, 0-9, _, -)",
|
||||||
|
layer.name
|
||||||
|
),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if layer.scope.trim().is_empty() {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: format!("layer '{}' has an empty scope", layer.name),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for duplicate layer names
|
||||||
|
{
|
||||||
|
let mut seen = std::collections::HashSet::new();
|
||||||
|
for layer in &memory_layers {
|
||||||
|
if !seen.insert(&layer.name) {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "MEMORY_LAYERS".to_string(),
|
||||||
|
message: format!("duplicate layer name '{}'", layer.name),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let user_tokens: Option<HashMap<String, UserTokenConfig>> =
|
||||||
|
match optional_env("GATEWAY_USER_TOKENS")? {
|
||||||
|
Some(json_str) => {
|
||||||
|
let tokens: HashMap<String, UserTokenConfig> = serde_json::from_str(
|
||||||
|
&json_str,
|
||||||
|
)
|
||||||
|
.map_err(|e| ConfigError::InvalidValue {
|
||||||
|
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||||
|
message: format!(
|
||||||
|
"must be valid JSON object mapping tokens to user configs: {e}"
|
||||||
|
),
|
||||||
|
})?;
|
||||||
|
if tokens.is_empty() {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||||
|
message:
|
||||||
|
"token map is empty — remove the variable to use single-user mode"
|
||||||
|
.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
for (tok, cfg) in &tokens {
|
||||||
|
if cfg.user_id.trim().is_empty() {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||||
|
message: format!(
|
||||||
|
"token '{}...' has an empty user_id",
|
||||||
|
&tok[..tok.len().min(8)]
|
||||||
|
),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(tokens)
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
let workspace_read_scopes: Vec<String> = optional_env("WORKSPACE_READ_SCOPES")?
|
||||||
|
.map(|s| {
|
||||||
|
s.split(',')
|
||||||
|
.map(|s| s.trim().to_string())
|
||||||
|
.filter(|s| !s.is_empty())
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
for scope in &workspace_read_scopes {
|
||||||
|
if scope.len() > 128 {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "WORKSPACE_READ_SCOPES".to_string(),
|
||||||
|
message: format!("scope '{}...' exceeds 128 characters", &scope[..32]),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
Some(GatewayConfig {
|
Some(GatewayConfig {
|
||||||
host: optional_env("GATEWAY_HOST")?
|
host: optional_env("GATEWAY_HOST")?
|
||||||
.or_else(|| cs.gateway_host.clone())
|
.or_else(|| cs.gateway_host.clone())
|
||||||
@@ -126,6 +259,9 @@ impl ChannelsConfig {
|
|||||||
auth_token: optional_env("GATEWAY_AUTH_TOKEN")?
|
auth_token: optional_env("GATEWAY_AUTH_TOKEN")?
|
||||||
.or_else(|| cs.gateway_auth_token.clone()),
|
.or_else(|| cs.gateway_auth_token.clone()),
|
||||||
user_id,
|
user_id,
|
||||||
|
workspace_read_scopes,
|
||||||
|
memory_layers,
|
||||||
|
user_tokens,
|
||||||
})
|
})
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -281,6 +417,9 @@ mod tests {
|
|||||||
port: 3000,
|
port: 3000,
|
||||||
auth_token: Some("tok-abc".to_string()),
|
auth_token: Some("tok-abc".to_string()),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
memory_layers: vec![],
|
||||||
|
user_tokens: None,
|
||||||
};
|
};
|
||||||
assert_eq!(cfg.host, "127.0.0.1");
|
assert_eq!(cfg.host, "127.0.0.1");
|
||||||
assert_eq!(cfg.port, 3000);
|
assert_eq!(cfg.port, 3000);
|
||||||
@@ -295,6 +434,9 @@ mod tests {
|
|||||||
port: 3001,
|
port: 3001,
|
||||||
auth_token: None,
|
auth_token: None,
|
||||||
user_id: "anon".to_string(),
|
user_id: "anon".to_string(),
|
||||||
|
workspace_read_scopes: vec![],
|
||||||
|
memory_layers: vec![],
|
||||||
|
user_tokens: None,
|
||||||
};
|
};
|
||||||
assert!(cfg.auth_token.is_none());
|
assert!(cfg.auth_token.is_none());
|
||||||
}
|
}
|
||||||
|
|||||||
+12
-4
@@ -406,7 +406,7 @@ impl LlmConfig {
|
|||||||
// Resolve extra headers
|
// Resolve extra headers
|
||||||
let extra_headers = if let Some(env_var) = extra_headers_env {
|
let extra_headers = if let Some(env_var) = extra_headers_env {
|
||||||
optional_env(env_var)?
|
optional_env(env_var)?
|
||||||
.map(|val| parse_extra_headers(&val))
|
.map(|val| parse_extra_headers_with_key(&val, env_var))
|
||||||
.transpose()?
|
.transpose()?
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
} else {
|
} else {
|
||||||
@@ -475,7 +475,10 @@ impl LlmConfig {
|
|||||||
///
|
///
|
||||||
/// Format: `Key1:Value1,Key2:Value2` (colon-separated, not `=`, because
|
/// Format: `Key1:Value1,Key2:Value2` (colon-separated, not `=`, because
|
||||||
/// header values often contain `=`).
|
/// header values often contain `=`).
|
||||||
fn parse_extra_headers(val: &str) -> Result<Vec<(String, String)>, ConfigError> {
|
fn parse_extra_headers_with_key(
|
||||||
|
val: &str,
|
||||||
|
env_var_name: &str,
|
||||||
|
) -> Result<Vec<(String, String)>, ConfigError> {
|
||||||
if val.trim().is_empty() {
|
if val.trim().is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
@@ -488,14 +491,14 @@ fn parse_extra_headers(val: &str) -> Result<Vec<(String, String)>, ConfigError>
|
|||||||
}
|
}
|
||||||
let Some((key, value)) = pair.split_once(':') else {
|
let Some((key, value)) = pair.split_once(':') else {
|
||||||
return Err(ConfigError::InvalidValue {
|
return Err(ConfigError::InvalidValue {
|
||||||
key: "LLM_EXTRA_HEADERS".to_string(),
|
key: env_var_name.to_string(),
|
||||||
message: format!("malformed header entry '{}', expected Key:Value", pair),
|
message: format!("malformed header entry '{}', expected Key:Value", pair),
|
||||||
});
|
});
|
||||||
};
|
};
|
||||||
let key = key.trim();
|
let key = key.trim();
|
||||||
if key.is_empty() {
|
if key.is_empty() {
|
||||||
return Err(ConfigError::InvalidValue {
|
return Err(ConfigError::InvalidValue {
|
||||||
key: "LLM_EXTRA_HEADERS".to_string(),
|
key: env_var_name.to_string(),
|
||||||
message: format!("empty header name in entry '{}'", pair),
|
message: format!("empty header name in entry '{}'", pair),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -536,6 +539,11 @@ mod tests {
|
|||||||
use crate::settings::Settings;
|
use crate::settings::Settings;
|
||||||
use crate::testing::credentials::*;
|
use crate::testing::credentials::*;
|
||||||
|
|
||||||
|
/// Convenience wrapper for tests — uses "TEST_HEADERS" as the env var name.
|
||||||
|
fn parse_extra_headers(val: &str) -> Result<Vec<(String, String)>, ConfigError> {
|
||||||
|
parse_extra_headers_with_key(val, "TEST_HEADERS")
|
||||||
|
}
|
||||||
|
|
||||||
/// Clear all openai-compatible-related env vars.
|
/// Clear all openai-compatible-related env vars.
|
||||||
fn clear_openai_compatible_env() {
|
fn clear_openai_compatible_env() {
|
||||||
// SAFETY: Only called under ENV_MUTEX in tests.
|
// SAFETY: Only called under ENV_MUTEX in tests.
|
||||||
|
|||||||
+5
-7
@@ -312,13 +312,11 @@ impl Config {
|
|||||||
let tunnel = TunnelConfig::resolve(settings)?;
|
let tunnel = TunnelConfig::resolve(settings)?;
|
||||||
let channels = ChannelsConfig::resolve(settings, &owner_id)?;
|
let channels = ChannelsConfig::resolve(settings, &owner_id)?;
|
||||||
|
|
||||||
// Resolve workspace config using the gateway user_id for default layers.
|
// Resolve the startup workspace against the durable owner scope. The
|
||||||
let workspace_user_id = channels
|
// gateway may expose a distinct sender identity, but the base runtime
|
||||||
.gateway
|
// workspace stays owner-scoped and per-user gateway workspaces are
|
||||||
.as_ref()
|
// handled separately by WorkspacePool.
|
||||||
.map(|gw| gw.user_id.as_str())
|
let workspace = WorkspaceConfig::resolve(&owner_id)?;
|
||||||
.unwrap_or("default");
|
|
||||||
let workspace = WorkspaceConfig::resolve(workspace_user_id)?;
|
|
||||||
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
owner_id: owner_id.clone(),
|
owner_id: owner_id.clone(),
|
||||||
|
|||||||
@@ -230,6 +230,49 @@ impl JobStore for LibSqlBackend {
|
|||||||
Ok(jobs)
|
Ok(jobs)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_agent_jobs_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||||
|
let conn = self.connect().await?;
|
||||||
|
let mut rows = conn
|
||||||
|
.query(
|
||||||
|
r#"
|
||||||
|
SELECT id, title, status, user_id, failure_reason,
|
||||||
|
created_at, started_at, completed_at
|
||||||
|
FROM agent_jobs WHERE source = 'direct' AND user_id = ?1
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
"#,
|
||||||
|
params![user_id],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||||
|
|
||||||
|
let mut jobs = Vec::new();
|
||||||
|
while let Some(row) = rows
|
||||||
|
.next()
|
||||||
|
.await
|
||||||
|
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||||
|
{
|
||||||
|
let id_str = get_text(&row, 0);
|
||||||
|
let Ok(id) = id_str.parse() else {
|
||||||
|
tracing::warn!("Skipping agent job with invalid UUID: {}", id_str);
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
jobs.push(AgentJobRecord {
|
||||||
|
id,
|
||||||
|
title: get_text(&row, 1),
|
||||||
|
status: get_text(&row, 2),
|
||||||
|
user_id: get_text(&row, 3),
|
||||||
|
failure_reason: get_opt_text(&row, 4),
|
||||||
|
created_at: get_ts(&row, 5),
|
||||||
|
started_at: get_opt_ts(&row, 6),
|
||||||
|
completed_at: get_opt_ts(&row, 7),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(jobs)
|
||||||
|
}
|
||||||
|
|
||||||
async fn get_agent_job_failure_reason(
|
async fn get_agent_job_failure_reason(
|
||||||
&self,
|
&self,
|
||||||
id: Uuid,
|
id: Uuid,
|
||||||
@@ -277,6 +320,32 @@ impl JobStore for LibSqlBackend {
|
|||||||
Ok(summary)
|
Ok(summary)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn agent_job_summary_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<AgentJobSummary, DatabaseError> {
|
||||||
|
let conn = self.connect().await?;
|
||||||
|
let mut rows = conn
|
||||||
|
.query(
|
||||||
|
"SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'direct' AND user_id = ?1 GROUP BY status",
|
||||||
|
params![user_id],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||||
|
|
||||||
|
let mut summary = AgentJobSummary::default();
|
||||||
|
while let Some(row) = rows
|
||||||
|
.next()
|
||||||
|
.await
|
||||||
|
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||||
|
{
|
||||||
|
let status = get_text(&row, 0);
|
||||||
|
let count = get_i64(&row, 1) as usize;
|
||||||
|
summary.add_count(&status, count);
|
||||||
|
}
|
||||||
|
Ok(summary)
|
||||||
|
}
|
||||||
|
|
||||||
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> {
|
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> {
|
||||||
let conn = self.connect().await?;
|
let conn = self.connect().await?;
|
||||||
let duration_ms = action.duration.as_millis() as i64;
|
let duration_ms = action.duration.as_millis() as i64;
|
||||||
|
|||||||
@@ -462,6 +462,56 @@ impl RoutineStore for LibSqlBackend {
|
|||||||
Ok(counts)
|
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(
|
async fn link_routine_run_to_job(
|
||||||
&self,
|
&self,
|
||||||
run_id: Uuid,
|
run_id: Uuid,
|
||||||
|
|||||||
@@ -409,7 +409,15 @@ pub trait JobStore: Send + Sync {
|
|||||||
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError>;
|
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError>;
|
||||||
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError>;
|
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError>;
|
||||||
async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError>;
|
async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError>;
|
||||||
|
async fn list_agent_jobs_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Vec<AgentJobRecord>, DatabaseError>;
|
||||||
async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError>;
|
async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError>;
|
||||||
|
async fn agent_job_summary_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<AgentJobSummary, DatabaseError>;
|
||||||
/// Get the failure reason for a single agent job (O(1) lookup).
|
/// Get the failure reason for a single agent job (O(1) lookup).
|
||||||
async fn get_agent_job_failure_reason(&self, id: Uuid)
|
async fn get_agent_job_failure_reason(&self, id: Uuid)
|
||||||
-> Result<Option<String>, DatabaseError>;
|
-> Result<Option<String>, DatabaseError>;
|
||||||
@@ -520,6 +528,15 @@ pub trait RoutineStore: Send + Sync {
|
|||||||
&self,
|
&self,
|
||||||
routine_ids: &[Uuid],
|
routine_ids: &[Uuid],
|
||||||
) -> Result<HashMap<Uuid, i64>, DatabaseError>;
|
) -> 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(
|
async fn link_routine_run_to_job(
|
||||||
&self,
|
&self,
|
||||||
run_id: Uuid,
|
run_id: Uuid,
|
||||||
|
|||||||
@@ -249,10 +249,24 @@ impl JobStore for PgBackend {
|
|||||||
self.store.list_agent_jobs().await
|
self.store.list_agent_jobs().await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_agent_jobs_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||||
|
self.store.list_agent_jobs_for_user(user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError> {
|
async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError> {
|
||||||
self.store.agent_job_summary().await
|
self.store.agent_job_summary().await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn agent_job_summary_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<AgentJobSummary, DatabaseError> {
|
||||||
|
self.store.agent_job_summary_for_user(user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn get_agent_job_failure_reason(
|
async fn get_agent_job_failure_reason(
|
||||||
&self,
|
&self,
|
||||||
id: Uuid,
|
id: Uuid,
|
||||||
@@ -496,6 +510,14 @@ impl RoutineStore for PgBackend {
|
|||||||
.await
|
.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(
|
async fn link_routine_run_to_job(
|
||||||
&self,
|
&self,
|
||||||
run_id: Uuid,
|
run_id: Uuid,
|
||||||
|
|||||||
+361
-283
File diff suppressed because it is too large
Load Diff
@@ -842,6 +842,38 @@ impl Store {
|
|||||||
.collect())
|
.collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn list_agent_jobs_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||||
|
let conn = self.conn().await?;
|
||||||
|
let rows = conn
|
||||||
|
.query(
|
||||||
|
r#"
|
||||||
|
SELECT id, title, status, user_id, failure_reason,
|
||||||
|
created_at, started_at, completed_at
|
||||||
|
FROM agent_jobs WHERE source = 'direct' AND user_id = $1
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
"#,
|
||||||
|
&[&user_id],
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(rows
|
||||||
|
.iter()
|
||||||
|
.map(|r| AgentJobRecord {
|
||||||
|
id: r.get("id"),
|
||||||
|
title: r.get("title"),
|
||||||
|
status: r.get("status"),
|
||||||
|
user_id: r.get::<_, Option<String>>("user_id").unwrap_or_default(),
|
||||||
|
created_at: r.get("created_at"),
|
||||||
|
started_at: r.get("started_at"),
|
||||||
|
completed_at: r.get("completed_at"),
|
||||||
|
failure_reason: r.get("failure_reason"),
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
/// Get the failure reason for a single agent job.
|
/// Get the failure reason for a single agent job.
|
||||||
pub async fn get_agent_job_failure_reason(
|
pub async fn get_agent_job_failure_reason(
|
||||||
&self,
|
&self,
|
||||||
@@ -875,6 +907,27 @@ impl Store {
|
|||||||
}
|
}
|
||||||
Ok(summary)
|
Ok(summary)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn agent_job_summary_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<AgentJobSummary, DatabaseError> {
|
||||||
|
let conn = self.conn().await?;
|
||||||
|
let rows = conn
|
||||||
|
.query(
|
||||||
|
"SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'direct' AND user_id = $1 GROUP BY status",
|
||||||
|
&[&user_id],
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let mut summary = AgentJobSummary::default();
|
||||||
|
for row in &rows {
|
||||||
|
let status: String = row.get("status");
|
||||||
|
let count: i64 = row.get("cnt");
|
||||||
|
summary.add_count(&status, count as usize);
|
||||||
|
}
|
||||||
|
Ok(summary)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ==================== Job Events ====================
|
// ==================== Job Events ====================
|
||||||
@@ -1350,6 +1403,40 @@ impl Store {
|
|||||||
Ok(counts)
|
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.
|
/// Link a routine run to a dispatched job.
|
||||||
pub async fn link_routine_run_to_job(
|
pub async fn link_routine_run_to_job(
|
||||||
&self,
|
&self,
|
||||||
|
|||||||
@@ -575,6 +575,7 @@ fn extract_response_content(response: &AnthropicResponse) -> (Option<String>, Ve
|
|||||||
id: id.clone(),
|
id: id.clone(),
|
||||||
name: name.clone(),
|
name: name.clone(),
|
||||||
arguments: input.clone(),
|
arguments: input.clone(),
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -623,6 +624,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "test"}),
|
arguments: serde_json::json!({"q": "test"}),
|
||||||
|
reasoning: None,
|
||||||
}];
|
}];
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
ChatMessage::user("Search for test"),
|
ChatMessage::user("Search for test"),
|
||||||
|
|||||||
@@ -522,6 +522,7 @@ fn extract_content_blocks(
|
|||||||
id: tu.tool_use_id().to_string(),
|
id: tu.tool_use_id().to_string(),
|
||||||
name: tu.name().to_string(),
|
name: tu.name().to_string(),
|
||||||
arguments: document_to_json(tu.input()),
|
arguments: document_to_json(tu.input()),
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// Ignore reasoning, citations, images, etc.
|
// Ignore reasoning, citations, images, etc.
|
||||||
@@ -759,11 +760,13 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({"text": "hi"}),
|
arguments: serde_json::json!({"text": "hi"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let tc2 = crate::llm::provider::ToolCall {
|
let tc2 = crate::llm::provider::ToolCall {
|
||||||
id: "call_2".to_string(),
|
id: "call_2".to_string(),
|
||||||
name: "time".to_string(),
|
name: "time".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
@@ -802,6 +805,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
@@ -825,6 +829,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
@@ -989,11 +994,13 @@ mod tests {
|
|||||||
id: "call_abc".to_string(),
|
id: "call_abc".to_string(),
|
||||||
name: "get_weather".to_string(),
|
name: "get_weather".to_string(),
|
||||||
arguments: serde_json::json!({"city": "NYC"}),
|
arguments: serde_json::json!({"city": "NYC"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let tc2 = crate::llm::provider::ToolCall {
|
let tc2 = crate::llm::provider::ToolCall {
|
||||||
id: "call_def".to_string(),
|
id: "call_def".to_string(),
|
||||||
name: "get_time".to_string(),
|
name: "get_time".to_string(),
|
||||||
arguments: serde_json::json!({"tz": "EST"}),
|
arguments: serde_json::json!({"tz": "EST"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
|
|||||||
@@ -732,6 +732,7 @@ impl LlmProvider for CodexChatGptProvider {
|
|||||||
id: tc.call_id,
|
id: tc.call_id,
|
||||||
name: tc.name,
|
name: tc.name,
|
||||||
arguments: args,
|
arguments: args,
|
||||||
|
reasoning: None,
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -825,6 +826,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: json!({"query": "rust"}),
|
arguments: json!({"query": "rust"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking...".into()), vec![tc]);
|
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking...".into()), vec![tc]);
|
||||||
let items = CodexChatGptProvider::message_to_input_items(&msg);
|
let items = CodexChatGptProvider::message_to_input_items(&msg);
|
||||||
|
|||||||
@@ -1898,6 +1898,7 @@ impl GeminiOauthProvider {
|
|||||||
id,
|
id,
|
||||||
name,
|
name,
|
||||||
arguments: args,
|
arguments: args,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+21
-52
@@ -107,14 +107,21 @@ impl GithubCopilotProvider {
|
|||||||
body: &impl Serialize,
|
body: &impl Serialize,
|
||||||
) -> Result<R, LlmError> {
|
) -> Result<R, LlmError> {
|
||||||
let url = self.api_url();
|
let url = self.api_url();
|
||||||
// Map token exchange failures to RequestFailed (retryable) rather than
|
// Distinguish permanent auth errors (non-retryable) from transient
|
||||||
// AuthFailed (non-retryable), since transient network errors during
|
// network failures (retryable) so RetryProvider handles them correctly.
|
||||||
// exchange should be retried by RetryProvider.
|
|
||||||
let token = self.token_manager.get_token().await.map_err(|e| {
|
let token = self.token_manager.get_token().await.map_err(|e| {
|
||||||
tracing::warn!(error = %e, "Copilot: token exchange failed");
|
tracing::warn!(error = %e, "Copilot: token exchange failed");
|
||||||
LlmError::RequestFailed {
|
match &e {
|
||||||
provider: "github_copilot".to_string(),
|
crate::llm::github_copilot_auth::GithubCopilotAuthError::AccessDenied
|
||||||
reason: format!("Token exchange failed: {e}"),
|
| crate::llm::github_copilot_auth::GithubCopilotAuthError::Expired => {
|
||||||
|
LlmError::AuthFailed {
|
||||||
|
provider: "github_copilot".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => LlmError::RequestFailed {
|
||||||
|
provider: "github_copilot".to_string(),
|
||||||
|
reason: format!("Token exchange failed: {e}"),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
@@ -157,54 +164,14 @@ impl GithubCopilotProvider {
|
|||||||
);
|
);
|
||||||
|
|
||||||
if status.as_u16() == 401 {
|
if status.as_u16() == 401 {
|
||||||
// Invalidate the cached session token and retry once with a
|
// Invalidate the cached session token so the next attempt
|
||||||
// fresh exchange — stale tokens are the most common 401 cause.
|
// (driven by RetryProvider) gets a fresh one. We don't retry
|
||||||
tracing::warn!("Copilot: 401 Unauthorized — invalidating session token, retrying");
|
// inline to avoid nested retries with the outer RetryProvider.
|
||||||
|
tracing::warn!("Copilot: 401 Unauthorized — invalidating session token for retry");
|
||||||
self.token_manager.invalidate().await;
|
self.token_manager.invalidate().await;
|
||||||
let fresh = self.token_manager.get_token().await.map_err(|e| {
|
return Err(LlmError::RequestFailed {
|
||||||
tracing::warn!(error = %e, "Copilot: re-exchange after 401 failed");
|
|
||||||
LlmError::RequestFailed {
|
|
||||||
provider: "github_copilot".to_string(),
|
|
||||||
reason: format!("Token re-exchange after 401 failed: {e}"),
|
|
||||||
}
|
|
||||||
})?;
|
|
||||||
let mut retry_req = self
|
|
||||||
.client
|
|
||||||
.post(&url)
|
|
||||||
.bearer_auth(fresh.expose_secret())
|
|
||||||
.header("Content-Type", "application/json");
|
|
||||||
for (key, value) in &self.extra_headers {
|
|
||||||
retry_req = retry_req.header(key.as_str(), value.as_str());
|
|
||||||
}
|
|
||||||
let retry =
|
|
||||||
retry_req
|
|
||||||
.json(body)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map_err(|e| LlmError::RequestFailed {
|
|
||||||
provider: "github_copilot".to_string(),
|
|
||||||
reason: format!("Retry after 401 failed: {e}"),
|
|
||||||
})?;
|
|
||||||
if retry.status().is_success() {
|
|
||||||
let text = retry.text().await.map_err(|e| LlmError::RequestFailed {
|
|
||||||
provider: "github_copilot".to_string(),
|
|
||||||
reason: format!("Failed to read retry response body: {e}"),
|
|
||||||
})?;
|
|
||||||
return serde_json::from_str(&text).map_err(|e| {
|
|
||||||
let truncated = crate::agent::truncate_for_preview(&text, 512);
|
|
||||||
LlmError::InvalidResponse {
|
|
||||||
provider: "github_copilot".to_string(),
|
|
||||||
reason: format!("JSON parse error: {e}. Raw: {truncated}"),
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
let retry_status = retry.status();
|
|
||||||
tracing::warn!(
|
|
||||||
status = %retry_status,
|
|
||||||
"Copilot: 401 retry also failed"
|
|
||||||
);
|
|
||||||
return Err(LlmError::AuthFailed {
|
|
||||||
provider: "github_copilot".to_string(),
|
provider: "github_copilot".to_string(),
|
||||||
|
reason: "HTTP 401 Unauthorized".to_string(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
if status.as_u16() == 429 {
|
if status.as_u16() == 429 {
|
||||||
@@ -629,6 +596,7 @@ fn extract_choice_content(choice: &OpenAiChoice) -> (Option<String>, Vec<ToolCal
|
|||||||
name: tc.function.name.clone(),
|
name: tc.function.name.clone(),
|
||||||
arguments: serde_json::from_str(&tc.function.arguments)
|
arguments: serde_json::from_str(&tc.function.arguments)
|
||||||
.unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
|
||||||
|
reasoning: None,
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
})
|
})
|
||||||
@@ -661,6 +629,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "test"}),
|
arguments: serde_json::json!({"q": "test"}),
|
||||||
|
reasoning: None,
|
||||||
}];
|
}];
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
ChatMessage::user("Search"),
|
ChatMessage::user("Search"),
|
||||||
|
|||||||
+75
-2
@@ -463,8 +463,15 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||||
let mut raw_messages = req.messages;
|
let mut raw_messages = req.messages;
|
||||||
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
||||||
let messages: Vec<ChatCompletionMessage> =
|
let raw: Vec<ChatCompletionMessage> = raw_messages.into_iter().map(|m| m.into()).collect();
|
||||||
raw_messages.into_iter().map(|m| m.into()).collect();
|
|
||||||
|
// NEAR AI rejects `role:"tool"` messages even on text-only completion paths.
|
||||||
|
// Apply the same flattening used by complete_with_tools().
|
||||||
|
let messages = if self.flatten_tool_messages {
|
||||||
|
flatten_tool_messages(raw)
|
||||||
|
} else {
|
||||||
|
raw
|
||||||
|
};
|
||||||
|
|
||||||
let request = ChatCompletionRequest {
|
let request = ChatCompletionRequest {
|
||||||
model,
|
model,
|
||||||
@@ -580,6 +587,7 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
id: tc.id,
|
id: tc.id,
|
||||||
name: tc.function.name,
|
name: tc.function.name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -1173,11 +1181,13 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "list_issues".to_string(),
|
name: "list_issues".to_string(),
|
||||||
arguments: serde_json::json!({"owner": "foo", "repo": "bar"}),
|
arguments: serde_json::json!({"owner": "foo", "repo": "bar"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
ToolCall {
|
ToolCall {
|
||||||
id: "call_2".to_string(),
|
id: "call_2".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
];
|
];
|
||||||
|
|
||||||
@@ -1210,6 +1220,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "test".to_string(),
|
name: "test".to_string(),
|
||||||
arguments: serde_json::json!({"key": "value"}),
|
arguments: serde_json::json!({"key": "value"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
||||||
let chat_msg: ChatCompletionMessage = msg.into();
|
let chat_msg: ChatCompletionMessage = msg.into();
|
||||||
@@ -1453,6 +1464,7 @@ mod tests {
|
|||||||
id: tc.id,
|
id: tc.id,
|
||||||
name: tc.function.name,
|
name: tc.function.name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -1502,6 +1514,7 @@ mod tests {
|
|||||||
id: tc.id,
|
id: tc.id,
|
||||||
name: tc.function.name,
|
name: tc.function.name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -2124,6 +2137,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "test".to_string(),
|
name: "test".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
}],
|
}],
|
||||||
);
|
);
|
||||||
let chat_msg: ChatCompletionMessage = msg.into();
|
let chat_msg: ChatCompletionMessage = msg.into();
|
||||||
@@ -2193,6 +2207,65 @@ mod tests {
|
|||||||
assert_eq!(deserialized.function.arguments, r#"{"city":"London"}"#);
|
assert_eq!(deserialized.function.arguments, r#"{"city":"London"}"#);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// -- flatten_tool_messages in complete() path ----------------------------
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_flatten_applied_on_text_only_path() {
|
||||||
|
// Verify that flatten_tool_messages converts tool-role messages to user
|
||||||
|
// messages (mirrors the complete_with_tools path).
|
||||||
|
let messages = vec![
|
||||||
|
ChatCompletionMessage {
|
||||||
|
role: "user".to_string(),
|
||||||
|
content: Some(MessageContent::Text("run it".to_string())),
|
||||||
|
tool_call_id: None,
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
ChatCompletionMessage {
|
||||||
|
role: "tool".to_string(),
|
||||||
|
content: Some(MessageContent::Text("ok".to_string())),
|
||||||
|
tool_call_id: Some("call_1".to_string()),
|
||||||
|
name: Some("run_cmd".to_string()),
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
let flattened = flatten_tool_messages(messages);
|
||||||
|
assert_eq!(flattened.len(), 2);
|
||||||
|
assert_eq!(flattened[1].role, "user");
|
||||||
|
let text = flattened[1]
|
||||||
|
.content
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|c| c.as_text())
|
||||||
|
.unwrap();
|
||||||
|
assert!(text.contains("run_cmd"), "should reference tool name");
|
||||||
|
assert!(text.contains("ok"), "should include tool result");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_no_flatten_when_no_tool_messages() {
|
||||||
|
// When there are no tool-role messages, flatten_tool_messages is a no-op.
|
||||||
|
let messages = vec![
|
||||||
|
ChatCompletionMessage {
|
||||||
|
role: "user".to_string(),
|
||||||
|
content: Some(MessageContent::Text("hi".to_string())),
|
||||||
|
tool_call_id: None,
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
ChatCompletionMessage {
|
||||||
|
role: "assistant".to_string(),
|
||||||
|
content: Some(MessageContent::Text("hello".to_string())),
|
||||||
|
tool_call_id: None,
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
let result = flatten_tool_messages(messages);
|
||||||
|
// No tool messages → unchanged roles
|
||||||
|
assert_eq!(result[0].role, "user");
|
||||||
|
assert_eq!(result[1].role, "assistant");
|
||||||
|
}
|
||||||
|
|
||||||
// -- api_url edge cases ---------------------------------------------------
|
// -- api_url edge cases ---------------------------------------------------
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -625,6 +625,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
|
|||||||
id: state.call_id,
|
id: state.call_id,
|
||||||
name: state.name,
|
name: state.name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
} else {
|
} else {
|
||||||
// Fallback: extract directly from the item
|
// Fallback: extract directly from the item
|
||||||
@@ -650,6 +651,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
|
|||||||
id: call_id,
|
id: call_id,
|
||||||
name,
|
name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -727,6 +729,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
|
|||||||
id: state.call_id,
|
id: state.call_id,
|
||||||
name: state.name,
|
name: state.name,
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -822,11 +825,13 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
ToolCall {
|
ToolCall {
|
||||||
id: "call_2".to_string(),
|
id: "call_2".to_string(),
|
||||||
name: "read".to_string(),
|
name: "read".to_string(),
|
||||||
arguments: serde_json::json!({"path": "/tmp"}),
|
arguments: serde_json::json!({"path": "/tmp"}),
|
||||||
|
reasoning: None,
|
||||||
},
|
},
|
||||||
];
|
];
|
||||||
let msg =
|
let msg =
|
||||||
|
|||||||
@@ -231,6 +231,10 @@ pub struct ToolCall {
|
|||||||
pub id: String,
|
pub id: String,
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub arguments: serde_json::Value,
|
pub arguments: serde_json::Value,
|
||||||
|
/// Optional reasoning for why this tool was chosen — supplied by the provider
|
||||||
|
/// or derived from the shared response content as a fallback.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub reasoning: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Generate a tool-call ID that satisfies all providers.
|
/// Generate a tool-call ID that satisfies all providers.
|
||||||
@@ -637,6 +641,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let mut messages = vec![
|
let mut messages = vec![
|
||||||
ChatMessage::user("hello"),
|
ChatMessage::user("hello"),
|
||||||
@@ -680,6 +685,7 @@ mod tests {
|
|||||||
id: "call_1".to_string(),
|
id: "call_1".to_string(),
|
||||||
name: "echo".to_string(),
|
name: "echo".to_string(),
|
||||||
arguments: serde_json::json!({}),
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let mut messages = vec![
|
let mut messages = vec![
|
||||||
ChatMessage::user("test"),
|
ChatMessage::user("test"),
|
||||||
@@ -705,11 +711,13 @@ mod tests {
|
|||||||
id: "call_sel_1".to_string(),
|
id: "call_sel_1".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "test"}),
|
arguments: serde_json::json!({"q": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let tc2 = ToolCall {
|
let tc2 = ToolCall {
|
||||||
id: "call_sel_2".to_string(),
|
id: "call_sel_2".to_string(),
|
||||||
name: "http".to_string(),
|
name: "http".to_string(),
|
||||||
arguments: serde_json::json!({"url": "https://example.com"}),
|
arguments: serde_json::json!({"url": "https://example.com"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let mut messages = vec![
|
let mut messages = vec![
|
||||||
ChatMessage::system("You are a helpful assistant."),
|
ChatMessage::system("You are a helpful assistant."),
|
||||||
|
|||||||
+85
-12
@@ -525,17 +525,35 @@ impl Reasoning {
|
|||||||
|
|
||||||
let response = self.llm.complete_with_tools(request).await?;
|
let response = self.llm.complete_with_tools(request).await?;
|
||||||
|
|
||||||
let reasoning = response.content.unwrap_or_default();
|
let shared_reasoning = response
|
||||||
|
.content
|
||||||
|
.map(|c| {
|
||||||
|
let pre_truncated = truncate_at_tool_tags(&c);
|
||||||
|
clean_response(&pre_truncated)
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
let selections: Vec<ToolSelection> = response
|
let selections: Vec<ToolSelection> = response
|
||||||
.tool_calls
|
.tool_calls
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|tool_call| ToolSelection {
|
.map(|tool_call| {
|
||||||
tool_name: tool_call.name,
|
// Prefer per-tool reasoning if the provider supplied it,
|
||||||
parameters: tool_call.arguments,
|
// otherwise fall back to the shared response content.
|
||||||
reasoning: reasoning.clone(),
|
let rationale = tool_call
|
||||||
alternatives: vec![],
|
.reasoning
|
||||||
tool_call_id: tool_call.id,
|
.map(|r| {
|
||||||
|
let pre_truncated = truncate_at_tool_tags(&r);
|
||||||
|
clean_response(&pre_truncated)
|
||||||
|
})
|
||||||
|
.filter(|r| !r.trim().is_empty())
|
||||||
|
.unwrap_or_else(|| shared_reasoning.clone());
|
||||||
|
ToolSelection {
|
||||||
|
tool_name: tool_call.name,
|
||||||
|
parameters: tool_call.arguments,
|
||||||
|
reasoning: rationale,
|
||||||
|
alternatives: vec![],
|
||||||
|
tool_call_id: tool_call.id,
|
||||||
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -664,13 +682,36 @@ Respond in JSON format:
|
|||||||
|
|
||||||
// If there were tool calls, return them for execution
|
// If there were tool calls, return them for execution
|
||||||
if !response.tool_calls.is_empty() {
|
if !response.tool_calls.is_empty() {
|
||||||
|
let narrative = response.content.map(|c| {
|
||||||
|
let pre_truncated = truncate_at_tool_tags(&c);
|
||||||
|
clean_response(&pre_truncated)
|
||||||
|
});
|
||||||
|
// Populate per-tool reasoning from the shared narrative when the
|
||||||
|
// provider did not supply per-tool rationale.
|
||||||
|
let tool_calls: Vec<ToolCall> = response
|
||||||
|
.tool_calls
|
||||||
|
.into_iter()
|
||||||
|
.map(|mut tc| {
|
||||||
|
if tc.reasoning.as_ref().is_none_or(|r| r.trim().is_empty()) {
|
||||||
|
tc.reasoning = narrative.as_ref().filter(|n| !n.is_empty()).cloned();
|
||||||
|
} else {
|
||||||
|
// Clean provider-supplied per-tool reasoning the same way
|
||||||
|
// we clean the shared narrative (strip thinking/tool tags).
|
||||||
|
tc.reasoning = tc
|
||||||
|
.reasoning
|
||||||
|
.map(|r| {
|
||||||
|
let pre_truncated = truncate_at_tool_tags(&r);
|
||||||
|
clean_response(&pre_truncated)
|
||||||
|
})
|
||||||
|
.filter(|r| !r.trim().is_empty());
|
||||||
|
}
|
||||||
|
tc
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
return Ok(RespondOutput {
|
return Ok(RespondOutput {
|
||||||
result: RespondResult::ToolCalls {
|
result: RespondResult::ToolCalls {
|
||||||
tool_calls: response.tool_calls,
|
tool_calls,
|
||||||
content: response.content.map(|c| {
|
content: narrative,
|
||||||
let pre_truncated = truncate_at_tool_tags(&c);
|
|
||||||
clean_response(&pre_truncated)
|
|
||||||
}),
|
|
||||||
},
|
},
|
||||||
usage,
|
usage,
|
||||||
});
|
});
|
||||||
@@ -1350,6 +1391,7 @@ fn recover_tool_calls_from_content(
|
|||||||
),
|
),
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -1364,6 +1406,7 @@ fn recover_tool_calls_from_content(
|
|||||||
),
|
),
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
arguments: serde_json::Value::Object(Default::default()),
|
arguments: serde_json::Value::Object(Default::default()),
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1401,6 +1444,7 @@ fn recover_tool_calls_from_content(
|
|||||||
),
|
),
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
remaining = &args_start[bracket_end + 1..];
|
remaining = &args_start[bracket_end + 1..];
|
||||||
continue;
|
continue;
|
||||||
@@ -1412,6 +1456,7 @@ fn recover_tool_calls_from_content(
|
|||||||
id: super::provider::generate_tool_call_id(calls.len(), RECOVERED_TOOL_CALL_SEED),
|
id: super::provider::generate_tool_call_id(calls.len(), RECOVERED_TOOL_CALL_SEED),
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
arguments: serde_json::Value::Object(Default::default()),
|
arguments: serde_json::Value::Object(Default::default()),
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
remaining = after_name;
|
remaining = after_name;
|
||||||
}
|
}
|
||||||
@@ -3145,4 +3190,32 @@ That's my plan."#;
|
|||||||
"Text <function_call>{}</function_call> middle "
|
"Text <function_call>{}</function_call> middle "
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Verify that reasoning normalization strips thinking tags and tool tags
|
||||||
|
/// from per-tool reasoning, matching the cleaning applied to shared reasoning.
|
||||||
|
#[test]
|
||||||
|
fn test_reasoning_normalization_strips_thinking_tags() {
|
||||||
|
let raw = "<thinking>Let me consider...</thinking>Search memory for prior context";
|
||||||
|
let pre_truncated = truncate_at_tool_tags(raw);
|
||||||
|
let cleaned = clean_response(&pre_truncated);
|
||||||
|
assert!(!cleaned.contains("<thinking>"));
|
||||||
|
assert!(cleaned.contains("Search memory"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_reasoning_normalization_strips_tool_tags() {
|
||||||
|
let raw = "Calling search <tool_call>{\"name\": \"search\"}";
|
||||||
|
let pre_truncated = truncate_at_tool_tags(raw);
|
||||||
|
let cleaned = clean_response(&pre_truncated);
|
||||||
|
assert!(!cleaned.contains("<tool_call>"));
|
||||||
|
assert!(cleaned.contains("Calling search"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_reasoning_normalization_empty_after_cleaning() {
|
||||||
|
let raw = "<thinking>internal only</thinking>";
|
||||||
|
let pre_truncated = truncate_at_tool_tags(raw);
|
||||||
|
let cleaned = clean_response(&pre_truncated);
|
||||||
|
assert!(cleaned.trim().is_empty());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -490,6 +490,7 @@ fn extract_response(
|
|||||||
id: tc.id.clone(),
|
id: tc.id.clone(),
|
||||||
name: tc.function.name.clone(),
|
name: tc.function.name.clone(),
|
||||||
arguments: tc.function.arguments.clone(),
|
arguments: tc.function.arguments.clone(),
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// Reasoning and Image variants are not mapped to IronClaw types
|
// Reasoning and Image variants are not mapped to IronClaw types
|
||||||
@@ -880,6 +881,7 @@ mod tests {
|
|||||||
id: "Xt7mK9pQ2".to_string(),
|
id: "Xt7mK9pQ2".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
|
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
|
||||||
let messages = vec![msg];
|
let messages = vec![msg];
|
||||||
@@ -997,6 +999,7 @@ mod tests {
|
|||||||
id: "".to_string(),
|
id: "".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
||||||
let (_preamble, history) = convert_messages(&messages);
|
let (_preamble, history) = convert_messages(&messages);
|
||||||
@@ -1028,6 +1031,7 @@ mod tests {
|
|||||||
id: " ".to_string(),
|
id: " ".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
||||||
let (_preamble, history) = convert_messages(&messages);
|
let (_preamble, history) = convert_messages(&messages);
|
||||||
@@ -1061,6 +1065,7 @@ mod tests {
|
|||||||
id: "".to_string(),
|
id: "".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"query": "test"}),
|
arguments: serde_json::json!({"query": "test"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
||||||
let tool_result_msg = ChatMessage {
|
let tool_result_msg = ChatMessage {
|
||||||
@@ -1380,11 +1385,13 @@ mod tests {
|
|||||||
id: "call_a".to_string(),
|
id: "call_a".to_string(),
|
||||||
name: "search".to_string(),
|
name: "search".to_string(),
|
||||||
arguments: serde_json::json!({"q": "rust"}),
|
arguments: serde_json::json!({"q": "rust"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let tc2 = IronToolCall {
|
let tc2 = IronToolCall {
|
||||||
id: "call_b".to_string(),
|
id: "call_b".to_string(),
|
||||||
name: "fetch".to_string(),
|
name: "fetch".to_string(),
|
||||||
arguments: serde_json::json!({"url": "https://example.com"}),
|
arguments: serde_json::json!({"url": "https://example.com"}),
|
||||||
|
reasoning: None,
|
||||||
};
|
};
|
||||||
let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]);
|
let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]);
|
||||||
let result_a = ChatMessage::tool_result("call_a", "search", "search results");
|
let result_a = ChatMessage::tool_result("call_a", "search", "search results");
|
||||||
|
|||||||
+67
-15
@@ -142,6 +142,11 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
init_cli_tracing();
|
init_cli_tracing();
|
||||||
return ironclaw::cli::run_logs_command(logs_cmd.clone(), cli.config.as_deref()).await;
|
return ironclaw::cli::run_logs_command(logs_cmd.clone(), cli.config.as_deref()).await;
|
||||||
}
|
}
|
||||||
|
Some(Command::Models(models_cmd)) => {
|
||||||
|
init_cli_tracing();
|
||||||
|
return ironclaw::cli::run_models_command(models_cmd.clone(), cli.config.as_deref())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
Some(Command::Doctor) => {
|
Some(Command::Doctor) => {
|
||||||
init_cli_tracing();
|
init_cli_tracing();
|
||||||
return ironclaw::cli::run_doctor_command().await;
|
return ironclaw::cli::run_doctor_command().await;
|
||||||
@@ -584,15 +589,47 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
// ── Gateway channel ────────────────────────────────────────────────
|
// ── Gateway channel ────────────────────────────────────────────────
|
||||||
|
|
||||||
let mut gateway_url: Option<String> = None;
|
let mut gateway_url: Option<String> = None;
|
||||||
let mut sse_sender: Option<
|
let mut sse_manager: Option<std::sync::Arc<ironclaw::channels::web::sse::SseManager>> = None;
|
||||||
tokio::sync::broadcast::Sender<ironclaw::channels::web::types::SseEvent>,
|
|
||||||
> = None;
|
|
||||||
if let Some(ref gw_config) = config.channels.gateway {
|
if let Some(ref gw_config) = config.channels.gateway {
|
||||||
let mut gw =
|
// Build multi-user auth state if user_tokens is configured, else single-user.
|
||||||
GatewayChannel::new(gw_config.clone()).with_llm_provider(Arc::clone(&components.llm));
|
let mut gw = if let Some(ref user_tokens) = gw_config.user_tokens {
|
||||||
|
use ironclaw::channels::web::auth::{MultiAuthState, UserIdentity};
|
||||||
|
let tokens = user_tokens
|
||||||
|
.iter()
|
||||||
|
.map(|(token, cfg)| {
|
||||||
|
(
|
||||||
|
token.clone(),
|
||||||
|
UserIdentity {
|
||||||
|
user_id: cfg.user_id.clone(),
|
||||||
|
workspace_read_scopes: cfg.workspace_read_scopes.clone(),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let auth = MultiAuthState::multi(tokens);
|
||||||
|
GatewayChannel::new_multi_auth(gw_config.clone(), auth)
|
||||||
|
} 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 {
|
if let Some(ref ws) = components.workspace {
|
||||||
gw = gw.with_workspace(Arc::clone(ws));
|
gw = gw.with_workspace(Arc::clone(ws));
|
||||||
}
|
}
|
||||||
|
// Create per-user workspace pool for multi-user mode.
|
||||||
|
if let Some(ref db) = components.db {
|
||||||
|
let emb_cache_config = ironclaw::workspace::EmbeddingCacheConfig {
|
||||||
|
max_entries: config.embeddings.cache_size,
|
||||||
|
};
|
||||||
|
let pool = Arc::new(ironclaw::channels::web::server::WorkspacePool::new(
|
||||||
|
Arc::clone(db),
|
||||||
|
components.embeddings.clone(),
|
||||||
|
emb_cache_config,
|
||||||
|
config.search.clone(),
|
||||||
|
config.workspace.clone(),
|
||||||
|
));
|
||||||
|
gw = gw.with_workspace_pool(pool);
|
||||||
|
}
|
||||||
gw = gw.with_session_manager(Arc::clone(&session_manager));
|
gw = gw.with_session_manager(Arc::clone(&session_manager));
|
||||||
gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster));
|
gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster));
|
||||||
gw = gw.with_log_level_handle(Arc::clone(&log_level_handle));
|
gw = gw.with_log_level_handle(Arc::clone(&log_level_handle));
|
||||||
@@ -643,8 +680,12 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
let mut rx = tx.subscribe();
|
let mut rx = tx.subscribe();
|
||||||
let gw_state = Arc::clone(gw.state());
|
let gw_state = Arc::clone(gw.state());
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
while let Ok((_job_id, event)) = rx.recv().await {
|
while let Ok((_job_id, user_id, event)) = rx.recv().await {
|
||||||
gw_state.sse.broadcast(event);
|
if user_id.is_empty() {
|
||||||
|
gw_state.sse.broadcast(event);
|
||||||
|
} else {
|
||||||
|
gw_state.sse.broadcast_for_user(&user_id, event);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -686,7 +727,7 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
// Capture SSE sender and routine engine slot before moving gw into channels.
|
// Capture SSE sender and routine engine slot before moving gw into channels.
|
||||||
// IMPORTANT: This must come after all `with_*` calls since `rebuild_state`
|
// IMPORTANT: This must come after all `with_*` calls since `rebuild_state`
|
||||||
// creates a new SseManager, which would orphan this sender.
|
// creates a new SseManager, which would orphan this sender.
|
||||||
sse_sender = Some(gw.state().sse.sender());
|
sse_manager = Some(Arc::clone(&gw.state().sse));
|
||||||
channel_names.push("gateway".to_string());
|
channel_names.push("gateway".to_string());
|
||||||
channels.add(Box::new(gw)).await;
|
channels.add(Box::new(gw)).await;
|
||||||
}
|
}
|
||||||
@@ -749,6 +790,14 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
.register_message_tools(Arc::clone(&channels), components.extension_manager.clone())
|
.register_message_tools(Arc::clone(&channels), components.extension_manager.clone())
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
// Default user ID for extension operations (single-user mode).
|
||||||
|
let ext_user_id = config
|
||||||
|
.channels
|
||||||
|
.gateway
|
||||||
|
.as_ref()
|
||||||
|
.map(|g| g.user_id.clone())
|
||||||
|
.unwrap_or_else(|| "default".to_string());
|
||||||
|
|
||||||
// Wire up channel runtime for hot-activation of WASM channels.
|
// Wire up channel runtime for hot-activation of WASM channels.
|
||||||
if let Some(ref ext_mgr) = components.extension_manager
|
if let Some(ref ext_mgr) = components.extension_manager
|
||||||
&& let Some((rt, ps, router)) = wasm_channel_runtime_state.take()
|
&& let Some((rt, ps, router)) = wasm_channel_runtime_state.take()
|
||||||
@@ -769,12 +818,14 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
|
|
||||||
// Auto-activate WASM channels that were active in a previous session.
|
// Auto-activate WASM channels that were active in a previous session.
|
||||||
// Relay channels are handled separately below via restore_relay_channels().
|
// Relay channels are handled separately below via restore_relay_channels().
|
||||||
let persisted = ext_mgr.load_persisted_active_channels().await;
|
let persisted = ext_mgr.load_persisted_active_channels(&ext_user_id).await;
|
||||||
for name in &persisted {
|
for name in &persisted {
|
||||||
if active_at_startup.contains(name) || ext_mgr.is_relay_channel(name).await {
|
if active_at_startup.contains(name)
|
||||||
|
|| ext_mgr.is_relay_channel(name, &ext_user_id).await
|
||||||
|
{
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
match ext_mgr.activate(name).await {
|
match ext_mgr.activate(name, &ext_user_id).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
channel = %name,
|
channel = %name,
|
||||||
@@ -799,14 +850,14 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
ext_mgr
|
ext_mgr
|
||||||
.set_relay_channel_manager(Arc::clone(&channels))
|
.set_relay_channel_manager(Arc::clone(&channels))
|
||||||
.await;
|
.await;
|
||||||
ext_mgr.restore_relay_channels().await;
|
ext_mgr.restore_relay_channels(&ext_user_id).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Wire SSE sender into extension manager for broadcasting status events.
|
// Wire SSE sender into extension manager for broadcasting status events.
|
||||||
if let Some(ref ext_mgr) = components.extension_manager
|
if let Some(ref ext_mgr) = components.extension_manager
|
||||||
&& let Some(ref sender) = sse_sender
|
&& let Some(ref sse) = sse_manager
|
||||||
{
|
{
|
||||||
ext_mgr.set_sse_sender(sender.clone()).await;
|
ext_mgr.set_sse_sender(Arc::clone(sse)).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Snapshot memory for trace recording before the agent starts
|
// Snapshot memory for trace recording before the agent starts
|
||||||
@@ -844,7 +895,7 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
skills_config: config.skills.clone(),
|
skills_config: config.skills.clone(),
|
||||||
hooks: components.hooks,
|
hooks: components.hooks,
|
||||||
cost_guard: components.cost_guard,
|
cost_guard: components.cost_guard,
|
||||||
sse_tx: sse_sender,
|
sse_tx: sse_manager,
|
||||||
http_interceptor,
|
http_interceptor,
|
||||||
transcription: config.transcription.create_provider().map(|p| {
|
transcription: config.transcription.create_provider().map(|p| {
|
||||||
Arc::new(ironclaw::llm::transcription::TranscriptionMiddleware::new(
|
Arc::new(ironclaw::llm::transcription::TranscriptionMiddleware::new(
|
||||||
@@ -862,6 +913,7 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
ironclaw::agent::routine_engine::SandboxReadiness::DockerUnavailable
|
ironclaw::agent::routine_engine::SandboxReadiness::DockerUnavailable
|
||||||
},
|
},
|
||||||
builder: components.builder,
|
builder: components.builder,
|
||||||
|
llm_backend: config.llm.backend.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let channels_for_warnings = Arc::clone(&channels);
|
let channels_for_warnings = Arc::clone(&channels);
|
||||||
|
|||||||
+79
-17
@@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize};
|
|||||||
use tokio::sync::{Mutex, broadcast};
|
use tokio::sync::{Mutex, broadcast};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::ToolDecisionDto;
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||||
@@ -25,6 +25,7 @@ use crate::worker::api::{
|
|||||||
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
||||||
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
||||||
};
|
};
|
||||||
|
use ironclaw_common::AppEvent;
|
||||||
|
|
||||||
/// A follow-up prompt queued for a Claude Code bridge.
|
/// A follow-up prompt queued for a Claude Code bridge.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -40,7 +41,8 @@ pub struct OrchestratorState {
|
|||||||
pub job_manager: Arc<ContainerJobManager>,
|
pub job_manager: Arc<ContainerJobManager>,
|
||||||
pub token_store: TokenStore,
|
pub token_store: TokenStore,
|
||||||
/// Broadcast channel for job events (consumed by the web gateway SSE).
|
/// Broadcast channel for job events (consumed by the web gateway SSE).
|
||||||
pub job_event_tx: Option<broadcast::Sender<(Uuid, SseEvent)>>,
|
/// Tuple: (job_id, user_id, event).
|
||||||
|
pub job_event_tx: Option<broadcast::Sender<(Uuid, String, AppEvent)>>,
|
||||||
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
|
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
|
||||||
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
|
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
|
||||||
/// Database handle for persisting job events.
|
/// Database handle for persisting job events.
|
||||||
@@ -49,6 +51,9 @@ pub struct OrchestratorState {
|
|||||||
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// User ID for secret lookups (single-tenant, typically "default").
|
/// User ID for secret lookups (single-tenant, typically "default").
|
||||||
pub user_id: String,
|
pub user_id: String,
|
||||||
|
/// In-memory cache of job_id → user_id for SSE scoping. Populated when
|
||||||
|
/// sandbox jobs are created, avoiding a DB round-trip on every job event.
|
||||||
|
pub job_owner_cache: Arc<std::sync::RwLock<HashMap<Uuid, String>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The orchestrator's internal API server.
|
/// The orchestrator's internal API server.
|
||||||
@@ -273,10 +278,10 @@ async fn job_event_handler(
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Convert to SSE event and broadcast
|
// Convert to app event and broadcast
|
||||||
let job_id_str = job_id.to_string();
|
let job_id_str = job_id.to_string();
|
||||||
let sse_event = match payload.event_type.as_str() {
|
let app_event = match payload.event_type.as_str() {
|
||||||
"message" => SseEvent::JobMessage {
|
"message" => AppEvent::JobMessage {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
role: payload
|
role: payload
|
||||||
.data
|
.data
|
||||||
@@ -291,7 +296,7 @@ async fn job_event_handler(
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
},
|
},
|
||||||
"tool_use" => SseEvent::JobToolUse {
|
"tool_use" => AppEvent::JobToolUse {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: payload
|
tool_name: payload
|
||||||
.data
|
.data
|
||||||
@@ -305,7 +310,7 @@ async fn job_event_handler(
|
|||||||
.cloned()
|
.cloned()
|
||||||
.unwrap_or(serde_json::Value::Null),
|
.unwrap_or(serde_json::Value::Null),
|
||||||
},
|
},
|
||||||
"tool_result" => SseEvent::JobToolResult {
|
"tool_result" => AppEvent::JobToolResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: payload
|
tool_name: payload
|
||||||
.data
|
.data
|
||||||
@@ -320,7 +325,7 @@ async fn job_event_handler(
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
},
|
},
|
||||||
"result" => SseEvent::JobResult {
|
"result" => AppEvent::JobResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
status: payload
|
status: payload
|
||||||
.data
|
.data
|
||||||
@@ -340,7 +345,21 @@ async fn job_event_handler(
|
|||||||
// gain context/memory tracking capabilities.
|
// gain context/memory tracking capabilities.
|
||||||
fallback_deliverable: payload.data.get("fallback_deliverable").cloned(),
|
fallback_deliverable: payload.data.get("fallback_deliverable").cloned(),
|
||||||
},
|
},
|
||||||
_ => SseEvent::JobStatus {
|
"reasoning" => {
|
||||||
|
let narrative = payload
|
||||||
|
.data
|
||||||
|
.get("narrative")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
let decisions = ToolDecisionDto::from_json_array(&payload.data["decisions"]);
|
||||||
|
AppEvent::JobReasoning {
|
||||||
|
job_id: job_id_str,
|
||||||
|
narrative,
|
||||||
|
decisions,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => AppEvent::JobStatus {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
message: payload
|
message: payload
|
||||||
.data
|
.data
|
||||||
@@ -351,9 +370,45 @@ async fn job_event_handler(
|
|||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
// Broadcast via the channel (if configured)
|
// Broadcast via the channel (if configured).
|
||||||
|
// Look up the job owner from the in-memory cache (populated at job creation).
|
||||||
if let Some(ref tx) = state.job_event_tx {
|
if let Some(ref tx) = state.job_event_tx {
|
||||||
let _ = tx.send((job_id, sse_event));
|
let cached_uid = state
|
||||||
|
.job_owner_cache
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
|
.get(&job_id)
|
||||||
|
.cloned();
|
||||||
|
|
||||||
|
let user_id = match cached_uid {
|
||||||
|
Some(uid) => uid,
|
||||||
|
None => {
|
||||||
|
// Cache miss: fall back to DB lookup and populate cache.
|
||||||
|
let uid = match state.store.as_ref() {
|
||||||
|
Some(store) => store
|
||||||
|
.get_sandbox_job(job_id)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.map(|j| j.user_id),
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
if let Some(ref uid) = uid {
|
||||||
|
state
|
||||||
|
.job_owner_cache
|
||||||
|
.write()
|
||||||
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
|
.insert(job_id, uid.clone());
|
||||||
|
}
|
||||||
|
uid.unwrap_or_default()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if user_id.is_empty() {
|
||||||
|
let _ = tx.send((job_id, String::new(), app_event));
|
||||||
|
} else {
|
||||||
|
let _ = tx.send((job_id, user_id, app_event));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(StatusCode::OK)
|
Ok(StatusCode::OK)
|
||||||
@@ -480,6 +535,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -709,6 +765,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: Some(secrets_store),
|
secrets_store: Some(secrets_store),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
let router = OrchestratorApi::router(state);
|
let router = OrchestratorApi::router(state);
|
||||||
@@ -744,6 +801,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -769,10 +827,12 @@ mod tests {
|
|||||||
let resp = router.oneshot(req).await.unwrap();
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
assert_eq!(resp.status(), StatusCode::OK);
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
let (recv_id, event) = rx.recv().await.unwrap();
|
let (recv_id, recv_uid, event) = rx.recv().await.unwrap();
|
||||||
assert_eq!(recv_id, job_id);
|
assert_eq!(recv_id, job_id);
|
||||||
|
// No store configured, so user_id falls back to empty string.
|
||||||
|
assert_eq!(recv_uid, "");
|
||||||
match event {
|
match event {
|
||||||
SseEvent::JobMessage {
|
AppEvent::JobMessage {
|
||||||
job_id: jid,
|
job_id: jid,
|
||||||
role,
|
role,
|
||||||
content,
|
content,
|
||||||
@@ -799,6 +859,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -824,9 +885,9 @@ mod tests {
|
|||||||
let resp = router.oneshot(req).await.unwrap();
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
assert_eq!(resp.status(), StatusCode::OK);
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
let (_recv_id, event) = rx.recv().await.unwrap();
|
let (_recv_id, _recv_uid, event) = rx.recv().await.unwrap();
|
||||||
match event {
|
match event {
|
||||||
SseEvent::JobToolUse { tool_name, .. } => {
|
AppEvent::JobToolUse { tool_name, .. } => {
|
||||||
assert_eq!(tool_name, "shell");
|
assert_eq!(tool_name, "shell");
|
||||||
}
|
}
|
||||||
other => panic!("Expected JobToolUse, got {:?}", other),
|
other => panic!("Expected JobToolUse, got {:?}", other),
|
||||||
@@ -847,6 +908,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -869,9 +931,9 @@ mod tests {
|
|||||||
let resp = router.oneshot(req).await.unwrap();
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
assert_eq!(resp.status(), StatusCode::OK);
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
let (_recv_id, event) = rx.recv().await.unwrap();
|
let (_recv_id, _recv_uid, event) = rx.recv().await.unwrap();
|
||||||
// Unknown event types fall through to JobStatus
|
// Unknown event types fall through to JobStatus
|
||||||
assert!(matches!(event, SseEvent::JobStatus { .. }));
|
assert!(matches!(event, AppEvent::JobStatus { .. }));
|
||||||
}
|
}
|
||||||
|
|
||||||
// -- Status update test --
|
// -- Status update test --
|
||||||
|
|||||||
@@ -46,10 +46,10 @@ use std::sync::Arc;
|
|||||||
use tokio::sync::{Mutex, broadcast};
|
use tokio::sync::{Mutex, broadcast};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::web::types::SseEvent;
|
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
|
use ironclaw_common::AppEvent;
|
||||||
|
|
||||||
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
||||||
/// variable, falling back to 50051.
|
/// variable, falling back to 50051.
|
||||||
@@ -63,7 +63,7 @@ fn resolve_orchestrator_port() -> u16 {
|
|||||||
/// Result of orchestrator setup, containing all handles needed by the agent.
|
/// Result of orchestrator setup, containing all handles needed by the agent.
|
||||||
pub struct OrchestratorSetup {
|
pub struct OrchestratorSetup {
|
||||||
pub container_job_manager: Option<Arc<ContainerJobManager>>,
|
pub container_job_manager: Option<Arc<ContainerJobManager>>,
|
||||||
pub job_event_tx: Option<broadcast::Sender<(Uuid, SseEvent)>>,
|
pub job_event_tx: Option<broadcast::Sender<(Uuid, String, AppEvent)>>,
|
||||||
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<api::PendingPrompt>>>>,
|
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<api::PendingPrompt>>>>,
|
||||||
pub docker_status: crate::sandbox::DockerStatus,
|
pub docker_status: crate::sandbox::DockerStatus,
|
||||||
}
|
}
|
||||||
@@ -134,6 +134,7 @@ pub async fn setup_orchestrator(
|
|||||||
store: db.cloned(),
|
store: db.cloned(),
|
||||||
secrets_store: secrets_store.cloned(),
|
secrets_store: secrets_store.cloned(),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
job_owner_cache: Arc::new(std::sync::RwLock::new(std::collections::HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
|
|||||||
+108
@@ -1297,6 +1297,92 @@ mod tests {
|
|||||||
assert_eq!(loaded.heartbeat.interval_secs, 900);
|
assert_eq!(loaded.heartbeat.interval_secs, 900);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Regression: /model writes a single key ("selected_model") to the DB via
|
||||||
|
/// set_setting(). On restart, get_all_settings() returns ALL keys including
|
||||||
|
/// wizard-written defaults. The single-key update must survive the full
|
||||||
|
/// from_db_map() round trip.
|
||||||
|
#[test]
|
||||||
|
fn db_single_key_model_update_survives_roundtrip() {
|
||||||
|
// Step 1: Wizard writes full settings to DB (including selected_model
|
||||||
|
// from initial setup).
|
||||||
|
let wizard_settings = Settings {
|
||||||
|
llm_backend: Some("nearai".to_string()),
|
||||||
|
selected_model: Some("old-wizard-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let mut db: std::collections::HashMap<String, serde_json::Value> =
|
||||||
|
wizard_settings.to_db_map();
|
||||||
|
|
||||||
|
// Step 2: User runs /model new-model — persist_selected_model writes
|
||||||
|
// a single key, overwriting the wizard value.
|
||||||
|
db.insert(
|
||||||
|
"selected_model".to_string(),
|
||||||
|
serde_json::Value::String("new-model".to_string()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Step 3: On restart, from_db_map() rebuilds Settings from the full
|
||||||
|
// DB map.
|
||||||
|
let restored = Settings::from_db_map(&db);
|
||||||
|
assert_eq!(
|
||||||
|
restored.selected_model,
|
||||||
|
Some("new-model".to_string()),
|
||||||
|
"/model change must survive DB round trip"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Regression: TOML overlay must not clobber a DB-persisted selected_model
|
||||||
|
/// when the TOML file matches the DB. This is the normal case after /model
|
||||||
|
/// successfully writes to both DB and TOML.
|
||||||
|
#[test]
|
||||||
|
fn toml_overlay_preserves_matching_model() {
|
||||||
|
// DB settings with new model from /model command.
|
||||||
|
let mut db_settings = Settings {
|
||||||
|
llm_backend: Some("nearai".to_string()),
|
||||||
|
selected_model: Some("new-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
// TOML also updated by /model command to the same value.
|
||||||
|
let toml_settings = Settings {
|
||||||
|
selected_model: Some("new-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
db_settings.merge_from(&toml_settings);
|
||||||
|
assert_eq!(
|
||||||
|
db_settings.selected_model,
|
||||||
|
Some("new-model".to_string()),
|
||||||
|
"TOML overlay must not clobber matching model"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Regression: when /model updates DB but TOML write fails, a stale TOML
|
||||||
|
/// file would overwrite the DB value. This test documents the priority:
|
||||||
|
/// TOML > DB (by design). persist_selected_model MUST update the TOML.
|
||||||
|
#[test]
|
||||||
|
fn stale_toml_overwrites_db_model() {
|
||||||
|
// DB has the new model from /model.
|
||||||
|
let mut db_settings = Settings {
|
||||||
|
selected_model: Some("new-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
// TOML still has the old model (write failed or was not attempted).
|
||||||
|
let stale_toml = Settings {
|
||||||
|
selected_model: Some("old-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
db_settings.merge_from(&stale_toml);
|
||||||
|
// This documents the current priority: TOML wins over DB.
|
||||||
|
// The fix in persist_selected_model ensures TOML is always updated.
|
||||||
|
assert_eq!(
|
||||||
|
db_settings.selected_model,
|
||||||
|
Some("old-model".to_string()),
|
||||||
|
"TOML overlay has higher priority than DB (by design)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
/// Regression test: /model command must persist selected_model to TOML config.
|
/// Regression test: /model command must persist selected_model to TOML config.
|
||||||
/// Prior to the fix, `set_model()` only changed the in-memory provider and the
|
/// Prior to the fix, `set_model()` only changed the in-memory provider and the
|
||||||
/// choice was lost on restart.
|
/// choice was lost on restart.
|
||||||
@@ -1322,6 +1408,28 @@ mod tests {
|
|||||||
assert_eq!(reloaded.selected_model, Some("new-model".to_string()));
|
assert_eq!(reloaded.selected_model, Some("new-model".to_string()));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Regression: /model must create config.toml when it doesn't exist, so the
|
||||||
|
/// model survives restarts. Previously the Ok(None) case was a no-op.
|
||||||
|
#[test]
|
||||||
|
fn toml_created_when_missing_for_model_persist() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let path = dir.path().join("config.toml");
|
||||||
|
|
||||||
|
// No config.toml yet (fresh install, no wizard).
|
||||||
|
assert!(Settings::load_toml(&path).unwrap().is_none());
|
||||||
|
|
||||||
|
// Simulate what persist_selected_model now does for the Ok(None) case.
|
||||||
|
let settings = Settings {
|
||||||
|
selected_model: Some("new-model".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
settings.save_toml(&path).unwrap();
|
||||||
|
|
||||||
|
// Verify the model survived.
|
||||||
|
let loaded = Settings::load_toml(&path).unwrap().unwrap();
|
||||||
|
assert_eq!(loaded.selected_model, Some("new-model".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn toml_missing_file_returns_none() {
|
fn toml_missing_file_returns_none() {
|
||||||
let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml"));
|
let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml"));
|
||||||
|
|||||||
@@ -563,6 +563,7 @@ impl TestHarnessBuilder {
|
|||||||
document_extraction: None,
|
document_extraction: None,
|
||||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||||
builder: None,
|
builder: None,
|
||||||
|
llm_backend: "nearai".to_string(),
|
||||||
};
|
};
|
||||||
|
|
||||||
TestHarness {
|
TestHarness {
|
||||||
|
|||||||
@@ -130,7 +130,7 @@ impl Tool for ToolInstallTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -150,7 +150,7 @@ impl Tool for ToolInstallTool {
|
|||||||
|
|
||||||
let result = self
|
let result = self
|
||||||
.manager
|
.manager
|
||||||
.install(name, url, kind_hint)
|
.install(name, url, kind_hint, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
@@ -205,7 +205,7 @@ impl Tool for ToolAuthTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -213,13 +213,13 @@ impl Tool for ToolAuthTool {
|
|||||||
|
|
||||||
let result = self
|
let result = self
|
||||||
.manager
|
.manager
|
||||||
.auth(name)
|
.auth(name, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
// Auto-activate after successful auth so tools are available immediately
|
// Auto-activate after successful auth so tools are available immediately
|
||||||
if result.is_authenticated() {
|
if result.is_authenticated() {
|
||||||
match self.manager.activate(name).await {
|
match self.manager.activate(name, &ctx.user_id).await {
|
||||||
Ok(activate_result) => {
|
Ok(activate_result) => {
|
||||||
let output = serde_json::json!({
|
let output = serde_json::json!({
|
||||||
"status": "authenticated_and_activated",
|
"status": "authenticated_and_activated",
|
||||||
@@ -304,13 +304,13 @@ impl Tool for ToolActivateTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
let name = require_str(¶ms, "name")?;
|
let name = require_str(¶ms, "name")?;
|
||||||
|
|
||||||
match self.manager.activate(name).await {
|
match self.manager.activate(name, &ctx.user_id).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
let output = serde_json::to_value(&result)
|
let output = serde_json::to_value(&result)
|
||||||
.unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"}));
|
.unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"}));
|
||||||
@@ -329,12 +329,12 @@ impl Tool for ToolActivateTool {
|
|||||||
|
|
||||||
// Activation failed due to missing auth; initiate auth flow
|
// Activation failed due to missing auth; initiate auth flow
|
||||||
// so the agent loop can show the auth card.
|
// so the agent loop can show the auth card.
|
||||||
match self.manager.auth(name).await {
|
match self.manager.auth(name, &ctx.user_id).await {
|
||||||
Ok(auth_result) if auth_result.is_authenticated() => {
|
Ok(auth_result) if auth_result.is_authenticated() => {
|
||||||
// Auth succeeded (e.g. env var was set); retry activation.
|
// Auth succeeded (e.g. env var was set); retry activation.
|
||||||
let result = self
|
let result = self
|
||||||
.manager
|
.manager
|
||||||
.activate(name)
|
.activate(name, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
let output = serde_json::to_value(&result).unwrap_or_else(
|
let output = serde_json::to_value(&result).unwrap_or_else(
|
||||||
@@ -404,7 +404,7 @@ impl Tool for ToolListTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -425,7 +425,7 @@ impl Tool for ToolListTool {
|
|||||||
|
|
||||||
let extensions = self
|
let extensions = self
|
||||||
.manager
|
.manager
|
||||||
.list(kind_filter, include_available)
|
.list(kind_filter, include_available, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
@@ -477,7 +477,7 @@ impl Tool for ToolRemoveTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -485,7 +485,7 @@ impl Tool for ToolRemoveTool {
|
|||||||
|
|
||||||
let message = self
|
let message = self
|
||||||
.manager
|
.manager
|
||||||
.remove(name)
|
.remove(name, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
@@ -541,7 +541,7 @@ impl Tool for ToolUpgradeTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -549,7 +549,7 @@ impl Tool for ToolUpgradeTool {
|
|||||||
|
|
||||||
let result = self
|
let result = self
|
||||||
.manager
|
.manager
|
||||||
.upgrade(name)
|
.upgrade(name, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
@@ -603,7 +603,7 @@ impl Tool for ExtensionInfoTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -611,7 +611,7 @@ impl Tool for ExtensionInfoTool {
|
|||||||
|
|
||||||
let info = self
|
let info = self
|
||||||
.manager
|
.manager
|
||||||
.extension_info(name)
|
.extension_info(name, &ctx.user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
|||||||
@@ -17,7 +17,6 @@ use uuid::Uuid;
|
|||||||
|
|
||||||
use crate::bootstrap::ironclaw_base_dir;
|
use crate::bootstrap::ironclaw_base_dir;
|
||||||
use crate::channels::IncomingMessage;
|
use crate::channels::IncomingMessage;
|
||||||
use crate::channels::web::types::SseEvent;
|
|
||||||
use crate::context::{ContextManager, JobContext, JobState};
|
use crate::context::{ContextManager, JobContext, JobState};
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::history::SandboxJobRecord;
|
use crate::history::SandboxJobRecord;
|
||||||
@@ -25,6 +24,7 @@ use crate::orchestrator::auth::CredentialGrant;
|
|||||||
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
|
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||||
|
use ironclaw_common::AppEvent;
|
||||||
|
|
||||||
/// Lazy scheduler reference, filled after Agent::new creates the Scheduler.
|
/// Lazy scheduler reference, filled after Agent::new creates the Scheduler.
|
||||||
///
|
///
|
||||||
@@ -85,7 +85,7 @@ pub struct CreateJobTool {
|
|||||||
job_manager: Option<Arc<ContainerJobManager>>,
|
job_manager: Option<Arc<ContainerJobManager>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<Arc<dyn Database>>,
|
||||||
/// Broadcast sender for job events (used to subscribe a monitor).
|
/// Broadcast sender for job events (used to subscribe a monitor).
|
||||||
event_tx: Option<tokio::sync::broadcast::Sender<(Uuid, SseEvent)>>,
|
event_tx: Option<tokio::sync::broadcast::Sender<(Uuid, String, AppEvent)>>,
|
||||||
/// Injection channel for pushing messages into the agent loop.
|
/// Injection channel for pushing messages into the agent loop.
|
||||||
inject_tx: Option<tokio::sync::mpsc::Sender<IncomingMessage>>,
|
inject_tx: Option<tokio::sync::mpsc::Sender<IncomingMessage>>,
|
||||||
/// Encrypted secrets store for validating credential grants.
|
/// Encrypted secrets store for validating credential grants.
|
||||||
@@ -120,7 +120,7 @@ impl CreateJobTool {
|
|||||||
/// monitor that forwards Claude Code output to the main agent loop.
|
/// monitor that forwards Claude Code output to the main agent loop.
|
||||||
pub fn with_monitor_deps(
|
pub fn with_monitor_deps(
|
||||||
mut self,
|
mut self,
|
||||||
event_tx: tokio::sync::broadcast::Sender<(Uuid, SseEvent)>,
|
event_tx: tokio::sync::broadcast::Sender<(Uuid, String, AppEvent)>,
|
||||||
inject_tx: tokio::sync::mpsc::Sender<IncomingMessage>,
|
inject_tx: tokio::sync::mpsc::Sender<IncomingMessage>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
self.event_tx = Some(event_tx);
|
self.event_tx = Some(event_tx);
|
||||||
|
|||||||
+283
-45
@@ -21,6 +21,35 @@ use crate::context::JobContext;
|
|||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||||
use crate::workspace::{Workspace, paths};
|
use crate::workspace::{Workspace, paths};
|
||||||
|
|
||||||
|
// ── WorkspaceResolver ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Resolves a workspace for a given user ID.
|
||||||
|
///
|
||||||
|
/// In single-user mode, always returns the same workspace.
|
||||||
|
/// In multi-tenant mode, creates per-user workspaces on demand.
|
||||||
|
#[async_trait]
|
||||||
|
pub trait WorkspaceResolver: Send + Sync {
|
||||||
|
async fn resolve(&self, user_id: &str) -> Arc<Workspace>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns a fixed workspace regardless of user ID (single-user mode).
|
||||||
|
pub struct FixedWorkspaceResolver {
|
||||||
|
workspace: Arc<Workspace>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FixedWorkspaceResolver {
|
||||||
|
pub fn new(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self { workspace }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl WorkspaceResolver for FixedWorkspaceResolver {
|
||||||
|
async fn resolve(&self, _user_id: &str) -> Arc<Workspace> {
|
||||||
|
Arc::clone(&self.workspace)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Detect paths that are clearly local filesystem references, not workspace-memory docs.
|
/// Detect paths that are clearly local filesystem references, not workspace-memory docs.
|
||||||
///
|
///
|
||||||
/// Examples:
|
/// Examples:
|
||||||
@@ -62,13 +91,20 @@ fn map_write_err(e: crate::error::WorkspaceError) -> ToolError {
|
|||||||
/// The agent should call this tool before answering questions about
|
/// The agent should call this tool before answering questions about
|
||||||
/// prior work, decisions, preferences, or any historical context.
|
/// prior work, decisions, preferences, or any historical context.
|
||||||
pub struct MemorySearchTool {
|
pub struct MemorySearchTool {
|
||||||
workspace: Arc<Workspace>,
|
resolver: Arc<dyn WorkspaceResolver>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemorySearchTool {
|
impl MemorySearchTool {
|
||||||
/// Create a new memory search tool.
|
/// Create a new memory search tool with a workspace resolver.
|
||||||
pub fn new(workspace: Arc<Workspace>) -> Self {
|
pub fn new(resolver: Arc<dyn WorkspaceResolver>) -> Self {
|
||||||
Self { workspace }
|
Self { resolver }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create from a fixed workspace (backward compatibility).
|
||||||
|
pub fn from_workspace(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self {
|
||||||
|
resolver: Arc::new(FixedWorkspaceResolver::new(workspace)),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -107,7 +143,7 @@ impl Tool for MemorySearchTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -119,8 +155,8 @@ impl Tool for MemorySearchTool {
|
|||||||
.unwrap_or(5)
|
.unwrap_or(5)
|
||||||
.min(20) as usize;
|
.min(20) as usize;
|
||||||
|
|
||||||
let results = self
|
let workspace = self.resolver.resolve(&ctx.user_id).await;
|
||||||
.workspace
|
let results = workspace
|
||||||
.search(query, limit)
|
.search(query, limit)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(format!("Search failed: {}", e)))?;
|
.map_err(|e| ToolError::ExecutionFailed(format!("Search failed: {}", e)))?;
|
||||||
@@ -151,13 +187,20 @@ impl Tool for MemorySearchTool {
|
|||||||
/// Use this to persist important information that should be remembered
|
/// Use this to persist important information that should be remembered
|
||||||
/// across sessions: decisions, preferences, facts, lessons learned.
|
/// across sessions: decisions, preferences, facts, lessons learned.
|
||||||
pub struct MemoryWriteTool {
|
pub struct MemoryWriteTool {
|
||||||
workspace: Arc<Workspace>,
|
resolver: Arc<dyn WorkspaceResolver>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryWriteTool {
|
impl MemoryWriteTool {
|
||||||
/// Create a new memory write tool.
|
/// Create a new memory write tool with a workspace resolver.
|
||||||
pub fn new(workspace: Arc<Workspace>) -> Self {
|
pub fn new(resolver: Arc<dyn WorkspaceResolver>) -> Self {
|
||||||
Self { workspace }
|
Self { resolver }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create from a fixed workspace (backward compatibility).
|
||||||
|
pub fn from_workspace(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self {
|
||||||
|
resolver: Arc::new(FixedWorkspaceResolver::new(workspace)),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -231,19 +274,21 @@ impl Tool for MemoryWriteTool {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let workspace = self.resolver.resolve(&ctx.user_id).await;
|
||||||
|
|
||||||
// Bootstrap target: clear BOOTSTRAP.md to mark first-run ritual complete.
|
// Bootstrap target: clear BOOTSTRAP.md to mark first-run ritual complete.
|
||||||
// Handled early because it accepts empty content (unlike other targets).
|
// Handled early because it accepts empty content (unlike other targets).
|
||||||
if target == "bootstrap" {
|
if target == "bootstrap" {
|
||||||
// Write empty content to effectively disable the bootstrap injection.
|
// Write empty content to effectively disable the bootstrap injection.
|
||||||
// system_prompt_for_context() skips empty files.
|
// system_prompt_for_context() skips empty files.
|
||||||
self.workspace
|
workspace
|
||||||
.write(paths::BOOTSTRAP, "")
|
.write(paths::BOOTSTRAP, "")
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
|
|
||||||
// Also set the in-memory flag so BOOTSTRAP.md injection stops
|
// Also set the in-memory flag so BOOTSTRAP.md injection stops
|
||||||
// immediately without waiting for a restart.
|
// immediately without waiting for a restart.
|
||||||
self.workspace.mark_bootstrap_completed();
|
workspace.mark_bootstrap_completed();
|
||||||
|
|
||||||
let output = serde_json::json!({
|
let output = serde_json::json!({
|
||||||
"status": "cleared",
|
"status": "cleared",
|
||||||
@@ -289,12 +334,12 @@ impl Tool for MemoryWriteTool {
|
|||||||
// Otherwise, use default workspace methods (which include injection scanning).
|
// Otherwise, use default workspace methods (which include injection scanning).
|
||||||
let layer_result = if let Some(layer_name) = layer {
|
let layer_result = if let Some(layer_name) = layer {
|
||||||
let result = if append {
|
let result = if append {
|
||||||
self.workspace
|
workspace
|
||||||
.append_to_layer(layer_name, &resolved_path, content, force)
|
.append_to_layer(layer_name, &resolved_path, content, force)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?
|
.map_err(map_write_err)?
|
||||||
} else {
|
} else {
|
||||||
self.workspace
|
workspace
|
||||||
.write_to_layer(layer_name, &resolved_path, content, force)
|
.write_to_layer(layer_name, &resolved_path, content, force)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?
|
.map_err(map_write_err)?
|
||||||
@@ -307,31 +352,33 @@ impl Tool for MemoryWriteTool {
|
|||||||
match target {
|
match target {
|
||||||
"memory" => {
|
"memory" => {
|
||||||
if append {
|
if append {
|
||||||
self.workspace
|
workspace
|
||||||
.append_memory(content)
|
.append_memory(content)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
} else {
|
} else {
|
||||||
self.workspace
|
workspace
|
||||||
.write(paths::MEMORY, content)
|
.write(paths::MEMORY, content)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"daily_log" => {
|
"daily_log" => {
|
||||||
self.workspace
|
let tz = crate::timezone::parse_timezone(&ctx.user_timezone)
|
||||||
|
.unwrap_or(chrono_tz::Tz::UTC);
|
||||||
|
workspace
|
||||||
.append_daily_log_tz(content, tz)
|
.append_daily_log_tz(content, tz)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
if append {
|
if append {
|
||||||
self.workspace
|
workspace
|
||||||
.append(&resolved_path, content)
|
.append(&resolved_path, content)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
} else {
|
} else {
|
||||||
self.workspace
|
workspace
|
||||||
.write(&resolved_path, content)
|
.write(&resolved_path, content)
|
||||||
.await
|
.await
|
||||||
.map_err(map_write_err)?;
|
.map_err(map_write_err)?;
|
||||||
@@ -361,12 +408,12 @@ impl Tool for MemoryWriteTool {
|
|||||||
};
|
};
|
||||||
let mut synced_docs: Vec<&str> = Vec::new();
|
let mut synced_docs: Vec<&str> = Vec::new();
|
||||||
if normalized_path == paths::PROFILE {
|
if normalized_path == paths::PROFILE {
|
||||||
match self.workspace.sync_profile_documents().await {
|
match workspace.sync_profile_documents().await {
|
||||||
Ok(true) => {
|
Ok(true) => {
|
||||||
tracing::info!("profile write: synced USER.md + assistant-directives.md");
|
tracing::info!("profile write: synced USER.md + assistant-directives.md");
|
||||||
synced_docs.extend_from_slice(&[paths::USER, paths::ASSISTANT_DIRECTIVES]);
|
synced_docs.extend_from_slice(&[paths::USER, paths::ASSISTANT_DIRECTIVES]);
|
||||||
|
|
||||||
self.workspace.mark_bootstrap_completed();
|
workspace.mark_bootstrap_completed();
|
||||||
let toml_path = crate::settings::Settings::default_toml_path();
|
let toml_path = crate::settings::Settings::default_toml_path();
|
||||||
if let Ok(Some(mut settings)) = crate::settings::Settings::load_toml(&toml_path)
|
if let Ok(Some(mut settings)) = crate::settings::Settings::load_toml(&toml_path)
|
||||||
&& !settings.profile_onboarding_completed
|
&& !settings.profile_onboarding_completed
|
||||||
@@ -416,13 +463,20 @@ impl Tool for MemoryWriteTool {
|
|||||||
///
|
///
|
||||||
/// Use this to read the full content of any file in the workspace.
|
/// Use this to read the full content of any file in the workspace.
|
||||||
pub struct MemoryReadTool {
|
pub struct MemoryReadTool {
|
||||||
workspace: Arc<Workspace>,
|
resolver: Arc<dyn WorkspaceResolver>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryReadTool {
|
impl MemoryReadTool {
|
||||||
/// Create a new memory read tool.
|
/// Create a new memory read tool with a workspace resolver.
|
||||||
pub fn new(workspace: Arc<Workspace>) -> Self {
|
pub fn new(resolver: Arc<dyn WorkspaceResolver>) -> Self {
|
||||||
Self { workspace }
|
Self { resolver }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create from a fixed workspace (backward compatibility).
|
||||||
|
pub fn from_workspace(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self {
|
||||||
|
resolver: Arc::new(FixedWorkspaceResolver::new(workspace)),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -456,7 +510,7 @@ impl Tool for MemoryReadTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -470,8 +524,8 @@ impl Tool for MemoryReadTool {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
let doc = self
|
let workspace = self.resolver.resolve(&ctx.user_id).await;
|
||||||
.workspace
|
let doc = workspace
|
||||||
.read(path)
|
.read(path)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?;
|
.map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?;
|
||||||
@@ -495,20 +549,27 @@ impl Tool for MemoryReadTool {
|
|||||||
///
|
///
|
||||||
/// Returns a hierarchical view of files and directories with configurable depth.
|
/// Returns a hierarchical view of files and directories with configurable depth.
|
||||||
pub struct MemoryTreeTool {
|
pub struct MemoryTreeTool {
|
||||||
workspace: Arc<Workspace>,
|
resolver: Arc<dyn WorkspaceResolver>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryTreeTool {
|
impl MemoryTreeTool {
|
||||||
/// Create a new memory tree tool.
|
/// Create a new memory tree tool with a workspace resolver.
|
||||||
pub fn new(workspace: Arc<Workspace>) -> Self {
|
pub fn new(resolver: Arc<dyn WorkspaceResolver>) -> Self {
|
||||||
Self { workspace }
|
Self { resolver }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create from a fixed workspace (backward compatibility).
|
||||||
|
pub fn from_workspace(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self {
|
||||||
|
resolver: Arc::new(FixedWorkspaceResolver::new(workspace)),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Recursively build tree structure.
|
/// Recursively build tree structure.
|
||||||
///
|
///
|
||||||
/// Returns a compact format where directories end with `/` and may have children.
|
/// Returns a compact format where directories end with `/` and may have children.
|
||||||
async fn build_tree(
|
async fn build_tree(
|
||||||
&self,
|
workspace: &Arc<Workspace>,
|
||||||
path: &str,
|
path: &str,
|
||||||
current_depth: usize,
|
current_depth: usize,
|
||||||
max_depth: usize,
|
max_depth: usize,
|
||||||
@@ -517,8 +578,7 @@ impl MemoryTreeTool {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let entries = self
|
let entries = workspace
|
||||||
.workspace
|
|
||||||
.list(path)
|
.list(path)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(format!("Tree failed: {}", e)))?;
|
.map_err(|e| ToolError::ExecutionFailed(format!("Tree failed: {}", e)))?;
|
||||||
@@ -533,8 +593,13 @@ impl MemoryTreeTool {
|
|||||||
};
|
};
|
||||||
|
|
||||||
if entry.is_directory && current_depth < max_depth {
|
if entry.is_directory && current_depth < max_depth {
|
||||||
let children =
|
let children = Box::pin(Self::build_tree(
|
||||||
Box::pin(self.build_tree(&entry.path, current_depth + 1, max_depth)).await?;
|
workspace,
|
||||||
|
&entry.path,
|
||||||
|
current_depth + 1,
|
||||||
|
max_depth,
|
||||||
|
))
|
||||||
|
.await?;
|
||||||
if children.is_empty() {
|
if children.is_empty() {
|
||||||
result.push(serde_json::Value::String(display_path));
|
result.push(serde_json::Value::String(display_path));
|
||||||
} else {
|
} else {
|
||||||
@@ -584,7 +649,7 @@ impl Tool for MemoryTreeTool {
|
|||||||
async fn execute(
|
async fn execute(
|
||||||
&self,
|
&self,
|
||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
_ctx: &JobContext,
|
ctx: &JobContext,
|
||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
@@ -596,7 +661,8 @@ impl Tool for MemoryTreeTool {
|
|||||||
.unwrap_or(1)
|
.unwrap_or(1)
|
||||||
.clamp(1, 10) as usize;
|
.clamp(1, 10) as usize;
|
||||||
|
|
||||||
let tree = self.build_tree(path, 1, depth).await?;
|
let workspace = self.resolver.resolve(&ctx.user_id).await;
|
||||||
|
let tree = Self::build_tree(&workspace, path, 1, depth).await?;
|
||||||
|
|
||||||
// Compact output: just the tree array
|
// Compact output: just the tree array
|
||||||
Ok(ToolOutput::success(
|
Ok(ToolOutput::success(
|
||||||
@@ -650,7 +716,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_memory_search_schema() {
|
fn test_memory_search_schema() {
|
||||||
let workspace = make_test_workspace();
|
let workspace = make_test_workspace();
|
||||||
let tool = MemorySearchTool::new(workspace);
|
let tool = MemorySearchTool::from_workspace(workspace);
|
||||||
|
|
||||||
assert_eq!(tool.name(), "memory_search");
|
assert_eq!(tool.name(), "memory_search");
|
||||||
assert!(!tool.requires_sanitization());
|
assert!(!tool.requires_sanitization());
|
||||||
@@ -668,7 +734,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_memory_write_schema() {
|
fn test_memory_write_schema() {
|
||||||
let workspace = make_test_workspace();
|
let workspace = make_test_workspace();
|
||||||
let tool = MemoryWriteTool::new(workspace);
|
let tool = MemoryWriteTool::from_workspace(workspace);
|
||||||
|
|
||||||
assert_eq!(tool.name(), "memory_write");
|
assert_eq!(tool.name(), "memory_write");
|
||||||
|
|
||||||
@@ -681,7 +747,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_memory_read_schema() {
|
fn test_memory_read_schema() {
|
||||||
let workspace = make_test_workspace();
|
let workspace = make_test_workspace();
|
||||||
let tool = MemoryReadTool::new(workspace);
|
let tool = MemoryReadTool::from_workspace(workspace);
|
||||||
|
|
||||||
assert_eq!(tool.name(), "memory_read");
|
assert_eq!(tool.name(), "memory_read");
|
||||||
|
|
||||||
@@ -698,7 +764,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_memory_tree_schema() {
|
fn test_memory_tree_schema() {
|
||||||
let workspace = make_test_workspace();
|
let workspace = make_test_workspace();
|
||||||
let tool = MemoryTreeTool::new(workspace);
|
let tool = MemoryTreeTool::from_workspace(workspace);
|
||||||
|
|
||||||
assert_eq!(tool.name(), "memory_tree");
|
assert_eq!(tool.name(), "memory_tree");
|
||||||
|
|
||||||
@@ -711,7 +777,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_memory_write_rejects_injection_to_identity_file() {
|
async fn test_memory_write_rejects_injection_to_identity_file() {
|
||||||
let workspace = make_test_workspace();
|
let workspace = make_test_workspace();
|
||||||
let tool = MemoryWriteTool::new(workspace);
|
let tool = MemoryWriteTool::from_workspace(workspace);
|
||||||
let ctx = JobContext::default();
|
let ctx = JobContext::default();
|
||||||
|
|
||||||
let params = serde_json::json!({
|
let params = serde_json::json!({
|
||||||
@@ -733,4 +799,176 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Regression tests for per-user workspace scoping (multi-tenant mode).
|
||||||
|
// See: https://github.com/nearai/ironclaw/pull/1118
|
||||||
|
// Bug: memory tools used a single startup workspace regardless of which
|
||||||
|
// user was chatting. Fix: resolve workspace per-request via JobContext.user_id.
|
||||||
|
|
||||||
|
#[cfg(feature = "postgres")]
|
||||||
|
mod resolver_tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn make_test_workspace_for_user(user_id: &str) -> Arc<Workspace> {
|
||||||
|
Arc::new(Workspace::new(
|
||||||
|
user_id,
|
||||||
|
deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new(
|
||||||
|
tokio_postgres::Config::new(),
|
||||||
|
tokio_postgres::NoTls,
|
||||||
|
))
|
||||||
|
.build()
|
||||||
|
.unwrap(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_fixed_workspace_resolver_ignores_user_id() {
|
||||||
|
let ws = make_test_workspace_for_user("alice");
|
||||||
|
let resolver = FixedWorkspaceResolver::new(Arc::clone(&ws));
|
||||||
|
|
||||||
|
let ws_alice = resolver.resolve("alice").await;
|
||||||
|
let ws_bob = resolver.resolve("bob").await;
|
||||||
|
|
||||||
|
// Both should return the exact same Arc (pointer equality)
|
||||||
|
assert!(Arc::ptr_eq(&ws_alice, &ws_bob));
|
||||||
|
assert_eq!(ws_alice.user_id(), "alice");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tracking resolver that records which user_ids were requested.
|
||||||
|
struct TrackingWorkspaceResolver {
|
||||||
|
inner: FixedWorkspaceResolver,
|
||||||
|
resolved_users: std::sync::Mutex<Vec<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TrackingWorkspaceResolver {
|
||||||
|
fn new(workspace: Arc<Workspace>) -> Self {
|
||||||
|
Self {
|
||||||
|
inner: FixedWorkspaceResolver::new(workspace),
|
||||||
|
resolved_users: std::sync::Mutex::new(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolved_users(&self) -> Vec<String> {
|
||||||
|
self.resolved_users.lock().unwrap().clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl WorkspaceResolver for TrackingWorkspaceResolver {
|
||||||
|
async fn resolve(&self, user_id: &str) -> Arc<Workspace> {
|
||||||
|
self.resolved_users
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.push(user_id.to_string());
|
||||||
|
self.inner.resolve(user_id).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_memory_search_uses_job_context_user_id() {
|
||||||
|
let ws = make_test_workspace_for_user("default");
|
||||||
|
let tracker = Arc::new(TrackingWorkspaceResolver::new(ws));
|
||||||
|
let tool = MemorySearchTool::new(tracker.clone() as Arc<dyn WorkspaceResolver>);
|
||||||
|
|
||||||
|
// Execute with user_id "alice"
|
||||||
|
let ctx_alice = JobContext::with_user("alice", "test", "test");
|
||||||
|
let params = serde_json::json!({"query": "test"});
|
||||||
|
// The search will fail (no real DB) but we only care about resolver call
|
||||||
|
let _ = tool.execute(params, &ctx_alice).await;
|
||||||
|
|
||||||
|
// Execute with user_id "bob"
|
||||||
|
let ctx_bob = JobContext::with_user("bob", "test", "test");
|
||||||
|
let params = serde_json::json!({"query": "test"});
|
||||||
|
let _ = tool.execute(params, &ctx_bob).await;
|
||||||
|
|
||||||
|
let resolved = tracker.resolved_users();
|
||||||
|
assert_eq!(resolved, vec!["alice", "bob"]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_memory_write_uses_job_context_user_id() {
|
||||||
|
let ws = make_test_workspace_for_user("default");
|
||||||
|
let tracker = Arc::new(TrackingWorkspaceResolver::new(ws));
|
||||||
|
let tool = MemoryWriteTool::new(tracker.clone() as Arc<dyn WorkspaceResolver>);
|
||||||
|
|
||||||
|
// Execute with user_id "alice"
|
||||||
|
let ctx_alice = JobContext::with_user("alice", "test", "test");
|
||||||
|
let params = serde_json::json!({
|
||||||
|
"content": "remember this",
|
||||||
|
"target": "daily_log",
|
||||||
|
});
|
||||||
|
let _ = tool.execute(params, &ctx_alice).await;
|
||||||
|
|
||||||
|
// Execute with user_id "bob"
|
||||||
|
let ctx_bob = JobContext::with_user("bob", "test", "test");
|
||||||
|
let params = serde_json::json!({
|
||||||
|
"content": "remember that",
|
||||||
|
"target": "daily_log",
|
||||||
|
});
|
||||||
|
let _ = tool.execute(params, &ctx_bob).await;
|
||||||
|
|
||||||
|
let resolved = tracker.resolved_users();
|
||||||
|
assert_eq!(resolved, vec!["alice", "bob"]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "libsql")]
|
||||||
|
mod per_user_resolver_tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
async fn make_test_db() -> Arc<dyn crate::db::Database> {
|
||||||
|
use crate::db::libsql::LibSqlBackend;
|
||||||
|
let temp_dir = tempfile::tempdir().expect("tempdir");
|
||||||
|
let db_path = temp_dir.path().join("resolver_test.db");
|
||||||
|
let backend = LibSqlBackend::new_local(&db_path)
|
||||||
|
.await
|
||||||
|
.expect("LibSqlBackend");
|
||||||
|
<LibSqlBackend as crate::db::Database>::run_migrations(&backend)
|
||||||
|
.await
|
||||||
|
.expect("migrations");
|
||||||
|
// Leak the tempdir so it outlives the test (cleaned up on process exit).
|
||||||
|
std::mem::forget(temp_dir);
|
||||||
|
Arc::new(backend)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_resolver_returns_different_workspaces() {
|
||||||
|
let db = make_test_db().await;
|
||||||
|
|
||||||
|
let pool = crate::channels::web::server::WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
crate::workspace::EmbeddingCacheConfig::default(),
|
||||||
|
crate::config::WorkspaceSearchConfig::default(),
|
||||||
|
crate::config::WorkspaceConfig::default(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let ws_alice = pool.resolve("alice").await;
|
||||||
|
let ws_bob = pool.resolve("bob").await;
|
||||||
|
|
||||||
|
// Different user IDs should get different workspaces
|
||||||
|
assert_eq!(ws_alice.user_id(), "alice");
|
||||||
|
assert_eq!(ws_bob.user_id(), "bob");
|
||||||
|
assert!(!Arc::ptr_eq(&ws_alice, &ws_bob));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_workspace_pool_resolver_caches_workspace() {
|
||||||
|
let db = make_test_db().await;
|
||||||
|
|
||||||
|
let pool = crate::channels::web::server::WorkspacePool::new(
|
||||||
|
db,
|
||||||
|
None,
|
||||||
|
crate::workspace::EmbeddingCacheConfig::default(),
|
||||||
|
crate::config::WorkspaceSearchConfig::default(),
|
||||||
|
crate::config::WorkspaceConfig::default(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let ws1 = pool.resolve("alice").await;
|
||||||
|
let ws2 = pool.resolve("alice").await;
|
||||||
|
|
||||||
|
// Same user_id should return the same cached Arc (pointer equality)
|
||||||
|
assert!(Arc::ptr_eq(&ws1, &ws2));
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ mod file;
|
|||||||
mod http;
|
mod http;
|
||||||
mod job;
|
mod job;
|
||||||
mod json;
|
mod json;
|
||||||
mod memory;
|
pub mod memory;
|
||||||
mod message;
|
mod message;
|
||||||
pub mod path_utils;
|
pub mod path_utils;
|
||||||
mod restart;
|
mod restart;
|
||||||
|
|||||||
@@ -140,7 +140,8 @@ fn execution_properties() -> Value {
|
|||||||
},
|
},
|
||||||
"use_tools": {
|
"use_tools": {
|
||||||
"type": "boolean",
|
"type": "boolean",
|
||||||
"description": "Only applies to lightweight mode. When true, safe non-approval tools are available."
|
"default": true,
|
||||||
|
"description": "Only applies to lightweight mode. New lightweight routines default this to true; when enabled, the routine can use the owner's live autonomous tool scope."
|
||||||
},
|
},
|
||||||
"max_tool_rounds": {
|
"max_tool_rounds": {
|
||||||
"type": "integer",
|
"type": "integer",
|
||||||
@@ -290,7 +291,7 @@ fn routine_request_discovery_schema() -> Value {
|
|||||||
fn lightweight_execution_variant() -> Value {
|
fn lightweight_execution_variant() -> Value {
|
||||||
serde_json::json!({
|
serde_json::json!({
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"description": "Default lightweight execution. Applies when execution is omitted or execution.mode='lightweight'.",
|
"description": "Default lightweight execution. Applies when execution is omitted or execution.mode='lightweight'. New lightweight routines default to tools enabled unless execution.use_tools=false is set.",
|
||||||
"properties": {
|
"properties": {
|
||||||
"mode": {
|
"mode": {
|
||||||
"type": "string",
|
"type": "string",
|
||||||
@@ -304,7 +305,8 @@ fn lightweight_execution_variant() -> Value {
|
|||||||
},
|
},
|
||||||
"use_tools": {
|
"use_tools": {
|
||||||
"type": "boolean",
|
"type": "boolean",
|
||||||
"description": "When true, safe non-approval tools are available."
|
"default": true,
|
||||||
|
"description": "Defaults to true for new lightweight routines. When enabled, the routine can use the owner's live autonomous tool scope."
|
||||||
},
|
},
|
||||||
"max_tool_rounds": {
|
"max_tool_rounds": {
|
||||||
"type": "integer",
|
"type": "integer",
|
||||||
@@ -335,7 +337,7 @@ fn full_job_execution_variant() -> Value {
|
|||||||
fn execution_discovery_schema() -> Value {
|
fn execution_discovery_schema() -> Value {
|
||||||
serde_json::json!({
|
serde_json::json!({
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"description": "Optional execution settings. Omit this block for the default lightweight mode.",
|
"description": "Optional execution settings. Omit this block for the default lightweight mode with tools enabled.",
|
||||||
"properties": execution_properties(),
|
"properties": execution_properties(),
|
||||||
"oneOf": [
|
"oneOf": [
|
||||||
lightweight_execution_variant(),
|
lightweight_execution_variant(),
|
||||||
@@ -408,7 +410,8 @@ fn routine_create_tool_summary() -> ToolDiscoverySummary {
|
|||||||
"execution.mode='full_job' uses the owner's live autonomous tool scope and ignores use_tools, max_tool_rounds, and context_paths.".into(),
|
"execution.mode='full_job' uses the owner's live autonomous tool scope and ignores use_tools, max_tool_rounds, and context_paths.".into(),
|
||||||
],
|
],
|
||||||
notes: vec![
|
notes: vec![
|
||||||
"Omitting execution defaults to lightweight mode.".into(),
|
"Omitting execution defaults to lightweight mode with tools enabled.".into(),
|
||||||
|
"Set execution.use_tools=false to keep a new lightweight routine text-only.".into(),
|
||||||
"Omitting delivery.user falls back to the owner's last-seen notification target.".into(),
|
"Omitting delivery.user falls back to the owner's last-seen notification target.".into(),
|
||||||
"advanced.cooldown_secs defaults to 300.".into(),
|
"advanced.cooldown_secs defaults to 300.".into(),
|
||||||
"Legacy flat aliases are still accepted for compatibility, but grouped fields are preferred.".into(),
|
"Legacy flat aliases are still accepted for compatibility, but grouped fields are preferred.".into(),
|
||||||
@@ -605,7 +608,8 @@ fn routine_create_schema(include_compatibility_aliases: bool) -> Value {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn routine_create_parameters_schema() -> Value {
|
pub(crate) fn routine_create_parameters_schema() -> Value {
|
||||||
routine_create_schema(false)
|
static CACHE: OnceLock<Value> = OnceLock::new();
|
||||||
|
CACHE.get_or_init(|| routine_create_schema(false)).clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn routine_create_discovery_schema() -> Value {
|
fn routine_create_discovery_schema() -> Value {
|
||||||
@@ -852,11 +856,15 @@ fn parse_execution_mode(value: Option<String>) -> Result<NormalizedExecutionMode
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn parse_routine_execution(params: &Value) -> Result<NormalizedExecutionRequest, ToolError> {
|
fn parse_routine_execution(
|
||||||
|
params: &Value,
|
||||||
|
default_use_tools: bool,
|
||||||
|
) -> Result<NormalizedExecutionRequest, ToolError> {
|
||||||
let mode = parse_execution_mode(string_field(params, "execution", "mode", &["action_type"]))?;
|
let mode = parse_execution_mode(string_field(params, "execution", "mode", &["action_type"]))?;
|
||||||
let context_paths =
|
let context_paths =
|
||||||
string_array_field(params, "execution", "context_paths", &["context_paths"]);
|
string_array_field(params, "execution", "context_paths", &["context_paths"]);
|
||||||
let use_tools = bool_field(params, "execution", "use_tools", &["use_tools"]).unwrap_or(false);
|
let use_tools =
|
||||||
|
bool_field(params, "execution", "use_tools", &["use_tools"]).unwrap_or(default_use_tools);
|
||||||
let max_tool_rounds = u64_field(params, "execution", "max_tool_rounds", &["max_tool_rounds"])
|
let max_tool_rounds = u64_field(params, "execution", "max_tool_rounds", &["max_tool_rounds"])
|
||||||
.unwrap_or(3)
|
.unwrap_or(3)
|
||||||
.clamp(1, crate::agent::routine::MAX_TOOL_ROUNDS_LIMIT as u64)
|
.clamp(1, crate::agent::routine::MAX_TOOL_ROUNDS_LIMIT as u64)
|
||||||
@@ -888,7 +896,7 @@ fn parse_routine_create_request(
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string();
|
.to_string();
|
||||||
let trigger = parse_routine_trigger(params)?;
|
let trigger = parse_routine_trigger(params)?;
|
||||||
let execution = parse_routine_execution(params)?;
|
let execution = parse_routine_execution(params, true)?;
|
||||||
let delivery = parse_routine_delivery(params);
|
let delivery = parse_routine_delivery(params);
|
||||||
let cooldown_secs =
|
let cooldown_secs =
|
||||||
u64_field(params, "advanced", "cooldown_secs", &["cooldown_secs"]).unwrap_or(300);
|
u64_field(params, "advanced", "cooldown_secs", &["cooldown_secs"]).unwrap_or(300);
|
||||||
@@ -1007,7 +1015,8 @@ fn event_emit_schema(include_source_alias: bool) -> Value {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn event_emit_parameters_schema() -> Value {
|
pub(crate) fn event_emit_parameters_schema() -> Value {
|
||||||
event_emit_schema(false)
|
static CACHE: OnceLock<Value> = OnceLock::new();
|
||||||
|
CACHE.get_or_init(|| event_emit_schema(false)).clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn event_emit_discovery_schema() -> Value {
|
fn event_emit_discovery_schema() -> Value {
|
||||||
@@ -1863,6 +1872,56 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_lightweight_create_with_tools_enabled_by_default() {
|
||||||
|
let params = serde_json::json!({
|
||||||
|
"name": "manual-check",
|
||||||
|
"prompt": "Inspect the repo for issues.",
|
||||||
|
"request": {
|
||||||
|
"kind": "manual"
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let parsed = parse_routine_create_request(¶ms).expect("parse default lightweight");
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
matches!(parsed.execution.mode, NormalizedExecutionMode::Lightweight),
|
||||||
|
"expected lightweight execution mode",
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
parsed.execution.use_tools,
|
||||||
|
"new lightweight routines should default use_tools=true",
|
||||||
|
);
|
||||||
|
assert_eq!(parsed.execution.max_tool_rounds, 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_lightweight_create_with_explicit_tools_disabled() {
|
||||||
|
let params = serde_json::json!({
|
||||||
|
"name": "manual-check",
|
||||||
|
"prompt": "Inspect the repo for issues.",
|
||||||
|
"request": {
|
||||||
|
"kind": "manual"
|
||||||
|
},
|
||||||
|
"execution": {
|
||||||
|
"use_tools": false
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let parsed =
|
||||||
|
parse_routine_create_request(¶ms).expect("parse lightweight with tools disabled");
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
matches!(parsed.execution.mode, NormalizedExecutionMode::Lightweight),
|
||||||
|
"expected lightweight execution mode",
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!parsed.execution.use_tools,
|
||||||
|
"explicit use_tools=false should be preserved",
|
||||||
|
);
|
||||||
|
assert_eq!(parsed.execution.max_tool_rounds, 3);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parses_context_paths_with_trim_drop_empty_and_stable_dedupe() {
|
fn parses_context_paths_with_trim_drop_empty_and_stable_dedupe() {
|
||||||
let params = serde_json::json!({
|
let params = serde_json::json!({
|
||||||
@@ -2201,6 +2260,20 @@ mod tests {
|
|||||||
.any(|rule| rule.contains("request.kind='cron'")),
|
.any(|rule| rule.contains("request.kind='cron'")),
|
||||||
"summary should explain cron requirement",
|
"summary should explain cron requirement",
|
||||||
);
|
);
|
||||||
|
assert!(
|
||||||
|
summary
|
||||||
|
.notes
|
||||||
|
.iter()
|
||||||
|
.any(|note| note.contains("lightweight mode with tools enabled")),
|
||||||
|
"summary should mention the new lightweight default",
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
summary
|
||||||
|
.notes
|
||||||
|
.iter()
|
||||||
|
.any(|note| note.contains("execution.use_tools=false")),
|
||||||
|
"summary should mention the text-only opt-out",
|
||||||
|
);
|
||||||
assert!(
|
assert!(
|
||||||
summary
|
summary
|
||||||
.notes
|
.notes
|
||||||
|
|||||||
+28
-6
@@ -334,15 +334,37 @@ impl ToolRegistry {
|
|||||||
tracing::debug!("Registered 5 development tools");
|
tracing::debug!("Registered 5 development tools");
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Register memory tools with a workspace.
|
/// Register memory tools with a workspace resolver.
|
||||||
|
///
|
||||||
|
/// Memory tools require a workspace resolver for persistence. Call this after
|
||||||
|
/// `register_builtin_tools()` if you have a workspace available.
|
||||||
|
pub fn register_memory_tools_with_resolver(
|
||||||
|
&self,
|
||||||
|
resolver: Arc<dyn crate::tools::builtin::memory::WorkspaceResolver>,
|
||||||
|
) {
|
||||||
|
self.register_sync(Arc::new(MemorySearchTool::new(Arc::clone(&resolver))));
|
||||||
|
self.register_sync(Arc::new(MemoryWriteTool::new(Arc::clone(&resolver))));
|
||||||
|
self.register_sync(Arc::new(MemoryReadTool::new(Arc::clone(&resolver))));
|
||||||
|
self.register_sync(Arc::new(MemoryTreeTool::new(resolver)));
|
||||||
|
|
||||||
|
tracing::debug!("Registered 4 memory tools");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register memory tools with a fixed workspace (backward compatibility).
|
||||||
///
|
///
|
||||||
/// Memory tools require a workspace for persistence. Call this after
|
/// Memory tools require a workspace for persistence. Call this after
|
||||||
/// `register_builtin_tools()` if you have a workspace available.
|
/// `register_builtin_tools()` if you have a workspace available.
|
||||||
pub fn register_memory_tools(&self, workspace: Arc<Workspace>) {
|
pub fn register_memory_tools(&self, workspace: Arc<Workspace>) {
|
||||||
self.register_sync(Arc::new(MemorySearchTool::new(Arc::clone(&workspace))));
|
self.register_sync(Arc::new(MemorySearchTool::from_workspace(Arc::clone(
|
||||||
self.register_sync(Arc::new(MemoryWriteTool::new(Arc::clone(&workspace))));
|
&workspace,
|
||||||
self.register_sync(Arc::new(MemoryReadTool::new(Arc::clone(&workspace))));
|
))));
|
||||||
self.register_sync(Arc::new(MemoryTreeTool::new(workspace)));
|
self.register_sync(Arc::new(MemoryWriteTool::from_workspace(Arc::clone(
|
||||||
|
&workspace,
|
||||||
|
))));
|
||||||
|
self.register_sync(Arc::new(MemoryReadTool::from_workspace(Arc::clone(
|
||||||
|
&workspace,
|
||||||
|
))));
|
||||||
|
self.register_sync(Arc::new(MemoryTreeTool::from_workspace(workspace)));
|
||||||
|
|
||||||
tracing::debug!("Registered 4 memory tools");
|
tracing::debug!("Registered 4 memory tools");
|
||||||
}
|
}
|
||||||
@@ -361,7 +383,7 @@ impl ToolRegistry {
|
|||||||
job_manager: Option<Arc<ContainerJobManager>>,
|
job_manager: Option<Arc<ContainerJobManager>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<Arc<dyn Database>>,
|
||||||
job_event_tx: Option<
|
job_event_tx: Option<
|
||||||
tokio::sync::broadcast::Sender<(uuid::Uuid, crate::channels::web::types::SseEvent)>,
|
tokio::sync::broadcast::Sender<(uuid::Uuid, String, ironclaw_common::AppEvent)>,
|
||||||
>,
|
>,
|
||||||
inject_tx: Option<tokio::sync::mpsc::Sender<crate::channels::IncomingMessage>>,
|
inject_tx: Option<tokio::sync::mpsc::Sender<crate::channels::IncomingMessage>>,
|
||||||
prompt_queue: Option<PromptQueue>,
|
prompt_queue: Option<PromptQueue>,
|
||||||
|
|||||||
@@ -47,12 +47,6 @@ pub struct CapabilitiesFile {
|
|||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub description: Option<String>,
|
pub description: Option<String>,
|
||||||
|
|
||||||
/// JSON Schema for the tool's input parameters.
|
|
||||||
/// Used as the `Tool::parameters_schema()` return value.
|
|
||||||
/// If omitted, a permissive fallback is used (with a warning).
|
|
||||||
#[serde(default)]
|
|
||||||
pub parameters: Option<serde_json::Value>,
|
|
||||||
|
|
||||||
/// Extension version (semver).
|
/// Extension version (semver).
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub version: Option<String>,
|
pub version: Option<String>,
|
||||||
@@ -103,9 +97,6 @@ pub struct CapabilitiesFile {
|
|||||||
|
|
||||||
/// Maximum length for the description field to prevent memory abuse.
|
/// Maximum length for the description field to prevent memory abuse.
|
||||||
const MAX_DESCRIPTION_CHARS: usize = 4096;
|
const MAX_DESCRIPTION_CHARS: usize = 4096;
|
||||||
/// Maximum serialized size of the parameters schema JSON.
|
|
||||||
const MAX_PARAMETERS_SCHEMA_BYTES: usize = 64 * 1024;
|
|
||||||
|
|
||||||
impl CapabilitiesFile {
|
impl CapabilitiesFile {
|
||||||
/// Parse from JSON string.
|
/// Parse from JSON string.
|
||||||
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
|
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
|
||||||
@@ -135,18 +126,6 @@ impl CapabilitiesFile {
|
|||||||
);
|
);
|
||||||
self.description = Some(truncated.to_string());
|
self.description = Some(truncated.to_string());
|
||||||
}
|
}
|
||||||
// Drop oversized parameters schema (issue #977)
|
|
||||||
if let Some(ref params) = self.parameters {
|
|
||||||
let size = params.to_string().len();
|
|
||||||
if size > MAX_PARAMETERS_SCHEMA_BYTES {
|
|
||||||
tracing::warn!(
|
|
||||||
"Capabilities parameters schema dropped ({} bytes exceeds {} limit)",
|
|
||||||
size,
|
|
||||||
MAX_PARAMETERS_SCHEMA_BYTES,
|
|
||||||
);
|
|
||||||
self.parameters = None;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Merge nested `capabilities` wrapper into top-level fields.
|
/// Merge nested `capabilities` wrapper into top-level fields.
|
||||||
@@ -171,7 +150,6 @@ impl CapabilitiesFile {
|
|||||||
if let Some(inner) = self.capabilities.take() {
|
if let Some(inner) = self.capabilities.take() {
|
||||||
let inner = inner.resolve_nested_inner(depth + 1);
|
let inner = inner.resolve_nested_inner(depth + 1);
|
||||||
self.description = self.description.or(inner.description);
|
self.description = self.description.or(inner.description);
|
||||||
self.parameters = self.parameters.or(inner.parameters);
|
|
||||||
self.http = self.http.or(inner.http);
|
self.http = self.http.or(inner.http);
|
||||||
self.secrets = self.secrets.or(inner.secrets);
|
self.secrets = self.secrets.or(inner.secrets);
|
||||||
self.tool_invoke = self.tool_invoke.or(inner.tool_invoke);
|
self.tool_invoke = self.tool_invoke.or(inner.tool_invoke);
|
||||||
@@ -1424,26 +1402,12 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Tool description and parameters schema ──────────────────────────
|
// ── Tool description ────────────────────────────────────────────────
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_parse_description_and_parameters() {
|
fn test_parse_description() {
|
||||||
let json = r#"{
|
let json = r#"{
|
||||||
"description": "Search the web using Brave Search API",
|
"description": "Search the web using Brave Search API"
|
||||||
"parameters": {
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"query": {
|
|
||||||
"type": "string",
|
|
||||||
"description": "Search query"
|
|
||||||
},
|
|
||||||
"count": {
|
|
||||||
"type": "integer",
|
|
||||||
"description": "Number of results"
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"required": ["query"]
|
|
||||||
}
|
|
||||||
}"#;
|
}"#;
|
||||||
|
|
||||||
let caps = CapabilitiesFile::from_json(json).unwrap();
|
let caps = CapabilitiesFile::from_json(json).unwrap();
|
||||||
@@ -1451,28 +1415,10 @@ mod tests {
|
|||||||
caps.description.as_deref(),
|
caps.description.as_deref(),
|
||||||
Some("Search the web using Brave Search API")
|
Some("Search the web using Brave Search API")
|
||||||
);
|
);
|
||||||
let params = caps.parameters.unwrap();
|
|
||||||
assert_eq!(params["type"], "object");
|
|
||||||
assert!(params["properties"]["query"].is_object());
|
|
||||||
assert_eq!(params["required"][0], "query");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_parse_description_only() {
|
fn test_parse_without_description() {
|
||||||
let json = r#"{
|
|
||||||
"description": "A tool without explicit parameters schema"
|
|
||||||
}"#;
|
|
||||||
|
|
||||||
let caps = CapabilitiesFile::from_json(json).unwrap();
|
|
||||||
assert_eq!(
|
|
||||||
caps.description.as_deref(),
|
|
||||||
Some("A tool without explicit parameters schema")
|
|
||||||
);
|
|
||||||
assert!(caps.parameters.is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_parse_without_description_or_parameters() {
|
|
||||||
let json = r#"{
|
let json = r#"{
|
||||||
"http": {
|
"http": {
|
||||||
"allowlist": [{ "host": "api.example.com" }]
|
"allowlist": [{ "host": "api.example.com" }]
|
||||||
@@ -1484,24 +1430,28 @@ mod tests {
|
|||||||
caps.description.is_none(),
|
caps.description.is_none(),
|
||||||
"description should be None when not provided"
|
"description should be None when not provided"
|
||||||
);
|
);
|
||||||
assert!(
|
}
|
||||||
caps.parameters.is_none(),
|
|
||||||
"parameters should be None when not provided"
|
#[test]
|
||||||
);
|
fn test_parameters_field_silently_ignored() {
|
||||||
|
// Backward compat: old capabilities files with "parameters" still parse.
|
||||||
|
let json = r#"{
|
||||||
|
"description": "A tool",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": { "action": { "type": "string" } }
|
||||||
|
}
|
||||||
|
}"#;
|
||||||
|
|
||||||
|
let caps = CapabilitiesFile::from_json(json).unwrap();
|
||||||
|
assert_eq!(caps.description.as_deref(), Some("A tool"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_resolve_nested_description_promoted() {
|
fn test_resolve_nested_description_promoted() {
|
||||||
let json = r#"{
|
let json = r#"{
|
||||||
"capabilities": {
|
"capabilities": {
|
||||||
"description": "Inner tool description",
|
"description": "Inner tool description"
|
||||||
"parameters": {
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"input": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["input"]
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}"#;
|
}"#;
|
||||||
|
|
||||||
@@ -1511,10 +1461,6 @@ mod tests {
|
|||||||
Some("Inner tool description"),
|
Some("Inner tool description"),
|
||||||
"description should be promoted from inner capabilities"
|
"description should be promoted from inner capabilities"
|
||||||
);
|
);
|
||||||
assert!(
|
|
||||||
caps.parameters.is_some(),
|
|
||||||
"parameters should be promoted from inner capabilities"
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1564,32 +1510,4 @@ mod tests {
|
|||||||
desc.len()
|
desc.len()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Regression test for issue #977: oversized parameters schema is dropped.
|
|
||||||
#[test]
|
|
||||||
fn test_oversized_parameters_schema_dropped() {
|
|
||||||
// Build a parameters schema larger than MAX_PARAMETERS_SCHEMA_BYTES
|
|
||||||
let mut properties = serde_json::Map::new();
|
|
||||||
for i in 0..2000 {
|
|
||||||
properties.insert(
|
|
||||||
format!("field_{i}"),
|
|
||||||
serde_json::json!({
|
|
||||||
"type": "string",
|
|
||||||
"description": "x".repeat(50)
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": properties,
|
|
||||||
});
|
|
||||||
let json = serde_json::json!({
|
|
||||||
"parameters": schema,
|
|
||||||
});
|
|
||||||
let caps = CapabilitiesFile::from_json(&json.to_string()).unwrap();
|
|
||||||
assert!(
|
|
||||||
caps.parameters.is_none(),
|
|
||||||
"oversized parameters schema should be dropped"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
+177
-58
@@ -123,73 +123,51 @@ impl WasmToolLoader {
|
|||||||
}
|
}
|
||||||
let wasm_bytes = fs::read(wasm_path).await?;
|
let wasm_bytes = fs::read(wasm_path).await?;
|
||||||
|
|
||||||
// Read capabilities (optional) and extract OAuth refresh config,
|
// Read capabilities (optional) and extract OAuth refresh config
|
||||||
// tool description, and parameter schema.
|
// and tool description. Parameter schema is auto-derived from the
|
||||||
let (capabilities, oauth_refresh, description, schema) =
|
// WASM module's schema() export (see WasmToolSchemas::compact_schema).
|
||||||
if let Some(cap_path) = capabilities_path {
|
let (capabilities, oauth_refresh, description) = if let Some(cap_path) = capabilities_path {
|
||||||
if cap_path.exists() {
|
if cap_path.exists() {
|
||||||
let cap_bytes = fs::read(cap_path).await?;
|
let cap_bytes = fs::read(cap_path).await?;
|
||||||
let cap_file = CapabilitiesFile::from_bytes(&cap_bytes)
|
let cap_file = CapabilitiesFile::from_bytes(&cap_bytes)
|
||||||
.map_err(|e| WasmLoadError::InvalidCapabilities(e.to_string()))?;
|
.map_err(|e| WasmLoadError::InvalidCapabilities(e.to_string()))?;
|
||||||
cap_file.validate(name);
|
cap_file.validate(name);
|
||||||
|
|
||||||
// Check WIT version compatibility
|
// Check WIT version compatibility
|
||||||
check_wit_version_compat(
|
check_wit_version_compat(
|
||||||
name,
|
name,
|
||||||
cap_file.wit_version.as_deref(),
|
cap_file.wit_version.as_deref(),
|
||||||
crate::tools::wasm::WIT_TOOL_VERSION,
|
crate::tools::wasm::WIT_TOOL_VERSION,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
let caps = cap_file.to_capabilities();
|
let caps = cap_file.to_capabilities();
|
||||||
let oauth = resolve_oauth_refresh_config(&cap_file);
|
let oauth = resolve_oauth_refresh_config(&cap_file);
|
||||||
let desc = cap_file.description.clone();
|
let desc = cap_file.description.clone();
|
||||||
// Validate parameters schema before accepting it.
|
if desc.is_none() {
|
||||||
let params = cap_file.parameters.clone().and_then(|p| {
|
|
||||||
let errors = crate::tools::validate_tool_schema(&p, name);
|
|
||||||
if errors.is_empty() {
|
|
||||||
Some(p)
|
|
||||||
} else {
|
|
||||||
tracing::warn!(
|
|
||||||
tool = name,
|
|
||||||
?errors,
|
|
||||||
"Invalid parameters schema in capabilities.json, \
|
|
||||||
using permissive fallback"
|
|
||||||
);
|
|
||||||
None
|
|
||||||
}
|
|
||||||
});
|
|
||||||
if desc.is_none() {
|
|
||||||
tracing::warn!(
|
|
||||||
tool = name,
|
|
||||||
path = %cap_path.display(),
|
|
||||||
"Capabilities file missing \"description\" field; \
|
|
||||||
tool will use generic fallback description"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
if params.is_none() && cap_file.parameters.is_none() {
|
|
||||||
tracing::warn!(
|
|
||||||
tool = name,
|
|
||||||
path = %cap_path.display(),
|
|
||||||
"Capabilities file missing \"parameters\" field; \
|
|
||||||
tool will accept any JSON object (permissive fallback)"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
(caps, oauth, desc, params)
|
|
||||||
} else {
|
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
|
tool = name,
|
||||||
path = %cap_path.display(),
|
path = %cap_path.display(),
|
||||||
"Capabilities file not found, using default (no permissions)"
|
"Capabilities file missing \"description\" field; \
|
||||||
|
tool will use generic fallback description"
|
||||||
);
|
);
|
||||||
(Capabilities::default(), None, None, None)
|
|
||||||
}
|
}
|
||||||
|
(caps, oauth, desc)
|
||||||
} else {
|
} else {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
tool = name,
|
tool = name,
|
||||||
"No capabilities file for WASM tool; \
|
path = %cap_path.display(),
|
||||||
tool will use generic fallback description and accept any JSON object"
|
"Capabilities file not found, using default (no permissions)"
|
||||||
);
|
);
|
||||||
(Capabilities::default(), None, None, None)
|
(Capabilities::default(), None, None)
|
||||||
};
|
}
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
tool = name,
|
||||||
|
"No capabilities file for WASM tool; \
|
||||||
|
tool will use generic fallback description"
|
||||||
|
);
|
||||||
|
(Capabilities::default(), None, None)
|
||||||
|
};
|
||||||
|
|
||||||
// Register the tool
|
// Register the tool
|
||||||
self.registry
|
self.registry
|
||||||
@@ -200,7 +178,7 @@ impl WasmToolLoader {
|
|||||||
capabilities,
|
capabilities,
|
||||||
limits: None,
|
limits: None,
|
||||||
description: description.as_deref(),
|
description: description.as_deref(),
|
||||||
schema,
|
schema: None,
|
||||||
secrets_store: self.secrets_store.clone(),
|
secrets_store: self.secrets_store.clone(),
|
||||||
oauth_refresh,
|
oauth_refresh,
|
||||||
})
|
})
|
||||||
@@ -440,6 +418,7 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
|||||||
let oauth = auth.oauth.as_ref()?;
|
let oauth = auth.oauth.as_ref()?;
|
||||||
|
|
||||||
let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name);
|
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
|
let client_id = oauth
|
||||||
.client_id
|
.client_id
|
||||||
@@ -462,11 +441,21 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
|||||||
.and_then(|env| std::env::var(env).ok())
|
.and_then(|env| std::env::var(env).ok())
|
||||||
})
|
})
|
||||||
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
|
.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 {
|
Some(OAuthRefreshConfig {
|
||||||
token_url: oauth.token_url.clone(),
|
token_url: oauth.token_url.clone(),
|
||||||
client_id,
|
client_id,
|
||||||
client_secret,
|
client_secret,
|
||||||
|
exchange_proxy_url,
|
||||||
|
gateway_token,
|
||||||
secret_name: auth.secret_name.clone(),
|
secret_name: auth.secret_name.clone(),
|
||||||
provider: auth.provider.clone(),
|
provider: auth.provider.clone(),
|
||||||
})
|
})
|
||||||
@@ -733,9 +722,44 @@ mod tests {
|
|||||||
|
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
use crate::config::helpers::lock_env;
|
||||||
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
|
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
|
||||||
use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools};
|
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]
|
#[test]
|
||||||
fn wit_version_compat_none_is_ok() {
|
fn wit_version_compat_none_is_ok() {
|
||||||
// Pre-versioning extensions (no wit_version declared) should always pass
|
// Pre-versioning extensions (no wit_version declared) should always pass
|
||||||
@@ -893,6 +917,8 @@ mod tests {
|
|||||||
config.client_secret,
|
config.client_secret,
|
||||||
Some(TEST_OAUTH_CLIENT_SECRET.to_string())
|
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.secret_name, "google_oauth_token");
|
||||||
assert_eq!(config.provider, Some("google".to_string()));
|
assert_eq!(config.provider, Some("google".to_string()));
|
||||||
}
|
}
|
||||||
@@ -953,6 +979,10 @@ mod tests {
|
|||||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
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
|
// google_oauth_token should fall back to built-in credentials
|
||||||
let caps = CapabilitiesFile {
|
let caps = CapabilitiesFile {
|
||||||
auth: Some(AuthCapabilitySchema {
|
auth: Some(AuthCapabilitySchema {
|
||||||
@@ -974,6 +1004,95 @@ mod tests {
|
|||||||
let config = config.unwrap();
|
let config = config.unwrap();
|
||||||
assert!(!config.client_id.is_empty());
|
assert!(!config.client_id.is_empty());
|
||||||
assert!(config.client_secret.is_some());
|
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"));
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------
|
// ---------------------------------------------------------------
|
||||||
|
|||||||
+686
-80
@@ -19,7 +19,7 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
|||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
|
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
|
||||||
use crate::safety::LeakDetector;
|
use crate::safety::LeakDetector;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::{DecryptedSecret, SecretsStore};
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||||
use crate::tools::wasm::capabilities::Capabilities;
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::credential_injector::{
|
use crate::tools::wasm::credential_injector::{
|
||||||
@@ -44,6 +44,7 @@ wasmtime::component::bindgen!({
|
|||||||
});
|
});
|
||||||
|
|
||||||
// Alias the export interface types for convenience.
|
// Alias the export interface types for convenience.
|
||||||
|
use crate::cli::oauth_defaults;
|
||||||
use exports::near::agent::tool as wit_tool;
|
use exports::near::agent::tool as wit_tool;
|
||||||
|
|
||||||
/// Configuration needed to refresh an expired OAuth access token.
|
/// Configuration needed to refresh an expired OAuth access token.
|
||||||
@@ -59,6 +60,10 @@ pub struct OAuthRefreshConfig {
|
|||||||
pub client_id: String,
|
pub client_id: String,
|
||||||
/// OAuth client_secret (optional, some providers use PKCE without a secret).
|
/// OAuth client_secret (optional, some providers use PKCE without a secret).
|
||||||
pub client_secret: Option<String>,
|
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").
|
/// Secret name of the access token (e.g., "google_oauth_token").
|
||||||
/// The refresh token lives at `{secret_name}_refresh_token`.
|
/// The refresh token lives at `{secret_name}_refresh_token`.
|
||||||
pub secret_name: String,
|
pub secret_name: String,
|
||||||
@@ -656,12 +661,125 @@ impl WasmToolSchemas {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn new(discovery: serde_json::Value) -> Self {
|
fn new(discovery: serde_json::Value) -> Self {
|
||||||
|
let advertised = Self::compact_schema(&discovery);
|
||||||
Self {
|
Self {
|
||||||
advertised: Self::permissive_schema(),
|
advertised,
|
||||||
discovery,
|
discovery,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Derive a compact advertised schema from the full discovery schema.
|
||||||
|
///
|
||||||
|
/// Collects properties from top-level `properties` and from
|
||||||
|
/// `oneOf`/`anyOf`/`allOf` variants. Keeps only properties that are in
|
||||||
|
/// the top-level `required` array or carry an `enum`/`const` constraint.
|
||||||
|
/// For properties defined via `const` across multiple variants (e.g.
|
||||||
|
/// `"action": {"const": "get_repo"}` in each `oneOf` branch), the `const`
|
||||||
|
/// values are merged into a single `enum` array.
|
||||||
|
///
|
||||||
|
/// Variant-level `required` fields (e.g. `owner`, `repo` required within
|
||||||
|
/// each `oneOf` variant but not top-level) are intentionally omitted from
|
||||||
|
/// the compact schema — the LLM can discover them via
|
||||||
|
/// `tool_info(detail: "schema")`.
|
||||||
|
///
|
||||||
|
/// At most `MAX_COMPACT_PROPERTIES` properties are collected to bound
|
||||||
|
/// allocations from adversarial schemas.
|
||||||
|
fn compact_schema(discovery: &serde_json::Value) -> serde_json::Value {
|
||||||
|
const MAX_COMPACT_PROPERTIES: usize = 100;
|
||||||
|
|
||||||
|
let required: std::collections::HashSet<String> = discovery
|
||||||
|
.get("required")
|
||||||
|
.and_then(|r| r.as_array())
|
||||||
|
.map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|v| v.as_str().map(String::from))
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
// Collect properties from top-level and oneOf/anyOf/allOf variants.
|
||||||
|
// For properties with `const` across variants, merge into an `enum`.
|
||||||
|
let mut all_properties = serde_json::Map::new();
|
||||||
|
// Track const values per property to merge into enum.
|
||||||
|
let mut const_values: std::collections::HashMap<String, Vec<serde_json::Value>> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
|
|
||||||
|
if let Some(props) = discovery.get("properties").and_then(|p| p.as_object()) {
|
||||||
|
for (k, v) in props {
|
||||||
|
if all_properties.len() >= MAX_COMPACT_PROPERTIES {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
all_properties.insert(k.clone(), v.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for key in ["oneOf", "anyOf", "allOf"] {
|
||||||
|
if let Some(variants) = discovery.get(key).and_then(|v| v.as_array()) {
|
||||||
|
for variant in variants {
|
||||||
|
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
||||||
|
for (k, v) in props {
|
||||||
|
if all_properties.len() >= MAX_COMPACT_PROPERTIES
|
||||||
|
&& !all_properties.contains_key(k)
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Track const values for merging into enum.
|
||||||
|
if let Some(c) = v.get("const") {
|
||||||
|
const_values.entry(k.clone()).or_default().push(c.clone());
|
||||||
|
}
|
||||||
|
all_properties.entry(k.clone()).or_insert_with(|| v.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Merge collected const values into enum arrays.
|
||||||
|
for (name, values) in &const_values {
|
||||||
|
if values.len() > 1
|
||||||
|
&& let Some(prop) = all_properties.get_mut(name)
|
||||||
|
{
|
||||||
|
let mut merged = prop.clone();
|
||||||
|
if let Some(obj) = merged.as_object_mut() {
|
||||||
|
obj.remove("const");
|
||||||
|
obj.insert("enum".to_string(), serde_json::Value::Array(values.clone()));
|
||||||
|
}
|
||||||
|
*prop = merged;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if all_properties.is_empty() {
|
||||||
|
return Self::permissive_schema();
|
||||||
|
}
|
||||||
|
|
||||||
|
let kept: serde_json::Map<String, serde_json::Value> = all_properties
|
||||||
|
.into_iter()
|
||||||
|
.filter(|(name, prop)| {
|
||||||
|
required.contains(name) || prop.get("enum").is_some() || prop.get("const").is_some()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if kept.is_empty() {
|
||||||
|
return Self::permissive_schema();
|
||||||
|
}
|
||||||
|
|
||||||
|
let kept_required: Vec<serde_json::Value> = required
|
||||||
|
.iter()
|
||||||
|
.filter(|name| kept.contains_key(name.as_str()))
|
||||||
|
.map(|name| serde_json::Value::String(name.clone()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut result = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": kept,
|
||||||
|
"additionalProperties": true,
|
||||||
|
});
|
||||||
|
if !kept_required.is_empty() {
|
||||||
|
result["required"] = serde_json::Value::Array(kept_required);
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
fn with_override(&self, schema: serde_json::Value) -> Self {
|
fn with_override(&self, schema: serde_json::Value) -> Self {
|
||||||
Self {
|
Self {
|
||||||
advertised: schema.clone(),
|
advertised: schema.clone(),
|
||||||
@@ -1097,6 +1215,53 @@ async fn refresh_oauth_token(
|
|||||||
user_id: &str,
|
user_id: &str,
|
||||||
config: &OAuthRefreshConfig,
|
config: &OAuthRefreshConfig,
|
||||||
) -> bool {
|
) -> 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.
|
// SSRF defense: token_url comes from the tool's capabilities file.
|
||||||
if !config.token_url.starts_with("https://") {
|
if !config.token_url.starts_with("https://") {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -1114,19 +1279,6 @@ async fn refresh_oauth_token(
|
|||||||
return false;
|
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()
|
let client = match reqwest::Client::builder()
|
||||||
.timeout(Duration::from_secs(15))
|
.timeout(Duration::from_secs(15))
|
||||||
.redirect(reqwest::redirect::Policy::none())
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
@@ -1139,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![
|
let mut params = vec![
|
||||||
("grant_type", "refresh_token".to_string()),
|
("grant_type", "refresh_token".to_string()),
|
||||||
("refresh_token", refresh_secret.expose().to_string()),
|
("refresh_token", refresh_secret.expose().to_string()),
|
||||||
@@ -1174,22 +1330,55 @@ async fn refresh_oauth_token(
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
let token_response = match token_data.get("access_token").and_then(|v| v.as_str()) {
|
||||||
let new_access_token = match token_data.get("access_token").and_then(|v| v.as_str()) {
|
Some(access_token) => oauth_defaults::OAuthTokenResponse {
|
||||||
Some(t) => t,
|
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 => {
|
None => {
|
||||||
tracing::warn!("Token refresh response missing access_token field");
|
tracing::warn!("Token refresh response missing access_token field");
|
||||||
return false;
|
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 =
|
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 {
|
if let Some(ref provider) = config.provider {
|
||||||
access_params = access_params.with_provider(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);
|
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in as i64);
|
||||||
access_params = access_params.with_expiry(expires_at);
|
access_params = access_params.with_expiry(expires_at);
|
||||||
}
|
}
|
||||||
@@ -1199,10 +1388,8 @@ async fn refresh_oauth_token(
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Store rotated refresh token if the provider sent a new one
|
if let Some(new_refresh) = token_response.refresh_token.as_deref() {
|
||||||
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);
|
||||||
let mut refresh_params =
|
|
||||||
crate::secrets::CreateSecretParams::new(&refresh_name, new_refresh);
|
|
||||||
if let Some(ref provider) = config.provider {
|
if let Some(ref provider) = config.provider {
|
||||||
refresh_params = refresh_params.with_provider(provider);
|
refresh_params = refresh_params.with_provider(provider);
|
||||||
}
|
}
|
||||||
@@ -1551,9 +1738,18 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::net::SocketAddr;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
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 uuid::Uuid;
|
||||||
|
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
@@ -1643,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]
|
#[test]
|
||||||
fn test_wrapper_creation() {
|
fn test_wrapper_creation() {
|
||||||
// This test verifies the runtime can be created
|
// This test verifies the runtime can be created
|
||||||
@@ -1655,7 +1940,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_advertised_schema_stays_permissive_until_sidecar_override() {
|
async fn test_advertised_schema_auto_compacted_from_discovery() {
|
||||||
let discovery_schema = serde_json::json!({
|
let discovery_schema = serde_json::json!({
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
@@ -1675,42 +1960,7 @@ mod tests {
|
|||||||
wrapper.schemas = super::WasmToolSchemas::new(discovery_schema.clone());
|
wrapper.schemas = super::WasmToolSchemas::new(discovery_schema.clone());
|
||||||
wrapper.description = "Search documents".to_string();
|
wrapper.description = "Search documents".to_string();
|
||||||
|
|
||||||
// Advertised schema stays permissive; discovery holds the typed schema
|
// Advertised schema is auto-compacted: keeps required props, drops optional
|
||||||
assert_eq!(
|
|
||||||
wrapper.parameters_schema(),
|
|
||||||
serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {},
|
|
||||||
"additionalProperties": true
|
|
||||||
})
|
|
||||||
);
|
|
||||||
assert_eq!(wrapper.discovery_schema(), discovery_schema);
|
|
||||||
|
|
||||||
// Raw description is clean — no tool_info hint baked in
|
|
||||||
assert!(!wrapper.description().contains("tool_info"));
|
|
||||||
|
|
||||||
// But schema() composes the hint at display time when advertised is permissive
|
|
||||||
let schema = wrapper.schema();
|
|
||||||
assert!(
|
|
||||||
schema.description.contains("tool_info"),
|
|
||||||
"schema().description should contain tool_info hint: {}",
|
|
||||||
schema.description
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
schema.description.contains("include_schema: true"),
|
|
||||||
"hint should mention include_schema: true: {}",
|
|
||||||
schema.description
|
|
||||||
);
|
|
||||||
|
|
||||||
// After sidecar override, both schemas match and hint disappears
|
|
||||||
let wrapper = wrapper.with_schema(serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"query": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["query"]
|
|
||||||
}));
|
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
wrapper.parameters_schema(),
|
wrapper.parameters_schema(),
|
||||||
serde_json::json!({
|
serde_json::json!({
|
||||||
@@ -1718,20 +1968,143 @@ mod tests {
|
|||||||
"properties": {
|
"properties": {
|
||||||
"query": { "type": "string" }
|
"query": { "type": "string" }
|
||||||
},
|
},
|
||||||
"required": ["query"]
|
"required": ["query"],
|
||||||
|
"additionalProperties": true
|
||||||
})
|
})
|
||||||
);
|
);
|
||||||
assert_eq!(wrapper.discovery_schema(), wrapper.parameters_schema());
|
// Discovery retains the full schema
|
||||||
|
assert_eq!(wrapper.discovery_schema(), discovery_schema);
|
||||||
|
|
||||||
// With typed schema, schema() should NOT include tool_info hint
|
// Compacted schema has typed properties, so no tool_info hint needed
|
||||||
let schema = wrapper.schema();
|
let schema = wrapper.schema();
|
||||||
assert!(
|
assert!(
|
||||||
!schema.description.contains("tool_info"),
|
!schema.description.contains("tool_info"),
|
||||||
"schema().description should not contain tool_info hint when typed: {}",
|
"schema().description should not contain tool_info hint when auto-compacted: {}",
|
||||||
schema.description
|
schema.description
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compact_schema_keeps_required_and_enum_properties() {
|
||||||
|
let schema = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"action": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["list", "get", "create"],
|
||||||
|
"description": "The operation"
|
||||||
|
},
|
||||||
|
"query": { "type": "string" },
|
||||||
|
"limit": { "type": "integer" },
|
||||||
|
"format": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["json", "csv"]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["action"]
|
||||||
|
});
|
||||||
|
|
||||||
|
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||||
|
let props = compacted["properties"].as_object().unwrap();
|
||||||
|
|
||||||
|
// action: required + enum → kept
|
||||||
|
assert!(props.contains_key("action"));
|
||||||
|
// format: has enum → kept
|
||||||
|
assert!(props.contains_key("format"));
|
||||||
|
// query: not required, no enum → dropped
|
||||||
|
assert!(!props.contains_key("query"));
|
||||||
|
// limit: not required, no enum → dropped
|
||||||
|
assert!(!props.contains_key("limit"));
|
||||||
|
// additionalProperties lets the LLM still pass dropped props
|
||||||
|
assert_eq!(compacted["additionalProperties"], true);
|
||||||
|
assert_eq!(compacted["required"], serde_json::json!(["action"]));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compact_schema_falls_back_to_permissive_when_empty() {
|
||||||
|
// No required, no enum → permissive fallback
|
||||||
|
let schema = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": { "type": "string" },
|
||||||
|
"limit": { "type": "integer" }
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||||
|
assert!(compacted["properties"].as_object().unwrap().is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compact_schema_handles_no_properties() {
|
||||||
|
let schema = serde_json::json!({ "type": "object" });
|
||||||
|
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||||
|
assert!(compacted["properties"].as_object().unwrap().is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compact_schema_handles_oneof_variants() {
|
||||||
|
// GitHub-style schema: oneOf with no top-level properties, const per variant
|
||||||
|
let schema = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"required": ["action"],
|
||||||
|
"oneOf": [
|
||||||
|
{
|
||||||
|
"properties": {
|
||||||
|
"action": { "const": "get_repo" },
|
||||||
|
"owner": { "type": "string" },
|
||||||
|
"repo": { "type": "string" }
|
||||||
|
},
|
||||||
|
"required": ["action", "owner", "repo"]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"properties": {
|
||||||
|
"action": { "const": "list_issues" },
|
||||||
|
"owner": { "type": "string" },
|
||||||
|
"repo": { "type": "string" },
|
||||||
|
"state": { "type": "string", "enum": ["open", "closed", "all"] }
|
||||||
|
},
|
||||||
|
"required": ["action", "owner", "repo"]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
});
|
||||||
|
|
||||||
|
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||||
|
let props = compacted["properties"].as_object().unwrap();
|
||||||
|
|
||||||
|
// action: required + const values merged into enum → kept
|
||||||
|
let action = &props["action"];
|
||||||
|
assert!(
|
||||||
|
action.get("enum").is_some(),
|
||||||
|
"action const values should be merged into enum: {action}"
|
||||||
|
);
|
||||||
|
let action_enum = action["enum"].as_array().unwrap();
|
||||||
|
assert!(
|
||||||
|
action_enum.contains(&serde_json::json!("get_repo")),
|
||||||
|
"enum should contain get_repo"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
action_enum.contains(&serde_json::json!("list_issues")),
|
||||||
|
"enum should contain list_issues"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
action.get("const").is_none(),
|
||||||
|
"const should be removed after merging into enum"
|
||||||
|
);
|
||||||
|
|
||||||
|
// state: has enum → kept
|
||||||
|
assert!(
|
||||||
|
props.contains_key("state"),
|
||||||
|
"state should be kept (has enum)"
|
||||||
|
);
|
||||||
|
// owner/repo: not in top-level required, no enum → intentionally dropped
|
||||||
|
// (variant-level required is omitted; discoverable via tool_info)
|
||||||
|
assert!(!props.contains_key("owner"), "owner should be dropped");
|
||||||
|
assert!(!props.contains_key("repo"), "repo should be dropped");
|
||||||
|
assert_eq!(compacted["additionalProperties"], true);
|
||||||
|
assert_eq!(compacted["required"], serde_json::json!(["action"]));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_capabilities_default() {
|
fn test_capabilities_default() {
|
||||||
let caps = Capabilities::default();
|
let caps = Capabilities::default();
|
||||||
@@ -1893,8 +2266,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_bearer() {
|
async fn test_resolve_host_credentials_bearer() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{
|
use crate::secrets::{
|
||||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||||
};
|
};
|
||||||
@@ -1940,8 +2311,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_owner_scope_bearer() {
|
async fn test_resolve_host_credentials_owner_scope_bearer() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{
|
use crate::secrets::{
|
||||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||||
};
|
};
|
||||||
@@ -1987,8 +2356,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_execute_resolves_host_credentials_from_owner_scope_context() {
|
async fn test_execute_resolves_host_credentials_from_owner_scope_context() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||||
use crate::tools::wasm::capabilities::HttpCapability;
|
use crate::tools::wasm::capabilities::HttpCapability;
|
||||||
|
|
||||||
@@ -2038,8 +2405,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_missing_secret() {
|
async fn test_resolve_host_credentials_missing_secret() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||||
use crate::tools::wasm::capabilities::HttpCapability;
|
use crate::tools::wasm::capabilities::HttpCapability;
|
||||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||||
@@ -2071,8 +2436,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_skips_refresh_when_not_expired() {
|
async fn test_resolve_host_credentials_skips_refresh_when_not_expired() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{
|
use crate::secrets::{
|
||||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||||
};
|
};
|
||||||
@@ -2114,6 +2477,8 @@ mod tests {
|
|||||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||||
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
||||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.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(),
|
secret_name: "google_oauth_token".to_string(),
|
||||||
provider: Some("google".to_string()),
|
provider: Some("google".to_string()),
|
||||||
};
|
};
|
||||||
@@ -2130,8 +2495,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_skips_refresh_no_config() {
|
async fn test_resolve_host_credentials_skips_refresh_no_config() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{
|
use crate::secrets::{
|
||||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||||
};
|
};
|
||||||
@@ -2175,8 +2538,6 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_resolve_host_credentials_skips_refresh_no_expires_at() {
|
async fn test_resolve_host_credentials_skips_refresh_no_expires_at() {
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use crate::secrets::{
|
use crate::secrets::{
|
||||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||||
};
|
};
|
||||||
@@ -2216,6 +2577,8 @@ mod tests {
|
|||||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||||
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
||||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.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(),
|
secret_name: "google_oauth_token".to_string(),
|
||||||
provider: Some("google".to_string()),
|
provider: Some("google".to_string()),
|
||||||
};
|
};
|
||||||
@@ -2230,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]
|
#[test]
|
||||||
fn test_is_private_ip_v4() {
|
fn test_is_private_ip_v4() {
|
||||||
use std::net::IpAddr;
|
use std::net::IpAddr;
|
||||||
|
|||||||
@@ -111,10 +111,23 @@ impl Tunnel for CloudflareTunnel {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Drain stderr in the background to prevent SIGPIPE/buffer stalls.
|
if let Ok(mut guard) = self.url.write() {
|
||||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
*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 {
|
if let Some(stdout) = stdout {
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let mut out_reader = tokio::io::BufReader::new(stdout).lines();
|
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;
|
let mut guard = self.proc.lock().await;
|
||||||
*guard = Some(TunnelProcess { child });
|
*guard = Some(TunnelProcess {
|
||||||
|
child,
|
||||||
|
_pipe_drain: Some(drain_handle),
|
||||||
|
});
|
||||||
|
|
||||||
Ok(public_url)
|
Ok(public_url)
|
||||||
}
|
}
|
||||||
|
|||||||
+18
-5
@@ -73,6 +73,7 @@ impl Tunnel for CustomTunnel {
|
|||||||
let stderr = child.stderr.take();
|
let stderr = child.stderr.take();
|
||||||
|
|
||||||
let mut public_url = format!("http://{local_host}:{local_port}");
|
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()
|
if self.url_pattern.is_some()
|
||||||
&& let Some(stdout) = stdout
|
&& let Some(stdout) = stdout
|
||||||
@@ -103,17 +104,26 @@ impl Tunnel for CustomTunnel {
|
|||||||
Err(_) => {}
|
Err(_) => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Drain remaining stdout to prevent SIGPIPE/buffer stalls.
|
// We took ownership of the process's stdout pipe above to parse the
|
||||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
// 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 {
|
} 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 {
|
tokio::spawn(async move {
|
||||||
let mut reader = tokio::io::BufReader::new(stdout).lines();
|
let mut reader = tokio::io::BufReader::new(stdout).lines();
|
||||||
while let Ok(Some(_)) = reader.next_line().await {}
|
while let Ok(Some(_)) = reader.next_line().await {}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Drain stderr silently.
|
// Drain stderr to prevent SIGPIPE/buffer stalls.
|
||||||
if let Some(stderr) = stderr {
|
if let Some(stderr) = stderr {
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let mut reader = tokio::io::BufReader::new(stderr).lines();
|
let mut reader = tokio::io::BufReader::new(stderr).lines();
|
||||||
@@ -126,7 +136,10 @@ impl Tunnel for CustomTunnel {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let mut guard = self.proc.lock().await;
|
let mut guard = self.proc.lock().await;
|
||||||
*guard = Some(TunnelProcess { child });
|
*guard = Some(TunnelProcess {
|
||||||
|
child,
|
||||||
|
_pipe_drain: drain_handle,
|
||||||
|
});
|
||||||
|
|
||||||
Ok(public_url)
|
Ok(public_url)
|
||||||
}
|
}
|
||||||
|
|||||||
+126
-16
@@ -66,6 +66,10 @@ pub trait Tunnel: Send + Sync {
|
|||||||
/// Wraps a spawned tunnel child process.
|
/// Wraps a spawned tunnel child process.
|
||||||
pub(crate) struct TunnelProcess {
|
pub(crate) struct TunnelProcess {
|
||||||
pub child: tokio::process::Child,
|
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>>>;
|
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 ───────────────────────────────────────
|
// ── 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.
|
/// Start a managed tunnel if configured and no static URL is already set.
|
||||||
///
|
///
|
||||||
/// Returns the (potentially mutated) config with `tunnel.public_url` set,
|
/// Returns the (potentially mutated) config with `tunnel.public_url` set,
|
||||||
@@ -201,28 +221,17 @@ pub async fn start_managed_tunnel(
|
|||||||
return (config, None);
|
return (config, None);
|
||||||
};
|
};
|
||||||
|
|
||||||
let gateway_port = config
|
let (tunnel_host, tunnel_port) = resolve_tunnel_target(&config.channels);
|
||||||
.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");
|
|
||||||
|
|
||||||
match create_tunnel(provider_config) {
|
match create_tunnel(provider_config) {
|
||||||
Ok(Some(tunnel)) => {
|
Ok(Some(tunnel)) => {
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
"Starting {} tunnel on {}:{}...",
|
"Starting {} tunnel on {}:{}...",
|
||||||
tunnel.name(),
|
tunnel.name(),
|
||||||
gateway_host,
|
tunnel_host,
|
||||||
gateway_port
|
tunnel_port
|
||||||
);
|
);
|
||||||
match tunnel.start(gateway_host, gateway_port).await {
|
match tunnel.start(tunnel_host, tunnel_port).await {
|
||||||
Ok(url) => {
|
Ok(url) => {
|
||||||
tracing::debug!("Tunnel started: {}", url);
|
tracing::debug!("Tunnel started: {}", url);
|
||||||
config.tunnel.public_url = Some(url);
|
config.tunnel.public_url = Some(url);
|
||||||
@@ -383,10 +392,111 @@ mod tests {
|
|||||||
|
|
||||||
{
|
{
|
||||||
let mut guard = proc.lock().await;
|
let mut guard = proc.lock().await;
|
||||||
*guard = Some(TunnelProcess { child });
|
*guard = Some(TunnelProcess {
|
||||||
|
child,
|
||||||
|
_pipe_drain: None,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
kill_shared(&proc).await.unwrap();
|
kill_shared(&proc).await.unwrap();
|
||||||
assert!(proc.lock().await.is_none());
|
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
|
if let Ok(mut guard) = self.url.write() {
|
||||||
// to stdout; the pipe must be consumed to prevent SIGPIPE/buffer stalls.
|
*guard = Some(public_url.clone());
|
||||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
}
|
||||||
|
|
||||||
// Drain stderr silently — with --log stdout all meaningful output goes
|
// We took ownership of ngrok's stdout pipe above to parse the URL.
|
||||||
// to stdout; stderr only needs to be consumed to prevent pipe stalls.
|
// 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 {
|
if let Some(stderr) = stderr {
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let mut err_reader = tokio::io::BufReader::new(stderr).lines();
|
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;
|
let mut guard = self.proc.lock().await;
|
||||||
*guard = Some(TunnelProcess { child });
|
*guard = Some(TunnelProcess {
|
||||||
|
child,
|
||||||
|
_pipe_drain: Some(drain_handle),
|
||||||
|
});
|
||||||
|
|
||||||
Ok(public_url)
|
Ok(public_url)
|
||||||
}
|
}
|
||||||
|
|||||||
+51
-1
@@ -1,5 +1,7 @@
|
|||||||
//! Shared utility functions used across the codebase.
|
//! Shared utility functions used across the codebase.
|
||||||
|
|
||||||
|
use crate::llm::{ChatMessage, Role};
|
||||||
|
|
||||||
/// Find the largest valid UTF-8 char boundary at or before `pos`.
|
/// Find the largest valid UTF-8 char boundary at or before `pos`.
|
||||||
///
|
///
|
||||||
/// Polyfill for `str::floor_char_boundary` (nightly-only). Use when
|
/// Polyfill for `str::floor_char_boundary` (nightly-only). Use when
|
||||||
@@ -16,6 +18,17 @@ pub fn floor_char_boundary(s: &str, pos: usize) -> usize {
|
|||||||
i
|
i
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Ensure the last message in `messages` is a user-role message.
|
||||||
|
///
|
||||||
|
/// NEAR AI rejects conversations that don't end with a user message;
|
||||||
|
/// Claude 4.6 rejects assistant prefill. Call this before any LLM
|
||||||
|
/// completion request to satisfy both requirements.
|
||||||
|
pub fn ensure_ends_with_user_message(messages: &mut Vec<ChatMessage>) {
|
||||||
|
if !matches!(messages.last(), Some(m) if m.role == Role::User) {
|
||||||
|
messages.push(ChatMessage::user("Continue."));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Check if an LLM response explicitly signals that a job/task is complete.
|
/// Check if an LLM response explicitly signals that a job/task is complete.
|
||||||
///
|
///
|
||||||
/// Uses phrase-level matching to avoid false positives from bare words like
|
/// Uses phrase-level matching to avoid false positives from bare words like
|
||||||
@@ -72,7 +85,8 @@ pub fn llm_signals_completion(response: &str) -> bool {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use crate::util::{floor_char_boundary, llm_signals_completion};
|
use crate::llm::ChatMessage;
|
||||||
|
use crate::util::{ensure_ends_with_user_message, floor_char_boundary, llm_signals_completion};
|
||||||
|
|
||||||
// ── floor_char_boundary ──
|
// ── floor_char_boundary ──
|
||||||
|
|
||||||
@@ -103,6 +117,42 @@ mod tests {
|
|||||||
assert_eq!(floor_char_boundary("", 5), 0);
|
assert_eq!(floor_char_boundary("", 5), 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── ensure_ends_with_user_message ──
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ensure_user_message_injects_when_empty() {
|
||||||
|
let mut msgs: Vec<ChatMessage> = vec![];
|
||||||
|
ensure_ends_with_user_message(&mut msgs);
|
||||||
|
assert_eq!(msgs.len(), 1);
|
||||||
|
assert_eq!(msgs[0].role, crate::llm::Role::User);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ensure_user_message_injects_after_assistant() {
|
||||||
|
let mut msgs = vec![ChatMessage::user("hi"), ChatMessage::assistant("hello")];
|
||||||
|
ensure_ends_with_user_message(&mut msgs);
|
||||||
|
assert_eq!(msgs.len(), 3);
|
||||||
|
assert_eq!(msgs[2].role, crate::llm::Role::User);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ensure_user_message_injects_after_tool_result() {
|
||||||
|
let mut msgs = vec![
|
||||||
|
ChatMessage::user("run tool"),
|
||||||
|
ChatMessage::tool_result("call_1", "my_tool", "result"),
|
||||||
|
];
|
||||||
|
ensure_ends_with_user_message(&mut msgs);
|
||||||
|
assert_eq!(msgs.len(), 3);
|
||||||
|
assert_eq!(msgs[2].role, crate::llm::Role::User);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ensure_user_message_no_op_when_already_user() {
|
||||||
|
let mut msgs = vec![ChatMessage::user("hello")];
|
||||||
|
ensure_ends_with_user_message(&mut msgs);
|
||||||
|
assert_eq!(msgs.len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
// ── llm_signals_completion ──
|
// ── llm_signals_completion ──
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -151,7 +151,7 @@ Job: {}
|
|||||||
Description: {}
|
Description: {}
|
||||||
|
|
||||||
You have tools for shell commands, file operations, and code editing.
|
You have tools for shell commands, file operations, and code editing.
|
||||||
Work independently to complete this job. Report when done."#,
|
Work independently to complete this job. When finished, your final message MUST include the phrase "The job is complete" to signal termination."#,
|
||||||
job.title, job.description
|
job.title, job.description
|
||||||
)));
|
)));
|
||||||
|
|
||||||
@@ -373,6 +373,10 @@ impl LoopDelegate for ContainerDelegate {
|
|||||||
// Poll for follow-up prompts from the user
|
// Poll for follow-up prompts from the user
|
||||||
self.poll_and_inject_prompt(reason_ctx).await;
|
self.poll_and_inject_prompt(reason_ctx).await;
|
||||||
|
|
||||||
|
// Claude 4.6 rejects assistant prefill; NEAR AI rejects any non-user-ending
|
||||||
|
// conversation. Ensure the last message is user-role before calling the LLM.
|
||||||
|
crate::util::ensure_ends_with_user_message(&mut reason_ctx.messages);
|
||||||
|
|
||||||
// Refresh tools (in case WASM tools were built)
|
// Refresh tools (in case WASM tools were built)
|
||||||
reason_ctx.available_tools = self.tools.tool_definitions().await;
|
reason_ctx.available_tools = self.tools.tool_definitions().await;
|
||||||
|
|
||||||
|
|||||||
+82
-11
@@ -18,7 +18,7 @@ use crate::agent::agentic_loop::{
|
|||||||
};
|
};
|
||||||
use crate::agent::scheduler::WorkerMessage;
|
use crate::agent::scheduler::WorkerMessage;
|
||||||
use crate::agent::task::TaskOutput;
|
use crate::agent::task::TaskOutput;
|
||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::ToolDecisionDto;
|
||||||
use crate::context::{ContextManager, JobState};
|
use crate::context::{ContextManager, JobState};
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
@@ -33,6 +33,7 @@ use crate::tools::rate_limiter::RateLimitResult;
|
|||||||
use crate::tools::{
|
use crate::tools::{
|
||||||
ApprovalContext, ToolRegistry, autonomous_unavailable_error, prepare_tool_params, redact_params,
|
ApprovalContext, ToolRegistry, autonomous_unavailable_error, prepare_tool_params, redact_params,
|
||||||
};
|
};
|
||||||
|
use ironclaw_common::AppEvent;
|
||||||
|
|
||||||
/// Shared dependencies for worker execution.
|
/// Shared dependencies for worker execution.
|
||||||
///
|
///
|
||||||
@@ -48,8 +49,8 @@ pub struct WorkerDeps {
|
|||||||
pub hooks: Arc<HookRegistry>,
|
pub hooks: Arc<HookRegistry>,
|
||||||
pub timeout: Duration,
|
pub timeout: Duration,
|
||||||
pub use_planning: bool,
|
pub use_planning: bool,
|
||||||
/// SSE broadcast sender for live job event streaming to the web gateway.
|
/// Broadcast sender for live job event streaming to the web gateway.
|
||||||
pub sse_tx: Option<tokio::sync::broadcast::Sender<SseEvent>>,
|
pub sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||||
/// Approval context for tool execution. When `None`, all non-`Never` tools are
|
/// Approval context for tool execution. When `None`, all non-`Never` tools are
|
||||||
/// blocked (legacy behavior). When `Some`, the context determines which tools
|
/// blocked (legacy behavior). When `Some`, the context determines which tools
|
||||||
/// are pre-approved for autonomous execution.
|
/// are pre-approved for autonomous execution.
|
||||||
@@ -138,10 +139,10 @@ impl Worker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Broadcast SSE for live web UI updates
|
// Broadcast SSE for live web UI updates
|
||||||
if let Some(ref tx) = self.deps.sse_tx {
|
if let Some(ref sse) = self.deps.sse_tx {
|
||||||
let job_id_str = job_id.to_string();
|
let job_id_str = job_id.to_string();
|
||||||
let event = match event_type {
|
let event = match event_type {
|
||||||
"message" => Some(SseEvent::JobMessage {
|
"message" => Some(AppEvent::JobMessage {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
role: data
|
role: data
|
||||||
.get("role")
|
.get("role")
|
||||||
@@ -154,7 +155,7 @@ impl Worker {
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
}),
|
}),
|
||||||
"tool_use" => Some(SseEvent::JobToolUse {
|
"tool_use" => Some(AppEvent::JobToolUse {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: data
|
tool_name: data
|
||||||
.get("tool_name")
|
.get("tool_name")
|
||||||
@@ -166,7 +167,7 @@ impl Worker {
|
|||||||
.cloned()
|
.cloned()
|
||||||
.unwrap_or(serde_json::Value::Null),
|
.unwrap_or(serde_json::Value::Null),
|
||||||
}),
|
}),
|
||||||
"tool_result" => Some(SseEvent::JobToolResult {
|
"tool_result" => Some(AppEvent::JobToolResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: data
|
tool_name: data
|
||||||
.get("tool_name")
|
.get("tool_name")
|
||||||
@@ -179,7 +180,7 @@ impl Worker {
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
}),
|
}),
|
||||||
"status" => Some(SseEvent::JobStatus {
|
"status" => Some(AppEvent::JobStatus {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
message: data
|
message: data
|
||||||
.get("message")
|
.get("message")
|
||||||
@@ -187,7 +188,7 @@ impl Worker {
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
}),
|
}),
|
||||||
"result" => Some(SseEvent::JobResult {
|
"result" => Some(AppEvent::JobResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
status: data
|
status: data
|
||||||
.get("status")
|
.get("status")
|
||||||
@@ -200,10 +201,23 @@ impl Worker {
|
|||||||
.map(|s| s.to_string()),
|
.map(|s| s.to_string()),
|
||||||
fallback_deliverable: data.get("fallback_deliverable").cloned(),
|
fallback_deliverable: data.get("fallback_deliverable").cloned(),
|
||||||
}),
|
}),
|
||||||
|
"reasoning" => {
|
||||||
|
let narrative = data
|
||||||
|
.get("narrative")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
let decisions = ToolDecisionDto::from_json_array(&data["decisions"]);
|
||||||
|
Some(AppEvent::JobReasoning {
|
||||||
|
job_id: job_id_str,
|
||||||
|
narrative,
|
||||||
|
decisions,
|
||||||
|
})
|
||||||
|
}
|
||||||
_ => None,
|
_ => None,
|
||||||
};
|
};
|
||||||
if let Some(event) = event {
|
if let Some(event) = event {
|
||||||
let _ = tx.send(event);
|
sse.broadcast(event);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -897,6 +911,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
|||||||
id: selection.tool_call_id.clone(),
|
id: selection.tool_call_id.clone(),
|
||||||
name: selection.tool_name.clone(),
|
name: selection.tool_name.clone(),
|
||||||
arguments: selection.parameters.clone(),
|
arguments: selection.parameters.clone(),
|
||||||
|
reasoning: if action.reasoning.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(action.reasoning.clone())
|
||||||
|
},
|
||||||
}],
|
}],
|
||||||
));
|
));
|
||||||
|
|
||||||
@@ -1232,6 +1251,11 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
|
|||||||
) -> Option<LoopOutcome> {
|
) -> Option<LoopOutcome> {
|
||||||
// Refresh tool definitions so newly built tools become visible
|
// Refresh tool definitions so newly built tools become visible
|
||||||
reason_ctx.available_tools = self.worker.tools().tool_definitions().await;
|
reason_ctx.available_tools = self.worker.tools().tool_definitions().await;
|
||||||
|
|
||||||
|
// Claude 4.6 rejects assistant prefill; NEAR AI rejects any non-user-ending
|
||||||
|
// conversation. Ensure the last message is user-role before calling the LLM.
|
||||||
|
crate::util::ensure_ends_with_user_message(&mut reason_ctx.messages);
|
||||||
|
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1352,6 +1376,48 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Emit reasoning event if any tool calls carry reasoning.
|
||||||
|
// Sanitize narrative and per-tool rationale through SafetyLayer
|
||||||
|
// (parity with ChatDelegate in dispatcher.rs).
|
||||||
|
let sanitized_narrative = content
|
||||||
|
.as_deref()
|
||||||
|
.filter(|c| !c.trim().is_empty())
|
||||||
|
.map(|c| {
|
||||||
|
self.worker
|
||||||
|
.deps
|
||||||
|
.safety
|
||||||
|
.sanitize_tool_output("job_narrative", c)
|
||||||
|
.content
|
||||||
|
})
|
||||||
|
.filter(|c| !c.trim().is_empty())
|
||||||
|
.unwrap_or_default();
|
||||||
|
let decisions: Vec<serde_json::Value> = tool_calls
|
||||||
|
.iter()
|
||||||
|
.filter_map(|tc| {
|
||||||
|
tc.reasoning.as_ref().map(|r| {
|
||||||
|
let sanitized = self
|
||||||
|
.worker
|
||||||
|
.deps
|
||||||
|
.safety
|
||||||
|
.sanitize_tool_output("tool_rationale", r)
|
||||||
|
.content;
|
||||||
|
serde_json::json!({
|
||||||
|
"tool_name": tc.name,
|
||||||
|
"rationale": sanitized,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
if !decisions.is_empty() {
|
||||||
|
self.worker.log_event(
|
||||||
|
"reasoning",
|
||||||
|
serde_json::json!({
|
||||||
|
"narrative": sanitized_narrative,
|
||||||
|
"decisions": decisions,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Add assistant message with tool_calls (OpenAI protocol)
|
// Add assistant message with tool_calls (OpenAI protocol)
|
||||||
reason_ctx
|
reason_ctx
|
||||||
.messages
|
.messages
|
||||||
@@ -1366,7 +1432,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
|
|||||||
.map(|tc| ToolSelection {
|
.map(|tc| ToolSelection {
|
||||||
tool_name: tc.name.clone(),
|
tool_name: tc.name.clone(),
|
||||||
parameters: tc.arguments.clone(),
|
parameters: tc.arguments.clone(),
|
||||||
reasoning: String::new(),
|
reasoning: tc.reasoning.clone().unwrap_or_default(),
|
||||||
alternatives: vec![],
|
alternatives: vec![],
|
||||||
tool_call_id: tc.id.clone(),
|
tool_call_id: tc.id.clone(),
|
||||||
})
|
})
|
||||||
@@ -1419,6 +1485,11 @@ fn selections_to_tool_calls(selections: &[ToolSelection]) -> Vec<ToolCall> {
|
|||||||
id: s.tool_call_id.clone(),
|
id: s.tool_call_id.clone(),
|
||||||
name: s.tool_name.clone(),
|
name: s.tool_name.clone(),
|
||||||
arguments: s.parameters.clone(),
|
arguments: s.parameters.clone(),
|
||||||
|
reasoning: if s.reasoning.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(s.reasoning.clone())
|
||||||
|
},
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,14 +3,13 @@
|
|||||||
//! Avoids redundant HTTP calls for identical texts by caching embeddings
|
//! Avoids redundant HTTP calls for identical texts by caching embeddings
|
||||||
//! in memory keyed by `SHA-256(model_name + "\0" + text)`.
|
//! in memory keyed by `SHA-256(model_name + "\0" + text)`.
|
||||||
//!
|
//!
|
||||||
//! Follows the same cache pattern as `llm::response_cache::CachedProvider`:
|
//! Uses `lru::LruCache` for O(1) insertion, lookup, and eviction.
|
||||||
//! `HashMap` + `last_accessed` tracking + manual LRU eviction.
|
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::num::NonZeroUsize;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
use std::time::Instant;
|
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use lru::LruCache;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
|
|
||||||
use crate::workspace::embeddings::{EmbeddingError, EmbeddingProvider};
|
use crate::workspace::embeddings::{EmbeddingError, EmbeddingProvider};
|
||||||
@@ -22,8 +21,7 @@ pub struct EmbeddingCacheConfig {
|
|||||||
///
|
///
|
||||||
/// Approximate raw embedding payload: `max_entries × dimension × 4 bytes`.
|
/// Approximate raw embedding payload: `max_entries × dimension × 4 bytes`.
|
||||||
/// At 10,000 entries × 1536 floats ≈ 58 MB (payload only; actual memory
|
/// At 10,000 entries × 1536 floats ≈ 58 MB (payload only; actual memory
|
||||||
/// is higher due to HashMap buckets, `[u8; 32]` hash keys, `Vec`/`Instant`
|
/// is higher due to per-entry overhead in the linked-list LRU).
|
||||||
/// per-entry overhead).
|
|
||||||
pub max_entries: usize,
|
pub max_entries: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -35,11 +33,6 @@ impl Default for EmbeddingCacheConfig {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct CacheEntry {
|
|
||||||
embedding: Vec<f32>,
|
|
||||||
last_accessed: Instant,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Embedding provider wrapper that caches results in memory.
|
/// Embedding provider wrapper that caches results in memory.
|
||||||
///
|
///
|
||||||
/// Thread-safe via `std::sync::Mutex`. The lock is **never held**
|
/// Thread-safe via `std::sync::Mutex`. The lock is **never held**
|
||||||
@@ -47,8 +40,7 @@ struct CacheEntry {
|
|||||||
/// so a synchronous mutex is cheaper than `tokio::sync::Mutex`.
|
/// so a synchronous mutex is cheaper than `tokio::sync::Mutex`.
|
||||||
pub struct CachedEmbeddingProvider {
|
pub struct CachedEmbeddingProvider {
|
||||||
inner: Arc<dyn EmbeddingProvider>,
|
inner: Arc<dyn EmbeddingProvider>,
|
||||||
cache: Mutex<HashMap<[u8; 32], CacheEntry>>,
|
cache: Mutex<LruCache<[u8; 32], Vec<f32>>>,
|
||||||
config: EmbeddingCacheConfig,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl CachedEmbeddingProvider {
|
impl CachedEmbeddingProvider {
|
||||||
@@ -56,19 +48,18 @@ impl CachedEmbeddingProvider {
|
|||||||
///
|
///
|
||||||
/// `config.max_entries` is clamped to at least 1.
|
/// `config.max_entries` is clamped to at least 1.
|
||||||
pub fn new(inner: Arc<dyn EmbeddingProvider>, config: EmbeddingCacheConfig) -> Self {
|
pub fn new(inner: Arc<dyn EmbeddingProvider>, config: EmbeddingCacheConfig) -> Self {
|
||||||
let config = EmbeddingCacheConfig {
|
let max_entries = config.max_entries.max(1);
|
||||||
max_entries: config.max_entries.max(1),
|
if max_entries > 100_000 {
|
||||||
};
|
|
||||||
if config.max_entries > 100_000 {
|
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
max_entries = config.max_entries,
|
max_entries,
|
||||||
"Embedding cache size exceeds 100,000 entries; memory usage may be significant"
|
"Embedding cache size exceeds 100,000 entries; memory usage may be significant"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
// safety: max_entries >= 1 due to .max(1) above
|
||||||
|
let cap = NonZeroUsize::new(max_entries).expect("clamped to >= 1"); // safety: always >= 1
|
||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
cache: Mutex::new(HashMap::with_capacity(config.max_entries.min(1024))),
|
cache: Mutex::new(LruCache::new(cap)),
|
||||||
config,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -100,49 +91,6 @@ impl CachedEmbeddingProvider {
|
|||||||
hasher.update(text.as_bytes());
|
hasher.update(text.as_bytes());
|
||||||
hasher.finalize().into()
|
hasher.finalize().into()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Evict the least-recently-used entry if at capacity (single-entry path).
|
|
||||||
// TODO: O(n) scan per eviction. If max_entries grows large, switch to
|
|
||||||
// an ordered data structure (e.g. `IndexMap` with swap_remove, or a
|
|
||||||
// linked-list LRU like the `lru` crate).
|
|
||||||
fn evict_lru(cache: &mut HashMap<[u8; 32], CacheEntry>, max_entries: usize) {
|
|
||||||
while cache.len() >= max_entries {
|
|
||||||
let oldest_key = cache
|
|
||||||
.iter()
|
|
||||||
.min_by_key(|(_, entry)| entry.last_accessed)
|
|
||||||
.map(|(k, _)| *k);
|
|
||||||
|
|
||||||
if let Some(k) = oldest_key {
|
|
||||||
cache.remove(&k);
|
|
||||||
} else {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Evict the `k` oldest entries in O(n) average time via partial selection.
|
|
||||||
///
|
|
||||||
/// Used by `embed_batch` to avoid the O(n×m) cost of calling
|
|
||||||
/// `evict_lru` per insert.
|
|
||||||
fn evict_k_oldest(cache: &mut HashMap<[u8; 32], CacheEntry>, k: usize) {
|
|
||||||
if k == 0 || cache.is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if k >= cache.len() {
|
|
||||||
cache.clear();
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
// Partial selection: find the k oldest in O(n) average via
|
|
||||||
// select_nth_unstable_by_key, then remove the first k entries.
|
|
||||||
let mut entries: Vec<([u8; 32], Instant)> = cache
|
|
||||||
.iter()
|
|
||||||
.map(|(key, entry)| (*key, entry.last_accessed))
|
|
||||||
.collect();
|
|
||||||
entries.select_nth_unstable_by_key(k - 1, |(_, t)| *t);
|
|
||||||
for (key, _) in entries.into_iter().take(k) {
|
|
||||||
cache.remove(&key);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -162,39 +110,32 @@ impl EmbeddingProvider for CachedEmbeddingProvider {
|
|||||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||||
let key = self.cache_key(text);
|
let key = self.cache_key(text);
|
||||||
|
|
||||||
// Check cache (short critical section)
|
// Check cache (short critical section). LruCache::get promotes the
|
||||||
|
// entry to most-recently-used automatically.
|
||||||
{
|
{
|
||||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
if let Some(entry) = guard.get_mut(&key) {
|
if let Some(embedding) = guard.get(&key) {
|
||||||
entry.last_accessed = Instant::now();
|
|
||||||
tracing::trace!("embedding cache hit");
|
tracing::trace!("embedding cache hit");
|
||||||
return Ok(entry.embedding.clone());
|
return Ok(embedding.clone());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Lock released before HTTP call.
|
// Lock released before HTTP call.
|
||||||
// NOTE: Thundering herd — multiple concurrent callers with the same
|
// NOTE: Thundering herd — multiple concurrent callers with the same
|
||||||
// uncached key will each call the inner provider. This is acceptable:
|
// uncached key will each call the inner provider. This is acceptable:
|
||||||
// embeddings are idempotent and the last writer wins in the HashMap.
|
// embeddings are idempotent and the last writer wins in the LruCache.
|
||||||
|
|
||||||
let embedding = self.inner.embed(text).await?;
|
let embedding = self.inner.embed(text).await?;
|
||||||
|
|
||||||
// Store result. Re-check under lock: another concurrent caller may
|
// Store result under lock. Re-check first: another concurrent caller
|
||||||
// have inserted this key while the lock was released for the HTTP call.
|
// may have already cached this key while the lock was released.
|
||||||
{
|
{
|
||||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
if let Some(entry) = guard.get_mut(&key) {
|
if guard.get(&key).is_some() {
|
||||||
// Thundering herd — another caller already cached it.
|
// Thundering herd — another caller beat us. LruCache::get
|
||||||
// Just touch timestamp; skip the clone.
|
// already promoted it to most-recently-used; skip the clone.
|
||||||
entry.last_accessed = Instant::now();
|
tracing::trace!("embedding cache: concurrent insert, skipping clone");
|
||||||
} else {
|
} else {
|
||||||
Self::evict_lru(&mut guard, self.config.max_entries);
|
guard.push(key, embedding.clone());
|
||||||
guard.insert(
|
|
||||||
key,
|
|
||||||
CacheEntry {
|
|
||||||
embedding: embedding.clone(),
|
|
||||||
last_accessed: Instant::now(),
|
|
||||||
},
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -214,11 +155,9 @@ impl EmbeddingProvider for CachedEmbeddingProvider {
|
|||||||
|
|
||||||
{
|
{
|
||||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
let now = Instant::now();
|
|
||||||
for (i, key) in keys.iter().enumerate() {
|
for (i, key) in keys.iter().enumerate() {
|
||||||
if let Some(entry) = guard.get_mut(key) {
|
if let Some(embedding) = guard.get(key) {
|
||||||
entry.last_accessed = now;
|
results[i] = Some(embedding.clone());
|
||||||
results[i] = Some(entry.embedding.clone());
|
|
||||||
} else {
|
} else {
|
||||||
miss_indices.push(i);
|
miss_indices.push(i);
|
||||||
}
|
}
|
||||||
@@ -228,7 +167,6 @@ impl EmbeddingProvider for CachedEmbeddingProvider {
|
|||||||
|
|
||||||
if miss_indices.is_empty() {
|
if miss_indices.is_empty() {
|
||||||
tracing::trace!(count = texts.len(), "embedding batch: all cache hits");
|
tracing::trace!(count = texts.len(), "embedding batch: all cache hits");
|
||||||
// All slots populated from cache hits
|
|
||||||
return results
|
return results
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
@@ -260,29 +198,18 @@ impl EmbeddingProvider for CachedEmbeddingProvider {
|
|||||||
"embedding batch: partial cache"
|
"embedding batch: partial cache"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Cache FIRST (clone only the cacheable subset), then move originals
|
// Cache only the last `cap` new embeddings — caching more than the
|
||||||
// into results. This avoids cloning capacity-skipped embeddings entirely.
|
// cache capacity wastes clone work on entries that are immediately evicted.
|
||||||
{
|
{
|
||||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
let cacheable = miss_indices.len().min(self.config.max_entries);
|
let cap = guard.cap().get();
|
||||||
let skip = miss_indices.len() - cacheable;
|
let skip = miss_indices.len().saturating_sub(cap);
|
||||||
let need_to_evict = (guard.len() + cacheable).saturating_sub(self.config.max_entries);
|
|
||||||
if need_to_evict > 0 {
|
|
||||||
Self::evict_k_oldest(&mut guard, need_to_evict);
|
|
||||||
}
|
|
||||||
let now = Instant::now();
|
|
||||||
for (&orig_idx, emb) in miss_indices[skip..].iter().zip(&new_embeddings[skip..]) {
|
for (&orig_idx, emb) in miss_indices[skip..].iter().zip(&new_embeddings[skip..]) {
|
||||||
guard.insert(
|
guard.push(keys[orig_idx], emb.clone());
|
||||||
keys[orig_idx],
|
|
||||||
CacheEntry {
|
|
||||||
embedding: emb.clone(),
|
|
||||||
last_accessed: now,
|
|
||||||
},
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Move originals into results (zero-copy for all, including cached ones).
|
// Move originals into results (zero-copy).
|
||||||
for (orig_idx, emb) in miss_indices.iter().copied().zip(new_embeddings) {
|
for (orig_idx, emb) in miss_indices.iter().copied().zip(new_embeddings) {
|
||||||
results[orig_idx] = Some(emb);
|
results[orig_idx] = Some(emb);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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_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_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_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`
|
## `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). |
|
| `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. |
|
| `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. |
|
| `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. |
|
| `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. |
|
||||||
|
|
||||||
### Function-scoped fixtures
|
### Function-scoped fixtures
|
||||||
@@ -100,6 +102,8 @@ EMBEDDING_ENABLED=false, SKILLS_ENABLED=true
|
|||||||
ONBOARD_COMPLETED=true # prevents setup wizard
|
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.
|
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`)
|
## 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 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:
|
To add a new canned response:
|
||||||
```python
|
```python
|
||||||
# In mock_llm.py
|
# In mock_llm.py
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user