Compare commits

...
Author SHA1 Message Date
Claude 702af21aeb style: fix cargo fmt line-length violation in process_approval
[skip-regression-check]

https://claude.ai/code/session_01Ey7iG1mvBHeAWq6iEccmjB
2026-03-25 17:14:58 +00:00
Zaki d9d48fdc72 Merge branch 'fix/approval-thread-safety' of https://github.com/nearai/ironclaw into fix/approval-thread-safety 2026-03-25 09:59:22 -07:00
ZakiandClaude Opus 4.6 71e200af74 merge: resolve conflicts with staging (identity-based tool recording)
Combines this PR's explicit match/error handling for missing threads
with staging's identity-based record_tool_result_for/record_tool_error_for
methods from the reasoning PR.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-25 09:59:18 -07:00
41ed0a0f98 feat(agent): thread per-tool reasoning through provider, session, and all surfaces (#1513)
* feat(agent): thread per-tool reasoning from LLM through to REPL, HTTP, SSE, and DB

Add end-to-end agent reasoning summaries so users can see *why* the
agent chose specific tools, not just what it did.

- Add `reasoning: Option<String>` to `ToolCall` (all providers)
- Populate from LLM response content in `Reasoning::respond_with_tools`
  and `select_tools`, with per-tool override when providers supply it
- Extend `Turn` with `narrative` and `TurnToolCall` with `rationale` +
  `tool_call_id` for identity-based result matching
- Persist reasoning in DB via existing tool_calls JSON (no migration)
- Add `StatusUpdate::ReasoningUpdate` and `SseEvent::ReasoningUpdate` +
  `SseEvent::JobReasoning` for real-time streaming
- Emit reasoning events in both chat dispatcher and worker job path
- Add `/reasoning [N|all]` command for inspecting turn reasoning
- Surface `narrative` and `rationale` in HTTP `/api/chat/history`

Based on the design from #361 and #456, reconstructed cleanly with
Option<String> to minimize blast radius (vs mandatory String that broke
compilation in #456).

Closes #456

Co-Authored-By: panosAthDBX <[email protected]>
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address PR review feedback from Gemini and Copilot

- Fix `_ => Ok(None)` in agent_loop.rs to avoid accidental shutdown
- Fix fallback in record_tool_result_for/record_tool_error_for to use
  first pending call instead of last_mut (parallel execution safety)
- Include per-tool decisions in WASM channel reasoning messages
- Apply truncate_at_tool_tags + clean_response to shared_reasoning in
  select_tools (parity with respond_with_tools)
- Persist turn-level narrative to DB in tool_calls JSON wrapper
- Parse both old (array) and new (object) tool_calls formats in
  build_turns_from_db_messages for backward compatibility
- Populate reasoning from action.reasoning in execute_plan ToolCalls

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address second round of review comments + merge fixes

- Add reasoning: None to new github_copilot.rs ToolCall sites (from staging merge)
- Run cargo fmt on 4 files with formatting diffs
- Truncate narrative to 1000 chars before DB persistence
- Clone turn data and drop session lock in /reasoning command
- Extract ToolDecisionDto::from_json_array shared helper (deduplicate
  worker/job.rs and orchestrator/api.rs)
- Add unit tests for wrapped tool_calls JSON format with narrative

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address third round of review comments (Copilot + serrrfirat)

- Reword ToolCall.reasoning docstring to reflect provider-supplied or
  fallback contract
- Sanitize narrative through SafetyLayer before storage/emission
- Clean per-tool reasoning via truncate_at_tool_tags + clean_response
  in select_tools (parity with shared reasoning)
- Convert 4 approval-path recording sites in thread_ops.rs to
  identity-based record_tool_result_for/record_tool_error_for
- Preserve tool_call_id and reasoning through restore_from_messages
- Fix has_result/has_error to reject JSON null values
- Truncate tool_call_id to 128 chars before DB persistence
- Add 4 unit tests for record_tool_result_for/error_for edge cases

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address zmanian review — sanitize JobDelegate reasoning + warn on dropped results

- Sanitize narrative and per-tool rationale through SafetyLayer in
  JobDelegate reasoning events (parity with ChatDelegate)
- Add tracing::warn when record_tool_result_for/error_for drops a
  result because no matching or pending tool call exists
- Add 3 unit tests for reasoning normalization (thinking tags,
  tool tags, empty-after-cleaning)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address 4 remaining unreplied review comments

- Clean per-tool reasoning in respond_with_tools via truncate_at_tool_tags
  + clean_response (parity with select_tools)
- Handle wrapped JSON format in rebuild_chat_messages_from_db so cold
  hydration works after persist_tool_calls format change
- Update persist_tool_calls doc comment to describe new JSON shape
- Sanitize per-tool rationale through SafetyLayer in ChatDelegate before
  emission and storage (parity with JobDelegate)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address zmanian review round 2

- Add tracing::debug on fallback-to-pending path in record_tool_result_for
  and record_tool_error_for (item 1)
- Add comment explaining why /reasoning is special-cased in agent_loop.rs
  (item 4)
- Items 2 (narrative persistence), 3 (rationale sanitization), and 5
  (catch-all fix) were already addressed in prior commits

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: panosAthDBX <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-25 08:35:41 -07:00
6daa2f155f fix: ensure LLM calls always end with user message (closes #763) (#1259)
* fix: ensure LLM calls always end with user message (closes #763)

Claude 4.6 models (claude-sonnet-4-6, claude-opus-4-6) no longer support
assistant message prefill — any LLM call where the conversation ends on an
assistant message is rejected with HTTP 400 "This model does not support
assistant message prefill".

The same root cause also triggers NEAR AI's "No user query found in messages"
400 error for the routine engine path.

Two fixes:

1. src/worker/container.rs — before_llm_call()
   After poll_and_inject_prompt(), if no user follow-up arrived and
   handle_text_response() left an assistant message at the end of the
   conversation, inject a sentinel "Continue." user message before
   the next LLM call.

2. src/agent/routine_engine.rs — execute_lightweight_with_tools()
   Before the force_text final completion call, ensure messages end
   with a user-role message. Tool result messages (Role::Tool) satisfy
   Anthropic but not NEAR AI; assistant messages satisfy neither.

Also updates the worker system prompt to instruct the agent to include
the phrase "The job is complete" in its final message, so the agentic
loop can detect termination reliably.

Tested with claude-sonnet-4-6 and claude-opus-4-6.
Workaround: ANTHROPIC_MODEL=claude-sonnet-4-20250514 (still supports prefill).

* fix: broaden sentinel guard to any non-user message (per review)

Gemini suggested the Role::Assistant check in before_llm_call() is too
specific. Changed to !Role::User to match the routine_engine.rs fix and
cover tool results too.

* fix: address zmanian review — JobDelegate sentinel, shared helper, NearAI complete() flattening

- Extract ensure_ends_with_user_message() to src/util.rs with 4 unit tests
  (empty list, after assistant, after tool result, no-op when already user)
- Add sentinel guard to JobDelegate::before_llm_call() in src/worker/job.rs
  so scheduler jobs (CreateJob / /job path) no longer hit Claude 4.6 / NEAR AI 400s
- Replace inline guards in ContainerDelegate and routine_engine.rs with the
  shared helper — all 3 call sites now use one implementation
- Fix complete() in nearai_chat.rs to apply flatten_tool_messages when
  flatten_tool_messages=true — previously only complete_with_tools() flattened,
  so force_text paths could still send role:"tool" messages to NEAR AI
- Update stale comment in container.rs: "assistant message" → "non-user message"
- Add flatten tests in nearai_chat.rs covering the complete() path

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* ci: fix fmt and tar advisory

---------

Co-authored-by: Jacob Lasky <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
Co-authored-by: firat.sertgoz <[email protected]>
2026-03-25 10:31:44 +03:00
706c3a1b47 refactor: extract AppEvent to crates/ironclaw_common (#1615)
* refactor: extract AppEvent to crates/ironclaw_common

SseEvent was defined in src/channels/web/types.rs but imported by 12+
modules across agent, orchestrator, worker, tools, and extensions — it
had become the application-wide event protocol, not a web transport
concern.

Create crates/ironclaw_common as a shared workspace crate and move the
enum there as AppEvent.  Also move the truncate_preview utility which
was similarly leaked from the web gateway into agent modules.

- New crate: crates/ironclaw_common (AppEvent, truncate_preview)
- Rename SseEvent → AppEvent, from_sse_event → from_app_event
- web/types.rs re-exports AppEvent for internal gateway use
- web/util.rs re-exports truncate_preview
- Wire format unchanged (serde renames are on variants, not the enum)

Aligned with the event bus direction on refactor/architectural-hardening
where DomainEvent (≡ AppEvent) is wrapped in a SystemEvent envelope.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor: add AppEvent::event_type() helper, deduplicate match blocks

Address Gemini review: extract the variant→string match into a single
method on AppEvent, replacing the duplicated 22-arm matches in sse.rs
and types.rs.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor: rename leftover sse vars/tests to match AppEvent rename

Address Copilot review: rename sse_event vars to app_event in
orchestrator/api.rs and ws.rs, rename test functions from
test_ws_server_from_sse_* to test_ws_server_from_app_event_*, and
update stale SSE comments.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor: add Deserialize to AppEvent, round-trip test, fix stale comments

Address zmanian review:
- Add Deserialize derive to AppEvent so downstream consumers can
  deserialize incoming events
- Add event_type_matches_serde_type_field test that round-trips every
  variant through serde and asserts event_type() matches the serialized
  "type" field — catches drift between serde renames and the manual match
- Add round_trip_deserialize test for basic Serialize/Deserialize parity
- Update remaining "SSE" references in comments across server.rs,
  manager.rs, ws_gateway_integration.rs, and worker/job.rs

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-24 23:02:46 -07:00
656151783c feat(cli): show credential auth status in tool info (#1572)
* feat(cli): show credential auth status in `tool info`

`ironclaw tool info` now checks the secrets store and shows whether
each required credential is configured or missing, consolidated into
a single Auth section that deduplicates across http.credentials,
auth, and setup.required_secrets. Secrets already shown in Auth are
filtered from the Secrets section to avoid redundancy.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(cli): address review feedback on tool info auth status

- Fix clippy collapsible-if by using `if let` + `&&`
- Use HashMap<String, usize> for O(1) dedup instead of HashSet + linear scan
- Add --user flag to `tool info` for checking non-default user credentials
- Show "? unknown" on secrets store errors instead of silently reporting missing
- Surface secrets store init failure via eprintln instead of silent .ok()
- Sort auth entries by secret name for deterministic output

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(cli): only filter secrets when auth section renders, add regression test

When the secrets store fails to initialize, the Auth section is not
rendered. Previously, secret names were still filtered from the Secrets
section, causing credential names to disappear entirely. Now secrets
are only filtered when the Auth section will actually be displayed.

Adds test verifying auth secret deduplication across auth, setup, and
http.credentials sections, plus secrets store existence checks.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor(cli): extract collect_auth_secrets helper, always render Auth section

Address review feedback:
- Extract dedup logic into `collect_auth_secrets()` so the test exercises
  the same code path as production (not a re-implementation)
- Always render the Auth section when auth secrets exist, showing
  "? unknown" status when the secrets store is unavailable instead of
  hiding credential names entirely
- Lazily init secrets store only when capabilities contain auth secrets,
  avoiding spurious warnings for tools with no auth
- Add test for empty capabilities edge case

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* style(cli): move HashMap/HashSet imports to top of file

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(cli): use correct tagged JSON format for credential location in test

The CredentialLocationSchema uses serde tagged enum format
({"type": "bearer"}), not a bare string ("AuthorizationBearer").

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-24 23:01:19 -07:00
Henry ParkandGitHub 82822d7b25 fix: restore owner-scoped gateway startup (#1625)
* fix: restore owner-scoped gateway startup

* fix: split gateway owner and sender scope

* fix: keep multi-user gateway sender identity

* test: cover gateway sender scope regression

* test: harden e2e startup teardown race

* fix: align gateway owner scope across auth modes
2026-03-24 16:11:53 -07:00
Henry ParkandGitHub dcb2d89e3a Fix hosted OAuth refresh via proxy (#1602)
* Fix hosted OAuth refresh via proxy

* Address OAuth refresh review feedback

* Address new OAuth refresh review comments

* Address additional OAuth refresh review feedback

* Harden proxy exchange redirects
2026-03-24 13:51:30 -07:00
Zaki ManianandGitHub f3da30a454 perf(agent): optimize approval thread resolution (UUID parsing + lock contention) (#1592) 2026-03-24 11:48:30 -07:00
Claude ba4645b192 Merge remote-tracking branch 'origin/staging' into fix/approval-thread-safety 2026-03-24 18:23:27 +00:00
Pierre LE GUENandGitHub 5901451603 fix: remove stale stream_token gate from channel-relay activation (#1623)
* fix: remove stale stream_token gate from channel-relay activation

The relay architecture now uses instance-scoped bearer auth + webhook
callbacks, not streaming. The `relay:<name>:stream_token` secret was
never written by the current OAuth flow, so activation always failed
with AuthRequired.

Replace stream_token with the team_id setting (already stored by the
OAuth callback) as the persistent "auth completed" marker:

- is_relay_channel(): check team_id setting instead of stream_token secret
- activate_channel_relay(): gate on team_id emptiness, not stream_token
- removal flow: delete team_id setting + oauth_state secret
- configure(): return empty allowed-secrets set (relay is OAuth-only)
- configure_token(): return AuthRequired (no manual token entry)
- list(): surface activation_error for relay channels (was hardcoded None)
- Clean up stale comments referencing stream_token / "stored token"
- Update test to match OAuth-only model (no secrets to pass)

Made-with: Cursor

* fix: address CI and review feedback

- Fix pre-existing tunnel/mod.rs test compilation (missing GatewayConfig
  fields: memory_layers, user_tokens, workspace_read_scopes)
- Log warnings on failed team_id/oauth_state cleanup during removal
  instead of silently ignoring errors (gemini review)
- Also delete legacy stream_token secret during removal for backward
  compatibility with pre-webhook installs (codex review)

Made-with: Cursor
2026-03-24 10:49:13 -07:00
Zaki ManianandGitHub d59e2a1d4b Merge branch 'staging' into fix/approval-thread-safety 2026-03-24 10:11:38 -07:00
ZakiandClaude Opus 4.6 bafe945c90 fix: convert remaining let-chains to nested ifs for MSRV compat
Refactor four let-chain patterns (`if let ... && ...`) in
process_approval() to use nested `if`/`if let` blocks. Let-chains
require `#![feature(let_chains)]` which is not available on our MSRV.
Add `#[allow(clippy::collapsible_if)]` with a comment explaining why.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-24 09:14:29 -07:00
d3d517fd67 fix(agent): case-insensitive channel match and user_id filter for event triggers (#1211)
* fix(agent): case-insensitive channel match and user_id filter for event triggers (#1051, #1076)

Event-triggered routines had two bugs preventing them from firing:

1. Channel comparison was case-sensitive (e.g., "Telegram" != "telegram"),
   while emit_system_event already used eq_ignore_ascii_case. Fixed to match.

2. No user_id scoping — routines from any user were evaluated against every
   message. Added ownership check so routines only fire for their owner's
   messages.

Also adds periodic event cache refresh (every ~60s) in the cron ticker so
web/CLI mutations are picked up without requiring the tool path. Upgrades
skip-reason logging from trace to debug for debuggability.

Closes #1051
Refs #1076

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: correct refresh_every from 6 to 4 to match 15s default interval

The default cron_check_interval_secs is 15s, not 10s. With refresh_every=6,
the cache would refresh every 90s instead of the intended ~60s. Fix to 4
ticks (4 * 15s = 60s).

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(agent): address #1211 review -- extract routine_matches_message, fix refresh interval

Extract user/channel filter logic from check_event_triggers into a
standalone pure function routine_matches_message(). Rewrite tests to
call this function directly with controlled Routine and IncomingMessage
values, so they exercise the real code path and would catch a revert.

Add test_no_channel_filter_matches_any_channel for the None channel case.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* ci: re-trigger CI with latest changes

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: add missing IncomingMessage fields in test helper

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(agent): address review -- time-based refresh, trace-level user mismatch, scope guard (#1211)

- Use tokio::time::Instant for cache refresh instead of tick counting
- Downgrade user-mismatch log to trace to reduce noise
- Add early return false for non-Event triggers in routine_matches_message
- Fix doc comment to say 'user scope' instead of 'message sender'

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* style: run cargo fmt on agent_loop.rs

https://claude.ai/code/session_01ABGWibdKVQ3b6pEKtxPPkM

* fix(agent): resolve clippy warnings for unused binding and needless borrow

Fix unused `content` variable in event trigger guard (use `content: _`)
and remove redundant `&` on `message` which was already a reference.

https://claude.ai/code/session_01PzBK21BbUAuZbrfLpoz4Xb

* fix(test): update check_event_triggers call sites to new single-arg signature

The staging merge brought e2e_routine_heartbeat tests that still used
the old 3-argument check_event_triggers(user_id, channel, content)
signature. Updated all 11 call sites to pass &IncomingMessage directly.

[skip-regression-check]

https://claude.ai/code/session_012GrkTDrtDFkpJos2hkgTcE

* fix(agent): address review feedback on event trigger handling

- Use post-hook content for event trigger matching so BeforeInbound
  hooks that rewrite input are respected
- Set MissedTickBehavior::Skip on cron ticker to avoid burst catch-up
  after delays

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* style: cargo fmt

https://claude.ai/code/session_01Va9wwvATNWFAx35GG7Zek7

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: firat.sertgoz <[email protected]>
2026-03-24 10:44:25 +01:00
01678be61d fix(routines): normalize status display across web and CLI (#1469)
* fix(routines): normalize status display across web and CLI surfaces (#1319)

- Use Display (lowercase) instead of Debug (PascalCase) for RunStatus serialization in web handler
- Update JavaScript status class mapping to match lowercase values from the API
- Enrich CLI `routines list` to show running/attention states by querying last run status

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(routines): address review -- batch last-run query, consistent status, simplify ternary (#1319)

- Parallelize last-run lookups with join_all to avoid N+1 sequential queries
- Normalize status in /api/routines/{id}/runs handler to match lowercase convention
- Remove redundant 'running' check in app.js runStatusClass logic

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(db): replace N+1 last-run-status queries with batch method

The CLI routines list was firing a separate list_routine_runs query per
routine to determine each one's last run status. For large routine sets
this overwhelms the connection pool.

Add batch_get_last_run_status to the Database trait with implementations
for both PostgreSQL (DISTINCT ON + ORDER BY) and libSQL (correlated
subquery + in-memory filter). Update the CLI to call the batch method
once instead of N times.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* style: cargo fmt

https://claude.ai/code/session_01Va9wwvATNWFAx35GG7Zek7

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-24 10:41:33 +01:00
fb3548956b fix(tunnel): managed tunnels target wrong port and die from SIGPIPE (#1093)
* fix(tunnel): target webhook server port instead of gateway port

start_managed_tunnel() always used the gateway port (3000) for the
tunnel target. Webhook routes live on the webhook server (HTTP_PORT,
default 8080), not the gateway. The old code never read
config.channels.http — no configuration could work around this.

Extracts resolve_tunnel_target() with regression tests.

* fix(tunnel): prevent SIGPIPE and fix default port fallback

Two fixes for managed tunnel subprocess lifetime:

1. After extracting the public URL from stdout/stderr, the pipe reader
   was dropped (Rust ownership). The tunnel binary's next log write hit
   the closed pipe and got SIGPIPE — killing it silently. Fix: drain
   pipes in background tasks stored in TunnelProcess. Storing without
   reading isn't enough — the OS pipe buffer fills up and the process
   blocks instead.

2. When neither HTTP_PORT nor gateway is configured, the tunnel fell
   back to 127.0.0.1:3000. But the webhook server defaults to
   0.0.0.0:8080 in this case. Now the tunnel matches that fallback.

Affects ngrok (stdout), cloudflare (stderr), and custom (stdout).
Tailscale uses a daemon and is not affected by SIGPIPE.

* fix(tunnel): simplify drain loops and suppress CI false positives

Simplify `while let Ok(Ok(Some(line)))` drain pattern to
`while let Ok(Some(line))` — the extra Ok wrapper was unnecessary.

Add `// safety: test-only` to assert_eq! lines in test module to
suppress the "No panics in production code" CI check which greps
the diff without understanding Rust's #[cfg(test)] module boundaries.

---------

Co-authored-by: firat.sertgoz <[email protected]>
2026-03-24 08:46:22 +01:00
Claude 170566ab93 style: run cargo fmt
https://claude.ai/code/session_01Mdiz3XwyZcjqMkqicaynGs
2026-03-23 14:12:34 +00:00
ZakiandClaude Opus 4.6 9170b28f7b test(agent): strengthen regression test for missing thread error handling
Replace shallow assertion-only test with one that exercises the actual
match-based error detection pattern used in process_approval()'s
rejection and state-setting paths.

Addresses Gemini review feedback on #1579.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-23 06:12:29 -07:00
ZakiandClaude Opus 4.6 cda0df3d7f fix(agent): return errors when approval thread disappears (#1487)
Replace silent if-let-Some patterns with explicit match arms that log
errors and return error responses when threads are not found during
approval processing. Critical state mutations (complete turn, clear
approval, set Processing, await approval) return errors. Auxiliary
operations (record tool result) log errors but continue since the tool
already executed.

Closes #1487

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-23 06:12:29 -07:00
ZakiandClaude Opus 4.6 75f7bea0fd fix(security): eliminate TOCTOU race in approval request_id check (#1486)
Hold session lock for the entire take-verify sequence so pending approval
cannot be lost if a concurrent operation modifies the thread between
take and restore.

Closes #1486

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-23 06:10:49 -07:00
77 changed files with 4791 additions and 991 deletions
Generated
+9
View File
@@ -3428,6 +3428,7 @@ dependencies = [
"hyper-util", "hyper-util",
"iana-time-zone", "iana-time-zone",
"insta", "insta",
"ironclaw_common",
"ironclaw_safety", "ironclaw_safety",
"json5", "json5",
"libsql", "libsql",
@@ -3485,6 +3486,14 @@ dependencies = [
"zip", "zip",
] ]
[[package]]
name = "ironclaw_common"
version = "0.1.0"
dependencies = [
"serde",
"serde_json",
]
[[package]] [[package]]
name = "ironclaw_safety" name = "ironclaw_safety"
version = "0.1.0" version = "0.1.0"
+4 -1
View File
@@ -1,5 +1,5 @@
[workspace] [workspace]
members = [".", "crates/ironclaw_safety"] members = [".", "crates/ironclaw_common", "crates/ironclaw_safety"]
exclude = [ exclude = [
"channels-src/discord", "channels-src/discord",
"channels-src/telegram", "channels-src/telegram",
@@ -100,6 +100,9 @@ tower-http = { version = "0.6", features = ["trace", "cors", "set-header"] }
# Cron scheduling for routines # Cron scheduling for routines
cron = "0.13" cron = "0.13"
# Shared types
ironclaw_common = { path = "crates/ironclaw_common", version = "0.1.0" }
# Safety/sanitization # Safety/sanitization
ironclaw_safety = { path = "crates/ironclaw_safety", version = "0.1.0" } ironclaw_safety = { path = "crates/ironclaw_safety", version = "0.1.0" }
regex = "1" regex = "1"
+18
View File
@@ -0,0 +1,18 @@
[package]
name = "ironclaw_common"
version = "0.1.0"
edition = "2024"
rust-version = "1.92"
description = "Shared types and utilities for the IronClaw workspace"
authors = ["NEAR AI <[email protected]>"]
license = "MIT OR Apache-2.0"
homepage = "https://github.com/nearai/ironclaw"
repository = "https://github.com/nearai/ironclaw"
publish = false
[package.metadata.dist]
dist = false
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
+393
View File
@@ -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");
}
}
+7
View File
@@ -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;
+100
View File
@@ -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>"));
}
}
+21 -4
View File
@@ -1055,10 +1055,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
} }
@@ -1139,9 +1140,9 @@ 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 // Use post-hook content so that BeforeInbound hooks that rewrite
.check_event_triggers(&message.user_id, &message.channel, content) // input are respected by event trigger matching.
.await; let fired = engine.check_event_triggers(message, content).await;
if fired > 0 { if fired > 0 {
tracing::debug!( tracing::debug!(
channel = %message.channel, channel = %message.channel,
@@ -1249,6 +1250,22 @@ 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)))
}
_ => Ok(Some(String::new())),
};
}
// Authorization checks (including restart channel check) are enforced in handle_system_command // Authorization checks (including restart channel check) are enforced in handle_system_command
self.handle_system_command(&command, &args, &message.channel) self.handle_system_command(&command, &args, &message.channel)
.await .await
+1
View File
@@ -414,6 +414,7 @@ mod tests {
id: "call_1".to_string(), id: "call_1".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}; };
let delegate = MockDelegate::new(vec![ let delegate = MockDelegate::new(vec![
tool_calls_output(vec![tool_call]), tool_calls_output(vec![tool_call]),
+89
View File
@@ -465,6 +465,94 @@ impl Agent {
} }
} }
/// Handle `/reasoning [N|all]` — show reasoning history for the active thread.
pub(super) async fn handle_reasoning_command(
&self,
args: &[String],
session: &Arc<Mutex<Session>>,
thread_id: Uuid,
) -> SubmissionResult {
// Clone the turn data we need, then drop the session lock.
let turns_snapshot: Vec<(
usize,
Option<String>,
Vec<crate::agent::session::TurnToolCall>,
)>;
{
let sess = session.lock().await;
let thread = match sess.threads.get(&thread_id) {
Some(t) => t,
None => return SubmissionResult::error("No active thread."),
};
if thread.turns.is_empty() {
return SubmissionResult::ok_with_message("No turns yet.");
}
// Parse argument: default=last turn, "all"=all turns, N=specific turn (1-based).
let selected: Vec<&crate::agent::session::Turn> = match args.first().map(|s| s.as_str())
{
Some("all") => thread.turns.iter().collect(),
Some(n) => match n.parse::<usize>() {
Ok(0) => return SubmissionResult::error("Turn numbers start at 1."),
Ok(num) if num > thread.turns.len() => {
return SubmissionResult::error(format!(
"Turn {} does not exist (max: {}).",
num,
thread.turns.len()
));
}
Ok(num) => vec![&thread.turns[num - 1]],
Err(_) => return SubmissionResult::error("Usage: /reasoning [N|all]"),
},
None => {
// Default: last turn that has tool calls
match thread.turns.iter().rev().find(|t| !t.tool_calls.is_empty()) {
Some(t) => vec![t],
None => {
return SubmissionResult::ok_with_message("No turns with tool calls.");
}
}
}
};
turns_snapshot = selected
.into_iter()
.map(|t| (t.turn_number, t.narrative.clone(), t.tool_calls.clone()))
.collect();
}
// Session lock is now dropped — format output without holding it.
let mut output = String::new();
for (turn_number, narrative, tool_calls) in &turns_snapshot {
output.push_str(&format!("--- Turn {} ---\n", turn_number + 1));
if let Some(narrative) = narrative {
output.push_str(&format!("Reasoning: {}\n", narrative));
}
if tool_calls.is_empty() {
output.push_str(" (no tool calls)\n");
} else {
for tc in tool_calls {
let status = if tc.error.is_some() {
"error"
} else if tc.result.is_some() {
"ok"
} else {
"pending"
};
output.push_str(&format!(" {} [{}]", tc.name, status));
if let Some(ref rationale) = tc.rationale {
output.push_str(&format!("{}", rationale));
}
output.push('\n');
}
}
output.push('\n');
}
SubmissionResult::response(output.trim_end())
}
/// Handle system commands that bypass thread-state checks entirely. /// Handle system commands that bypass thread-state checks entirely.
pub(super) async fn handle_system_command( pub(super) async fn handle_system_command(
&self, &self,
@@ -480,6 +568,7 @@ impl Agent {
" /version Show version info\n", " /version Show version info\n",
" /tools List available tools\n", " /tools List available tools\n",
" /debug Toggle debug mode\n", " /debug Toggle debug mode\n",
" /reasoning [N|all] Show agent reasoning for turns\n",
" /ping Connectivity check\n", " /ping Connectivity check\n",
"\n", "\n",
"Jobs:\n", "Jobs:\n",
+79 -5
View File
@@ -420,6 +420,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
content: Option<String>, content: Option<String>,
reason_ctx: &mut ReasoningContext, reason_ctx: &mut ReasoningContext,
) -> Result<Option<LoopOutcome>, Error> { ) -> Result<Option<LoopOutcome>, Error> {
// Extract and sanitize the narrative before consuming `content`.
let narrative = content
.as_deref()
.filter(|c| !c.trim().is_empty())
.map(|c| {
let sanitized = self
.agent
.safety()
.sanitize_tool_output("agent_narrative", c);
sanitized.content
})
.filter(|c| !c.trim().is_empty());
// Add the assistant message with tool_calls to context. // Add the assistant message with tool_calls to context.
// OpenAI protocol requires this before tool-result messages. // OpenAI protocol requires this before tool-result messages.
reason_ctx reason_ctx
@@ -440,6 +453,41 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
) )
.await; .await;
// Build per-tool decisions for the reasoning update.
// Sanitize each rationale through SafetyLayer (parity with JobDelegate).
let decisions: Vec<crate::channels::ToolDecision> = tool_calls
.iter()
.filter_map(|tc| {
tc.reasoning.as_ref().map(|r| {
let sanitized = self
.agent
.safety()
.sanitize_tool_output("tool_rationale", r)
.content;
crate::channels::ToolDecision {
tool_name: tc.name.clone(),
rationale: sanitized,
}
})
})
.collect();
// Emit reasoning update to channels.
if narrative.is_some() || !decisions.is_empty() {
let _ = self
.agent
.channels
.send_status(
&self.message.channel,
StatusUpdate::ReasoningUpdate {
narrative: narrative.clone().unwrap_or_default(),
decisions: decisions.clone(),
},
&self.message.metadata,
)
.await;
}
// Record tool calls in the thread with sensitive params redacted. // Record tool calls in the thread with sensitive params redacted.
{ {
let mut redacted_args: Vec<serde_json::Value> = Vec::with_capacity(tool_calls.len()); let mut redacted_args: Vec<serde_json::Value> = Vec::with_capacity(tool_calls.len());
@@ -455,8 +503,23 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
if let Some(thread) = sess.threads.get_mut(&self.thread_id) if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut() && let Some(turn) = thread.last_turn_mut()
{ {
// Set turn-level narrative.
if turn.narrative.is_none() {
turn.narrative = narrative;
}
for (tc, safe_args) in tool_calls.iter().zip(redacted_args) { for (tc, safe_args) in tool_calls.iter().zip(redacted_args) {
turn.record_tool_call(&tc.name, safe_args); let sanitized_rationale = tc.reasoning.as_ref().map(|r| {
self.agent
.safety()
.sanitize_tool_output("tool_rationale", r)
.content
});
turn.record_tool_call_with_reasoning(
&tc.name,
safe_args,
sanitized_rationale,
Some(tc.id.clone()),
);
} }
} }
} }
@@ -726,7 +789,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
if let Some(thread) = sess.threads.get_mut(&self.thread_id) if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut() && let Some(turn) = thread.last_turn_mut()
{ {
turn.record_tool_error(error_msg.clone()); turn.record_tool_error_for(&tc.id, error_msg.clone());
} }
} }
reason_ctx reason_ctx
@@ -852,16 +915,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
Err(e) => format!("Tool '{}' failed: {}", tc.name, e), Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
}; };
// Record sanitized result in thread // Record sanitized result in thread (identity-based matching).
{ {
let mut sess = self.session.lock().await; let mut sess = self.session.lock().await;
if let Some(thread) = sess.threads.get_mut(&self.thread_id) if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut() && let Some(turn) = thread.last_turn_mut()
{ {
if is_tool_error { if is_tool_error {
turn.record_tool_error(result_content.clone()); turn.record_tool_error_for(&tc.id, result_content.clone());
} else { } else {
turn.record_tool_result(serde_json::json!(result_content)); turn.record_tool_result_for(
&tc.id,
serde_json::json!(result_content),
);
} }
} }
} }
@@ -1462,11 +1528,13 @@ mod tests {
id: "call_2".to_string(), id: "call_2".to_string(),
name: "http".to_string(), name: "http".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}), arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
}, },
ToolCall { ToolCall {
id: "call_3".to_string(), id: "call_3".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({"message": "done"}), arguments: serde_json::json!({"message": "done"}),
reasoning: None,
}, },
], ],
user_timezone: None, user_timezone: None,
@@ -1652,6 +1720,7 @@ mod tests {
id: "call_1".to_string(), id: "call_1".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({"message": "hi"}), arguments: serde_json::json!({"message": "hi"}),
reasoning: None,
}], }],
), ),
ChatMessage::tool_result("call_1", "echo", "hi"), ChatMessage::tool_result("call_1", "echo", "hi"),
@@ -1744,11 +1813,13 @@ mod tests {
id: "c1".to_string(), id: "c1".to_string(),
name: "http".to_string(), name: "http".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}, },
ToolCall { ToolCall {
id: "c2".to_string(), id: "c2".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}, },
], ],
), ),
@@ -1782,6 +1853,7 @@ mod tests {
id: "c1".to_string(), id: "c1".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}], }],
), ),
ChatMessage::tool_result("c1", "echo", "done"), ChatMessage::tool_result("c1", "echo", "done"),
@@ -1912,6 +1984,7 @@ mod tests {
id: crate::llm::generate_tool_call_id(0, 0), id: crate::llm::generate_tool_call_id(0, 0),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({"message": "looping"}), arguments: serde_json::json!({"message": "looping"}),
reasoning: None,
}], }],
input_tokens: 0, input_tokens: 0,
output_tokens: 5, output_tokens: 5,
@@ -2065,6 +2138,7 @@ mod tests {
id: crate::llm::generate_tool_call_id(0, 0), id: crate::llm::generate_tool_call_id(0, 0),
name: "nonexistent_tool".to_string(), name: "nonexistent_tool".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}], }],
input_tokens: 0, input_tokens: 0,
output_tokens: 5, output_tokens: 5,
+24 -24
View File
@@ -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,
+190 -20
View File
@@ -24,7 +24,7 @@ use crate::agent::Scheduler;
use crate::agent::routine::{ use crate::agent::routine::{
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire, NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
}; };
use crate::channels::OutgoingResponse; use crate::channels::{IncomingMessage, OutgoingResponse};
use crate::config::RoutineConfig; use crate::config::RoutineConfig;
use crate::context::{JobContext, JobState}; use crate::context::{JobContext, JobState};
use crate::db::Database; use crate::db::Database;
@@ -56,6 +56,40 @@ pub enum SandboxReadiness {
DockerUnavailable, DockerUnavailable,
} }
/// Check whether an event-triggered routine's user/channel filters match an
/// incoming message.
///
/// Returns `true` if:
/// - The routine has an `Event` trigger (non-Event routines always return `false`)
/// - The routine's `user_id` matches the message's user scope
/// - The routine's channel filter (if any) matches the message channel
/// case-insensitively
///
/// This is a pure function extracted from `check_event_triggers` so the
/// filter logic can be unit-tested without async infrastructure.
pub(crate) fn routine_matches_message(routine: &Routine, message: &IncomingMessage) -> bool {
// Only Event-triggered routines can match incoming messages.
if !matches!(routine.trigger, Trigger::Event { .. }) {
return false;
}
// User ownership filter — only fire routines scoped to this user.
if routine.user_id != message.user_id {
return false;
}
// Channel filter (case-insensitive, matching emit_system_event behavior)
if let Trigger::Event {
channel: Some(ch), ..
} = &routine.trigger
&& !ch.eq_ignore_ascii_case(&message.channel)
{
return false;
}
true
}
/// The routine execution engine. /// The routine execution engine.
pub struct RoutineEngine { pub struct RoutineEngine {
config: RoutineConfig, config: RoutineConfig,
@@ -167,10 +201,7 @@ 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 {
/// Accepts only the three fields needed for matching (user scope, channel,
/// message content) so callers never need to clone a full `IncomingMessage`.
pub async fn check_event_triggers(&self, user_id: &str, channel: &str, content: &str) -> usize {
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.
@@ -208,16 +239,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 +267,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;
} }
@@ -1502,7 +1541,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);
@@ -1781,6 +1823,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;
@@ -1788,7 +1837,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();
}
} }
}) })
} }
@@ -1854,7 +1907,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]
@@ -2052,6 +2111,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![
+193 -2
View File
@@ -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
View File
@@ -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"
);
}
} }
+11
View File
@@ -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 {
+314 -103
View File
@@ -16,12 +16,12 @@ use crate::agent::dispatcher::{
}; };
use crate::agent::session::{MAX_PENDING_MESSAGES, PendingApproval, Session, ThreadState}; use crate::agent::session::{MAX_PENDING_MESSAGES, PendingApproval, Session, ThreadState};
use crate::agent::submission::SubmissionResult; use crate::agent::submission::SubmissionResult;
use crate::channels::web::util::truncate_preview;
use crate::channels::{IncomingMessage, StatusUpdate}; use crate::channels::{IncomingMessage, StatusUpdate};
use crate::context::JobContext; use crate::context::JobContext;
use crate::error::Error; use crate::error::Error;
use crate::llm::{ChatMessage, ToolCall}; use crate::llm::{ChatMessage, ToolCall};
use crate::tools::redact_params; use crate::tools::redact_params;
use ironclaw_common::truncate_preview;
const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID."; const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID.";
@@ -513,10 +513,10 @@ impl Agent {
}; };
thread.complete_turn(&response); thread.complete_turn(&response);
let (turn_number, tool_calls) = thread let (turn_number, tool_calls, narrative) = thread
.turns .turns
.last() .last()
.map(|t| (t.turn_number, t.tool_calls.clone())) .map(|t| (t.turn_number, t.tool_calls.clone(), t.narrative.clone()))
.unwrap_or_default(); .unwrap_or_default();
let _ = self let _ = self
.channels .channels
@@ -534,6 +534,7 @@ impl Agent {
&message.user_id, &message.user_id,
turn_number, turn_number,
&tool_calls, &tool_calls,
narrative.as_deref(),
) )
.await; .await;
self.persist_assistant_response( self.persist_assistant_response(
@@ -725,7 +726,9 @@ impl Agent {
/// ///
/// Stored between the user and assistant messages so that /// Stored between the user and assistant messages so that
/// `build_turns_from_db_messages` can reconstruct the tool call history. /// `build_turns_from_db_messages` can reconstruct the tool call history.
/// Content is a JSON array of tool call summaries. /// Content is a JSON object: `{ "calls": [...], "narrative": "..." }`.
/// The `calls` array contains tool call summaries with optional `rationale`
/// and `tool_call_id` fields. Legacy rows may be plain JSON arrays.
pub(super) async fn persist_tool_calls( pub(super) async fn persist_tool_calls(
&self, &self,
thread_id: Uuid, thread_id: Uuid,
@@ -733,6 +736,7 @@ impl Agent {
user_id: &str, user_id: &str,
turn_number: usize, turn_number: usize,
tool_calls: &[crate::agent::session::TurnToolCall], tool_calls: &[crate::agent::session::TurnToolCall],
narrative: Option<&str>,
) { ) {
if tool_calls.is_empty() { if tool_calls.is_empty() {
return; return;
@@ -767,11 +771,30 @@ impl Agent {
if let Some(ref error) = tc.error { if let Some(ref error) = tc.error {
obj["error"] = serde_json::Value::String(truncate_preview(error, 200)); obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
} }
if let Some(ref rationale) = tc.rationale {
obj["rationale"] = serde_json::Value::String(truncate_preview(rationale, 500));
}
if let Some(ref tool_call_id) = tc.tool_call_id {
obj["tool_call_id"] =
serde_json::Value::String(truncate_preview(tool_call_id, 128));
}
obj obj
}) })
.collect(); .collect();
let content = match serde_json::to_string(&summaries) { // Wrap in an object with optional narrative so it can be reconstructed.
// safety: no byte-index slicing here; comment describes JSON shape
let wrapper = if let Some(n) = narrative {
serde_json::json!({
"narrative": truncate_preview(n, 1000),
"calls": summaries,
})
} else {
serde_json::json!({
"calls": summaries,
})
};
let content = match serde_json::to_string(&wrapper) {
Ok(c) => c, Ok(c) => c,
Err(e) => { Err(e) => {
tracing::warn!("Failed to serialize tool calls: {}", e); tracing::warn!("Failed to serialize tool calls: {}", e);
@@ -945,6 +968,10 @@ impl Agent {
} }
/// Process an approval or rejection of a pending tool execution. /// Process an approval or rejection of a pending tool execution.
// Nested `if` blocks are intentional: collapsing them would produce
// `if let … && …` (let-chains), which require `#![feature(let_chains)]`
// and are not available on our MSRV.
#[allow(clippy::collapsible_if)]
pub(super) async fn process_approval( pub(super) async fn process_approval(
&self, &self,
message: &IncomingMessage, message: &IncomingMessage,
@@ -954,7 +981,10 @@ impl Agent {
approved: bool, approved: bool,
always: bool, always: bool,
) -> Result<SubmissionResult, Error> { ) -> Result<SubmissionResult, Error> {
// Get pending approval for this thread // Get pending approval for this thread.
// The take-verify sequence is atomic under a single lock acquisition
// to prevent a TOCTOU race where a concurrent operation could modify
// or delete the thread between take and restore (#1486).
let pending = { let pending = {
let mut sess = session.lock().await; let mut sess = session.lock().await;
let thread = sess let thread = sess
@@ -972,33 +1002,31 @@ impl Agent {
return Ok(SubmissionResult::ok_with_message("")); return Ok(SubmissionResult::ok_with_message(""));
} }
thread.take_pending_approval() let taken = match thread.take_pending_approval() {
}; Some(p) => p,
None => {
tracing::debug!(
%thread_id,
"Ignoring stale approval: no pending approval found"
);
return Ok(SubmissionResult::ok_with_message(""));
}
};
let pending = match pending { // Verify request ID while still holding the lock — atomic with take
Some(p) => p, if let Some(req_id) = request_id {
None => { if req_id != taken.request_id {
tracing::debug!( // Restore atomically under same lock
%thread_id, thread.await_approval(taken);
"Ignoring stale approval: no pending approval found" return Ok(SubmissionResult::error(
); "Request ID mismatch. Use the correct request ID.",
return Ok(SubmissionResult::ok_with_message("")); ));
}
} }
};
// Verify request ID if provided taken
if let Some(req_id) = request_id // Lock dropped here — pending approval validated
&& req_id != pending.request_id };
{
// Put it back and return error
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.await_approval(pending);
}
return Ok(SubmissionResult::error(
"Request ID mismatch. Use the correct request ID.",
));
}
if approved { if approved {
// If always, add to auto-approved set // If always, add to auto-approved set
@@ -1015,8 +1043,19 @@ impl Agent {
// Reset thread state to processing // Reset thread state to processing
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) { match sess.threads.get_mut(&thread_id) {
thread.state = ThreadState::Processing; Some(thread) => {
thread.state = ThreadState::Processing;
}
None => {
tracing::error!(
%thread_id,
"Thread disappeared while setting state to Processing during approval"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
} }
} }
@@ -1067,20 +1106,20 @@ impl Agent {
) )
.await; .await;
if let Ok(ref output) = tool_result if let Ok(ref output) = tool_result {
&& !output.is_empty() if !output.is_empty() {
{ let _ = self
let _ = self .channels
.channels .send_status(
.send_status( &message.channel,
&message.channel, StatusUpdate::ToolResult {
StatusUpdate::ToolResult { name: pending.tool_name.clone(),
name: pending.tool_name.clone(), preview: output.clone(),
preview: output.clone(), },
}, &message.metadata,
&message.metadata, )
) .await;
.await; }
} }
// Build context including the tool result // Build context including the tool result
@@ -1100,13 +1139,27 @@ impl Agent {
// Record sanitized result in thread // Record sanitized result in thread
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) match sess.threads.get_mut(&thread_id) {
&& let Some(turn) = thread.last_turn_mut() Some(thread) => {
{ if 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(
} else { &pending.tool_call_id,
turn.record_tool_result(serde_json::json!(result_content)); result_content.clone(),
);
} else {
turn.record_tool_result_for(
&pending.tool_call_id,
serde_json::json!(result_content),
);
}
}
}
None => {
tracing::error!(
%thread_id,
"Thread disappeared while recording tool result during approval"
);
} }
} }
} }
@@ -1325,20 +1378,20 @@ impl Agent {
let mut deferred_auth: Option<String> = None; let mut deferred_auth: Option<String> = None;
for (tc, deferred_result) in exec_results { for (tc, deferred_result) in exec_results {
if let Ok(ref output) = deferred_result if let Ok(ref output) = deferred_result {
&& !output.is_empty() if !output.is_empty() {
{ let _ = self
let _ = self .channels
.channels .send_status(
.send_status( &message.channel,
&message.channel, StatusUpdate::ToolResult {
StatusUpdate::ToolResult { name: tc.name.clone(),
name: tc.name.clone(), preview: output.clone(),
preview: output.clone(), },
}, &message.metadata,
&message.metadata, )
) .await;
.await; }
} }
// Sanitize first, then record the cleaned version in thread. // Sanitize first, then record the cleaned version in thread.
@@ -1354,32 +1407,45 @@ impl Agent {
// Record sanitized result in thread // Record sanitized result in thread
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) match sess.threads.get_mut(&thread_id) {
&& let Some(turn) = thread.last_turn_mut() Some(thread) => {
{ if 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),
);
}
}
}
None => {
tracing::error!(
%thread_id,
tool_name = %tc.name,
"Thread disappeared while recording deferred tool result during approval"
);
} }
} }
} }
// Auth detection — defer return until all results are recorded // Auth detection — defer return until all results are recorded
if deferred_auth.is_none() if deferred_auth.is_none() {
&& let Some((ext_name, instructions)) = if let Some((ext_name, instructions)) =
check_auth_required(&tc.name, &deferred_result) check_auth_required(&tc.name, &deferred_result)
{ {
self.handle_auth_intercept( self.handle_auth_intercept(
&session, &session,
thread_id, thread_id,
message, message,
&deferred_result, &deferred_result,
ext_name, ext_name,
instructions.clone(), instructions.clone(),
) )
.await; .await;
deferred_auth = Some(instructions); deferred_auth = Some(instructions);
}
} }
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content)); context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
@@ -1413,8 +1479,19 @@ impl Agent {
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) { match sess.threads.get_mut(&thread_id) {
thread.await_approval(new_pending); Some(thread) => {
thread.await_approval(new_pending);
}
None => {
tracing::error!(
%thread_id,
"Thread disappeared while setting up deferred tool approval"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
} }
} }
@@ -1459,10 +1536,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 +1548,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(
@@ -1546,17 +1624,28 @@ impl Agent {
); );
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) { match sess.threads.get_mut(&thread_id) {
thread.clear_pending_approval(); Some(thread) => {
thread.complete_turn(&rejection); thread.clear_pending_approval();
// User message already persisted at turn start; save rejection response thread.complete_turn(&rejection);
self.persist_assistant_response( // User message already persisted at turn start; save rejection response
thread_id, self.persist_assistant_response(
&message.channel, thread_id,
&message.user_id, &message.channel,
&rejection, &message.user_id,
) &rejection,
.await; )
.await;
}
None => {
tracing::error!(
%thread_id,
"Thread disappeared during approval rejection"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
} }
} }
@@ -1816,7 +1905,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 +1941,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();
@@ -2098,6 +2204,70 @@ mod tests {
} }
} }
#[tokio::test]
async fn test_approval_on_missing_thread_should_error() {
// Regression for #1487: when a thread disappears from the session
// during approval processing, the code must return a visible error
// rather than silently succeeding.
//
// We can't call process_approval() directly (requires full Agent),
// so we simulate the exact code pattern used in the rejection and
// state-setting paths: lock session, match on get_mut, verify the
// None arm produces an error.
use crate::agent::session::{Session, Thread, ThreadState};
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;
let thread_id = Uuid::new_v4();
let session_id = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("test-user")));
// Scenario 1: Thread never existed
{
let sess = session.lock().await;
let result = match sess.threads.get(&thread_id) {
Some(_) => Ok("processed"),
None => Err("Internal error: thread no longer exists"),
};
assert!(result.is_err());
assert_eq!(
result.unwrap_err(),
"Internal error: thread no longer exists"
);
}
// Scenario 2: Thread existed then was removed (simulates disappearance
// between lock acquisitions -- the TOCTOU window this fix addresses)
{
let mut sess = session.lock().await;
let mut thread = Thread::with_id(thread_id, session_id);
thread.start_turn("pending approval");
thread.state = ThreadState::AwaitingApproval;
sess.threads.insert(thread_id, thread);
}
{
let mut sess = session.lock().await;
// Simulate thread disappearing (e.g., pruned by another task)
sess.threads.remove(&thread_id);
// The rejection path must detect this and return an error
let result = match sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.clear_pending_approval();
thread.complete_turn("rejected");
Ok("rejection persisted")
}
None => Err("Internal error: thread no longer exists"),
};
assert!(result.is_err());
assert_eq!(
result.unwrap_err(),
"Internal error: thread no longer exists"
);
}
}
#[test] #[test]
fn test_queue_cap_rejects_at_capacity() { fn test_queue_cap_rejects_at_capacity() {
use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState}; use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState};
@@ -2204,6 +2374,47 @@ mod tests {
assert!(t.pending_messages.is_empty()); assert!(t.pending_messages.is_empty());
} }
#[test]
fn test_approval_request_id_mismatch_restores_pending() {
// Regression test for #1486: after a request_id mismatch, the pending
// approval must still be intact (take + verify + restore is atomic).
use crate::agent::session::{PendingApproval, Thread, ThreadState};
use uuid::Uuid;
let session_id = Uuid::new_v4();
let thread_id = Uuid::new_v4();
let mut thread = Thread::with_id(thread_id, session_id);
let correct_request_id = Uuid::new_v4();
let pending = PendingApproval {
request_id: correct_request_id,
tool_name: "shell".to_string(),
parameters: serde_json::json!({}),
display_parameters: serde_json::json!({}),
description: "test".to_string(),
tool_call_id: "call_0".to_string(),
context_messages: vec![],
deferred_tool_calls: vec![],
user_timezone: None,
allow_always: true,
};
thread.await_approval(pending);
assert_eq!(thread.state, ThreadState::AwaitingApproval);
// Simulate: take, verify mismatch, restore -- all must be atomic
let taken = thread.take_pending_approval().unwrap();
assert_eq!(taken.request_id, correct_request_id);
// On mismatch, restore
thread.await_approval(taken);
// Must still be in AwaitingApproval with pending intact
assert_eq!(thread.state, ThreadState::AwaitingApproval);
assert!(thread.pending_approval.is_some());
assert_eq!(
thread.pending_approval.as_ref().unwrap().request_id,
correct_request_id
);
}
// Helper function to extract the approval message without needing a full Agent instance // Helper function to extract the approval message without needing a full Agent instance
fn extract_approval_message( fn extract_approval_message(
session: &crate::agent::session::Session, session: &crate::agent::session::Session,
+1 -7
View File
@@ -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,
+16
View File
@@ -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
View File
@@ -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;
+14
View File
@@ -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.
@@ -841,6 +842,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
} }
+14
View File
@@ -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,
}
}
}) })
} }
+5 -3
View File
@@ -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();
+2 -2
View File
@@ -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,
+52 -19
View File
@@ -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,
}, },
+2
View File
@@ -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);
+32 -28
View File
@@ -345,8 +345,10 @@ pub struct GatewayState {
pub job_manager: Option<Arc<ContainerJobManager>>, pub job_manager: Option<Arc<ContainerJobManager>>,
/// Prompt queue for Claude Code follow-up prompts. /// Prompt queue for Claude Code follow-up prompts.
pub prompt_queue: Option<PromptQueue>, pub prompt_queue: Option<PromptQueue>,
/// Default user ID (fallback for non-request contexts like heartbeat/routines). /// Durable owner scope for persistence and unauthenticated callback flows.
pub default_user_id: String, pub owner_id: String,
/// Default sender/routing identity for gateway-originated messages.
pub default_sender_id: String,
/// Shutdown signal sender. /// Shutdown signal sender.
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>, pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
/// WebSocket connection tracker. /// WebSocket connection tracker.
@@ -775,7 +777,7 @@ async fn oauth_callback_handler(
error = %error, error = %error,
"OAuth callback received with malformed state" "OAuth callback received with malformed state"
); );
clear_auth_mode(&state, &state.default_user_id).await; clear_auth_mode(&state, &state.owner_id).await;
return oauth_error_page("IronClaw"); return oauth_error_page("IronClaw");
} }
}; };
@@ -811,7 +813,7 @@ async fn oauth_callback_handler(
if let Some(ref sse) = flow.sse_manager { if let Some(ref sse) = flow.sse_manager {
sse.broadcast_for_user( sse.broadcast_for_user(
&flow.user_id, &flow.user_id,
SseEvent::AuthCompleted { AppEvent::AuthCompleted {
extension_name: flow.extension_name.clone(), extension_name: flow.extension_name.clone(),
success: false, success: false,
message: "OAuth flow expired. Please try again.".to_string(), message: "OAuth flow expired. Please try again.".to_string(),
@@ -949,11 +951,11 @@ async fn oauth_callback_handler(
message message
}; };
// Broadcast SSE event to notify the web UI // Broadcast event to notify the web UI
if let Some(ref sse) = flow.sse_manager { if let Some(ref sse) = flow.sse_manager {
sse.broadcast_for_user( sse.broadcast_for_user(
&flow.user_id, &flow.user_id,
SseEvent::AuthCompleted { AppEvent::AuthCompleted {
extension_name: flow.extension_name, extension_name: flow.extension_name,
success, success,
message: final_message.clone(), message: final_message.clone(),
@@ -1136,7 +1138,7 @@ async fn slack_relay_oauth_callback_handler(
let state_key = format!("relay:{}:oauth_state", DEFAULT_RELAY_NAME); let state_key = format!("relay:{}:oauth_state", DEFAULT_RELAY_NAME);
let stored_state = match ext_mgr let stored_state = match ext_mgr
.secrets() .secrets()
.get_decrypted(&state.default_user_id, &state_key) .get_decrypted(&state.owner_id, &state_key)
.await .await
{ {
Ok(secret) => secret.expose().to_string(), Ok(secret) => secret.expose().to_string(),
@@ -1160,10 +1162,7 @@ async fn slack_relay_oauth_callback_handler(
} }
// Delete the nonce (one-time use) // Delete the nonce (one-time use)
let _ = ext_mgr let _ = ext_mgr.secrets().delete(&state.owner_id, &state_key).await;
.secrets()
.delete(&state.default_user_id, &state_key)
.await;
let result: Result<(), String> = async { let result: Result<(), String> = async {
let store = state.store.as_ref().ok_or_else(|| { let store = state.store.as_ref().ok_or_else(|| {
@@ -1174,16 +1173,12 @@ async fn slack_relay_oauth_callback_handler(
// Store team_id in settings // Store team_id in settings
let team_id_key = format!("relay:{}:team_id", DEFAULT_RELAY_NAME); let team_id_key = format!("relay:{}:team_id", DEFAULT_RELAY_NAME);
let _ = store let _ = store
.set_setting( .set_setting(&state.owner_id, &team_id_key, &serde_json::json!(team_id))
&state.default_user_id,
&team_id_key,
&serde_json::json!(team_id),
)
.await; .await;
// Activate the relay channel // Activate the relay channel
ext_mgr ext_mgr
.activate_stored_relay(DEFAULT_RELAY_NAME, &state.default_user_id) .activate_stored_relay(DEFAULT_RELAY_NAME, &state.owner_id)
.await .await
.map_err(|e| format!("Failed to activate relay channel: {}", e))?; .map_err(|e| format!("Failed to activate relay channel: {}", e))?;
@@ -1202,8 +1197,8 @@ async fn slack_relay_oauth_callback_handler(
} }
}; };
// Broadcast SSE event to notify the web UI // Broadcast event to notify the web UI
state.sse.broadcast(SseEvent::AuthCompleted { state.sse.broadcast(AppEvent::AuthCompleted {
extension_name: DEFAULT_RELAY_NAME.to_string(), extension_name: DEFAULT_RELAY_NAME.to_string(),
success, success,
message: message.clone(), message: message.clone(),
@@ -1303,6 +1298,9 @@ async fn chat_send_handler(
} }
let mut msg = IncomingMessage::new("gateway", &user.user_id, &req.content); let mut msg = IncomingMessage::new("gateway", &user.user_id, &req.content);
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
msg = msg.with_sender_id(&state.default_sender_id);
}
// Prefer timezone from JSON body, fall back to X-Timezone header // Prefer timezone from JSON body, fall back to X-Timezone header
let tz = req let tz = req
.timezone .timezone
@@ -1404,6 +1402,9 @@ async fn chat_approval_handler(
})?; })?;
let mut msg = IncomingMessage::new("gateway", &user.user_id, content); let mut msg = IncomingMessage::new("gateway", &user.user_id, content);
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
msg = msg.with_sender_id(&state.default_sender_id);
}
if let Some(ref thread_id) = req.thread_id { if let Some(ref thread_id) = req.thread_id {
msg = msg.with_thread(thread_id); msg = msg.with_thread(thread_id);
@@ -1470,7 +1471,7 @@ 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,
@@ -1483,7 +1484,7 @@ 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,
@@ -1492,7 +1493,7 @@ async fn chat_auth_token_handler(
} else { } else {
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: false, success: false,
message: result.message, message: result.message,
@@ -1508,7 +1509,7 @@ 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,
@@ -1724,8 +1725,10 @@ 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();
@@ -2476,7 +2479,7 @@ async fn extensions_setup_submit_handler(
// auth card or setup modal that was triggered by tool_auth/tool_activate. // auth card or setup modal that was triggered by tool_auth/tool_activate.
state.sse.broadcast_for_user( state.sse.broadcast_for_user(
&user.user_id, &user.user_id,
SseEvent::AuthCompleted { AppEvent::AuthCompleted {
extension_name: name.clone(), extension_name: name.clone(),
success: result.activated, success: result.activated,
message: resp.message.clone(), message: resp.message.clone(),
@@ -2572,7 +2575,7 @@ 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,
@@ -2976,7 +2979,8 @@ mod tests {
store: None, store: None,
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
default_user_id: "test".to_string(), owner_id: "test".to_string(),
default_sender_id: "test".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: None, ws_tracker: None,
llm_provider: None, llm_provider: None,
@@ -3167,7 +3171,7 @@ mod tests {
Ok(Ok(scoped)) Ok(Ok(scoped))
if matches!( if matches!(
scoped.event, scoped.event,
crate::channels::web::types::SseEvent::AuthRequired { .. } crate::channels::web::types::AppEvent::AuthRequired { .. }
) => ) =>
{ {
panic!("verification responses should not emit auth_required SSE events") panic!("verification responses should not emit auth_required SSE events")
@@ -3449,7 +3453,7 @@ mod tests {
assert_eq!(resp.status(), StatusCode::OK); assert_eq!(resp.status(), StatusCode::OK);
match receiver.recv().await.expect("auth_completed event").event { match receiver.recv().await.expect("auth_completed event").event {
crate::channels::web::types::SseEvent::AuthCompleted { crate::channels::web::types::AppEvent::AuthCompleted {
extension_name, extension_name,
success, success,
message, message,
+19 -42
View File
@@ -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(); // safety: test-only let e = bob.next().await.unwrap(); // safety: test-only
assert!(matches!(e, SseEvent::Heartbeat)); // safety: test assertion assert!(matches!(e, AppEvent::Heartbeat)); // safety: test assertion
} }
} }
+3 -3
View File
@@ -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>'
+2 -1
View File
@@ -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,
+38 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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);
}
} }
+6 -5
View File
@@ -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,
+400 -24
View File
@@ -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`
@@ -661,6 +685,48 @@ pub struct ProxyTokenExchangeRequest<'a> {
pub extra_token_params: &'a HashMap<String, String>, pub extra_token_params: &'a HashMap<String, String>,
} }
pub struct ProxyRefreshTokenRequest<'a> {
pub proxy_url: &'a str,
pub gateway_token: &'a str,
pub token_url: &'a str,
pub client_id: &'a str,
pub client_secret: Option<&'a str>,
pub refresh_token: &'a str,
pub provider: Option<&'a str>,
}
fn oauth_token_response_from_json(
token_data: serde_json::Value,
access_token_field: &str,
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
let access_token = token_data
.get(access_token_field)
.and_then(|v| v.as_str())
.ok_or_else(|| {
let fields: Vec<&str> = token_data
.as_object()
.map(|o| o.keys().map(|k| k.as_str()).collect())
.unwrap_or_default();
OAuthCallbackError::Io(format!(
"No '{}' field in proxy response (fields present: {:?})",
access_token_field, fields
))
})?
.to_string();
let refresh_token = token_data
.get("refresh_token")
.and_then(|v| v.as_str())
.map(String::from);
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
Ok(OAuthTokenResponse {
access_token,
refresh_token,
expires_in,
})
}
/// Exchange an OAuth authorization code via the platform's token exchange proxy. /// Exchange an OAuth authorization code via the platform's token exchange proxy.
/// ///
/// Authenticated via the gateway auth token (Bearer header). The caller may /// Authenticated via the gateway auth token (Bearer header). The caller may
@@ -682,6 +748,7 @@ pub async fn exchange_via_proxy(
let client = reqwest::Client::builder() let client = reqwest::Client::builder()
.timeout(Duration::from_secs(60)) .timeout(Duration::from_secs(60))
.redirect(reqwest::redirect::Policy::none())
.build() .build()
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?; .map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
let mut params = vec![ let mut params = vec![
@@ -724,41 +791,350 @@ pub async fn exchange_via_proxy(
.json() .json()
.await .await
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?; .map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?;
oauth_token_response_from_json(token_data, request.access_token_field)
}
let access_token = token_data /// Refresh an OAuth access token via the platform's token refresh proxy.
.get(request.access_token_field) ///
.and_then(|v| v.as_str()) /// Authenticated via the gateway auth token (Bearer header). The caller may
.ok_or_else(|| { /// either rely on proxy-side secret lookup or forward a `client_secret` when
let fields: Vec<&str> = token_data /// the provider requires it.
.as_object() pub async fn refresh_token_via_proxy(
.map(|o| o.keys().map(|k| k.as_str()).collect()) request: ProxyRefreshTokenRequest<'_>,
.unwrap_or_default(); ) -> Result<OAuthTokenResponse, OAuthCallbackError> {
OAuthCallbackError::Io(format!( if request.gateway_token.is_empty() {
"No '{}' field in proxy response (fields present: {:?})", return Err(OAuthCallbackError::Io(
request.access_token_field, fields "Gateway auth token is required for proxy token refresh".to_string(),
)) ));
})? }
.to_string();
let refresh_token = token_data let refresh_url = format!("{}/oauth/refresh", request.proxy_url.trim_end_matches('/'));
.get("refresh_token") let client = reqwest::Client::builder()
.and_then(|v| v.as_str()) .timeout(Duration::from_secs(15))
.map(String::from); .redirect(reqwest::redirect::Policy::none())
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64()); .build()
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
Ok(OAuthTokenResponse { let mut params = vec![
access_token, ("refresh_token", request.refresh_token.to_string()),
refresh_token, ("token_url", request.token_url.to_string()),
expires_in, ("client_id", request.client_id.to_string()),
}) ];
if let Some(secret) = request.client_secret {
params.push(("client_secret", secret.to_string()));
}
if let Some(provider) = request.provider {
params.push(("provider", provider.to_string()));
}
let response = client
.post(&refresh_url)
.bearer_auth(request.gateway_token)
.form(&params)
.send()
.await
.map_err(|e| {
OAuthCallbackError::Io(format!("Token refresh proxy request failed: {}", e))
})?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(OAuthCallbackError::Io(format!(
"Token refresh proxy failed: {} - {}",
status, body
)));
}
let token_data: serde_json::Value = response
.json()
.await
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?;
oauth_token_response_from_json(token_data, "access_token")
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use axum::extract::{Form, State};
use axum::http::HeaderMap;
use axum::response::Redirect;
use axum::routing::post;
use axum::{Json, Router};
use serde_json::json;
use tokio::net::TcpListener;
use tokio::sync::{Mutex, oneshot};
use crate::cli::oauth_defaults::{ use crate::cli::oauth_defaults::{
builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html, builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html,
}; };
use crate::config::helpers::lock_env; use crate::config::helpers::lock_env;
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
#[derive(Clone, Debug, PartialEq, Eq)]
struct RecordedProxyRequest {
authorization: Option<String>,
form: HashMap<String, String>,
}
#[derive(Clone)]
struct MockProxyState {
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
exchange_redirect_target: String,
refresh_redirect_target: String,
}
struct MockProxyServer {
addr: SocketAddr,
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
shutdown_tx: Option<oneshot::Sender<()>>,
server_task: Option<tokio::task::JoinHandle<()>>,
}
impl MockProxyServer {
async fn start() -> Self {
async fn exchange_handler(
State(state): State<MockProxyState>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
state.requests.lock().await.push(RecordedProxyRequest {
authorization: headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(str::to_string),
form,
});
Json(json!({
"access_token": "proxy-access-token",
"refresh_token": "proxy-refresh-token",
"expires_in": 7200
}))
}
async fn refresh_handler(
State(state): State<MockProxyState>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
state.requests.lock().await.push(RecordedProxyRequest {
authorization: headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(str::to_string),
form,
});
Json(json!({
"access_token": "proxy-access-token",
"refresh_token": "proxy-refresh-token",
"expires_in": 7200
}))
}
async fn exchange_redirect_handler(State(state): State<MockProxyState>) -> Redirect {
Redirect::temporary(&state.exchange_redirect_target)
}
async fn refresh_redirect_handler(State(state): State<MockProxyState>) -> Redirect {
Redirect::temporary(&state.refresh_redirect_target)
}
let requests = Arc::new(Mutex::new(Vec::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind mock proxy");
let addr = listener.local_addr().expect("read mock proxy addr");
let exchange_redirect_target = format!("http://{addr}/oauth/exchange");
let refresh_redirect_target = format!("http://{addr}/oauth/refresh");
let app = Router::new()
.route("/oauth/exchange", post(exchange_handler))
.route("/oauth/refresh", post(refresh_handler))
.route("/redirect/oauth/exchange", post(exchange_redirect_handler))
.route("/redirect/oauth/refresh", post(refresh_redirect_handler))
.with_state(MockProxyState {
requests: Arc::clone(&requests),
exchange_redirect_target,
refresh_redirect_target,
});
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let server_task = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
Self {
addr,
requests,
shutdown_tx: Some(shutdown_tx),
server_task: Some(server_task),
}
}
fn base_url(&self) -> String {
format!("http://{}", self.addr)
}
fn redirecting_base_url(&self) -> String {
format!("{}/redirect", self.base_url())
}
async fn requests(&self) -> Vec<RecordedProxyRequest> {
self.requests.lock().await.clone()
}
async fn shutdown(mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
if let Some(task) = self.server_task.take() {
let _ = task.await;
}
}
}
impl Drop for MockProxyServer {
fn drop(&mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
if let Some(task) = self.server_task.take() {
task.abort();
}
}
}
#[test]
fn test_hosted_proxy_client_secret_suppresses_builtin_secret() {
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
let client_secret = Some(builtin.client_secret.to_string());
let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true);
assert_eq!(result, None);
}
#[test]
fn test_hosted_proxy_client_secret_preserves_explicit_secret() {
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
let client_secret = Some("hosted-server-secret".to_string());
let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true);
assert_eq!(result, client_secret);
}
#[tokio::test]
async fn test_refresh_token_via_proxy_sends_auth_and_form() {
let server = MockProxyServer::start().await;
let response = super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest {
proxy_url: &server.base_url(),
gateway_token: "gateway-test-token",
token_url: "https://oauth2.googleapis.com/token",
client_id: TEST_OAUTH_CLIENT_ID,
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
refresh_token: "refresh-token-123",
provider: Some("google"),
})
.await
.expect("proxy refresh succeeds");
assert_eq!(response.access_token, "proxy-access-token");
assert_eq!(
response.refresh_token.as_deref(),
Some("proxy-refresh-token")
);
assert_eq!(response.expires_in, Some(7200));
let requests = server.requests().await;
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].authorization.as_deref(),
Some("Bearer gateway-test-token")
);
assert_eq!(
requests[0].form.get("token_url").map(String::as_str),
Some("https://oauth2.googleapis.com/token")
);
assert_eq!(
requests[0].form.get("client_id").map(String::as_str),
Some(TEST_OAUTH_CLIENT_ID)
);
assert_eq!(
requests[0].form.get("client_secret").map(String::as_str),
Some(TEST_OAUTH_CLIENT_SECRET)
);
assert_eq!(
requests[0].form.get("refresh_token").map(String::as_str),
Some("refresh-token-123")
);
assert_eq!(
requests[0].form.get("provider").map(String::as_str),
Some("google")
);
server.shutdown().await;
}
#[tokio::test]
async fn test_exchange_via_proxy_does_not_follow_redirects() {
let server = MockProxyServer::start().await;
let error = match super::exchange_via_proxy(super::ProxyTokenExchangeRequest {
proxy_url: &server.redirecting_base_url(),
gateway_token: "gateway-test-token",
code: "auth-code-123",
redirect_uri: "http://localhost:3000/oauth/callback",
token_url: "https://oauth2.googleapis.com/token",
client_id: TEST_OAUTH_CLIENT_ID,
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
access_token_field: "access_token",
code_verifier: Some("code-verifier-123"),
extra_token_params: &HashMap::new(),
})
.await
{
Ok(_) => panic!("redirected proxy exchange should fail"),
Err(error) => error,
};
assert!(error.to_string().contains("307"));
assert!(server.requests().await.is_empty());
server.shutdown().await;
}
#[tokio::test]
async fn test_refresh_token_via_proxy_does_not_follow_redirects() {
let server = MockProxyServer::start().await;
let error = match super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest {
proxy_url: &server.redirecting_base_url(),
gateway_token: "gateway-test-token",
token_url: "https://oauth2.googleapis.com/token",
client_id: TEST_OAUTH_CLIENT_ID,
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
refresh_token: "refresh-token-123",
provider: Some("google"),
})
.await
{
Ok(_) => panic!("redirected proxy refresh should fail"),
Err(error) => error,
};
assert!(error.to_string().contains("307"));
assert!(server.requests().await.is_empty());
server.shutdown().await;
}
#[test] #[test]
fn test_is_loopback_host() { fn test_is_loopback_host() {
+19 -8
View File
@@ -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
+270 -16
View File
@@ -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());
}
} }
+5 -7
View File
@@ -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(),
+50
View File
@@ -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,
+9
View File
@@ -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,
+8
View File
@@ -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,
+64 -79
View File
@@ -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() {
@@ -891,24 +875,27 @@ 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
} }
/// 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 +1118,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 +1415,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.is_relay_channel(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 +1432,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 +1615,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))
@@ -3179,7 +3183,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(),
@@ -3284,7 +3288,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 +3314,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,13 +4185,13 @@ 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) // Check if already authenticated (team_id setting exists)
if self.is_relay_channel(name, user_id).await { if self.is_relay_channel(name, user_id).await {
return Ok(AuthResult::authenticated(name, ExtensionKind::ChannelRelay)); return Ok(AuthResult::authenticated(name, ExtensionKind::ChannelRelay));
} }
@@ -4233,19 +4237,9 @@ 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);
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 store
.get_setting(user_id, &team_id_key) .get_setting(user_id, &team_id_key)
@@ -4258,6 +4252,10 @@ impl ExtensionManager {
String::new() String::new()
}; };
if team_id.is_empty() {
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()?;
@@ -4367,11 +4365,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);
} }
@@ -4999,11 +4997,7 @@ impl ExtensionManager {
names.insert(server.token_secret_name()); names.insert(server.token_secret_name());
(names, Vec::new()) (names, Vec::new())
} }
ExtensionKind::ChannelRelay => { ExtensionKind::ChannelRelay => (std::collections::HashSet::new(), Vec::new()),
let mut names = std::collections::HashSet::new();
names.insert(format!("relay:{}:stream_token", name));
(names, Vec::new())
}
}; };
let allowed_fields: std::collections::HashSet<String> = let allowed_fields: std::collections::HashSet<String> =
@@ -5434,7 +5428,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 +5698,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 +7039,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);
} }
@@ -7862,19 +7858,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 +7876,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 +7885,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 +7950,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 +7963,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 +7978,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 +7990,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 +8005,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()),
+34
View File
@@ -1403,6 +1403,40 @@ impl Store {
Ok(counts) Ok(counts)
} }
/// Batch-load the most recent run status for multiple routines in a single query.
/// Uses a window function to pick only the latest run per routine.
#[cfg(feature = "postgres")]
pub async fn batch_get_last_run_status(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
if routine_ids.is_empty() {
return Ok(HashMap::new());
}
let conn = self.conn().await?;
let rows = conn
.query(
"SELECT DISTINCT ON (routine_id) routine_id, status
FROM routine_runs
WHERE routine_id = ANY($1)
ORDER BY routine_id, started_at DESC",
&[&routine_ids],
)
.await?;
let mut statuses = HashMap::new();
for row in rows {
let id: Uuid = row.get("routine_id");
let status_str: String = row.get("status");
if let std::result::Result::Ok(status) = status_str.parse::<RunStatus>() {
statuses.insert(id, status);
}
}
Ok(statuses)
}
/// Link a routine run to a dispatched job. /// Link a routine run to a dispatched job.
pub async fn link_routine_run_to_job( pub async fn link_routine_run_to_job(
&self, &self,
+2
View File
@@ -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"),
+7
View File
@@ -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![
+2
View File
@@ -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);
+1
View File
@@ -1898,6 +1898,7 @@ impl GeminiOauthProvider {
id, id,
name, name,
arguments: args, arguments: args,
reasoning: None,
}); });
} }
} }
+2
View File
@@ -596,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()
}) })
@@ -628,6 +629,7 @@ mod tests {
id: "call_1".to_string(), id: "call_1".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}), arguments: serde_json::json!({"q": "test"}),
reasoning: None,
}]; }];
let messages = vec![ let messages = vec![
ChatMessage::user("Search"), ChatMessage::user("Search"),
+75 -2
View File
@@ -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]
+5
View File
@@ -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 =
+8
View File
@@ -231,6 +231,10 @@ pub struct ToolCall {
pub id: String, pub id: String,
pub name: String, pub name: String,
pub arguments: serde_json::Value, pub arguments: serde_json::Value,
/// Optional reasoning for why this tool was chosen — supplied by the provider
/// or derived from the shared response content as a fallback.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
} }
/// Generate a tool-call ID that satisfies all providers. /// Generate a tool-call ID that satisfies all providers.
@@ -637,6 +641,7 @@ mod tests {
id: "call_1".to_string(), id: "call_1".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}; };
let mut messages = vec![ let mut messages = vec![
ChatMessage::user("hello"), ChatMessage::user("hello"),
@@ -680,6 +685,7 @@ mod tests {
id: "call_1".to_string(), id: "call_1".to_string(),
name: "echo".to_string(), name: "echo".to_string(),
arguments: serde_json::json!({}), arguments: serde_json::json!({}),
reasoning: None,
}; };
let mut messages = vec![ let mut messages = vec![
ChatMessage::user("test"), ChatMessage::user("test"),
@@ -705,11 +711,13 @@ mod tests {
id: "call_sel_1".to_string(), id: "call_sel_1".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}), arguments: serde_json::json!({"q": "test"}),
reasoning: None,
}; };
let tc2 = ToolCall { let tc2 = ToolCall {
id: "call_sel_2".to_string(), id: "call_sel_2".to_string(),
name: "http".to_string(), name: "http".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}), arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
}; };
let mut messages = vec![ let mut messages = vec![
ChatMessage::system("You are a helpful assistant."), ChatMessage::system("You are a helpful assistant."),
+85 -12
View File
@@ -525,17 +525,35 @@ impl Reasoning {
let response = self.llm.complete_with_tools(request).await?; let response = self.llm.complete_with_tools(request).await?;
let reasoning = response.content.unwrap_or_default(); let shared_reasoning = response
.content
.map(|c| {
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
})
.unwrap_or_default();
let selections: Vec<ToolSelection> = response let selections: Vec<ToolSelection> = response
.tool_calls .tool_calls
.into_iter() .into_iter()
.map(|tool_call| ToolSelection { .map(|tool_call| {
tool_name: tool_call.name, // Prefer per-tool reasoning if the provider supplied it,
parameters: tool_call.arguments, // otherwise fall back to the shared response content.
reasoning: reasoning.clone(), let rationale = tool_call
alternatives: vec![], .reasoning
tool_call_id: tool_call.id, .map(|r| {
let pre_truncated = truncate_at_tool_tags(&r);
clean_response(&pre_truncated)
})
.filter(|r| !r.trim().is_empty())
.unwrap_or_else(|| shared_reasoning.clone());
ToolSelection {
tool_name: tool_call.name,
parameters: tool_call.arguments,
reasoning: rationale,
alternatives: vec![],
tool_call_id: tool_call.id,
}
}) })
.collect(); .collect();
@@ -664,13 +682,36 @@ Respond in JSON format:
// If there were tool calls, return them for execution // If there were tool calls, return them for execution
if !response.tool_calls.is_empty() { if !response.tool_calls.is_empty() {
let narrative = response.content.map(|c| {
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
});
// Populate per-tool reasoning from the shared narrative when the
// provider did not supply per-tool rationale.
let tool_calls: Vec<ToolCall> = response
.tool_calls
.into_iter()
.map(|mut tc| {
if tc.reasoning.as_ref().is_none_or(|r| r.trim().is_empty()) {
tc.reasoning = narrative.as_ref().filter(|n| !n.is_empty()).cloned();
} else {
// Clean provider-supplied per-tool reasoning the same way
// we clean the shared narrative (strip thinking/tool tags).
tc.reasoning = tc
.reasoning
.map(|r| {
let pre_truncated = truncate_at_tool_tags(&r);
clean_response(&pre_truncated)
})
.filter(|r| !r.trim().is_empty());
}
tc
})
.collect();
return Ok(RespondOutput { return Ok(RespondOutput {
result: RespondResult::ToolCalls { result: RespondResult::ToolCalls {
tool_calls: response.tool_calls, tool_calls,
content: response.content.map(|c| { content: narrative,
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
}),
}, },
usage, usage,
}); });
@@ -1350,6 +1391,7 @@ fn recover_tool_calls_from_content(
), ),
name: name.to_string(), name: name.to_string(),
arguments, arguments,
reasoning: None,
}); });
continue; continue;
} }
@@ -1364,6 +1406,7 @@ fn recover_tool_calls_from_content(
), ),
name: name.to_string(), name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()), arguments: serde_json::Value::Object(Default::default()),
reasoning: None,
}); });
} }
} }
@@ -1401,6 +1444,7 @@ fn recover_tool_calls_from_content(
), ),
name: name.to_string(), name: name.to_string(),
arguments, arguments,
reasoning: None,
}); });
remaining = &args_start[bracket_end + 1..]; remaining = &args_start[bracket_end + 1..];
continue; continue;
@@ -1412,6 +1456,7 @@ fn recover_tool_calls_from_content(
id: super::provider::generate_tool_call_id(calls.len(), RECOVERED_TOOL_CALL_SEED), id: super::provider::generate_tool_call_id(calls.len(), RECOVERED_TOOL_CALL_SEED),
name: name.to_string(), name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()), arguments: serde_json::Value::Object(Default::default()),
reasoning: None,
}); });
remaining = after_name; remaining = after_name;
} }
@@ -3145,4 +3190,32 @@ That's my plan."#;
"Text <function_call>{}</function_call> middle " "Text <function_call>{}</function_call> middle "
); );
} }
/// Verify that reasoning normalization strips thinking tags and tool tags
/// from per-tool reasoning, matching the cleaning applied to shared reasoning.
#[test]
fn test_reasoning_normalization_strips_thinking_tags() {
let raw = "<thinking>Let me consider...</thinking>Search memory for prior context";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(!cleaned.contains("<thinking>"));
assert!(cleaned.contains("Search memory"));
}
#[test]
fn test_reasoning_normalization_strips_tool_tags() {
let raw = "Calling search <tool_call>{\"name\": \"search\"}";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(!cleaned.contains("<tool_call>"));
assert!(cleaned.contains("Calling search"));
}
#[test]
fn test_reasoning_normalization_empty_after_cleaning() {
let raw = "<thinking>internal only</thinking>";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(cleaned.trim().is_empty());
}
} }
+7
View File
@@ -490,6 +490,7 @@ fn extract_response(
id: tc.id.clone(), id: tc.id.clone(),
name: tc.function.name.clone(), name: tc.function.name.clone(),
arguments: tc.function.arguments.clone(), arguments: tc.function.arguments.clone(),
reasoning: None,
}); });
} }
// Reasoning and Image variants are not mapped to IronClaw types // Reasoning and Image variants are not mapped to IronClaw types
@@ -880,6 +881,7 @@ mod tests {
id: "Xt7mK9pQ2".to_string(), id: "Xt7mK9pQ2".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}), arguments: serde_json::json!({"query": "test"}),
reasoning: None,
}; };
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]); let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
let messages = vec![msg]; let messages = vec![msg];
@@ -997,6 +999,7 @@ mod tests {
id: "".to_string(), id: "".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}), arguments: serde_json::json!({"query": "test"}),
reasoning: None,
}; };
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])]; let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
let (_preamble, history) = convert_messages(&messages); let (_preamble, history) = convert_messages(&messages);
@@ -1028,6 +1031,7 @@ mod tests {
id: " ".to_string(), id: " ".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}), arguments: serde_json::json!({"query": "test"}),
reasoning: None,
}; };
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])]; let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
let (_preamble, history) = convert_messages(&messages); let (_preamble, history) = convert_messages(&messages);
@@ -1061,6 +1065,7 @@ mod tests {
id: "".to_string(), id: "".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}), arguments: serde_json::json!({"query": "test"}),
reasoning: None,
}; };
let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]); let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
let tool_result_msg = ChatMessage { let tool_result_msg = ChatMessage {
@@ -1380,11 +1385,13 @@ mod tests {
id: "call_a".to_string(), id: "call_a".to_string(),
name: "search".to_string(), name: "search".to_string(),
arguments: serde_json::json!({"q": "rust"}), arguments: serde_json::json!({"q": "rust"}),
reasoning: None,
}; };
let tc2 = IronToolCall { let tc2 = IronToolCall {
id: "call_b".to_string(), id: "call_b".to_string(),
name: "fetch".to_string(), name: "fetch".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}), arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
}; };
let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]); let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]);
let result_a = ChatMessage::tool_result("call_a", "search", "search results"); let result_a = ChatMessage::tool_result("call_a", "search", "search results");
+1
View File
@@ -611,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));
+29 -14
View File
@@ -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.
@@ -277,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
@@ -295,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
@@ -309,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
@@ -324,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
@@ -344,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
@@ -390,9 +405,9 @@ async fn job_event_handler(
}; };
if user_id.is_empty() { if user_id.is_empty() {
let _ = tx.send((job_id, String::new(), sse_event)); let _ = tx.send((job_id, String::new(), app_event));
} else { } else {
let _ = tx.send((job_id, user_id, sse_event)); let _ = tx.send((job_id, user_id, app_event));
} }
} }
@@ -817,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,
@@ -872,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),
@@ -918,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 --
+2 -2
View File
@@ -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,
} }
+3 -3
View File
@@ -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);
+1 -5
View File
@@ -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>,
+141
View File
@@ -418,6 +418,7 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
let oauth = auth.oauth.as_ref()?; let oauth = auth.oauth.as_ref()?;
let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name); let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name);
let exchange_proxy_url = crate::cli::oauth_defaults::exchange_proxy_url();
let client_id = oauth let client_id = oauth
.client_id .client_id
@@ -440,11 +441,21 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
.and_then(|env| std::env::var(env).ok()) .and_then(|env| std::env::var(env).ok())
}) })
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string())); .or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
let client_secret = crate::cli::oauth_defaults::hosted_proxy_client_secret(
&client_secret,
builtin.as_ref(),
exchange_proxy_url.is_some(),
);
let gateway_token = crate::config::helpers::env_or_override("GATEWAY_AUTH_TOKEN")
.map(|token| token.trim().to_string())
.filter(|token| !token.is_empty());
Some(OAuthRefreshConfig { Some(OAuthRefreshConfig {
token_url: oauth.token_url.clone(), token_url: oauth.token_url.clone(),
client_id, client_id,
client_secret, client_secret,
exchange_proxy_url,
gateway_token,
secret_name: auth.secret_name.clone(), secret_name: auth.secret_name.clone(),
provider: auth.provider.clone(), provider: auth.provider.clone(),
}) })
@@ -711,9 +722,44 @@ mod tests {
use tempfile::TempDir; use tempfile::TempDir;
use crate::config::helpers::lock_env;
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET}; use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools}; use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools};
/// Restores a test-scoped env var override on drop.
struct EnvVarGuard {
key: String,
previous: Option<String>,
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
// SAFETY: Tests use lock_env() to serialize environment access.
unsafe {
if let Some(ref value) = self.previous {
std::env::set_var(&self.key, value);
} else {
std::env::remove_var(&self.key);
}
}
}
}
fn set_env_var(key: &str, value: Option<&str>) -> EnvVarGuard {
let previous = std::env::var(key).ok();
// SAFETY: Tests use lock_env() to serialize environment access.
unsafe {
match value {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
}
EnvVarGuard {
key: key.to_string(),
previous,
}
}
#[test] #[test]
fn wit_version_compat_none_is_ok() { fn wit_version_compat_none_is_ok() {
// Pre-versioning extensions (no wit_version declared) should always pass // Pre-versioning extensions (no wit_version declared) should always pass
@@ -871,6 +917,8 @@ mod tests {
config.client_secret, config.client_secret,
Some(TEST_OAUTH_CLIENT_SECRET.to_string()) Some(TEST_OAUTH_CLIENT_SECRET.to_string())
); );
assert_eq!(config.exchange_proxy_url, None);
assert_eq!(config.gateway_token, None);
assert_eq!(config.secret_name, "google_oauth_token"); assert_eq!(config.secret_name, "google_oauth_token");
assert_eq!(config.provider, Some("google".to_string())); assert_eq!(config.provider, Some("google".to_string()));
} }
@@ -931,6 +979,10 @@ mod tests {
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
}; };
let _guard = lock_env();
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", None);
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
// google_oauth_token should fall back to built-in credentials // google_oauth_token should fall back to built-in credentials
let caps = CapabilitiesFile { let caps = CapabilitiesFile {
auth: Some(AuthCapabilitySchema { auth: Some(AuthCapabilitySchema {
@@ -952,6 +1004,95 @@ mod tests {
let config = config.unwrap(); let config = config.unwrap();
assert!(!config.client_id.is_empty()); assert!(!config.client_id.is_empty());
assert!(config.client_secret.is_some()); assert!(config.client_secret.is_some());
assert_eq!(config.exchange_proxy_url, None);
assert_eq!(config.gateway_token, None);
}
#[test]
fn test_resolve_oauth_refresh_config_hosted_proxy_populates_env_and_suppresses_builtin_secret()
{
use crate::tools::wasm::capabilities_schema::{
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
};
let _guard = lock_env();
let _proxy_guard = set_env_var(
"IRONCLAW_OAUTH_EXCHANGE_URL",
Some("https://compose-api.example.com"),
);
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
let _client_id_guard =
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
let caps = CapabilitiesFile {
auth: Some(AuthCapabilitySchema {
secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()),
oauth: Some(OAuthConfigSchema {
authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
token_url: "https://oauth2.googleapis.com/token".to_string(),
client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()),
..Default::default()
}),
..Default::default()
}),
..Default::default()
};
let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config");
assert_eq!(config.client_id, "hosted-google-client-id");
assert_eq!(config.client_secret, None);
assert_eq!(
config.exchange_proxy_url.as_deref(),
Some("https://compose-api.example.com")
);
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
}
#[test]
fn test_resolve_oauth_refresh_config_hosted_proxy_preserves_explicit_secret() {
use crate::tools::wasm::capabilities_schema::{
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
};
let _guard = lock_env();
let _proxy_guard = set_env_var(
"IRONCLAW_OAUTH_EXCHANGE_URL",
Some("https://compose-api.example.com"),
);
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
let _client_id_guard =
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
let _client_secret_guard =
set_env_var("GOOGLE_OAUTH_CLIENT_SECRET", Some("hosted-server-secret"));
let caps = CapabilitiesFile {
auth: Some(AuthCapabilitySchema {
secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()),
oauth: Some(OAuthConfigSchema {
authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
token_url: "https://oauth2.googleapis.com/token".to_string(),
client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()),
client_secret_env: Some("GOOGLE_OAUTH_CLIENT_SECRET".to_string()),
..Default::default()
}),
..Default::default()
}),
..Default::default()
};
let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config");
assert_eq!(config.client_id, "hosted-google-client-id");
assert_eq!(
config.client_secret.as_deref(),
Some("hosted-server-secret")
);
assert_eq!(
config.exchange_proxy_url.as_deref(),
Some("https://compose-api.example.com")
);
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
} }
// --------------------------------------------------------------- // ---------------------------------------------------------------
+443 -38
View File
@@ -19,7 +19,7 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
use crate::context::JobContext; use crate::context::JobContext;
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor}; use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
use crate::safety::LeakDetector; use crate::safety::LeakDetector;
use crate::secrets::SecretsStore; use crate::secrets::{DecryptedSecret, SecretsStore};
use crate::tools::tool::{Tool, ToolError, ToolOutput}; use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::wasm::capabilities::Capabilities; use crate::tools::wasm::capabilities::Capabilities;
use crate::tools::wasm::credential_injector::{ use crate::tools::wasm::credential_injector::{
@@ -44,6 +44,7 @@ wasmtime::component::bindgen!({
}); });
// Alias the export interface types for convenience. // Alias the export interface types for convenience.
use crate::cli::oauth_defaults;
use exports::near::agent::tool as wit_tool; use exports::near::agent::tool as wit_tool;
/// Configuration needed to refresh an expired OAuth access token. /// Configuration needed to refresh an expired OAuth access token.
@@ -59,6 +60,10 @@ pub struct OAuthRefreshConfig {
pub client_id: String, pub client_id: String,
/// OAuth client_secret (optional, some providers use PKCE without a secret). /// OAuth client_secret (optional, some providers use PKCE without a secret).
pub client_secret: Option<String>, pub client_secret: Option<String>,
/// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080").
pub exchange_proxy_url: Option<String>,
/// Gateway auth token for authenticating with the hosted OAuth proxy.
pub gateway_token: Option<String>,
/// Secret name of the access token (e.g., "google_oauth_token"). /// Secret name of the access token (e.g., "google_oauth_token").
/// The refresh token lives at `{secret_name}_refresh_token`. /// The refresh token lives at `{secret_name}_refresh_token`.
pub secret_name: String, pub secret_name: String,
@@ -1210,6 +1215,53 @@ async fn refresh_oauth_token(
user_id: &str, user_id: &str,
config: &OAuthRefreshConfig, config: &OAuthRefreshConfig,
) -> bool { ) -> bool {
let refresh_name = format!("{}_refresh_token", config.secret_name);
if let Some(proxy_url) = config.exchange_proxy_url.as_deref() {
let Some(gateway_token) = config.gateway_token.as_deref() else {
tracing::warn!(
"OAuth refresh proxy is configured, but no gateway auth token is available"
);
return false;
};
// In hosted mode, the configured exchange proxy owns the outbound token
// refresh and validation policy for the provider token_url. Direct-mode
// HTTPS/private-IP checks remain in place for self-hosted refreshes below.
let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await {
Some(secret) => secret,
None => return false,
};
let token_response = match oauth_defaults::refresh_token_via_proxy(
oauth_defaults::ProxyRefreshTokenRequest {
proxy_url,
gateway_token,
token_url: &config.token_url,
client_id: &config.client_id,
client_secret: config.client_secret.as_deref(),
refresh_token: refresh_secret.expose(),
provider: config.provider.as_deref(),
},
)
.await
{
Ok(response) => response,
Err(error) => {
tracing::warn!(error = %error, "OAuth token refresh via proxy failed");
return false;
}
};
return persist_refreshed_oauth_tokens(
store,
user_id,
config,
&refresh_name,
token_response,
)
.await;
}
// SSRF defense: token_url comes from the tool's capabilities file. // SSRF defense: token_url comes from the tool's capabilities file.
if !config.token_url.starts_with("https://") { if !config.token_url.starts_with("https://") {
tracing::warn!( tracing::warn!(
@@ -1227,19 +1279,6 @@ async fn refresh_oauth_token(
return false; return false;
} }
let refresh_name = format!("{}_refresh_token", config.secret_name);
let refresh_secret = match store.get_decrypted(user_id, &refresh_name).await {
Ok(s) => s,
Err(e) => {
tracing::debug!(
secret_name = %refresh_name,
error = %e,
"No refresh token available, skipping token refresh"
);
return false;
}
};
let client = match reqwest::Client::builder() let client = match reqwest::Client::builder()
.timeout(Duration::from_secs(15)) .timeout(Duration::from_secs(15))
.redirect(reqwest::redirect::Policy::none()) .redirect(reqwest::redirect::Policy::none())
@@ -1252,6 +1291,10 @@ async fn refresh_oauth_token(
} }
}; };
let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await {
Some(secret) => secret,
None => return false,
};
let mut params = vec![ let mut params = vec![
("grant_type", "refresh_token".to_string()), ("grant_type", "refresh_token".to_string()),
("refresh_token", refresh_secret.expose().to_string()), ("refresh_token", refresh_secret.expose().to_string()),
@@ -1287,22 +1330,55 @@ async fn refresh_oauth_token(
return false; return false;
} }
}; };
let token_response = match token_data.get("access_token").and_then(|v| v.as_str()) {
let new_access_token = match token_data.get("access_token").and_then(|v| v.as_str()) { Some(access_token) => oauth_defaults::OAuthTokenResponse {
Some(t) => t, access_token: access_token.to_string(),
refresh_token: token_data
.get("refresh_token")
.and_then(|v| v.as_str())
.map(str::to_string),
expires_in: token_data.get("expires_in").and_then(|v| v.as_u64()),
},
None => { None => {
tracing::warn!("Token refresh response missing access_token field"); tracing::warn!("Token refresh response missing access_token field");
return false; return false;
} }
}; };
// Store the new access token with expiry persist_refreshed_oauth_tokens(store, user_id, config, &refresh_name, token_response).await
}
async fn load_oauth_refresh_secret(
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
refresh_name: &str,
) -> Option<DecryptedSecret> {
match store.get_decrypted(user_id, refresh_name).await {
Ok(secret) => Some(secret),
Err(error) => {
tracing::debug!(
secret_name = %refresh_name,
error = %error,
"No refresh token available, skipping token refresh"
);
None
}
}
}
async fn persist_refreshed_oauth_tokens(
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
config: &OAuthRefreshConfig,
refresh_name: &str,
token_response: oauth_defaults::OAuthTokenResponse,
) -> bool {
let mut access_params = let mut access_params =
crate::secrets::CreateSecretParams::new(&config.secret_name, new_access_token); crate::secrets::CreateSecretParams::new(&config.secret_name, &token_response.access_token);
if let Some(ref provider) = config.provider { if let Some(ref provider) = config.provider {
access_params = access_params.with_provider(provider); access_params = access_params.with_provider(provider);
} }
if let Some(expires_in) = token_data.get("expires_in").and_then(|v| v.as_u64()) { if let Some(expires_in) = token_response.expires_in {
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in as i64); let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in as i64);
access_params = access_params.with_expiry(expires_at); access_params = access_params.with_expiry(expires_at);
} }
@@ -1312,10 +1388,8 @@ async fn refresh_oauth_token(
return false; return false;
} }
// Store rotated refresh token if the provider sent a new one if let Some(new_refresh) = token_response.refresh_token.as_deref() {
if let Some(new_refresh) = token_data.get("refresh_token").and_then(|v| v.as_str()) { let mut refresh_params = crate::secrets::CreateSecretParams::new(refresh_name, new_refresh);
let mut refresh_params =
crate::secrets::CreateSecretParams::new(&refresh_name, new_refresh);
if let Some(ref provider) = config.provider { if let Some(ref provider) = config.provider {
refresh_params = refresh_params.with_provider(provider); refresh_params = refresh_params.with_provider(provider);
} }
@@ -1664,9 +1738,18 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use async_trait::async_trait; use async_trait::async_trait;
use axum::extract::{Form, State};
use axum::http::HeaderMap;
use axum::routing::post;
use axum::{Json, Router};
use serde_json::json;
use tokio::net::TcpListener;
use tokio::sync::{Mutex as AsyncMutex, oneshot};
use uuid::Uuid; use uuid::Uuid;
use crate::context::JobContext; use crate::context::JobContext;
@@ -1756,6 +1839,95 @@ mod tests {
} }
} }
#[derive(Clone, Debug, PartialEq, Eq)]
struct RecordedProxyRequest {
authorization: Option<String>,
form: HashMap<String, String>,
}
struct MockProxyServer {
addr: SocketAddr,
requests: Arc<AsyncMutex<Vec<RecordedProxyRequest>>>,
shutdown_tx: Option<oneshot::Sender<()>>,
server_task: Option<tokio::task::JoinHandle<()>>,
}
impl MockProxyServer {
async fn start() -> Self {
async fn refresh_handler(
State(requests): State<Arc<AsyncMutex<Vec<RecordedProxyRequest>>>>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
requests.lock().await.push(RecordedProxyRequest {
authorization: headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(str::to_string),
form,
});
Json(json!({
"access_token": "mock-refreshed-access-token",
"refresh_token": "mock-rotated-refresh-token",
"expires_in": 3600
}))
}
let requests = Arc::new(AsyncMutex::new(Vec::new()));
let app = Router::new()
.route("/oauth/refresh", post(refresh_handler))
.with_state(Arc::clone(&requests));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind mock proxy");
let addr = listener.local_addr().expect("read mock proxy addr");
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let server_task = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
Self {
addr,
requests,
shutdown_tx: Some(shutdown_tx),
server_task: Some(server_task),
}
}
fn base_url(&self) -> String {
format!("http://{}", self.addr)
}
async fn requests(&self) -> Vec<RecordedProxyRequest> {
self.requests.lock().await.clone()
}
async fn shutdown(mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
if let Some(task) = self.server_task.take() {
let _ = task.await;
}
}
}
impl Drop for MockProxyServer {
fn drop(&mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
if let Some(task) = self.server_task.take() {
task.abort();
}
}
}
#[test] #[test]
fn test_wrapper_creation() { fn test_wrapper_creation() {
// This test verifies the runtime can be created // This test verifies the runtime can be created
@@ -2094,8 +2266,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_bearer() { async fn test_resolve_host_credentials_bearer() {
use std::collections::HashMap;
use crate::secrets::{ use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
}; };
@@ -2141,8 +2311,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_owner_scope_bearer() { async fn test_resolve_host_credentials_owner_scope_bearer() {
use std::collections::HashMap;
use crate::secrets::{ use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
}; };
@@ -2188,8 +2356,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_execute_resolves_host_credentials_from_owner_scope_context() { async fn test_execute_resolves_host_credentials_from_owner_scope_context() {
use std::collections::HashMap;
use crate::secrets::{CredentialLocation, CredentialMapping}; use crate::secrets::{CredentialLocation, CredentialMapping};
use crate::tools::wasm::capabilities::HttpCapability; use crate::tools::wasm::capabilities::HttpCapability;
@@ -2239,8 +2405,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_missing_secret() { async fn test_resolve_host_credentials_missing_secret() {
use std::collections::HashMap;
use crate::secrets::{CredentialLocation, CredentialMapping}; use crate::secrets::{CredentialLocation, CredentialMapping};
use crate::tools::wasm::capabilities::HttpCapability; use crate::tools::wasm::capabilities::HttpCapability;
use crate::tools::wasm::wrapper::resolve_host_credentials; use crate::tools::wasm::wrapper::resolve_host_credentials;
@@ -2272,8 +2436,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_skips_refresh_when_not_expired() { async fn test_resolve_host_credentials_skips_refresh_when_not_expired() {
use std::collections::HashMap;
use crate::secrets::{ use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
}; };
@@ -2315,6 +2477,8 @@ mod tests {
token_url: "https://oauth2.googleapis.com/token".to_string(), token_url: "https://oauth2.googleapis.com/token".to_string(),
client_id: TEST_OAUTH_CLIENT_ID.to_string(), client_id: TEST_OAUTH_CLIENT_ID.to_string(),
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()), client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
exchange_proxy_url: None,
gateway_token: None,
secret_name: "google_oauth_token".to_string(), secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()), provider: Some("google".to_string()),
}; };
@@ -2331,8 +2495,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_skips_refresh_no_config() { async fn test_resolve_host_credentials_skips_refresh_no_config() {
use std::collections::HashMap;
use crate::secrets::{ use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
}; };
@@ -2376,8 +2538,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resolve_host_credentials_skips_refresh_no_expires_at() { async fn test_resolve_host_credentials_skips_refresh_no_expires_at() {
use std::collections::HashMap;
use crate::secrets::{ use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
}; };
@@ -2417,6 +2577,8 @@ mod tests {
token_url: "https://oauth2.googleapis.com/token".to_string(), token_url: "https://oauth2.googleapis.com/token".to_string(),
client_id: TEST_OAUTH_CLIENT_ID.to_string(), client_id: TEST_OAUTH_CLIENT_ID.to_string(),
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()), client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
exchange_proxy_url: None,
gateway_token: None,
secret_name: "google_oauth_token".to_string(), secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()), provider: Some("google".to_string()),
}; };
@@ -2431,6 +2593,249 @@ mod tests {
); );
} }
#[tokio::test]
async fn test_resolve_host_credentials_refreshes_via_proxy_without_direct_token_url_validation()
{
use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
};
use crate::tools::wasm::capabilities::HttpCapability;
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
let proxy = MockProxyServer::start().await;
let store = test_secrets_store();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token", "expired-access-token")
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
)
.await
.unwrap();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
)
.await
.unwrap();
let mut credentials = HashMap::new();
credentials.insert(
"google_oauth_token".to_string(),
CredentialMapping {
secret_name: "google_oauth_token".to_string(),
location: CredentialLocation::AuthorizationBearer,
host_patterns: vec!["www.googleapis.com".to_string()],
},
);
let caps = Capabilities {
http: Some(HttpCapability {
credentials,
..Default::default()
}),
..Default::default()
};
let oauth_config = OAuthRefreshConfig {
token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(),
client_id: "hosted-google-client-id".to_string(),
client_secret: None,
exchange_proxy_url: Some(proxy.base_url()),
gateway_token: Some("gateway-test-token".to_string()),
secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()),
};
let resolved =
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
assert_eq!(resolved.len(), 1);
assert_eq!(
resolved[0].headers.get("Authorization"),
Some(&"Bearer mock-refreshed-access-token".to_string())
);
let access_secret = store.get("user1", "google_oauth_token").await.unwrap();
assert!(
access_secret
.expires_at
.expect("refreshed access token expiry")
> chrono::Utc::now()
);
let access_value = store
.get_decrypted("user1", "google_oauth_token")
.await
.unwrap();
assert_eq!(access_value.expose(), "mock-refreshed-access-token");
let refresh_value = store
.get_decrypted("user1", "google_oauth_token_refresh_token")
.await
.unwrap();
assert_eq!(refresh_value.expose(), "mock-rotated-refresh-token");
let requests = proxy.requests().await;
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].authorization.as_deref(),
Some("Bearer gateway-test-token")
);
assert_eq!(
requests[0].form.get("client_id").map(String::as_str),
Some("hosted-google-client-id")
);
assert_eq!(
requests[0].form.get("token_url").map(String::as_str),
Some("http://127.0.0.1:9/provider-token-endpoint")
);
assert_eq!(
requests[0].form.get("refresh_token").map(String::as_str),
Some("stored-refresh-token")
);
assert_eq!(
requests[0].form.get("provider").map(String::as_str),
Some("google")
);
assert!(!requests[0].form.contains_key("client_secret"));
proxy.shutdown().await;
}
#[tokio::test]
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() {
use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
};
use crate::tools::wasm::capabilities::HttpCapability;
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
let store = RecordingSecretsStore::new();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token", "expired-access-token")
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
)
.await
.unwrap();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
)
.await
.unwrap();
let mut credentials = HashMap::new();
credentials.insert(
"google_oauth_token".to_string(),
CredentialMapping {
secret_name: "google_oauth_token".to_string(),
location: CredentialLocation::AuthorizationBearer,
host_patterns: vec!["www.googleapis.com".to_string()],
},
);
let caps = Capabilities {
http: Some(HttpCapability {
credentials,
..Default::default()
}),
..Default::default()
};
let oauth_config = OAuthRefreshConfig {
token_url: "https://oauth2.googleapis.com/token".to_string(),
client_id: "hosted-google-client-id".to_string(),
client_secret: None,
exchange_proxy_url: Some("https://compose-api.example.com".to_string()),
gateway_token: None,
secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()),
};
let resolved =
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
assert!(resolved.is_empty());
let lookups = store.decrypted_lookups();
assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string())));
assert!(!lookups.contains(&(
"user1".to_string(),
"google_oauth_token_refresh_token".to_string(),
)));
}
#[tokio::test]
async fn test_resolve_host_credentials_skips_refresh_token_lookup_for_invalid_direct_token_url()
{
use crate::secrets::{
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
};
use crate::tools::wasm::capabilities::HttpCapability;
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
let store = RecordingSecretsStore::new();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token", "expired-access-token")
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
)
.await
.unwrap();
store
.create(
"user1",
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
)
.await
.unwrap();
let mut credentials = HashMap::new();
credentials.insert(
"google_oauth_token".to_string(),
CredentialMapping {
secret_name: "google_oauth_token".to_string(),
location: CredentialLocation::AuthorizationBearer,
host_patterns: vec!["www.googleapis.com".to_string()],
},
);
let caps = Capabilities {
http: Some(HttpCapability {
credentials,
..Default::default()
}),
..Default::default()
};
let oauth_config = OAuthRefreshConfig {
token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(),
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
exchange_proxy_url: None,
gateway_token: None,
secret_name: "google_oauth_token".to_string(),
provider: Some("google".to_string()),
};
let resolved =
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
assert!(resolved.is_empty());
let lookups = store.decrypted_lookups();
assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string())));
assert!(!lookups.contains(&(
"user1".to_string(),
"google_oauth_token_refresh_token".to_string(),
)));
}
#[test] #[test]
fn test_is_private_ip_v4() { fn test_is_private_ip_v4() {
use std::net::IpAddr; use std::net::IpAddr;
+20 -8
View File
@@ -111,10 +111,23 @@ impl Tunnel for CloudflareTunnel {
} }
} }
// Drain stderr in the background to prevent SIGPIPE/buffer stalls. if let Ok(mut guard) = self.url.write() {
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} }); *guard = Some(public_url.clone());
}
// Drain stdout silently. // We took ownership of cloudflared's stderr pipe above to parse the URL.
// cloudflared continues writing logs for its entire lifetime. If we drop
// the reader, the pipe closes and cloudflared gets SIGPIPE on its next
// write. We can't just store the reader without reading — the OS pipe
// buffer fills up and cloudflared blocks. So we drain it in a background
// task. The task exits naturally when cloudflared is killed (EOF).
let drain_handle = tokio::spawn(async move {
while let Ok(Some(line)) = reader.next_line().await {
tracing::trace!("cloudflared: {line}");
}
});
// Drain stdout silently to prevent SIGPIPE/buffer stalls.
if let Some(stdout) = stdout { if let Some(stdout) = stdout {
tokio::spawn(async move { tokio::spawn(async move {
let mut out_reader = tokio::io::BufReader::new(stdout).lines(); let mut out_reader = tokio::io::BufReader::new(stdout).lines();
@@ -122,12 +135,11 @@ impl Tunnel for CloudflareTunnel {
}); });
} }
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await; let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child }); *guard = Some(TunnelProcess {
child,
_pipe_drain: Some(drain_handle),
});
Ok(public_url) Ok(public_url)
} }
+18 -5
View File
@@ -73,6 +73,7 @@ impl Tunnel for CustomTunnel {
let stderr = child.stderr.take(); let stderr = child.stderr.take();
let mut public_url = format!("http://{local_host}:{local_port}"); let mut public_url = format!("http://{local_host}:{local_port}");
let mut drain_handle: Option<tokio::task::JoinHandle<()>> = None;
if self.url_pattern.is_some() if self.url_pattern.is_some()
&& let Some(stdout) = stdout && let Some(stdout) = stdout
@@ -103,17 +104,26 @@ impl Tunnel for CustomTunnel {
Err(_) => {} Err(_) => {}
} }
} }
// Drain remaining stdout to prevent SIGPIPE/buffer stalls. // We took ownership of the process's stdout pipe above to parse the
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} }); // URL. The process may continue writing to stdout for its lifetime.
// If we drop the reader, the pipe closes and the process gets SIGPIPE.
// We can't just store the reader without reading — the OS pipe buffer
// fills up and the process blocks. So we drain it in a background task.
// The task exits naturally when the process is killed (EOF).
drain_handle = Some(tokio::spawn(async move {
while let Ok(Some(line)) = reader.next_line().await {
tracing::trace!("custom-tunnel: {line}");
}
}));
} else if let Some(stdout) = stdout { } else if let Some(stdout) = stdout {
// No url_pattern: still drain stdout to prevent pipe stalls. // No url_pattern: still drain stdout to prevent SIGPIPE/buffer stalls.
tokio::spawn(async move { tokio::spawn(async move {
let mut reader = tokio::io::BufReader::new(stdout).lines(); let mut reader = tokio::io::BufReader::new(stdout).lines();
while let Ok(Some(_)) = reader.next_line().await {} while let Ok(Some(_)) = reader.next_line().await {}
}); });
} }
// Drain stderr silently. // Drain stderr to prevent SIGPIPE/buffer stalls.
if let Some(stderr) = stderr { if let Some(stderr) = stderr {
tokio::spawn(async move { tokio::spawn(async move {
let mut reader = tokio::io::BufReader::new(stderr).lines(); let mut reader = tokio::io::BufReader::new(stderr).lines();
@@ -126,7 +136,10 @@ impl Tunnel for CustomTunnel {
} }
let mut guard = self.proc.lock().await; let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child }); *guard = Some(TunnelProcess {
child,
_pipe_drain: drain_handle,
});
Ok(public_url) Ok(public_url)
} }
+126 -16
View File
@@ -66,6 +66,10 @@ pub trait Tunnel: Send + Sync {
/// Wraps a spawned tunnel child process. /// Wraps a spawned tunnel child process.
pub(crate) struct TunnelProcess { pub(crate) struct TunnelProcess {
pub child: tokio::process::Child, pub child: tokio::process::Child,
/// Background task that drains the process's output pipe (stdout or stderr).
/// Must stay alive or the process dies (SIGPIPE from closed pipe) or hangs
/// (OS pipe buffer fills up, blocking the process's writes).
pub _pipe_drain: Option<tokio::task::JoinHandle<()>>,
} }
pub(crate) type SharedProcess = Arc<Mutex<Option<TunnelProcess>>>; pub(crate) type SharedProcess = Arc<Mutex<Option<TunnelProcess>>>;
@@ -182,6 +186,22 @@ pub fn create_tunnel(config: &TunnelProviderConfig) -> Result<Option<Box<dyn Tun
// ── Managed tunnel startup ─────────────────────────────────────── // ── Managed tunnel startup ───────────────────────────────────────
/// Determine which local address the tunnel should forward traffic to.
///
/// Prefers the webhook server (`HTTP_PORT`) since that's where webhook routes
/// (Telegram, etc.) are served. Falls back to the gateway port if configured,
/// otherwise defaults to 0.0.0.0:8080 (the same fallback the webhook server
/// uses in main.rs when no HTTP config is present).
fn resolve_tunnel_target(channels: &crate::config::ChannelsConfig) -> (&str, u16) {
if let Some(ref http) = channels.http {
return (http.host.as_str(), http.port);
}
if let Some(ref gw) = channels.gateway {
return (gw.host.as_str(), gw.port);
}
("0.0.0.0", 8080)
}
/// Start a managed tunnel if configured and no static URL is already set. /// Start a managed tunnel if configured and no static URL is already set.
/// ///
/// Returns the (potentially mutated) config with `tunnel.public_url` set, /// Returns the (potentially mutated) config with `tunnel.public_url` set,
@@ -201,28 +221,17 @@ pub async fn start_managed_tunnel(
return (config, None); return (config, None);
}; };
let gateway_port = config let (tunnel_host, tunnel_port) = resolve_tunnel_target(&config.channels);
.channels
.gateway
.as_ref()
.map(|g| g.port)
.unwrap_or(3000);
let gateway_host = config
.channels
.gateway
.as_ref()
.map(|g| g.host.as_str())
.unwrap_or("127.0.0.1");
match create_tunnel(provider_config) { match create_tunnel(provider_config) {
Ok(Some(tunnel)) => { Ok(Some(tunnel)) => {
tracing::debug!( tracing::debug!(
"Starting {} tunnel on {}:{}...", "Starting {} tunnel on {}:{}...",
tunnel.name(), tunnel.name(),
gateway_host, tunnel_host,
gateway_port tunnel_port
); );
match tunnel.start(gateway_host, gateway_port).await { match tunnel.start(tunnel_host, tunnel_port).await {
Ok(url) => { Ok(url) => {
tracing::debug!("Tunnel started: {}", url); tracing::debug!("Tunnel started: {}", url);
config.tunnel.public_url = Some(url); config.tunnel.public_url = Some(url);
@@ -383,10 +392,111 @@ mod tests {
{ {
let mut guard = proc.lock().await; let mut guard = proc.lock().await;
*guard = Some(TunnelProcess { child }); *guard = Some(TunnelProcess {
child,
_pipe_drain: None,
});
} }
kill_shared(&proc).await.unwrap(); kill_shared(&proc).await.unwrap();
assert!(proc.lock().await.is_none()); assert!(proc.lock().await.is_none());
} }
// ── Port selection regression tests ──────────────────────────────
fn base_channels() -> crate::config::ChannelsConfig {
crate::config::ChannelsConfig {
cli: crate::config::CliConfig { enabled: false },
http: None,
gateway: None,
signal: None,
wasm_channels_dir: std::env::temp_dir().join("ironclaw-test-channels"),
wasm_channels_enabled: false,
wasm_channel_owner_ids: std::collections::HashMap::new(),
}
}
fn channels_with_http(host: &str, port: u16) -> crate::config::ChannelsConfig {
let mut c = base_channels();
c.http = Some(crate::config::HttpConfig {
host: host.to_string(),
port,
webhook_secret: None,
user_id: "test".to_string(),
});
c.gateway = Some(crate::config::GatewayConfig {
host: "127.0.0.1".to_string(),
port: 3000,
auth_token: None,
user_id: "test".to_string(),
workspace_read_scopes: Vec::new(),
memory_layers: Vec::new(),
user_tokens: None,
});
c
}
fn channels_gateway_only(host: &str, port: u16) -> crate::config::ChannelsConfig {
let mut c = base_channels();
c.gateway = Some(crate::config::GatewayConfig {
host: host.to_string(),
port,
auth_token: None,
user_id: "test".to_string(),
workspace_read_scopes: Vec::new(),
memory_layers: Vec::new(),
user_tokens: None,
});
c
}
fn channels_neither() -> crate::config::ChannelsConfig {
base_channels()
}
#[test]
fn tunnel_target_prefers_http_port() {
let channels = channels_with_http("0.0.0.0", 8080);
let (host, port) = resolve_tunnel_target(&channels);
assert_eq!(host, "0.0.0.0"); // safety: test-only
assert_eq!(port, 8080); // safety: test-only
}
#[test]
fn tunnel_target_falls_back_to_gateway() {
let channels = channels_gateway_only("10.0.0.1", 4000);
let (host, port) = resolve_tunnel_target(&channels);
assert_eq!(host, "10.0.0.1"); // safety: test-only
assert_eq!(port, 4000); // safety: test-only
}
#[test]
fn tunnel_target_defaults_to_webhook_fallback() {
let channels = channels_neither();
let (host, port) = resolve_tunnel_target(&channels);
// Matches the webhook server's hardcoded fallback in main.rs
assert_eq!(host, "0.0.0.0"); // safety: test-only
assert_eq!(port, 8080); // safety: test-only
}
#[test]
fn tunnel_target_http_takes_priority_over_gateway() {
let channels = channels_with_http("192.168.1.1", 9090);
let (host, port) = resolve_tunnel_target(&channels);
// Should use HTTP config, not gateway's 127.0.0.1:3000
assert_eq!(host, "192.168.1.1"); // safety: test-only
assert_eq!(port, 9090); // safety: test-only
}
#[test]
fn tunnel_target_no_http_no_gateway_matches_webhook_fallback() {
// When HTTP_PORT is not set and gateway is not configured (e.g. WASM
// channels exist but no explicit HTTP config), the webhook server in
// main.rs binds to 0.0.0.0:8080 as a hardcoded fallback. The tunnel
// must target the same address so webhook traffic reaches the right
// server.
let channels = channels_neither();
let (host, port) = resolve_tunnel_target(&channels);
assert_eq!((host, port), ("0.0.0.0", 8080)); // safety: test-only
}
} }
+21 -10
View File
@@ -110,12 +110,24 @@ impl Tunnel for NgrokTunnel {
} }
} }
// Drain stdout silently — ngrok only emits low-level connection events if let Ok(mut guard) = self.url.write() {
// to stdout; the pipe must be consumed to prevent SIGPIPE/buffer stalls. *guard = Some(public_url.clone());
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} }); }
// Drain stderr silently — with --log stdout all meaningful output goes // We took ownership of ngrok's stdout pipe above to parse the URL.
// to stdout; stderr only needs to be consumed to prevent pipe stalls. // ngrok continues writing logs to stdout for its entire lifetime.
// If we drop the reader, the pipe closes and ngrok gets SIGPIPE on
// its next write → process dies. We can't just store the reader
// without reading — the OS pipe buffer (~64KB) fills up and ngrok
// blocks. So we drain it in a background task. The task exits
// naturally when ngrok is killed (EOF on the pipe).
let drain_handle = tokio::spawn(async move {
while let Ok(Some(line)) = reader.next_line().await {
tracing::trace!("ngrok: {line}");
}
});
// Drain stderr silently to prevent SIGPIPE/buffer stalls.
if let Some(stderr) = stderr { if let Some(stderr) = stderr {
tokio::spawn(async move { tokio::spawn(async move {
let mut err_reader = tokio::io::BufReader::new(stderr).lines(); let mut err_reader = tokio::io::BufReader::new(stderr).lines();
@@ -123,12 +135,11 @@ impl Tunnel for NgrokTunnel {
}); });
} }
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await; let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child }); *guard = Some(TunnelProcess {
child,
_pipe_drain: Some(drain_handle),
});
Ok(public_url) Ok(public_url)
} }
+51 -1
View File
@@ -1,5 +1,7 @@
//! Shared utility functions used across the codebase. //! Shared utility functions used across the codebase.
use crate::llm::{ChatMessage, Role};
/// Find the largest valid UTF-8 char boundary at or before `pos`. /// Find the largest valid UTF-8 char boundary at or before `pos`.
/// ///
/// Polyfill for `str::floor_char_boundary` (nightly-only). Use when /// Polyfill for `str::floor_char_boundary` (nightly-only). Use when
@@ -16,6 +18,17 @@ pub fn floor_char_boundary(s: &str, pos: usize) -> usize {
i i
} }
/// Ensure the last message in `messages` is a user-role message.
///
/// NEAR AI rejects conversations that don't end with a user message;
/// Claude 4.6 rejects assistant prefill. Call this before any LLM
/// completion request to satisfy both requirements.
pub fn ensure_ends_with_user_message(messages: &mut Vec<ChatMessage>) {
if !matches!(messages.last(), Some(m) if m.role == Role::User) {
messages.push(ChatMessage::user("Continue."));
}
}
/// Check if an LLM response explicitly signals that a job/task is complete. /// Check if an LLM response explicitly signals that a job/task is complete.
/// ///
/// Uses phrase-level matching to avoid false positives from bare words like /// Uses phrase-level matching to avoid false positives from bare words like
@@ -72,7 +85,8 @@ pub fn llm_signals_completion(response: &str) -> bool {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use crate::util::{floor_char_boundary, llm_signals_completion}; use crate::llm::ChatMessage;
use crate::util::{ensure_ends_with_user_message, floor_char_boundary, llm_signals_completion};
// ── floor_char_boundary ── // ── floor_char_boundary ──
@@ -103,6 +117,42 @@ mod tests {
assert_eq!(floor_char_boundary("", 5), 0); assert_eq!(floor_char_boundary("", 5), 0);
} }
// ── ensure_ends_with_user_message ──
#[test]
fn ensure_user_message_injects_when_empty() {
let mut msgs: Vec<ChatMessage> = vec![];
ensure_ends_with_user_message(&mut msgs);
assert_eq!(msgs.len(), 1);
assert_eq!(msgs[0].role, crate::llm::Role::User);
}
#[test]
fn ensure_user_message_injects_after_assistant() {
let mut msgs = vec![ChatMessage::user("hi"), ChatMessage::assistant("hello")];
ensure_ends_with_user_message(&mut msgs);
assert_eq!(msgs.len(), 3);
assert_eq!(msgs[2].role, crate::llm::Role::User);
}
#[test]
fn ensure_user_message_injects_after_tool_result() {
let mut msgs = vec![
ChatMessage::user("run tool"),
ChatMessage::tool_result("call_1", "my_tool", "result"),
];
ensure_ends_with_user_message(&mut msgs);
assert_eq!(msgs.len(), 3);
assert_eq!(msgs[2].role, crate::llm::Role::User);
}
#[test]
fn ensure_user_message_no_op_when_already_user() {
let mut msgs = vec![ChatMessage::user("hello")];
ensure_ends_with_user_message(&mut msgs);
assert_eq!(msgs.len(), 1);
}
// ── llm_signals_completion ── // ── llm_signals_completion ──
#[test] #[test]
+5 -1
View File
@@ -151,7 +151,7 @@ Job: {}
Description: {} Description: {}
You have tools for shell commands, file operations, and code editing. You have tools for shell commands, file operations, and code editing.
Work independently to complete this job. Report when done."#, Work independently to complete this job. When finished, your final message MUST include the phrase "The job is complete" to signal termination."#,
job.title, job.description job.title, job.description
))); )));
@@ -373,6 +373,10 @@ impl LoopDelegate for ContainerDelegate {
// Poll for follow-up prompts from the user // Poll for follow-up prompts from the user
self.poll_and_inject_prompt(reason_ctx).await; self.poll_and_inject_prompt(reason_ctx).await;
// Claude 4.6 rejects assistant prefill; NEAR AI rejects any non-user-ending
// conversation. Ensure the last message is user-role before calling the LLM.
crate::util::ensure_ends_with_user_message(&mut reason_ctx.messages);
// Refresh tools (in case WASM tools were built) // Refresh tools (in case WASM tools were built)
reason_ctx.available_tools = self.tools.tool_definitions().await; reason_ctx.available_tools = self.tools.tool_definitions().await;
+79 -8
View File
@@ -18,7 +18,7 @@ use crate::agent::agentic_loop::{
}; };
use crate::agent::scheduler::WorkerMessage; use crate::agent::scheduler::WorkerMessage;
use crate::agent::task::TaskOutput; use crate::agent::task::TaskOutput;
use crate::channels::web::types::SseEvent; use crate::channels::web::types::ToolDecisionDto;
use crate::context::{ContextManager, JobState}; use crate::context::{ContextManager, JobState};
use crate::db::Database; use crate::db::Database;
use crate::error::Error; use crate::error::Error;
@@ -33,6 +33,7 @@ use crate::tools::rate_limiter::RateLimitResult;
use crate::tools::{ use crate::tools::{
ApprovalContext, ToolRegistry, autonomous_unavailable_error, prepare_tool_params, redact_params, ApprovalContext, ToolRegistry, autonomous_unavailable_error, prepare_tool_params, redact_params,
}; };
use ironclaw_common::AppEvent;
/// Shared dependencies for worker execution. /// Shared dependencies for worker execution.
/// ///
@@ -48,7 +49,7 @@ pub struct WorkerDeps {
pub hooks: Arc<HookRegistry>, pub hooks: Arc<HookRegistry>,
pub timeout: Duration, pub timeout: Duration,
pub use_planning: bool, pub use_planning: bool,
/// SSE manager for live job event streaming to the web gateway. /// Broadcast sender for live job event streaming to the web gateway.
pub sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>, pub sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
/// Approval context for tool execution. When `None`, all non-`Never` tools are /// Approval context for tool execution. When `None`, all non-`Never` tools are
/// blocked (legacy behavior). When `Some`, the context determines which tools /// blocked (legacy behavior). When `Some`, the context determines which tools
@@ -141,7 +142,7 @@ impl Worker {
if let Some(ref sse) = self.deps.sse_tx { if let Some(ref sse) = self.deps.sse_tx {
let job_id_str = job_id.to_string(); let job_id_str = job_id.to_string();
let event = match event_type { let event = match event_type {
"message" => Some(SseEvent::JobMessage { "message" => Some(AppEvent::JobMessage {
job_id: job_id_str, job_id: job_id_str,
role: data role: data
.get("role") .get("role")
@@ -154,7 +155,7 @@ impl Worker {
.unwrap_or("") .unwrap_or("")
.to_string(), .to_string(),
}), }),
"tool_use" => Some(SseEvent::JobToolUse { "tool_use" => Some(AppEvent::JobToolUse {
job_id: job_id_str, job_id: job_id_str,
tool_name: data tool_name: data
.get("tool_name") .get("tool_name")
@@ -166,7 +167,7 @@ impl Worker {
.cloned() .cloned()
.unwrap_or(serde_json::Value::Null), .unwrap_or(serde_json::Value::Null),
}), }),
"tool_result" => Some(SseEvent::JobToolResult { "tool_result" => Some(AppEvent::JobToolResult {
job_id: job_id_str, job_id: job_id_str,
tool_name: data tool_name: data
.get("tool_name") .get("tool_name")
@@ -179,7 +180,7 @@ impl Worker {
.unwrap_or("") .unwrap_or("")
.to_string(), .to_string(),
}), }),
"status" => Some(SseEvent::JobStatus { "status" => Some(AppEvent::JobStatus {
job_id: job_id_str, job_id: job_id_str,
message: data message: data
.get("message") .get("message")
@@ -187,7 +188,7 @@ impl Worker {
.unwrap_or("") .unwrap_or("")
.to_string(), .to_string(),
}), }),
"result" => Some(SseEvent::JobResult { "result" => Some(AppEvent::JobResult {
job_id: job_id_str, job_id: job_id_str,
status: data status: data
.get("status") .get("status")
@@ -200,6 +201,19 @@ impl Worker {
.map(|s| s.to_string()), .map(|s| s.to_string()),
fallback_deliverable: data.get("fallback_deliverable").cloned(), fallback_deliverable: data.get("fallback_deliverable").cloned(),
}), }),
"reasoning" => {
let narrative = data
.get("narrative")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let decisions = ToolDecisionDto::from_json_array(&data["decisions"]);
Some(AppEvent::JobReasoning {
job_id: job_id_str,
narrative,
decisions,
})
}
_ => None, _ => None,
}; };
if let Some(event) = event { if let Some(event) = event {
@@ -897,6 +911,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
id: selection.tool_call_id.clone(), id: selection.tool_call_id.clone(),
name: selection.tool_name.clone(), name: selection.tool_name.clone(),
arguments: selection.parameters.clone(), arguments: selection.parameters.clone(),
reasoning: if action.reasoning.is_empty() {
None
} else {
Some(action.reasoning.clone())
},
}], }],
)); ));
@@ -1232,6 +1251,11 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
) -> Option<LoopOutcome> { ) -> Option<LoopOutcome> {
// Refresh tool definitions so newly built tools become visible // Refresh tool definitions so newly built tools become visible
reason_ctx.available_tools = self.worker.tools().tool_definitions().await; reason_ctx.available_tools = self.worker.tools().tool_definitions().await;
// Claude 4.6 rejects assistant prefill; NEAR AI rejects any non-user-ending
// conversation. Ensure the last message is user-role before calling the LLM.
crate::util::ensure_ends_with_user_message(&mut reason_ctx.messages);
None None
} }
@@ -1352,6 +1376,48 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
); );
} }
// Emit reasoning event if any tool calls carry reasoning.
// Sanitize narrative and per-tool rationale through SafetyLayer
// (parity with ChatDelegate in dispatcher.rs).
let sanitized_narrative = content
.as_deref()
.filter(|c| !c.trim().is_empty())
.map(|c| {
self.worker
.deps
.safety
.sanitize_tool_output("job_narrative", c)
.content
})
.filter(|c| !c.trim().is_empty())
.unwrap_or_default();
let decisions: Vec<serde_json::Value> = tool_calls
.iter()
.filter_map(|tc| {
tc.reasoning.as_ref().map(|r| {
let sanitized = self
.worker
.deps
.safety
.sanitize_tool_output("tool_rationale", r)
.content;
serde_json::json!({
"tool_name": tc.name,
"rationale": sanitized,
})
})
})
.collect();
if !decisions.is_empty() {
self.worker.log_event(
"reasoning",
serde_json::json!({
"narrative": sanitized_narrative,
"decisions": decisions,
}),
);
}
// Add assistant message with tool_calls (OpenAI protocol) // Add assistant message with tool_calls (OpenAI protocol)
reason_ctx reason_ctx
.messages .messages
@@ -1366,7 +1432,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
.map(|tc| ToolSelection { .map(|tc| ToolSelection {
tool_name: tc.name.clone(), tool_name: tc.name.clone(),
parameters: tc.arguments.clone(), parameters: tc.arguments.clone(),
reasoning: String::new(), reasoning: tc.reasoning.clone().unwrap_or_default(),
alternatives: vec![], alternatives: vec![],
tool_call_id: tc.id.clone(), tool_call_id: tc.id.clone(),
}) })
@@ -1419,6 +1485,11 @@ fn selections_to_tool_calls(selections: &[ToolSelection]) -> Vec<ToolCall> {
id: s.tool_call_id.clone(), id: s.tool_call_id.clone(),
name: s.tool_name.clone(), name: s.tool_name.clone(),
arguments: s.parameters.clone(), arguments: s.parameters.clone(),
reasoning: if s.reasoning.is_empty() {
None
} else {
Some(s.reasoning.clone())
},
}) })
.collect() .collect()
} }
+191
View File
@@ -0,0 +1,191 @@
//! Tests for batch_get_last_run_status (#1469 N+1 fix).
//!
//! Verifies:
//! 1. Empty input returns empty map
//! 2. Returns the most recent run status per routine
//! 3. Routines with no runs are omitted from result
//! 4. Multiple routines with different statuses are correctly returned
#[cfg(feature = "libsql")]
mod tests {
use std::sync::Arc;
use chrono::{Duration, Utc};
use uuid::Uuid;
use ironclaw::agent::routine::{
Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
};
use ironclaw::db::Database;
async fn create_test_db() -> (Arc<dyn Database>, tempfile::TempDir) {
use ironclaw::db::libsql::LibSqlBackend;
let temp_dir = tempfile::tempdir().expect("tempdir");
let db_path = temp_dir.path().join("test.db");
let backend = LibSqlBackend::new_local(&db_path)
.await
.expect("LibSqlBackend");
backend.run_migrations().await.expect("migrations");
let db: Arc<dyn Database> = Arc::new(backend);
(db, temp_dir)
}
fn make_routine(id: Uuid) -> Routine {
Routine {
id,
name: format!("test-routine-{}", id),
description: "Test routine".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Manual,
action: RoutineAction::FullJob {
title: "Test job".to_string(),
description: "Test description".to_string(),
max_iterations: 5,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent: 1,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
}
}
fn make_run(
routine_id: Uuid,
status: RunStatus,
started_at: chrono::DateTime<chrono::Utc>,
) -> RoutineRun {
RoutineRun {
id: Uuid::new_v4(),
routine_id,
trigger_type: "manual".to_string(),
trigger_detail: None,
started_at,
completed_at: if status == RunStatus::Running {
None
} else {
Some(Utc::now())
},
status,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
}
}
#[tokio::test]
async fn test_batch_get_last_run_status_empty_input() {
let (db, _tmp) = create_test_db().await;
let result = db
.batch_get_last_run_status(&[])
.await
.expect("batch query");
assert!(result.is_empty());
}
#[tokio::test]
async fn test_batch_get_last_run_status_returns_latest() {
let (db, _tmp) = create_test_db().await;
let routine_id = Uuid::new_v4();
db.create_routine(&make_routine(routine_id))
.await
.expect("create routine");
// Create an older run with Ok status
let older_run = make_run(routine_id, RunStatus::Ok, Utc::now() - Duration::hours(2));
db.create_routine_run(&older_run)
.await
.expect("create older run");
db.complete_routine_run(older_run.id, RunStatus::Ok, None, None)
.await
.expect("complete older run");
// Create a newer run with Attention status
let newer_run = make_run(
routine_id,
RunStatus::Attention,
Utc::now() - Duration::hours(1),
);
db.create_routine_run(&newer_run)
.await
.expect("create newer run");
db.complete_routine_run(newer_run.id, RunStatus::Attention, None, None)
.await
.expect("complete newer run");
let result = db
.batch_get_last_run_status(&[routine_id])
.await
.expect("batch query");
assert_eq!(result.get(&routine_id), Some(&RunStatus::Attention));
}
#[tokio::test]
async fn test_batch_get_last_run_status_omits_routines_without_runs() {
let (db, _tmp) = create_test_db().await;
let with_runs = Uuid::new_v4();
let without_runs = Uuid::new_v4();
db.create_routine(&make_routine(with_runs))
.await
.expect("create routine");
db.create_routine(&make_routine(without_runs))
.await
.expect("create routine");
let run = make_run(with_runs, RunStatus::Ok, Utc::now());
db.create_routine_run(&run).await.expect("create run");
db.complete_routine_run(run.id, RunStatus::Ok, None, None)
.await
.expect("complete run");
let result = db
.batch_get_last_run_status(&[with_runs, without_runs])
.await
.expect("batch query");
assert_eq!(result.get(&with_runs), Some(&RunStatus::Ok));
assert_eq!(result.get(&without_runs), None);
}
#[tokio::test]
async fn test_batch_get_last_run_status_multiple_routines() {
let (db, _tmp) = create_test_db().await;
let r1 = Uuid::new_v4();
let r2 = Uuid::new_v4();
db.create_routine(&make_routine(r1))
.await
.expect("create r1");
db.create_routine(&make_routine(r2))
.await
.expect("create r2");
let run1 = make_run(r1, RunStatus::Running, Utc::now());
db.create_routine_run(&run1).await.expect("create run1");
let run2 = make_run(r2, RunStatus::Failed, Utc::now());
db.create_routine_run(&run2).await.expect("create run2");
db.complete_routine_run(run2.id, RunStatus::Failed, None, None)
.await
.expect("complete run2");
let result = db
.batch_get_last_run_status(&[r1, r2])
.await
.expect("batch query");
assert_eq!(result.get(&r1), Some(&RunStatus::Running));
assert_eq!(result.get(&r2), Some(&RunStatus::Failed));
}
}
+9
View File
@@ -53,6 +53,7 @@ HEADED=1 pytest scenarios/
| `test_skills.py` | Skills tab UI visibility, ClawHub search (skipped if registry unreachable), install + remove lifecycle | | `test_skills.py` | Skills tab UI visibility, ClawHub search (skipped if registry unreachable), install + remove lifecycle |
| `test_sse_reconnect.py` | SSE reconnects after programmatic `eventSource.close()` + `connectSSE()`; history is reloaded after reconnect | | `test_sse_reconnect.py` | SSE reconnects after programmatic `eventSource.close()` + `connectSSE()`; history is reloaded after reconnect |
| `test_tool_approval.py` | Approval card appears, buttons disable on approve/deny, parameters toggle via `page.evaluate("showApproval(...)")`; the waiting-approval regression uses a real HTTP tool call | | `test_tool_approval.py` | Approval card appears, buttons disable on approve/deny, parameters toggle via `page.evaluate("showApproval(...)")`; the waiting-approval regression uses a real HTTP tool call |
| `test_oauth_refresh.py` | Hosted Gmail OAuth regression: complete setup via `/oauth/callback`, expire the stored access token in libSQL, trigger a real `gmail` tool call through `/api/chat/send`, and verify refresh goes through the mock `/oauth/refresh` proxy without forwarding `client_secret` |
## `helpers.py` ## `helpers.py`
@@ -75,6 +76,7 @@ All fixtures are defined in `tests/e2e/conftest.py`. Running `pytest scenarios/`
| `ironclaw_binary` | Checks `target/debug/ironclaw`; if absent, runs `cargo build --no-default-features --features libsql` (timeout 600s). | | `ironclaw_binary` | Checks `target/debug/ironclaw`; if absent, runs `cargo build --no-default-features --features libsql` (timeout 600s). |
| `mock_llm_server` | Starts `mock_llm.py --port 0`, reads the assigned port from stdout, waits for `/v1/models` to return 200. Yields the base URL. | | `mock_llm_server` | Starts `mock_llm.py --port 0`, reads the assigned port from stdout, waits for `/v1/models` to return 200. Yields the base URL. |
| `ironclaw_server` | Starts the ironclaw binary with a minimal env (see below), waits for `/api/health` (timeout 60s). Yields the base URL. On teardown sends **SIGINT** (not SIGTERM) so the tokio ctrl_c handler triggers a graceful shutdown and LLVM coverage data is flushed. | | `ironclaw_server` | Starts the ironclaw binary with a minimal env (see below), waits for `/api/health` (timeout 60s). Yields the base URL. On teardown sends **SIGINT** (not SIGTERM) so the tokio ctrl_c handler triggers a graceful shutdown and LLVM coverage data is flushed. |
| `hosted_oauth_refresh_server` | Starts a second ironclaw instance with a dedicated libSQL DB and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id`, while still pointing `IRONCLAW_OAUTH_EXCHANGE_URL` at `mock_llm.py`. Yields a dict with `base_url`, `db_path`, `gateway_user_id`, and `mock_llm_url` for the hosted refresh regression scenario. |
| `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. | | `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. |
### Function-scoped fixtures ### Function-scoped fixtures
@@ -100,6 +102,8 @@ EMBEDDING_ENABLED=false, SKILLS_ENABLED=true
ONBOARD_COMPLETED=true # prevents setup wizard ONBOARD_COMPLETED=true # prevents setup wizard
``` ```
The `hosted_oauth_refresh_server` fixture uses the same baseline, but with its own DB/home tempdirs and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id` so hosted OAuth flows exercise proxy credential injection instead of the baked-in desktop Google app.
The binary is also started with `--no-onboard`. Coverage env vars (`CARGO_LLVM_COV*`, `LLVM_*`, `CARGO_ENCODED_RUSTFLAGS`, `CARGO_INCREMENTAL`) are forwarded from the outer environment when present. The binary is also started with `--no-onboard`. Coverage env vars (`CARGO_LLVM_COV*`, `LLVM_*`, `CARGO_ENCODED_RUSTFLAGS`, `CARGO_INCREMENTAL`) are forwarded from the outer environment when present.
## Mock LLM (`mock_llm.py`) ## Mock LLM (`mock_llm.py`)
@@ -113,6 +117,11 @@ python mock_llm.py --port 0
It serves `POST /v1/chat/completions` (streaming + non-streaming) and `GET /v1/models`. Responses are pattern-matched from `CANNED_RESPONSES` against the last user message. Unmatched messages return `"I understand your request."`. The model name reported is always `"mock-model"`. It serves `POST /v1/chat/completions` (streaming + non-streaming) and `GET /v1/models`. Responses are pattern-matched from `CANNED_RESPONSES` against the last user message. Unmatched messages return `"I understand your request."`. The model name reported is always `"mock-model"`.
It also hosts OAuth test endpoints:
- `POST /oauth/exchange` for hosted auth-code exchange
- `POST /oauth/refresh` for hosted refresh-token exchange
- `GET /__mock/oauth/state` and `POST /__mock/oauth/reset` so HTTP E2E scenarios can assert exact proxy payloads and reset counters between setup and refresh assertions
To add a new canned response: To add a new canned response:
```python ```python
# In mock_llm.py # In mock_llm.py
+169 -34
View File
@@ -112,6 +112,39 @@ def _reserve_loopback_sockets(count: int) -> list[socket.socket]:
sock.close() sock.close()
raise raise
async def _stop_process(
proc: asyncio.subprocess.Process, *, sig: int | None = None, timeout: float
) -> None:
"""Signal a subprocess and wait briefly without masking exit races."""
if proc.returncode is not None:
return
try:
if sig is None:
proc.kill()
else:
proc.send_signal(sig)
except ProcessLookupError:
try:
await asyncio.wait_for(proc.wait(), timeout=timeout)
except asyncio.TimeoutError:
pass
return
try:
await asyncio.wait_for(proc.wait(), timeout=timeout)
except asyncio.TimeoutError:
pass
def _forward_coverage_env(env: dict[str, str]) -> None:
"""Forward cargo-llvm-cov env vars into child processes when present."""
cov_env_prefixes = ("CARGO_LLVM_COV", "LLVM_")
cov_env_extras = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
for key, val in os.environ.items():
if key.startswith(cov_env_prefixes) or key in cov_env_extras:
env[key] = val
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def ironclaw_binary(): def ironclaw_binary():
@@ -264,14 +297,7 @@ async def ironclaw_server(
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
} }
# Forward LLVM coverage instrumentation env vars when present _forward_coverage_env(env)
# (allows cargo-llvm-cov to collect profraw data from E2E runs).
# Use prefix matching to stay resilient to cargo-llvm-cov changes.
COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_")
COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
for key, val in os.environ.items():
if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS:
env[key] = val
proc = await asyncio.create_subprocess_exec( proc = await asyncio.create_subprocess_exec(
ironclaw_binary, "--no-onboard", ironclaw_binary, "--no-onboard",
stdin=asyncio.subprocess.DEVNULL, stdin=asyncio.subprocess.DEVNULL,
@@ -279,35 +305,145 @@ async def ironclaw_server(
stderr=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
env=env, env=env,
) )
startup_kill_attempted = False
base_url = f"http://127.0.0.1:{gateway_port}" base_url = f"http://127.0.0.1:{gateway_port}"
try: try:
await wait_for_ready(f"{base_url}/api/health", timeout=60) await wait_for_ready(f"{base_url}/api/health", timeout=60)
yield base_url yield base_url
except TimeoutError: except TimeoutError:
# Dump stderr so CI logs show why the server failed to start # Dump stderr so CI logs show why the server failed to start
if proc.returncode is None:
startup_kill_attempted = True
await _stop_process(proc, timeout=2)
returncode = proc.returncode returncode = proc.returncode
stderr_bytes = b"" stderr_bytes = b""
if proc.stderr: if proc.stderr:
try: try:
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2) stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
except (asyncio.TimeoutError, Exception): except asyncio.TimeoutError:
pass pass
stderr_text = stderr_bytes.decode("utf-8", errors="replace") stderr_text = stderr_bytes.decode("utf-8", errors="replace")
proc.kill()
pytest.fail( pytest.fail(
f"ironclaw server failed to start on port {gateway_port} " f"ironclaw server failed to start on port {gateway_port} "
f"(returncode={returncode}).\nstderr:\n{stderr_text}" f"(returncode={returncode}).\nstderr:\n{stderr_text}"
) )
finally: finally:
if proc.returncode is None: if proc.returncode is None:
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a if startup_kill_attempted:
# graceful shutdown. This lets the LLVM coverage runtime run its await _stop_process(proc, timeout=2)
# atexit handler and flush .profraw files for cargo-llvm-cov. else:
proc.send_signal(signal.SIGINT) # Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
try: # graceful shutdown. This lets the LLVM coverage runtime run its
await asyncio.wait_for(proc.wait(), timeout=10) # atexit handler and flush .profraw files for cargo-llvm-cov.
except asyncio.TimeoutError: await _stop_process(proc, sig=signal.SIGINT, timeout=10)
proc.kill() if proc.returncode is None:
await _stop_process(proc, timeout=2)
@pytest.fixture(scope="session")
async def hosted_oauth_refresh_server(
ironclaw_binary,
mock_llm_server,
wasm_tools_dir,
):
"""Start a hosted-mode ironclaw instance for OAuth refresh regression tests."""
reserved = _reserve_loopback_sockets(2)
db_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-db-")
home_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-home-")
try:
gateway_port = reserved[0].getsockname()[1]
http_port = reserved[1].getsockname()[1]
for sock in reserved:
if sock.fileno() != -1:
sock.close()
db_path = os.path.join(db_tmpdir.name, "hosted-oauth-refresh.db")
home_dir = home_tmpdir.name
env = {
"PATH": os.environ.get("PATH", "/usr/bin:/bin"),
"HOME": home_dir,
"IRONCLAW_BASE_DIR": os.path.join(home_dir, ".ironclaw"),
"RUST_LOG": "ironclaw=info",
"RUST_BACKTRACE": "1",
"IRONCLAW_OWNER_ID": OWNER_SCOPE_ID,
"GATEWAY_ENABLED": "true",
"GATEWAY_HOST": "127.0.0.1",
"GATEWAY_PORT": str(gateway_port),
"GATEWAY_AUTH_TOKEN": AUTH_TOKEN,
"GATEWAY_USER_ID": OWNER_SCOPE_ID,
"HTTP_HOST": "127.0.0.1",
"HTTP_PORT": str(http_port),
"HTTP_WEBHOOK_SECRET": HTTP_WEBHOOK_SECRET,
"CLI_ENABLED": "false",
"LLM_BACKEND": "openai_compatible",
"LLM_BASE_URL": mock_llm_server,
"LLM_MODEL": "mock-model",
"DATABASE_BACKEND": "libsql",
"LIBSQL_PATH": db_path,
"SECRETS_MASTER_KEY": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
"SANDBOX_ENABLED": "false",
"SKILLS_ENABLED": "true",
"ROUTINES_ENABLED": "true",
"HEARTBEAT_ENABLED": "false",
"EMBEDDING_ENABLED": "false",
"WASM_ENABLED": "true",
"WASM_TOOLS_DIR": wasm_tools_dir,
"WASM_CHANNELS_DIR": _WASM_CHANNELS_TMPDIR.name,
"ONBOARD_COMPLETED": "true",
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
"GOOGLE_OAUTH_CLIENT_ID": "hosted-google-client-id",
}
_forward_coverage_env(env)
proc = await asyncio.create_subprocess_exec(
ironclaw_binary, "--no-onboard",
stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
)
startup_kill_attempted = False
base_url = f"http://127.0.0.1:{gateway_port}"
try:
await wait_for_ready(f"{base_url}/api/health", timeout=60)
yield {
"base_url": base_url,
"db_path": db_path,
"gateway_user_id": OWNER_SCOPE_ID,
"mock_llm_url": mock_llm_server,
}
except TimeoutError:
if proc.returncode is None:
startup_kill_attempted = True
await _stop_process(proc, timeout=2)
returncode = proc.returncode
stderr_bytes = b""
if proc.stderr:
try:
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
except asyncio.TimeoutError:
pass
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
pytest.fail(
f"hosted oauth refresh server failed to start on port {gateway_port} "
f"(returncode={returncode}).\nstderr:\n{stderr_text}"
)
finally:
if proc.returncode is None:
if startup_kill_attempted:
await _stop_process(proc, timeout=2)
else:
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
if proc.returncode is None:
await _stop_process(proc, timeout=2)
finally:
for sock in reserved:
if sock.fileno() != -1:
sock.close()
db_tmpdir.cleanup()
home_tmpdir.cleanup()
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
@@ -362,12 +498,7 @@ async def http_channel_server_without_secret(
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
} }
# Forward LLVM coverage instrumentation env vars when present _forward_coverage_env(env)
COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_")
COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
for key, val in os.environ.items():
if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS:
env[key] = val
proc = await asyncio.create_subprocess_exec( proc = await asyncio.create_subprocess_exec(
ironclaw_binary, "--no-onboard", ironclaw_binary, "--no-onboard",
stdin=asyncio.subprocess.DEVNULL, stdin=asyncio.subprocess.DEVNULL,
@@ -375,6 +506,7 @@ async def http_channel_server_without_secret(
stderr=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
env=env, env=env,
) )
startup_kill_attempted = False
gateway_url = f"http://127.0.0.1:{gateway_port}" gateway_url = f"http://127.0.0.1:{gateway_port}"
http_base_url = f"http://127.0.0.1:{http_port}" http_base_url = f"http://127.0.0.1:{http_port}"
try: try:
@@ -383,15 +515,17 @@ async def http_channel_server_without_secret(
yield http_base_url yield http_base_url
except TimeoutError: except TimeoutError:
# Dump stderr so CI logs show why the server failed to start # Dump stderr so CI logs show why the server failed to start
if proc.returncode is None:
startup_kill_attempted = True
await _stop_process(proc, timeout=2)
returncode = proc.returncode returncode = proc.returncode
stderr_bytes = b"" stderr_bytes = b""
if proc.stderr: if proc.stderr:
try: try:
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2) stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
except (asyncio.TimeoutError, Exception): except asyncio.TimeoutError:
pass pass
stderr_text = stderr_bytes.decode("utf-8", errors="replace") stderr_text = stderr_bytes.decode("utf-8", errors="replace")
proc.kill()
pytest.fail( pytest.fail(
f"ironclaw server without webhook secret failed to start on ports " f"ironclaw server without webhook secret failed to start on ports "
f"gateway={gateway_port}, http={http_port} " f"gateway={gateway_port}, http={http_port} "
@@ -399,14 +533,15 @@ async def http_channel_server_without_secret(
) )
finally: finally:
if proc.returncode is None: if proc.returncode is None:
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a if startup_kill_attempted:
# graceful shutdown. This lets the LLVM coverage runtime run its await _stop_process(proc, timeout=2)
# atexit handler and flush .profraw files for cargo-llvm-cov. else:
proc.send_signal(signal.SIGINT) # Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
try: # graceful shutdown. This lets the LLVM coverage runtime run its
await asyncio.wait_for(proc.wait(), timeout=10) # atexit handler and flush .profraw files for cargo-llvm-cov.
except asyncio.TimeoutError: await _stop_process(proc, sig=signal.SIGINT, timeout=10)
proc.kill() if proc.returncode is None:
await _stop_process(proc, timeout=2)
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
+61
View File
@@ -34,6 +34,15 @@ TOOL_CALL_PATTERNS = [
"body": {"label": m.group("label")}, "body": {"label": m.group("label")},
}, },
), ),
(
re.compile(r"check gmail unread|gmail unread", re.IGNORECASE),
"gmail",
lambda _: {
"action": "list_messages",
"query": "is:unread",
"max_results": 1,
},
),
(re.compile(r"what time|current time", re.IGNORECASE), "time", lambda _: {"operation": "now"}), (re.compile(r"what time|current time", re.IGNORECASE), "time", lambda _: {"operation": "now"}),
( (
re.compile( re.compile(
@@ -91,6 +100,15 @@ TOOL_CALL_PATTERNS = [
] ]
def _new_oauth_state() -> dict:
return {
"exchange_count": 0,
"refresh_count": 0,
"last_exchange": None,
"last_refresh": None,
}
def _last_user_content(messages: list[dict]) -> str: def _last_user_content(messages: list[dict]) -> str:
for msg in reversed(messages): for msg in reversed(messages):
if msg.get("role") == "user": if msg.get("role") == "user":
@@ -272,6 +290,12 @@ async def oauth_exchange(request: web.Request) -> web.Response:
specific token params such as RFC 8707 `resource` are forwarded here. specific token params such as RFC 8707 `resource` are forwarded here.
""" """
data = await request.post() data = await request.post()
oauth_state = request.app["oauth_state"]
oauth_state["exchange_count"] += 1
oauth_state["last_exchange"] = {
"authorization": request.headers.get("Authorization"),
"form": dict(data),
}
code = data.get("code", "") code = data.get("code", "")
access_token_field = data.get("access_token_field", "access_token") access_token_field = data.get("access_token_field", "access_token")
@@ -290,6 +314,39 @@ async def oauth_exchange(request: web.Request) -> web.Response:
}) })
async def oauth_refresh(request: web.Request) -> web.Response:
"""Mock OAuth token refresh proxy for hosted refresh E2E tests."""
data = await request.post()
oauth_state = request.app["oauth_state"]
oauth_state["refresh_count"] += 1
oauth_state["last_refresh"] = {
"authorization": request.headers.get("Authorization"),
"form": dict(data),
}
if request.headers.get("Authorization") != "Bearer e2e-test-token":
return web.json_response({"error": "invalid_gateway_auth"}, status=401)
if data.get("client_id") != "hosted-google-client-id":
return web.json_response({"error": "invalid_client_id"}, status=400)
if "client_secret" in data:
return web.json_response({"error": "unexpected_client_secret"}, status=400)
return web.json_response({
"access_token": "mock-refreshed-access-token",
"refresh_token": "mock-rotated-refresh-token",
"expires_in": 3600,
})
async def oauth_state_handler(request: web.Request) -> web.Response:
return web.json_response(request.app["oauth_state"])
async def oauth_reset(request: web.Request) -> web.Response:
request.app["oauth_state"] = _new_oauth_state()
return web.json_response({"ok": True})
async def models(_request: web.Request) -> web.Response: async def models(_request: web.Request) -> web.Response:
return web.json_response({ return web.json_response({
"object": "list", "object": "list",
@@ -424,12 +481,16 @@ def main():
parser.add_argument("--port", type=int, default=0) parser.add_argument("--port", type=int, default=0)
args = parser.parse_args() args = parser.parse_args()
app = web.Application() app = web.Application()
app["oauth_state"] = _new_oauth_state()
# Register both /v1/ and non-/v1/ paths (rig-core omits the /v1/ prefix) # Register both /v1/ and non-/v1/ paths (rig-core omits the /v1/ prefix)
app.router.add_post("/v1/chat/completions", chat_completions) app.router.add_post("/v1/chat/completions", chat_completions)
app.router.add_post("/chat/completions", chat_completions) app.router.add_post("/chat/completions", chat_completions)
app.router.add_get("/v1/models", models) app.router.add_get("/v1/models", models)
app.router.add_get("/models", models) app.router.add_get("/models", models)
app.router.add_post("/oauth/exchange", oauth_exchange) app.router.add_post("/oauth/exchange", oauth_exchange)
app.router.add_post("/oauth/refresh", oauth_refresh)
app.router.add_get("/__mock/oauth/state", oauth_state_handler)
app.router.add_post("/__mock/oauth/reset", oauth_reset)
# Mock MCP server endpoints # Mock MCP server endpoints
app.router.add_post("/mcp", mcp_endpoint) app.router.add_post("/mcp", mcp_endpoint)
app.router.add_post("/mcp-400", mcp_endpoint_400) app.router.add_post("/mcp-400", mcp_endpoint_400)
+227
View File
@@ -0,0 +1,227 @@
"""Hosted OAuth refresh HTTP regression test.
Runs a real ironclaw binary in hosted mode, expires a stored Gmail access
token in the libSQL database, triggers a real gmail tool call through the
chat API, and verifies that refresh uses the hosted proxy endpoint.
"""
import asyncio
import sqlite3
from datetime import datetime, timezone
from urllib.parse import parse_qs, urlparse
import httpx
from helpers import api_get, api_post
def _extract_state(auth_url: str) -> str:
parsed = urlparse(auth_url)
state = parse_qs(parsed.query).get("state", [None])[0]
assert state, f"auth_url should include state: {auth_url}"
return state
def _parse_timestamp(value: str | None) -> datetime | None:
if value is None:
return None
return datetime.fromisoformat(value.replace("Z", "+00:00"))
def _expire_access_token(db_path: str, user_id: str, secret_name: str) -> None:
with sqlite3.connect(db_path) as conn:
cursor = conn.execute(
"""
UPDATE secrets
SET expires_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-1 hour')
WHERE user_id = ?1 AND name = ?2
""",
(user_id, secret_name),
)
conn.commit()
assert cursor.rowcount == 1, f"Expected one secret row for {user_id}/{secret_name}"
def _find_secret_row(
db_path: str,
secret_name: str,
) -> tuple[str, str | None, str | None]:
with sqlite3.connect(db_path) as conn:
row = conn.execute(
"""
SELECT user_id, expires_at, updated_at
FROM secrets
WHERE name = ?1
ORDER BY updated_at DESC
LIMIT 1
""",
(secret_name,),
).fetchone()
assert row is not None, f"Missing secret row for {secret_name}"
return row[0], row[1], row[2]
async def _get_extension(base_url: str, name: str) -> dict | None:
response = await api_get(base_url, "/api/extensions", timeout=15)
response.raise_for_status()
for extension in response.json().get("extensions", []):
if extension["name"] == name:
return extension
return None
async def _reset_mock_oauth_state(mock_base_url: str) -> None:
async with httpx.AsyncClient() as client:
response = await client.post(f"{mock_base_url}/__mock/oauth/reset", timeout=10)
response.raise_for_status()
async def _get_mock_oauth_state(mock_base_url: str) -> dict:
async with httpx.AsyncClient() as client:
response = await client.get(f"{mock_base_url}/__mock/oauth/state", timeout=10)
response.raise_for_status()
return response.json()
async def _approve_pending_request(base_url: str, thread_id: str, request_id: str) -> None:
response = await api_post(
base_url,
"/api/chat/approval",
json={"request_id": request_id, "action": "approve", "thread_id": thread_id},
timeout=15,
)
assert response.status_code == 202, (
f"Approval submission failed: {response.status_code} {response.text[:400]}"
)
async def _wait_for_gmail_tool_call(base_url: str, thread_id: str, timeout: float = 30.0) -> dict:
approved_request_ids = set()
for _ in range(int(timeout * 2)):
response = await api_get(
base_url,
f"/api/chat/history?thread_id={thread_id}",
timeout=15,
)
response.raise_for_status()
history = response.json()
pending = history.get("pending_approval")
if pending and pending["request_id"] not in approved_request_ids:
await _approve_pending_request(base_url, thread_id, pending["request_id"])
approved_request_ids.add(pending["request_id"])
for turn in history.get("turns", []):
for tool_call in turn.get("tool_calls", []):
if tool_call.get("name") == "gmail":
return history
await asyncio.sleep(0.5)
raise AssertionError(f"Timed out waiting for gmail tool call in thread {thread_id}")
async def _wait_for_refresh_request(mock_base_url: str, timeout: float = 20.0) -> dict:
for _ in range(int(timeout * 2)):
state = await _get_mock_oauth_state(mock_base_url)
if state.get("refresh_count") == 1:
return state
await asyncio.sleep(0.5)
raise AssertionError("Timed out waiting for exactly one OAuth refresh request")
async def test_hosted_gmail_oauth_refresh_uses_proxy(hosted_oauth_refresh_server):
server = hosted_oauth_refresh_server["base_url"]
db_path = hosted_oauth_refresh_server["db_path"]
mock_base_url = hosted_oauth_refresh_server["mock_llm_url"]
install_response = await api_post(
server,
"/api/extensions/install",
json={"name": "gmail"},
timeout=180,
)
assert install_response.status_code == 200, install_response.text
assert install_response.json().get("success") is True
setup_response = await api_post(
server,
"/api/extensions/gmail/setup",
json={"secrets": {}},
timeout=30,
)
assert setup_response.status_code == 200, setup_response.text
setup_data = setup_response.json()
assert setup_data.get("success") is True, setup_data
auth_url = setup_data.get("auth_url")
assert auth_url, setup_data
auth_params = parse_qs(urlparse(auth_url).query)
assert auth_params.get("client_id") == ["hosted-google-client-id"]
async with httpx.AsyncClient() as client:
callback_response = await client.get(
f"{server}/oauth/callback",
params={"code": "mock_auth_code", "state": _extract_state(auth_url)},
timeout=30,
follow_redirects=True,
)
assert callback_response.status_code == 200, callback_response.text[:400]
callback_body = callback_response.text.lower()
assert "connected" in callback_body or "success" in callback_body
gmail = await _get_extension(server, "gmail")
assert gmail is not None, "gmail should be installed"
assert gmail["authenticated"] is True, gmail
assert "gmail" in gmail.get("tools", []), gmail
await _reset_mock_oauth_state(mock_base_url)
stored_user_id, expires_before, updated_before = _find_secret_row(
db_path, "google_oauth_token"
)
assert _parse_timestamp(expires_before) is not None
assert _parse_timestamp(updated_before) is not None
await asyncio.sleep(0.1)
_expire_access_token(db_path, stored_user_id, "google_oauth_token")
thread_response = await api_post(server, "/api/chat/thread/new", timeout=15)
assert thread_response.status_code == 200, thread_response.text
thread_id = thread_response.json()["id"]
send_response = await api_post(
server,
"/api/chat/send",
json={"content": "check gmail unread", "thread_id": thread_id},
timeout=30,
)
assert send_response.status_code == 202, send_response.text
history = await _wait_for_gmail_tool_call(server, thread_id)
assert any(
tool_call.get("name") == "gmail"
for turn in history.get("turns", [])
for tool_call in turn.get("tool_calls", [])
), history
oauth_state = await _wait_for_refresh_request(mock_base_url)
assert oauth_state["refresh_count"] == 1, oauth_state
last_refresh = oauth_state["last_refresh"]
assert last_refresh is not None, oauth_state
assert last_refresh["authorization"] == "Bearer e2e-test-token"
assert last_refresh["form"]["client_id"] == "hosted-google-client-id"
assert "client_secret" not in last_refresh["form"], last_refresh
refreshed_user_id, expires_after, updated_after = _find_secret_row(
db_path, "google_oauth_token"
)
assert refreshed_user_id == stored_user_id
expires_after_dt = _parse_timestamp(expires_after)
updated_after_dt = _parse_timestamp(updated_after)
updated_before_dt = _parse_timestamp(updated_before)
assert expires_after_dt is not None
assert updated_after_dt is not None
assert updated_before_dt is not None
assert expires_after_dt > datetime.now(timezone.utc)
assert updated_after_dt > updated_before_dt
+12 -32
View File
@@ -561,11 +561,7 @@ mod tests {
"deploy to production now", "deploy to production now",
); );
let fired = engine let fired = engine
.check_event_triggers( .check_event_triggers(&matching_msg, &matching_msg.content)
&matching_msg.user_id,
&matching_msg.channel,
&matching_msg.content,
)
.await; .await;
assert!( assert!(
fired >= 1, fired >= 1,
@@ -584,11 +580,7 @@ mod tests {
"check the staging environment", "check the staging environment",
); );
let fired_neg = engine let fired_neg = engine
.check_event_triggers( .check_event_triggers(&non_matching_msg, &non_matching_msg.content)
&non_matching_msg.user_id,
&non_matching_msg.channel,
&non_matching_msg.content,
)
.await; .await;
assert_eq!(fired_neg, 0, "Expected 0 routines fired on non-match"); assert_eq!(fired_neg, 0, "Expected 0 routines fired on non-match");
} }
@@ -652,7 +644,7 @@ mod tests {
"deploy to production now", "deploy to production now",
); );
let guest_fired = engine let guest_fired = engine
.check_event_triggers(&guest_msg.user_id, &guest_msg.channel, &guest_msg.content) .check_event_triggers(&guest_msg, &guest_msg.content)
.await; .await;
assert_eq!( assert_eq!(
guest_fired, 0, guest_fired, 0,
@@ -677,7 +669,7 @@ mod tests {
"deploy to production now", "deploy to production now",
); );
let owner_fired = engine let owner_fired = engine
.check_event_triggers(&owner_msg.user_id, &owner_msg.channel, &owner_msg.content) .check_event_triggers(&owner_msg, &owner_msg.content)
.await; .await;
assert!( assert!(
owner_fired >= 1, owner_fired >= 1,
@@ -906,9 +898,7 @@ mod tests {
"default", "default",
"test-cooldown trigger", "test-cooldown trigger",
); );
let fired1 = engine let fired1 = engine.check_event_triggers(&msg, &msg.content).await;
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await;
assert!(fired1 >= 1, "First fire should work"); assert!(fired1 >= 1, "First fire should work");
// Give spawn time, then update last_run_at to simulate recent execution. // Give spawn time, then update last_run_at to simulate recent execution.
@@ -923,9 +913,7 @@ mod tests {
engine.refresh_event_cache().await; engine.refresh_event_cache().await;
// Second fire should be blocked by cooldown. // Second fire should be blocked by cooldown.
let fired2 = engine let fired2 = engine.check_event_triggers(&msg, &msg.content).await;
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await;
assert_eq!(fired2, 0, "Second fire should be blocked by cooldown"); assert_eq!(fired2, 0, "Second fire should be blocked by cooldown");
} }
@@ -1095,9 +1083,7 @@ mod tests {
engine.refresh_event_cache().await; engine.refresh_event_cache().await;
let msg = IncomingMessage::new("test", "default", "DISABLE_ME"); let msg = IncomingMessage::new("test", "default", "DISABLE_ME");
let fired_before = engine let fired_before = engine.check_event_triggers(&msg, &msg.content).await;
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await;
assert!(fired_before >= 1, "Expected routine to fire before disable"); assert!(fired_before >= 1, "Expected routine to fire before disable");
// Simulate what routines_toggle_handler now does: update DB, then refresh. // Simulate what routines_toggle_handler now does: update DB, then refresh.
@@ -1106,9 +1092,7 @@ mod tests {
db.update_routine(&routine).await.expect("update_routine"); db.update_routine(&routine).await.expect("update_routine");
engine.refresh_event_cache().await; engine.refresh_event_cache().await;
let fired_after = engine let fired_after = engine.check_event_triggers(&msg, &msg.content).await;
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await;
assert_eq!( assert_eq!(
fired_after, 0, fired_after, 0,
"Disabled routine must not fire after cache refresh" "Disabled routine must not fire after cache refresh"
@@ -1134,10 +1118,7 @@ mod tests {
let msg = IncomingMessage::new("test", "default", "DELETE_ME"); let msg = IncomingMessage::new("test", "default", "DELETE_ME");
assert!( assert!(
engine engine.check_event_triggers(&msg, &msg.content).await >= 1,
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await
>= 1,
"Expected routine to fire before delete" "Expected routine to fire before delete"
); );
@@ -1146,9 +1127,7 @@ mod tests {
engine.refresh_event_cache().await; engine.refresh_event_cache().await;
assert_eq!( assert_eq!(
engine engine.check_event_triggers(&msg, &msg.content).await,
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
.await,
0, 0,
"Deleted routine must not fire after cache refresh" "Deleted routine must not fire after cache refresh"
); );
@@ -1462,8 +1441,9 @@ mod tests {
db.create_routine(&routine).await.expect("create_routine"); db.create_routine(&routine).await.expect("create_routine");
engine.refresh_event_cache().await; engine.refresh_event_cache().await;
let trigger_msg = IncomingMessage::new("test", "default", "owner-gate");
let fired = engine let fired = engine
.check_event_triggers("default", "test", "owner-gate") .check_event_triggers(&trigger_msg, &trigger_msg.content)
.await; .await;
assert_eq!(fired, 1, "expected one matching event routine"); assert_eq!(fired, 1, "expected one matching event routine");
+142 -27
View File
@@ -19,10 +19,13 @@ use axum::middleware;
use axum::routing::{get, post}; use axum::routing::{get, post};
use tower::ServiceExt; use tower::ServiceExt;
use ironclaw::channels::IncomingMessage;
use ironclaw::channels::web::auth::{ use ironclaw::channels::web::auth::{
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware, AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
}; };
use ironclaw::channels::web::server::{GatewayState, PerUserRateLimiter, RateLimiter}; use ironclaw::channels::web::server::{
GatewayState, PerUserRateLimiter, RateLimiter, start_server,
};
use ironclaw::channels::web::sse::SseManager; use ironclaw::channels::web::sse::SseManager;
use ironclaw::channels::web::test_helpers::TestGatewayBuilder; use ironclaw::channels::web::test_helpers::TestGatewayBuilder;
use ironclaw::channels::web::ws::WsConnectionTracker; use ironclaw::channels::web::ws::WsConnectionTracker;
@@ -37,6 +40,9 @@ const ALICE_TOKEN: &str = "tok-alice-secret";
const BOB_TOKEN: &str = "tok-bob-secret"; const BOB_TOKEN: &str = "tok-bob-secret";
const ALICE_USER_ID: &str = "alice"; const ALICE_USER_ID: &str = "alice";
const BOB_USER_ID: &str = "bob"; const BOB_USER_ID: &str = "bob";
const OWNER_TOKEN: &str = "tok-owner-secret";
const OWNER_SCOPE_ID: &str = "owner-scope";
const GATEWAY_SENDER_ID: &str = "gateway-sender";
/// Build a MultiAuthState with two users. /// Build a MultiAuthState with two users.
fn two_user_auth() -> MultiAuthState { fn two_user_auth() -> MultiAuthState {
@@ -301,7 +307,7 @@ fn per_user_rate_limiter_single_user_mode() {
#[tokio::test] #[tokio::test]
async fn sse_scoped_event_only_delivered_to_target_user() { async fn sse_scoped_event_only_delivered_to_target_user() {
use ironclaw::channels::web::types::SseEvent; use ironclaw_common::AppEvent;
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
let manager = SseManager::new(); let manager = SseManager::new();
@@ -319,34 +325,34 @@ async fn sse_scoped_event_only_delivered_to_target_user() {
// Send event scoped to alice // Send event scoped to alice
manager.broadcast_for_user( manager.broadcast_for_user(
ALICE_USER_ID, ALICE_USER_ID,
SseEvent::Status { AppEvent::Status {
message: "alice's event".to_string(), message: "alice's event".to_string(),
thread_id: None, thread_id: None,
}, },
); );
// Send global heartbeat (both should get it) // Send global heartbeat (both should get it)
manager.broadcast(SseEvent::Heartbeat); manager.broadcast(AppEvent::Heartbeat);
// Alice gets her scoped event first // Alice gets her scoped event first
let e = alice_stream.next().await.unwrap(); let e = alice_stream.next().await.unwrap();
match &e { match &e {
SseEvent::Status { message, .. } => assert_eq!(message, "alice's event"), AppEvent::Status { message, .. } => assert_eq!(message, "alice's event"),
_ => panic!("Expected Status, got {:?}", e), _ => panic!("Expected Status, got {:?}", e),
} }
// Alice also gets heartbeat // Alice also gets heartbeat
let e = alice_stream.next().await.unwrap(); let e = alice_stream.next().await.unwrap();
assert!(matches!(e, SseEvent::Heartbeat)); assert!(matches!(e, AppEvent::Heartbeat));
// Bob only gets the heartbeat (alice's event was filtered) // Bob only gets the heartbeat (alice's event was filtered)
let e = bob_stream.next().await.unwrap(); let e = bob_stream.next().await.unwrap();
assert!(matches!(e, SseEvent::Heartbeat)); assert!(matches!(e, AppEvent::Heartbeat));
} }
#[tokio::test] #[tokio::test]
async fn sse_global_event_delivered_to_all_users() { async fn sse_global_event_delivered_to_all_users() {
use ironclaw::channels::web::types::SseEvent; use ironclaw_common::AppEvent;
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
let manager = SseManager::new(); let manager = SseManager::new();
@@ -361,7 +367,7 @@ async fn sse_global_event_delivered_to_all_users() {
.expect("subscribe"), .expect("subscribe"),
); );
manager.broadcast(SseEvent::Status { manager.broadcast(AppEvent::Status {
message: "global announcement".to_string(), message: "global announcement".to_string(),
thread_id: None, thread_id: None,
}); });
@@ -369,7 +375,7 @@ async fn sse_global_event_delivered_to_all_users() {
let ea = alice.next().await.unwrap(); let ea = alice.next().await.unwrap();
let eb = bob.next().await.unwrap(); let eb = bob.next().await.unwrap();
match (&ea, &eb) { match (&ea, &eb) {
(SseEvent::Status { message: a, .. }, SseEvent::Status { message: b, .. }) => { (AppEvent::Status { message: a, .. }, AppEvent::Status { message: b, .. }) => {
assert_eq!(a, "global announcement"); assert_eq!(a, "global announcement");
assert_eq!(b, "global announcement"); assert_eq!(b, "global announcement");
} }
@@ -379,7 +385,7 @@ async fn sse_global_event_delivered_to_all_users() {
#[tokio::test] #[tokio::test]
async fn sse_user_b_event_not_visible_to_user_a() { async fn sse_user_b_event_not_visible_to_user_a() {
use ironclaw::channels::web::types::SseEvent; use ironclaw_common::AppEvent;
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
let manager = SseManager::new(); let manager = SseManager::new();
@@ -392,19 +398,19 @@ async fn sse_user_b_event_not_visible_to_user_a() {
// Send event for bob only // Send event for bob only
manager.broadcast_for_user( manager.broadcast_for_user(
BOB_USER_ID, BOB_USER_ID,
SseEvent::Response { AppEvent::Response {
content: "bob's secret".to_string(), content: "bob's secret".to_string(),
thread_id: "t1".to_string(), thread_id: "t1".to_string(),
}, },
); );
// Send heartbeat so alice has something to receive // Send heartbeat so alice has something to receive
manager.broadcast(SseEvent::Heartbeat); manager.broadcast(AppEvent::Heartbeat);
// Alice should only get heartbeat, not bob's response // Alice should only get heartbeat, not bob's response
let e = alice.next().await.unwrap(); let e = alice.next().await.unwrap();
assert!( assert!(
matches!(e, SseEvent::Heartbeat), matches!(e, AppEvent::Heartbeat),
"Expected Heartbeat, got {:?}", "Expected Heartbeat, got {:?}",
e e
); );
@@ -412,7 +418,7 @@ async fn sse_user_b_event_not_visible_to_user_a() {
#[tokio::test] #[tokio::test]
async fn sse_unscoped_subscriber_receives_all_events() { async fn sse_unscoped_subscriber_receives_all_events() {
use ironclaw::channels::web::types::SseEvent; use ironclaw_common::AppEvent;
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
let manager = SseManager::new(); let manager = SseManager::new();
@@ -421,19 +427,19 @@ async fn sse_unscoped_subscriber_receives_all_events() {
manager.broadcast_for_user( manager.broadcast_for_user(
ALICE_USER_ID, ALICE_USER_ID,
SseEvent::Status { AppEvent::Status {
message: "alice only".to_string(), message: "alice only".to_string(),
thread_id: None, thread_id: None,
}, },
); );
manager.broadcast_for_user( manager.broadcast_for_user(
BOB_USER_ID, BOB_USER_ID,
SseEvent::Status { AppEvent::Status {
message: "bob only".to_string(), message: "bob only".to_string(),
thread_id: None, thread_id: None,
}, },
); );
manager.broadcast(SseEvent::Heartbeat); manager.broadcast(AppEvent::Heartbeat);
// Unscoped subscriber gets ALL three events // Unscoped subscriber gets ALL three events
let e1 = stream.next().await.unwrap(); let e1 = stream.next().await.unwrap();
@@ -441,14 +447,14 @@ async fn sse_unscoped_subscriber_receives_all_events() {
let e3 = stream.next().await.unwrap(); let e3 = stream.next().await.unwrap();
match &e1 { match &e1 {
SseEvent::Status { message, .. } => assert_eq!(message, "alice only"), AppEvent::Status { message, .. } => assert_eq!(message, "alice only"),
_ => panic!("Expected alice's Status"), _ => panic!("Expected alice's Status"),
} }
match &e2 { match &e2 {
SseEvent::Status { message, .. } => assert_eq!(message, "bob only"), AppEvent::Status { message, .. } => assert_eq!(message, "bob only"),
_ => panic!("Expected bob's Status"), _ => panic!("Expected bob's Status"),
} }
assert!(matches!(e3, SseEvent::Heartbeat)); assert!(matches!(e3, AppEvent::Heartbeat));
} }
// =========================================================================== // ===========================================================================
@@ -537,7 +543,8 @@ fn gateway_state_has_multi_tenant_fields() {
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: None, scheduler: None,
default_user_id: "fallback".to_string(), // Multi-tenant: renamed from user_id owner_id: "fallback".to_string(),
default_sender_id: "fallback".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), 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,
@@ -553,7 +560,8 @@ fn gateway_state_has_multi_tenant_fields() {
active_config: Default::default(), active_config: Default::default(),
}; };
assert_eq!(state.default_user_id, "fallback"); assert_eq!(state.owner_id, "fallback");
assert_eq!(state.default_sender_id, "fallback");
assert!(state.workspace_pool.is_none()); assert!(state.workspace_pool.is_none());
} }
@@ -572,6 +580,69 @@ async fn start_multi_user_server() -> (SocketAddr, Arc<GatewayState>) {
.expect("Failed to start multi-user test server") .expect("Failed to start multi-user test server")
} }
async fn start_owner_scoped_sender_server() -> (
SocketAddr,
Arc<GatewayState>,
tokio::sync::mpsc::Receiver<IncomingMessage>,
) {
let (agent_tx, agent_rx) = tokio::sync::mpsc::channel(64);
let mut tokens = HashMap::new();
tokens.insert(
OWNER_TOKEN.to_string(),
UserIdentity {
user_id: OWNER_SCOPE_ID.to_string(),
workspace_read_scopes: Vec::new(),
},
);
tokens.insert(
BOB_TOKEN.to_string(),
UserIdentity {
user_id: BOB_USER_ID.to_string(),
workspace_read_scopes: Vec::new(),
},
);
let state = Arc::new(GatewayState {
msg_tx: tokio::sync::RwLock::new(Some(agent_tx)),
sse: Arc::new(SseManager::new()),
workspace: None,
workspace_pool: None,
session_manager: None,
log_broadcaster: None,
log_level_handle: None,
extension_manager: None,
tool_registry: None,
store: None,
job_manager: None,
prompt_queue: None,
scheduler: None,
owner_id: OWNER_SCOPE_ID.to_string(),
default_sender_id: GATEWAY_SENDER_ID.to_string(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
llm_provider: None,
skill_registry: None,
skill_catalog: None,
chat_rate_limiter: PerUserRateLimiter::new(30, 60),
oauth_rate_limiter: RateLimiter::new(10, 60),
webhook_rate_limiter: RateLimiter::new(10, 60),
registry_entries: Vec::new(),
cost_guard: None,
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
startup_time: std::time::Instant::now(),
active_config: Default::default(),
});
let auth = MultiAuthState::multi(tokens);
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
let bound = start_server(addr, state.clone(), auth)
.await
.expect("Failed to start owner-scoped sender test server");
(bound, state, agent_rx)
}
#[tokio::test] #[tokio::test]
async fn full_server_alice_can_access_protected_endpoint() { async fn full_server_alice_can_access_protected_endpoint() {
let (addr, _state) = start_multi_user_server().await; let (addr, _state) = start_multi_user_server().await;
@@ -677,6 +748,49 @@ async fn full_server_chat_send_accepted_for_alice() {
assert_eq!(msg.channel, "gateway"); assert_eq!(msg.channel, "gateway");
} }
#[tokio::test]
async fn full_server_chat_send_rewrites_sender_only_for_owner_scope_rebind() {
let (addr, _state, mut agent_rx) = start_owner_scoped_sender_server().await;
let client = reqwest::Client::new();
let owner_resp = client
.post(format!("http://{}/api/chat/send", addr))
.header("Authorization", format!("Bearer {}", OWNER_TOKEN))
.header("Content-Type", "application/json")
.body(r#"{"content":"hello from owner"}"#)
.send()
.await
.unwrap();
assert_eq!(owner_resp.status(), 202);
let owner_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
.await
.expect("Timed out waiting for owner message")
.expect("Agent channel closed");
assert_eq!(owner_msg.user_id, OWNER_SCOPE_ID);
assert_eq!(owner_msg.sender_id, GATEWAY_SENDER_ID);
assert_eq!(owner_msg.content, "hello from owner");
let other_resp = client
.post(format!("http://{}/api/chat/send", addr))
.header("Authorization", format!("Bearer {}", BOB_TOKEN))
.header("Content-Type", "application/json")
.body(r#"{"content":"hello from bob"}"#)
.send()
.await
.unwrap();
assert_eq!(other_resp.status(), 202);
let other_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
.await
.expect("Timed out waiting for non-owner message")
.expect("Agent channel closed");
assert_eq!(other_msg.user_id, BOB_USER_ID);
assert_eq!(other_msg.sender_id, BOB_USER_ID);
assert_eq!(other_msg.content, "hello from bob");
}
#[tokio::test] #[tokio::test]
async fn full_server_chat_send_rejected_without_auth() { async fn full_server_chat_send_rejected_without_auth() {
let (addr, _state) = start_multi_user_server().await; let (addr, _state) = start_multi_user_server().await;
@@ -767,7 +881,7 @@ async fn full_server_jobs_endpoint_rejected_without_auth() {
#[tokio::test] #[tokio::test]
async fn full_server_ws_multi_user_event_isolation() { async fn full_server_ws_multi_user_event_isolation() {
use futures::StreamExt; use futures::StreamExt;
use ironclaw::channels::web::types::SseEvent; use ironclaw_common::AppEvent;
use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::client::IntoClientRequest;
@@ -800,14 +914,14 @@ async fn full_server_ws_multi_user_event_isolation() {
// Broadcast an event scoped to Alice only // Broadcast an event scoped to Alice only
state.sse.broadcast_for_user( state.sse.broadcast_for_user(
ALICE_USER_ID, ALICE_USER_ID,
SseEvent::Status { AppEvent::Status {
message: "alice-only-event".to_string(), message: "alice-only-event".to_string(),
thread_id: None, thread_id: None,
}, },
); );
// Broadcast a global heartbeat so Bob has something to receive // Broadcast a global heartbeat so Bob has something to receive
state.sse.broadcast(SseEvent::Heartbeat); state.sse.broadcast(AppEvent::Heartbeat);
// Alice should get her scoped event // Alice should get her scoped event
let alice_msg = tokio::time::timeout(Duration::from_secs(2), alice_ws.next()) let alice_msg = tokio::time::timeout(Duration::from_secs(2), alice_ws.next())
@@ -888,7 +1002,8 @@ async fn start_multi_user_server_with_db() -> (
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: None, scheduler: None,
default_user_id: ALICE_USER_ID.to_string(), owner_id: ALICE_USER_ID.to_string(),
default_sender_id: ALICE_USER_ID.to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), 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,
+5 -2
View File
@@ -94,6 +94,7 @@ impl LlmProvider for MockLlmProvider {
id: "call_mock_001".to_string(), id: "call_mock_001".to_string(),
name: tool.name.clone(), name: tool.name.clone(),
arguments: serde_json::json!({"test": true}), arguments: serde_json::json!({"test": true}),
reasoning: None,
}], }],
input_tokens: 15, input_tokens: 15,
output_tokens: 8, output_tokens: 8,
@@ -203,7 +204,8 @@ async fn start_test_server_with_provider(
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: None, scheduler: None,
default_user_id: "test-user".to_string(), owner_id: "test-user".to_string(),
default_sender_id: "test-user".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())), ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
llm_provider: Some(llm_provider), llm_provider: Some(llm_provider),
@@ -701,7 +703,8 @@ async fn test_no_llm_provider_returns_503() {
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: None, scheduler: None,
default_user_id: "test-user".to_string(), owner_id: "test-user".to_string(),
default_sender_id: "test-user".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())), ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
llm_provider: None, // No LLM! llm_provider: None, // No LLM!
+2 -1
View File
@@ -226,7 +226,8 @@ impl GatewayWorkflowHarness {
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: Some(scheduler_slot.clone()), scheduler: Some(scheduler_slot.clone()),
default_user_id: user_id.clone(), owner_id: user_id.clone(),
default_sender_id: user_id.clone(),
shutdown_tx: tokio::sync::RwLock::new(None), shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())), ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
llm_provider: Some(Arc::clone(&components.llm)), llm_provider: Some(Arc::clone(&components.llm)),
+1
View File
@@ -566,6 +566,7 @@ impl LlmProvider for TraceLlm {
id: tc.id, id: tc.id,
name: tc.name, name: tc.name,
arguments: tc.arguments, arguments: tc.arguments,
reasoning: None,
}) })
.collect(); .collect();
Ok(ToolCompletionResponse { Ok(ToolCompletionResponse {
+11 -10
View File
@@ -5,7 +5,7 @@
//! - WebSocket upgrade with auth //! - WebSocket upgrade with auth
//! - Ping/pong //! - Ping/pong
//! - Client message → agent msg_tx //! - Client message → agent msg_tx
//! - Broadcast SSE event → WebSocket client //! - Broadcast AppEvent → WebSocket client
//! - Connection tracking (counter increment/decrement) //! - Connection tracking (counter increment/decrement)
//! - Gateway status endpoint //! - Gateway status endpoint
@@ -22,8 +22,8 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use ironclaw::channels::IncomingMessage; use ironclaw::channels::IncomingMessage;
use ironclaw::channels::web::server::{GatewayState, start_server}; use ironclaw::channels::web::server::{GatewayState, start_server};
use ironclaw::channels::web::sse::SseManager; use ironclaw::channels::web::sse::SseManager;
use ironclaw::channels::web::types::SseEvent;
use ironclaw::channels::web::ws::WsConnectionTracker; use ironclaw::channels::web::ws::WsConnectionTracker;
use ironclaw_common::AppEvent;
const AUTH_TOKEN: &str = "test-token-12345"; const AUTH_TOKEN: &str = "test-token-12345";
const TIMEOUT: Duration = Duration::from_secs(5); const TIMEOUT: Duration = Duration::from_secs(5);
@@ -51,7 +51,8 @@ async fn start_test_server() -> (
job_manager: None, job_manager: None,
prompt_queue: None, prompt_queue: None,
scheduler: None, scheduler: None,
default_user_id: "test-user".to_string(), owner_id: "test-user".to_string(),
default_sender_id: "test-user".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None), 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,
@@ -163,8 +164,8 @@ async fn test_ws_broadcast_event_received() {
// Give the connection a moment to fully establish // Give the connection a moment to fully establish
tokio::time::sleep(Duration::from_millis(50)).await; tokio::time::sleep(Duration::from_millis(50)).await;
// Broadcast an SSE event (simulates agent sending a response) // Broadcast an event (simulates agent sending a response)
state.sse.broadcast(SseEvent::Response { state.sse.broadcast(AppEvent::Response {
content: "agent says hi".to_string(), content: "agent says hi".to_string(),
thread_id: "t1".to_string(), thread_id: "t1".to_string(),
}); });
@@ -185,7 +186,7 @@ async fn test_ws_thinking_event() {
let mut ws = connect_ws(addr).await; let mut ws = connect_ws(addr).await;
tokio::time::sleep(Duration::from_millis(50)).await; tokio::time::sleep(Duration::from_millis(50)).await;
state.sse.broadcast(SseEvent::Thinking { state.sse.broadcast(AppEvent::Thinking {
message: "analyzing...".to_string(), message: "analyzing...".to_string(),
thread_id: None, thread_id: None,
}); });
@@ -310,22 +311,22 @@ async fn test_ws_multiple_events_in_sequence() {
tokio::time::sleep(Duration::from_millis(50)).await; tokio::time::sleep(Duration::from_millis(50)).await;
// Broadcast multiple events rapidly // Broadcast multiple events rapidly
state.sse.broadcast(SseEvent::Thinking { state.sse.broadcast(AppEvent::Thinking {
message: "step 1".to_string(), message: "step 1".to_string(),
thread_id: None, thread_id: None,
}); });
state.sse.broadcast(SseEvent::ToolStarted { state.sse.broadcast(AppEvent::ToolStarted {
name: "shell".to_string(), name: "shell".to_string(),
thread_id: None, thread_id: None,
}); });
state.sse.broadcast(SseEvent::ToolCompleted { state.sse.broadcast(AppEvent::ToolCompleted {
name: "shell".to_string(), name: "shell".to_string(),
success: true, success: true,
error: None, error: None,
parameters: None, parameters: None,
thread_id: None, thread_id: None,
}); });
state.sse.broadcast(SseEvent::Response { state.sse.broadcast(AppEvent::Response {
content: "done".to_string(), content: "done".to_string(),
thread_id: "t1".to_string(), thread_id: "t1".to_string(),
}); });