mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
Compare commits
95
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
52551f0ef4 | ||
|
|
db6450c47a | ||
|
|
018416d3f3 | ||
|
|
5b95d22218 | ||
|
|
dd0a0e10ab | ||
|
|
1d5777824c | ||
|
|
b92d333b0c | ||
|
|
00aba928e0 | ||
|
|
07c32b788c | ||
|
|
adf4e25c8f | ||
|
|
f40019aa4c | ||
|
|
9c63d189b7 | ||
|
|
ed4d92932a | ||
|
|
b3fbef5287 | ||
|
|
6b8a38e147 | ||
|
|
ab67f02886 | ||
|
|
f02345fd1f | ||
|
|
4c043bf057 | ||
|
|
0b4e7c761b | ||
|
|
cdc625566f | ||
|
|
bb24952622 | ||
|
|
ef37d705a1 | ||
|
|
b400c2a711 | ||
|
|
ea24d79ace | ||
|
|
8d632872fd | ||
|
|
4c5d961102 | ||
|
|
2b4e881a72 | ||
|
|
c0f33c37f7 | ||
|
|
5d714be354 | ||
|
|
3d43917cd0 | ||
|
|
f9dfb74800 | ||
|
|
cb01800f73 | ||
|
|
16aaea8d74 | ||
|
|
a19deb6812 | ||
|
|
2f80b7b0b8 | ||
|
|
2f47c611d4 | ||
|
|
1f8d901cf6 | ||
|
|
ad20a5ab4f | ||
|
|
e15c50ea2d | ||
|
|
d4e18020e2 | ||
|
|
a23d87fc00 | ||
|
|
c737fb0855 | ||
|
|
0145672f36 | ||
|
|
9fd5537a01 | ||
|
|
492d9d22c9 | ||
|
|
b8b88ab84e | ||
|
|
c98ec3fb18 | ||
|
|
189fa35e64 | ||
|
|
c5dce279e2 | ||
|
|
5a5ffe8d08 | ||
|
|
86d1143064 | ||
|
|
ab0ad948f3 | ||
|
|
c949521d8d | ||
|
|
0341fcc940 | ||
|
|
41ed0a0f98 | ||
|
|
67a025e2fa | ||
|
|
6daa2f155f | ||
|
|
706c3a1b47 | ||
|
|
656151783c | ||
|
|
82822d7b25 | ||
|
|
dcb2d89e3a | ||
|
|
f3da30a454 | ||
|
|
424b470c59 | ||
|
|
5901451603 | ||
|
|
d3d517fd67 | ||
|
|
01678be61d | ||
|
|
fb3548956b | ||
|
|
5847479fd8 | ||
|
|
3fdb187796 | ||
|
|
b441ebec02 | ||
|
|
ae370d7e2b | ||
|
|
fa51b9f52d | ||
|
|
98418b3ef0 | ||
|
|
74b2b4129e | ||
|
|
bb57e36e6d | ||
|
|
0194275792 | ||
|
|
ddf64e8485 | ||
|
|
bd6977e6a8 | ||
|
|
dea789cca9 | ||
|
|
d47b4b0346 | ||
|
|
91a241a3c7 | ||
|
|
d1d74d665a | ||
|
|
e077e1277d | ||
|
|
6fc8cc2f39 | ||
|
|
e031d8246b | ||
|
|
23263029f9 | ||
|
|
d5e08b95f9 | ||
|
|
e4d3200d80 | ||
|
|
7dc3c6d067 | ||
|
|
e1774e9ec0 | ||
|
|
e1d9827b21 | ||
|
|
e582166781 | ||
|
|
656d1f3e86 | ||
|
|
0e3aa4f806 | ||
|
|
44d16732a7 |
@@ -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
|
||||||
|
|||||||
+150
@@ -7,6 +7,156 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
## [0.23.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.22.0...ironclaw-v0.23.0) - 2026-03-27
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- complete multi-tenant isolation — phases 2–4 ([#1614](https://github.com/nearai/ironclaw/pull/1614))
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- *(routines)* recover delete name after failed update fallback ([#1108](https://github.com/nearai/ironclaw/pull/1108))
|
||||||
|
- *(mcp)* handle 202 Accepted and wire session manager for Streamable HTTP ([#1437](https://github.com/nearai/ironclaw/pull/1437))
|
||||||
|
- *(extensions)* channel-relay auth dead-end, observability, and URL override ([#1681](https://github.com/nearai/ironclaw/pull/1681))
|
||||||
|
- *(agent)* discard truncated tool calls when finish_reason == Length ([#1631](https://github.com/nearai/ironclaw/pull/1631)) ([#1632](https://github.com/nearai/ironclaw/pull/1632))
|
||||||
|
- *(llm)* filter XML tool-call recovery by context ([#1641](https://github.com/nearai/ironclaw/pull/1641))
|
||||||
|
|
||||||
|
### Other
|
||||||
|
|
||||||
|
- Support direct hosted OAuth callbacks with proxy auth token ([#1684](https://github.com/nearai/ironclaw/pull/1684))
|
||||||
|
|
||||||
|
## [0.22.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.21.0...ironclaw-v0.22.0) - 2026-03-25
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- *(agent)* thread per-tool reasoning through provider, session, and all surfaces ([#1513](https://github.com/nearai/ironclaw/pull/1513))
|
||||||
|
- *(cli)* show credential auth status in tool info ([#1572](https://github.com/nearai/ironclaw/pull/1572))
|
||||||
|
- multi-tenant auth with per-user workspace isolation ([#1118](https://github.com/nearai/ironclaw/pull/1118))
|
||||||
|
- *(cli)* add ironclaw models subcommands (list/status/set/set-provider) ([#1043](https://github.com/nearai/ironclaw/pull/1043))
|
||||||
|
- *(workspace)* multi-scope workspace reads ([#1117](https://github.com/nearai/ironclaw/pull/1117))
|
||||||
|
- *(ux)* complete UX overhaul — design system, onboarding, web polish ([#1277](https://github.com/nearai/ironclaw/pull/1277))
|
||||||
|
- *(gemini_oauth)* full Gemini CLI OAuth integration with Cloud Code API ([#1356](https://github.com/nearai/ironclaw/pull/1356))
|
||||||
|
- *(shell)* add Low/Medium/High risk levels for graduated command approval (closes #172) ([#368](https://github.com/nearai/ironclaw/pull/368))
|
||||||
|
- *(agent)* queue and merge messages during active turns ([#1412](https://github.com/nearai/ironclaw/pull/1412))
|
||||||
|
- *(cli)* add `ironclaw hooks list` subcommand ([#1023](https://github.com/nearai/ironclaw/pull/1023))
|
||||||
|
- *(extensions)* support text setup fields in web configure modal ([#496](https://github.com/nearai/ironclaw/pull/496))
|
||||||
|
- *(llm)* add GitHub Copilot as LLM provider ([#1512](https://github.com/nearai/ironclaw/pull/1512))
|
||||||
|
- *(workspace)* layered memory with sensitivity-based privacy redirect ([#1112](https://github.com/nearai/ironclaw/pull/1112))
|
||||||
|
- *(webhooks)* add public webhook trigger endpoint for routines ([#736](https://github.com/nearai/ironclaw/pull/736))
|
||||||
|
- *(llm)* Add OpenAI Codex (ChatGPT subscription) as LLM provider ([#1461](https://github.com/nearai/ironclaw/pull/1461))
|
||||||
|
- *(web)* add light theme with dark/light/system toggle ([#1457](https://github.com/nearai/ironclaw/pull/1457))
|
||||||
|
- *(agent)* activate stuck_threshold for time-based stuck job detection ([#1234](https://github.com/nearai/ironclaw/pull/1234))
|
||||||
|
- chat onboarding and routine advisor ([#927](https://github.com/nearai/ironclaw/pull/927))
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- ensure LLM calls always end with user message (closes #763) ([#1259](https://github.com/nearai/ironclaw/pull/1259))
|
||||||
|
- restore owner-scoped gateway startup ([#1625](https://github.com/nearai/ironclaw/pull/1625))
|
||||||
|
- remove stale stream_token gate from channel-relay activation ([#1623](https://github.com/nearai/ironclaw/pull/1623))
|
||||||
|
- *(agent)* case-insensitive channel match and user_id filter for event triggers ([#1211](https://github.com/nearai/ironclaw/pull/1211))
|
||||||
|
- *(routines)* normalize status display across web and CLI ([#1469](https://github.com/nearai/ironclaw/pull/1469))
|
||||||
|
- *(tunnel)* managed tunnels target wrong port and die from SIGPIPE ([#1093](https://github.com/nearai/ironclaw/pull/1093))
|
||||||
|
- *(agent)* persist /model selection to .env, TOML, and DB ([#1581](https://github.com/nearai/ironclaw/pull/1581))
|
||||||
|
- post-merge review sweep — 8 fixes across security, perf, and correctness ([#1550](https://github.com/nearai/ironclaw/pull/1550))
|
||||||
|
- generate Mistral-compatible 9-char alphanumeric tool call IDs ([#1242](https://github.com/nearai/ironclaw/pull/1242))
|
||||||
|
- *(mcp)* handle empty 202 notification acknowledgements ([#1539](https://github.com/nearai/ironclaw/pull/1539))
|
||||||
|
- *(tests)* eliminate env mutex poison cascade ([#1558](https://github.com/nearai/ironclaw/pull/1558))
|
||||||
|
- *(safety)* escape tool output XML content and remove misleading sanitized attr ([#1067](https://github.com/nearai/ironclaw/pull/1067))
|
||||||
|
- *(oauth)* reject malformed ic2.* states in decode_hosted_oauth_state ([#1441](https://github.com/nearai/ironclaw/pull/1441)) ([#1454](https://github.com/nearai/ironclaw/pull/1454))
|
||||||
|
- parameter coercion and validation for oneOf/anyOf/allOf schemas ([#1397](https://github.com/nearai/ironclaw/pull/1397))
|
||||||
|
- persist startup-loaded MCP clients in ExtensionManager ([#1509](https://github.com/nearai/ironclaw/pull/1509))
|
||||||
|
- *(deps)* patch rustls-webpki vulnerability (RUSTSEC-2026-0049)
|
||||||
|
- *(routines)* add missing extension_manager field in trigger_manual EngineContext
|
||||||
|
- *(ci)* serialize env-mutating OAuth wildcard tests with ENV_MUTEX ([#1280](https://github.com/nearai/ironclaw/pull/1280)) ([#1468](https://github.com/nearai/ironclaw/pull/1468))
|
||||||
|
- *(setup)* remove redundant LLM config and API keys from bootstrap .env ([#1448](https://github.com/nearai/ironclaw/pull/1448))
|
||||||
|
- resolve wasm broadcast merge conflicts with staging ([#395](https://github.com/nearai/ironclaw/pull/395)) ([#1460](https://github.com/nearai/ironclaw/pull/1460))
|
||||||
|
- skip credential validation for Bedrock backend ([#1011](https://github.com/nearai/ironclaw/pull/1011))
|
||||||
|
- register sandbox jobs in ContextManager for query tool visibility ([#1426](https://github.com/nearai/ironclaw/pull/1426))
|
||||||
|
- prefer execution-local message routing metadata ([#1449](https://github.com/nearai/ironclaw/pull/1449))
|
||||||
|
- *(security)* validate embedding base URLs to prevent SSRF ([#1221](https://github.com/nearai/ironclaw/pull/1221))
|
||||||
|
- f32→f64 precision artifact in temperature causes provider 400 errors ([#1450](https://github.com/nearai/ironclaw/pull/1450))
|
||||||
|
- *(routines)* surface errors when sandbox unavailable for full_job routines ([#769](https://github.com/nearai/ironclaw/pull/769))
|
||||||
|
- restore libSQL vector search with dynamic dimensions ([#1393](https://github.com/nearai/ironclaw/pull/1393))
|
||||||
|
- staging CI triage — consolidate retry parsing, fix flaky tests, add docs ([#1427](https://github.com/nearai/ironclaw/pull/1427))
|
||||||
|
|
||||||
|
### Other
|
||||||
|
|
||||||
|
- Merge branch 'main' into staging-promote/455f543b-23329172268
|
||||||
|
- Merge pull request #1655 from nearai/codex/fix-staging-promotion-1451-version-bumps
|
||||||
|
- Merge pull request #1499 from nearai/staging-promote/9603fefd-23364438978
|
||||||
|
- Fix libsql prompt scope regressions ([#1651](https://github.com/nearai/ironclaw/pull/1651))
|
||||||
|
- Normalize cron schedules on routine create ([#1648](https://github.com/nearai/ironclaw/pull/1648))
|
||||||
|
- Fix MCP lifecycle trace user scope ([#1646](https://github.com/nearai/ironclaw/pull/1646))
|
||||||
|
- Fix REPL single-message hang and cap CI test duration ([#1643](https://github.com/nearai/ironclaw/pull/1643))
|
||||||
|
- extract AppEvent to crates/ironclaw_common ([#1615](https://github.com/nearai/ironclaw/pull/1615))
|
||||||
|
- Fix hosted OAuth refresh via proxy ([#1602](https://github.com/nearai/ironclaw/pull/1602))
|
||||||
|
- *(agent)* optimize approval thread resolution (UUID parsing + lock contention) ([#1592](https://github.com/nearai/ironclaw/pull/1592))
|
||||||
|
- *(tools)* auto-compact WASM tool schemas, add descriptions, improve credential prompts ([#1525](https://github.com/nearai/ironclaw/pull/1525))
|
||||||
|
- Default new lightweight routines to tools-enabled ([#1573](https://github.com/nearai/ironclaw/pull/1573))
|
||||||
|
- Google OAuth URL broken when initiated from Telegram channel ([#1165](https://github.com/nearai/ironclaw/pull/1165))
|
||||||
|
- add gitcgr code graph badge ([#1563](https://github.com/nearai/ironclaw/pull/1563))
|
||||||
|
- Fix owner-scoped message routing fallbacks ([#1574](https://github.com/nearai/ironclaw/pull/1574))
|
||||||
|
- *(tools)* remove unconditional params clone in shared execution (fix #893) ([#926](https://github.com/nearai/ironclaw/pull/926))
|
||||||
|
- *(llm)* move transcription module into src/llm/ ([#1559](https://github.com/nearai/ironclaw/pull/1559))
|
||||||
|
- *(agent)* avoid preview allocations for non-truncated strings (fix #894) ([#924](https://github.com/nearai/ironclaw/pull/924))
|
||||||
|
- Expand AGENTS.md with coding agents guidance ([#1392](https://github.com/nearai/ironclaw/pull/1392))
|
||||||
|
- Fix CI approval flows and stale fixtures ([#1478](https://github.com/nearai/ironclaw/pull/1478))
|
||||||
|
- Use live owner tool scope for autonomous routines and jobs ([#1453](https://github.com/nearai/ironclaw/pull/1453))
|
||||||
|
- use Arc in embedding cache to avoid clones on miss path ([#1438](https://github.com/nearai/ironclaw/pull/1438))
|
||||||
|
- Add owner-scoped permissions for full-job routines ([#1440](https://github.com/nearai/ironclaw/pull/1440))
|
||||||
|
|
||||||
|
## [0.21.0](https://github.com/nearai/ironclaw/compare/v0.20.0...v0.21.0) - 2026-03-20
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- structured fallback deliverables for failed/stuck jobs ([#236](https://github.com/nearai/ironclaw/pull/236))
|
||||||
|
- LRU embedding cache for workspace search ([#1423](https://github.com/nearai/ironclaw/pull/1423))
|
||||||
|
- receive relay events via webhook callbacks ([#1254](https://github.com/nearai/ironclaw/pull/1254))
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- bump Feishu channel version for promotion
|
||||||
|
- *(approval)* make "always" auto-approve work for credentialed HTTP requests ([#1257](https://github.com/nearai/ironclaw/pull/1257))
|
||||||
|
- skip NEAR AI session check when backend is not nearai ([#1413](https://github.com/nearai/ironclaw/pull/1413))
|
||||||
|
|
||||||
|
### Other
|
||||||
|
|
||||||
|
- Make hosted OAuth and MCP auth generic ([#1375](https://github.com/nearai/ironclaw/pull/1375))
|
||||||
|
|
||||||
|
## [0.20.0](https://github.com/nearai/ironclaw/compare/v0.19.0...v0.20.0) - 2026-03-19
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- *(self-repair)* wire stuck_threshold, store, and builder ([#712](https://github.com/nearai/ironclaw/pull/712))
|
||||||
|
- *(testing)* add FaultInjector framework for StubLlm ([#1233](https://github.com/nearai/ironclaw/pull/1233))
|
||||||
|
- *(gateway)* unified settings page with subtabs ([#1191](https://github.com/nearai/ironclaw/pull/1191))
|
||||||
|
- upgrade MiniMax default model to M2.7 ([#1357](https://github.com/nearai/ironclaw/pull/1357))
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- navigate telegram E2E tests to channels subtab ([#1408](https://github.com/nearai/ironclaw/pull/1408))
|
||||||
|
- add missing `builder` field and update E2E extensions tab navigation ([#1400](https://github.com/nearai/ironclaw/pull/1400))
|
||||||
|
- remove debug_assert guards that panic on valid error paths ([#1385](https://github.com/nearai/ironclaw/pull/1385))
|
||||||
|
- address valid review comments from PR #1359 ([#1380](https://github.com/nearai/ironclaw/pull/1380))
|
||||||
|
- full_job routine runs stay running until linked job completion ([#1374](https://github.com/nearai/ironclaw/pull/1374))
|
||||||
|
- full_job routine concurrency tracks linked job lifetime ([#1372](https://github.com/nearai/ironclaw/pull/1372))
|
||||||
|
- remove -x from coverage pytest to prevent suite-blocking failures ([#1360](https://github.com/nearai/ironclaw/pull/1360))
|
||||||
|
- add debug_assert invariant guards to critical code paths ([#1312](https://github.com/nearai/ironclaw/pull/1312))
|
||||||
|
- *(mcp)* retry after missing session id errors ([#1355](https://github.com/nearai/ironclaw/pull/1355))
|
||||||
|
- *(telegram)* preserve polling after secret-blocked updates ([#1353](https://github.com/nearai/ironclaw/pull/1353))
|
||||||
|
- *(llm)* cap retry-after delays ([#1351](https://github.com/nearai/ironclaw/pull/1351))
|
||||||
|
- *(setup)* remove nonexistent webhook secret command hint ([#1349](https://github.com/nearai/ironclaw/pull/1349))
|
||||||
|
- Rate limiter returns retry after None instead of a duration ([#1269](https://github.com/nearai/ironclaw/pull/1269))
|
||||||
|
|
||||||
|
### Other
|
||||||
|
|
||||||
|
- bump telegram channel version to 0.2.5 ([#1410](https://github.com/nearai/ironclaw/pull/1410))
|
||||||
|
- *(ci)* enforce test requirement for state machine and resilience changes ([#1230](https://github.com/nearai/ironclaw/pull/1230)) ([#1304](https://github.com/nearai/ironclaw/pull/1304))
|
||||||
|
- Fix duplicate LLM responses for matched event routines ([#1275](https://github.com/nearai/ironclaw/pull/1275))
|
||||||
|
- add Japanese README ([#1306](https://github.com/nearai/ironclaw/pull/1306))
|
||||||
|
- *(ci)* add coverage gates via codecov.yml ([#1228](https://github.com/nearai/ironclaw/pull/1228)) ([#1291](https://github.com/nearai/ironclaw/pull/1291))
|
||||||
|
- Redesign routine create requests for LLMs ([#1147](https://github.com/nearai/ironclaw/pull/1147))
|
||||||
|
|
||||||
## [0.19.0](https://github.com/nearai/ironclaw/compare/v0.18.0...v0.19.0) - 2026-03-17
|
## [0.19.0](https://github.com/nearai/ironclaw/compare/v0.18.0...v0.19.0) - 2026-03-17
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|||||||
Generated
+14
-5
@@ -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.52.0",
|
"windows-sys 0.59.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -3390,7 +3390,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ironclaw"
|
name = "ironclaw"
|
||||||
version = "0.19.0"
|
version = "0.23.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"aes-gcm",
|
"aes-gcm",
|
||||||
"aho-corasick",
|
"aho-corasick",
|
||||||
@@ -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",
|
||||||
@@ -3486,8 +3487,16 @@ dependencies = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ironclaw_safety"
|
name = "ironclaw_common"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ironclaw_safety"
|
||||||
|
version = "0.2.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"aho-corasick",
|
"aho-corasick",
|
||||||
"regex",
|
"regex",
|
||||||
@@ -5472,7 +5481,7 @@ dependencies = [
|
|||||||
"errno",
|
"errno",
|
||||||
"libc",
|
"libc",
|
||||||
"linux-raw-sys 0.12.1",
|
"linux-raw-sys 0.12.1",
|
||||||
"windows-sys 0.52.0",
|
"windows-sys 0.59.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -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.52.0",
|
"windows-sys 0.59.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|||||||
+6
-3
@@ -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",
|
||||||
@@ -20,7 +20,7 @@ exclude = [
|
|||||||
|
|
||||||
[package]
|
[package]
|
||||||
name = "ironclaw"
|
name = "ironclaw"
|
||||||
version = "0.19.0"
|
version = "0.23.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
rust-version = "1.92"
|
rust-version = "1.92"
|
||||||
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
||||||
@@ -100,8 +100,11 @@ 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.2.0" }
|
||||||
regex = "1"
|
regex = "1"
|
||||||
aho-corasick = "1"
|
aho-corasick = "1"
|
||||||
|
|
||||||
|
|||||||
@@ -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,17 @@
|
|||||||
|
[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"
|
||||||
|
|
||||||
|
[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>"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "ironclaw_safety"
|
name = "ironclaw_safety"
|
||||||
version = "0.1.0"
|
version = "0.2.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
rust-version = "1.92"
|
rust-version = "1.92"
|
||||||
description = "Prompt injection defense, input validation, secret leak detection, and safety policy enforcement"
|
description = "Prompt injection defense, input validation, secret leak detection, and safety policy enforcement"
|
||||||
@@ -8,7 +8,6 @@ authors = ["NEAR AI <[email protected]>"]
|
|||||||
license = "MIT OR Apache-2.0"
|
license = "MIT OR Apache-2.0"
|
||||||
homepage = "https://github.com/nearai/ironclaw"
|
homepage = "https://github.com/nearai/ironclaw"
|
||||||
repository = "https://github.com/nearai/ironclaw"
|
repository = "https://github.com/nearai/ironclaw"
|
||||||
publish = false
|
|
||||||
|
|
||||||
[package.metadata.dist]
|
[package.metadata.dist]
|
||||||
dist = false
|
dist = false
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "feishu",
|
"name": "feishu",
|
||||||
"display_name": "Feishu / Lark Channel",
|
"display_name": "Feishu / Lark Channel",
|
||||||
"kind": "channel",
|
"kind": "channel",
|
||||||
"version": "0.1.1",
|
"version": "0.1.3",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Talk to your agent through a Feishu or Lark bot",
|
"description": "Talk to your agent through a Feishu or Lark bot",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -19,8 +19,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"sha256": "5fca74022264d1c8e78a0853766276f7ffa3cf0d8065b2f51ca10985acad4714",
|
"sha256": "a66ff0dafb67d2216d8161bb7e96e724a94acb0ab993b85d2782d30412f8fe94",
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/channel-feishu-0.1.1-wasm32-wasip2.tar.gz"
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/channel-feishu-0.1.3-wasm32-wasip2.tar.gz"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/channel-telegram-0.2.4-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/v0.20.0/channel-telegram-0.2.5-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "a7cb300ec1c946831cfceaa95c1dc8f30d0f42a3924f3cb5de8098821573f4b8"
|
"sha256": "1ef20a538f55b379e049356e4d6758006251846bc3365ceaa1c87eba8379a329"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "github",
|
"name": "github",
|
||||||
"display_name": "GitHub",
|
"display_name": "GitHub",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.1",
|
"version": "0.2.2",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "GitHub integration for issues, PRs, repos, and code search",
|
"description": "GitHub integration for issues, PRs, repos, and code search",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -19,8 +19,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-github-0.2.1-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-github-0.2.2-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "92c530b3ad172e2372d819744b5233f1d8f65768e26eb5a6c213eba3ce1de758"
|
"sha256": "70b55af593193d8fa495c0f702ea23284d83a624124f8a5f7564916ec5032c3f"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "gmail",
|
"name": "gmail",
|
||||||
"display_name": "Gmail",
|
"display_name": "Gmail",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Read, send, and manage Gmail messages and threads",
|
"description": "Read, send, and manage Gmail messages and threads",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/gmail-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-gmail-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "ee9574e02e92bc1d481f1310eb88afd99ee52bf6971074ab33bd76bf99b34b1d"
|
"sha256": "79025b40ee70ce1120acc4320bae50da095d7afb0ef67bd56d99b064b72ea779"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "google-calendar",
|
"name": "google-calendar",
|
||||||
"display_name": "Google Calendar",
|
"display_name": "Google Calendar",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Create, read, update, and delete Google Calendar events",
|
"description": "Create, read, update, and delete Google Calendar events",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-calendar-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-calendar-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "2fa47150ea222e787c122182ad6f4dfa2ffaf5fe490d05e8de887a76445f8d2d"
|
"sha256": "86bcc075010b08f5ab2f98f504cec1c6c9e0ca144857d185cbecf72a11f504bf"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "google-docs",
|
"name": "google-docs",
|
||||||
"display_name": "Google Docs",
|
"display_name": "Google Docs",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Create and edit Google Docs documents",
|
"description": "Create and edit Google Docs documents",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-docs-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-docs-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "40e134a1c1564f832ca861c3396895d4e33ec67b99313fc1f97baf8d971423a9"
|
"sha256": "39d476029764949498a53a6a223f9952b5f4df151be7b8b19bf3fe4d401a57cd"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "google-drive",
|
"name": "google-drive",
|
||||||
"display_name": "Google Drive",
|
"display_name": "Google Drive",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Upload, download, search, and manage Google Drive files and folders",
|
"description": "Upload, download, search, and manage Google Drive files and folders",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-drive-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-drive-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "002a341a1d58125563a7c69561b26fbc2629b04ea723cade744102bdc0fbb71f"
|
"sha256": "6e9a700fab93865c852af718666af64c5b534ad6a419fb4b736e07740188f494"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "google-sheets",
|
"name": "google-sheets",
|
||||||
"display_name": "Google Sheets",
|
"display_name": "Google Sheets",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Read and write Google Sheets spreadsheet data",
|
"description": "Read and write Google Sheets spreadsheet data",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-sheets-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-sheets-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "8aa2c9d52f033edea3a6c2311b0ec694ccb6d0a54ef07e94d72bf8be1ce8009a"
|
"sha256": "1f8c381799a916be83263cac9d497d52946e21b1b588592a3a42ca94a73b7051"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "google-slides",
|
"name": "google-slides",
|
||||||
"display_name": "Google Slides",
|
"display_name": "Google Slides",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Create and edit Google Slides presentations",
|
"description": "Create and edit Google Slides presentations",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -17,8 +17,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-slides-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-slides-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "e931a97d4fd0b0b938e464dc7c7f2be6ea6b4d1508f5ea3cd931d44db23f05f5"
|
"sha256": "e2528be5da02f1b8cfc8ee9b0cdd849516c53d412e2f75c6175b3bded7f512cb"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "llm-context",
|
"name": "llm-context",
|
||||||
"display_name": "LLM Context",
|
"display_name": "LLM Context",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.1.0",
|
"version": "0.1.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Fetch pre-extracted web content from Brave Search for grounding LLM answers (RAG, fact-checking)",
|
"description": "Fetch pre-extracted web content from Brave Search for grounding LLM answers (RAG, fact-checking)",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -21,8 +21,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-llm-context-0.1.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-llm-context-0.1.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "d9ced2b1226b879135891e0ee40e072c7c95412e1b2462925a23853e1f92497e"
|
"sha256": "9b19e2fd05dbbbe3c8bd55309a91db09124e8415eb0f767828b6e10b55771e63"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "slack-tool",
|
"name": "slack-tool",
|
||||||
"display_name": "Slack Tool",
|
"display_name": "Slack Tool",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Your agent uses Slack to post and read messages in your workspace",
|
"description": "Your agent uses Slack to post and read messages in your workspace",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -17,8 +17,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-slack-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-slack-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "ccfb0415d7a04f9497726c712d15216de36e86f498b849101283c017f5ab4efb"
|
"sha256": "927519e5b7734beeb022d3b8bbd152e0e6b9f67c9452a8ad47809d3c4221a137"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "telegram-mtproto",
|
"name": "telegram-mtproto",
|
||||||
"display_name": "Telegram Tool",
|
"display_name": "Telegram Tool",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.0",
|
"version": "0.2.1",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Your agent uses your Telegram account to read and send messages",
|
"description": "Your agent uses your Telegram account to read and send messages",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-telegram-0.2.0-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-telegram-0.2.1-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "c17065ca41fae5f2a7c43b36144686718cd310a2f22442313bb1aa82bbad0ae4"
|
"sha256": "1e57d0755fc9c7b3ec013d079f30168898b484a6919f9edd105f0cd80131c1cd"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
"name": "web-search",
|
"name": "web-search",
|
||||||
"display_name": "Web Search",
|
"display_name": "Web Search",
|
||||||
"kind": "tool",
|
"kind": "tool",
|
||||||
"version": "0.2.1",
|
"version": "0.2.2",
|
||||||
"wit_version": "0.3.0",
|
"wit_version": "0.3.0",
|
||||||
"description": "Search the web using Brave Search API",
|
"description": "Search the web using Brave Search API",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
@@ -18,8 +18,8 @@
|
|||||||
},
|
},
|
||||||
"artifacts": {
|
"artifacts": {
|
||||||
"wasm32-wasip2": {
|
"wasm32-wasip2": {
|
||||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-web-search-0.2.1-wasm32-wasip2.tar.gz",
|
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-web-search-0.2.2-wasm32-wasip2.tar.gz",
|
||||||
"sha256": "bad275ca4ec314adea5241d6b92c44ccf9cebcbca8e30ba2493cc0bcb4b57218"
|
"sha256": "47382b50c1ea7525b20d59dc02fab04e336d018665826c2f24710bdf460779ae"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"auth_summary": {
|
"auth_summary": {
|
||||||
|
|||||||
@@ -1,7 +1,2 @@
|
|||||||
[workspace]
|
[workspace]
|
||||||
git_release_enable = false
|
git_release_enable = false
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "ironclaw_safety"
|
|
||||||
publish = false
|
|
||||||
release = false
|
|
||||||
|
|||||||
+226
-48
@@ -13,9 +13,10 @@ use futures::StreamExt;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::context_monitor::ContextMonitor;
|
use crate::agent::context_monitor::ContextMonitor;
|
||||||
use crate::agent::heartbeat::spawn_heartbeat;
|
use crate::agent::heartbeat::{spawn_heartbeat, spawn_multi_user_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,11 @@ 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,
|
||||||
|
/// Per-tenant rate limiting registry (lazily creates rate state per user).
|
||||||
|
pub tenant_rates: Arc<crate::tenant::TenantRateRegistry>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The main agent that coordinates all components.
|
/// The main agent that coordinates all components.
|
||||||
@@ -231,12 +246,15 @@ impl Agent {
|
|||||||
SchedulerDeps {
|
SchedulerDeps {
|
||||||
tools: deps.tools.clone(),
|
tools: deps.tools.clone(),
|
||||||
extension_manager: deps.extension_manager.clone(),
|
extension_manager: deps.extension_manager.clone(),
|
||||||
store: deps.store.clone(),
|
store: deps
|
||||||
|
.store
|
||||||
|
.as_ref()
|
||||||
|
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db))),
|
||||||
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));
|
||||||
@@ -312,6 +330,50 @@ impl Agent {
|
|||||||
&self.deps.cost_guard
|
&self.deps.cost_guard
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Build a tenant-scoped execution context for the given user.
|
||||||
|
///
|
||||||
|
/// This is the standard entry point for per-user operations. The returned
|
||||||
|
/// [`TenantCtx`] provides a [`TenantScope`] that auto-binds `user_id` on
|
||||||
|
/// every database operation and a per-user rate limiter.
|
||||||
|
pub(super) async fn tenant_ctx(&self, user_id: &str) -> crate::tenant::TenantCtx {
|
||||||
|
let rate = self.deps.tenant_rates.get_or_create(user_id).await;
|
||||||
|
|
||||||
|
let store = self
|
||||||
|
.deps
|
||||||
|
.store
|
||||||
|
.as_ref()
|
||||||
|
.map(|db| crate::tenant::TenantScope::new(user_id, Arc::clone(db)));
|
||||||
|
|
||||||
|
// Reuse the owner workspace if user matches, otherwise create per-user.
|
||||||
|
let workspace = match &self.deps.workspace {
|
||||||
|
Some(ws) if ws.user_id() == user_id => Some(Arc::clone(ws)),
|
||||||
|
_ => self
|
||||||
|
.deps
|
||||||
|
.store
|
||||||
|
.as_ref()
|
||||||
|
.map(|db| Arc::new(Workspace::new_with_db(user_id, Arc::clone(db)))),
|
||||||
|
};
|
||||||
|
|
||||||
|
crate::tenant::TenantCtx::new(
|
||||||
|
user_id,
|
||||||
|
store,
|
||||||
|
workspace,
|
||||||
|
Arc::clone(&self.deps.cost_guard),
|
||||||
|
rate,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get an admin-scoped database accessor for cross-tenant operations.
|
||||||
|
///
|
||||||
|
/// Only for system-level components (heartbeat, routine engine, self-repair,
|
||||||
|
/// scheduler). Handler code should use [`tenant_ctx()`](Self::tenant_ctx) instead.
|
||||||
|
pub(super) fn admin_store(&self) -> Option<crate::tenant::AdminScope> {
|
||||||
|
self.deps
|
||||||
|
.store
|
||||||
|
.as_ref()
|
||||||
|
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db)))
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn skill_registry(&self) -> Option<&Arc<std::sync::RwLock<SkillRegistry>>> {
|
pub(super) fn skill_registry(&self) -> Option<&Arc<std::sync::RwLock<SkillRegistry>>> {
|
||||||
self.deps.skill_registry.as_ref()
|
self.deps.skill_registry.as_ref()
|
||||||
}
|
}
|
||||||
@@ -397,8 +459,8 @@ impl Agent {
|
|||||||
self.config.stuck_threshold,
|
self.config.stuck_threshold,
|
||||||
self.config.max_repair_attempts,
|
self.config.max_repair_attempts,
|
||||||
);
|
);
|
||||||
if let Some(ref store) = self.deps.store {
|
if let Some(admin) = self.admin_store() {
|
||||||
self_repair = self_repair.with_store(Arc::clone(store));
|
self_repair = self_repair.with_store(admin);
|
||||||
}
|
}
|
||||||
if let Some(ref builder) = self.deps.builder {
|
if let Some(ref builder) = self.deps.builder {
|
||||||
self_repair = self_repair.with_builder(Arc::clone(builder), Arc::clone(self.tools()));
|
self_repair = self_repair.with_builder(Arc::clone(builder), Arc::clone(self.tools()));
|
||||||
@@ -505,6 +567,7 @@ impl Agent {
|
|||||||
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
|
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
|
||||||
config.quiet_hours_start = hb_config.quiet_hours_start;
|
config.quiet_hours_start = hb_config.quiet_hours_start;
|
||||||
config.quiet_hours_end = hb_config.quiet_hours_end;
|
config.quiet_hours_end = hb_config.quiet_hours_end;
|
||||||
|
config.multi_tenant = hb_config.multi_tenant;
|
||||||
config.timezone = hb_config
|
config.timezone = hb_config
|
||||||
.timezone
|
.timezone
|
||||||
.clone()
|
.clone()
|
||||||
@@ -534,30 +597,52 @@ impl Agent {
|
|||||||
.await;
|
.await;
|
||||||
let notify_user = heartbeat_notify_user;
|
let notify_user = heartbeat_notify_user;
|
||||||
let channels = self.channels.clone();
|
let channels = self.channels.clone();
|
||||||
|
let is_multi_tenant = hb_config.multi_tenant;
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
while let Some(response) = notify_rx.recv().await {
|
while let Some(response) = notify_rx.recv().await {
|
||||||
|
// In multi-tenant mode, extract the owning user_id from
|
||||||
|
// the response metadata so notifications reach the
|
||||||
|
// correct user rather than the agent's owner.
|
||||||
|
// This intentionally overrides the configured notify_target
|
||||||
|
// because each user's heartbeat should notify that user.
|
||||||
|
let effective_user = if is_multi_tenant {
|
||||||
|
response
|
||||||
|
.metadata
|
||||||
|
.get("owner_id")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(String::from)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
// Try the configured channel first, fall back to
|
// Try the configured channel first, fall back to
|
||||||
// broadcasting on all channels.
|
// broadcasting on all channels.
|
||||||
let targeted_ok = if let Some(ref channel) = notify_channel
|
let targeted_ok = if let Some(ref channel) = notify_channel {
|
||||||
&& let Some(ref user) = notify_target
|
let target = effective_user.as_deref().or(notify_target.as_deref());
|
||||||
{
|
if let Some(user) = target {
|
||||||
channels
|
channels
|
||||||
.broadcast(channel, user, response.clone())
|
.broadcast(channel, user, response.clone())
|
||||||
.await
|
.await
|
||||||
.is_ok()
|
.is_ok()
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
false
|
false
|
||||||
};
|
};
|
||||||
|
|
||||||
if !targeted_ok && let Some(ref user) = notify_user {
|
if !targeted_ok {
|
||||||
let results = channels.broadcast_all(user, response).await;
|
let fallback = effective_user.as_deref().or(notify_user.as_deref());
|
||||||
for (ch, result) in results {
|
if let Some(user) = fallback {
|
||||||
if let Err(e) = result {
|
let results = channels.broadcast_all(user, response).await;
|
||||||
tracing::warn!(
|
for (ch, result) in results {
|
||||||
"Failed to broadcast heartbeat to {}: {}",
|
if let Err(e) = result {
|
||||||
ch,
|
tracing::warn!(
|
||||||
e
|
"Failed to broadcast heartbeat to {}: {}",
|
||||||
);
|
ch,
|
||||||
|
e
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -570,14 +655,29 @@ impl Agent {
|
|||||||
.map(|h| h.to_workspace_config())
|
.map(|h| h.to_workspace_config())
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
|
|
||||||
Some(spawn_heartbeat(
|
if config.multi_tenant {
|
||||||
config,
|
if let Some(admin) = self.admin_store() {
|
||||||
hygiene,
|
Some(spawn_multi_user_heartbeat(
|
||||||
workspace.clone(),
|
config,
|
||||||
self.cheap_llm().clone(),
|
hygiene,
|
||||||
Some(notify_tx),
|
self.cheap_llm().clone(),
|
||||||
self.store().map(Arc::clone),
|
Some(notify_tx),
|
||||||
))
|
admin,
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
tracing::warn!("Multi-tenant heartbeat requires a database store");
|
||||||
|
None
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Some(spawn_heartbeat(
|
||||||
|
config,
|
||||||
|
hygiene,
|
||||||
|
workspace.clone(),
|
||||||
|
self.cheap_llm().clone(),
|
||||||
|
Some(notify_tx),
|
||||||
|
self.admin_store(),
|
||||||
|
))
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
tracing::warn!("Heartbeat enabled but no workspace available");
|
tracing::warn!("Heartbeat enabled but no workspace available");
|
||||||
None
|
None
|
||||||
@@ -599,7 +699,7 @@ impl Agent {
|
|||||||
|
|
||||||
let engine = Arc::new(RoutineEngine::new(
|
let engine = Arc::new(RoutineEngine::new(
|
||||||
rt_config.clone(),
|
rt_config.clone(),
|
||||||
Arc::clone(store),
|
crate::tenant::AdminScope::new(Arc::clone(store)),
|
||||||
self.llm().clone(),
|
self.llm().clone(),
|
||||||
Arc::clone(workspace),
|
Arc::clone(workspace),
|
||||||
notify_tx,
|
notify_tx,
|
||||||
@@ -1052,10 +1152,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 +1237,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,15 +1252,30 @@ 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()))
|
||||||
|
};
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Build per-tenant execution context once; threaded through all handlers.
|
||||||
|
let tenant = self.tenant_ctx(&message.user_id).await;
|
||||||
|
|
||||||
|
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 } => {
|
||||||
let mut result = self
|
let mut result = self
|
||||||
.process_user_input(message, session.clone(), thread_id, &content)
|
.process_user_input(
|
||||||
|
message,
|
||||||
|
tenant.clone(),
|
||||||
|
session.clone(),
|
||||||
|
thread_id,
|
||||||
|
&content,
|
||||||
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
// Drain any messages queued during processing.
|
// Drain any messages queued during processing.
|
||||||
@@ -1221,7 +1342,13 @@ impl Agent {
|
|||||||
let mut queued_msg = message.clone();
|
let mut queued_msg = message.clone();
|
||||||
queued_msg.attachments.clear();
|
queued_msg.attachments.clear();
|
||||||
result = self
|
result = self
|
||||||
.process_user_input(&queued_msg, session.clone(), thread_id, &next_content)
|
.process_user_input(
|
||||||
|
&queued_msg,
|
||||||
|
tenant.clone(),
|
||||||
|
session.clone(),
|
||||||
|
thread_id,
|
||||||
|
&next_content,
|
||||||
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
// If processing failed, re-queue the drained content so it
|
// If processing failed, re-queue the drained content so it
|
||||||
@@ -1246,8 +1373,30 @@ 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, &tenant)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
Submission::Undo => self.process_undo(session, thread_id).await,
|
Submission::Undo => self.process_undo(session, thread_id).await,
|
||||||
@@ -1260,12 +1409,9 @@ impl Agent {
|
|||||||
Submission::Summarize => self.process_summarize(session, thread_id).await,
|
Submission::Summarize => self.process_summarize(session, thread_id).await,
|
||||||
Submission::Suggest => self.process_suggest(session, thread_id).await,
|
Submission::Suggest => self.process_suggest(session, thread_id).await,
|
||||||
Submission::JobStatus { job_id } => {
|
Submission::JobStatus { job_id } => {
|
||||||
self.process_job_status(&message.user_id, job_id.as_deref())
|
self.process_job_status(&tenant, job_id.as_deref()).await
|
||||||
.await
|
|
||||||
}
|
|
||||||
Submission::JobCancel { job_id } => {
|
|
||||||
self.process_job_cancel(&message.user_id, &job_id).await
|
|
||||||
}
|
}
|
||||||
|
Submission::JobCancel { job_id } => self.process_job_cancel(&tenant, &job_id).await,
|
||||||
Submission::Quit => return Ok(None),
|
Submission::Quit => return Ok(None),
|
||||||
Submission::SwitchThread { thread_id: target } => {
|
Submission::SwitchThread { thread_id: target } => {
|
||||||
self.process_switch_thread(message, target).await
|
self.process_switch_thread(message, target).await
|
||||||
@@ -1305,7 +1451,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 +1486,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 +1648,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
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+126
-1
@@ -10,7 +10,7 @@ use std::borrow::Cow;
|
|||||||
|
|
||||||
use crate::agent::session::PendingApproval;
|
use crate::agent::session::PendingApproval;
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult};
|
use crate::llm::{ChatMessage, FinishReason, Reasoning, ReasoningContext, RespondResult};
|
||||||
|
|
||||||
/// Signal from the delegate indicating how the loop should proceed.
|
/// Signal from the delegate indicating how the loop should proceed.
|
||||||
pub enum LoopSignal {
|
pub enum LoopSignal {
|
||||||
@@ -134,6 +134,9 @@ pub async fn run_agentic_loop(
|
|||||||
config: &AgenticLoopConfig,
|
config: &AgenticLoopConfig,
|
||||||
) -> Result<LoopOutcome, Error> {
|
) -> Result<LoopOutcome, Error> {
|
||||||
let mut consecutive_tool_intent_nudges: u32 = 0;
|
let mut consecutive_tool_intent_nudges: u32 = 0;
|
||||||
|
// Accumulates across all iterations (not reset by text responses) so
|
||||||
|
// non-consecutive truncations still escalate to force_text.
|
||||||
|
let mut truncation_count: u32 = 0;
|
||||||
|
|
||||||
for iteration in 1..=config.max_iterations {
|
for iteration in 1..=config.max_iterations {
|
||||||
// Check for external signals (stop, cancellation, user messages)
|
// Check for external signals (stop, cancellation, user messages)
|
||||||
@@ -215,7 +218,35 @@ pub async fn run_agentic_loop(
|
|||||||
tool_calls,
|
tool_calls,
|
||||||
content,
|
content,
|
||||||
} => {
|
} => {
|
||||||
|
// If the response was truncated, tool call parameters are likely
|
||||||
|
// incomplete. Discard them and tell the LLM to try a different
|
||||||
|
// approach rather than executing malformed tool calls.
|
||||||
|
if output.finish_reason == FinishReason::Length {
|
||||||
|
truncation_count += 1;
|
||||||
|
let names: Vec<&str> = tool_calls.iter().map(|tc| tc.name.as_str()).collect();
|
||||||
|
tracing::warn!(
|
||||||
|
iteration,
|
||||||
|
tools = ?names,
|
||||||
|
truncation_count,
|
||||||
|
"Discarding truncated tool calls (finish_reason=Length)"
|
||||||
|
);
|
||||||
|
if let Some(ref text) = content {
|
||||||
|
reason_ctx.messages.push(ChatMessage::assistant(text));
|
||||||
|
}
|
||||||
|
reason_ctx
|
||||||
|
.messages
|
||||||
|
.push(ChatMessage::user(crate::llm::TRUNCATED_TOOL_CALL_NOTICE));
|
||||||
|
// After repeated truncations, force text-only mode so the LLM
|
||||||
|
// stops attempting tool calls it can't fit in the output budget.
|
||||||
|
if truncation_count >= 3 {
|
||||||
|
reason_ctx.force_text = true;
|
||||||
|
}
|
||||||
|
delegate.after_iteration(iteration).await;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
consecutive_tool_intent_nudges = 0;
|
consecutive_tool_intent_nudges = 0;
|
||||||
|
truncation_count = 0;
|
||||||
|
|
||||||
if let Some(outcome) = delegate
|
if let Some(outcome) = delegate
|
||||||
.execute_tool_calls(tool_calls, content, reason_ctx)
|
.execute_tool_calls(tool_calls, content, reason_ctx)
|
||||||
@@ -271,6 +302,7 @@ mod tests {
|
|||||||
RespondOutput {
|
RespondOutput {
|
||||||
result: RespondResult::Text(text.to_string()),
|
result: RespondResult::Text(text.to_string()),
|
||||||
usage: zero_usage(),
|
usage: zero_usage(),
|
||||||
|
finish_reason: FinishReason::Stop,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -281,6 +313,7 @@ mod tests {
|
|||||||
content: None,
|
content: None,
|
||||||
},
|
},
|
||||||
usage: zero_usage(),
|
usage: zero_usage(),
|
||||||
|
finish_reason: FinishReason::ToolUse,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -414,6 +447,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]),
|
||||||
@@ -621,4 +655,95 @@ mod tests {
|
|||||||
let result = truncate_for_preview("café", 4);
|
let result = truncate_for_preview("café", 4);
|
||||||
assert_eq!(result, "caf...");
|
assert_eq!(result, "caf...");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_truncated_tool_calls_discarded_on_length() {
|
||||||
|
let truncated_tool_call = ToolCall {
|
||||||
|
id: "call_1".to_string(),
|
||||||
|
name: "memory_write".to_string(),
|
||||||
|
arguments: serde_json::json!({}), // empty — truncated
|
||||||
|
reasoning: None,
|
||||||
|
};
|
||||||
|
let truncated_output = RespondOutput {
|
||||||
|
result: RespondResult::ToolCalls {
|
||||||
|
tool_calls: vec![truncated_tool_call],
|
||||||
|
content: Some("I'll write the report.".to_string()),
|
||||||
|
},
|
||||||
|
usage: zero_usage(),
|
||||||
|
finish_reason: FinishReason::Length, // response was truncated
|
||||||
|
};
|
||||||
|
let delegate = MockDelegate::new(vec![truncated_output, text_output("Summarized it.")]);
|
||||||
|
let reasoning = stub_reasoning();
|
||||||
|
let mut ctx = ReasoningContext::new();
|
||||||
|
let config = AgenticLoopConfig {
|
||||||
|
max_iterations: 5,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Tool calls should NOT have been executed
|
||||||
|
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
|
||||||
|
// The loop should have continued and returned the text response
|
||||||
|
assert!(matches!(outcome, LoopOutcome::Response(ref t) if t == "Summarized it."));
|
||||||
|
// A truncation notice should have been injected into context
|
||||||
|
assert!(
|
||||||
|
ctx.messages
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.role == crate::llm::Role::User && m.content.contains("truncated")),
|
||||||
|
"Should inject truncation notice into context"
|
||||||
|
);
|
||||||
|
// The partial assistant content should have been preserved
|
||||||
|
assert!(
|
||||||
|
ctx.messages
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.role == crate::llm::Role::Assistant
|
||||||
|
&& m.content.contains("write the report")),
|
||||||
|
"Should preserve partial assistant content"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_repeated_truncations_force_text_mode() {
|
||||||
|
let make_truncated = || RespondOutput {
|
||||||
|
result: RespondResult::ToolCalls {
|
||||||
|
tool_calls: vec![ToolCall {
|
||||||
|
id: "call_1".to_string(),
|
||||||
|
name: "memory_write".to_string(),
|
||||||
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
|
}],
|
||||||
|
content: None,
|
||||||
|
},
|
||||||
|
usage: zero_usage(),
|
||||||
|
finish_reason: FinishReason::Length,
|
||||||
|
};
|
||||||
|
// Three truncated responses, then a text response
|
||||||
|
let delegate = MockDelegate::new(vec![
|
||||||
|
make_truncated(),
|
||||||
|
make_truncated(),
|
||||||
|
make_truncated(),
|
||||||
|
text_output("Gave up on tool calls."),
|
||||||
|
]);
|
||||||
|
let reasoning = stub_reasoning();
|
||||||
|
let mut ctx = ReasoningContext::new();
|
||||||
|
let config = AgenticLoopConfig {
|
||||||
|
max_iterations: 5,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(matches!(outcome, LoopOutcome::Response(_)));
|
||||||
|
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
|
||||||
|
// After 3 truncations, force_text should be set
|
||||||
|
assert!(
|
||||||
|
ctx.force_text,
|
||||||
|
"Should escalate to force_text after repeated truncations"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+224
-54
@@ -33,6 +33,7 @@ impl Agent {
|
|||||||
&self,
|
&self,
|
||||||
intent: MessageIntent,
|
intent: MessageIntent,
|
||||||
message: &IncomingMessage,
|
message: &IncomingMessage,
|
||||||
|
tenant: &crate::tenant::TenantCtx,
|
||||||
) -> Result<SubmissionResult, Error> {
|
) -> Result<SubmissionResult, Error> {
|
||||||
// Send thinking status for non-trivial operations
|
// Send thinking status for non-trivial operations
|
||||||
if let MessageIntent::CreateJob { .. } = &intent {
|
if let MessageIntent::CreateJob { .. } = &intent {
|
||||||
@@ -52,24 +53,18 @@ impl Agent {
|
|||||||
description,
|
description,
|
||||||
category,
|
category,
|
||||||
} => {
|
} => {
|
||||||
self.handle_create_job(&message.user_id, title, description, category)
|
self.handle_create_job(tenant, title, description, category)
|
||||||
.await?
|
.await?
|
||||||
}
|
}
|
||||||
MessageIntent::CheckJobStatus { job_id } => {
|
MessageIntent::CheckJobStatus { job_id } => {
|
||||||
self.handle_check_status(&message.user_id, job_id).await?
|
self.handle_check_status(tenant, job_id).await?
|
||||||
}
|
|
||||||
MessageIntent::CancelJob { job_id } => {
|
|
||||||
self.handle_cancel_job(&message.user_id, &job_id).await?
|
|
||||||
}
|
|
||||||
MessageIntent::ListJobs { filter } => {
|
|
||||||
self.handle_list_jobs(&message.user_id, filter).await?
|
|
||||||
}
|
|
||||||
MessageIntent::HelpJob { job_id } => {
|
|
||||||
self.handle_help_job(&message.user_id, &job_id).await?
|
|
||||||
}
|
}
|
||||||
|
MessageIntent::CancelJob { job_id } => self.handle_cancel_job(tenant, &job_id).await?,
|
||||||
|
MessageIntent::ListJobs { filter } => self.handle_list_jobs(tenant, filter).await?,
|
||||||
|
MessageIntent::HelpJob { job_id } => self.handle_help_job(tenant, &job_id).await?,
|
||||||
MessageIntent::Command { command, args } => {
|
MessageIntent::Command { command, args } => {
|
||||||
match self
|
match self
|
||||||
.handle_command(&command, &args, &message.channel)
|
.handle_command(&command, &args, &message.channel, tenant)
|
||||||
.await?
|
.await?
|
||||||
{
|
{
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -83,14 +78,14 @@ impl Agent {
|
|||||||
|
|
||||||
async fn handle_create_job(
|
async fn handle_create_job(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
tenant: &crate::tenant::TenantCtx,
|
||||||
title: String,
|
title: String,
|
||||||
description: String,
|
description: String,
|
||||||
category: Option<String>,
|
category: Option<String>,
|
||||||
) -> Result<String, Error> {
|
) -> Result<String, Error> {
|
||||||
let job_id = self
|
let job_id = self
|
||||||
.scheduler
|
.scheduler
|
||||||
.dispatch_job(user_id, &title, &description, None)
|
.dispatch_job(tenant.user_id(), &title, &description, None)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
// Set the dedicated category field (not stored in metadata)
|
// Set the dedicated category field (not stored in metadata)
|
||||||
@@ -113,7 +108,7 @@ impl Agent {
|
|||||||
|
|
||||||
async fn handle_check_status(
|
async fn handle_check_status(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
tenant: &crate::tenant::TenantCtx,
|
||||||
job_id: Option<String>,
|
job_id: Option<String>,
|
||||||
) -> Result<String, Error> {
|
) -> Result<String, Error> {
|
||||||
match job_id {
|
match job_id {
|
||||||
@@ -122,7 +117,8 @@ impl Agent {
|
|||||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||||
|
|
||||||
// Try DB first for persistent state, fall back to ContextManager.
|
// Try DB first for persistent state, fall back to ContextManager.
|
||||||
if let Some(store) = self.store()
|
// TenantScope.get_job() auto-filters by ownership — no manual check needed.
|
||||||
|
if let Some(store) = tenant.store()
|
||||||
&& let Ok(Some(ctx)) = store.get_job(uuid).await
|
&& let Ok(Some(ctx)) = store.get_job(uuid).await
|
||||||
{
|
{
|
||||||
return Ok(format!(
|
return Ok(format!(
|
||||||
@@ -138,7 +134,7 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let ctx = self.context_manager.get_context(uuid).await?;
|
let ctx = self.context_manager.get_context(uuid).await?;
|
||||||
if ctx.user_id != user_id {
|
if ctx.user_id != tenant.user_id() {
|
||||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -155,7 +151,8 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
// Show summary from DB for consistency with Jobs tab.
|
// Show summary from DB for consistency with Jobs tab.
|
||||||
if let Some(store) = self.store() {
|
// TenantScope methods auto-scope to user — no user_id parameter needed.
|
||||||
|
if let Some(store) = tenant.store() {
|
||||||
let mut total = 0;
|
let mut total = 0;
|
||||||
let mut in_progress = 0;
|
let mut in_progress = 0;
|
||||||
let mut completed = 0;
|
let mut completed = 0;
|
||||||
@@ -183,7 +180,7 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Fallback to ContextManager if no DB.
|
// Fallback to ContextManager if no DB.
|
||||||
let summary = self.context_manager.summary_for(user_id).await;
|
let summary = self.context_manager.summary_for(tenant.user_id()).await;
|
||||||
Ok(format!(
|
Ok(format!(
|
||||||
"Jobs summary: Total: {} In Progress: {} Completed: {} Failed: {} Stuck: {}",
|
"Jobs summary: Total: {} In Progress: {} Completed: {} Failed: {} Stuck: {}",
|
||||||
summary.total,
|
summary.total,
|
||||||
@@ -196,19 +193,24 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
|
async fn handle_cancel_job(
|
||||||
|
&self,
|
||||||
|
tenant: &crate::tenant::TenantCtx,
|
||||||
|
job_id: &str,
|
||||||
|
) -> Result<String, Error> {
|
||||||
let uuid = Uuid::parse_str(job_id)
|
let uuid = Uuid::parse_str(job_id)
|
||||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||||
|
|
||||||
let ctx = self.context_manager.get_context(uuid).await?;
|
let ctx = self.context_manager.get_context(uuid).await?;
|
||||||
if ctx.user_id != user_id {
|
if ctx.user_id != tenant.user_id() {
|
||||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
self.scheduler.stop(uuid).await?;
|
self.scheduler.stop(uuid).await?;
|
||||||
|
|
||||||
// Also update DB so the Jobs tab reflects cancellation immediately.
|
// Also update DB so the Jobs tab reflects cancellation immediately.
|
||||||
if let Some(store) = self.store()
|
// Use TenantScope — ownership already verified above.
|
||||||
|
if let Some(store) = tenant.store()
|
||||||
&& let Err(e) = store
|
&& let Err(e) = store
|
||||||
.update_job_status(uuid, JobState::Cancelled, Some("Cancelled by user"))
|
.update_job_status(uuid, JobState::Cancelled, Some("Cancelled by user"))
|
||||||
.await
|
.await
|
||||||
@@ -221,11 +223,12 @@ impl Agent {
|
|||||||
|
|
||||||
async fn handle_list_jobs(
|
async fn handle_list_jobs(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
tenant: &crate::tenant::TenantCtx,
|
||||||
_filter: Option<String>,
|
_filter: Option<String>,
|
||||||
) -> Result<String, Error> {
|
) -> Result<String, Error> {
|
||||||
// List from DB for consistency with Jobs tab.
|
// List from DB for consistency with Jobs tab.
|
||||||
if let Some(store) = self.store() {
|
// TenantScope methods auto-scope to user.
|
||||||
|
if let Some(store) = tenant.store() {
|
||||||
let agent_jobs = match store.list_agent_jobs().await {
|
let agent_jobs = match store.list_agent_jobs().await {
|
||||||
Ok(jobs) => jobs,
|
Ok(jobs) => jobs,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -256,7 +259,7 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Fallback to ContextManager if no DB.
|
// Fallback to ContextManager if no DB.
|
||||||
let jobs = self.context_manager.all_jobs_for(user_id).await;
|
let jobs = self.context_manager.all_jobs_for(tenant.user_id()).await;
|
||||||
if jobs.is_empty() {
|
if jobs.is_empty() {
|
||||||
return Ok("No jobs found.".to_string());
|
return Ok("No jobs found.".to_string());
|
||||||
}
|
}
|
||||||
@@ -270,12 +273,16 @@ impl Agent {
|
|||||||
Ok(output)
|
Ok(output)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
|
async fn handle_help_job(
|
||||||
|
&self,
|
||||||
|
tenant: &crate::tenant::TenantCtx,
|
||||||
|
job_id: &str,
|
||||||
|
) -> Result<String, Error> {
|
||||||
let uuid = Uuid::parse_str(job_id)
|
let uuid = Uuid::parse_str(job_id)
|
||||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||||
|
|
||||||
let ctx = self.context_manager.get_context(uuid).await?;
|
let ctx = self.context_manager.get_context(uuid).await?;
|
||||||
if ctx.user_id != user_id {
|
if ctx.user_id != tenant.user_id() {
|
||||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -308,11 +315,11 @@ impl Agent {
|
|||||||
/// Show job status inline — either all jobs (no id) or a specific job.
|
/// Show job status inline — either all jobs (no id) or a specific job.
|
||||||
pub(super) async fn process_job_status(
|
pub(super) async fn process_job_status(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
tenant: &crate::tenant::TenantCtx,
|
||||||
job_id: Option<&str>,
|
job_id: Option<&str>,
|
||||||
) -> Result<SubmissionResult, Error> {
|
) -> Result<SubmissionResult, Error> {
|
||||||
match self
|
match self
|
||||||
.handle_check_status(user_id, job_id.map(|s| s.to_string()))
|
.handle_check_status(tenant, job_id.map(|s| s.to_string()))
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(text) => Ok(SubmissionResult::response(text)),
|
Ok(text) => Ok(SubmissionResult::response(text)),
|
||||||
@@ -323,10 +330,10 @@ impl Agent {
|
|||||||
/// Cancel a job by ID.
|
/// Cancel a job by ID.
|
||||||
pub(super) async fn process_job_cancel(
|
pub(super) async fn process_job_cancel(
|
||||||
&self,
|
&self,
|
||||||
user_id: &str,
|
tenant: &crate::tenant::TenantCtx,
|
||||||
job_id: &str,
|
job_id: &str,
|
||||||
) -> Result<SubmissionResult, Error> {
|
) -> Result<SubmissionResult, Error> {
|
||||||
match self.handle_cancel_job(user_id, job_id).await {
|
match self.handle_cancel_job(tenant, job_id).await {
|
||||||
Ok(text) => Ok(SubmissionResult::response(text)),
|
Ok(text) => Ok(SubmissionResult::response(text)),
|
||||||
Err(e) => Ok(SubmissionResult::error(format!("Cancel error: {}", e))),
|
Err(e) => Ok(SubmissionResult::error(format!("Cancel error: {}", e))),
|
||||||
}
|
}
|
||||||
@@ -465,12 +472,101 @@ 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,
|
||||||
command: &str,
|
command: &str,
|
||||||
args: &[String],
|
args: &[String],
|
||||||
channel: &str,
|
channel: &str,
|
||||||
|
tenant: &crate::tenant::TenantCtx,
|
||||||
) -> Result<SubmissionResult, Error> {
|
) -> Result<SubmissionResult, Error> {
|
||||||
match command {
|
match command {
|
||||||
"help" => Ok(SubmissionResult::response(concat!(
|
"help" => Ok(SubmissionResult::response(concat!(
|
||||||
@@ -480,6 +576,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",
|
||||||
@@ -663,19 +760,32 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
match self.llm().set_model(requested) {
|
if self.config.multi_tenant {
|
||||||
Ok(()) => {
|
// Multi-tenant: only persist to per-user DB settings.
|
||||||
// Persist the model choice so it survives restarts.
|
// Do NOT call set_model() on the shared provider — that
|
||||||
self.persist_selected_model(requested).await;
|
// would change the default for all users. The per-request
|
||||||
Ok(SubmissionResult::response(format!(
|
// model_override in the dispatcher reads from the same
|
||||||
"Switched model to: {}",
|
// "selected_model" setting and applies it per-user.
|
||||||
requested
|
self.persist_selected_model(tenant, requested).await;
|
||||||
)))
|
Ok(SubmissionResult::response(format!(
|
||||||
|
"Model preference set to: {} (per-user)",
|
||||||
|
requested
|
||||||
|
)))
|
||||||
|
} else {
|
||||||
|
match self.llm().set_model(requested) {
|
||||||
|
Ok(()) => {
|
||||||
|
// Persist the model choice so it survives restarts.
|
||||||
|
self.persist_selected_model(tenant, requested).await;
|
||||||
|
Ok(SubmissionResult::response(format!(
|
||||||
|
"Switched model to: {}",
|
||||||
|
requested
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
Err(e) => Ok(SubmissionResult::error(format!(
|
||||||
|
"Failed to switch model: {}",
|
||||||
|
e
|
||||||
|
))),
|
||||||
}
|
}
|
||||||
Err(e) => Ok(SubmissionResult::error(format!(
|
|
||||||
"Failed to switch model: {}",
|
|
||||||
e
|
|
||||||
))),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -817,10 +927,14 @@ impl Agent {
|
|||||||
command: &str,
|
command: &str,
|
||||||
args: &[String],
|
args: &[String],
|
||||||
channel: &str,
|
channel: &str,
|
||||||
|
tenant: &crate::tenant::TenantCtx,
|
||||||
) -> Result<Option<String>, Error> {
|
) -> Result<Option<String>, Error> {
|
||||||
// System commands are now handled directly via Submission::SystemCommand,
|
// System commands are now handled directly via Submission::SystemCommand,
|
||||||
// but the router may still send us unknown /commands.
|
// but the router may still send us unknown /commands.
|
||||||
match self.handle_system_command(command, args, channel).await? {
|
match self
|
||||||
|
.handle_system_command(command, args, channel, tenant)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
SubmissionResult::Response { content } => Ok(Some(content)),
|
SubmissionResult::Response { content } => Ok(Some(content)),
|
||||||
SubmissionResult::Ok { message } => Ok(message),
|
SubmissionResult::Ok { message } => Ok(message),
|
||||||
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
||||||
@@ -832,21 +946,69 @@ impl Agent {
|
|||||||
///
|
///
|
||||||
/// Best-effort: logs warnings on failure but does not propagate errors,
|
/// Best-effort: logs warnings on failure but does not propagate errors,
|
||||||
/// since the in-memory model switch already succeeded.
|
/// since the in-memory model switch already succeeded.
|
||||||
async fn persist_selected_model(&self, model: &str) {
|
///
|
||||||
// 1. Persist to DB if available.
|
/// In multi-tenant mode, only the per-user DB setting is written — global
|
||||||
if let Some(store) = self.store() {
|
/// .env and TOML files are shared across users and must not be mutated.
|
||||||
|
async fn persist_selected_model(&self, tenant: &crate::tenant::TenantCtx, model: &str) {
|
||||||
|
// 1. Persist to DB if available (per-user scoped via TenantScope).
|
||||||
|
if let Some(store) = tenant.store() {
|
||||||
let value = serde_json::Value::String(model.to_string());
|
let value = serde_json::Value::String(model.to_string());
|
||||||
if let Err(e) = store
|
if let Err(e) = store.set_setting("selected_model", &value).await {
|
||||||
.set_setting(self.owner_id(), "selected_model", &value)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
tracing::warn!("Failed to persist model to DB: {}", e);
|
tracing::warn!("Failed to persist model to DB: {}", e);
|
||||||
|
} else {
|
||||||
|
tracing::debug!(
|
||||||
|
user_id = tenant.user_id(),
|
||||||
|
"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. In multi-tenant mode, skip .env/TOML writes — these are global
|
||||||
|
// files shared by all users. The per-user DB setting is sufficient.
|
||||||
|
if self.config.multi_tenant {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. 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 +1018,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 +1035,7 @@ impl Agent {
|
|||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
tracing::warn!("Model TOML persistence task failed: {}", e);
|
tracing::warn!("Model persistence task failed: {}", e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+236
-3
@@ -21,6 +21,9 @@ pub struct CostGuardConfig {
|
|||||||
pub max_cost_per_day_cents: Option<u64>,
|
pub max_cost_per_day_cents: Option<u64>,
|
||||||
/// Maximum LLM calls per hour. None = unlimited.
|
/// Maximum LLM calls per hour. None = unlimited.
|
||||||
pub max_actions_per_hour: Option<u64>,
|
pub max_actions_per_hour: Option<u64>,
|
||||||
|
/// Maximum spend per user per day in cents. None = unlimited.
|
||||||
|
/// Applied independently per user alongside the global budget.
|
||||||
|
pub max_cost_per_user_per_day_cents: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Error returned when a cost limit is exceeded.
|
/// Error returned when a cost limit is exceeded.
|
||||||
@@ -30,6 +33,12 @@ pub enum CostLimitExceeded {
|
|||||||
DailyBudget { spent_cents: u64, limit_cents: u64 },
|
DailyBudget { spent_cents: u64, limit_cents: u64 },
|
||||||
/// Hourly action rate limit reached.
|
/// Hourly action rate limit reached.
|
||||||
HourlyRate { actions: u64, limit: u64 },
|
HourlyRate { actions: u64, limit: u64 },
|
||||||
|
/// Per-user daily spending cap reached.
|
||||||
|
UserDailyBudget {
|
||||||
|
user_id: String,
|
||||||
|
spent_cents: u64,
|
||||||
|
limit_cents: u64,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
impl std::fmt::Display for CostLimitExceeded {
|
impl std::fmt::Display for CostLimitExceeded {
|
||||||
@@ -49,6 +58,17 @@ impl std::fmt::Display for CostLimitExceeded {
|
|||||||
"Hourly action limit exceeded: {} actions of {} allowed per hour",
|
"Hourly action limit exceeded: {} actions of {} allowed per hour",
|
||||||
actions, limit
|
actions, limit
|
||||||
),
|
),
|
||||||
|
Self::UserDailyBudget {
|
||||||
|
user_id,
|
||||||
|
spent_cents,
|
||||||
|
limit_cents,
|
||||||
|
} => write!(
|
||||||
|
f,
|
||||||
|
"User '{}' daily cost limit exceeded: spent ${:.2} of ${:.2} allowed",
|
||||||
|
user_id,
|
||||||
|
*spent_cents as f64 / 100.0,
|
||||||
|
*limit_cents as f64 / 100.0
|
||||||
|
),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -78,6 +98,9 @@ pub struct CostGuard {
|
|||||||
|
|
||||||
/// Per-model token usage since startup.
|
/// Per-model token usage since startup.
|
||||||
model_tokens: Mutex<HashMap<String, ModelTokens>>,
|
model_tokens: Mutex<HashMap<String, ModelTokens>>,
|
||||||
|
|
||||||
|
/// Per-user daily cost tracking. Each entry resets independently at midnight UTC.
|
||||||
|
per_user_daily_cost: Mutex<HashMap<String, DailyCost>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
struct DailyCost {
|
struct DailyCost {
|
||||||
@@ -97,6 +120,7 @@ impl CostGuard {
|
|||||||
action_window: Mutex::new(VecDeque::new()),
|
action_window: Mutex::new(VecDeque::new()),
|
||||||
budget_exceeded: AtomicBool::new(false),
|
budget_exceeded: AtomicBool::new(false),
|
||||||
model_tokens: Mutex::new(HashMap::new()),
|
model_tokens: Mutex::new(HashMap::new()),
|
||||||
|
per_user_daily_cost: Mutex::new(HashMap::new()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -203,6 +227,11 @@ impl CostGuard {
|
|||||||
daily.reset_date = today;
|
daily.reset_date = today;
|
||||||
self.budget_exceeded.store(false, Ordering::Relaxed);
|
self.budget_exceeded.store(false, Ordering::Relaxed);
|
||||||
tracing::info!("Cost guard: daily counter reset for {}", today);
|
tracing::info!("Cost guard: daily counter reset for {}", today);
|
||||||
|
|
||||||
|
// Prune per-user entries from previous days to prevent
|
||||||
|
// unbounded HashMap growth in long-lived deployments.
|
||||||
|
let mut per_user = self.per_user_daily_cost.lock().await;
|
||||||
|
per_user.retain(|_, entry| entry.reset_date == today);
|
||||||
}
|
}
|
||||||
daily.total += cost;
|
daily.total += cost;
|
||||||
|
|
||||||
@@ -248,6 +277,85 @@ impl CostGuard {
|
|||||||
cost
|
cost
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Record an LLM call with per-user attribution.
|
||||||
|
///
|
||||||
|
/// Delegates to `record_llm_call` for global tracking, then additionally
|
||||||
|
/// records the cost against the user's daily budget.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub async fn record_llm_call_for_user(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
model: &str,
|
||||||
|
input_tokens: u32,
|
||||||
|
output_tokens: u32,
|
||||||
|
cache_read_input_tokens: u32,
|
||||||
|
cache_creation_input_tokens: u32,
|
||||||
|
cache_read_discount: Decimal,
|
||||||
|
cache_write_multiplier: Decimal,
|
||||||
|
cost_per_token: Option<(Decimal, Decimal)>,
|
||||||
|
) -> Decimal {
|
||||||
|
let cost = self
|
||||||
|
.record_llm_call(
|
||||||
|
model,
|
||||||
|
input_tokens,
|
||||||
|
output_tokens,
|
||||||
|
cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens,
|
||||||
|
cache_read_discount,
|
||||||
|
cache_write_multiplier,
|
||||||
|
cost_per_token,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Track per-user daily cost
|
||||||
|
{
|
||||||
|
let today = chrono::Utc::now().date_naive();
|
||||||
|
let mut per_user = self.per_user_daily_cost.lock().await;
|
||||||
|
let entry = per_user
|
||||||
|
.entry(user_id.to_string())
|
||||||
|
.or_insert_with(|| DailyCost {
|
||||||
|
total: Decimal::ZERO,
|
||||||
|
reset_date: today,
|
||||||
|
});
|
||||||
|
if today != entry.reset_date {
|
||||||
|
entry.total = Decimal::ZERO;
|
||||||
|
entry.reset_date = today;
|
||||||
|
}
|
||||||
|
entry.total += cost;
|
||||||
|
}
|
||||||
|
|
||||||
|
cost
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether the next action is allowed for a specific user.
|
||||||
|
///
|
||||||
|
/// Checks the global limits first (via `check_allowed`), then additionally
|
||||||
|
/// checks the per-user daily budget if configured.
|
||||||
|
pub async fn check_allowed_for_user(&self, user_id: &str) -> Result<(), CostLimitExceeded> {
|
||||||
|
// Check global limits first
|
||||||
|
self.check_allowed().await?;
|
||||||
|
|
||||||
|
// Check per-user daily budget
|
||||||
|
if let Some(limit_cents) = self.config.max_cost_per_user_per_day_cents {
|
||||||
|
let today = chrono::Utc::now().date_naive();
|
||||||
|
let per_user = self.per_user_daily_cost.lock().await;
|
||||||
|
if let Some(entry) = per_user.get(user_id)
|
||||||
|
&& entry.reset_date == today
|
||||||
|
{
|
||||||
|
let spent_cents = to_cents(entry.total);
|
||||||
|
if spent_cents >= limit_cents {
|
||||||
|
return Err(CostLimitExceeded::UserDailyBudget {
|
||||||
|
user_id: user_id.to_string(),
|
||||||
|
spent_cents,
|
||||||
|
limit_cents,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Current daily spend in USD (as Decimal).
|
/// Current daily spend in USD (as Decimal).
|
||||||
pub async fn daily_spend(&self) -> Decimal {
|
pub async fn daily_spend(&self) -> Decimal {
|
||||||
let daily = self.daily_cost.lock().await;
|
let daily = self.daily_cost.lock().await;
|
||||||
@@ -259,6 +367,16 @@ impl CostGuard {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Current daily spend for a specific user in USD (as Decimal).
|
||||||
|
pub async fn daily_spend_for_user(&self, user_id: &str) -> Decimal {
|
||||||
|
let today = chrono::Utc::now().date_naive();
|
||||||
|
let per_user = self.per_user_daily_cost.lock().await;
|
||||||
|
match per_user.get(user_id) {
|
||||||
|
Some(entry) if entry.reset_date == today => entry.total,
|
||||||
|
_ => Decimal::ZERO,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Number of actions in the current hourly window.
|
/// Number of actions in the current hourly window.
|
||||||
pub async fn actions_this_hour(&self) -> u64 {
|
pub async fn actions_this_hour(&self) -> u64 {
|
||||||
let mut window = self.action_window.lock().await;
|
let mut window = self.action_window.lock().await;
|
||||||
@@ -314,7 +432,7 @@ mod tests {
|
|||||||
async fn test_daily_budget_enforcement() {
|
async fn test_daily_budget_enforcement() {
|
||||||
let guard = CostGuard::new(CostGuardConfig {
|
let guard = CostGuard::new(CostGuardConfig {
|
||||||
max_cost_per_day_cents: Some(1), // $0.01 limit
|
max_cost_per_day_cents: Some(1), // $0.01 limit
|
||||||
max_actions_per_hour: None,
|
..CostGuardConfig::default()
|
||||||
});
|
});
|
||||||
|
|
||||||
// First call allowed
|
// First call allowed
|
||||||
@@ -350,8 +468,8 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_hourly_rate_enforcement() {
|
async fn test_hourly_rate_enforcement() {
|
||||||
let guard = CostGuard::new(CostGuardConfig {
|
let guard = CostGuard::new(CostGuardConfig {
|
||||||
max_cost_per_day_cents: None,
|
|
||||||
max_actions_per_hour: Some(3),
|
max_actions_per_hour: Some(3),
|
||||||
|
..CostGuardConfig::default()
|
||||||
});
|
});
|
||||||
|
|
||||||
// First 3 actions allowed
|
// First 3 actions allowed
|
||||||
@@ -633,8 +751,8 @@ mod tests {
|
|||||||
// A fresh CostGuard with rate limits should not panic even if
|
// A fresh CostGuard with rate limits should not panic even if
|
||||||
// checked_sub returns None (simulating short uptime).
|
// checked_sub returns None (simulating short uptime).
|
||||||
let guard = CostGuard::new(CostGuardConfig {
|
let guard = CostGuard::new(CostGuardConfig {
|
||||||
max_cost_per_day_cents: None,
|
|
||||||
max_actions_per_hour: Some(100),
|
max_actions_per_hour: Some(100),
|
||||||
|
..CostGuardConfig::default()
|
||||||
});
|
});
|
||||||
|
|
||||||
// These must not panic regardless of system uptime
|
// These must not panic regardless of system uptime
|
||||||
@@ -656,4 +774,119 @@ mod tests {
|
|||||||
let result = Instant::now().checked_sub(std::time::Duration::MAX);
|
let result = Instant::now().checked_sub(std::time::Duration::MAX);
|
||||||
assert!(result.is_none());
|
assert!(result.is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_per_user_daily_budget_enforcement() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig {
|
||||||
|
max_cost_per_day_cents: None,
|
||||||
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
|
||||||
|
});
|
||||||
|
|
||||||
|
// Both users initially allowed
|
||||||
|
assert!(guard.check_allowed_for_user("alice").await.is_ok());
|
||||||
|
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||||
|
|
||||||
|
// Alice makes an expensive call
|
||||||
|
guard
|
||||||
|
.record_llm_call_for_user(
|
||||||
|
"alice",
|
||||||
|
"gpt-4o",
|
||||||
|
10_000,
|
||||||
|
10_000,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Alice should be blocked, Bob should still be allowed
|
||||||
|
let result = guard.check_allowed_for_user("alice").await;
|
||||||
|
assert!(result.is_err());
|
||||||
|
match result.unwrap_err() {
|
||||||
|
CostLimitExceeded::UserDailyBudget {
|
||||||
|
user_id,
|
||||||
|
limit_cents,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
assert_eq!(user_id, "alice");
|
||||||
|
assert_eq!(limit_cents, 1);
|
||||||
|
}
|
||||||
|
other => panic!("Expected UserDailyBudget, got {:?}", other),
|
||||||
|
}
|
||||||
|
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_per_user_daily_spend_tracking() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
assert_eq!(guard.daily_spend_for_user("alice").await, Decimal::ZERO);
|
||||||
|
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
|
||||||
|
|
||||||
|
let cost = guard
|
||||||
|
.record_llm_call_for_user(
|
||||||
|
"alice",
|
||||||
|
"gpt-4o",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert_eq!(guard.daily_spend_for_user("alice").await, cost);
|
||||||
|
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
|
||||||
|
// Global spend should also be tracked
|
||||||
|
assert_eq!(guard.daily_spend().await, cost);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_per_user_budget_independent_of_global() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig {
|
||||||
|
max_cost_per_day_cents: Some(100_000), // $1000 global limit
|
||||||
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
|
||||||
|
});
|
||||||
|
|
||||||
|
// User hits their personal limit
|
||||||
|
guard
|
||||||
|
.record_llm_call_for_user(
|
||||||
|
"alice",
|
||||||
|
"gpt-4o",
|
||||||
|
10_000,
|
||||||
|
10_000,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Alice blocked by per-user limit, not global
|
||||||
|
assert!(guard.check_allowed_for_user("alice").await.is_err());
|
||||||
|
// Global limit is far from reached
|
||||||
|
assert!(guard.check_allowed().await.is_ok());
|
||||||
|
// Bob is unaffected
|
||||||
|
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_user_cost_limit_display() {
|
||||||
|
let limit = CostLimitExceeded::UserDailyBudget {
|
||||||
|
user_id: "alice".to_string(),
|
||||||
|
spent_cents: 150,
|
||||||
|
limit_cents: 100,
|
||||||
|
};
|
||||||
|
let msg = limit.to_string();
|
||||||
|
assert!(msg.contains("alice"));
|
||||||
|
assert!(msg.contains("$1.50"));
|
||||||
|
assert!(msg.contains("$1.00"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+165
-19
@@ -42,6 +42,7 @@ impl Agent {
|
|||||||
pub(super) async fn run_agentic_loop(
|
pub(super) async fn run_agentic_loop(
|
||||||
&self,
|
&self,
|
||||||
message: &IncomingMessage,
|
message: &IncomingMessage,
|
||||||
|
tenant: crate::tenant::TenantCtx,
|
||||||
session: Arc<Mutex<Session>>,
|
session: Arc<Mutex<Session>>,
|
||||||
thread_id: Uuid,
|
thread_id: Uuid,
|
||||||
initial_messages: Vec<ChatMessage>,
|
initial_messages: Vec<ChatMessage>,
|
||||||
@@ -63,7 +64,12 @@ impl Agent {
|
|||||||
);
|
);
|
||||||
|
|
||||||
let system_prompt = if let Some(ws) = self.workspace() {
|
let system_prompt = if let Some(ws) = self.workspace() {
|
||||||
match ws
|
let scoped_workspace = if ws.user_id() == message.user_id {
|
||||||
|
Arc::clone(ws)
|
||||||
|
} else {
|
||||||
|
Arc::new(ws.scoped_to_user(&message.user_id))
|
||||||
|
};
|
||||||
|
match scoped_workspace
|
||||||
.system_prompt_for_context_tz(is_group_chat, user_tz)
|
.system_prompt_for_context_tz(is_group_chat, user_tz)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
@@ -163,6 +169,7 @@ impl Agent {
|
|||||||
|
|
||||||
let delegate = ChatDelegate {
|
let delegate = ChatDelegate {
|
||||||
agent: self,
|
agent: self,
|
||||||
|
tenant,
|
||||||
session: session.clone(),
|
session: session.clone(),
|
||||||
thread_id,
|
thread_id,
|
||||||
message,
|
message,
|
||||||
@@ -235,6 +242,7 @@ impl Agent {
|
|||||||
/// auth intercept, and cost tracking.
|
/// auth intercept, and cost tracking.
|
||||||
struct ChatDelegate<'a> {
|
struct ChatDelegate<'a> {
|
||||||
agent: &'a Agent,
|
agent: &'a Agent,
|
||||||
|
tenant: crate::tenant::TenantCtx,
|
||||||
session: Arc<Mutex<Session>>,
|
session: Arc<Mutex<Session>>,
|
||||||
thread_id: Uuid,
|
thread_id: Uuid,
|
||||||
message: &'a IncomingMessage,
|
message: &'a IncomingMessage,
|
||||||
@@ -298,6 +306,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
|
|
||||||
// Update context for this iteration
|
// Update context for this iteration
|
||||||
reason_ctx.available_tools = tool_defs;
|
reason_ctx.available_tools = tool_defs;
|
||||||
|
// Preserve force_text if already set (e.g. by truncation escalation).
|
||||||
|
let force_text = force_text || reason_ctx.force_text;
|
||||||
reason_ctx.system_prompt = Some(if force_text {
|
reason_ctx.system_prompt = Some(if force_text {
|
||||||
self.cached_prompt_no_tools.clone()
|
self.cached_prompt_no_tools.clone()
|
||||||
} else {
|
} else {
|
||||||
@@ -331,8 +341,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
reason_ctx: &mut ReasoningContext,
|
reason_ctx: &mut ReasoningContext,
|
||||||
iteration: usize,
|
iteration: usize,
|
||||||
) -> Result<crate::llm::RespondOutput, Error> {
|
) -> Result<crate::llm::RespondOutput, Error> {
|
||||||
// Enforce cost guardrails before the LLM call
|
// Enforce cost guardrails before the LLM call (global + per-user)
|
||||||
if let Err(limit) = self.agent.cost_guard().check_allowed().await {
|
if let Err(limit) = self.tenant.check_cost_allowed().await {
|
||||||
return Err(crate::error::LlmError::InvalidResponse {
|
return Err(crate::error::LlmError::InvalidResponse {
|
||||||
provider: "agent".to_string(),
|
provider: "agent".to_string(),
|
||||||
reason: limit.to_string(),
|
reason: limit.to_string(),
|
||||||
@@ -340,6 +350,21 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
.into());
|
.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Apply per-user model override from settings (first iteration only
|
||||||
|
// to avoid repeated DB lookups within the same agentic loop).
|
||||||
|
// Uses "selected_model" — the same key the /model command persists to
|
||||||
|
// via SettingsStore (per-user scoped via TenantScope).
|
||||||
|
if iteration == 0
|
||||||
|
&& let Some(store) = self.tenant.store()
|
||||||
|
&& let Ok(Some(value)) = store.get_setting("selected_model").await
|
||||||
|
&& let Some(model) = value.as_str()
|
||||||
|
{
|
||||||
|
let model = model.trim();
|
||||||
|
if !model.is_empty() {
|
||||||
|
reason_ctx.model_override = Some(model.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let output = match reasoning.respond_with_tools(reason_ctx).await {
|
let output = match reasoning.respond_with_tools(reason_ctx).await {
|
||||||
Ok(output) => output,
|
Ok(output) => output,
|
||||||
Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => {
|
Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => {
|
||||||
@@ -374,13 +399,22 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
Err(e) => return Err(e.into()),
|
Err(e) => return Err(e.into()),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Record cost and track token usage
|
// Record cost and track token usage (global + per-user).
|
||||||
let model_name = self.agent.llm().active_model_name();
|
// When a model override is active, use the override name for attribution
|
||||||
|
// and let CostGuard look up pricing via costs::model_cost() instead of
|
||||||
|
// using the default provider's cost_per_token (which reflects the wrong model).
|
||||||
|
let (model_name, cost_per_token) = if let Some(ref ovr) = reason_ctx.model_override {
|
||||||
|
(ovr.clone(), None)
|
||||||
|
} else {
|
||||||
|
(
|
||||||
|
self.agent.llm().active_model_name(),
|
||||||
|
Some(self.agent.llm().cost_per_token()),
|
||||||
|
)
|
||||||
|
};
|
||||||
let read_discount = self.agent.llm().cache_read_discount();
|
let read_discount = self.agent.llm().cache_read_discount();
|
||||||
let write_multiplier = self.agent.llm().cache_write_multiplier();
|
let write_multiplier = self.agent.llm().cache_write_multiplier();
|
||||||
let call_cost = self
|
let call_cost = self
|
||||||
.agent
|
.tenant
|
||||||
.cost_guard()
|
|
||||||
.record_llm_call(
|
.record_llm_call(
|
||||||
&model_name,
|
&model_name,
|
||||||
output.usage.input_tokens,
|
output.usage.input_tokens,
|
||||||
@@ -389,7 +423,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
|||||||
output.usage.cache_creation_input_tokens,
|
output.usage.cache_creation_input_tokens,
|
||||||
read_discount,
|
read_discount,
|
||||||
write_multiplier,
|
write_multiplier,
|
||||||
Some(self.agent.llm().cost_per_token()),
|
cost_per_token,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
@@ -420,6 +454,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 +487,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 +537,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 +823,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 +949,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 +1198,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 +1333,8 @@ 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(),
|
||||||
|
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -1240,10 +1350,14 @@ mod tests {
|
|||||||
allow_local_tools: false,
|
allow_local_tools: false,
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
max_tool_iterations: 50,
|
max_tool_iterations: 50,
|
||||||
auto_approve_tools: false,
|
auto_approve_tools: false,
|
||||||
default_timezone: "UTC".to_string(),
|
default_timezone: "UTC".to_string(),
|
||||||
max_tokens_per_job: 0,
|
max_tokens_per_job: 0,
|
||||||
|
multi_tenant: false,
|
||||||
|
max_llm_concurrent_per_user: None,
|
||||||
|
max_jobs_concurrent_per_user: None,
|
||||||
},
|
},
|
||||||
deps,
|
deps,
|
||||||
Arc::new(ChannelManager::new()),
|
Arc::new(ChannelManager::new()),
|
||||||
@@ -1453,11 +1567,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 +1759,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 +1852,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 +1892,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 +2023,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 +2177,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 +2214,8 @@ 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(),
|
||||||
|
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -2107,10 +2231,14 @@ mod tests {
|
|||||||
allow_local_tools: false,
|
allow_local_tools: false,
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
max_tool_iterations,
|
max_tool_iterations,
|
||||||
auto_approve_tools: true,
|
auto_approve_tools: true,
|
||||||
default_timezone: "UTC".to_string(),
|
default_timezone: "UTC".to_string(),
|
||||||
max_tokens_per_job: 0,
|
max_tokens_per_job: 0,
|
||||||
|
multi_tenant: false,
|
||||||
|
max_llm_concurrent_per_user: None,
|
||||||
|
max_jobs_concurrent_per_user: None,
|
||||||
},
|
},
|
||||||
deps,
|
deps,
|
||||||
Arc::new(ChannelManager::new()),
|
Arc::new(ChannelManager::new()),
|
||||||
@@ -2145,13 +2273,14 @@ mod tests {
|
|||||||
|
|
||||||
let message = IncomingMessage::new("test", "test-user", "do something");
|
let message = IncomingMessage::new("test", "test-user", "do something");
|
||||||
let initial_messages = vec![ChatMessage::user("do something")];
|
let initial_messages = vec![ChatMessage::user("do something")];
|
||||||
|
let tenant = agent.tenant_ctx("test-user").await;
|
||||||
|
|
||||||
// The dispatcher must terminate within 5 seconds. If there is an
|
// The dispatcher must terminate within 5 seconds. If there is an
|
||||||
// infinite loop bug (e.g., index not advancing on tool failure), the
|
// infinite loop bug (e.g., index not advancing on tool failure), the
|
||||||
// timeout will fire and the test will fail.
|
// timeout will fire and the test will fail.
|
||||||
let result = tokio::time::timeout(
|
let result = tokio::time::timeout(
|
||||||
Duration::from_secs(5),
|
Duration::from_secs(5),
|
||||||
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
|
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -2212,6 +2341,8 @@ 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(),
|
||||||
|
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||||
};
|
};
|
||||||
|
|
||||||
Agent::new(
|
Agent::new(
|
||||||
@@ -2227,10 +2358,14 @@ mod tests {
|
|||||||
allow_local_tools: false,
|
allow_local_tools: false,
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
max_tool_iterations: max_iter,
|
max_tool_iterations: max_iter,
|
||||||
auto_approve_tools: true,
|
auto_approve_tools: true,
|
||||||
default_timezone: "UTC".to_string(),
|
default_timezone: "UTC".to_string(),
|
||||||
max_tokens_per_job: 0,
|
max_tokens_per_job: 0,
|
||||||
|
multi_tenant: false,
|
||||||
|
max_llm_concurrent_per_user: None,
|
||||||
|
max_jobs_concurrent_per_user: None,
|
||||||
},
|
},
|
||||||
deps,
|
deps,
|
||||||
Arc::new(ChannelManager::new()),
|
Arc::new(ChannelManager::new()),
|
||||||
@@ -2250,13 +2385,14 @@ mod tests {
|
|||||||
|
|
||||||
let message = IncomingMessage::new("test", "test-user", "keep calling tools");
|
let message = IncomingMessage::new("test", "test-user", "keep calling tools");
|
||||||
let initial_messages = vec![ChatMessage::user("keep calling tools")];
|
let initial_messages = vec![ChatMessage::user("keep calling tools")];
|
||||||
|
let tenant = agent.tenant_ctx("test-user").await;
|
||||||
|
|
||||||
// Even with an LLM that always wants to call tools, the dispatcher
|
// Even with an LLM that always wants to call tools, the dispatcher
|
||||||
// must terminate within the timeout thanks to force_text at
|
// must terminate within the timeout thanks to force_text at
|
||||||
// max_tool_iterations.
|
// max_tool_iterations.
|
||||||
let result = tokio::time::timeout(
|
let result = tokio::time::timeout(
|
||||||
Duration::from_secs(5),
|
Duration::from_secs(5),
|
||||||
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
|
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -2345,6 +2481,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>";
|
||||||
|
|||||||
+184
-7
@@ -31,8 +31,8 @@ use chrono_tz::Tz;
|
|||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
use crate::channels::OutgoingResponse;
|
use crate::channels::OutgoingResponse;
|
||||||
use crate::db::Database;
|
|
||||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||||
|
use crate::tenant::AdminScope;
|
||||||
use crate::workspace::Workspace;
|
use crate::workspace::Workspace;
|
||||||
use crate::workspace::hygiene::HygieneConfig;
|
use crate::workspace::hygiene::HygieneConfig;
|
||||||
|
|
||||||
@@ -57,6 +57,9 @@ pub struct HeartbeatConfig {
|
|||||||
pub quiet_hours_end: Option<u32>,
|
pub quiet_hours_end: Option<u32>,
|
||||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||||
pub timezone: Option<String>,
|
pub timezone: Option<String>,
|
||||||
|
/// When true, cycle through all users with routines instead of
|
||||||
|
/// running heartbeat for a single user. Requires a database store.
|
||||||
|
pub multi_tenant: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for HeartbeatConfig {
|
impl Default for HeartbeatConfig {
|
||||||
@@ -71,6 +74,7 @@ impl Default for HeartbeatConfig {
|
|||||||
quiet_hours_start: None,
|
quiet_hours_start: None,
|
||||||
quiet_hours_end: None,
|
quiet_hours_end: None,
|
||||||
timezone: None,
|
timezone: None,
|
||||||
|
multi_tenant: false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -178,7 +182,7 @@ pub struct HeartbeatRunner {
|
|||||||
workspace: Arc<Workspace>,
|
workspace: Arc<Workspace>,
|
||||||
llm: Arc<dyn LlmProvider>,
|
llm: Arc<dyn LlmProvider>,
|
||||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<AdminScope>,
|
||||||
consecutive_failures: u32,
|
consecutive_failures: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -207,8 +211,8 @@ impl HeartbeatRunner {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the database store for persistent heartbeat conversations.
|
/// Set the admin-scoped database store for persistent heartbeat conversations.
|
||||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
pub fn with_store(mut self, store: AdminScope) -> Self {
|
||||||
self.store = Some(store);
|
self.store = Some(store);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
@@ -396,7 +400,7 @@ impl HeartbeatRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Send a notification about heartbeat findings.
|
/// Send a notification about heartbeat findings.
|
||||||
async fn send_notification(&self, message: &str) {
|
pub(crate) async fn send_notification(&self, message: &str) {
|
||||||
let Some(ref tx) = self.response_tx else {
|
let Some(ref tx) = self.response_tx else {
|
||||||
tracing::debug!("No response channel configured for heartbeat notifications");
|
tracing::debug!("No response channel configured for heartbeat notifications");
|
||||||
return;
|
return;
|
||||||
@@ -493,7 +497,7 @@ pub fn spawn_heartbeat(
|
|||||||
workspace: Arc<Workspace>,
|
workspace: Arc<Workspace>,
|
||||||
llm: Arc<dyn LlmProvider>,
|
llm: Arc<dyn LlmProvider>,
|
||||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<AdminScope>,
|
||||||
) -> tokio::task::JoinHandle<()> {
|
) -> tokio::task::JoinHandle<()> {
|
||||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
||||||
if let Some(tx) = response_tx {
|
if let Some(tx) = response_tx {
|
||||||
@@ -508,6 +512,179 @@ pub fn spawn_heartbeat(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Spawn a multi-user heartbeat runner that cycles through all users that
|
||||||
|
/// own routines (enabled or not). Each tick, it queries the DB for distinct
|
||||||
|
/// user_ids, creates a per-user workspace, and runs a heartbeat check for
|
||||||
|
/// each user concurrently. Per-user failure counts are tracked independently.
|
||||||
|
pub fn spawn_multi_user_heartbeat(
|
||||||
|
config: HeartbeatConfig,
|
||||||
|
hygiene_config: HygieneConfig,
|
||||||
|
llm: Arc<dyn LlmProvider>,
|
||||||
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||||
|
store: AdminScope,
|
||||||
|
) -> tokio::task::JoinHandle<()> {
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if !config.enabled {
|
||||||
|
tracing::info!("Multi-user heartbeat is disabled");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut tick_interval = if config.fire_at.is_none() {
|
||||||
|
let mut iv = tokio::time::interval(config.interval);
|
||||||
|
iv.tick().await; // skip immediate tick
|
||||||
|
Some(iv)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
// Track consecutive failures per user so we can disable heartbeat
|
||||||
|
// for persistently-failing users (same semantics as single-user mode).
|
||||||
|
let mut user_failures: std::collections::HashMap<String, u32> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
|
|
||||||
|
tracing::info!("Starting multi-user heartbeat loop");
|
||||||
|
|
||||||
|
loop {
|
||||||
|
if let Some(fire_at) = config.fire_at {
|
||||||
|
let sleep_dur = duration_until_next_fire(fire_at, config.resolved_tz());
|
||||||
|
tokio::time::sleep(sleep_dur).await;
|
||||||
|
} else if let Some(ref mut iv) = tick_interval {
|
||||||
|
iv.tick().await;
|
||||||
|
}
|
||||||
|
|
||||||
|
if config.is_quiet_hours() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get distinct user_ids from routines
|
||||||
|
let user_ids = match store.list_all_routines().await {
|
||||||
|
Ok(routines) => {
|
||||||
|
let mut ids: Vec<String> = routines
|
||||||
|
.iter()
|
||||||
|
.map(|r| r.user_id.clone())
|
||||||
|
.collect::<std::collections::HashSet<_>>()
|
||||||
|
.into_iter()
|
||||||
|
.collect();
|
||||||
|
ids.sort();
|
||||||
|
ids
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Multi-user heartbeat: failed to list routines: {}", e);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Run user heartbeats concurrently so one slow LLM call doesn't
|
||||||
|
// block others. Cap concurrency to avoid flooding the LLM provider.
|
||||||
|
const MAX_CONCURRENT_HEARTBEATS: usize = 8;
|
||||||
|
let mut join_set = tokio::task::JoinSet::new();
|
||||||
|
|
||||||
|
for user_id in &user_ids {
|
||||||
|
// Skip users that have exceeded max_failures
|
||||||
|
let failures = user_failures.get(user_id).copied().unwrap_or(0);
|
||||||
|
if failures >= config.max_failures {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let workspace = Arc::new(Workspace::new_with_db(user_id, Arc::clone(store.db())));
|
||||||
|
|
||||||
|
// Run memory hygiene per user (same as single-user heartbeat).
|
||||||
|
let hygiene_ws = Arc::clone(&workspace);
|
||||||
|
let hygiene_cfg = hygiene_config.clone();
|
||||||
|
let hygiene_user = user_id.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let report =
|
||||||
|
crate::workspace::hygiene::run_if_due(&hygiene_ws, &hygiene_cfg).await;
|
||||||
|
if report.had_work() {
|
||||||
|
tracing::info!(
|
||||||
|
user_id = hygiene_user,
|
||||||
|
daily_logs_deleted = report.daily_logs_deleted,
|
||||||
|
conversation_docs_deleted = report.conversation_docs_deleted,
|
||||||
|
"multi-user heartbeat: memory hygiene deleted stale documents"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Drain completed tasks to stay within the concurrency cap.
|
||||||
|
while join_set.len() >= MAX_CONCURRENT_HEARTBEATS {
|
||||||
|
if let Some(join_result) = join_set.join_next().await {
|
||||||
|
collect_heartbeat_result(join_result, &mut user_failures, &config);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let uid = user_id.clone();
|
||||||
|
let cfg = config.clone();
|
||||||
|
let hyg = hygiene_config.clone();
|
||||||
|
let llm_clone = llm.clone();
|
||||||
|
let tx = response_tx.clone();
|
||||||
|
let admin = store.clone();
|
||||||
|
|
||||||
|
join_set.spawn(async move {
|
||||||
|
let mut runner = HeartbeatRunner::new(cfg, hyg, workspace, llm_clone);
|
||||||
|
if let Some(tx) = tx {
|
||||||
|
runner = runner.with_response_channel(tx);
|
||||||
|
}
|
||||||
|
runner = runner.with_store(admin);
|
||||||
|
|
||||||
|
let result = runner.check_heartbeat().await;
|
||||||
|
if let HeartbeatResult::NeedsAttention(msg) = &result {
|
||||||
|
runner.send_notification(msg).await;
|
||||||
|
}
|
||||||
|
(uid, result)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Collect remaining results and update failure counts
|
||||||
|
while let Some(join_result) = join_set.join_next().await {
|
||||||
|
collect_heartbeat_result(join_result, &mut user_failures, &config);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process a single JoinSet result from the multi-user heartbeat loop.
|
||||||
|
fn collect_heartbeat_result(
|
||||||
|
join_result: Result<(String, HeartbeatResult), tokio::task::JoinError>,
|
||||||
|
user_failures: &mut std::collections::HashMap<String, u32>,
|
||||||
|
config: &HeartbeatConfig,
|
||||||
|
) {
|
||||||
|
let (uid, result) = match join_result {
|
||||||
|
Ok(pair) => pair,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Multi-user heartbeat task panicked: {}", e);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match result {
|
||||||
|
HeartbeatResult::Ok => {
|
||||||
|
tracing::trace!(user_id = uid, "Multi-user heartbeat OK");
|
||||||
|
user_failures.remove(&uid);
|
||||||
|
}
|
||||||
|
HeartbeatResult::NeedsAttention(_) => {
|
||||||
|
tracing::info!(user_id = uid, "Multi-user heartbeat needs attention");
|
||||||
|
user_failures.remove(&uid);
|
||||||
|
}
|
||||||
|
HeartbeatResult::Skipped => {}
|
||||||
|
HeartbeatResult::Failed(err) => {
|
||||||
|
let count = user_failures.entry(uid.clone()).or_insert(0);
|
||||||
|
*count += 1;
|
||||||
|
tracing::error!(
|
||||||
|
user_id = uid,
|
||||||
|
consecutive_failures = *count,
|
||||||
|
"Multi-user heartbeat failed: {}",
|
||||||
|
err
|
||||||
|
);
|
||||||
|
if *count >= config.max_failures {
|
||||||
|
tracing::error!(
|
||||||
|
user_id = uid,
|
||||||
|
"Multi-user heartbeat disabled for user after {} consecutive failures",
|
||||||
|
count
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -726,7 +903,7 @@ mod tests {
|
|||||||
Arc<crate::workspace::Workspace>,
|
Arc<crate::workspace::Workspace>,
|
||||||
Arc<dyn crate::llm::LlmProvider>,
|
Arc<dyn crate::llm::LlmProvider>,
|
||||||
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
|
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
|
||||||
Option<Arc<dyn crate::db::Database>>,
|
Option<AdminScope>,
|
||||||
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
|
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
|
||||||
let _ = _fn_ptr;
|
let _ = _fn_ptr;
|
||||||
}
|
}
|
||||||
|
|||||||
+24
-24
@@ -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, String, 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, String, 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>>,
|
||||||
@@ -74,7 +74,7 @@ pub fn spawn_job_monitor_with_context(
|
|||||||
}
|
}
|
||||||
|
|
||||||
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, String, 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,7 @@ 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, _user_id, SseEvent::JobResult { status, .. }))
|
Ok((ev_job_id, _user_id, AppEvent::JobResult { status, .. }))
|
||||||
if ev_job_id == job_id =>
|
if ev_job_id == job_id =>
|
||||||
{
|
{
|
||||||
let target = if status == "completed" {
|
let target = if status == "completed" {
|
||||||
@@ -229,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, String, 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();
|
||||||
@@ -240,7 +240,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobMessage {
|
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(),
|
||||||
@@ -262,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, String, 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();
|
||||||
@@ -274,7 +274,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
other_job_id,
|
other_job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobMessage {
|
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(),
|
||||||
@@ -293,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, String, 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();
|
||||||
@@ -304,7 +304,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobResult {
|
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,
|
||||||
@@ -329,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, String, 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();
|
||||||
@@ -340,7 +340,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobToolUse {
|
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"}),
|
||||||
@@ -353,7 +353,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobMessage {
|
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(),
|
||||||
@@ -409,7 +409,7 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, String, 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(
|
||||||
@@ -425,7 +425,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobResult {
|
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,
|
||||||
@@ -458,7 +458,7 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, String, 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(
|
||||||
@@ -474,7 +474,7 @@ mod tests {
|
|||||||
.send((
|
.send((
|
||||||
job_id,
|
job_id,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobResult {
|
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,
|
||||||
@@ -507,14 +507,14 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let (event_tx, _) = broadcast::channel::<(Uuid, String, 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,
|
||||||
"test-user".to_string(),
|
"test-user".to_string(),
|
||||||
SseEvent::JobResult {
|
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,
|
||||||
|
|||||||
+3
-1
@@ -36,7 +36,9 @@ pub(crate) use agent_loop::truncate_for_preview;
|
|||||||
pub use agent_loop::{Agent, AgentDeps};
|
pub use agent_loop::{Agent, AgentDeps};
|
||||||
pub use compaction::{CompactionResult, ContextCompactor};
|
pub use compaction::{CompactionResult, ContextCompactor};
|
||||||
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
|
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
|
||||||
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
|
pub use heartbeat::{
|
||||||
|
HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat, spawn_multi_user_heartbeat,
|
||||||
|
};
|
||||||
pub use router::{MessageIntent, Router};
|
pub use router::{MessageIntent, Router};
|
||||||
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
|
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
|
||||||
pub use routine_engine::{RoutineEngine, SandboxReadiness};
|
pub use routine_engine::{RoutineEngine, SandboxReadiness};
|
||||||
|
|||||||
+297
-39
@@ -18,21 +18,22 @@ 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::error::RoutineError;
|
use crate::error::RoutineError;
|
||||||
use crate::extensions::ExtensionManager;
|
use crate::extensions::ExtensionManager;
|
||||||
use crate::llm::{
|
use crate::llm::{
|
||||||
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
|
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
|
||||||
};
|
};
|
||||||
|
use crate::tenant::AdminScope;
|
||||||
use crate::tools::{
|
use crate::tools::{
|
||||||
ToolError, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_message,
|
ToolError, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_message,
|
||||||
prepare_tool_params,
|
prepare_tool_params,
|
||||||
@@ -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,10 +62,44 @@ 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,
|
||||||
store: Arc<dyn Database>,
|
store: AdminScope,
|
||||||
llm: Arc<dyn LlmProvider>,
|
llm: Arc<dyn LlmProvider>,
|
||||||
workspace: Arc<Workspace>,
|
workspace: Arc<Workspace>,
|
||||||
/// Sender for notifications (routed to channel manager).
|
/// Sender for notifications (routed to channel manager).
|
||||||
@@ -88,7 +128,7 @@ impl RoutineEngine {
|
|||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn new(
|
pub fn new(
|
||||||
config: RoutineConfig,
|
config: RoutineConfig,
|
||||||
store: Arc<dyn Database>,
|
store: AdminScope,
|
||||||
llm: Arc<dyn LlmProvider>,
|
llm: Arc<dyn LlmProvider>,
|
||||||
workspace: Arc<Workspace>,
|
workspace: Arc<Workspace>,
|
||||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||||
@@ -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.
|
||||||
@@ -698,12 +782,22 @@ impl RoutineEngine {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Per-user workspace (same pattern as spawn_fire).
|
||||||
|
let routine_workspace = if routine.user_id == self.workspace.user_id() {
|
||||||
|
self.workspace.clone()
|
||||||
|
} else {
|
||||||
|
Arc::new(Workspace::new_with_db(
|
||||||
|
&routine.user_id,
|
||||||
|
Arc::clone(self.store.db()),
|
||||||
|
))
|
||||||
|
};
|
||||||
|
|
||||||
// Execute inline for manual triggers (caller wants to wait)
|
// Execute inline for manual triggers (caller wants to wait)
|
||||||
let engine = EngineContext {
|
let engine = EngineContext {
|
||||||
config: self.config.clone(),
|
config: self.config.clone(),
|
||||||
store: self.store.clone(),
|
store: self.store.clone(),
|
||||||
llm: self.llm.clone(),
|
llm: self.llm.clone(),
|
||||||
workspace: self.workspace.clone(),
|
workspace: routine_workspace,
|
||||||
notify_tx: self.notify_tx.clone(),
|
notify_tx: self.notify_tx.clone(),
|
||||||
running_count: self.running_count.clone(),
|
running_count: self.running_count.clone(),
|
||||||
scheduler: self.scheduler.clone(),
|
scheduler: self.scheduler.clone(),
|
||||||
@@ -806,7 +900,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,
|
||||||
@@ -821,11 +920,23 @@ impl RoutineEngine {
|
|||||||
created_at: Utc::now(),
|
created_at: Utc::now(),
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Use per-user workspace so each routine executes in the correct
|
||||||
|
// user's context. Fall back to the engine-wide workspace when the
|
||||||
|
// routine belongs to the same user (avoids unnecessary allocation).
|
||||||
|
let routine_workspace = if routine.user_id == self.workspace.user_id() {
|
||||||
|
self.workspace.clone()
|
||||||
|
} else {
|
||||||
|
Arc::new(Workspace::new_with_db(
|
||||||
|
&routine.user_id,
|
||||||
|
Arc::clone(self.store.db()),
|
||||||
|
))
|
||||||
|
};
|
||||||
|
|
||||||
let engine = EngineContext {
|
let engine = EngineContext {
|
||||||
config: self.config.clone(),
|
config: self.config.clone(),
|
||||||
store: self.store.clone(),
|
store: self.store.clone(),
|
||||||
llm: self.llm.clone(),
|
llm: self.llm.clone(),
|
||||||
workspace: self.workspace.clone(),
|
workspace: routine_workspace,
|
||||||
notify_tx: self.notify_tx.clone(),
|
notify_tx: self.notify_tx.clone(),
|
||||||
running_count: self.running_count.clone(),
|
running_count: self.running_count.clone(),
|
||||||
scheduler: self.scheduler.clone(),
|
scheduler: self.scheduler.clone(),
|
||||||
@@ -843,7 +954,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 {
|
||||||
@@ -878,7 +989,7 @@ impl RoutineEngine {
|
|||||||
/// an active state (Pending/InProgress/Stuck). Maps the final `JobState` to
|
/// an active state (Pending/InProgress/Stuck). Maps the final `JobState` to
|
||||||
/// a `RunStatus` for the routine run.
|
/// a `RunStatus` for the routine run.
|
||||||
struct FullJobWatcher {
|
struct FullJobWatcher {
|
||||||
store: Arc<dyn Database>,
|
store: AdminScope,
|
||||||
job_id: Uuid,
|
job_id: Uuid,
|
||||||
routine_name: String,
|
routine_name: String,
|
||||||
}
|
}
|
||||||
@@ -889,7 +1000,7 @@ impl FullJobWatcher {
|
|||||||
/// Safety ceiling: 24 hours, derived from POLL_INTERVAL.
|
/// Safety ceiling: 24 hours, derived from POLL_INTERVAL.
|
||||||
const MAX_POLLS: u32 = (24 * 60 * 60) / Self::POLL_INTERVAL.as_secs() as u32;
|
const MAX_POLLS: u32 = (24 * 60 * 60) / Self::POLL_INTERVAL.as_secs() as u32;
|
||||||
|
|
||||||
fn new(store: Arc<dyn Database>, job_id: Uuid, routine_name: String) -> Self {
|
fn new(store: AdminScope, job_id: Uuid, routine_name: String) -> Self {
|
||||||
Self {
|
Self {
|
||||||
store,
|
store,
|
||||||
job_id,
|
job_id,
|
||||||
@@ -961,7 +1072,7 @@ impl FullJobWatcher {
|
|||||||
/// Shared context passed to the execution function.
|
/// Shared context passed to the execution function.
|
||||||
struct EngineContext {
|
struct EngineContext {
|
||||||
config: RoutineConfig,
|
config: RoutineConfig,
|
||||||
store: Arc<dyn Database>,
|
store: AdminScope,
|
||||||
llm: Arc<dyn LlmProvider>,
|
llm: Arc<dyn LlmProvider>,
|
||||||
workspace: Arc<Workspace>,
|
workspace: Arc<Workspace>,
|
||||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||||
@@ -1305,6 +1416,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 +1447,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 +1566,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 +1613,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 +1895,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 +1909,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 +1979,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 +2183,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![
|
||||||
|
|||||||
+12
-9
@@ -9,15 +9,14 @@ 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::error::{Error, JobError};
|
use crate::error::{Error, JobError};
|
||||||
use crate::extensions::ExtensionManager;
|
use crate::extensions::ExtensionManager;
|
||||||
use crate::hooks::HookRegistry;
|
use crate::hooks::HookRegistry;
|
||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
use crate::safety::SafetyLayer;
|
use crate::safety::SafetyLayer;
|
||||||
|
use crate::tenant::AdminScope;
|
||||||
use crate::tools::{
|
use crate::tools::{
|
||||||
ApprovalContext, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_error,
|
ApprovalContext, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_error,
|
||||||
prepare_tool_params,
|
prepare_tool_params,
|
||||||
@@ -53,7 +52,7 @@ struct ScheduledSubtask {
|
|||||||
pub struct SchedulerDeps {
|
pub struct SchedulerDeps {
|
||||||
pub tools: Arc<ToolRegistry>,
|
pub tools: Arc<ToolRegistry>,
|
||||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||||
pub store: Option<Arc<dyn Database>>,
|
pub store: Option<AdminScope>,
|
||||||
pub hooks: Arc<HookRegistry>,
|
pub hooks: Arc<HookRegistry>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -65,10 +64,10 @@ pub struct Scheduler {
|
|||||||
safety: Arc<SafetyLayer>,
|
safety: Arc<SafetyLayer>,
|
||||||
tools: Arc<ToolRegistry>,
|
tools: Arc<ToolRegistry>,
|
||||||
extension_manager: Option<Arc<ExtensionManager>>,
|
extension_manager: Option<Arc<ExtensionManager>>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<AdminScope>,
|
||||||
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.
|
||||||
@@ -781,10 +780,14 @@ mod tests {
|
|||||||
allow_local_tools: true,
|
allow_local_tools: true,
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
max_tool_iterations: 10,
|
max_tool_iterations: 10,
|
||||||
auto_approve_tools: true,
|
auto_approve_tools: true,
|
||||||
default_timezone: "UTC".to_string(),
|
default_timezone: "UTC".to_string(),
|
||||||
max_tokens_per_job,
|
max_tokens_per_job,
|
||||||
|
multi_tenant: false,
|
||||||
|
max_llm_concurrent_per_user: None,
|
||||||
|
max_jobs_concurrent_per_user: None,
|
||||||
};
|
};
|
||||||
let cm = Arc::new(ContextManager::new(5));
|
let cm = Arc::new(ContextManager::new(5));
|
||||||
let llm: Arc<dyn LlmProvider> = Arc::new(StubLlm);
|
let llm: Arc<dyn LlmProvider> = Arc::new(StubLlm);
|
||||||
|
|||||||
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::context::{ContextManager, JobState};
|
use crate::context::{ContextManager, JobState};
|
||||||
use crate::db::Database;
|
|
||||||
use crate::error::RepairError;
|
use crate::error::RepairError;
|
||||||
|
use crate::tenant::AdminScope;
|
||||||
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
||||||
|
|
||||||
/// A job that has been detected as stuck.
|
/// A job that has been detected as stuck.
|
||||||
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
|
|||||||
/// Jobs in `InProgress` longer than this are treated as stuck.
|
/// Jobs in `InProgress` longer than this are treated as stuck.
|
||||||
stuck_threshold: Duration,
|
stuck_threshold: Duration,
|
||||||
max_repair_attempts: u32,
|
max_repair_attempts: u32,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<AdminScope>,
|
||||||
builder: Option<Arc<dyn SoftwareBuilder>>,
|
builder: Option<Arc<dyn SoftwareBuilder>>,
|
||||||
tools: Option<Arc<ToolRegistry>>,
|
tools: Option<Arc<ToolRegistry>>,
|
||||||
}
|
}
|
||||||
@@ -91,8 +91,8 @@ impl DefaultSelfRepair {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Add a Store for tool failure tracking.
|
/// Add an admin-scoped store for tool failure tracking.
|
||||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
pub fn with_store(mut self, store: AdminScope) -> Self {
|
||||||
self.store = Some(store);
|
self.store = Some(store);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
@@ -806,7 +806,7 @@ mod tests {
|
|||||||
// Create self-repair with zero threshold (detect immediately),
|
// Create self-repair with zero threshold (detect immediately),
|
||||||
// wired with store, builder, and tools.
|
// wired with store, builder, and tools.
|
||||||
let repair = DefaultSelfRepair::new(Arc::clone(&cm), Duration::from_secs(0), 3)
|
let repair = DefaultSelfRepair::new(Arc::clone(&cm), Duration::from_secs(0), 3)
|
||||||
.with_store(Arc::clone(&db))
|
.with_store(crate::tenant::AdminScope::new(Arc::clone(&db)))
|
||||||
.with_builder(
|
.with_builder(
|
||||||
Arc::clone(&builder) as Arc<dyn crate::tools::SoftwareBuilder>,
|
Arc::clone(&builder) as Arc<dyn crate::tools::SoftwareBuilder>,
|
||||||
tools,
|
tools,
|
||||||
|
|||||||
+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 {
|
||||||
|
|||||||
+69
-15
@@ -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.";
|
||||||
|
|
||||||
@@ -175,6 +175,7 @@ impl Agent {
|
|||||||
pub(super) async fn process_user_input(
|
pub(super) async fn process_user_input(
|
||||||
&self,
|
&self,
|
||||||
message: &IncomingMessage,
|
message: &IncomingMessage,
|
||||||
|
tenant: crate::tenant::TenantCtx,
|
||||||
session: Arc<Mutex<Session>>,
|
session: Arc<Mutex<Session>>,
|
||||||
thread_id: Uuid,
|
thread_id: Uuid,
|
||||||
content: &str,
|
content: &str,
|
||||||
@@ -351,7 +352,7 @@ impl Agent {
|
|||||||
|
|
||||||
if let Some(intent) = self.router.route_command(&temp_message) {
|
if let Some(intent) = self.router.route_command(&temp_message) {
|
||||||
// Explicit command like /status, /job, /list - handle directly
|
// Explicit command like /status, /job, /list - handle directly
|
||||||
return self.handle_job_or_command(intent, message).await;
|
return self.handle_job_or_command(intent, message, &tenant).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Natural language goes through the agentic loop
|
// Natural language goes through the agentic loop
|
||||||
@@ -462,7 +463,7 @@ impl Agent {
|
|||||||
|
|
||||||
// Run the agentic tool execution loop
|
// Run the agentic tool execution loop
|
||||||
let result = self
|
let result = self
|
||||||
.run_agentic_loop(message, session.clone(), thread_id, turn_messages)
|
.run_agentic_loop(message, tenant, session.clone(), thread_id, turn_messages)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
// Re-acquire lock and check if interrupted
|
// Re-acquire lock and check if interrupted
|
||||||
@@ -513,10 +514,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 +535,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 +727,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 +737,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 +772,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 +1128,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 +1385,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),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1444,7 +1474,13 @@ impl Agent {
|
|||||||
|
|
||||||
// Continue the agentic loop (a tool was already executed this turn)
|
// Continue the agentic loop (a tool was already executed this turn)
|
||||||
let result = self
|
let result = self
|
||||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages)
|
.run_agentic_loop(
|
||||||
|
message,
|
||||||
|
self.tenant_ctx(&message.user_id).await,
|
||||||
|
session.clone(),
|
||||||
|
thread_id,
|
||||||
|
context_messages,
|
||||||
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
// Handle the result
|
// Handle the result
|
||||||
@@ -1459,10 +1495,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 +1507,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(
|
||||||
@@ -1816,7 +1853,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 +1889,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();
|
||||||
|
|
||||||
|
|||||||
+13
-19
@@ -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,
|
||||||
@@ -344,8 +338,9 @@ impl AppBuilder {
|
|||||||
|
|
||||||
// Detect multi-tenant mode: when GATEWAY_USER_TOKENS is configured,
|
// Detect multi-tenant mode: when GATEWAY_USER_TOKENS is configured,
|
||||||
// each authenticated user needs their own workspace scope. Use
|
// each authenticated user needs their own workspace scope. Use
|
||||||
// PerUserWorkspaceResolver to create per-user workspaces on demand
|
// WorkspacePool (which implements WorkspaceResolver) to create
|
||||||
// instead of sharing the startup workspace across all users.
|
// per-user workspaces on demand instead of sharing the startup
|
||||||
|
// workspace across all users.
|
||||||
let is_multi_tenant = self
|
let is_multi_tenant = self
|
||||||
.config
|
.config
|
||||||
.channels
|
.channels
|
||||||
@@ -354,16 +349,14 @@ impl AppBuilder {
|
|||||||
.is_some_and(|gw| gw.user_tokens.is_some());
|
.is_some_and(|gw| gw.user_tokens.is_some());
|
||||||
|
|
||||||
if is_multi_tenant {
|
if is_multi_tenant {
|
||||||
let resolver = Arc::new(
|
let pool = Arc::new(crate::channels::web::server::WorkspacePool::new(
|
||||||
crate::tools::builtin::memory::PerUserWorkspaceResolver::new(
|
Arc::clone(db),
|
||||||
Arc::clone(db),
|
embeddings.clone(),
|
||||||
embeddings.clone(),
|
emb_cache_config,
|
||||||
emb_cache_config,
|
self.config.search.clone(),
|
||||||
self.config.search.clone(),
|
self.config.workspace.clone(),
|
||||||
self.config.workspace.clone(),
|
));
|
||||||
),
|
tools.register_memory_tools_with_resolver(pool);
|
||||||
);
|
|
||||||
tools.register_memory_tools_with_resolver(resolver);
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Memory tools configured with per-user workspace resolver (multi-tenant mode)"
|
"Memory tools configured with per-user workspace resolver (multi-tenant mode)"
|
||||||
);
|
);
|
||||||
@@ -887,6 +880,7 @@ impl AppBuilder {
|
|||||||
crate::agent::cost_guard::CostGuardConfig {
|
crate::agent::cost_guard::CostGuardConfig {
|
||||||
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
|
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
|
||||||
max_actions_per_hour: self.config.agent.max_actions_per_hour,
|
max_actions_per_hour: self.config.agent.max_actions_per_hour,
|
||||||
|
max_cost_per_user_per_day_cents: self.config.agent.max_cost_per_user_per_day_cents,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
|
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -122,18 +122,32 @@ impl RelayClient {
|
|||||||
/// instance_url in chat-api. IronClaw only passes an optional CSRF nonce
|
/// instance_url in chat-api. IronClaw only passes an optional CSRF nonce
|
||||||
/// for validating the callback — no URLs.
|
/// for validating the callback — no URLs.
|
||||||
pub async fn initiate_oauth(&self, state_nonce: Option<&str>) -> Result<String, RelayError> {
|
pub async fn initiate_oauth(&self, state_nonce: Option<&str>) -> Result<String, RelayError> {
|
||||||
|
let url = format!("{}/oauth/slack/auth", self.base_url);
|
||||||
|
tracing::debug!(relay_url = %url, "RelayClient::initiate_oauth: sending request");
|
||||||
let mut query: Vec<(&str, &str)> = vec![];
|
let mut query: Vec<(&str, &str)> = vec![];
|
||||||
if let Some(nonce) = state_nonce {
|
if let Some(nonce) = state_nonce {
|
||||||
query.push(("state_nonce", nonce));
|
query.push(("state_nonce", nonce));
|
||||||
}
|
}
|
||||||
let resp = self
|
let resp = self
|
||||||
.http
|
.http
|
||||||
.get(format!("{}/oauth/slack/auth", self.base_url))
|
.get(&url)
|
||||||
.bearer_auth(self.api_key.expose_secret())
|
.bearer_auth(self.api_key.expose_secret())
|
||||||
.query(&query)
|
.query(&query)
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| RelayError::Network(e.to_string()))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
relay_url = %url,
|
||||||
|
error = %e,
|
||||||
|
"RelayClient::initiate_oauth: network request failed"
|
||||||
|
);
|
||||||
|
RelayError::Network(e.to_string())
|
||||||
|
})?;
|
||||||
|
tracing::debug!(
|
||||||
|
relay_url = %url,
|
||||||
|
status = %resp.status(),
|
||||||
|
"RelayClient::initiate_oauth: received response"
|
||||||
|
);
|
||||||
|
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
if status.is_redirection() {
|
if status.is_redirection() {
|
||||||
@@ -224,20 +238,39 @@ impl RelayClient {
|
|||||||
method: &str,
|
method: &str,
|
||||||
body: serde_json::Value,
|
body: serde_json::Value,
|
||||||
) -> Result<serde_json::Value, RelayError> {
|
) -> Result<serde_json::Value, RelayError> {
|
||||||
|
let url = format!("{}/proxy/{}/{}", self.base_url, provider, method);
|
||||||
|
tracing::debug!(
|
||||||
|
relay_url = %url,
|
||||||
|
provider = %provider,
|
||||||
|
method = %method,
|
||||||
|
"RelayClient::proxy_provider: sending request"
|
||||||
|
);
|
||||||
let query: Vec<(&str, &str)> = vec![("team_id", team_id)];
|
let query: Vec<(&str, &str)> = vec![("team_id", team_id)];
|
||||||
let resp = self
|
let resp = self
|
||||||
.http
|
.http
|
||||||
.post(format!("{}/proxy/{}/{}", self.base_url, provider, method))
|
.post(&url)
|
||||||
.bearer_auth(self.api_key.expose_secret())
|
.bearer_auth(self.api_key.expose_secret())
|
||||||
.query(&query)
|
.query(&query)
|
||||||
.json(&body)
|
.json(&body)
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| RelayError::Network(e.to_string()))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
relay_url = %url,
|
||||||
|
error = %e,
|
||||||
|
"RelayClient::proxy_provider: network request failed"
|
||||||
|
);
|
||||||
|
RelayError::Network(e.to_string())
|
||||||
|
})?;
|
||||||
|
|
||||||
if !resp.status().is_success() {
|
if !resp.status().is_success() {
|
||||||
let status = resp.status().as_u16();
|
let status = resp.status().as_u16();
|
||||||
let body = resp.text().await.unwrap_or_default();
|
let body = resp.text().await.unwrap_or_default();
|
||||||
|
tracing::warn!(
|
||||||
|
relay_url = %url,
|
||||||
|
status = status,
|
||||||
|
"RelayClient::proxy_provider: channel-relay returned error"
|
||||||
|
);
|
||||||
return Err(RelayError::Api {
|
return Err(RelayError::Api {
|
||||||
status,
|
status,
|
||||||
message: body,
|
message: body,
|
||||||
@@ -255,23 +288,45 @@ impl RelayClient {
|
|||||||
/// 32-byte secret. Called once at activation time; the result is cached in the
|
/// 32-byte secret. Called once at activation time; the result is cached in the
|
||||||
/// extension manager so subsequent calls to `relay_signing_secret()` use it.
|
/// extension manager so subsequent calls to `relay_signing_secret()` use it.
|
||||||
pub async fn get_signing_secret(&self, team_id: &str) -> Result<Vec<u8>, RelayError> {
|
pub async fn get_signing_secret(&self, team_id: &str) -> Result<Vec<u8>, RelayError> {
|
||||||
|
let url = format!("{}/relay/signing-secret", self.base_url);
|
||||||
|
tracing::debug!(
|
||||||
|
relay_url = %url,
|
||||||
|
"RelayClient::get_signing_secret: fetching signing secret"
|
||||||
|
);
|
||||||
let resp = self
|
let resp = self
|
||||||
.http
|
.http
|
||||||
.get(format!("{}/relay/signing-secret", self.base_url))
|
.get(&url)
|
||||||
.bearer_auth(self.api_key.expose_secret())
|
.bearer_auth(self.api_key.expose_secret())
|
||||||
.query(&[("team_id", team_id)])
|
.query(&[("team_id", team_id)])
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| RelayError::Network(e.to_string()))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
relay_url = %url,
|
||||||
|
error = %e,
|
||||||
|
"RelayClient::get_signing_secret: network request failed"
|
||||||
|
);
|
||||||
|
RelayError::Network(e.to_string())
|
||||||
|
})?;
|
||||||
|
|
||||||
if !resp.status().is_success() {
|
if !resp.status().is_success() {
|
||||||
let status = resp.status().as_u16();
|
let status = resp.status().as_u16();
|
||||||
let body = resp.text().await.unwrap_or_default();
|
let body = resp.text().await.unwrap_or_default();
|
||||||
|
tracing::warn!(
|
||||||
|
relay_url = %url,
|
||||||
|
status = status,
|
||||||
|
body = %body,
|
||||||
|
"RelayClient::get_signing_secret: channel-relay returned error"
|
||||||
|
);
|
||||||
return Err(RelayError::Api {
|
return Err(RelayError::Api {
|
||||||
status,
|
status,
|
||||||
message: body,
|
message: body,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
tracing::debug!(
|
||||||
|
relay_url = %url,
|
||||||
|
"RelayClient::get_signing_secret: received successful response"
|
||||||
|
);
|
||||||
|
|
||||||
let body: serde_json::Value = resp
|
let body: serde_json::Value = resp
|
||||||
.json()
|
.json()
|
||||||
|
|||||||
+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,
|
||||||
|
}
|
||||||
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -175,7 +175,7 @@ pub async fn chat_auth_token_handler(
|
|||||||
if result.verification.is_some() {
|
if result.verification.is_some() {
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
&user.user_id,
|
&user.user_id,
|
||||||
SseEvent::AuthRequired {
|
AppEvent::AuthRequired {
|
||||||
extension_name: req.extension_name.clone(),
|
extension_name: req.extension_name.clone(),
|
||||||
instructions: Some(result.message),
|
instructions: Some(result.message),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
@@ -187,7 +187,7 @@ pub async fn chat_auth_token_handler(
|
|||||||
|
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
&user.user_id,
|
&user.user_id,
|
||||||
SseEvent::AuthCompleted {
|
AppEvent::AuthCompleted {
|
||||||
extension_name: req.extension_name.clone(),
|
extension_name: req.extension_name.clone(),
|
||||||
success: true,
|
success: true,
|
||||||
message: result.message,
|
message: result.message,
|
||||||
@@ -202,7 +202,7 @@ pub async fn chat_auth_token_handler(
|
|||||||
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
&user.user_id,
|
&user.user_id,
|
||||||
SseEvent::AuthRequired {
|
AppEvent::AuthRequired {
|
||||||
extension_name: req.extension_name.clone(),
|
extension_name: req.extension_name.clone(),
|
||||||
instructions: Some(msg.clone()),
|
instructions: Some(msg.clone()),
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
@@ -398,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();
|
||||||
|
|
||||||
|
|||||||
@@ -598,11 +598,23 @@ pub async fn jobs_prompt_handler(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Try agent job path: verify ownership, then send via scheduler.
|
// Try agent job path: verify ownership, then send via scheduler.
|
||||||
if let Some(ref store) = state.store
|
if let Some(ref store) = state.store {
|
||||||
&& let Ok(Some(agent_job)) = store.get_job(job_id).await
|
match store.get_job(job_id).await {
|
||||||
&& agent_job.user_id != user.user_id
|
Ok(Some(agent_job)) => {
|
||||||
{
|
if agent_job.user_id != user.user_id {
|
||||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
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((
|
||||||
|
|||||||
@@ -0,0 +1,225 @@
|
|||||||
|
//! Memory/workspace API handlers.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use axum::{
|
||||||
|
Json,
|
||||||
|
extract::{Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use crate::channels::web::auth::{AuthenticatedUser, UserIdentity};
|
||||||
|
use crate::channels::web::server::GatewayState;
|
||||||
|
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)]
|
||||||
|
pub struct TreeQuery {
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub depth: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn memory_tree_handler(
|
||||||
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
|
Query(_query): Query<TreeQuery>,
|
||||||
|
) -> Result<Json<MemoryTreeResponse>, (StatusCode, String)> {
|
||||||
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
|
|
||||||
|
// Build tree from list_all (flat list of all paths)
|
||||||
|
let all_paths = workspace
|
||||||
|
.list_all()
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
// Collect unique directories and files
|
||||||
|
let mut entries: Vec<TreeEntry> = Vec::new();
|
||||||
|
let mut seen_dirs: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||||
|
|
||||||
|
for path in &all_paths {
|
||||||
|
// Add parent directories
|
||||||
|
let parts: Vec<&str> = path.split('/').collect();
|
||||||
|
for i in 0..parts.len().saturating_sub(1) {
|
||||||
|
let dir_path = parts[..=i].join("/");
|
||||||
|
if seen_dirs.insert(dir_path.clone()) {
|
||||||
|
entries.push(TreeEntry {
|
||||||
|
path: dir_path,
|
||||||
|
is_dir: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Add the file itself
|
||||||
|
entries.push(TreeEntry {
|
||||||
|
path: path.clone(),
|
||||||
|
is_dir: false,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
entries.sort_by(|a, b| a.path.cmp(&b.path));
|
||||||
|
|
||||||
|
Ok(Json(MemoryTreeResponse { entries }))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
pub struct ListQuery {
|
||||||
|
pub path: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn memory_list_handler(
|
||||||
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
|
Query(query): Query<ListQuery>,
|
||||||
|
) -> Result<Json<MemoryListResponse>, (StatusCode, String)> {
|
||||||
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
|
|
||||||
|
let path = query.path.as_deref().unwrap_or("");
|
||||||
|
let entries = workspace
|
||||||
|
.list(path)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
let list_entries: Vec<ListEntry> = entries
|
||||||
|
.iter()
|
||||||
|
.map(|e| ListEntry {
|
||||||
|
name: e.path.rsplit('/').next().unwrap_or(&e.path).to_string(),
|
||||||
|
path: e.path.clone(),
|
||||||
|
is_dir: e.is_directory,
|
||||||
|
updated_at: e.updated_at.map(|dt| dt.to_rfc3339()),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(MemoryListResponse {
|
||||||
|
path: path.to_string(),
|
||||||
|
entries: list_entries,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
pub struct ReadQuery {
|
||||||
|
pub path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn memory_read_handler(
|
||||||
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
|
Query(query): Query<ReadQuery>,
|
||||||
|
) -> Result<Json<MemoryReadResponse>, (StatusCode, String)> {
|
||||||
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
|
|
||||||
|
let doc = workspace
|
||||||
|
.read(&query.path)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::NOT_FOUND, e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(Json(MemoryReadResponse {
|
||||||
|
path: query.path,
|
||||||
|
content: doc.content,
|
||||||
|
updated_at: Some(doc.updated_at.to_rfc3339()),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn memory_write_handler(
|
||||||
|
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(
|
||||||
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
AuthenticatedUser(user): AuthenticatedUser,
|
||||||
|
Json(req): Json<MemorySearchRequest>,
|
||||||
|
) -> Result<Json<MemorySearchResponse>, (StatusCode, String)> {
|
||||||
|
let workspace = resolve_workspace(&state, &user).await?;
|
||||||
|
|
||||||
|
let limit = req.limit.unwrap_or(10);
|
||||||
|
let results = workspace
|
||||||
|
.search(&req.query, limit)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
|
let hits: Vec<SearchHit> = results
|
||||||
|
.iter()
|
||||||
|
.map(|r| SearchHit {
|
||||||
|
path: r.document_id.to_string(),
|
||||||
|
content: r.content.clone(),
|
||||||
|
score: r.score as f64,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(MemorySearchResponse { results: hits }))
|
||||||
|
}
|
||||||
@@ -3,6 +3,7 @@
|
|||||||
//! Each module groups related endpoint handlers by domain.
|
//! Each module groups related endpoint handlers by domain.
|
||||||
|
|
||||||
pub mod jobs;
|
pub mod jobs;
|
||||||
|
pub mod memory;
|
||||||
pub mod routines;
|
pub mod routines;
|
||||||
pub mod skills;
|
pub mod skills;
|
||||||
|
|
||||||
|
|||||||
@@ -114,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,
|
||||||
@@ -324,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,
|
||||||
|
|||||||
@@ -54,10 +54,37 @@ fn validate_webhook_secret(
|
|||||||
///
|
///
|
||||||
/// This endpoint is **public** (no gateway auth token required) but protected
|
/// This endpoint is **public** (no gateway auth token required) but protected
|
||||||
/// by the per-routine webhook secret sent via the `X-Webhook-Secret` header.
|
/// by the per-routine webhook secret sent via the `X-Webhook-Secret` header.
|
||||||
|
///
|
||||||
|
/// **Single-user/backward-compatible**: looks up routines by path across all
|
||||||
|
/// users. For multi-tenant isolation, use the user-scoped endpoint at
|
||||||
|
/// `/api/webhooks/u/{user_id}/{path}` instead.
|
||||||
pub async fn webhook_trigger_handler(
|
pub async fn webhook_trigger_handler(
|
||||||
State(state): State<Arc<GatewayState>>,
|
State(state): State<Arc<GatewayState>>,
|
||||||
Path(path): Path<String>,
|
Path(path): Path<String>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
|
fire_webhook_inner(state, &path, None, &headers).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle incoming webhook POST to `/api/webhooks/u/{user_id}/{path}`.
|
||||||
|
///
|
||||||
|
/// User-scoped variant for multi-tenant deployments. The `user_id` in the URL
|
||||||
|
/// restricts the routine lookup to that user only, preventing cross-user
|
||||||
|
/// webhook triggering even when paths collide.
|
||||||
|
pub async fn webhook_trigger_user_scoped_handler(
|
||||||
|
State(state): State<Arc<GatewayState>>,
|
||||||
|
Path((user_id, path)): Path<(String, String)>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
|
fire_webhook_inner(state, &path, Some(&user_id), &headers).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Shared webhook logic for both scoped and unscoped endpoints.
|
||||||
|
async fn fire_webhook_inner(
|
||||||
|
state: Arc<GatewayState>,
|
||||||
|
path: &str,
|
||||||
|
user_id: Option<&str>,
|
||||||
|
headers: &HeaderMap,
|
||||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
// Rate limit check
|
// Rate limit check
|
||||||
if !state.webhook_rate_limiter.check() {
|
if !state.webhook_rate_limiter.check() {
|
||||||
@@ -72,9 +99,9 @@ pub async fn webhook_trigger_handler(
|
|||||||
"Database not available".to_string(),
|
"Database not available".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
// Targeted query instead of loading all routines
|
// Targeted query — when user_id is provided, restrict to that user's routines
|
||||||
let routine = store
|
let routine = store
|
||||||
.get_webhook_routine_by_path(&path)
|
.get_webhook_routine_by_path(path, user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||||
.ok_or((
|
.ok_or((
|
||||||
@@ -99,7 +126,7 @@ pub async fn webhook_trigger_handler(
|
|||||||
))?
|
))?
|
||||||
};
|
};
|
||||||
|
|
||||||
let run_id = engine.fire_webhook(routine.id, &path).await.map_err(|e| {
|
let run_id = engine.fire_webhook(routine.id, path).await.map_err(|e| {
|
||||||
let status = match &e {
|
let status = match &e {
|
||||||
crate::error::RoutineError::NotFound { .. } => StatusCode::NOT_FOUND,
|
crate::error::RoutineError::NotFound { .. } => StatusCode::NOT_FOUND,
|
||||||
crate::error::RoutineError::Disabled { .. }
|
crate::error::RoutineError::Disabled { .. }
|
||||||
|
|||||||
+52
-19
@@ -58,7 +58,7 @@ use self::log_layer::{LogBroadcaster, LogLevelHandle};
|
|||||||
use self::auth::MultiAuthState;
|
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 {
|
||||||
@@ -98,7 +98,8 @@ impl GatewayChannel {
|
|||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
default_user_id: config.user_id.clone(),
|
owner_id: config.user_id.clone(),
|
||||||
|
default_sender_id: config.user_id.clone(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
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,
|
||||||
@@ -121,6 +122,22 @@ impl GatewayChannel {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Rebind the single-user auth identity to the durable owner scope while
|
||||||
|
/// preserving the configured gateway sender/routing identity.
|
||||||
|
pub fn with_owner_scope(mut self, owner_id: impl Into<String>) -> Self {
|
||||||
|
let owner_id = owner_id.into();
|
||||||
|
let single_user_token = if self.config.user_tokens.is_none() {
|
||||||
|
self.auth.first_token().map(ToOwned::to_owned)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
if let Some(token) = single_user_token {
|
||||||
|
self.auth = MultiAuthState::single(token, owner_id.clone());
|
||||||
|
}
|
||||||
|
self.rebuild_state(|s| s.owner_id = owner_id);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
/// Create a gateway channel with a pre-built multi-user auth state.
|
/// Create a gateway channel with a pre-built multi-user auth state.
|
||||||
pub fn new_multi_auth(config: GatewayConfig, auth: MultiAuthState) -> Self {
|
pub fn new_multi_auth(config: GatewayConfig, auth: MultiAuthState) -> Self {
|
||||||
let state = Arc::new(GatewayState {
|
let state = Arc::new(GatewayState {
|
||||||
@@ -137,7 +154,8 @@ impl GatewayChannel {
|
|||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
default_user_id: config.user_id.clone(),
|
owner_id: config.user_id.clone(),
|
||||||
|
default_sender_id: config.user_id.clone(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
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,
|
||||||
@@ -177,7 +195,8 @@ 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(),
|
||||||
default_user_id: self.state.default_user_id.clone(),
|
owner_id: self.state.owner_id.clone(),
|
||||||
|
default_sender_id: self.state.default_sender_id.clone(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
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(),
|
||||||
@@ -367,7 +386,7 @@ impl Channel for GatewayChannel {
|
|||||||
|
|
||||||
self.state.sse.broadcast_for_user(
|
self.state.sse.broadcast_for_user(
|
||||||
&msg.user_id,
|
&msg.user_id,
|
||||||
SseEvent::Response {
|
AppEvent::Response {
|
||||||
content: response.content,
|
content: response.content,
|
||||||
thread_id,
|
thread_id,
|
||||||
},
|
},
|
||||||
@@ -386,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(),
|
||||||
},
|
},
|
||||||
@@ -399,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(),
|
||||||
},
|
},
|
||||||
@@ -423,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,
|
||||||
@@ -434,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,
|
||||||
@@ -448,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,
|
||||||
@@ -458,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,
|
||||||
@@ -512,7 +545,7 @@ impl Channel for GatewayChannel {
|
|||||||
};
|
};
|
||||||
self.state.sse.broadcast_for_user(
|
self.state.sse.broadcast_for_user(
|
||||||
user_id,
|
user_id,
|
||||||
SseEvent::Response {
|
AppEvent::Response {
|
||||||
content: response.content,
|
content: response.content,
|
||||||
thread_id,
|
thread_id,
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -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(
|
||||||
@@ -954,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);
|
||||||
|
|||||||
+571
-258
File diff suppressed because it is too large
Load Diff
+20
-43
@@ -11,7 +11,7 @@ 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.
|
||||||
@@ -25,7 +25,7 @@ const MAX_CONNECTIONS: u64 = 100;
|
|||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub(crate) struct ScopedEvent {
|
pub(crate) struct ScopedEvent {
|
||||||
pub(crate) user_id: Option<String>,
|
pub(crate) user_id: Option<String>,
|
||||||
pub(crate) event: SseEvent,
|
pub(crate) event: AppEvent,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Manages SSE broadcast to all connected browser tabs.
|
/// Manages SSE broadcast to all connected browser tabs.
|
||||||
@@ -75,7 +75,7 @@ impl SseManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Broadcast an event to all connected clients (global/unscoped).
|
/// Broadcast an event to all connected clients (global/unscoped).
|
||||||
pub fn broadcast(&self, event: SseEvent) {
|
pub fn broadcast(&self, event: AppEvent) {
|
||||||
let _ = self.tx.send(ScopedEvent {
|
let _ = self.tx.send(ScopedEvent {
|
||||||
user_id: None,
|
user_id: None,
|
||||||
event,
|
event,
|
||||||
@@ -86,7 +86,7 @@ impl SseManager {
|
|||||||
///
|
///
|
||||||
/// Only subscribers for this user_id (or unscoped subscribers) will
|
/// Only subscribers for this user_id (or unscoped subscribers) will
|
||||||
/// receive the event.
|
/// receive the event.
|
||||||
pub fn broadcast_for_user(&self, user_id: &str, event: SseEvent) {
|
pub fn broadcast_for_user(&self, user_id: &str, event: AppEvent) {
|
||||||
let _ = self.tx.send(ScopedEvent {
|
let _ = self.tx.send(ScopedEvent {
|
||||||
user_id: Some(user_id.to_string()),
|
user_id: Some(user_id.to_string()),
|
||||||
event,
|
event,
|
||||||
@@ -108,7 +108,7 @@ impl SseManager {
|
|||||||
pub fn subscribe_raw(
|
pub fn subscribe_raw(
|
||||||
&self,
|
&self,
|
||||||
user_id: Option<String>,
|
user_id: Option<String>,
|
||||||
) -> Option<impl Stream<Item = SseEvent> + Send + 'static + use<>> {
|
) -> 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);
|
||||||
@@ -186,30 +186,7 @@ impl SseManager {
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
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::ApprovalNeeded { .. } => "approval_needed",
|
|
||||||
SseEvent::AuthRequired { .. } => "auth_required",
|
|
||||||
SseEvent::AuthCompleted { .. } => "auth_completed",
|
|
||||||
SseEvent::Error { .. } => "error",
|
|
||||||
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",
|
|
||||||
};
|
|
||||||
Some(Ok(Event::default().event(event_type).data(data)))
|
Some(Ok(Event::default().event(event_type).data(data)))
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -272,7 +249,7 @@ 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]
|
||||||
@@ -280,14 +257,14 @@ mod tests {
|
|||||||
let manager = SseManager::new();
|
let manager = SseManager::new();
|
||||||
let mut stream = Box::pin(manager.subscribe_raw(None).expect("should 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 = stream.next().await.unwrap();
|
let event = stream.next().await.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"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -299,14 +276,14 @@ mod tests {
|
|||||||
|
|
||||||
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"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -329,12 +306,12 @@ mod tests {
|
|||||||
let mut s2 = Box::pin(manager.subscribe_raw(None).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);
|
||||||
@@ -373,25 +350,25 @@ mod tests {
|
|||||||
// Send event scoped to alice
|
// Send event scoped to alice
|
||||||
manager.broadcast_for_user(
|
manager.broadcast_for_user(
|
||||||
"alice",
|
"alice",
|
||||||
SseEvent::Status {
|
AppEvent::Status {
|
||||||
message: "alice only".to_string(),
|
message: "alice only".to_string(),
|
||||||
thread_id: None,
|
thread_id: None,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
|
||||||
// Send global event
|
// Send global event
|
||||||
manager.broadcast(SseEvent::Heartbeat);
|
manager.broadcast(AppEvent::Heartbeat);
|
||||||
|
|
||||||
// Alice gets her scoped event
|
// Alice gets her scoped event
|
||||||
let e = alice.next().await.unwrap();
|
let e = alice.next().await.unwrap();
|
||||||
assert!(matches!(e, SseEvent::Status { .. }));
|
assert!(matches!(e, AppEvent::Status { .. }));
|
||||||
|
|
||||||
// Alice also gets the global heartbeat
|
// Alice also gets the global heartbeat
|
||||||
let e = alice.next().await.unwrap();
|
let e = alice.next().await.unwrap();
|
||||||
assert!(matches!(e, SseEvent::Heartbeat));
|
assert!(matches!(e, AppEvent::Heartbeat));
|
||||||
|
|
||||||
// Bob only gets the global heartbeat (alice's event was filtered)
|
// Bob only gets the global heartbeat (alice's event was filtered)
|
||||||
let e = bob.next().await.unwrap();
|
let e = bob.next().await.unwrap(); // safety: test-only
|
||||||
assert!(matches!(e, SseEvent::Heartbeat));
|
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>'
|
||||||
|
|||||||
@@ -76,7 +76,8 @@ impl TestGatewayBuilder {
|
|||||||
store: None,
|
store: None,
|
||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
default_user_id: self.user_id,
|
owner_id: self.user_id.clone(),
|
||||||
|
default_sender_id: self.user_id,
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
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,
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ use axum::routing::{delete, get, post};
|
|||||||
use tower::ServiceExt;
|
use tower::ServiceExt;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::channels::web::GatewayChannel;
|
||||||
use crate::channels::web::auth::{
|
use crate::channels::web::auth::{
|
||||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||||
};
|
};
|
||||||
@@ -23,6 +24,7 @@ use crate::channels::web::server::{
|
|||||||
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
||||||
};
|
};
|
||||||
use crate::channels::web::sse::SseManager;
|
use crate::channels::web::sse::SseManager;
|
||||||
|
use crate::config::GatewayConfig;
|
||||||
|
|
||||||
// ── Helpers ────────────────────────────────────────────────────────────
|
// ── Helpers ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -64,7 +66,8 @@ fn build_state(
|
|||||||
store,
|
store,
|
||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue,
|
prompt_queue,
|
||||||
default_user_id: "test".to_string(),
|
owner_id: "test".to_string(),
|
||||||
|
default_sender_id: "test".to_string(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||||
ws_tracker: None,
|
ws_tracker: None,
|
||||||
llm_provider: None,
|
llm_provider: None,
|
||||||
@@ -82,6 +85,40 @@ fn build_state(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn gateway_config() -> GatewayConfig {
|
||||||
|
GatewayConfig {
|
||||||
|
host: "127.0.0.1".to_string(),
|
||||||
|
port: 3000,
|
||||||
|
auth_token: Some("gateway-auth".to_string()),
|
||||||
|
user_id: "gateway-sender".to_string(),
|
||||||
|
workspace_read_scopes: Vec::new(),
|
||||||
|
memory_layers: Vec::new(),
|
||||||
|
user_tokens: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn with_owner_scope_updates_gateway_owner_scope_in_multi_user_mode() {
|
||||||
|
let mut gateway = GatewayChannel::new(gateway_config());
|
||||||
|
gateway.auth = two_user_auth();
|
||||||
|
gateway.config.user_tokens = Some(HashMap::new());
|
||||||
|
let gateway = gateway.with_owner_scope("owner-scope");
|
||||||
|
|
||||||
|
assert_eq!(gateway.state.owner_id, "owner-scope");
|
||||||
|
assert_eq!(gateway.state.default_sender_id, "gateway-sender");
|
||||||
|
|
||||||
|
let alice = gateway
|
||||||
|
.auth
|
||||||
|
.authenticate("tok-alice")
|
||||||
|
.expect("alice token should remain valid");
|
||||||
|
let bob = gateway
|
||||||
|
.auth
|
||||||
|
.authenticate("tok-bob")
|
||||||
|
.expect("bob token should remain valid");
|
||||||
|
assert_eq!(alice.user_id, "alice");
|
||||||
|
assert_eq!(bob.user_id, "bob");
|
||||||
|
}
|
||||||
|
|
||||||
/// Create a libSQL-backed test database in a temporary directory.
|
/// Create a libSQL-backed test database in a temporary directory.
|
||||||
///
|
///
|
||||||
/// Returns the database and a `TempDir` guard — the database file is
|
/// Returns the database and a `TempDir` guard — the database file is
|
||||||
|
|||||||
+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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -97,7 +97,7 @@ pub async fn handle_ws_connection(
|
|||||||
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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -275,7 +275,7 @@ async fn handle_client_message(
|
|||||||
if result.verification.is_some() {
|
if result.verification.is_some() {
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
user_id,
|
user_id,
|
||||||
crate::channels::web::types::SseEvent::AuthRequired {
|
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,
|
||||||
@@ -286,7 +286,7 @@ async fn handle_client_message(
|
|||||||
crate::channels::web::server::clear_auth_mode(state, user_id).await;
|
crate::channels::web::server::clear_auth_mode(state, user_id).await;
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
user_id,
|
user_id,
|
||||||
crate::channels::web::types::SseEvent::AuthCompleted {
|
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||||
extension_name,
|
extension_name,
|
||||||
success: true,
|
success: true,
|
||||||
message: result.message,
|
message: result.message,
|
||||||
@@ -299,7 +299,7 @@ async fn handle_client_message(
|
|||||||
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
|
||||||
state.sse.broadcast_for_user(
|
state.sse.broadcast_for_user(
|
||||||
user_id,
|
user_id,
|
||||||
crate::channels::web::types::SseEvent::AuthRequired {
|
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,
|
||||||
@@ -520,7 +520,8 @@ mod tests {
|
|||||||
job_manager: None,
|
job_manager: None,
|
||||||
prompt_queue: None,
|
prompt_queue: None,
|
||||||
scheduler: None,
|
scheduler: None,
|
||||||
default_user_id: "test".to_string(),
|
owner_id: "test".to_string(),
|
||||||
|
default_sender_id: "test".to_string(),
|
||||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
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,
|
||||||
|
|||||||
+582
-27
@@ -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`
|
||||||
@@ -449,7 +473,8 @@ pub struct PendingOAuthFlow {
|
|||||||
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||||
/// SSE broadcast manager for notifying the web UI.
|
/// SSE broadcast manager for notifying the web UI.
|
||||||
pub sse_manager: Option<Arc<crate::channels::web::sse::SseManager>>,
|
pub sse_manager: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||||
|
/// Kept as `gateway_token` for public API compatibility.
|
||||||
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.
|
||||||
/// Used for provider-specific requirements such as RFC 8707 `resource`.
|
/// Used for provider-specific requirements such as RFC 8707 `resource`.
|
||||||
@@ -472,6 +497,12 @@ impl std::fmt::Debug for PendingOAuthFlow {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl PendingOAuthFlow {
|
||||||
|
pub fn oauth_proxy_auth_token(&self) -> Option<&str> {
|
||||||
|
self.gateway_token.as_deref()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter.
|
/// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter.
|
||||||
pub type PendingOAuthRegistry = Arc<RwLock<HashMap<String, PendingOAuthFlow>>>;
|
pub type PendingOAuthRegistry = Arc<RwLock<HashMap<String, PendingOAuthFlow>>>;
|
||||||
|
|
||||||
@@ -505,6 +536,22 @@ pub fn exchange_proxy_url() -> Option<String> {
|
|||||||
.filter(|url| !url.is_empty())
|
.filter(|url| !url.is_empty())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the configured OAuth proxy auth token, if any.
|
||||||
|
///
|
||||||
|
/// New hosted infra can inject a dedicated shared proxy secret via
|
||||||
|
/// `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`. Existing hosted instances continue to
|
||||||
|
/// work by falling back to `GATEWAY_AUTH_TOKEN`.
|
||||||
|
pub fn oauth_proxy_auth_token() -> Option<String> {
|
||||||
|
fn normalized_env_value(key: &str) -> Option<String> {
|
||||||
|
crate::config::helpers::env_or_override(key)
|
||||||
|
.map(|value| value.trim().to_string())
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
normalized_env_value("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN")
|
||||||
|
.or_else(|| normalized_env_value("GATEWAY_AUTH_TOKEN"))
|
||||||
|
}
|
||||||
|
|
||||||
/// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout).
|
/// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout).
|
||||||
pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300);
|
pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300);
|
||||||
|
|
||||||
@@ -650,6 +697,8 @@ pub fn strip_instance_prefix(state: &str) -> &str {
|
|||||||
|
|
||||||
pub struct ProxyTokenExchangeRequest<'a> {
|
pub struct ProxyTokenExchangeRequest<'a> {
|
||||||
pub proxy_url: &'a str,
|
pub proxy_url: &'a str,
|
||||||
|
/// OAuth proxy auth token.
|
||||||
|
/// Kept as `gateway_token` for public API compatibility.
|
||||||
pub gateway_token: &'a str,
|
pub gateway_token: &'a str,
|
||||||
pub token_url: &'a str,
|
pub token_url: &'a str,
|
||||||
pub client_id: &'a str,
|
pub client_id: &'a str,
|
||||||
@@ -661,9 +710,53 @@ 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,
|
||||||
|
/// OAuth proxy auth token.
|
||||||
|
/// Kept as `gateway_token` for public API compatibility.
|
||||||
|
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 an OAuth proxy auth token (Bearer header). The caller may
|
||||||
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||||
/// the provider requires it.
|
/// the provider requires it.
|
||||||
///
|
///
|
||||||
@@ -675,13 +768,14 @@ pub async fn exchange_via_proxy(
|
|||||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||||
if request.gateway_token.is_empty() {
|
if request.gateway_token.is_empty() {
|
||||||
return Err(OAuthCallbackError::Io(
|
return Err(OAuthCallbackError::Io(
|
||||||
"Gateway auth token is required for proxy token exchange".to_string(),
|
"OAuth proxy auth token is required for proxy token exchange".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let exchange_url = format!("{}/oauth/exchange", request.proxy_url.trim_end_matches('/'));
|
let exchange_url = format!("{}/oauth/exchange", request.proxy_url.trim_end_matches('/'));
|
||||||
|
|
||||||
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 +818,454 @@ 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 an OAuth proxy 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
|
"OAuth proxy 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();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct EnvVarGuard {
|
||||||
|
key: &'static str,
|
||||||
|
original: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for EnvVarGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||||
|
unsafe {
|
||||||
|
if let Some(ref value) = self.original {
|
||||||
|
std::env::set_var(self.key, value);
|
||||||
|
} else {
|
||||||
|
std::env::remove_var(self.key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard {
|
||||||
|
let original = std::env::var(key).ok();
|
||||||
|
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||||
|
unsafe {
|
||||||
|
if let Some(value) = value {
|
||||||
|
std::env::set_var(key, value);
|
||||||
|
} else {
|
||||||
|
std::env::remove_var(key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
EnvVarGuard { key, original }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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_exchange_via_proxy_sends_auth_and_form() {
|
||||||
|
let server = MockProxyServer::start().await;
|
||||||
|
let mut extra_token_params = HashMap::new();
|
||||||
|
extra_token_params.insert("resource".to_string(), "https://mcp.notion.com".to_string());
|
||||||
|
|
||||||
|
let response = super::exchange_via_proxy(super::ProxyTokenExchangeRequest {
|
||||||
|
proxy_url: &server.base_url(),
|
||||||
|
gateway_token: "shared-oauth-proxy-secret",
|
||||||
|
code: "auth-code-123",
|
||||||
|
redirect_uri: "https://oauth.example.com/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: &extra_token_params,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("proxy exchange 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 shared-oauth-proxy-secret")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[0].form.get("code").map(String::as_str),
|
||||||
|
Some("auth-code-123")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[0].form.get("redirect_uri").map(String::as_str),
|
||||||
|
Some("https://oauth.example.com/oauth/callback")
|
||||||
|
);
|
||||||
|
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("access_token_field")
|
||||||
|
.map(String::as_str),
|
||||||
|
Some("access_token")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[0].form.get("code_verifier").map(String::as_str),
|
||||||
|
Some("code-verifier-123")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[0].form.get("resource").map(String::as_str),
|
||||||
|
Some("https://mcp.notion.com")
|
||||||
|
);
|
||||||
|
|
||||||
|
server.shutdown().await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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() {
|
||||||
@@ -1159,6 +1666,54 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_oauth_proxy_auth_token_prefers_dedicated_env() {
|
||||||
|
let _guard = lock_env();
|
||||||
|
let _proxy_guard = set_env_var(
|
||||||
|
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||||
|
Some("shared-proxy-secret"),
|
||||||
|
);
|
||||||
|
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||||
|
Some("shared-proxy-secret")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_oauth_proxy_auth_token_falls_back_to_gateway_token() {
|
||||||
|
let _guard = lock_env();
|
||||||
|
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||||
|
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||||
|
Some("gateway-token")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_oauth_proxy_auth_token_whitespace_dedicated_env_falls_back_to_gateway_token() {
|
||||||
|
let _guard = lock_env();
|
||||||
|
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", Some(" "));
|
||||||
|
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||||
|
Some("gateway-token")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_oauth_proxy_auth_token_returns_none_when_unset() {
|
||||||
|
let _guard = lock_env();
|
||||||
|
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||||
|
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||||
|
|
||||||
|
assert_eq!(crate::cli::oauth_defaults::oauth_proxy_auth_token(), None);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_strip_instance_prefix_with_colon() {
|
fn test_strip_instance_prefix_with_colon() {
|
||||||
use crate::cli::oauth_defaults::strip_instance_prefix;
|
use crate::cli::oauth_defaults::strip_instance_prefix;
|
||||||
|
|||||||
+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:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+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());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+20
-1
@@ -1,6 +1,6 @@
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use crate::config::helpers::{parse_bool_env, parse_option_env, parse_optional_env};
|
use crate::config::helpers::{optional_env, parse_bool_env, parse_option_env, parse_optional_env};
|
||||||
use crate::error::ConfigError;
|
use crate::error::ConfigError;
|
||||||
use crate::settings::Settings;
|
use crate::settings::Settings;
|
||||||
|
|
||||||
@@ -23,6 +23,8 @@ pub struct AgentConfig {
|
|||||||
pub max_cost_per_day_cents: Option<u64>,
|
pub max_cost_per_day_cents: Option<u64>,
|
||||||
/// Maximum LLM/tool actions per hour. None = unlimited.
|
/// Maximum LLM/tool actions per hour. None = unlimited.
|
||||||
pub max_actions_per_hour: Option<u64>,
|
pub max_actions_per_hour: Option<u64>,
|
||||||
|
/// Maximum daily LLM spend per user in cents. None = unlimited.
|
||||||
|
pub max_cost_per_user_per_day_cents: Option<u64>,
|
||||||
/// Maximum tool-call iterations per agentic loop invocation. Default 50.
|
/// Maximum tool-call iterations per agentic loop invocation. Default 50.
|
||||||
pub max_tool_iterations: usize,
|
pub max_tool_iterations: usize,
|
||||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||||
@@ -31,6 +33,13 @@ pub struct AgentConfig {
|
|||||||
pub default_timezone: String,
|
pub default_timezone: String,
|
||||||
/// Maximum tokens per job (0 = unlimited).
|
/// Maximum tokens per job (0 = unlimited).
|
||||||
pub max_tokens_per_job: u64,
|
pub max_tokens_per_job: u64,
|
||||||
|
/// Whether the deployment is multi-tenant (multiple users sharing one
|
||||||
|
/// instance). Auto-detected from GATEWAY_USER_TOKENS presence.
|
||||||
|
pub multi_tenant: bool,
|
||||||
|
/// Maximum concurrent LLM calls per user. None = use default (4).
|
||||||
|
pub max_llm_concurrent_per_user: Option<usize>,
|
||||||
|
/// Maximum concurrent jobs per user. None = use default (3).
|
||||||
|
pub max_jobs_concurrent_per_user: Option<usize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl AgentConfig {
|
impl AgentConfig {
|
||||||
@@ -49,10 +58,14 @@ impl AgentConfig {
|
|||||||
allow_local_tools: true,
|
allow_local_tools: true,
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
max_tool_iterations: 10,
|
max_tool_iterations: 10,
|
||||||
auto_approve_tools: true,
|
auto_approve_tools: true,
|
||||||
default_timezone: "UTC".to_string(),
|
default_timezone: "UTC".to_string(),
|
||||||
max_tokens_per_job: 0,
|
max_tokens_per_job: 0,
|
||||||
|
multi_tenant: false,
|
||||||
|
max_llm_concurrent_per_user: None,
|
||||||
|
max_jobs_concurrent_per_user: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -87,6 +100,7 @@ impl AgentConfig {
|
|||||||
allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?,
|
allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?,
|
||||||
max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?,
|
max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?,
|
||||||
max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?,
|
max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?,
|
||||||
|
max_cost_per_user_per_day_cents: parse_option_env("MAX_COST_PER_USER_PER_DAY_CENTS")?,
|
||||||
max_tool_iterations: parse_optional_env(
|
max_tool_iterations: parse_optional_env(
|
||||||
"AGENT_MAX_TOOL_ITERATIONS",
|
"AGENT_MAX_TOOL_ITERATIONS",
|
||||||
settings.agent.max_tool_iterations,
|
settings.agent.max_tool_iterations,
|
||||||
@@ -112,6 +126,11 @@ impl AgentConfig {
|
|||||||
"AGENT_MAX_TOKENS_PER_JOB",
|
"AGENT_MAX_TOKENS_PER_JOB",
|
||||||
settings.agent.max_tokens_per_job,
|
settings.agent.max_tokens_per_job,
|
||||||
)?,
|
)?,
|
||||||
|
// Auto-detected from GATEWAY_USER_TOKENS presence. Not a separate
|
||||||
|
// knob — multi-tenant mode is always implied by configuring user tokens.
|
||||||
|
multi_tenant: optional_env("GATEWAY_USER_TOKENS")?.is_some(),
|
||||||
|
max_llm_concurrent_per_user: parse_option_env("TENANT_MAX_LLM_CONCURRENT")?,
|
||||||
|
max_jobs_concurrent_per_user: parse_option_env("TENANT_MAX_JOBS_CONCURRENT")?,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -21,6 +21,9 @@ pub struct HeartbeatConfig {
|
|||||||
pub quiet_hours_end: Option<u32>,
|
pub quiet_hours_end: Option<u32>,
|
||||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||||
pub timezone: Option<String>,
|
pub timezone: Option<String>,
|
||||||
|
/// When true, cycle through all users with routines. Auto-detected from
|
||||||
|
/// GATEWAY_USER_TOKENS or set explicitly via HEARTBEAT_MULTI_TENANT.
|
||||||
|
pub multi_tenant: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for HeartbeatConfig {
|
impl Default for HeartbeatConfig {
|
||||||
@@ -34,6 +37,7 @@ impl Default for HeartbeatConfig {
|
|||||||
quiet_hours_start: None,
|
quiet_hours_start: None,
|
||||||
quiet_hours_end: None,
|
quiet_hours_end: None,
|
||||||
timezone: None,
|
timezone: None,
|
||||||
|
multi_tenant: false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -101,6 +105,12 @@ impl HeartbeatConfig {
|
|||||||
}
|
}
|
||||||
tz
|
tz
|
||||||
},
|
},
|
||||||
|
// Auto-detect multi-tenant mode from GATEWAY_USER_TOKENS presence,
|
||||||
|
// or allow explicit override via HEARTBEAT_MULTI_TENANT.
|
||||||
|
multi_tenant: parse_bool_env(
|
||||||
|
"HEARTBEAT_MULTI_TENANT",
|
||||||
|
optional_env("GATEWAY_USER_TOKENS")?.is_some(),
|
||||||
|
)?,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+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(),
|
||||||
|
|||||||
@@ -192,6 +192,9 @@ pub struct JobContext {
|
|||||||
/// but subsequent tools (e.g., `json`) may need the full output. This
|
/// but subsequent tools (e.g., `json`) may need the full output. This
|
||||||
/// stash stores the complete, unsanitized output so tools can reference
|
/// stash stores the complete, unsanitized output so tools can reference
|
||||||
/// previous results by ID via `$tool_call_id` parameter syntax.
|
/// previous results by ID via `$tool_call_id` parameter syntax.
|
||||||
|
///
|
||||||
|
/// Also used for cross-tool implicit state (keys prefixed with `__`) such
|
||||||
|
/// as `__routine_last_name` for fallback recovery in routine tool chains.
|
||||||
#[serde(skip)]
|
#[serde(skip)]
|
||||||
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
||||||
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
||||||
|
|||||||
@@ -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,
|
||||||
@@ -480,10 +530,24 @@ impl RoutineStore for LibSqlBackend {
|
|||||||
async fn get_webhook_routine_by_path(
|
async fn get_webhook_routine_by_path(
|
||||||
&self,
|
&self,
|
||||||
path: &str,
|
path: &str,
|
||||||
|
user_id: Option<&str>,
|
||||||
) -> Result<Option<Routine>, DatabaseError> {
|
) -> Result<Option<Routine>, DatabaseError> {
|
||||||
let conn = self.connect().await?;
|
let conn = self.connect().await?;
|
||||||
let mut rows = conn
|
let mut rows = if let Some(uid) = user_id {
|
||||||
.query(
|
conn.query(
|
||||||
|
&format!(
|
||||||
|
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
|
||||||
|
AND user_id = ?2 \
|
||||||
|
AND (json_extract(trigger_config, '$.path') = ?1 \
|
||||||
|
OR (json_extract(trigger_config, '$.path') IS NULL AND CAST(id AS TEXT) = ?1))",
|
||||||
|
ROUTINE_COLUMNS
|
||||||
|
),
|
||||||
|
params![path, uid],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||||
|
} else {
|
||||||
|
conn.query(
|
||||||
&format!(
|
&format!(
|
||||||
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
|
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
|
||||||
AND (json_extract(trigger_config, '$.path') = ?1 \
|
AND (json_extract(trigger_config, '$.path') = ?1 \
|
||||||
@@ -493,7 +557,8 @@ impl RoutineStore for LibSqlBackend {
|
|||||||
params![path],
|
params![path],
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||||
|
};
|
||||||
|
|
||||||
match rows
|
match rows
|
||||||
.next()
|
.next()
|
||||||
|
|||||||
@@ -528,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,
|
||||||
@@ -536,6 +545,7 @@ pub trait RoutineStore: Send + Sync {
|
|||||||
async fn get_webhook_routine_by_path(
|
async fn get_webhook_routine_by_path(
|
||||||
&self,
|
&self,
|
||||||
path: &str,
|
path: &str,
|
||||||
|
user_id: Option<&str>,
|
||||||
) -> Result<Option<Routine>, DatabaseError>;
|
) -> Result<Option<Routine>, DatabaseError>;
|
||||||
|
|
||||||
/// List routine runs that were dispatched as full_job but have not yet
|
/// List routine runs that were dispatched as full_job but have not yet
|
||||||
|
|||||||
+10
-1
@@ -510,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,
|
||||||
@@ -521,8 +529,9 @@ impl RoutineStore for PgBackend {
|
|||||||
async fn get_webhook_routine_by_path(
|
async fn get_webhook_routine_by_path(
|
||||||
&self,
|
&self,
|
||||||
path: &str,
|
path: &str,
|
||||||
|
user_id: Option<&str>,
|
||||||
) -> Result<Option<Routine>, DatabaseError> {
|
) -> Result<Option<Routine>, DatabaseError> {
|
||||||
self.store.get_webhook_routine_by_path(path).await
|
self.store.get_webhook_routine_by_path(path, user_id).await
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||||
|
|||||||
+455
-110
@@ -53,22 +53,6 @@ struct HostedOAuthFlowStart {
|
|||||||
flow: crate::cli::oauth_defaults::PendingOAuthFlow,
|
flow: crate::cli::oauth_defaults::PendingOAuthFlow,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn hosted_proxy_client_secret(
|
|
||||||
client_secret: &Option<String>,
|
|
||||||
builtin: Option<&crate::cli::oauth_defaults::OAuthCredentials>,
|
|
||||||
exchange_proxy_configured: bool,
|
|
||||||
) -> Option<String> {
|
|
||||||
if !exchange_proxy_configured {
|
|
||||||
return client_secret.clone();
|
|
||||||
}
|
|
||||||
|
|
||||||
let builtin_secret = builtin.map(|credentials| credentials.client_secret);
|
|
||||||
match (client_secret, builtin_secret) {
|
|
||||||
(Some(resolved), Some(baked_in)) if resolved == baked_in => None,
|
|
||||||
_ => client_secret.clone(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn normalize_oauth_callback_path(path: &str) -> String {
|
fn normalize_oauth_callback_path(path: &str) -> String {
|
||||||
let trimmed_path = path.trim_end_matches('/');
|
let trimmed_path = path.trim_end_matches('/');
|
||||||
if trimmed_path.is_empty() {
|
if trimmed_path.is_empty() {
|
||||||
@@ -419,9 +403,10 @@ pub struct ExtensionManager {
|
|||||||
/// when running in gateway mode, consumed by the web gateway's
|
/// when running in gateway mode, consumed by the web gateway's
|
||||||
/// `/oauth/callback` handler.
|
/// `/oauth/callback` handler.
|
||||||
pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry,
|
pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry,
|
||||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||||
/// Read once at construction from `GATEWAY_AUTH_TOKEN` env var.
|
/// Resolved once at construction from `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`,
|
||||||
gateway_token: Option<String>,
|
/// then `GATEWAY_AUTH_TOKEN` as a backward-compatible fallback.
|
||||||
|
oauth_proxy_auth_token: Option<String>,
|
||||||
/// Relay config captured at startup. Used by `auth_channel_relay` and
|
/// Relay config captured at startup. Used by `auth_channel_relay` and
|
||||||
/// `activate_channel_relay` instead of re-reading env vars.
|
/// `activate_channel_relay` instead of re-reading env vars.
|
||||||
relay_config: Option<crate::config::RelayConfig>,
|
relay_config: Option<crate::config::RelayConfig>,
|
||||||
@@ -551,7 +536,7 @@ impl ExtensionManager {
|
|||||||
activation_errors: RwLock::new(HashMap::new()),
|
activation_errors: RwLock::new(HashMap::new()),
|
||||||
sse_manager: RwLock::new(None),
|
sse_manager: RwLock::new(None),
|
||||||
pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(),
|
pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(),
|
||||||
gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(),
|
oauth_proxy_auth_token: crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||||
relay_config: crate::config::RelayConfig::from_env(),
|
relay_config: crate::config::RelayConfig::from_env(),
|
||||||
relay_event_tx: Arc::new(tokio::sync::Mutex::new(None)),
|
relay_event_tx: Arc::new(tokio::sync::Mutex::new(None)),
|
||||||
relay_signing_secret_cache: Arc::new(std::sync::Mutex::new(None)),
|
relay_signing_secret_cache: Arc::new(std::sync::Mutex::new(None)),
|
||||||
@@ -675,6 +660,66 @@ impl ExtensionManager {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Resolve the relay URL override for an extension from settings.
|
||||||
|
///
|
||||||
|
/// Returns `Some(url)` if a non-empty per-extension `relay_url` override is
|
||||||
|
/// set for the given extension; otherwise returns `None` and callers should
|
||||||
|
/// fall back to the env-level `RelayConfig`.
|
||||||
|
///
|
||||||
|
/// Uses `self.user_id` (owner scope) for consistency with `configure()`,
|
||||||
|
/// which also writes setting_path fields under the owner scope.
|
||||||
|
///
|
||||||
|
/// The override is validated: only `http` / `https` schemes are accepted
|
||||||
|
/// and the URL must not contain userinfo (embedded credentials). This
|
||||||
|
/// prevents a malicious override from exfiltrating the instance-wide relay
|
||||||
|
/// API key to an attacker-controlled host.
|
||||||
|
async fn effective_relay_url(&self, name: &str) -> Option<String> {
|
||||||
|
if let Some(ref store) = self.store {
|
||||||
|
let key = format!("extensions.{name}.relay_url");
|
||||||
|
if let Ok(Some(v)) = store.get_setting(&self.user_id, &key).await {
|
||||||
|
let url = v
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.trim().to_string())
|
||||||
|
.filter(|s| !s.is_empty());
|
||||||
|
if let Some(ref u) = url {
|
||||||
|
// Validate the override to prevent API-key exfiltration:
|
||||||
|
// only allow http(s) with no embedded credentials.
|
||||||
|
match url::Url::parse(u) {
|
||||||
|
Ok(parsed)
|
||||||
|
if (parsed.scheme() == "http" || parsed.scheme() == "https")
|
||||||
|
&& parsed.username().is_empty()
|
||||||
|
&& parsed.password().is_none() =>
|
||||||
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url_host = %parsed.host_str().unwrap_or("unknown"),
|
||||||
|
"effective_relay_url: using per-extension override from settings"
|
||||||
|
);
|
||||||
|
return url;
|
||||||
|
}
|
||||||
|
Ok(parsed) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
scheme = %parsed.scheme(),
|
||||||
|
has_userinfo = !parsed.username().is_empty() || parsed.password().is_some(),
|
||||||
|
"effective_relay_url: rejecting override — \
|
||||||
|
only http/https without embedded credentials is allowed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"effective_relay_url: rejecting override — invalid URL"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
/// Get the shared relay event sender for the webhook endpoint.
|
/// Get the shared relay event sender for the webhook endpoint.
|
||||||
pub fn relay_event_tx(
|
pub fn relay_event_tx(
|
||||||
&self,
|
&self,
|
||||||
@@ -891,24 +936,67 @@ impl ExtensionManager {
|
|||||||
*self.relay_channel_manager.write().await = Some(channel_manager);
|
*self.relay_channel_manager.write().await = Some(channel_manager);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Check if a channel name corresponds to a relay extension (has stored stream token
|
/// Check if a channel name corresponds to a relay extension (has stored team_id
|
||||||
/// or is tracked in the installed relay extensions set).
|
/// or is tracked in the installed relay extensions set).
|
||||||
pub async fn is_relay_channel(&self, name: &str, user_id: &str) -> bool {
|
pub async fn is_relay_channel(&self, name: &str, user_id: &str) -> bool {
|
||||||
// Check in-memory installed set first (supports no-store mode)
|
// Check in-memory installed set first (supports no-store mode)
|
||||||
if self.installed_relay_extensions.read().await.contains(name) {
|
if self.installed_relay_extensions.read().await.contains(name) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
// Then check for stored stream token
|
// Check for stored team_id (persisted across restarts by the OAuth callback)
|
||||||
self.secrets
|
if let Some(ref store) = self.store {
|
||||||
.exists(user_id, &format!("relay:{}:stream_token", name))
|
let key = format!("relay:{}:team_id", name);
|
||||||
.await
|
if let Ok(Some(v)) = store.get_setting(user_id, &key).await {
|
||||||
.unwrap_or(false)
|
return v.as_str().is_some_and(|s| !s.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether a stored `team_id` setting exists for the given relay extension.
|
||||||
|
///
|
||||||
|
/// Unlike [`is_relay_channel`], this does **not** consult the in-memory
|
||||||
|
/// `installed_relay_extensions` set — it only looks at the persistent settings
|
||||||
|
/// store. This distinction matters for `auth_channel_relay`: an extension can
|
||||||
|
/// be *installed* (present in the in-memory set) but not yet *authenticated*
|
||||||
|
/// (no OAuth completed, no team_id stored).
|
||||||
|
async fn has_stored_team_id(&self, name: &str, _user_id: &str) -> bool {
|
||||||
|
if let Some(ref store) = self.store {
|
||||||
|
let key = format!("relay:{}:team_id", name);
|
||||||
|
// Use owner scope (self.user_id) for consistency: the OAuth callback
|
||||||
|
// stores team_id under state.owner_id which maps to self.user_id.
|
||||||
|
match store.get_setting(&self.user_id, &key).await {
|
||||||
|
Ok(Some(v)) => {
|
||||||
|
let has_id = v.as_str().is_some_and(|s| !s.is_empty());
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
has_team_id = has_id,
|
||||||
|
"has_stored_team_id: checked store"
|
||||||
|
);
|
||||||
|
return has_id;
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
"has_stored_team_id: no team_id setting found"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"has_stored_team_id: failed to read from settings store"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
false
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Restore persisted relay channels after startup.
|
/// Restore persisted relay channels after startup.
|
||||||
///
|
///
|
||||||
/// Loads the persisted active channel list, filters to relay types (those with
|
/// Loads the persisted active channel list, filters to relay types (those with
|
||||||
/// a stored stream token), and activates each via `activate_stored_relay()`.
|
/// a stored team_id setting), and activates each via `activate_stored_relay()`.
|
||||||
/// Skips channels that are already active.
|
/// Skips channels that are already active.
|
||||||
///
|
///
|
||||||
/// Call this only after `set_relay_channel_manager()` or `set_channel_runtime()`.
|
/// Call this only after `set_relay_channel_manager()` or `set_channel_runtime()`.
|
||||||
@@ -1131,7 +1219,7 @@ impl ExtensionManager {
|
|||||||
/// Broadcast an extension status change to the web UI via SSE.
|
/// Broadcast an extension status change to the web UI via SSE.
|
||||||
async fn broadcast_extension_status(&self, name: &str, status: &str, message: Option<&str>) {
|
async fn broadcast_extension_status(&self, name: &str, status: &str, message: Option<&str>) {
|
||||||
if let Some(ref sse) = *self.sse_manager.read().await {
|
if let Some(ref sse) = *self.sse_manager.read().await {
|
||||||
sse.broadcast(crate::channels::web::types::SseEvent::ExtensionStatus {
|
sse.broadcast(ironclaw_common::AppEvent::ExtensionStatus {
|
||||||
extension_name: name.to_string(),
|
extension_name: name.to_string(),
|
||||||
status: status.to_string(),
|
status: status.to_string(),
|
||||||
message: message.map(|m| m.to_string()),
|
message: message.map(|m| m.to_string()),
|
||||||
@@ -1428,9 +1516,11 @@ impl ExtensionManager {
|
|||||||
if kind_filter.is_none() || kind_filter == Some(ExtensionKind::ChannelRelay) {
|
if kind_filter.is_none() || kind_filter == Some(ExtensionKind::ChannelRelay) {
|
||||||
let installed = self.installed_relay_extensions.read().await;
|
let installed = self.installed_relay_extensions.read().await;
|
||||||
let active_names = self.active_channel_names.read().await;
|
let active_names = self.active_channel_names.read().await;
|
||||||
|
let errors = self.activation_errors.read().await;
|
||||||
for name in installed.iter() {
|
for name in installed.iter() {
|
||||||
let active = active_names.contains(name);
|
let active = active_names.contains(name);
|
||||||
let has_token = self.is_relay_channel(name, user_id).await;
|
let authenticated = self.has_stored_team_id(name, user_id).await;
|
||||||
|
let activation_error = errors.get(name).cloned();
|
||||||
let registry_entry = self
|
let registry_entry = self
|
||||||
.registry
|
.registry
|
||||||
.get_with_kind(name, Some(ExtensionKind::ChannelRelay))
|
.get_with_kind(name, Some(ExtensionKind::ChannelRelay))
|
||||||
@@ -1443,13 +1533,13 @@ impl ExtensionManager {
|
|||||||
display_name,
|
display_name,
|
||||||
description,
|
description,
|
||||||
url: None,
|
url: None,
|
||||||
authenticated: has_token,
|
authenticated,
|
||||||
active,
|
active,
|
||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: false,
|
needs_setup: false,
|
||||||
has_auth: true,
|
has_auth: true,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error: None,
|
activation_error,
|
||||||
version: None,
|
version: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -1626,7 +1716,22 @@ impl ExtensionManager {
|
|||||||
self.persist_active_channels(user_id).await;
|
self.persist_active_channels(user_id).await;
|
||||||
self.activation_errors.write().await.remove(name);
|
self.activation_errors.write().await.remove(name);
|
||||||
|
|
||||||
// Remove stored stream token
|
// Remove stored team_id setting and clean up secrets
|
||||||
|
if let Some(ref store) = self.store
|
||||||
|
&& let Err(e) = store
|
||||||
|
.delete_setting(user_id, &format!("relay:{}:team_id", name))
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!(error = %e, name, "Failed to delete relay team_id setting on removal");
|
||||||
|
}
|
||||||
|
if let Err(e) = self
|
||||||
|
.secrets
|
||||||
|
.delete(user_id, &format!("relay:{}:oauth_state", name))
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!(error = %e, name, "Failed to delete relay oauth_state secret on removal");
|
||||||
|
}
|
||||||
|
// Clean up legacy stream_token secret from pre-webhook installs
|
||||||
let _ = self
|
let _ = self
|
||||||
.secrets
|
.secrets
|
||||||
.delete(user_id, &format!("relay:{}:stream_token", name))
|
.delete(user_id, &format!("relay:{}:stream_token", name))
|
||||||
@@ -2684,7 +2789,7 @@ impl ExtensionManager {
|
|||||||
user_id: user_id.to_string(),
|
user_id: user_id.to_string(),
|
||||||
secrets: Arc::clone(&self.secrets),
|
secrets: Arc::clone(&self.secrets),
|
||||||
sse_manager: self.sse_manager.read().await.clone(),
|
sse_manager: self.sse_manager.read().await.clone(),
|
||||||
gateway_token: self.gateway_token.clone(),
|
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||||
token_exchange_extra_params,
|
token_exchange_extra_params,
|
||||||
client_id_secret_name: if server.oauth.is_none() {
|
client_id_secret_name: if server.oauth.is_none() {
|
||||||
Some(server.client_id_secret_name())
|
Some(server.client_id_secret_name())
|
||||||
@@ -3179,7 +3284,7 @@ impl ExtensionManager {
|
|||||||
// apps. Sending the desktop secret would cause a client_id/secret
|
// apps. Sending the desktop secret would cause a client_id/secret
|
||||||
// mismatch because the container's GOOGLE_OAUTH_CLIENT_ID is the web
|
// mismatch because the container's GOOGLE_OAUTH_CLIENT_ID is the web
|
||||||
// app, not the desktop app.
|
// app, not the desktop app.
|
||||||
let proxy_client_secret = hosted_proxy_client_secret(
|
let proxy_client_secret = oauth_defaults::hosted_proxy_client_secret(
|
||||||
&client_secret,
|
&client_secret,
|
||||||
builtin.as_ref(),
|
builtin.as_ref(),
|
||||||
oauth_defaults::exchange_proxy_url().is_some(),
|
oauth_defaults::exchange_proxy_url().is_some(),
|
||||||
@@ -3201,7 +3306,7 @@ impl ExtensionManager {
|
|||||||
user_id: user_id.to_string(),
|
user_id: user_id.to_string(),
|
||||||
secrets: Arc::clone(&self.secrets),
|
secrets: Arc::clone(&self.secrets),
|
||||||
sse_manager: self.sse_manager.read().await.clone(),
|
sse_manager: self.sse_manager.read().await.clone(),
|
||||||
gateway_token: self.gateway_token.clone(),
|
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||||
client_id_secret_name: None,
|
client_id_secret_name: None,
|
||||||
created_at: std::time::Instant::now(),
|
created_at: std::time::Instant::now(),
|
||||||
@@ -3284,7 +3389,7 @@ impl ExtensionManager {
|
|||||||
}
|
}
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
// Broadcast SSE event
|
// Broadcast auth result event
|
||||||
let (success, message) = match result {
|
let (success, message) = match result {
|
||||||
Ok(()) => (true, format!("{} authenticated successfully", display_name)),
|
Ok(()) => (true, format!("{} authenticated successfully", display_name)),
|
||||||
Err(ref e) => (
|
Err(ref e) => (
|
||||||
@@ -3310,7 +3415,7 @@ impl ExtensionManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if let Some(ref sse) = sse_manager {
|
if let Some(ref sse) = sse_manager {
|
||||||
sse.broadcast(crate::channels::web::types::SseEvent::AuthCompleted {
|
sse.broadcast(ironclaw_common::AppEvent::AuthCompleted {
|
||||||
extension_name: ext_name,
|
extension_name: ext_name,
|
||||||
success,
|
success,
|
||||||
message,
|
message,
|
||||||
@@ -4181,26 +4286,75 @@ impl ExtensionManager {
|
|||||||
///
|
///
|
||||||
/// For Slack: initiates OAuth flow (redirect-based).
|
/// For Slack: initiates OAuth flow (redirect-based).
|
||||||
/// For Telegram: accepts a bot token, registers it with channel-relay,
|
/// For Telegram: accepts a bot token, registers it with channel-relay,
|
||||||
/// and stores the returned stream token.
|
/// and stores the team_id setting.
|
||||||
async fn auth_channel_relay(
|
async fn auth_channel_relay(
|
||||||
&self,
|
&self,
|
||||||
name: &str,
|
name: &str,
|
||||||
user_id: &str,
|
user_id: &str,
|
||||||
) -> Result<AuthResult, ExtensionError> {
|
) -> Result<AuthResult, ExtensionError> {
|
||||||
// Check if already authenticated (stream token exists)
|
tracing::debug!(
|
||||||
if self.is_relay_channel(name, user_id).await {
|
extension = %name,
|
||||||
|
user_id = %user_id,
|
||||||
|
"auth_channel_relay: starting"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Check if already authenticated by looking for a stored team_id.
|
||||||
|
// We intentionally skip the `installed_relay_extensions` in-memory set
|
||||||
|
// here because that set only tracks *installed* extensions — an extension
|
||||||
|
// can be installed (via registry) but not yet authenticated (no OAuth
|
||||||
|
// completed). Checking just `is_relay_channel()` would short-circuit
|
||||||
|
// to "authenticated" even when no team_id exists, preventing the OAuth
|
||||||
|
// flow from being offered to the user.
|
||||||
|
if self.has_stored_team_id(name, user_id).await {
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
"auth_channel_relay: already authenticated (team_id in store)"
|
||||||
|
);
|
||||||
return Ok(AuthResult::authenticated(name, ExtensionKind::ChannelRelay));
|
return Ok(AuthResult::authenticated(name, ExtensionKind::ChannelRelay));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
"auth_channel_relay: no stored team_id, initiating OAuth"
|
||||||
|
);
|
||||||
|
|
||||||
// Use relay config captured at startup
|
// Use relay config captured at startup
|
||||||
let relay_config = self.relay_config()?;
|
let relay_config = self.relay_config().map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"auth_channel_relay: relay config not available — \
|
||||||
|
CHANNEL_RELAY_URL and CHANNEL_RELAY_API_KEY must be set"
|
||||||
|
);
|
||||||
|
e
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Allow per-extension URL override from settings
|
||||||
|
let effective_url = self
|
||||||
|
.effective_relay_url(name)
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|| relay_config.url.clone());
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
"auth_channel_relay: creating relay client for OAuth"
|
||||||
|
);
|
||||||
|
|
||||||
let client = crate::channels::relay::RelayClient::new(
|
let client = crate::channels::relay::RelayClient::new(
|
||||||
relay_config.url.clone(),
|
effective_url.clone(),
|
||||||
relay_config.api_key.clone(),
|
relay_config.api_key.clone(),
|
||||||
relay_config.request_timeout_secs,
|
relay_config.request_timeout_secs,
|
||||||
)
|
)
|
||||||
.map_err(|e| ExtensionError::Config(e.to_string()))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
error = %e,
|
||||||
|
"auth_channel_relay: failed to create relay HTTP client"
|
||||||
|
);
|
||||||
|
ExtensionError::Config(e.to_string())
|
||||||
|
})?;
|
||||||
|
|
||||||
// Generate CSRF nonce — IronClaw validates this on the callback to ensure
|
// Generate CSRF nonce — IronClaw validates this on the callback to ensure
|
||||||
// the OAuth completion is legitimate. Channel-relay embeds it in the signed
|
// the OAuth completion is legitimate. Channel-relay embeds it in the signed
|
||||||
@@ -4212,18 +4366,44 @@ impl ExtensionManager {
|
|||||||
self.secrets
|
self.secrets
|
||||||
.create(user_id, CreateSecretParams::new(&state_key, &state_nonce))
|
.create(user_id, CreateSecretParams::new(&state_key, &state_nonce))
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ExtensionError::AuthFailed(format!("Failed to store OAuth state: {e}")))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"auth_channel_relay: failed to store OAuth state nonce"
|
||||||
|
);
|
||||||
|
ExtensionError::AuthFailed(format!("Failed to store OAuth state: {e}"))
|
||||||
|
})?;
|
||||||
|
|
||||||
// Channel-relay derives all URLs from trusted instance_url in chat-api.
|
// Channel-relay derives all URLs from trusted instance_url in chat-api.
|
||||||
// We only pass the nonce for CSRF validation on the callback.
|
// We only pass the nonce for CSRF validation on the callback.
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
"auth_channel_relay: calling initiate_oauth on channel-relay"
|
||||||
|
);
|
||||||
match client.initiate_oauth(Some(&state_nonce)).await {
|
match client.initiate_oauth(Some(&state_nonce)).await {
|
||||||
Ok(auth_url) => Ok(AuthResult::awaiting_authorization(
|
Ok(auth_url) => {
|
||||||
name,
|
tracing::info!(
|
||||||
ExtensionKind::ChannelRelay,
|
extension = %name,
|
||||||
auth_url,
|
"auth_channel_relay: OAuth URL obtained, awaiting user authorization"
|
||||||
"redirect".to_string(),
|
);
|
||||||
)),
|
Ok(AuthResult::awaiting_authorization(
|
||||||
Err(e) => Err(ExtensionError::AuthFailed(e.to_string())),
|
name,
|
||||||
|
ExtensionKind::ChannelRelay,
|
||||||
|
auth_url,
|
||||||
|
"redirect".to_string(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
error = %e,
|
||||||
|
"auth_channel_relay: initiate_oauth call to channel-relay failed"
|
||||||
|
);
|
||||||
|
Err(ExtensionError::AuthFailed(e.to_string()))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4233,46 +4413,112 @@ impl ExtensionManager {
|
|||||||
name: &str,
|
name: &str,
|
||||||
user_id: &str,
|
user_id: &str,
|
||||||
) -> Result<ActivateResult, ExtensionError> {
|
) -> Result<ActivateResult, ExtensionError> {
|
||||||
let token_key = format!("relay:{}:stream_token", name);
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
user_id = %user_id,
|
||||||
|
"activate_channel_relay: starting"
|
||||||
|
);
|
||||||
|
|
||||||
let team_id_key = format!("relay:{}:team_id", name);
|
let team_id_key = format!("relay:{}:team_id", name);
|
||||||
|
|
||||||
// Check if we have a stream token
|
// Get team_id from settings (stored by the OAuth callback)
|
||||||
// Verify auth: stream token must exist (even though we don't use it in this constructor path)
|
|
||||||
let _stream_token = match self.secrets.get_decrypted(user_id, &token_key).await {
|
|
||||||
Ok(secret) => secret.expose().to_string(),
|
|
||||||
Err(_) => {
|
|
||||||
return Err(ExtensionError::AuthRequired);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// Get team_id from settings
|
|
||||||
let team_id = if let Some(ref store) = self.store {
|
let team_id = if let Some(ref store) = self.store {
|
||||||
store
|
match store.get_setting(user_id, &team_id_key).await {
|
||||||
.get_setting(user_id, &team_id_key)
|
Ok(Some(v)) => {
|
||||||
.await
|
let id = v.as_str().map(|s| s.to_string()).unwrap_or_default();
|
||||||
.ok()
|
tracing::debug!(
|
||||||
.flatten()
|
extension = %name,
|
||||||
.and_then(|v| v.as_str().map(|s| s.to_string()))
|
team_id_empty = id.is_empty(),
|
||||||
.unwrap_or_default()
|
"activate_channel_relay: loaded team_id from store"
|
||||||
|
);
|
||||||
|
id
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
setting_key = %team_id_key,
|
||||||
|
"activate_channel_relay: no team_id in settings store"
|
||||||
|
);
|
||||||
|
String::new()
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"activate_channel_relay: failed to read team_id from settings store"
|
||||||
|
);
|
||||||
|
String::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
"activate_channel_relay: no settings store available"
|
||||||
|
);
|
||||||
String::new()
|
String::new()
|
||||||
};
|
};
|
||||||
|
|
||||||
|
if team_id.is_empty() {
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
"activate_channel_relay: team_id is empty, returning AuthRequired"
|
||||||
|
);
|
||||||
|
return Err(ExtensionError::AuthRequired);
|
||||||
|
}
|
||||||
|
|
||||||
// Use relay config captured at startup
|
// Use relay config captured at startup
|
||||||
let relay_config = self.relay_config()?;
|
let relay_config = self.relay_config().map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
error = %e,
|
||||||
|
"activate_channel_relay: relay config not available"
|
||||||
|
);
|
||||||
|
e
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Allow per-extension URL override from settings
|
||||||
|
let effective_url = self
|
||||||
|
.effective_relay_url(name)
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|| relay_config.url.clone());
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
"activate_channel_relay: relay config loaded"
|
||||||
|
);
|
||||||
|
|
||||||
let instance_id = self.relay_instance_id(relay_config, user_id);
|
let instance_id = self.relay_instance_id(relay_config, user_id);
|
||||||
|
|
||||||
let client = crate::channels::relay::RelayClient::new(
|
let client = crate::channels::relay::RelayClient::new(
|
||||||
relay_config.url.clone(),
|
effective_url.clone(),
|
||||||
relay_config.api_key.clone(),
|
relay_config.api_key.clone(),
|
||||||
relay_config.request_timeout_secs,
|
relay_config.request_timeout_secs,
|
||||||
)
|
)
|
||||||
.map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?;
|
.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
error = %e,
|
||||||
|
"activate_channel_relay: failed to create relay HTTP client"
|
||||||
|
);
|
||||||
|
ExtensionError::ActivationFailed(e.to_string())
|
||||||
|
})?;
|
||||||
|
|
||||||
// Fetch the per-instance signing secret from channel-relay.
|
// Fetch the per-instance signing secret from channel-relay.
|
||||||
// This must succeed — there is no fallback.
|
// This must succeed — there is no fallback.
|
||||||
|
tracing::debug!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
"activate_channel_relay: fetching signing secret from channel-relay"
|
||||||
|
);
|
||||||
let signing_secret = client.get_signing_secret(&team_id).await.map_err(|e| {
|
let signing_secret = client.get_signing_secret(&team_id).await.map_err(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
relay_url = %effective_url,
|
||||||
|
error = %e,
|
||||||
|
"activate_channel_relay: failed to fetch signing secret from channel-relay"
|
||||||
|
);
|
||||||
ExtensionError::Config(format!("Failed to fetch relay signing secret: {e}"))
|
ExtensionError::Config(format!("Failed to fetch relay signing secret: {e}"))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
@@ -4291,16 +4537,29 @@ impl ExtensionManager {
|
|||||||
// Hot-add to channel manager
|
// Hot-add to channel manager
|
||||||
let cm_guard = self.relay_channel_manager.read().await;
|
let cm_guard = self.relay_channel_manager.read().await;
|
||||||
let channel_mgr = cm_guard.as_ref().ok_or_else(|| {
|
let channel_mgr = cm_guard.as_ref().ok_or_else(|| {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
"activate_channel_relay: channel manager not initialized"
|
||||||
|
);
|
||||||
ExtensionError::ActivationFailed("Channel manager not initialized".to_string())
|
ExtensionError::ActivationFailed("Channel manager not initialized".to_string())
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
channel_mgr
|
channel_mgr.hot_add(Box::new(channel)).await.map_err(|e| {
|
||||||
.hot_add(Box::new(channel))
|
tracing::warn!(
|
||||||
.await
|
extension = %name,
|
||||||
.map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?;
|
error = %e,
|
||||||
|
"activate_channel_relay: hot_add to channel manager failed"
|
||||||
|
);
|
||||||
|
ExtensionError::ActivationFailed(e.to_string())
|
||||||
|
})?;
|
||||||
|
|
||||||
if let Ok(mut cache) = self.relay_signing_secret_cache.lock() {
|
if let Ok(mut cache) = self.relay_signing_secret_cache.lock() {
|
||||||
*cache = Some(signing_secret);
|
*cache = Some(signing_secret);
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
"activate_channel_relay: failed to cache signing secret (mutex poisoned)"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Store the event sender so the web gateway's relay webhook endpoint can push events
|
// Store the event sender so the web gateway's relay webhook endpoint can push events
|
||||||
@@ -4318,6 +4577,12 @@ impl ExtensionManager {
|
|||||||
self.broadcast_extension_status(name, "active", Some(&status_msg))
|
self.broadcast_extension_status(name, "active", Some(&status_msg))
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
extension = %name,
|
||||||
|
instance_id = %instance_id,
|
||||||
|
"activate_channel_relay: relay channel activated successfully"
|
||||||
|
);
|
||||||
|
|
||||||
Ok(ActivateResult {
|
Ok(ActivateResult {
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
kind: ExtensionKind::ChannelRelay,
|
kind: ExtensionKind::ChannelRelay,
|
||||||
@@ -4367,11 +4632,11 @@ impl ExtensionManager {
|
|||||||
return Ok(ExtensionKind::WasmChannel);
|
return Ok(ExtensionKind::WasmChannel);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check channel-relay extensions (installed in memory or has stored token)
|
// Check channel-relay extensions (installed in memory or has stored team_id)
|
||||||
if self.installed_relay_extensions.read().await.contains(name) {
|
if self.installed_relay_extensions.read().await.contains(name) {
|
||||||
return Ok(ExtensionKind::ChannelRelay);
|
return Ok(ExtensionKind::ChannelRelay);
|
||||||
}
|
}
|
||||||
// Also check if there's a stored stream token (persisted across restarts)
|
// Also check if there's a stored team_id setting (persisted across restarts)
|
||||||
if self.is_relay_channel(name, user_id).await {
|
if self.is_relay_channel(name, user_id).await {
|
||||||
return Ok(ExtensionKind::ChannelRelay);
|
return Ok(ExtensionKind::ChannelRelay);
|
||||||
}
|
}
|
||||||
@@ -4597,6 +4862,41 @@ impl ExtensionManager {
|
|||||||
}
|
}
|
||||||
Ok(ExtensionSetupSchema { secrets, fields })
|
Ok(ExtensionSetupSchema { secrets, fields })
|
||||||
}
|
}
|
||||||
|
ExtensionKind::ChannelRelay => {
|
||||||
|
let relay_url_key = format!("extensions.{name}.relay_url");
|
||||||
|
let current_url = if let Some(ref store) = self.store {
|
||||||
|
match store.get_setting(&self.user_id, &relay_url_key).await {
|
||||||
|
Ok(value_opt) => value_opt
|
||||||
|
.and_then(|v| v.as_str().map(|s| s.to_string()))
|
||||||
|
.filter(|s| !s.is_empty()),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
extension = %name,
|
||||||
|
setting_key = %relay_url_key,
|
||||||
|
error = %e,
|
||||||
|
"get_setup_schema: failed to read relay_url from settings"
|
||||||
|
);
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let env_url = self.relay_config.as_ref().map(|c| c.url.as_str());
|
||||||
|
Ok(ExtensionSetupSchema {
|
||||||
|
secrets: Vec::new(),
|
||||||
|
fields: vec![crate::channels::web::types::SetupFieldInfo {
|
||||||
|
name: "relay_url".to_string(),
|
||||||
|
prompt: format!(
|
||||||
|
"Channel-relay service URL (leave empty to use env default{})",
|
||||||
|
env_url.map(|u| format!(": {u}")).unwrap_or_default()
|
||||||
|
),
|
||||||
|
optional: true,
|
||||||
|
provided: current_url.is_some(),
|
||||||
|
input_type: crate::tools::wasm::ToolSetupFieldInputType::Text,
|
||||||
|
}],
|
||||||
|
})
|
||||||
|
}
|
||||||
_ => Ok(ExtensionSetupSchema {
|
_ => Ok(ExtensionSetupSchema {
|
||||||
secrets: Vec::new(),
|
secrets: Vec::new(),
|
||||||
fields: Vec::new(),
|
fields: Vec::new(),
|
||||||
@@ -5000,9 +5300,15 @@ impl ExtensionManager {
|
|||||||
(names, Vec::new())
|
(names, Vec::new())
|
||||||
}
|
}
|
||||||
ExtensionKind::ChannelRelay => {
|
ExtensionKind::ChannelRelay => {
|
||||||
let mut names = std::collections::HashSet::new();
|
let relay_fields = vec![crate::tools::wasm::ToolFieldSetupSchema {
|
||||||
names.insert(format!("relay:{}:stream_token", name));
|
name: "relay_url".to_string(),
|
||||||
(names, Vec::new())
|
prompt: "Channel-relay service URL override".to_string(),
|
||||||
|
optional: true,
|
||||||
|
setting_path: Some(format!("extensions.{name}.relay_url")),
|
||||||
|
input_type: crate::tools::wasm::ToolSetupFieldInputType::Text,
|
||||||
|
restart_required: false,
|
||||||
|
}];
|
||||||
|
(std::collections::HashSet::new(), relay_fields)
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -5094,13 +5400,28 @@ impl ExtensionManager {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
let trimmed = field_value.trim();
|
let trimmed = field_value.trim();
|
||||||
|
let field_def = setup_field_defs.get(field_name);
|
||||||
|
|
||||||
|
// Empty value on an optional field with a setting_path: clear the
|
||||||
|
// stored override so the system reverts to the env/default value.
|
||||||
if trimmed.is_empty() {
|
if trimmed.is_empty() {
|
||||||
|
if let Some(def) = field_def
|
||||||
|
&& def.optional
|
||||||
|
{
|
||||||
|
stored_fields.remove(field_name);
|
||||||
|
if let Some(setting_path) = &def.setting_path {
|
||||||
|
Self::validate_setup_setting_path(name, setting_path)?;
|
||||||
|
if let Some(store) = self.store.as_ref() {
|
||||||
|
let _ = store.delete_setting(&self.user_id, setting_path).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
stored_fields.insert(field_name.clone(), trimmed.to_string());
|
stored_fields.insert(field_name.clone(), trimmed.to_string());
|
||||||
|
|
||||||
if let Some(field_def) = setup_field_defs.get(field_name) {
|
if let Some(field_def) = field_def {
|
||||||
if field_def.restart_required {
|
if field_def.restart_required {
|
||||||
restart_required = true;
|
restart_required = true;
|
||||||
}
|
}
|
||||||
@@ -5434,7 +5755,9 @@ impl ExtensionManager {
|
|||||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||||
server.token_secret_name()
|
server.token_secret_name()
|
||||||
}
|
}
|
||||||
ExtensionKind::ChannelRelay => format!("relay:{}:stream_token", name),
|
ExtensionKind::ChannelRelay => {
|
||||||
|
return Err(ExtensionError::AuthRequired);
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut secrets = std::collections::HashMap::new();
|
let mut secrets = std::collections::HashMap::new();
|
||||||
@@ -5702,7 +6025,7 @@ mod tests {
|
|||||||
use crate::extensions::manager::{
|
use crate::extensions::manager::{
|
||||||
ChannelRuntimeState, FallbackDecision, TelegramBindingData, TelegramBindingResult,
|
ChannelRuntimeState, FallbackDecision, TelegramBindingData, TelegramBindingResult,
|
||||||
TelegramOwnerBindingState, build_wasm_channel_runtime_config_updates,
|
TelegramOwnerBindingState, build_wasm_channel_runtime_config_updates,
|
||||||
combine_install_errors, fallback_decision, hosted_proxy_client_secret, infer_kind_from_url,
|
combine_install_errors, fallback_decision, infer_kind_from_url,
|
||||||
normalize_hosted_callback_url, send_telegram_text_message,
|
normalize_hosted_callback_url, send_telegram_text_message,
|
||||||
telegram_message_matches_verification_code,
|
telegram_message_matches_verification_code,
|
||||||
};
|
};
|
||||||
@@ -7043,7 +7366,7 @@ mod tests {
|
|||||||
let dir = tempfile::tempdir().expect("temp dir");
|
let dir = tempfile::tempdir().expect("temp dir");
|
||||||
let mgr = make_test_manager(None, dir.path().to_path_buf());
|
let mgr = make_test_manager(None, dir.path().to_path_buf());
|
||||||
|
|
||||||
// No token stored → not a relay channel
|
// No store configured, no team_id → not a relay channel
|
||||||
assert!(!mgr.is_relay_channel("slack-relay", "test").await);
|
assert!(!mgr.is_relay_channel("slack-relay", "test").await);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -7062,6 +7385,39 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Regression: installed-but-not-authenticated relay must NOT short-circuit
|
||||||
|
/// `auth_channel_relay()` to "authenticated". Previously, `auth_channel_relay`
|
||||||
|
/// called `is_relay_channel()` which checked the in-memory
|
||||||
|
/// `installed_relay_extensions` set; that returned `true` even when no team_id
|
||||||
|
/// existed in the store, so the OAuth URL was never offered.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_auth_channel_relay_installed_without_team_id_is_not_authenticated() {
|
||||||
|
let dir = tempfile::tempdir().expect("temp dir");
|
||||||
|
let mgr = make_test_manager(None, dir.path().to_path_buf());
|
||||||
|
|
||||||
|
// Mark as installed (simulates clicking Install in the UI)
|
||||||
|
mgr.installed_relay_extensions
|
||||||
|
.write()
|
||||||
|
.await
|
||||||
|
.insert("slack-relay".to_string());
|
||||||
|
|
||||||
|
// Without a stored team_id, auth should NOT return authenticated.
|
||||||
|
// It should fail because relay config is missing (no CHANNEL_RELAY_URL),
|
||||||
|
// but the key assertion is that it does NOT return Ok(authenticated).
|
||||||
|
let result = mgr.auth_channel_relay("slack-relay", "test").await;
|
||||||
|
match result {
|
||||||
|
Ok(ref auth_result) if auth_result.is_authenticated() => {
|
||||||
|
panic!(
|
||||||
|
"auth_channel_relay returned authenticated for installed-but-no-team-id relay; \
|
||||||
|
expected either an OAuth URL or a config error"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Config error (no relay URL) or awaiting_authorization — both are correct
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_remove_relay_shuts_down_via_relay_channel_manager() {
|
async fn test_remove_relay_shuts_down_via_relay_channel_manager() {
|
||||||
// Regression: remove() only checked channel_runtime for shutdown, missing
|
// Regression: remove() only checked channel_runtime for shutdown, missing
|
||||||
@@ -7862,19 +8218,13 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.insert("test-relay".to_string());
|
.insert("test-relay".to_string());
|
||||||
|
|
||||||
// configure() should dispatch to activate_channel_relay(), not
|
// configure() with empty secrets should dispatch to
|
||||||
// activate_wasm_channel(). Both will fail (no runtime configured),
|
// activate_channel_relay(), not activate_wasm_channel(). Relay auth
|
||||||
// but the error should be about relay config, not WASM channels.
|
// is OAuth-only so there are no manual secrets to pass.
|
||||||
let mut secrets = std::collections::HashMap::new();
|
|
||||||
secrets.insert(
|
|
||||||
"relay:test-relay:stream_token".to_string(),
|
|
||||||
"tok".to_string(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let result = mgr
|
let result = mgr
|
||||||
.configure(
|
.configure(
|
||||||
"test-relay",
|
"test-relay",
|
||||||
&secrets,
|
&std::collections::HashMap::new(),
|
||||||
&std::collections::HashMap::new(),
|
&std::collections::HashMap::new(),
|
||||||
"test",
|
"test",
|
||||||
)
|
)
|
||||||
@@ -7886,7 +8236,6 @@ mod tests {
|
|||||||
);
|
);
|
||||||
|
|
||||||
let result = result.unwrap();
|
let result = result.unwrap();
|
||||||
// Activation will fail (no relay config), but secrets should still be stored
|
|
||||||
assert!(
|
assert!(
|
||||||
!result.activated,
|
!result.activated,
|
||||||
"activation should fail without relay config"
|
"activation should fail without relay config"
|
||||||
@@ -7896,15 +8245,6 @@ mod tests {
|
|||||||
"error should not mention WASM — got: {}",
|
"error should not mention WASM — got: {}",
|
||||||
result.message
|
result.message
|
||||||
);
|
);
|
||||||
|
|
||||||
// Verify the secret was stored
|
|
||||||
assert!(
|
|
||||||
mgr.secrets
|
|
||||||
.exists("test", "relay:test-relay:stream_token")
|
|
||||||
.await
|
|
||||||
.unwrap_or(false),
|
|
||||||
"configure should have stored the relay stream token"
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
#[test]
|
#[test]
|
||||||
fn test_validation_failed_is_distinct_error_variant() {
|
fn test_validation_failed_is_distinct_error_variant() {
|
||||||
@@ -7970,7 +8310,8 @@ mod tests {
|
|||||||
let builtin_ref = builtin.as_ref();
|
let builtin_ref = builtin.as_ref();
|
||||||
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
||||||
|
|
||||||
let result = hosted_proxy_client_secret(&secret, builtin_ref, true);
|
let result =
|
||||||
|
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, true);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result, None,
|
result, None,
|
||||||
"built-in desktop secret must be suppressed when the exchange proxy is configured"
|
"built-in desktop secret must be suppressed when the exchange proxy is configured"
|
||||||
@@ -7982,7 +8323,8 @@ mod tests {
|
|||||||
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
||||||
let secret = Some("user-entered-custom-secret".to_string());
|
let secret = Some("user-entered-custom-secret".to_string());
|
||||||
|
|
||||||
let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
let result =
|
||||||
|
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result,
|
result,
|
||||||
Some("user-entered-custom-secret".to_string()),
|
Some("user-entered-custom-secret".to_string()),
|
||||||
@@ -7996,7 +8338,8 @@ mod tests {
|
|||||||
let builtin_ref = builtin.as_ref();
|
let builtin_ref = builtin.as_ref();
|
||||||
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
||||||
|
|
||||||
let result = hosted_proxy_client_secret(&secret, builtin_ref, false);
|
let result =
|
||||||
|
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, false);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result, secret,
|
result, secret,
|
||||||
"built-in secret must be kept when the callback will exchange directly"
|
"built-in secret must be kept when the callback will exchange directly"
|
||||||
@@ -8007,7 +8350,8 @@ mod tests {
|
|||||||
fn test_proxy_client_secret_none_stays_none() {
|
fn test_proxy_client_secret_none_stays_none() {
|
||||||
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
||||||
|
|
||||||
let result = hosted_proxy_client_secret(&None, builtin.as_ref(), true);
|
let result =
|
||||||
|
crate::cli::oauth_defaults::hosted_proxy_client_secret(&None, builtin.as_ref(), true);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result, None,
|
result, None,
|
||||||
"None secret stays None even when the exchange proxy is configured"
|
"None secret stays None even when the exchange proxy is configured"
|
||||||
@@ -8021,7 +8365,8 @@ mod tests {
|
|||||||
assert!(builtin.is_none());
|
assert!(builtin.is_none());
|
||||||
|
|
||||||
let secret = Some("dcr-secret".to_string());
|
let secret = Some("dcr-secret".to_string());
|
||||||
let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
let result =
|
||||||
|
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result,
|
result,
|
||||||
Some("dcr-secret".to_string()),
|
Some("dcr-secret".to_string()),
|
||||||
|
|||||||
+47
-3
@@ -1162,15 +1162,25 @@ impl Store {
|
|||||||
pub async fn get_webhook_routine_by_path(
|
pub async fn get_webhook_routine_by_path(
|
||||||
&self,
|
&self,
|
||||||
path: &str,
|
path: &str,
|
||||||
|
user_id: Option<&str>,
|
||||||
) -> Result<Option<Routine>, DatabaseError> {
|
) -> Result<Option<Routine>, DatabaseError> {
|
||||||
let conn = self.conn().await?;
|
let conn = self.conn().await?;
|
||||||
let row = conn
|
let row = if let Some(uid) = user_id {
|
||||||
.query_opt(
|
conn.query_opt(
|
||||||
|
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
|
||||||
|
AND user_id = $2 \
|
||||||
|
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
|
||||||
|
&[&path, &uid],
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
} else {
|
||||||
|
conn.query_opt(
|
||||||
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
|
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
|
||||||
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
|
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
|
||||||
&[&path],
|
&[&path],
|
||||||
)
|
)
|
||||||
.await?;
|
.await?
|
||||||
|
};
|
||||||
row.as_ref().map(row_to_routine).transpose()
|
row.as_ref().map(row_to_routine).transpose()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1403,6 +1413,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,
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ pub mod service;
|
|||||||
pub mod settings;
|
pub mod settings;
|
||||||
pub mod setup;
|
pub mod setup;
|
||||||
pub mod skills;
|
pub mod skills;
|
||||||
|
pub mod tenant;
|
||||||
pub mod timezone;
|
pub mod timezone;
|
||||||
pub mod tools;
|
pub mod tools;
|
||||||
pub mod tracing_fmt;
|
pub mod tracing_fmt;
|
||||||
|
|||||||
@@ -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"),
|
||||||
|
|||||||
+2
-1
@@ -63,7 +63,8 @@ pub use provider::{
|
|||||||
};
|
};
|
||||||
pub use reasoning::{
|
pub use reasoning::{
|
||||||
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN,
|
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN,
|
||||||
TOOL_INTENT_NUDGE, TokenUsage, ToolSelection, is_silent_reply, llm_signals_tool_intent,
|
TOOL_INTENT_NUDGE, TRUNCATED_TOOL_CALL_NOTICE, TokenUsage, ToolSelection, is_silent_reply,
|
||||||
|
llm_signals_tool_intent,
|
||||||
};
|
};
|
||||||
pub use recording::RecordingLlm;
|
pub use recording::RecordingLlm;
|
||||||
pub use registry::{ProviderDefinition, ProviderProtocol, ProviderRegistry};
|
pub use registry::{ProviderDefinition, ProviderProtocol, ProviderRegistry};
|
||||||
|
|||||||
+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."),
|
||||||
|
|||||||
+296
-20
@@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize};
|
|||||||
use crate::llm::error::LlmError;
|
use crate::llm::error::LlmError;
|
||||||
|
|
||||||
use crate::llm::{
|
use crate::llm::{
|
||||||
ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest,
|
ChatMessage, CompletionRequest, FinishReason, LlmProvider, Role, ToolCall,
|
||||||
ToolDefinition,
|
ToolCompletionRequest, ToolDefinition,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Token the agent returns when it has nothing to say (e.g. in group chats).
|
/// Token the agent returns when it has nothing to say (e.g. in group chats).
|
||||||
@@ -23,6 +23,13 @@ You said you would perform an action, but you did not include any tool calls.\n\
|
|||||||
Do NOT describe what you intend to do — actually call the tool now.\n\
|
Do NOT describe what you intend to do — actually call the tool now.\n\
|
||||||
Use the tool_calls mechanism to invoke the appropriate tool.";
|
Use the tool_calls mechanism to invoke the appropriate tool.";
|
||||||
|
|
||||||
|
/// Notice injected when the LLM's response was truncated mid-tool-call,
|
||||||
|
/// causing incomplete parameters. Tells the LLM to try a different approach.
|
||||||
|
pub const TRUNCATED_TOOL_CALL_NOTICE: &str = "\
|
||||||
|
Your previous response was truncated while generating tool call parameters. \
|
||||||
|
The tool calls were discarded. Please try a different approach — \
|
||||||
|
summarize or transform the data instead of echoing it verbatim in a tool call.";
|
||||||
|
|
||||||
/// Seed value used as the second argument to `generate_tool_call_id` when
|
/// Seed value used as the second argument to `generate_tool_call_id` when
|
||||||
/// recovering tool calls from malformed LLM text responses. This must differ
|
/// recovering tool calls from malformed LLM text responses. This must differ
|
||||||
/// from the `0` seed used in `rig_adapter::normalized_tool_call_id` to avoid
|
/// from the `0` seed used in `rig_adapter::normalized_tool_call_id` to avoid
|
||||||
@@ -194,11 +201,17 @@ pub struct ReasoningContext {
|
|||||||
pub metadata: std::collections::HashMap<String, String>,
|
pub metadata: std::collections::HashMap<String, String>,
|
||||||
/// When true, force a text-only response (ignore available tools).
|
/// When true, force a text-only response (ignore available tools).
|
||||||
/// Used by the agentic loop to guarantee termination near the iteration limit.
|
/// Used by the agentic loop to guarantee termination near the iteration limit.
|
||||||
|
/// Sticky: once set, never cleared within a loop invocation. Callers must
|
||||||
|
/// create a fresh `ReasoningContext` per `run_agentic_loop()` call.
|
||||||
pub force_text: bool,
|
pub force_text: bool,
|
||||||
/// Pre-built system prompt. When set, `respond_with_tools` uses this directly
|
/// Pre-built system prompt. When set, `respond_with_tools` uses this directly
|
||||||
/// instead of calling `build_system_prompt_with_tools`. Allows callers to build
|
/// instead of calling `build_system_prompt_with_tools`. Allows callers to build
|
||||||
/// the prompt once and reuse it across iterations.
|
/// the prompt once and reuse it across iterations.
|
||||||
pub system_prompt: Option<String>,
|
pub system_prompt: Option<String>,
|
||||||
|
/// Per-user model override. When set, completion requests use this model
|
||||||
|
/// instead of the provider's default. Only effective with providers that
|
||||||
|
/// support per-request model overrides (e.g. NearAI).
|
||||||
|
pub model_override: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ReasoningContext {
|
impl ReasoningContext {
|
||||||
@@ -212,6 +225,7 @@ impl ReasoningContext {
|
|||||||
metadata: std::collections::HashMap::new(),
|
metadata: std::collections::HashMap::new(),
|
||||||
force_text: false,
|
force_text: false,
|
||||||
system_prompt: None,
|
system_prompt: None,
|
||||||
|
model_override: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -344,6 +358,7 @@ pub enum RespondResult {
|
|||||||
pub struct RespondOutput {
|
pub struct RespondOutput {
|
||||||
pub result: RespondResult,
|
pub result: RespondResult,
|
||||||
pub usage: TokenUsage,
|
pub usage: TokenUsage,
|
||||||
|
pub finish_reason: FinishReason,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reasoning engine for the agent.
|
/// Reasoning engine for the agent.
|
||||||
@@ -525,17 +540,46 @@ 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();
|
// If the response was truncated, tool call parameters are likely incomplete.
|
||||||
|
// Return empty so the caller can fall through to respond_with_tools() which
|
||||||
|
// has a larger output token budget.
|
||||||
|
if response.finish_reason == FinishReason::Length {
|
||||||
|
tracing::warn!(
|
||||||
|
"select_tools response truncated (finish_reason=Length), \
|
||||||
|
discarding potentially incomplete tool selections"
|
||||||
|
);
|
||||||
|
return Ok(vec![]);
|
||||||
|
}
|
||||||
|
|
||||||
|
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();
|
||||||
|
|
||||||
@@ -653,6 +697,9 @@ Respond in JSON format:
|
|||||||
.with_temperature(0.7)
|
.with_temperature(0.7)
|
||||||
.with_tool_choice("auto");
|
.with_tool_choice("auto");
|
||||||
request.metadata = context.metadata.clone();
|
request.metadata = context.metadata.clone();
|
||||||
|
if let Some(ref model) = context.model_override {
|
||||||
|
request.model = Some(model.clone());
|
||||||
|
}
|
||||||
|
|
||||||
let response = self.llm.complete_with_tools(request).await?;
|
let response = self.llm.complete_with_tools(request).await?;
|
||||||
let usage = TokenUsage {
|
let usage = TokenUsage {
|
||||||
@@ -664,15 +711,39 @@ 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,
|
||||||
|
finish_reason: response.finish_reason,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -700,6 +771,7 @@ Respond in JSON format:
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
usage,
|
usage,
|
||||||
|
finish_reason: response.finish_reason,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -725,6 +797,7 @@ Respond in JSON format:
|
|||||||
Ok(RespondOutput {
|
Ok(RespondOutput {
|
||||||
result: RespondResult::Text(final_text),
|
result: RespondResult::Text(final_text),
|
||||||
usage,
|
usage,
|
||||||
|
finish_reason: response.finish_reason,
|
||||||
})
|
})
|
||||||
} else {
|
} else {
|
||||||
// No tools, use simple completion
|
// No tools, use simple completion
|
||||||
@@ -732,6 +805,9 @@ Respond in JSON format:
|
|||||||
.with_max_tokens(4096)
|
.with_max_tokens(4096)
|
||||||
.with_temperature(0.7);
|
.with_temperature(0.7);
|
||||||
request.metadata = context.metadata.clone();
|
request.metadata = context.metadata.clone();
|
||||||
|
if let Some(ref model) = context.model_override {
|
||||||
|
request.model = Some(model.clone());
|
||||||
|
}
|
||||||
|
|
||||||
let response = self.llm.complete(request).await?;
|
let response = self.llm.complete(request).await?;
|
||||||
let pre_truncated = truncate_at_tool_tags(&response.content);
|
let pre_truncated = truncate_at_tool_tags(&response.content);
|
||||||
@@ -753,6 +829,7 @@ Respond in JSON format:
|
|||||||
cache_read_input_tokens: response.cache_read_input_tokens,
|
cache_read_input_tokens: response.cache_read_input_tokens,
|
||||||
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
||||||
},
|
},
|
||||||
|
finish_reason: response.finish_reason,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1293,6 +1370,49 @@ fn is_inside_code(pos: usize, regions: &[CodeRegion]) -> bool {
|
|||||||
regions.iter().any(|r| pos >= r.start && pos < r.end)
|
regions.iter().any(|r| pos >= r.start && pos < r.end)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Check whether a byte range overlaps any code region.
|
||||||
|
fn overlaps_code_region(start: usize, end: usize, regions: &[CodeRegion]) -> bool {
|
||||||
|
regions.iter().any(|r| start < r.end && end > r.start)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return the byte bounds of the line containing `pos`, excluding the trailing newline.
|
||||||
|
fn line_bounds(text: &str, pos: usize) -> (usize, usize) {
|
||||||
|
let start = text[..pos].rfind('\n').map_or(0, |idx| idx + 1);
|
||||||
|
let end = text[pos..].find('\n').map_or(text.len(), |idx| pos + idx);
|
||||||
|
(start, end)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Only recover XML-style tool calls when they are isolated content outside
|
||||||
|
/// markdown code and quote contexts. This avoids converting code examples or
|
||||||
|
/// quoted snippets into executable tool calls.
|
||||||
|
fn is_recoverable_tool_call_segment(
|
||||||
|
text: &str,
|
||||||
|
start: usize,
|
||||||
|
end: usize,
|
||||||
|
code_regions: &[CodeRegion],
|
||||||
|
) -> bool {
|
||||||
|
if overlaps_code_region(start, end, code_regions) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
let (first_line_start, first_line_end) = line_bounds(text, start);
|
||||||
|
let first_line = &text[first_line_start..first_line_end];
|
||||||
|
|
||||||
|
if first_line.trim_start().starts_with('>') {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
let (_, last_line_end) = line_bounds(text, end.saturating_sub(1));
|
||||||
|
let first_line_prefix = &text[first_line_start..start];
|
||||||
|
let last_line_suffix = &text[end..last_line_end];
|
||||||
|
|
||||||
|
if !first_line_prefix.trim().is_empty() || !last_line_suffix.trim().is_empty() {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
|
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
|
||||||
///
|
///
|
||||||
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
|
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
|
||||||
@@ -1312,6 +1432,7 @@ fn recover_tool_calls_from_content(
|
|||||||
) -> Vec<ToolCall> {
|
) -> Vec<ToolCall> {
|
||||||
let tool_names: std::collections::HashSet<&str> =
|
let tool_names: std::collections::HashSet<&str> =
|
||||||
available_tools.iter().map(|t| t.name.as_str()).collect();
|
available_tools.iter().map(|t| t.name.as_str()).collect();
|
||||||
|
let code_regions = find_code_regions(content);
|
||||||
let mut calls = Vec::new();
|
let mut calls = Vec::new();
|
||||||
|
|
||||||
for (open, close) in &[
|
for (open, close) in &[
|
||||||
@@ -1320,15 +1441,23 @@ fn recover_tool_calls_from_content(
|
|||||||
("<function_call>", "</function_call>"),
|
("<function_call>", "</function_call>"),
|
||||||
("<|function_call|>", "<|/function_call|>"),
|
("<|function_call|>", "<|/function_call|>"),
|
||||||
] {
|
] {
|
||||||
let mut remaining = content;
|
let mut search_from = 0;
|
||||||
while let Some(start) = remaining.find(open) {
|
while let Some(offset) = content[search_from..].find(open) {
|
||||||
|
let start = search_from + offset;
|
||||||
let inner_start = start + open.len();
|
let inner_start = start + open.len();
|
||||||
let after = &remaining[inner_start..];
|
let after = &content[inner_start..];
|
||||||
let Some(end) = after.find(close) else {
|
let Some(end_offset) = after.find(close) else {
|
||||||
break;
|
break;
|
||||||
};
|
};
|
||||||
let inner = after[..end].trim();
|
let end = inner_start + end_offset;
|
||||||
remaining = &after[end + close.len()..];
|
let segment_end = end + close.len();
|
||||||
|
search_from = segment_end;
|
||||||
|
|
||||||
|
if !is_recoverable_tool_call_segment(content, start, segment_end, &code_regions) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let inner = content[inner_start..end].trim();
|
||||||
|
|
||||||
if inner.is_empty() {
|
if inner.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
@@ -1350,6 +1479,7 @@ fn recover_tool_calls_from_content(
|
|||||||
),
|
),
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
arguments,
|
arguments,
|
||||||
|
reasoning: None,
|
||||||
});
|
});
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -1364,6 +1494,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 +1532,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 +1544,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;
|
||||||
}
|
}
|
||||||
@@ -2257,6 +2390,40 @@ That's my plan."#;
|
|||||||
assert_eq!(calls[0].name, "tool_list");
|
assert_eq!(calls[0].name, "tool_list");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_recover_tool_call_in_fenced_code_block_ignored() {
|
||||||
|
let tools = make_tools(&["tool_list"]);
|
||||||
|
let content = "Here is the XML format:\n\n```xml\n<tool_call>tool_list</tool_call>\n```";
|
||||||
|
let calls = recover_tool_calls_from_content(content, &tools);
|
||||||
|
assert!(calls.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_recover_tool_call_in_inline_code_ignored() {
|
||||||
|
let tools = make_tools(&["tool_list"]);
|
||||||
|
let content = "Use `<tool_call>tool_list</tool_call>` to illustrate the syntax.";
|
||||||
|
let calls = recover_tool_calls_from_content(content, &tools);
|
||||||
|
assert!(calls.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_recover_tool_call_in_blockquote_ignored() {
|
||||||
|
let tools = make_tools(&["tool_list"]);
|
||||||
|
let content = "The page replied:\n> <tool_call>tool_list</tool_call>";
|
||||||
|
let calls = recover_tool_calls_from_content(content, &tools);
|
||||||
|
assert!(calls.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_recover_multiline_json_tool_call_on_own_line() {
|
||||||
|
let tools = make_tools(&["memory_search"]);
|
||||||
|
let content = "Let me check.\n\n<tool_call>\n{\"name\": \"memory_search\", \"arguments\": {\"query\": \"test\"}}\n</tool_call>\n\nDone.";
|
||||||
|
let calls = recover_tool_calls_from_content(content, &tools);
|
||||||
|
assert_eq!(calls.len(), 1);
|
||||||
|
assert_eq!(calls[0].name, "memory_search");
|
||||||
|
assert_eq!(calls[0].arguments, serde_json::json!({"query": "test"}));
|
||||||
|
}
|
||||||
|
|
||||||
// ---- System prompt building tests (issue #565) ----
|
// ---- System prompt building tests (issue #565) ----
|
||||||
|
|
||||||
fn make_test_reasoning() -> Reasoning {
|
fn make_test_reasoning() -> Reasoning {
|
||||||
@@ -3145,4 +3312,113 @@ 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());
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- select_tools truncation guard ----
|
||||||
|
|
||||||
|
/// Mock provider that returns tool calls with a configurable finish_reason.
|
||||||
|
struct TruncatingLlm {
|
||||||
|
finish_reason: crate::llm::FinishReason,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl crate::llm::LlmProvider for TruncatingLlm {
|
||||||
|
fn model_name(&self) -> &str {
|
||||||
|
"truncating-stub"
|
||||||
|
}
|
||||||
|
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
|
||||||
|
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
|
||||||
|
}
|
||||||
|
async fn complete(
|
||||||
|
&self,
|
||||||
|
_request: crate::llm::CompletionRequest,
|
||||||
|
) -> Result<crate::llm::CompletionResponse, crate::llm::error::LlmError> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn complete_with_tools(
|
||||||
|
&self,
|
||||||
|
_request: crate::llm::ToolCompletionRequest,
|
||||||
|
) -> Result<crate::llm::ToolCompletionResponse, crate::llm::error::LlmError> {
|
||||||
|
Ok(crate::llm::ToolCompletionResponse {
|
||||||
|
content: Some("I'll write the report.".to_string()),
|
||||||
|
tool_calls: vec![ToolCall {
|
||||||
|
id: "call_1".to_string(),
|
||||||
|
name: "memory_write".to_string(),
|
||||||
|
arguments: serde_json::json!({}),
|
||||||
|
reasoning: None,
|
||||||
|
}],
|
||||||
|
input_tokens: 5000,
|
||||||
|
output_tokens: 1024,
|
||||||
|
finish_reason: self.finish_reason,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_select_tools_returns_empty_on_truncation() {
|
||||||
|
let llm = Arc::new(TruncatingLlm {
|
||||||
|
finish_reason: FinishReason::Length,
|
||||||
|
});
|
||||||
|
let reasoning = Reasoning::new(llm);
|
||||||
|
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
|
||||||
|
ctx.available_tools.push(ToolDefinition {
|
||||||
|
name: "memory_write".to_string(),
|
||||||
|
description: "Write to memory".to_string(),
|
||||||
|
parameters: serde_json::json!({"type": "object"}),
|
||||||
|
});
|
||||||
|
|
||||||
|
let selections = reasoning.select_tools(&ctx).await.unwrap();
|
||||||
|
assert!(
|
||||||
|
selections.is_empty(),
|
||||||
|
"Truncated tool selections should be discarded (got {} selections)",
|
||||||
|
selections.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_select_tools_returns_selections_when_not_truncated() {
|
||||||
|
let llm = Arc::new(TruncatingLlm {
|
||||||
|
finish_reason: FinishReason::ToolUse,
|
||||||
|
});
|
||||||
|
let reasoning = Reasoning::new(llm);
|
||||||
|
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
|
||||||
|
ctx.available_tools.push(ToolDefinition {
|
||||||
|
name: "memory_write".to_string(),
|
||||||
|
description: "Write to memory".to_string(),
|
||||||
|
parameters: serde_json::json!({"type": "object"}),
|
||||||
|
});
|
||||||
|
|
||||||
|
let selections = reasoning.select_tools(&ctx).await.unwrap();
|
||||||
|
assert_eq!(selections.len(), 1);
|
||||||
|
assert_eq!(selections[0].tool_name, "memory_write");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+39
-20
@@ -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
|
||||||
@@ -597,6 +598,30 @@ fn build_rig_request(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Inject a per-request model override into the rig request's `additional_params`.
|
||||||
|
///
|
||||||
|
/// Rig-core bakes the model name at construction time inside each provider's
|
||||||
|
/// `CompletionModel` implementation. The actual HTTP request body includes a
|
||||||
|
/// `model` field set by the provider. Rig-core's `#[serde(flatten)]` on
|
||||||
|
/// `additional_params` emits these fields AFTER the provider's own fields.
|
||||||
|
/// Most API servers (Python, Go) use last-key-wins when deserializing
|
||||||
|
/// duplicate JSON keys, so the injected `model` value takes effect.
|
||||||
|
fn inject_model_override(rig_req: &mut RigRequest, model_override: Option<&str>) {
|
||||||
|
let Some(model) = model_override else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
match rig_req.additional_params {
|
||||||
|
Some(ref mut params) => {
|
||||||
|
if let Some(obj) = params.as_object_mut() {
|
||||||
|
obj.insert("model".to_string(), serde_json::json!(model));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
rig_req.additional_params = Some(serde_json::json!({ "model": model }));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl<M> LlmProvider for RigAdapter<M>
|
impl<M> LlmProvider for RigAdapter<M>
|
||||||
where
|
where
|
||||||
@@ -631,15 +656,7 @@ where
|
|||||||
&self,
|
&self,
|
||||||
mut request: CompletionRequest,
|
mut request: CompletionRequest,
|
||||||
) -> Result<CompletionResponse, LlmError> {
|
) -> Result<CompletionResponse, LlmError> {
|
||||||
if let Some(requested_model) = request.model.as_deref()
|
let model_override = request.model.take();
|
||||||
&& requested_model != self.model_name.as_str()
|
|
||||||
{
|
|
||||||
tracing::warn!(
|
|
||||||
requested_model = requested_model,
|
|
||||||
active_model = %self.model_name,
|
|
||||||
"Per-request model override is not supported for this provider; using configured model"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
self.strip_unsupported_completion_params(&mut request);
|
self.strip_unsupported_completion_params(&mut request);
|
||||||
|
|
||||||
@@ -647,7 +664,7 @@ where
|
|||||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||||
let (preamble, history) = convert_messages(&messages);
|
let (preamble, history) = convert_messages(&messages);
|
||||||
|
|
||||||
let rig_req = build_rig_request(
|
let mut rig_req = build_rig_request(
|
||||||
preamble,
|
preamble,
|
||||||
history,
|
history,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
@@ -657,6 +674,8 @@ where
|
|||||||
self.cache_retention,
|
self.cache_retention,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
|
inject_model_override(&mut rig_req, model_override.as_deref());
|
||||||
|
|
||||||
let response =
|
let response =
|
||||||
self.model
|
self.model
|
||||||
.completion(rig_req)
|
.completion(rig_req)
|
||||||
@@ -694,15 +713,7 @@ where
|
|||||||
&self,
|
&self,
|
||||||
mut request: ToolCompletionRequest,
|
mut request: ToolCompletionRequest,
|
||||||
) -> Result<ToolCompletionResponse, LlmError> {
|
) -> Result<ToolCompletionResponse, LlmError> {
|
||||||
if let Some(requested_model) = request.model.as_deref()
|
let model_override = request.model.take();
|
||||||
&& requested_model != self.model_name.as_str()
|
|
||||||
{
|
|
||||||
tracing::warn!(
|
|
||||||
requested_model = requested_model,
|
|
||||||
active_model = %self.model_name,
|
|
||||||
"Per-request model override is not supported for this provider; using configured model"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
self.strip_unsupported_tool_params(&mut request);
|
self.strip_unsupported_tool_params(&mut request);
|
||||||
|
|
||||||
@@ -715,7 +726,7 @@ where
|
|||||||
let tools = convert_tools(&request.tools);
|
let tools = convert_tools(&request.tools);
|
||||||
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
||||||
|
|
||||||
let rig_req = build_rig_request(
|
let mut rig_req = build_rig_request(
|
||||||
preamble,
|
preamble,
|
||||||
history,
|
history,
|
||||||
tools,
|
tools,
|
||||||
@@ -725,6 +736,8 @@ where
|
|||||||
self.cache_retention,
|
self.cache_retention,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
|
inject_model_override(&mut rig_req, model_override.as_deref());
|
||||||
|
|
||||||
let response =
|
let response =
|
||||||
self.model
|
self.model
|
||||||
.completion(rig_req)
|
.completion(rig_req)
|
||||||
@@ -880,6 +893,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 +1011,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 +1043,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 +1077,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 +1397,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");
|
||||||
|
|||||||
+17
-18
@@ -590,8 +590,6 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
|
|
||||||
let mut gateway_url: Option<String> = None;
|
let mut gateway_url: Option<String> = None;
|
||||||
let mut sse_manager: Option<std::sync::Arc<ironclaw::channels::web::sse::SseManager>> = None;
|
let mut sse_manager: Option<std::sync::Arc<ironclaw::channels::web::sse::SseManager>> = None;
|
||||||
let mut _gateway_state: Option<std::sync::Arc<ironclaw::channels::web::server::GatewayState>> =
|
|
||||||
None;
|
|
||||||
if let Some(ref gw_config) = config.channels.gateway {
|
if let Some(ref gw_config) = config.channels.gateway {
|
||||||
// Build multi-user auth state if user_tokens is configured, else single-user.
|
// Build multi-user auth state if user_tokens is configured, else single-user.
|
||||||
let mut gw = if let Some(ref user_tokens) = gw_config.user_tokens {
|
let mut gw = if let Some(ref user_tokens) = gw_config.user_tokens {
|
||||||
@@ -613,6 +611,7 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
} else {
|
} else {
|
||||||
GatewayChannel::new(gw_config.clone())
|
GatewayChannel::new(gw_config.clone())
|
||||||
};
|
};
|
||||||
|
gw = gw.with_owner_scope(config.owner_id.clone());
|
||||||
gw = gw.with_llm_provider(Arc::clone(&components.llm));
|
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));
|
||||||
@@ -729,7 +728,6 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
// 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_manager = Some(Arc::clone(&gw.state().sse));
|
sse_manager = Some(Arc::clone(&gw.state().sse));
|
||||||
_gateway_state = Some(Arc::clone(gw.state()));
|
|
||||||
channel_names.push("gateway".to_string());
|
channel_names.push("gateway".to_string());
|
||||||
channels.add(Box::new(gw)).await;
|
channels.add(Box::new(gw)).await;
|
||||||
}
|
}
|
||||||
@@ -792,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()
|
||||||
@@ -812,12 +818,6 @@ 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 ext_user_id = config
|
|
||||||
.channels
|
|
||||||
.gateway
|
|
||||||
.as_ref()
|
|
||||||
.map(|g| g.user_id.clone())
|
|
||||||
.unwrap_or_else(|| "default".to_string());
|
|
||||||
let persisted = ext_mgr.load_persisted_active_channels(&ext_user_id).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)
|
if active_at_startup.contains(name)
|
||||||
@@ -850,20 +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;
|
||||||
let ext_user_id = config
|
|
||||||
.channels
|
|
||||||
.gateway
|
|
||||||
.as_ref()
|
|
||||||
.map(|g| g.user_id.clone())
|
|
||||||
.unwrap_or_else(|| "default".to_string());
|
|
||||||
ext_mgr.restore_relay_channels(&ext_user_id).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(sse) = sse_manager
|
&& let Some(ref sse) = sse_manager
|
||||||
{
|
{
|
||||||
ext_mgr.set_sse_sender(sse).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
|
||||||
@@ -901,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: None, // TODO: wire SseManager into scheduler (needs Sender<SseEvent> → Arc<SseManager> refactor)
|
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(
|
||||||
@@ -919,6 +913,11 @@ 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(),
|
||||||
|
tenant_rates: Arc::new(ironclaw::tenant::TenantRateRegistry::new(
|
||||||
|
config.agent.max_llm_concurrent_per_user.unwrap_or(4),
|
||||||
|
config.agent.max_jobs_concurrent_per_user.unwrap_or(3),
|
||||||
|
)),
|
||||||
};
|
};
|
||||||
|
|
||||||
let channels_for_warnings = Arc::clone(&channels);
|
let channels_for_warnings = Arc::clone(&channels);
|
||||||
|
|||||||
+69
-25
@@ -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)]
|
||||||
@@ -41,7 +42,7 @@ pub struct OrchestratorState {
|
|||||||
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).
|
||||||
/// Tuple: (job_id, user_id, event).
|
/// Tuple: (job_id, user_id, event).
|
||||||
pub job_event_tx: Option<broadcast::Sender<(Uuid, String, SseEvent)>>,
|
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.
|
||||||
@@ -50,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.
|
||||||
@@ -274,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
|
||||||
@@ -292,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
|
||||||
@@ -306,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
|
||||||
@@ -321,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
|
||||||
@@ -341,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
|
||||||
@@ -353,22 +371,43 @@ async fn job_event_handler(
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Broadcast via the channel (if configured).
|
// Broadcast via the channel (if configured).
|
||||||
// Look up the job owner so the gateway can scope delivery per-user.
|
// 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 user_id = match state.store.as_ref() {
|
let cached_uid = state
|
||||||
Some(store) => store
|
.job_owner_cache
|
||||||
.get_sandbox_job(job_id)
|
.read()
|
||||||
.await
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
.ok()
|
.get(&job_id)
|
||||||
.flatten()
|
.cloned();
|
||||||
.map(|j| j.user_id),
|
|
||||||
None => None,
|
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 let Some(uid) = user_id {
|
|
||||||
let _ = tx.send((job_id, uid, sse_event));
|
if user_id.is_empty() {
|
||||||
|
let _ = tx.send((job_id, String::new(), app_event));
|
||||||
} else {
|
} else {
|
||||||
// Fallback: broadcast globally (single-user mode or job not found).
|
let _ = tx.send((job_id, user_id, app_event));
|
||||||
let _ = tx.send((job_id, String::new(), sse_event));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -496,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())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -725,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);
|
||||||
@@ -760,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();
|
||||||
@@ -790,7 +832,7 @@ mod tests {
|
|||||||
// No store configured, so user_id falls back to empty string.
|
// No store configured, so user_id falls back to empty string.
|
||||||
assert_eq!(recv_uid, "");
|
assert_eq!(recv_uid, "");
|
||||||
match event {
|
match event {
|
||||||
SseEvent::JobMessage {
|
AppEvent::JobMessage {
|
||||||
job_id: jid,
|
job_id: jid,
|
||||||
role,
|
role,
|
||||||
content,
|
content,
|
||||||
@@ -817,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();
|
||||||
@@ -844,7 +887,7 @@ mod tests {
|
|||||||
|
|
||||||
let (_recv_id, _recv_uid, 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),
|
||||||
@@ -865,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();
|
||||||
@@ -889,7 +933,7 @@ mod tests {
|
|||||||
|
|
||||||
let (_recv_id, _recv_uid, 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, String, 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"));
|
||||||
|
|||||||
+906
@@ -0,0 +1,906 @@
|
|||||||
|
//! Compile-time tenant isolation.
|
||||||
|
//!
|
||||||
|
//! Provides two database access tiers:
|
||||||
|
//!
|
||||||
|
//! - **[`TenantScope`]** (default): All operations are bound to a single user.
|
||||||
|
//! ID-based lookups return `None` if the resource doesn't belong to this user.
|
||||||
|
//! This is the only way handler code should access the database.
|
||||||
|
//!
|
||||||
|
//! - **[`AdminScope`]**: Cross-tenant access for system-level operations
|
||||||
|
//! (heartbeat, routine engine, self-repair). Must be obtained explicitly via
|
||||||
|
//! [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
|
||||||
|
//!
|
||||||
|
//! [`TenantCtx`] bundles a `TenantScope` with workspace, cost guard, and
|
||||||
|
//! per-tenant rate limiting. Constructed once per request at the entry point
|
||||||
|
//! where a `user_id` becomes known.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use rust_decimal::Decimal;
|
||||||
|
use tokio::sync::{Semaphore, SemaphorePermit};
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::agent::BrokenTool;
|
||||||
|
use crate::agent::cost_guard::{CostGuard, CostLimitExceeded};
|
||||||
|
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||||
|
use crate::context::{ActionRecord, JobContext, JobState};
|
||||||
|
use crate::db::Database;
|
||||||
|
use crate::error::DatabaseError;
|
||||||
|
use crate::history::{
|
||||||
|
AgentJobRecord, AgentJobSummary, ConversationMessage, ConversationSummary, LlmCallRecord,
|
||||||
|
SandboxJobRecord, SandboxJobSummary, SettingRow,
|
||||||
|
};
|
||||||
|
use crate::workspace::Workspace;
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// TenantScope — scoped database access (default tier)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Scoped database view. All operations are bound to a single user.
|
||||||
|
///
|
||||||
|
/// This is the **only** way handler code should access the database.
|
||||||
|
/// ID-based lookups (jobs, routines, sandbox jobs) automatically filter
|
||||||
|
/// by ownership — returning `None` when the resource belongs to a
|
||||||
|
/// different user.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct TenantScope {
|
||||||
|
user_id: String,
|
||||||
|
inner: Arc<dyn Database>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TenantScope {
|
||||||
|
pub fn new(user_id: impl Into<String>, db: Arc<dyn Database>) -> Self {
|
||||||
|
Self {
|
||||||
|
user_id: user_id.into(),
|
||||||
|
inner: db,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn user_id(&self) -> &str {
|
||||||
|
&self.user_id
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Jobs ===
|
||||||
|
|
||||||
|
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||||
|
self.inner.list_agent_jobs_for_user(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError> {
|
||||||
|
self.inner.agent_job_summary_for_user(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch a job by ID, returning `None` if it doesn't belong to this user.
|
||||||
|
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||||
|
match self.inner.get_job(id).await? {
|
||||||
|
Some(ctx) if ctx.user_id == self.user_id => Ok(Some(ctx)),
|
||||||
|
_ => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_agent_job_failure_reason(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
) -> Result<Option<String>, DatabaseError> {
|
||||||
|
// Verify ownership first
|
||||||
|
if self.get_job(id).await?.is_none() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
self.inner.get_agent_job_failure_reason(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_job_status(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
status: JobState,
|
||||||
|
failure_reason: Option<&str>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
// Verify ownership before mutating
|
||||||
|
if self.get_job(id).await?.is_none() {
|
||||||
|
return Err(DatabaseError::NotFound {
|
||||||
|
entity: "job".to_string(),
|
||||||
|
id: id.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
self.inner
|
||||||
|
.update_job_status(id, status, failure_reason)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Sandbox jobs ===
|
||||||
|
|
||||||
|
pub async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||||
|
self.inner.list_sandbox_jobs_for_user(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
|
||||||
|
self.inner.sandbox_job_summary_for_user(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch a sandbox job by ID, returning `None` if it doesn't belong to this user.
|
||||||
|
pub async fn get_sandbox_job(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||||
|
match self.inner.get_sandbox_job(id).await? {
|
||||||
|
Some(job) if job.user_id == self.user_id => Ok(Some(job)),
|
||||||
|
_ => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn sandbox_job_belongs_to_user(&self, job_id: Uuid) -> Result<bool, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.sandbox_job_belongs_to_user(job_id, &self.user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Routines ===
|
||||||
|
|
||||||
|
pub async fn list_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||||
|
self.inner.list_routines(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_routine_by_name(&self, name: &str) -> Result<Option<Routine>, DatabaseError> {
|
||||||
|
self.inner.get_routine_by_name(&self.user_id, name).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch a routine by ID, returning `None` if it doesn't belong to this user.
|
||||||
|
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||||
|
match self.inner.get_routine(id).await? {
|
||||||
|
Some(r) if r.user_id == self.user_id => Ok(Some(r)),
|
||||||
|
_ => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||||
|
debug_assert_eq!(
|
||||||
|
routine.user_id, self.user_id,
|
||||||
|
"routine.user_id must match TenantScope user"
|
||||||
|
);
|
||||||
|
self.inner.create_routine(routine).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||||
|
// Verify ownership
|
||||||
|
if self.get_routine(routine.id).await?.is_none() {
|
||||||
|
return Err(DatabaseError::NotFound {
|
||||||
|
entity: "routine".to_string(),
|
||||||
|
id: routine.id.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
self.inner.update_routine(routine).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
|
||||||
|
// Verify ownership
|
||||||
|
if self.get_routine(id).await?.is_none() {
|
||||||
|
return Err(DatabaseError::NotFound {
|
||||||
|
entity: "routine".to_string(),
|
||||||
|
id: id.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
self.inner.delete_routine(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List routine runs, verifying the routine belongs to this user.
|
||||||
|
pub async fn list_routine_runs(
|
||||||
|
&self,
|
||||||
|
routine_id: Uuid,
|
||||||
|
limit: i64,
|
||||||
|
) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||||
|
// Verify routine ownership first
|
||||||
|
if self.get_routine(routine_id).await?.is_none() {
|
||||||
|
return Err(DatabaseError::NotFound {
|
||||||
|
entity: "routine".to_string(),
|
||||||
|
id: routine_id.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
self.inner.list_routine_runs(routine_id, limit).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_webhook_routine_by_path(
|
||||||
|
&self,
|
||||||
|
path: &str,
|
||||||
|
) -> Result<Option<Routine>, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_webhook_routine_by_path(path, Some(&self.user_id))
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Settings ===
|
||||||
|
|
||||||
|
pub async fn get_setting(&self, key: &str) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||||
|
self.inner.get_setting(&self.user_id, key).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_setting_full(&self, key: &str) -> Result<Option<SettingRow>, DatabaseError> {
|
||||||
|
self.inner.get_setting_full(&self.user_id, key).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn set_setting(
|
||||||
|
&self,
|
||||||
|
key: &str,
|
||||||
|
value: &serde_json::Value,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.set_setting(&self.user_id, key, value).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_setting(&self, key: &str) -> Result<bool, DatabaseError> {
|
||||||
|
self.inner.delete_setting(&self.user_id, key).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_settings(&self) -> Result<Vec<SettingRow>, DatabaseError> {
|
||||||
|
self.inner.list_settings(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_all_settings(
|
||||||
|
&self,
|
||||||
|
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
|
||||||
|
self.inner.get_all_settings(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn set_all_settings(
|
||||||
|
&self,
|
||||||
|
settings: &HashMap<String, serde_json::Value>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.set_all_settings(&self.user_id, settings).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn has_settings(&self) -> Result<bool, DatabaseError> {
|
||||||
|
self.inner.has_settings(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Conversations ===
|
||||||
|
|
||||||
|
pub async fn create_conversation(
|
||||||
|
&self,
|
||||||
|
channel: &str,
|
||||||
|
thread_id: Option<&str>,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.create_conversation(channel, &self.user_id, thread_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn ensure_conversation(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
channel: &str,
|
||||||
|
thread_id: Option<&str>,
|
||||||
|
) -> Result<bool, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.ensure_conversation(id, channel, &self.user_id, thread_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_conversations_with_preview(
|
||||||
|
&self,
|
||||||
|
channel: &str,
|
||||||
|
limit: i64,
|
||||||
|
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.list_conversations_with_preview(&self.user_id, channel, limit)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_conversations_all_channels(
|
||||||
|
&self,
|
||||||
|
limit: i64,
|
||||||
|
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.list_conversations_all_channels(&self.user_id, limit)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_or_create_routine_conversation(
|
||||||
|
&self,
|
||||||
|
routine_id: Uuid,
|
||||||
|
routine_name: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_or_create_routine_conversation(routine_id, routine_name, &self.user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_or_create_heartbeat_conversation(&self) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_or_create_heartbeat_conversation(&self.user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_or_create_assistant_conversation(
|
||||||
|
&self,
|
||||||
|
channel: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_or_create_assistant_conversation(&self.user_id, channel)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn conversation_belongs_to_user(
|
||||||
|
&self,
|
||||||
|
conversation_id: Uuid,
|
||||||
|
) -> Result<bool, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.conversation_belongs_to_user(conversation_id, &self.user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a message to a conversation owned by this tenant.
|
||||||
|
///
|
||||||
|
/// Verifies the conversation belongs to this user before adding.
|
||||||
|
pub async fn add_conversation_message(
|
||||||
|
&self,
|
||||||
|
conversation_id: Uuid,
|
||||||
|
role: &str,
|
||||||
|
content: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.add_conversation_message(conversation_id, role, content)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.touch_conversation(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_conversation_messages(
|
||||||
|
&self,
|
||||||
|
conversation_id: Uuid,
|
||||||
|
) -> Result<Vec<ConversationMessage>, DatabaseError> {
|
||||||
|
self.inner.list_conversation_messages(conversation_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_conversation_messages_paginated(
|
||||||
|
&self,
|
||||||
|
conversation_id: Uuid,
|
||||||
|
before: Option<DateTime<Utc>>,
|
||||||
|
limit: i64,
|
||||||
|
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.list_conversation_messages_paginated(conversation_id, before, limit)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn create_conversation_with_metadata(
|
||||||
|
&self,
|
||||||
|
channel: &str,
|
||||||
|
metadata: &serde_json::Value,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.create_conversation_with_metadata(channel, &self.user_id, metadata)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_conversation_metadata_field(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
key: &str,
|
||||||
|
value: &serde_json::Value,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.update_conversation_metadata_field(id, key, value)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_conversation_metadata(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||||
|
self.inner.get_conversation_metadata(id).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// AdminScope — explicit cross-tenant access
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Cross-tenant database access for system-level operations.
|
||||||
|
///
|
||||||
|
/// **Not** available through [`TenantCtx`] — must be obtained explicitly via
|
||||||
|
/// [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
|
||||||
|
///
|
||||||
|
/// Used by: heartbeat enumeration, routine engine scheduling, self-repair,
|
||||||
|
/// scheduler job persistence, worker status updates.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AdminScope {
|
||||||
|
inner: Arc<dyn Database>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AdminScope {
|
||||||
|
pub fn new(db: Arc<dyn Database>) -> Self {
|
||||||
|
Self { inner: db }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Access the raw Database trait object.
|
||||||
|
///
|
||||||
|
/// Prefer using the typed methods on AdminScope instead. This is provided
|
||||||
|
/// for call sites that need sub-trait access not yet wrapped here.
|
||||||
|
pub fn db(&self) -> &Arc<dyn Database> {
|
||||||
|
&self.inner
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Routine engine ===
|
||||||
|
|
||||||
|
pub async fn list_all_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||||
|
self.inner.list_all_routines().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||||
|
self.inner.list_event_routines().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||||
|
self.inner.list_due_cron_routines().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||||
|
self.inner.list_dispatched_routine_runs().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn count_running_routine_runs_batch(
|
||||||
|
&self,
|
||||||
|
routine_ids: &[Uuid],
|
||||||
|
) -> Result<HashMap<Uuid, i64>, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.count_running_routine_runs_batch(routine_ids)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn batch_get_last_run_status(
|
||||||
|
&self,
|
||||||
|
routine_ids: &[Uuid],
|
||||||
|
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
|
||||||
|
self.inner.batch_get_last_run_status(routine_ids).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
|
||||||
|
self.inner.count_running_routine_runs(routine_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_routine_runtime(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
last_run_at: DateTime<Utc>,
|
||||||
|
next_fire_at: Option<DateTime<Utc>>,
|
||||||
|
run_count: u64,
|
||||||
|
consecutive_failures: u32,
|
||||||
|
state: &serde_json::Value,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.update_routine_runtime(
|
||||||
|
id,
|
||||||
|
last_run_at,
|
||||||
|
next_fire_at,
|
||||||
|
run_count,
|
||||||
|
consecutive_failures,
|
||||||
|
state,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.create_routine_run(run).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn complete_routine_run(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
status: RunStatus,
|
||||||
|
result_summary: Option<&str>,
|
||||||
|
tokens_used: Option<i32>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.complete_routine_run(id, status, result_summary, tokens_used)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn link_routine_run_to_job(
|
||||||
|
&self,
|
||||||
|
run_id: Uuid,
|
||||||
|
job_id: Uuid,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.link_routine_run_to_job(run_id, job_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||||
|
self.inner.get_routine(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.update_routine(routine).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Self-repair ===
|
||||||
|
|
||||||
|
pub async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
|
||||||
|
self.inner.get_stuck_jobs().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
|
||||||
|
self.inner.get_broken_tools(threshold).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_tool_failure(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
error_message: &str,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.record_tool_failure(tool_name, error_message)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.mark_tool_repaired(tool_name).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.increment_repair_attempts(tool_name).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Sandbox housekeeping ===
|
||||||
|
|
||||||
|
pub async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
|
||||||
|
self.inner.cleanup_stale_sandbox_jobs().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_sandbox_job(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||||
|
self.inner.get_sandbox_job(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.save_sandbox_job(job).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_sandbox_job_status(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
status: &str,
|
||||||
|
success: Option<bool>,
|
||||||
|
message: Option<&str>,
|
||||||
|
started_at: Option<DateTime<Utc>>,
|
||||||
|
completed_at: Option<DateTime<Utc>>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.update_sandbox_job_mode(id, mode).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
|
||||||
|
self.inner.get_sandbox_job_mode(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save_job_event(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
event_type: &str,
|
||||||
|
data: &serde_json::Value,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.save_job_event(job_id, event_type, data).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_job_events(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
limit: Option<i64>,
|
||||||
|
) -> Result<Vec<crate::history::JobEventRecord>, DatabaseError> {
|
||||||
|
self.inner.list_job_events(job_id, limit).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Job persistence (scheduler, worker) ===
|
||||||
|
|
||||||
|
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||||
|
self.inner.get_job(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.save_job(ctx).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_job_status(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
status: JobState,
|
||||||
|
failure_reason: Option<&str>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.update_job_status(id, status, failure_reason)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.mark_job_stuck(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||||
|
self.inner.list_agent_jobs().await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_agent_job_failure_reason(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
) -> Result<Option<String>, DatabaseError> {
|
||||||
|
self.inner.get_agent_job_failure_reason(id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === LLM call recording ===
|
||||||
|
|
||||||
|
pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner.record_llm_call(record).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save_action(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
action: &ActionRecord,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner.save_action(job_id, action).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
|
||||||
|
self.inner.get_job_actions(job_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Estimation ===
|
||||||
|
|
||||||
|
pub async fn save_estimation_snapshot(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
category: &str,
|
||||||
|
tool_names: &[String],
|
||||||
|
estimated_cost: Decimal,
|
||||||
|
estimated_time_secs: i32,
|
||||||
|
estimated_value: Decimal,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.save_estimation_snapshot(
|
||||||
|
job_id,
|
||||||
|
category,
|
||||||
|
tool_names,
|
||||||
|
estimated_cost,
|
||||||
|
estimated_time_secs,
|
||||||
|
estimated_value,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn update_estimation_actuals(
|
||||||
|
&self,
|
||||||
|
id: Uuid,
|
||||||
|
actual_cost: Decimal,
|
||||||
|
actual_time_secs: i32,
|
||||||
|
actual_value: Option<Decimal>,
|
||||||
|
) -> Result<(), DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Conversations (admin context) ===
|
||||||
|
|
||||||
|
pub async fn add_conversation_message(
|
||||||
|
&self,
|
||||||
|
conversation_id: Uuid,
|
||||||
|
role: &str,
|
||||||
|
content: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.add_conversation_message(conversation_id, role, content)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_or_create_routine_conversation(
|
||||||
|
&self,
|
||||||
|
routine_id: Uuid,
|
||||||
|
routine_name: &str,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_or_create_routine_conversation(routine_id, routine_name, user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_or_create_heartbeat_conversation(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
) -> Result<Uuid, DatabaseError> {
|
||||||
|
self.inner
|
||||||
|
.get_or_create_heartbeat_conversation(user_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// TenantRateState / TenantRateRegistry — per-user concurrency
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Per-tenant concurrency limits.
|
||||||
|
pub struct TenantRateState {
|
||||||
|
/// Limits concurrent LLM calls for this user.
|
||||||
|
pub llm_semaphore: Arc<Semaphore>,
|
||||||
|
/// Limits concurrent jobs for this user.
|
||||||
|
pub job_semaphore: Arc<Semaphore>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TenantRateState {
|
||||||
|
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
llm_semaphore: Arc::new(Semaphore::new(max_llm_concurrent)),
|
||||||
|
job_semaphore: Arc::new(Semaphore::new(max_job_concurrent)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Registry that lazily creates per-tenant rate state.
|
||||||
|
///
|
||||||
|
/// Uses `tokio::sync::RwLock<HashMap>` (consistent with the rest of the
|
||||||
|
/// codebase — no DashMap dependency).
|
||||||
|
pub struct TenantRateRegistry {
|
||||||
|
state: tokio::sync::RwLock<HashMap<String, Arc<TenantRateState>>>,
|
||||||
|
max_llm_concurrent: usize,
|
||||||
|
max_job_concurrent: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TenantRateRegistry {
|
||||||
|
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
state: tokio::sync::RwLock::new(HashMap::new()),
|
||||||
|
max_llm_concurrent,
|
||||||
|
max_job_concurrent,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get or lazily create rate state for a user.
|
||||||
|
pub async fn get_or_create(&self, user_id: &str) -> Arc<TenantRateState> {
|
||||||
|
// Fast path: read lock
|
||||||
|
{
|
||||||
|
let map = self.state.read().await;
|
||||||
|
if let Some(s) = map.get(user_id) {
|
||||||
|
return Arc::clone(s);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Slow path: write lock with double-check
|
||||||
|
let mut map = self.state.write().await;
|
||||||
|
if let Some(s) = map.get(user_id) {
|
||||||
|
return Arc::clone(s);
|
||||||
|
}
|
||||||
|
let s = Arc::new(TenantRateState::new(
|
||||||
|
self.max_llm_concurrent,
|
||||||
|
self.max_job_concurrent,
|
||||||
|
));
|
||||||
|
map.insert(user_id.to_string(), Arc::clone(&s));
|
||||||
|
s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// TenantCtx — per-request tenant execution context
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Per-request tenant execution context.
|
||||||
|
///
|
||||||
|
/// Bundles a [`TenantScope`] (scoped DB access), workspace, cost guard,
|
||||||
|
/// and per-tenant rate limiting. Constructed once per request via
|
||||||
|
/// [`AgentDeps::tenant_ctx()`](crate::agent::AgentDeps::tenant_ctx).
|
||||||
|
///
|
||||||
|
/// `Clone + Send + Sync` — safe to store on `ChatDelegate` without lifetime issues.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct TenantCtx {
|
||||||
|
user_id: String,
|
||||||
|
store: Option<TenantScope>,
|
||||||
|
workspace: Option<Arc<Workspace>>,
|
||||||
|
cost_guard: Arc<CostGuard>,
|
||||||
|
rate: Arc<TenantRateState>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TenantCtx {
|
||||||
|
pub fn new(
|
||||||
|
user_id: impl Into<String>,
|
||||||
|
store: Option<TenantScope>,
|
||||||
|
workspace: Option<Arc<Workspace>>,
|
||||||
|
cost_guard: Arc<CostGuard>,
|
||||||
|
rate: Arc<TenantRateState>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
user_id: user_id.into(),
|
||||||
|
store,
|
||||||
|
workspace,
|
||||||
|
cost_guard,
|
||||||
|
rate,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn user_id(&self) -> &str {
|
||||||
|
&self.user_id
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn store(&self) -> Option<&TenantScope> {
|
||||||
|
self.store.as_ref()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn workspace(&self) -> Option<&Arc<Workspace>> {
|
||||||
|
self.workspace.as_ref()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn cost_guard(&self) -> &CostGuard {
|
||||||
|
&self.cost_guard
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check cost limits for this tenant (global + per-user).
|
||||||
|
pub async fn check_cost_allowed(&self) -> Result<(), CostLimitExceeded> {
|
||||||
|
self.cost_guard.check_allowed_for_user(&self.user_id).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record an LLM call for this tenant.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub async fn record_llm_call(
|
||||||
|
&self,
|
||||||
|
model: &str,
|
||||||
|
input_tokens: u32,
|
||||||
|
output_tokens: u32,
|
||||||
|
cache_read_input_tokens: u32,
|
||||||
|
cache_creation_input_tokens: u32,
|
||||||
|
cache_read_discount: Decimal,
|
||||||
|
cache_write_multiplier: Decimal,
|
||||||
|
cost_per_token: Option<(Decimal, Decimal)>,
|
||||||
|
) -> Decimal {
|
||||||
|
self.cost_guard
|
||||||
|
.record_llm_call_for_user(
|
||||||
|
&self.user_id,
|
||||||
|
model,
|
||||||
|
input_tokens,
|
||||||
|
output_tokens,
|
||||||
|
cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens,
|
||||||
|
cache_read_discount,
|
||||||
|
cache_write_multiplier,
|
||||||
|
cost_per_token,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Acquire an LLM concurrency permit for this tenant.
|
||||||
|
pub async fn acquire_llm_permit(&self) -> Result<SemaphorePermit<'_>, crate::error::Error> {
|
||||||
|
self.rate.llm_semaphore.acquire().await.map_err(|_| {
|
||||||
|
crate::error::Error::Config(crate::error::ConfigError::InvalidValue {
|
||||||
|
key: "llm_semaphore".to_string(),
|
||||||
|
message: "semaphore closed".to_string(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_rate_registry_returns_same_state_for_same_user() {
|
||||||
|
let registry = TenantRateRegistry::new(4, 3);
|
||||||
|
let a1 = registry.get_or_create("alice").await;
|
||||||
|
let a2 = registry.get_or_create("alice").await;
|
||||||
|
assert!(Arc::ptr_eq(&a1, &a2));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_rate_registry_different_users_get_different_state() {
|
||||||
|
let registry = TenantRateRegistry::new(4, 3);
|
||||||
|
let alice = registry.get_or_create("alice").await;
|
||||||
|
let bob = registry.get_or_create("bob").await;
|
||||||
|
assert!(!Arc::ptr_eq(&alice, &bob));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -532,6 +532,7 @@ impl TestHarnessBuilder {
|
|||||||
let cost_guard = Arc::new(CostGuard::new(CostGuardConfig {
|
let cost_guard = Arc::new(CostGuard::new(CostGuardConfig {
|
||||||
max_cost_per_day_cents: None,
|
max_cost_per_day_cents: None,
|
||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
|
max_cost_per_user_per_day_cents: None,
|
||||||
}));
|
}));
|
||||||
|
|
||||||
let channel = if self.stub_channel {
|
let channel = if self.stub_channel {
|
||||||
@@ -563,6 +564,8 @@ 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(),
|
||||||
|
tenant_rates: std::sync::Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||||
};
|
};
|
||||||
|
|
||||||
TestHarness {
|
TestHarness {
|
||||||
|
|||||||
@@ -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, String, 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, String, 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);
|
||||||
|
|||||||
@@ -12,12 +12,10 @@
|
|||||||
//! Use `memory_write` to persist important facts that should be remembered
|
//! Use `memory_write` to persist important facts that should be remembered
|
||||||
//! across sessions.
|
//! across sessions.
|
||||||
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use tokio::sync::RwLock;
|
|
||||||
|
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||||
@@ -52,79 +50,6 @@ impl WorkspaceResolver for FixedWorkspaceResolver {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Creates per-user workspaces on demand, caching them for reuse.
|
|
||||||
///
|
|
||||||
/// Used in multi-tenant mode where each authenticated user gets their own
|
|
||||||
/// workspace scope. The workspace is constructed with the same configuration
|
|
||||||
/// (embeddings, search config, memory layers) as the startup workspace.
|
|
||||||
pub struct PerUserWorkspaceResolver {
|
|
||||||
db: Arc<dyn crate::db::Database>,
|
|
||||||
embeddings: Option<Arc<dyn crate::workspace::EmbeddingProvider>>,
|
|
||||||
embedding_cache_config: crate::workspace::EmbeddingCacheConfig,
|
|
||||||
search_config: crate::config::WorkspaceSearchConfig,
|
|
||||||
workspace_config: crate::config::WorkspaceConfig,
|
|
||||||
cache: RwLock<HashMap<String, Arc<Workspace>>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl PerUserWorkspaceResolver {
|
|
||||||
pub fn new(
|
|
||||||
db: Arc<dyn crate::db::Database>,
|
|
||||||
embeddings: Option<Arc<dyn crate::workspace::EmbeddingProvider>>,
|
|
||||||
embedding_cache_config: crate::workspace::EmbeddingCacheConfig,
|
|
||||||
search_config: crate::config::WorkspaceSearchConfig,
|
|
||||||
workspace_config: crate::config::WorkspaceConfig,
|
|
||||||
) -> Self {
|
|
||||||
Self {
|
|
||||||
db,
|
|
||||||
embeddings,
|
|
||||||
embedding_cache_config,
|
|
||||||
search_config,
|
|
||||||
workspace_config,
|
|
||||||
cache: RwLock::new(HashMap::new()),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_workspace(&self, user_id: &str) -> Arc<Workspace> {
|
|
||||||
let mut ws = Workspace::new_with_db(user_id, Arc::clone(&self.db))
|
|
||||||
.with_search_config(&self.search_config);
|
|
||||||
|
|
||||||
if let Some(ref emb) = self.embeddings {
|
|
||||||
ws = ws.with_embeddings_cached(Arc::clone(emb), self.embedding_cache_config.clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
if !self.workspace_config.read_scopes.is_empty() {
|
|
||||||
ws = ws.with_additional_read_scopes(self.workspace_config.read_scopes.clone());
|
|
||||||
}
|
|
||||||
ws = ws.with_memory_layers(self.workspace_config.memory_layers.clone());
|
|
||||||
|
|
||||||
Arc::new(ws)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[async_trait]
|
|
||||||
impl WorkspaceResolver for PerUserWorkspaceResolver {
|
|
||||||
async fn resolve(&self, user_id: &str) -> Arc<Workspace> {
|
|
||||||
// Fast path: read lock
|
|
||||||
{
|
|
||||||
let cache = self.cache.read().await;
|
|
||||||
if let Some(ws) = cache.get(user_id) {
|
|
||||||
return Arc::clone(ws);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Slow path: write lock, double-check
|
|
||||||
let mut cache = self.cache.write().await;
|
|
||||||
if let Some(ws) = cache.get(user_id) {
|
|
||||||
return Arc::clone(ws);
|
|
||||||
}
|
|
||||||
|
|
||||||
let ws = self.build_workspace(user_id);
|
|
||||||
cache.insert(user_id.to_string(), Arc::clone(&ws));
|
|
||||||
tracing::debug!(user_id = user_id, "Created per-user workspace");
|
|
||||||
ws
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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:
|
||||||
@@ -1007,10 +932,10 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_per_user_workspace_resolver_returns_different_workspaces() {
|
async fn test_workspace_pool_resolver_returns_different_workspaces() {
|
||||||
let db = make_test_db().await;
|
let db = make_test_db().await;
|
||||||
|
|
||||||
let resolver = PerUserWorkspaceResolver::new(
|
let pool = crate::channels::web::server::WorkspacePool::new(
|
||||||
db,
|
db,
|
||||||
None,
|
None,
|
||||||
crate::workspace::EmbeddingCacheConfig::default(),
|
crate::workspace::EmbeddingCacheConfig::default(),
|
||||||
@@ -1018,8 +943,8 @@ mod tests {
|
|||||||
crate::config::WorkspaceConfig::default(),
|
crate::config::WorkspaceConfig::default(),
|
||||||
);
|
);
|
||||||
|
|
||||||
let ws_alice = resolver.resolve("alice").await;
|
let ws_alice = pool.resolve("alice").await;
|
||||||
let ws_bob = resolver.resolve("bob").await;
|
let ws_bob = pool.resolve("bob").await;
|
||||||
|
|
||||||
// Different user IDs should get different workspaces
|
// Different user IDs should get different workspaces
|
||||||
assert_eq!(ws_alice.user_id(), "alice");
|
assert_eq!(ws_alice.user_id(), "alice");
|
||||||
@@ -1028,10 +953,10 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_per_user_workspace_resolver_caches_workspace() {
|
async fn test_workspace_pool_resolver_caches_workspace() {
|
||||||
let db = make_test_db().await;
|
let db = make_test_db().await;
|
||||||
|
|
||||||
let resolver = PerUserWorkspaceResolver::new(
|
let pool = crate::channels::web::server::WorkspacePool::new(
|
||||||
db,
|
db,
|
||||||
None,
|
None,
|
||||||
crate::workspace::EmbeddingCacheConfig::default(),
|
crate::workspace::EmbeddingCacheConfig::default(),
|
||||||
@@ -1039,8 +964,8 @@ mod tests {
|
|||||||
crate::config::WorkspaceConfig::default(),
|
crate::config::WorkspaceConfig::default(),
|
||||||
);
|
);
|
||||||
|
|
||||||
let ws1 = resolver.resolve("alice").await;
|
let ws1 = pool.resolve("alice").await;
|
||||||
let ws2 = resolver.resolve("alice").await;
|
let ws2 = pool.resolve("alice").await;
|
||||||
|
|
||||||
// Same user_id should return the same cached Arc (pointer equality)
|
// Same user_id should return the same cached Arc (pointer equality)
|
||||||
assert!(Arc::ptr_eq(&ws1, &ws2));
|
assert!(Arc::ptr_eq(&ws1, &ws2));
|
||||||
|
|||||||
+133
-14
@@ -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 {
|
||||||
@@ -646,6 +650,23 @@ pub(crate) fn routine_update_parameters_schema() -> Value {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const ROUTINE_LAST_NAME_STASH_KEY: &str = "__routine_last_name";
|
||||||
|
|
||||||
|
async fn stash_last_routine_name(ctx: &JobContext, name: &str) {
|
||||||
|
ctx.tool_output_stash
|
||||||
|
.write()
|
||||||
|
.await
|
||||||
|
.insert(ROUTINE_LAST_NAME_STASH_KEY.to_string(), name.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn restore_last_routine_name(ctx: &JobContext) -> Option<String> {
|
||||||
|
ctx.tool_output_stash
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.get(ROUTINE_LAST_NAME_STASH_KEY)
|
||||||
|
.cloned()
|
||||||
|
}
|
||||||
|
|
||||||
fn nested_object<'a>(params: &'a Value, field: &str) -> Option<&'a Map<String, Value>> {
|
fn nested_object<'a>(params: &'a Value, field: &str) -> Option<&'a Map<String, Value>> {
|
||||||
params.get(field).and_then(Value::as_object)
|
params.get(field).and_then(Value::as_object)
|
||||||
}
|
}
|
||||||
@@ -852,11 +873,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 +913,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);
|
||||||
@@ -907,7 +932,7 @@ fn parse_routine_create_request(
|
|||||||
fn build_routine_trigger(trigger: &NormalizedTriggerRequest) -> Trigger {
|
fn build_routine_trigger(trigger: &NormalizedTriggerRequest) -> Trigger {
|
||||||
match trigger {
|
match trigger {
|
||||||
NormalizedTriggerRequest::Cron { schedule, timezone } => Trigger::Cron {
|
NormalizedTriggerRequest::Cron { schedule, timezone } => Trigger::Cron {
|
||||||
schedule: schedule.clone(),
|
schedule: normalize_cron_expression(schedule),
|
||||||
timezone: timezone.clone(),
|
timezone: timezone.clone(),
|
||||||
},
|
},
|
||||||
NormalizedTriggerRequest::Manual => Trigger::Manual,
|
NormalizedTriggerRequest::Manual => Trigger::Manual,
|
||||||
@@ -1007,7 +1032,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 {
|
||||||
@@ -1084,6 +1110,7 @@ impl Tool for RoutineCreateTool {
|
|||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
let normalized = parse_routine_create_request(¶ms)?;
|
let normalized = parse_routine_create_request(¶ms)?;
|
||||||
|
stash_last_routine_name(ctx, &normalized.name).await;
|
||||||
let trigger = build_routine_trigger(&normalized.trigger);
|
let trigger = build_routine_trigger(&normalized.trigger);
|
||||||
let action =
|
let action =
|
||||||
build_routine_action(&normalized.name, &normalized.prompt, &normalized.execution);
|
build_routine_action(&normalized.name, &normalized.prompt, &normalized.execution);
|
||||||
@@ -1265,6 +1292,7 @@ impl Tool for RoutineUpdateTool {
|
|||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
let name = require_str(¶ms, "name")?;
|
let name = require_str(¶ms, "name")?;
|
||||||
|
stash_last_routine_name(ctx, name).await;
|
||||||
|
|
||||||
let mut routine = self
|
let mut routine = self
|
||||||
.store
|
.store
|
||||||
@@ -1402,11 +1430,24 @@ impl Tool for RoutineDeleteTool {
|
|||||||
) -> 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 = if let Some(name) = params.get("name").and_then(|v| v.as_str()) {
|
||||||
|
if name.trim().is_empty() {
|
||||||
|
return Err(ToolError::InvalidParameters(
|
||||||
|
"'name' parameter cannot be empty".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
name.to_string()
|
||||||
|
} else {
|
||||||
|
restore_last_routine_name(ctx).await.ok_or_else(|| {
|
||||||
|
ToolError::InvalidParameters(
|
||||||
|
"missing 'name' parameter and no previous routine target to infer".to_string(),
|
||||||
|
)
|
||||||
|
})?
|
||||||
|
};
|
||||||
|
|
||||||
let routine = self
|
let routine = self
|
||||||
.store
|
.store
|
||||||
.get_routine_by_name(&ctx.user_id, name)
|
.get_routine_by_name(&ctx.user_id, &name)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
|
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
|
||||||
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
|
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
|
||||||
@@ -1421,7 +1462,7 @@ impl Tool for RoutineDeleteTool {
|
|||||||
self.engine.refresh_event_cache().await;
|
self.engine.refresh_event_cache().await;
|
||||||
|
|
||||||
let result = serde_json::json!({
|
let result = serde_json::json!({
|
||||||
"name": name,
|
"name": &name,
|
||||||
"deleted": deleted,
|
"deleted": deleted,
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -1827,6 +1868,20 @@ mod tests {
|
|||||||
assert_eq!(parsed.cooldown_secs, 30);
|
assert_eq!(parsed.cooldown_secs, 30);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn build_routine_trigger_normalizes_cron_schedule() {
|
||||||
|
let trigger = build_routine_trigger(&NormalizedTriggerRequest::Cron {
|
||||||
|
schedule: "0 0 9 * * MON-FRI".to_string(),
|
||||||
|
timezone: Some("UTC".to_string()),
|
||||||
|
});
|
||||||
|
|
||||||
|
assert!(matches!(
|
||||||
|
trigger,
|
||||||
|
Trigger::Cron { schedule, timezone }
|
||||||
|
if schedule == "0 0 9 * * MON-FRI *" && timezone.as_deref() == Some("UTC")
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parses_grouped_message_event_with_tools() {
|
fn parses_grouped_message_event_with_tools() {
|
||||||
let params = serde_json::json!({
|
let params = serde_json::json!({
|
||||||
@@ -1863,6 +1918,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 +2306,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
|
||||||
|
|||||||
+19
-1
@@ -117,6 +117,11 @@ impl McpClient {
|
|||||||
/// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`.
|
/// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`.
|
||||||
///
|
///
|
||||||
/// Returns an error if the config uses a non-HTTP transport.
|
/// Returns an error if the config uses a non-HTTP transport.
|
||||||
|
///
|
||||||
|
/// **Note:** The session manager is NOT wired into the transport. For
|
||||||
|
/// production use, prefer `create_client_from_config()` which constructs
|
||||||
|
/// the transport with session tracking.
|
||||||
|
#[cfg(test)]
|
||||||
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
||||||
if !matches!(
|
if !matches!(
|
||||||
config.effective_transport(),
|
config.effective_transport(),
|
||||||
@@ -214,7 +219,14 @@ impl McpClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Attach a session manager for Streamable HTTP session tracking.
|
/// Attach a session manager to the **client** only.
|
||||||
|
///
|
||||||
|
/// **Warning:** This does NOT wire the session manager into the underlying
|
||||||
|
/// `HttpMcpTransport`, so the transport will not capture `Mcp-Session-Id`
|
||||||
|
/// from responses. For production use, construct the transport with
|
||||||
|
/// `HttpMcpTransport::with_session_manager()` and pass it to
|
||||||
|
/// `new_with_transport()` instead. See `create_client_from_config()`.
|
||||||
|
#[cfg(test)]
|
||||||
pub fn with_session_manager(mut self, session_manager: Arc<McpSessionManager>) -> Self {
|
pub fn with_session_manager(mut self, session_manager: Arc<McpSessionManager>) -> Self {
|
||||||
self.session_manager = Some(session_manager);
|
self.session_manager = Some(session_manager);
|
||||||
self
|
self
|
||||||
@@ -235,6 +247,12 @@ impl McpClient {
|
|||||||
self.session_manager.is_some()
|
self.session_manager.is_some()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Get the underlying transport (test-only).
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) fn transport(&self) -> &Arc<dyn McpTransport> {
|
||||||
|
&self.transport
|
||||||
|
}
|
||||||
|
|
||||||
/// Get the next request ID.
|
/// Get the next request ID.
|
||||||
fn next_request_id(&self) -> u64 {
|
fn next_request_id(&self) -> u64 {
|
||||||
self.next_id.fetch_add(1, Ordering::SeqCst)
|
self.next_id.fetch_add(1, Ordering::SeqCst)
|
||||||
|
|||||||
+101
-16
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
|||||||
|
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::mcp::config::{EffectiveTransport, McpServerConfig};
|
use crate::tools::mcp::config::{EffectiveTransport, McpServerConfig};
|
||||||
|
use crate::tools::mcp::http_transport::HttpMcpTransport;
|
||||||
use crate::tools::mcp::{McpClient, McpProcessManager, McpSessionManager, McpTransport};
|
use crate::tools::mcp::{McpClient, McpProcessManager, McpSessionManager, McpTransport};
|
||||||
|
|
||||||
/// Error returned when MCP client creation fails.
|
/// Error returned when MCP client creation fails.
|
||||||
@@ -78,33 +79,37 @@ pub async fn create_client_from_config(
|
|||||||
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
||||||
}
|
}
|
||||||
EffectiveTransport::Http => {
|
EffectiveTransport::Http => {
|
||||||
|
// Authenticated (OAuth) path: tokens exist or server requires auth.
|
||||||
if let Some(ref secrets) = secrets {
|
if let Some(ref secrets) = secrets {
|
||||||
let has_tokens =
|
let has_tokens =
|
||||||
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
||||||
|
|
||||||
if has_tokens || server.requires_auth() {
|
if has_tokens || server.requires_auth() {
|
||||||
Ok(McpClient::new_authenticated(
|
return Ok(McpClient::new_authenticated(
|
||||||
server,
|
server,
|
||||||
Arc::clone(session_manager),
|
Arc::clone(session_manager),
|
||||||
Arc::clone(secrets),
|
Arc::clone(secrets),
|
||||||
user_id,
|
user_id,
|
||||||
))
|
));
|
||||||
} else {
|
|
||||||
Ok(McpClient::new_with_config(server)
|
|
||||||
.map_err(|e| McpFactoryError::InvalidConfig {
|
|
||||||
name: server_name.clone(),
|
|
||||||
reason: e.to_string(),
|
|
||||||
})?
|
|
||||||
.with_session_manager(Arc::clone(session_manager)))
|
|
||||||
}
|
}
|
||||||
} else {
|
|
||||||
Ok(McpClient::new_with_config(server)
|
|
||||||
.map_err(|e| McpFactoryError::InvalidConfig {
|
|
||||||
name: server_name,
|
|
||||||
reason: e.to_string(),
|
|
||||||
})?
|
|
||||||
.with_session_manager(Arc::clone(session_manager)))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Non-OAuth HTTP: wire the session manager into the *transport* so
|
||||||
|
// it captures `Mcp-Session-Id` from responses. Passing it only to
|
||||||
|
// the client (via `with_session_manager`) is not enough — the
|
||||||
|
// transport must know about it to read/write the header.
|
||||||
|
let transport = Arc::new(
|
||||||
|
HttpMcpTransport::new(server.url.clone(), server.name.clone())
|
||||||
|
.with_session_manager(Arc::clone(session_manager)),
|
||||||
|
);
|
||||||
|
Ok(McpClient::new_with_transport(
|
||||||
|
server.name.clone(),
|
||||||
|
transport,
|
||||||
|
Some(Arc::clone(session_manager)),
|
||||||
|
secrets,
|
||||||
|
user_id,
|
||||||
|
Some(server),
|
||||||
|
))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -134,4 +139,84 @@ mod tests {
|
|||||||
"non-OAuth HTTP clients must carry a session manager"
|
"non-OAuth HTTP clients must carry a session manager"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Regression test: the factory must wire the session manager into the
|
||||||
|
/// *transport*, not just the client. Otherwise the transport never
|
||||||
|
/// captures `Mcp-Session-Id` from responses and subsequent requests
|
||||||
|
/// lack the header, causing the server to reject them.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_factory_non_oauth_http_transport_captures_session_id() {
|
||||||
|
use axum::http::header::HeaderName;
|
||||||
|
use axum::{Router, http::StatusCode, response::IntoResponse, routing::post};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
const SESSION_ID: &str = "test-session-abc123";
|
||||||
|
|
||||||
|
async fn session_echo() -> impl IntoResponse {
|
||||||
|
let body = serde_json::json!({
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": 1,
|
||||||
|
"result": {}
|
||||||
|
})
|
||||||
|
.to_string();
|
||||||
|
(
|
||||||
|
StatusCode::OK,
|
||||||
|
[(
|
||||||
|
HeaderName::from_static("mcp-session-id"),
|
||||||
|
SESSION_ID.to_string(),
|
||||||
|
)],
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
let app = Router::new().route("/", post(session_echo));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let url = format!("http://127.0.0.1:{}", addr.port());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let server = McpServerConfig::new("session-test", &url);
|
||||||
|
let session_manager = Arc::new(McpSessionManager::new());
|
||||||
|
let process_manager = Arc::new(McpProcessManager::new());
|
||||||
|
|
||||||
|
let client = create_client_from_config(
|
||||||
|
server,
|
||||||
|
&session_manager,
|
||||||
|
&process_manager,
|
||||||
|
None,
|
||||||
|
"test-user",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("factory should succeed for HTTP config");
|
||||||
|
|
||||||
|
// Pre-create a session entry so that update_session_id has something to update.
|
||||||
|
// In production, the MCP initialize handshake calls get_or_create before responses arrive.
|
||||||
|
session_manager.get_or_create("session-test", &url).await;
|
||||||
|
|
||||||
|
// Send a request through the client's transport to trigger session capture.
|
||||||
|
use crate::tools::mcp::protocol::McpRequest;
|
||||||
|
let request = McpRequest {
|
||||||
|
jsonrpc: "2.0".to_string(),
|
||||||
|
id: Some(1),
|
||||||
|
method: "test".to_string(),
|
||||||
|
params: Some(serde_json::json!({})),
|
||||||
|
};
|
||||||
|
let headers = std::collections::HashMap::new();
|
||||||
|
client
|
||||||
|
.transport()
|
||||||
|
.send(&request, &headers)
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
|
||||||
|
// Verify the session manager captured the session ID from the response.
|
||||||
|
let captured = session_manager.get_session_id("session-test").await;
|
||||||
|
assert_eq!(
|
||||||
|
captured.as_deref(),
|
||||||
|
Some(SESSION_ID),
|
||||||
|
"transport must capture Mcp-Session-Id into session manager"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -494,6 +494,34 @@ mod tests {
|
|||||||
assert_eq!(echoed["authorization"], "Bearer oauth-token");
|
assert_eq!(echoed["authorization"], "Bearer oauth-token");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Regression test for #1436: 202 Accepted responses for notifications
|
||||||
|
/// were parsed as JSON, causing "Failed to parse MCP response" errors
|
||||||
|
/// that broke the MCP session handshake.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_wire_202_accepted_for_notification() {
|
||||||
|
use axum::{Router, http::StatusCode, routing::post};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
async fn accept_notification() -> StatusCode {
|
||||||
|
StatusCode::ACCEPTED
|
||||||
|
}
|
||||||
|
|
||||||
|
let app = Router::new().route("/", post(accept_notification));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let url = format!("http://127.0.0.1:{}", addr.port());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let transport = HttpMcpTransport::new(&url, "test-202");
|
||||||
|
let request = McpRequest::initialized_notification();
|
||||||
|
let response = transport.send(&request, &HashMap::new()).await.unwrap();
|
||||||
|
assert!(response.result.is_none());
|
||||||
|
assert!(response.error.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_wire_custom_auth_preserved_when_no_per_request_auth() {
|
async fn test_wire_custom_auth_preserved_when_no_per_request_auth() {
|
||||||
let (url, _handle) = spawn_echo_server().await;
|
let (url, _handle) = spawn_echo_server().await;
|
||||||
|
|||||||
@@ -383,11 +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<(
|
tokio::sync::broadcast::Sender<(uuid::Uuid, String, ironclaw_common::AppEvent)>,
|
||||||
uuid::Uuid,
|
|
||||||
String,
|
|
||||||
crate::channels::web::types::SseEvent,
|
|
||||||
)>,
|
|
||||||
>,
|
>,
|
||||||
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>,
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user