mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 07:30:11 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c1ca3bb91c | ||
|
|
e499795b8c | ||
|
|
5e1da4827a | ||
|
|
d04af5cd75 | ||
|
|
956037c4d3 | ||
|
|
68a1851c19 | ||
|
|
dfa105539b | ||
|
|
8929baf76a | ||
|
|
e07dfab449 | ||
|
|
6783cba4e4 | ||
|
|
63302ab406 | ||
|
|
7c553b0973 | ||
|
|
5e44185e48 | ||
|
|
72623c9e5b | ||
|
|
6895adbcc9 | ||
|
|
f1480f471b | ||
|
|
9db949746f | ||
|
|
61a123a746 | ||
|
|
0e981429ee | ||
|
|
1b38a64e15 | ||
|
|
2e5f8b60d5 | ||
|
|
f0a0642e7d | ||
|
|
ca8d5c6b5e | ||
|
|
9fed8453c7 | ||
|
|
eaef335db6 | ||
|
|
225af29db2 | ||
|
|
a53b2c10b5 | ||
|
|
408ae8a29a | ||
|
|
d9ff86d7e0 | ||
|
|
e843c18141 | ||
|
|
54e9206f0b | ||
|
|
5df0d13b59 | ||
|
|
bbb68f7490 | ||
|
|
b3dee13954 | ||
|
|
33ef0a6ea5 | ||
|
|
e0a43c81f9 | ||
|
|
bada79ba4a | ||
|
|
a70c89d9e3 |
@@ -0,0 +1,97 @@
|
||||
---
|
||||
description: Fetch a GitHub issue, create a branch, research the codebase, plan the fix, implement with tests, and commit
|
||||
disable-model-invocation: true
|
||||
allowed-tools: Bash(gh issue view:*), Bash(gh repo view:*), Bash(git fetch:*), Bash(git checkout:*), Bash(git status:*), Bash(git branch:*), Bash(git add:*), Bash(git commit:*), Bash(cargo fmt:*), Bash(cargo clippy:*), Bash(cargo test:*), Read, Edit, Write, Grep, Glob
|
||||
argument-hint: "<issue-number or github-issue-url>"
|
||||
---
|
||||
|
||||
# Fix GitHub Issue
|
||||
|
||||
## Step 1: Resolve the issue
|
||||
|
||||
Parse `$ARGUMENTS` to extract the issue number:
|
||||
- If it's a URL like `https://github.com/owner/repo/issues/42`, extract `42`.
|
||||
- If it's a bare number, use it directly.
|
||||
- If empty, stop and ask the user for an issue number.
|
||||
|
||||
Fetch the issue:
|
||||
|
||||
```
|
||||
gh issue view {number} --json title,body,labels,assignees,comments,state
|
||||
```
|
||||
|
||||
If the issue is closed, warn the user and ask if they still want to proceed.
|
||||
|
||||
## Step 2: Create a branch
|
||||
|
||||
Create a fresh branch off the latest main:
|
||||
|
||||
1. Fetch latest: `git fetch origin`
|
||||
2. Detect default branch: `gh repo view --json defaultBranchRef --jq .defaultBranchRef.name`
|
||||
3. Create and switch to a new branch: `git checkout -b fix/{number}-{short-slug} origin/{default-branch}`
|
||||
- `{short-slug}` is 3-5 words from the issue title, lowercase, hyphenated (e.g. `fix/42-idor-workspace-check`)
|
||||
|
||||
If the working tree has uncommitted changes, warn the user and stop. Do not stash or discard their work.
|
||||
|
||||
## Step 3: Understand the issue
|
||||
|
||||
Summarize the issue in 2-3 sentences. Identify:
|
||||
- **What's broken or missing** (the symptom or feature request)
|
||||
- **Acceptance criteria** (what "done" looks like, from the issue body or comments)
|
||||
- **Constraints** (mentioned technologies, backward compatibility, performance requirements)
|
||||
|
||||
If the issue is unclear or ambiguous, list the open questions. These will be addressed during planning.
|
||||
|
||||
## Step 4: Research the codebase
|
||||
|
||||
Before planning, gather context:
|
||||
|
||||
1. **Find relevant code** - Search for files, functions, types, and patterns mentioned in the issue. Read them in full.
|
||||
2. **Trace the flow** - If the issue is about a specific behavior, trace the code path from the entry point (route handler, CLI command, etc.) through to the relevant logic.
|
||||
3. **Check existing tests** - Find tests related to the affected code. Understand what's already covered.
|
||||
4. **Check for prior art** - Look for similar patterns in the codebase that solve analogous problems. Prefer consistency with existing patterns.
|
||||
|
||||
## Step 5: Enter planning mode
|
||||
|
||||
Enter planning mode to design the implementation. The plan MUST cover:
|
||||
|
||||
1. **Root cause** (for bugs) or **design approach** (for features)
|
||||
2. **Files to modify** with specific descriptions of what changes in each
|
||||
3. **New files** (if any) with justification for why they're needed
|
||||
4. **Tests to add** - every code path introduced or changed needs a test:
|
||||
- Happy path (expected input produces expected output)
|
||||
- Error paths (invalid input, missing data, permission denied)
|
||||
- Edge cases (empty collections, boundary values, concurrent access)
|
||||
5. **IronClaw-specific concerns**:
|
||||
- If the change touches persistence, both database backends must be updated (`postgres.rs` and `libsql_backend.rs`)
|
||||
- New `Database` trait methods need implementations in both backends
|
||||
- No `.unwrap()` or `.expect()` in production code
|
||||
- Use `crate::` imports, not `super::`
|
||||
- Error types via `thiserror` in `error.rs`
|
||||
6. **Migration or compatibility concerns** (if any)
|
||||
|
||||
Follow the project's CLAUDE.md guidance for architecture decisions.
|
||||
|
||||
Wait for user approval before implementing.
|
||||
|
||||
## Step 6: Implement
|
||||
|
||||
After the plan is approved:
|
||||
|
||||
1. Implement each change from the plan.
|
||||
2. Write all planned tests.
|
||||
3. Run IronClaw's full quality gate:
|
||||
- `cargo fmt`
|
||||
- `cargo clippy --all --benches --tests --examples --all-features` (zero warnings)
|
||||
- `cargo test --lib` (all tests pass)
|
||||
4. If any check fails, fix it before proceeding.
|
||||
|
||||
Note: Integration tests (`--test workspace_integration`) require PostgreSQL and are expected to fail locally. Only `--lib` test failures are blocking.
|
||||
|
||||
## Step 7: Commit and summarize
|
||||
|
||||
1. Commit with a descriptive message referencing the issue (e.g. `fix: prevent IDOR in function call outputs (#42)`).
|
||||
2. Summarize what was done:
|
||||
- Files changed with line references
|
||||
- Tests added and what they cover
|
||||
- Any follow-up work or open questions
|
||||
@@ -0,0 +1,81 @@
|
||||
---
|
||||
description: Respond to PR review comments — triage, plan fixes, implement after confirmation, push, and reply to reviewers
|
||||
disable-model-invocation: true
|
||||
allowed-tools: Bash(gh pr list:*), Bash(gh pr comment:*), Bash(gh api:*), Bash(gh repo view:*), Bash(git branch:*), Bash(git status:*), Bash(git add:*), Bash(git commit:*), Bash(git push:*), Bash(cargo fmt:*), Bash(cargo clippy:*), Bash(cargo test:*), Read, Edit, Write, Grep, Glob
|
||||
argument-hint: "[pr-number (optional, auto-detects from branch)]"
|
||||
---
|
||||
|
||||
# Review and Address PR Comments
|
||||
|
||||
## Step 1: Find the PR
|
||||
|
||||
If `$ARGUMENTS` is provided, use that as the PR number. Otherwise, detect the PR for the current branch:
|
||||
|
||||
```
|
||||
gh pr list --head $(git branch --show-current) --json number,title,url --jq '.[0]'
|
||||
```
|
||||
|
||||
If no PR is found, tell the user and stop.
|
||||
|
||||
## Step 2: Fetch all review comments
|
||||
|
||||
Resolve the repo owner and name:
|
||||
|
||||
```
|
||||
gh repo view --json owner,name --jq '"\(.owner.login)/\(.name)"'
|
||||
```
|
||||
|
||||
Fetch the full set of review comments (not issue-level comments):
|
||||
|
||||
```
|
||||
gh api --paginate repos/{owner}/{repo}/pulls/{number}/comments
|
||||
```
|
||||
|
||||
Also fetch the review summaries:
|
||||
|
||||
```
|
||||
gh api --paginate repos/{owner}/{repo}/pulls/{number}/reviews
|
||||
```
|
||||
|
||||
Deduplicate comments that appear multiple times (bots sometimes post the same finding under different IDs). Group by the actual issue being raised, not by comment ID.
|
||||
|
||||
## Step 3: Triage and plan
|
||||
|
||||
For each unique issue raised in the comments:
|
||||
|
||||
1. **Check if already addressed** - Read the current code at the referenced location. If a prior commit already fixed it, note it as "already resolved".
|
||||
2. **Assess validity** - Determine if the comment identifies a real problem or is a false positive. Be honest about false positives but explain why.
|
||||
3. **Classify severity** - Critical (security/data loss), High (bugs/broken behavior), Medium (correctness/robustness), Low (style/naming/nits).
|
||||
4. **Plan the fix** - For each valid unresolved issue, describe the specific code change needed.
|
||||
|
||||
Present the plan as a table to the user:
|
||||
|
||||
| # | Issue | File:Line | Severity | Status | Planned Fix |
|
||||
|---|-------|-----------|----------|--------|-------------|
|
||||
|
||||
Wait for user confirmation before proceeding to implementation.
|
||||
|
||||
## Step 4: Implement fixes
|
||||
|
||||
After user confirms:
|
||||
|
||||
1. Implement each fix in the plan.
|
||||
2. Run IronClaw's quality gate to verify nothing breaks:
|
||||
- `cargo fmt`
|
||||
- `cargo clippy --all --benches --tests --examples --all-features`
|
||||
- `cargo test --lib`
|
||||
3. Commit with a descriptive message referencing the PR review.
|
||||
4. Push to the branch.
|
||||
|
||||
## Step 5: Reply to comments
|
||||
|
||||
For each comment addressed, reply on the PR with a short message stating what was fixed and the commit SHA. For false positives or already-resolved items, reply explaining why no change was needed.
|
||||
|
||||
## Rules
|
||||
|
||||
- Never guess at code you haven't read. Always read the referenced file and line before assessing a comment.
|
||||
- Group duplicate comments (same issue reported by multiple bots) and reply to all of them.
|
||||
- Do not make changes beyond what the review comments ask for. Stay focused.
|
||||
- If a comment suggests a change you disagree with, present your reasoning to the user during the planning phase rather than silently ignoring it.
|
||||
- Follow IronClaw conventions: no `.unwrap()` in production code, use `crate::` imports, `thiserror` errors.
|
||||
- If changes touch persistence, verify both database backends are updated.
|
||||
@@ -0,0 +1,245 @@
|
||||
---
|
||||
description: Deep audit of the IronClaw crate for vulnerabilities, bugs, unfinished work, inconsistencies, and oversights
|
||||
disable-model-invocation: true
|
||||
allowed-tools: Bash(cargo fmt:*), Bash(cargo clippy:*), Bash(cargo test:*), Bash(cargo audit:*), Bash(git diff:*), Bash(git log:*), Bash(git show:*), Bash(wc:*), Read, Grep, Glob, Task
|
||||
argument-hint: "[path/to/crate]"
|
||||
---
|
||||
|
||||
# Rust Crate Audit
|
||||
|
||||
You are performing a thorough audit of a Rust crate. Your goal is to find every vulnerability, bug, unfinished piece of work, inconsistency, and oversight before it ships. Leave no stone unturned.
|
||||
|
||||
## Step 1: Locate the crate
|
||||
|
||||
Parse `$ARGUMENTS`:
|
||||
- If a path is provided, use it as the crate root.
|
||||
- If empty, use the current working directory.
|
||||
|
||||
Verify it's a valid Rust crate by checking for `Cargo.toml`. If not found, stop and ask the user.
|
||||
|
||||
## Step 2: Understand the crate
|
||||
|
||||
Read `Cargo.toml` to understand:
|
||||
- Crate name, version, edition
|
||||
- Dependencies (look for outdated, unmaintained, or suspicious crates)
|
||||
- Feature flags and their implications
|
||||
- Build scripts (`build.rs`) if any
|
||||
|
||||
Read `CLAUDE.md`, `README.md`, or top-level documentation if present to understand intent and architecture.
|
||||
|
||||
Read `src/lib.rs` or `src/main.rs` to get the module tree. Then read each module's `mod.rs` or top-level file to build a mental map of the crate's structure before diving into details.
|
||||
Read all Rust files (`src/*.rs`) to make sure everything is in context when you are reasoning.
|
||||
|
||||
## Step 3: Run the compiler's checks
|
||||
|
||||
Run these commands and capture output. Do NOT fix anything, just collect findings:
|
||||
|
||||
```
|
||||
cargo fmt --check 2>&1
|
||||
```
|
||||
|
||||
```
|
||||
cargo clippy --all --benches --tests --examples --all-features -- -W clippy::all -W clippy::pedantic -W clippy::nursery 2>&1
|
||||
```
|
||||
|
||||
```
|
||||
cargo test --lib 2>&1
|
||||
```
|
||||
|
||||
If any of these fail, record the failures as findings. If `cargo test` has ignored tests, note which ones and why.
|
||||
|
||||
Note: Integration tests (`--test workspace_integration`) require a PostgreSQL database and are expected to fail locally. Only report `--lib` test failures as blocking.
|
||||
|
||||
## Step 4: Scan for unfinished work
|
||||
|
||||
Search the entire `src/` tree for:
|
||||
|
||||
```
|
||||
todo!
|
||||
unimplemented!
|
||||
fixme
|
||||
FIXME
|
||||
TODO
|
||||
HACK
|
||||
XXX
|
||||
SAFETY:
|
||||
stub
|
||||
placeholder
|
||||
temporary
|
||||
```
|
||||
|
||||
For each match:
|
||||
- Is it in production code or test code?
|
||||
- Is it a genuine incomplete feature or a deliberate placeholder?
|
||||
- Is there a tracking issue referenced?
|
||||
- Could this panic at runtime?
|
||||
|
||||
Any `todo!()` or `unimplemented!()` in non-test code is **High severity** (runtime panic).
|
||||
|
||||
## Step 5: Audit for vulnerabilities and unsafe code
|
||||
|
||||
### 5a. Unsafe code
|
||||
|
||||
Search for all `unsafe` blocks. For each one:
|
||||
- Is the safety invariant documented with a `// SAFETY:` comment?
|
||||
- Is the invariant actually upheld by the surrounding code?
|
||||
- Could the unsafe block be replaced with a safe alternative?
|
||||
- Are there any pointer dereferences, transmutes, or FFI calls?
|
||||
|
||||
### 5b. Unwrap and panic paths
|
||||
|
||||
Search for `.unwrap()`, `.expect(`, `panic!`, `unreachable!` in non-test code. For each:
|
||||
- Can this actually panic in production?
|
||||
- Is there a code path that reaches this with None/Err?
|
||||
- Should it be replaced with proper error handling (`?`, `.ok()`, `.unwrap_or_default()`)?
|
||||
|
||||
IronClaw convention: `.unwrap()` and `.expect()` are banned in production code. Any occurrence outside `#[cfg(test)]` blocks is a **High severity** finding.
|
||||
|
||||
### 5c. SQL and injection vectors
|
||||
|
||||
Search for string formatting used in SQL queries, shell commands, or HTML:
|
||||
- `format!` used near `.execute(`, `.query(`, `Command::new(`
|
||||
- String interpolation in query construction vs parameterized queries
|
||||
- User input flowing into file paths (`Path::new`, `std::fs::`)
|
||||
|
||||
IronClaw has two database backends (PostgreSQL and libSQL). Check both for injection vectors.
|
||||
|
||||
### 5d. Cryptographic issues
|
||||
|
||||
If the crate uses crypto:
|
||||
- Are comparisons constant-time? (look for `==` on secrets/hashes vs `subtle::ConstantTimeEq`)
|
||||
- Is randomness from `OsRng` / `thread_rng` and not a fixed seed?
|
||||
- Are keys/secrets zeroized after use? (`secrecy`, `zeroize` crates)
|
||||
- Are deprecated algorithms used? (MD5, SHA1 for security, RC4, DES)
|
||||
|
||||
### 5e. Resource exhaustion
|
||||
|
||||
- Are there unbounded allocations? (`Vec` growing from user input without limits)
|
||||
- Are there unbounded loops? (retry loops without max attempts)
|
||||
- Are file reads bounded? (`std::fs::read_to_string` on user-provided paths)
|
||||
- Are timeouts set on all network operations?
|
||||
- Are there connection/resource leaks? (opened but never closed, missing `Drop`)
|
||||
|
||||
### 5f. Error handling
|
||||
|
||||
- Are errors swallowed silently? (`let _ = ...`, `.ok()` discarding errors that matter)
|
||||
- Do error types carry enough context to debug in production?
|
||||
- Are there error type mismatches? (returning generic `anyhow::Error` where a typed error would prevent confusion)
|
||||
- Is `thiserror` used consistently for error types (IronClaw convention)?
|
||||
|
||||
## Step 6: Check for inconsistencies
|
||||
|
||||
### 6a. Naming conventions
|
||||
|
||||
- Are types, functions, modules named consistently? (e.g., mixing `get_` and `fetch_`, `create_` and `new_`)
|
||||
- Do similar operations follow the same patterns?
|
||||
|
||||
### 6b. Duplicate or near-duplicate code
|
||||
|
||||
Look for:
|
||||
- Functions that do nearly the same thing with minor variations (candidates for generics or shared helpers)
|
||||
- Repeated error mapping patterns that should be extracted
|
||||
- Copy-pasted SQL queries or string templates with slight differences
|
||||
- Identical struct definitions or conversion logic in different modules
|
||||
|
||||
### 6c. API consistency
|
||||
|
||||
- Do similar functions take arguments in the same order?
|
||||
- Are return types consistent? (e.g., some functions return `Option<T>`, similar ones return `Result<T, E>`)
|
||||
- Are visibility modifiers consistent? (`pub` where it should be `pub(crate)`, or vice versa)
|
||||
|
||||
### 6d. Dead code and unused items
|
||||
|
||||
- Are there functions, structs, or modules that nothing references?
|
||||
- Are there `#[allow(dead_code)]` annotations that should be investigated?
|
||||
- Are there feature-gated items where the feature is never enabled?
|
||||
|
||||
### 6e. Import style
|
||||
|
||||
IronClaw convention: use `crate::` imports, not `super::`. Flag any `super::` imports in non-test code.
|
||||
|
||||
## Step 7: Inspect for change oversights
|
||||
|
||||
### 7a. Partial refactors
|
||||
|
||||
- Are there old patterns coexisting with new patterns?
|
||||
- Are there renamed types/functions where some call sites still use the old name via a compatibility alias?
|
||||
- Are there comments referencing behavior that no longer exists?
|
||||
|
||||
### 7b. Trait implementation gaps
|
||||
|
||||
- If a trait is defined, do all intended types implement it?
|
||||
- Are there `impl` blocks that look incomplete?
|
||||
- Are `Default` implementations sensible?
|
||||
|
||||
IronClaw key traits: `Database` (~60 methods), `Channel`, `Tool`, `LlmProvider`, `SuccessEvaluator`, `EmbeddingProvider`. If any new methods were added to `Database`, verify both `postgres.rs` and `libsql_backend.rs` implement them.
|
||||
|
||||
### 7c. Test coverage gaps
|
||||
|
||||
- Are there public functions without any test?
|
||||
- Are there error paths without tests?
|
||||
- Are there recently-changed functions where the tests still assert old behavior?
|
||||
|
||||
### 7d. Documentation drift
|
||||
|
||||
- Do doc comments match actual function behavior?
|
||||
- Are examples in doc comments still valid and compilable?
|
||||
|
||||
## Step 8: Dependency audit
|
||||
|
||||
Review `Cargo.toml` and `Cargo.lock`:
|
||||
- Are there duplicate versions of the same crate in the lock file? (potential version conflicts)
|
||||
- Are there dependencies with known security advisories? Run `cargo audit` to check (install with `cargo install cargo-audit` if not present).
|
||||
- Are there heavy dependencies used for trivial functionality?
|
||||
- Are dependency features minimal?
|
||||
|
||||
## Step 9: Present findings
|
||||
|
||||
Compile all findings into a structured report. Group by severity, then by category.
|
||||
|
||||
### Format
|
||||
|
||||
For each finding:
|
||||
|
||||
```
|
||||
### [Severity] Category: One-line summary
|
||||
|
||||
**Location:** `file_path:line_number`
|
||||
**Category:** Vulnerability | Bug | Unfinished | Inconsistency | Duplicate | Oversight | Style
|
||||
|
||||
**Description:**
|
||||
Detailed explanation of the issue, why it matters, and how it could manifest.
|
||||
|
||||
**Suggested fix:**
|
||||
Concrete suggestion with code if applicable.
|
||||
```
|
||||
|
||||
### Severity levels
|
||||
|
||||
- **Critical**: Security vulnerability, data loss, or crash in production
|
||||
- **High**: Bug that causes incorrect behavior, `todo!()`/`unimplemented!()` in prod code, or missing validation on trust boundaries
|
||||
- **Medium**: Inconsistency, duplicate code, incomplete error handling, missing tests for important paths
|
||||
- **Low**: Naming inconsistency, unnecessary complexity, documentation drift, minor dead code
|
||||
- **Nit**: Style preference, optional improvement
|
||||
|
||||
### Summary table
|
||||
|
||||
End with a summary table:
|
||||
|
||||
| # | Severity | Category | File:Line | Finding |
|
||||
|---|----------|----------|-----------|---------|
|
||||
|
||||
And a final tally: X Critical, Y High, Z Medium, W Low, V Nit.
|
||||
|
||||
## Rules
|
||||
|
||||
- Read every file before reporting on it. Never guess about code you haven't seen.
|
||||
- Be specific. "This might have issues" is worthless. "Line 42 calls `.unwrap()` on a `Result` that returns `Err` when the DB connection is dropped" is useful.
|
||||
- Distinguish certainty levels: "this IS a bug" vs "this COULD be a bug if X".
|
||||
- Don't invent problems to look thorough. If the code is solid, say so.
|
||||
- Focus on substance over style. Don't flag formatting unless it causes real confusion.
|
||||
- Respect existing project conventions (check CLAUDE.md). Don't flag patterns the project explicitly endorses.
|
||||
- When in doubt about severity, round up.
|
||||
- For large crates (>50 files), prioritize: core logic > public API > internal utilities > tests > examples.
|
||||
- Use the Task tool to parallelize file reading across modules when the crate is large.
|
||||
- Do NOT fix anything. This is a read-only audit. Report findings for the user to action.
|
||||
@@ -0,0 +1,170 @@
|
||||
---
|
||||
description: Paranoid architect review of a PR — fetches diff, reads changed files, deep review across 6 lenses, posts findings as GitHub comments
|
||||
disable-model-invocation: true
|
||||
allowed-tools: Bash(gh pr view:*), Bash(gh pr diff:*), Bash(gh pr comment:*), Bash(gh api:*), Bash(gh repo view:*), Bash(git diff:*), Bash(git log:*), Read, Grep, Glob
|
||||
argument-hint: "<pr-number or github-pr-url>"
|
||||
---
|
||||
|
||||
# Paranoid Architect Code Review
|
||||
|
||||
You are reviewing this PR as a paranoid architect. Your job is to find every bug, vulnerability, race condition, edge case, and undocumented assumption before it ships. Assume adversarial users, concurrent access, and Murphy's law.
|
||||
|
||||
## Step 1: Resolve the PR
|
||||
|
||||
Parse `$ARGUMENTS` to extract the PR number:
|
||||
- If it's a URL like `https://github.com/owner/repo/pull/123`, extract `123`.
|
||||
- If it's a bare number, use it directly.
|
||||
- If empty, stop and ask the user for a PR number.
|
||||
|
||||
Fetch PR metadata (including head commit SHA for posting line comments later):
|
||||
|
||||
```
|
||||
gh pr view {number} --json title,body,baseRefName,headRefName,headRefOid,files,additions,deletions
|
||||
```
|
||||
|
||||
Save the `headRefOid` value, you'll need it as `commit_id` in Step 6.
|
||||
|
||||
## Step 2: Load the full diff
|
||||
|
||||
```
|
||||
gh pr diff {number}
|
||||
```
|
||||
|
||||
Also get the list of changed files:
|
||||
|
||||
```
|
||||
gh pr diff {number} --name-only
|
||||
```
|
||||
|
||||
## Step 3: Read every changed file in full
|
||||
|
||||
For each changed file, read the ENTIRE current file (not just the diff hunks). You need surrounding context to catch:
|
||||
- Callers of modified functions that now behave differently
|
||||
- Trait/interface contracts that the change may violate
|
||||
- Invariants established elsewhere that the diff breaks
|
||||
|
||||
If the PR touches more than 20 files, still read all of them, but process in this priority order: service logic > routes/handlers > models/types > tests > docs. Batch reads in groups of ~20 if needed.
|
||||
|
||||
## Step 4: Deep review
|
||||
|
||||
Go through the changes with each of these lenses. For every finding, note the file, line range, severity, and a concrete description.
|
||||
|
||||
### IronClaw-specific checks
|
||||
|
||||
In addition to the general lenses below, check IronClaw conventions (see CLAUDE.md):
|
||||
- No `.unwrap()` or `.expect()` in production code (tests are fine)
|
||||
- Use `crate::` imports, not `super::`
|
||||
- Error types use `thiserror` in `error.rs`
|
||||
- If the change touches persistence, verify both database backends are updated (PostgreSQL in `postgres.rs` AND libSQL in `libsql_backend.rs`)
|
||||
- New tools must implement the `Tool` trait correctly and be registered in `registry.rs`
|
||||
- External tool output must pass through the safety layer
|
||||
|
||||
### 4a. Correctness and bugs
|
||||
|
||||
- Off-by-one errors, wrong comparison operators, inverted conditions
|
||||
- Unreachable code, dead branches, impossible match arms
|
||||
- Type confusion (mixing up IDs, using wrong enum variant)
|
||||
- Incorrect error propagation (swallowed errors, wrong error type/status code)
|
||||
- Broken invariants (e.g. uniqueness assumptions violated, ordering assumptions wrong)
|
||||
- Concurrency issues (TOCTOU, missing locks, race conditions between check and use)
|
||||
|
||||
### 4b. Edge cases and failure handling
|
||||
|
||||
- What happens with empty input, None/null, zero-length collections?
|
||||
- What happens when external services fail (DB down, HTTP timeout, malformed response)?
|
||||
- What happens at integer boundaries (overflow, underflow, i64::MAX)?
|
||||
- What happens with malformed or adversarial input (invalid UTF-8, huge payloads, deeply nested JSON)?
|
||||
- Are all error paths tested? Does every `?` propagation make sense?
|
||||
- Are partial failures handled (e.g. wrote to DB but failed to emit event)?
|
||||
|
||||
### 4c. Security (assume a malicious actor)
|
||||
|
||||
- **Authentication/Authorization bypass**: Can an unauthenticated user reach this? Can workspace A's user access workspace B's data? Are there IDOR vulnerabilities?
|
||||
- **Injection**: SQL injection via string interpolation? Command injection? Log injection? Header injection?
|
||||
- **Data leakage**: Are secrets, PII, or conversation content logged? Returned in error messages? Exposed in API responses?
|
||||
- **Resource exhaustion / DoS**: Can an attacker send unbounded input? Trigger expensive operations without rate limits? Cause OOM via large allocations?
|
||||
- **Financial abuse**: Can tokens/credits be consumed without being tracked? Can usage limits be bypassed?
|
||||
- **Replay / race conditions**: Can the same request be replayed for double-spend? Can concurrent requests bypass limits?
|
||||
- **Cryptographic issues**: Timing attacks on comparisons? Weak randomness? Missing HMAC verification?
|
||||
|
||||
### 4d. Test coverage
|
||||
|
||||
- Is every new public function/method tested?
|
||||
- Are error paths tested (not just happy paths)?
|
||||
- Are edge cases covered (empty input, boundary values, concurrent access)?
|
||||
- Do existing tests still make sense with the new changes, or do they assert stale behavior?
|
||||
- Are there integration/e2e tests for the full flow?
|
||||
- If a test is missing, describe exactly what test should be written.
|
||||
|
||||
### 4e. Documentation and assumptions
|
||||
|
||||
- Are new assumptions documented in comments? (e.g. "this field is always non-empty because X")
|
||||
- Are non-obvious algorithms or business rules explained?
|
||||
- Are API contracts (request/response shapes, error codes, status codes) documented?
|
||||
- Are there TODO/FIXME/HACK comments that should be tracked as issues?
|
||||
|
||||
### 4f. Architectural concerns
|
||||
|
||||
- Does this change follow existing patterns in the codebase, or does it introduce a new one without justification?
|
||||
- Are there unnecessary abstractions or premature generalizations?
|
||||
- Is there duplicated logic that should be extracted?
|
||||
- Are dependencies between modules clean, or does this create circular/tight coupling?
|
||||
- Will this change make future work harder?
|
||||
|
||||
## Step 5: Present findings
|
||||
|
||||
Summarize findings to the user as a table:
|
||||
|
||||
| # | Severity | Category | File:Line | Finding | Suggested Fix |
|
||||
|---|----------|----------|-----------|---------|---------------|
|
||||
|
||||
Severity levels:
|
||||
- **Critical**: Security vulnerability, data loss, or financial exploit
|
||||
- **High**: Bug that will cause incorrect behavior in production
|
||||
- **Medium**: Robustness issue, missing validation, or incomplete error handling
|
||||
- **Low**: Style, naming, documentation, or minor improvement
|
||||
- **Nit**: Optional suggestion, take-it-or-leave-it
|
||||
|
||||
Ask the user which findings to post as PR comments. Default: all Critical, High, and Medium.
|
||||
|
||||
## Step 6: Post comments on GitHub
|
||||
|
||||
Resolve the repo owner and name if not already known:
|
||||
|
||||
```
|
||||
gh repo view --json owner,name --jq '"\(.owner.login)/\(.name)"'
|
||||
```
|
||||
|
||||
For each approved finding, post a review comment on the PR at the specific file and line. Use the `headRefOid` from Step 1 as the `commit_id`:
|
||||
|
||||
```
|
||||
gh api repos/{owner}/{repo}/pulls/{number}/comments \
|
||||
-f body="..." \
|
||||
-f path="..." \
|
||||
-f commit_id="{headRefOid}" \
|
||||
-F line=... \
|
||||
-f side="RIGHT"
|
||||
```
|
||||
|
||||
For findings that span multiple locations or are architectural, post as a regular PR comment:
|
||||
|
||||
```
|
||||
gh pr comment {number} --body "..."
|
||||
```
|
||||
|
||||
Format each comment clearly:
|
||||
- Severity tag (e.g. `**High Severity**`)
|
||||
- One-line summary
|
||||
- Detailed explanation of the issue
|
||||
- Concrete suggestion for the fix (with code if possible)
|
||||
|
||||
## Rules
|
||||
|
||||
- Read every changed file in full before writing a single finding. Context matters.
|
||||
- Never post a comment about code you haven't actually read. Verify line numbers against the actual file.
|
||||
- Be specific. "This might have issues" is useless. "Line 42 returns 404 but should return 400 because X" is useful.
|
||||
- Distinguish between "this IS a bug" and "this COULD be a bug if X". Be honest about certainty.
|
||||
- Don't nitpick formatting or style unless it causes actual confusion. Focus on substance.
|
||||
- If the code is good and you find nothing, say so. Don't invent problems to look thorough.
|
||||
- Respect the project's CLAUDE.md privacy rules: never include customer data, secrets, or PII in comments.
|
||||
- When in doubt about severity, round up. It's cheaper to dismiss a false alarm than to miss a real bug.
|
||||
@@ -8,12 +8,13 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
profile: minimal
|
||||
components: rustfmt, clippy
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Check formatting
|
||||
run: |
|
||||
cargo fmt --all -- --check
|
||||
|
||||
@@ -24,6 +24,7 @@ jobs:
|
||||
- &install-rust
|
||||
name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
# Generating a GitHub token, so that PRs and tags created by
|
||||
# the release-plz-action can trigger actions workflows.
|
||||
- name: Generate GitHub token
|
||||
@@ -56,6 +57,7 @@ jobs:
|
||||
steps:
|
||||
- *checkout
|
||||
- *install-rust
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Run release-plz
|
||||
uses: release-plz/[email protected]
|
||||
with:
|
||||
|
||||
@@ -39,7 +39,6 @@ permissions:
|
||||
# If there's a prerelease-style suffix to the version, then the release(s)
|
||||
# will be marked as a prerelease.
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
tags:
|
||||
- '**[0-9]+.[0-9]+.[0-9]+*'
|
||||
@@ -282,43 +281,14 @@ jobs:
|
||||
|
||||
gh release create "${{ needs.plan.outputs.tag }}" --target "$RELEASE_COMMIT" $PRERELEASE_FLAG --title "$ANNOUNCEMENT_TITLE" --notes-file "$RUNNER_TEMP/notes.txt" artifacts/*
|
||||
|
||||
publish-npm:
|
||||
needs:
|
||||
- plan
|
||||
- host
|
||||
runs-on: "ubuntu-22.04"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PLAN: ${{ needs.plan.outputs.val }}
|
||||
if: ${{ !fromJson(needs.plan.outputs.val).announcement_is_prerelease || fromJson(needs.plan.outputs.val).publish_prereleases }}
|
||||
steps:
|
||||
- name: Fetch npm packages
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: npm/
|
||||
merge-multiple: true
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20.x'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
- run: |
|
||||
for release in $(echo "$PLAN" | jq --compact-output '.releases[] | select([.artifacts[] | endswith("-npm-package.tar.gz")] | any)'); do
|
||||
pkg=$(echo "$release" | jq '.artifacts[] | select(endswith("-npm-package.tar.gz"))' --raw-output)
|
||||
npm publish --access public "./npm/${pkg}"
|
||||
done
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
announce:
|
||||
needs:
|
||||
- plan
|
||||
- host
|
||||
- publish-npm
|
||||
# use "always() && ..." to allow us to wait for all publish jobs while
|
||||
# still allowing individual publish jobs to skip themselves (for prereleases).
|
||||
# "host" however must run to completion, no skipping allowed!
|
||||
if: ${{ always() && needs.host.result == 'success' && (needs.publish-npm.result == 'skipped' || needs.publish-npm.result == 'success') }}
|
||||
if: ${{ always() && needs.host.result == 'success' }}
|
||||
runs-on: "ubuntu-22.04"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
@@ -11,10 +11,11 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
profile: minimal
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Run Tests
|
||||
run: cargo test --all-features -- --nocapture
|
||||
|
||||
@@ -1,6 +1,15 @@
|
||||
|
||||
.env
|
||||
.env.local
|
||||
.env.*
|
||||
!.env.example
|
||||
|
||||
# Claude Code worktrees
|
||||
.claude/worktrees/
|
||||
|
||||
# Sidecar tool data
|
||||
.sidecar/
|
||||
.todos/
|
||||
|
||||
target/
|
||||
|
||||
|
||||
@@ -7,6 +7,67 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.5.0](https://github.com/nearai/ironclaw/compare/v0.4.0...v0.5.0) - 2026-02-17
|
||||
|
||||
### Added
|
||||
|
||||
- add cooldown management to FailoverProvider ([#114](https://github.com/nearai/ironclaw/pull/114))
|
||||
|
||||
## [0.4.0](https://github.com/nearai/ironclaw/compare/v0.3.0...v0.4.0) - 2026-02-17
|
||||
|
||||
### Added
|
||||
|
||||
- move per-invocation approval check into Tool trait ([#119](https://github.com/nearai/ironclaw/pull/119))
|
||||
- add polished boot screen on CLI startup ([#118](https://github.com/nearai/ironclaw/pull/118))
|
||||
- Add lifecycle hooks system with 6 interception points ([#18](https://github.com/nearai/ironclaw/pull/18))
|
||||
|
||||
### Other
|
||||
|
||||
- remove accidentally committed .sidecar and .todos directories ([#123](https://github.com/nearai/ironclaw/pull/123))
|
||||
|
||||
## [0.3.0](https://github.com/nearai/ironclaw/compare/v0.2.0...v0.3.0) - 2026-02-17
|
||||
|
||||
### Added
|
||||
|
||||
- direct api key and cheap model ([#116](https://github.com/nearai/ironclaw/pull/116))
|
||||
|
||||
## [0.2.0](https://github.com/nearai/ironclaw/compare/v0.1.3...v0.2.0) - 2026-02-16
|
||||
|
||||
### Added
|
||||
|
||||
- mark Ollama + OpenAI-compatible as implemented ([#102](https://github.com/nearai/ironclaw/pull/102))
|
||||
- multi-provider inference + libSQL onboarding selection ([#92](https://github.com/nearai/ironclaw/pull/92))
|
||||
- add multi-provider LLM failover with retry backoff ([#28](https://github.com/nearai/ironclaw/pull/28))
|
||||
- add libSQL/Turso embedded database backend ([#47](https://github.com/nearai/ironclaw/pull/47))
|
||||
- Move debug log truncation from agent loop to REPL channel ([#65](https://github.com/nearai/ironclaw/pull/65))
|
||||
|
||||
### Fixed
|
||||
|
||||
- shell destructive-command check bypassed by Value::Object arguments ([#72](https://github.com/nearai/ironclaw/pull/72))
|
||||
- propagate real tool_call_id instead of hardcoded placeholder ([#73](https://github.com/nearai/ironclaw/pull/73))
|
||||
- Fix wasm tool schemas and runtime ([#42](https://github.com/nearai/ironclaw/pull/42))
|
||||
- flatten tool messages for NEAR AI cloud-api compatibility ([#41](https://github.com/nearai/ironclaw/pull/41))
|
||||
- security hardening across all layers ([#35](https://github.com/nearai/ironclaw/pull/35))
|
||||
|
||||
### Other
|
||||
|
||||
- Explicitly enable cargo-dist caching for binary artifacts building
|
||||
- Skip building binary artifacts on every PR
|
||||
- add module specification rules to CLAUDE.md
|
||||
- add setup/onboarding specification (src/setup/README.md)
|
||||
- deduplicate tool code and remove dead stubs ([#98](https://github.com/nearai/ironclaw/pull/98))
|
||||
- Reformat architecture diagram in README ([#64](https://github.com/nearai/ironclaw/pull/64))
|
||||
- Add review discipline guidelines to CLAUDE.md ([#68](https://github.com/nearai/ironclaw/pull/68))
|
||||
- Bump MSRV to 1.92, add GCP deployment files ([#40](https://github.com/nearai/ironclaw/pull/40))
|
||||
- Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) ([#31](https://github.com/nearai/ironclaw/pull/31))
|
||||
|
||||
## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12
|
||||
|
||||
### Other
|
||||
|
||||
- Enabled builds caching during CI/CD
|
||||
- Disabled npm publishing as the name is already taken
|
||||
|
||||
## [0.1.2](https://github.com/nearai/ironclaw/compare/v0.1.1...v0.1.2) - 2026-02-12
|
||||
|
||||
### Other
|
||||
|
||||
@@ -151,6 +151,12 @@ src/
|
||||
│ ├── rate_limiter.rs # Per-tool rate limiting
|
||||
│ └── storage.rs # Linear memory persistence
|
||||
│
|
||||
├── db/ # Database abstraction layer
|
||||
│ ├── mod.rs # Database trait (~60 async methods)
|
||||
│ ├── postgres.rs # PostgreSQL backend (delegates to Store + Repository)
|
||||
│ ├── libsql_backend.rs # libSQL/Turso backend (embedded SQLite)
|
||||
│ └── libsql_migrations.rs # SQLite-dialect schema (idempotent)
|
||||
│
|
||||
├── workspace/ # Persistent memory system (OpenClaw-inspired)
|
||||
│ ├── mod.rs # Workspace struct, memory operations
|
||||
│ ├── document.rs # MemoryDocument, MemoryChunk, WorkspaceEntry
|
||||
@@ -192,8 +198,9 @@ When designing new features or systems, always prefer generic/extensible archite
|
||||
|
||||
### Error Handling
|
||||
- Use `thiserror` for error types in `error.rs`
|
||||
- Never use `.unwrap()` in production code (tests are fine)
|
||||
- Never use `.unwrap()` or `.expect()` in production code (tests are fine)
|
||||
- Map errors with context: `.map_err(|e| SomeError::Variant { reason: e.to_string() })?`
|
||||
- Before committing, grep for `.unwrap()` and `.expect(` in changed files to catch violations mechanically
|
||||
|
||||
### Async
|
||||
- All I/O is async with tokio
|
||||
@@ -201,6 +208,7 @@ When designing new features or systems, always prefer generic/extensible archite
|
||||
- Use `RwLock` for concurrent read/write access
|
||||
|
||||
### Traits for Extensibility
|
||||
- `Database` - Add new database backends (must implement all ~60 methods)
|
||||
- `Channel` - Add new input sources
|
||||
- `Tool` - Add new capabilities
|
||||
- `LlmProvider` - Add new LLM backends
|
||||
@@ -248,7 +256,12 @@ Pending -> InProgress -> Completed -> Submitted -> Accepted
|
||||
|
||||
Environment variables (see `.env.example`):
|
||||
```bash
|
||||
# Database backend (default: postgres)
|
||||
DATABASE_BACKEND=postgres # or "libsql" / "turso"
|
||||
DATABASE_URL=postgres://user:pass@localhost/ironclaw
|
||||
LIBSQL_PATH=~/.ironclaw/ironclaw.db # libSQL local path (default)
|
||||
# LIBSQL_URL=libsql://xxx.turso.io # Turso cloud (optional)
|
||||
# LIBSQL_AUTH_TOKEN=xxx # Required with LIBSQL_URL
|
||||
|
||||
# NEAR AI (required)
|
||||
NEARAI_SESSION_TOKEN=sess_...
|
||||
@@ -308,7 +321,51 @@ Session tokens have the format `sess_xxx` (37 characters). They are authenticate
|
||||
|
||||
## Database
|
||||
|
||||
Single migration in `migrations/V1__initial.sql`. Tables:
|
||||
IronClaw supports two database backends, selected at compile time via Cargo feature flags and at runtime via the `DATABASE_BACKEND` environment variable.
|
||||
|
||||
**IMPORTANT: All new features that touch persistence MUST support both backends.** Implement the operation as a method on the `Database` trait in `src/db/mod.rs`, then add the implementation in both `src/db/postgres.rs` (delegate to Store/Repository) and `src/db/libsql_backend.rs` (native SQL).
|
||||
|
||||
### Backends
|
||||
|
||||
| Backend | Feature Flag | Default | Use Case |
|
||||
|---------|-------------|---------|----------|
|
||||
| PostgreSQL | `postgres` (default) | Yes | Production, existing deployments |
|
||||
| libSQL/Turso | `libsql` | No | Zero-dependency local mode, edge, Turso cloud |
|
||||
|
||||
```bash
|
||||
# Build with PostgreSQL only (default)
|
||||
cargo build
|
||||
|
||||
# Build with libSQL only
|
||||
cargo build --no-default-features --features libsql
|
||||
|
||||
# Build with both backends available
|
||||
cargo build --features "postgres,libsql"
|
||||
```
|
||||
|
||||
### Database Trait
|
||||
|
||||
The `Database` trait (`src/db/mod.rs`) defines ~60 async methods covering all persistence:
|
||||
- Conversations, messages, metadata
|
||||
- Jobs, actions, LLM calls, estimation snapshots
|
||||
- Sandbox jobs, job events
|
||||
- Routines, routine runs
|
||||
- Tool failures, settings
|
||||
- Workspace: documents, chunks, hybrid search
|
||||
|
||||
Both backends implement this trait. PostgreSQL delegates to the existing `Store` + `Repository`. libSQL implements native SQLite-dialect SQL.
|
||||
|
||||
### Schema
|
||||
|
||||
**PostgreSQL:** `migrations/V1__initial.sql` (351 lines). Uses pgvector for embeddings, tsvector for FTS, PL/pgSQL functions. Managed by `refinery`.
|
||||
|
||||
**libSQL:** `src/db/libsql_migrations.rs` (consolidated schema, ~480 lines). Translates PG types:
|
||||
- `UUID` -> `TEXT`, `TIMESTAMPTZ` -> `TEXT` (ISO-8601), `JSONB` -> `TEXT`
|
||||
- `VECTOR(1536)` -> `F32_BLOB(1536)` with `libsql_vector_idx`
|
||||
- `tsvector`/`ts_rank_cd` -> FTS5 virtual table with sync triggers
|
||||
- PL/pgSQL functions -> SQLite triggers
|
||||
|
||||
**Tables (both backends):**
|
||||
|
||||
**Core:**
|
||||
- `conversations` - Multi-channel conversation tracking
|
||||
@@ -320,12 +377,41 @@ Single migration in `migrations/V1__initial.sql`. Tables:
|
||||
|
||||
**Workspace/Memory:**
|
||||
- `memory_documents` - Flexible path-based files (e.g., "context/vision.md", "daily/2024-01-15.md")
|
||||
- `memory_chunks` - Chunked content with FTS (tsvector) and vector (pgvector) indexes
|
||||
- `memory_chunks` - Chunked content with FTS and vector indexes
|
||||
- `heartbeat_state` - Periodic execution tracking
|
||||
|
||||
Requires pgvector extension: `CREATE EXTENSION IF NOT EXISTS vector;`
|
||||
**Other:**
|
||||
- `routines`, `routine_runs` - Scheduled/reactive execution
|
||||
- `settings` - Per-user key-value settings
|
||||
- `tool_failures` - Self-repair tracking
|
||||
- `secrets`, `wasm_tools`, `tool_capabilities` - Extension infrastructure
|
||||
|
||||
Run migrations: `refinery migrate -c refinery.toml`
|
||||
### Configuration
|
||||
|
||||
```bash
|
||||
# Backend selection (default: postgres)
|
||||
DATABASE_BACKEND=libsql
|
||||
|
||||
# PostgreSQL
|
||||
DATABASE_URL=postgres://user:pass@localhost/ironclaw
|
||||
|
||||
# libSQL (embedded)
|
||||
LIBSQL_PATH=~/.ironclaw/ironclaw.db # Default path
|
||||
|
||||
# libSQL (Turso cloud sync)
|
||||
LIBSQL_URL=libsql://your-db.turso.io
|
||||
LIBSQL_AUTH_TOKEN=your-token # Required when LIBSQL_URL is set
|
||||
```
|
||||
|
||||
### Current Limitations (libSQL backend)
|
||||
|
||||
- **Workspace/memory system** not yet wired through Database trait (requires Store migration)
|
||||
- **Secrets store** not yet available (still requires PostgresSecretsStore)
|
||||
- **Hybrid search** uses FTS5 only (vector search via libsql_vector_idx not yet implemented)
|
||||
- **Settings reload from DB** skipped (Config::from_db requires Store)
|
||||
- No incremental migration versioning (schema is CREATE IF NOT EXISTS, no ALTER TABLE support yet)
|
||||
- **No encryption at rest** -- The local SQLite database file stores conversation content, job data, workspace memory, and other application data in plaintext. Only secrets (API tokens, credentials) are encrypted via AES-256-GCM before storage. Users handling sensitive data should use full-disk encryption (FileVault, LUKS, BitLocker) or consider the PostgreSQL backend with TDE/encrypted storage.
|
||||
- **JSON merge patch vs path-targeted update** -- The libSQL backend uses RFC 7396 JSON Merge Patch (`json_patch`) for metadata updates, while PostgreSQL uses path-targeted `jsonb_set`. Merge patch replaces top-level keys entirely, which may drop nested keys not present in the patch. Callers should avoid relying on partial nested object updates in metadata fields.
|
||||
|
||||
## Safety Layer
|
||||
|
||||
@@ -387,6 +473,7 @@ Key test patterns:
|
||||
- ✅ **Claude Code mode** - Delegate jobs to Claude CLI inside containers
|
||||
- ✅ **Routines system** - Cron, event, webhook, and manual triggers with guardrails
|
||||
- ✅ **Extension management** - Install, auth, activate MCP/WASM extensions via CLI and web UI
|
||||
- ✅ **libSQL/Turso backend** - Database trait abstraction (`src/db/`), feature-gated dual backend support (postgres/libsql), embedded SQLite for zero-dependency local mode
|
||||
|
||||
## Adding a New Tool
|
||||
|
||||
@@ -543,6 +630,22 @@ RUST_LOG=ironclaw::agent=debug cargo run
|
||||
RUST_LOG=ironclaw=debug,tower_http=debug cargo run
|
||||
```
|
||||
|
||||
## Module Specifications
|
||||
|
||||
Some modules have a `README.md` that serves as the authoritative specification
|
||||
for that module's behavior. When modifying code in a module that has a spec:
|
||||
|
||||
1. **Read the spec first** before making changes
|
||||
2. **Code follows spec**: if the spec says X, the code must do X
|
||||
3. **Update both sides**: if you change behavior, update the spec to match;
|
||||
if you're implementing a spec change, update the code to match
|
||||
4. **Spec is the tiebreaker**: when code and spec disagree, the spec is correct
|
||||
(unless the spec is clearly outdated, in which case fix the spec first)
|
||||
|
||||
| Module | Spec File |
|
||||
|--------|-----------|
|
||||
| `src/setup/` | `src/setup/README.md` |
|
||||
|
||||
## Code Style
|
||||
|
||||
- Use `crate::` imports, not `super::`
|
||||
@@ -551,6 +654,37 @@ RUST_LOG=ironclaw=debug,tower_http=debug cargo run
|
||||
- Keep functions focused, extract helpers when logic is reused
|
||||
- Comments for non-obvious logic only
|
||||
|
||||
## Review & Fix Discipline
|
||||
|
||||
Hard-won lessons from code review -- follow these when fixing bugs or addressing review feedback.
|
||||
|
||||
### Fix the pattern, not just the instance
|
||||
When a reviewer flags a bug (e.g., TOCTOU race in INSERT + SELECT-back), search the entire codebase for all instances of that same pattern. A fix in `SecretsStore::create()` that doesn't also fix `WasmToolStore::store()` is half a fix.
|
||||
|
||||
### Propagate architectural fixes to satellite types
|
||||
If a core type changes its concurrency model (e.g., `LibSqlBackend` switches to connection-per-operation), every type that was handed a resource from the old model (e.g., `LibSqlSecretsStore`, `LibSqlWasmToolStore` holding a single `Connection`) must also be updated. Grep for the old type across the codebase.
|
||||
|
||||
### Schema translation is more than DDL
|
||||
When translating a database schema between backends (PostgreSQL to libSQL, etc.), check for:
|
||||
- **Indexes** -- diff `CREATE INDEX` statements between the two schemas
|
||||
- **Seed data** -- check for `INSERT INTO` in migrations (e.g., `leak_detection_patterns`)
|
||||
- **Semantic differences** -- document where SQL functions behave differently (e.g., `json_patch` vs `jsonb_set`)
|
||||
|
||||
### Feature flag testing
|
||||
When adding feature-gated code, test compilation with each feature in isolation:
|
||||
```bash
|
||||
cargo check # default features
|
||||
cargo check --no-default-features --features libsql # libsql only
|
||||
cargo check --all-features # all features
|
||||
```
|
||||
Dead code behind the wrong `#[cfg]` gate will only show up when building with a single feature.
|
||||
|
||||
### Mechanical verification before committing
|
||||
Run these checks on changed files before committing:
|
||||
- `grep -rnE '\.unwrap\(|\.expect\(' <files>` -- no panics in production
|
||||
- `grep -rn 'super::' <files>` -- use `crate::` imports
|
||||
- If you fixed a pattern bug, `grep` for other instances of that pattern across `src/`
|
||||
|
||||
## Workspace & Memory System
|
||||
|
||||
Inspired by [OpenClaw](https://github.com/openclaw/openclaw), the workspace provides persistent memory for agents with a flexible filesystem-like structure.
|
||||
@@ -625,7 +759,7 @@ Four tools for LLM use:
|
||||
|
||||
### Hybrid Search (RRF)
|
||||
|
||||
Combines full-text search (PostgreSQL `ts_rank_cd`) and vector similarity (pgvector cosine) using Reciprocal Rank Fusion:
|
||||
Combines full-text search and vector similarity using Reciprocal Rank Fusion:
|
||||
|
||||
```
|
||||
score(d) = Σ 1/(k + rank(d)) for each method where d appears
|
||||
@@ -633,6 +767,10 @@ score(d) = Σ 1/(k + rank(d)) for each method where d appears
|
||||
|
||||
Default k=60. Results from both methods are combined, with documents appearing in both getting boosted scores.
|
||||
|
||||
**Backend differences:**
|
||||
- **PostgreSQL:** `ts_rank_cd` for FTS, pgvector cosine distance for vectors, full RRF
|
||||
- **libSQL:** FTS5 for keyword search only (vector search via `libsql_vector_idx` not yet wired)
|
||||
|
||||
### Heartbeat System
|
||||
|
||||
Proactive periodic execution (default: 30 minutes):
|
||||
|
||||
Generated
+635
-66
File diff suppressed because it is too large
Load Diff
+42
-14
@@ -1,8 +1,16 @@
|
||||
[workspace]
|
||||
exclude = [
|
||||
"channels-src/telegram",
|
||||
"channels-src/slack",
|
||||
"channels-src/whatsapp",
|
||||
"tools-src/gmail",
|
||||
]
|
||||
|
||||
[package]
|
||||
name = "ironclaw"
|
||||
version = "0.1.2"
|
||||
version = "0.5.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.85"
|
||||
rust-version = "1.92"
|
||||
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
||||
authors = ["NEAR AI <[email protected]>"]
|
||||
license = "MIT OR Apache-2.0"
|
||||
@@ -22,17 +30,20 @@ tokio-stream = { version = "0.1", features = ["sync"] }
|
||||
futures = "0.3"
|
||||
|
||||
# HTTP client
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots", "stream"] }
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
|
||||
# Database
|
||||
deadpool-postgres = "0.14"
|
||||
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"] }
|
||||
postgres-types = { version = "0.2", features = ["with-serde_json-1"] }
|
||||
refinery = { version = "0.8", features = ["tokio-postgres"] }
|
||||
# Database - PostgreSQL (default, feature-gated)
|
||||
deadpool-postgres = { version = "0.14", optional = true }
|
||||
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"], optional = true }
|
||||
postgres-types = { version = "0.2", features = ["with-serde_json-1"], optional = true }
|
||||
refinery = { version = "0.8", features = ["tokio-postgres"], optional = true }
|
||||
|
||||
# Database - libSQL/Turso (optional embedded database)
|
||||
libsql = { version = "0.6", optional = true, default-features = false, features = ["core", "replication"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
@@ -48,7 +59,7 @@ dotenvy = "0.15"
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "db-tokio-postgres", "maths"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "maths"] }
|
||||
rust_decimal_macros = "1"
|
||||
|
||||
# Async traits
|
||||
@@ -81,7 +92,8 @@ fs4 = "0.6"
|
||||
# Secrecy for sensitive values
|
||||
secrecy = { version = "0.10", features = ["serde"] }
|
||||
|
||||
# URL encoding for OAuth flow
|
||||
# URL parsing and encoding
|
||||
url = "2"
|
||||
urlencoding = "2"
|
||||
|
||||
# Open URLs in browser
|
||||
@@ -89,7 +101,7 @@ open = "5"
|
||||
|
||||
# Vector embeddings for semantic search
|
||||
# The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres)
|
||||
pgvector = { version = "0.4", features = ["postgres"] }
|
||||
pgvector = { version = "0.4", features = ["postgres"], optional = true }
|
||||
|
||||
# WASM sandbox for untrusted tool execution
|
||||
wasmtime = { version = "28", features = ["component-model"] }
|
||||
@@ -102,6 +114,7 @@ hkdf = "0.12"
|
||||
sha2 = "0.10"
|
||||
blake3 = "1"
|
||||
rand = "0.8"
|
||||
subtle = "2" # Constant-time comparisons for token validation
|
||||
|
||||
# Multi-provider LLM support
|
||||
rig-core = "0.30"
|
||||
@@ -134,9 +147,22 @@ pretty_assertions = "1"
|
||||
tempfile = "3"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
default = ["postgres", "libsql"]
|
||||
postgres = [
|
||||
"dep:deadpool-postgres",
|
||||
"dep:tokio-postgres",
|
||||
"dep:postgres-types",
|
||||
"dep:refinery",
|
||||
"dep:pgvector",
|
||||
"rust_decimal/db-tokio-postgres",
|
||||
]
|
||||
libsql = ["dep:libsql"]
|
||||
integration = []
|
||||
|
||||
[[example]]
|
||||
name = "test_heartbeat"
|
||||
required-features = ["postgres"]
|
||||
|
||||
# The profile that 'cargo dist' will build with
|
||||
[profile.dist]
|
||||
inherits = "release"
|
||||
@@ -151,7 +177,7 @@ ci = "github"
|
||||
# The installers to generate for each app
|
||||
installers = ["shell", "powershell", "npm", "msi"]
|
||||
# Publish jobs to run in CI
|
||||
publish-jobs = ["npm"]
|
||||
publish-jobs = []
|
||||
# Target platforms to build apps for (Rust target-triple syntax)
|
||||
targets = [
|
||||
"aarch64-apple-darwin",
|
||||
@@ -165,11 +191,13 @@ windows-archive = ".tar.gz"
|
||||
# The archive format to use for non-windows builds (defaults .tar.xz)
|
||||
unix-archive = ".tar.gz"
|
||||
# Which actions to run on pull requests
|
||||
pr-run-mode = "upload"
|
||||
pr-run-mode = "skip"
|
||||
# Path that installers should place binaries in
|
||||
install-path = "CARGO_HOME"
|
||||
# Whether to install an updater program
|
||||
install-updater = true
|
||||
# Cache intermediate build artifacts to speed up the release pipelines
|
||||
cache-builds = true
|
||||
|
||||
[workspace.metadata.dist.github-custom-runners]
|
||||
aarch64-unknown-linux-gnu = "ubuntu-24.04-arm"
|
||||
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
# Multi-stage Dockerfile for the IronClaw agent (cloud deployment).
|
||||
#
|
||||
# Build:
|
||||
# docker build --platform linux/amd64 -t ironclaw:latest .
|
||||
#
|
||||
# Run:
|
||||
# docker run --env-file .env -p 3000:3000 ironclaw:latest
|
||||
|
||||
# Stage 1: Build
|
||||
FROM rust:1.92-slim-bookworm AS builder
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
pkg-config libssl-dev cmake gcc g++ \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy manifests first for layer caching
|
||||
COPY Cargo.toml Cargo.lock ./
|
||||
|
||||
# Copy source and build artifacts
|
||||
COPY src/ src/
|
||||
COPY migrations/ migrations/
|
||||
COPY wit/ wit/
|
||||
|
||||
RUN cargo build --release --bin ironclaw
|
||||
|
||||
# Stage 2: Runtime
|
||||
FROM debian:bookworm-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates libssl3 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=builder /app/target/release/ironclaw /usr/local/bin/ironclaw
|
||||
COPY --from=builder /app/migrations /app/migrations
|
||||
|
||||
# Non-root user
|
||||
RUN useradd -m -u 1000 -s /bin/bash ironclaw
|
||||
USER ironclaw
|
||||
|
||||
EXPOSE 3000
|
||||
|
||||
ENV RUST_LOG=ironclaw=info
|
||||
|
||||
ENTRYPOINT ["ironclaw"]
|
||||
+2
-2
@@ -9,7 +9,7 @@
|
||||
# The image includes common development tools so workers can build software,
|
||||
# run tests, and execute shell commands.
|
||||
|
||||
FROM rust:1.85-bookworm AS builder
|
||||
FROM rust:1.92-bookworm AS builder
|
||||
|
||||
WORKDIR /build
|
||||
COPY . .
|
||||
@@ -40,7 +40,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ENV RUSTUP_HOME=/usr/local/rustup \
|
||||
CARGO_HOME=/usr/local/cargo \
|
||||
PATH=/usr/local/cargo/bin:$PATH
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.85.0 \
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.92.0 \
|
||||
&& chmod -R a+r /usr/local/rustup /usr/local/cargo
|
||||
|
||||
# Install Claude Code CLI (for claude-bridge mode)
|
||||
|
||||
+15
-19
@@ -37,7 +37,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Session management/routing | ✅ | ✅ | SessionManager exists |
|
||||
| Configuration hot-reload | ✅ | ❌ | |
|
||||
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
|
||||
| OpenAI-compatible HTTP API | ✅ | ❌ | /v1/chat/completions |
|
||||
| OpenAI-compatible HTTP API | ✅ | ✅ | /v1/chat/completions |
|
||||
| Canvas hosting | ✅ | ❌ | Agent-driven UI |
|
||||
| Gateway lock (PID-based) | ✅ | ❌ | |
|
||||
| launchd/systemd integration | ✅ | ❌ | |
|
||||
@@ -112,7 +112,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| `pairing` | ✅ | ✅ | - | list/approve for channel DM pairing |
|
||||
| `nodes` | ✅ | ❌ | P3 | Device management |
|
||||
| `plugins` | ✅ | ❌ | P3 | Plugin management |
|
||||
| `hooks` | ✅ | ❌ | P2 | Lifecycle hooks |
|
||||
| `hooks` | ✅ | ✅ | P2 | Lifecycle hooks |
|
||||
| `cron` | ✅ | ❌ | P2 | Scheduled jobs |
|
||||
| `webhooks` | ✅ | ❌ | P3 | Webhook config |
|
||||
| `message send` | ✅ | ❌ | P2 | Send to channels |
|
||||
@@ -133,7 +133,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|---------|----------|----------|-------|
|
||||
| Pi agent runtime | ✅ | ➖ | IronClaw uses custom runtime |
|
||||
| RPC-based execution | ✅ | ✅ | Orchestrator/worker pattern |
|
||||
| Multi-provider failover | ✅ | ❌ | Provider fallback chains |
|
||||
| Multi-provider failover | ✅ | ✅ | `FailoverProvider` tries providers sequentially on retryable errors |
|
||||
| Per-sender sessions | ✅ | ✅ | |
|
||||
| Global sessions | ✅ | ❌ | Optional shared context |
|
||||
| Session pruning | ✅ | ❌ | Auto cleanup old sessions |
|
||||
@@ -164,7 +164,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| AWS Bedrock | ✅ | ❌ | P3 | |
|
||||
| Google Gemini | ✅ | ❌ | P3 | |
|
||||
| OpenRouter | ✅ | ❌ | P3 | |
|
||||
| Ollama (local) | ✅ | ❌ | P2 | Local models |
|
||||
| Ollama (local) | ✅ | ✅ | - | via `rig::providers::ollama` (full support) |
|
||||
| node-llama-cpp | ✅ | ➖ | - | N/A for Rust |
|
||||
| llama.cpp (native) | ❌ | 🔮 | P3 | Rust bindings |
|
||||
|
||||
@@ -173,8 +173,8 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| Auto-discovery | ✅ | ❌ | |
|
||||
| Failover chains | ✅ | ❌ | Provider fallback |
|
||||
| Cooldown management | ✅ | ❌ | Skip failed providers |
|
||||
| Failover chains | ✅ | ✅ | `FailoverProvider` with configurable `fallback_model` |
|
||||
| Cooldown management | ✅ | ✅ | Lock-free per-provider cooldown in `FailoverProvider` |
|
||||
| Per-session model override | ✅ | ✅ | Model selector in TUI |
|
||||
| Model selection UI | ✅ | ✅ | TUI keyboard shortcut |
|
||||
|
||||
@@ -323,14 +323,14 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Cron jobs | ✅ | ✅ | - | Routines with cron trigger |
|
||||
| Timezone support | ✅ | ✅ | - | Via cron expressions |
|
||||
| One-shot/recurring jobs | ✅ | ✅ | - | Manual + cron triggers |
|
||||
| `beforeInbound` hook | ✅ | ❌ | P2 | |
|
||||
| `beforeOutbound` hook | ✅ | ❌ | P2 | |
|
||||
| `beforeToolCall` hook | ✅ | ❌ | P2 | |
|
||||
| `beforeInbound` hook | ✅ | ✅ | P2 | |
|
||||
| `beforeOutbound` hook | ✅ | ✅ | P2 | |
|
||||
| `beforeToolCall` hook | ✅ | ✅ | P2 | |
|
||||
| `onMessage` hook | ✅ | ✅ | - | Routines with event trigger |
|
||||
| `onSessionStart` hook | ✅ | ❌ | P2 | |
|
||||
| `onSessionEnd` hook | ✅ | ❌ | P2 | |
|
||||
| `onSessionStart` hook | ✅ | ✅ | P2 | |
|
||||
| `onSessionEnd` hook | ✅ | ✅ | P2 | |
|
||||
| `transcribeAudio` hook | ✅ | ❌ | P3 | |
|
||||
| `transformResponse` hook | ✅ | ❌ | P2 | |
|
||||
| `transformResponse` hook | ✅ | ✅ | P2 | |
|
||||
| Bundled hooks | ✅ | ❌ | P2 | |
|
||||
| Plugin hooks | ✅ | ❌ | P3 | |
|
||||
| Workspace hooks | ✅ | ❌ | P2 | Inline code |
|
||||
@@ -419,15 +419,11 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
- ❌ Slack channel (real implementation)
|
||||
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
|
||||
- ❌ WhatsApp channel
|
||||
- ❌ Multi-provider failover
|
||||
- ❌ Hooks system (beforeInbound, beforeToolCall, etc.)
|
||||
- ✅ Multi-provider failover (`FailoverProvider` with retryable error classification)
|
||||
- ✅ Hooks system (beforeInbound, beforeToolCall, beforeOutbound, onSessionStart, onSessionEnd, transformResponse)
|
||||
|
||||
### P2 - Medium Priority
|
||||
- ❌ Cron job scheduling
|
||||
- ❌ Web Control UI
|
||||
- ❌ WebChat channel
|
||||
- 🚧 Media handling (caption support; no image/PDF processing)
|
||||
- ❌ CLI subcommands (config, status, memory, doctor)
|
||||
- ❌ Media handling (images, PDFs)
|
||||
- ❌ Ollama/local model support
|
||||
- ❌ Configuration hot-reload
|
||||
- ❌ Webhook trigger endpoint in web gateway
|
||||
|
||||
@@ -99,22 +99,6 @@ curl --proto '=https' --tlsv1.2 -LsSf https://github.com/nearai/ironclaw/release
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Run via npx (Node.js on Windows, Linux, macOS)</summary>
|
||||
|
||||
```sh
|
||||
npx ironclaw
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Use in package.json scripts (Node.js on Windows, Linux, macOS)</summary>
|
||||
|
||||
```sh
|
||||
npm install ironclaw
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Compile the source code (Cargo on Windows, Linux, macOS)</summary>
|
||||
|
||||
@@ -197,42 +181,42 @@ External content passes through multiple security layers:
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────────────┐
|
||||
│ Channels │
|
||||
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||
│ │ REPL │ │ HTTP │ │WASM Channels│ │ Web Gateway │ │
|
||||
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
|
||||
│ │ │ │ └──────┬──────┘ │
|
||||
│ └─────────┴──────────────┴────────────────┘ │
|
||||
│ │ │
|
||||
│ ┌─────────▼─────────┐ │
|
||||
│ │ Agent Loop │ Intent routing │
|
||||
│ └────┬─────────┬────┘ │
|
||||
│ │ │ │
|
||||
│ ┌──────────▼───┐ ┌──▼──────────────┐ │
|
||||
│ │ Scheduler │ │ Routines Engine │ │
|
||||
│ │(parallel jobs)│ │(cron, event, wh) │ │
|
||||
│ └──────┬───────┘ └────────┬─────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌─────────────┼───────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌───▼────┐ ┌────▼────────────────┐ │
|
||||
│ │ Local │ │ Orchestrator │ │
|
||||
│ │Workers │ │ ┌───────────────┐ │ │
|
||||
│ │(in-proc)│ │ │ Docker Sandbox│ │ │
|
||||
│ └───┬────┘ │ │ Containers │ │ │
|
||||
│ │ │ │ ┌───────────┐ │ │ │
|
||||
│ │ │ │ │Worker / CC│ │ │ │
|
||||
│ │ │ │ └───────────┘ │ │ │
|
||||
│ │ │ └───────────────┘ │ │
|
||||
│ │ └─────────┬───────────┘ │
|
||||
│ └──────────────────┤ │
|
||||
│ │ │
|
||||
│ ┌───────────▼──────────┐ │
|
||||
│ │ Tool Registry │ │
|
||||
│ │ Built-in, MCP, WASM │ │
|
||||
│ └──────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────────┘
|
||||
┌────────────────────────────────────────────────────────────────┐
|
||||
│ Channels │
|
||||
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||
│ │ REPL │ │ HTTP │ │WASM Channels│ │ Web Gateway │ │
|
||||
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
|
||||
│ │ │ │ └──────┬──────┘ │
|
||||
│ └─────────┴──────────────┴────────────────┘ │
|
||||
│ │ │
|
||||
│ ┌─────────▼─────────┐ │
|
||||
│ │ Agent Loop │ Intent routing │
|
||||
│ └────┬──────────┬───┘ │
|
||||
│ │ │ │
|
||||
│ ┌──────────▼────┐ ┌──▼───────────────┐ │
|
||||
│ │ Scheduler │ │ Routines Engine │ │
|
||||
│ │(parallel jobs)│ │(cron, event, wh) │ │
|
||||
│ └──────┬────────┘ └────────┬─────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌─────────────┼────────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌───▼─────┐ ┌────▼────────────────┐ │
|
||||
│ │ Local │ │ Orchestrator │ │
|
||||
│ │Workers │ │ ┌───────────────┐ │ │
|
||||
│ │(in-proc)│ │ │ Docker Sandbox│ │ │
|
||||
│ └───┬─────┘ │ │ Containers │ │ │
|
||||
│ │ │ │ ┌───────────┐ │ │ │
|
||||
│ │ │ │ │Worker / CC│ │ │ │
|
||||
│ │ │ │ └───────────┘ │ │ │
|
||||
│ │ │ └───────────────┘ │ │
|
||||
│ │ └─────────┬───────────┘ │
|
||||
│ └──────────────────┤ │
|
||||
│ │ │
|
||||
│ ┌───────────▼──────────┐ │
|
||||
│ │ Tool Registry │ │
|
||||
│ │ Built-in, MCP, WASM │ │
|
||||
│ └──────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Core Components
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
[package]
|
||||
name = "discord-channel"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "Discord channel for IronClaw"
|
||||
license = "MIT OR Apache-2.0"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
wit-bindgen = "0.41.0"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[profile.release]
|
||||
strip = true
|
||||
opt-level = "s"
|
||||
lto = true
|
||||
codegen-units = 1
|
||||
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
# Discord Channel for IronClaw
|
||||
|
||||
WASM channel for Discord integration - handle slash commands and button interactions via webhooks.
|
||||
|
||||
## Features
|
||||
|
||||
- **Slash Commands** - Process Discord slash commands
|
||||
- **Button Interactions** - Handle button clicks
|
||||
- **Thread Support** - Respond in threads
|
||||
- **DM Support** - Handle direct messages
|
||||
|
||||
## Setup
|
||||
|
||||
1. Create a Discord Application at <https://discord.com/developers/applications>
|
||||
2. Create a Bot and get the token
|
||||
3. Set up Interactions URL to point to your IronClaw instance
|
||||
4. Copy the Application ID and Public Key
|
||||
5. Store in IronClaw secrets:
|
||||
|
||||
```bash
|
||||
ironclaw secret set discord_bot_token YOUR_BOT_TOKEN
|
||||
```
|
||||
|
||||
**Note:** The `discord_bot_token` secret is the only value read directly by this
|
||||
Discord channel WASM component. The `discord_app_id` and `discord_public_key`
|
||||
secrets are used by the IronClaw host (for example, to verify Discord
|
||||
interaction signatures and manage slash command registration) and are not
|
||||
accessed from the WASM module itself.
|
||||
|
||||
## Discord Configuration
|
||||
|
||||
### Register Slash Commands
|
||||
|
||||
```bash
|
||||
curl -X POST \
|
||||
-H "Authorization: Bot YOUR_BOT_TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
https://discord.com/api/v10/applications/YOUR_APP_ID/commands \
|
||||
-d '{
|
||||
"name": "ask",
|
||||
"description": "Ask the AI agent",
|
||||
"options": [{
|
||||
"name": "question",
|
||||
"description": "Your question",
|
||||
"type": 3,
|
||||
"required": true
|
||||
}]
|
||||
}'
|
||||
```
|
||||
|
||||
### Set Interactions Endpoint
|
||||
|
||||
In your Discord app settings, set:
|
||||
|
||||
- Interactions Endpoint URL: `https://your-ironclaw.com/webhook/discord`
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Slash Command
|
||||
|
||||
User types: `/ask question: What is the weather?`
|
||||
|
||||
The agent receives:
|
||||
|
||||
```text
|
||||
User: @username
|
||||
Content: /ask question: What is the weather?
|
||||
```
|
||||
|
||||
### Button Click
|
||||
|
||||
When a user clicks a button in a message, the agent receives:
|
||||
|
||||
```text
|
||||
User: @username
|
||||
Content: [Button clicked] Original message content
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
If an internal error occurs (e.g., metadata serialization failure), the tool attempts to send an ephemeral message to the user:
|
||||
|
||||
```text
|
||||
❌ Internal Error: Failed to process command metadata.
|
||||
```
|
||||
|
||||
Check the host logs for detailed error information.
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Embeds
|
||||
|
||||
To send embeds, include an `embeds` array in the `metadata_json` field of the agent's response. The structure should match the Discord API `embed` object.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "Invalid Signature"
|
||||
|
||||
- Check that `discord_public_key` is set correctly in IronClaw secrets.
|
||||
- This validation happens on the host before reaching the WASM.
|
||||
|
||||
### "401 Unauthorized"
|
||||
|
||||
- Check that `discord_bot_token` is set correctly in IronClaw secrets.
|
||||
- Ensure the bot is added to the server.
|
||||
|
||||
### "Interaction Failed"
|
||||
|
||||
- The interaction might have timed out (Discord requires a response within 3 seconds).
|
||||
- The `interactions_endpoint_url` might be unreachable.
|
||||
|
||||
## Building
|
||||
|
||||
```bash
|
||||
cd channels-src/discord
|
||||
cargo build --target wasm32-wasi --release
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT/Apache-2.0
|
||||
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"type": "channel",
|
||||
"name": "discord",
|
||||
"description": "Discord Gateway/Webhook channel for handling slash commands, buttons, and messages",
|
||||
"capabilities": {
|
||||
"http": {
|
||||
"allowlist": [
|
||||
{ "host": "discord.com", "path_prefix": "/api/v10" }
|
||||
],
|
||||
"credentials": {
|
||||
"discord_bot_token": {
|
||||
"secret_name": "discord_bot_token",
|
||||
"location": { "type": "header", "header_name": "Authorization", "prefix": "Bot " },
|
||||
"host_patterns": ["discord.com"]
|
||||
}
|
||||
},
|
||||
"rate_limit": {
|
||||
"requests_per_minute": 60,
|
||||
"requests_per_hour": 3600
|
||||
}
|
||||
},
|
||||
"secrets": {
|
||||
"allowed_names": ["discord_bot_token", "discord_*"]
|
||||
},
|
||||
"channel": {
|
||||
"allowed_paths": ["/webhook/discord"],
|
||||
"allow_polling": false,
|
||||
"callback_timeout_secs": 45,
|
||||
"workspace_prefix": "channels/discord/",
|
||||
"emit_rate_limit": {
|
||||
"messages_per_minute": 100,
|
||||
"messages_per_hour": 5000
|
||||
}
|
||||
}
|
||||
},
|
||||
"config": {
|
||||
"require_signature_verification": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
//! Discord Gateway/Webhook channel for IronClaw.
|
||||
//!
|
||||
//! This WASM component implements the channel interface for handling Discord
|
||||
//! interactions via webhooks and sending messages back to Discord.
|
||||
//!
|
||||
//! # Features
|
||||
//!
|
||||
//! - URL verification for Discord interactions
|
||||
//! - Slash command handling
|
||||
//! - Message event parsing (@mentions, DMs)
|
||||
//! - Thread support for conversations
|
||||
//! - Response posting via Discord Web API
|
||||
//! - Automatic message truncation (> 2000 chars)
|
||||
//!
|
||||
//! # Security
|
||||
//!
|
||||
//! - Signature validation is handled by the host (webhook secrets)
|
||||
//! - Bot token is injected by host during HTTP requests
|
||||
//! - WASM never sees raw credentials
|
||||
|
||||
wit_bindgen::generate!({
|
||||
world: "sandboxed-channel",
|
||||
path: "../../wit/channel.wit",
|
||||
});
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use exports::near::agent::channel::{
|
||||
AgentResponse, ChannelConfig, Guest, HttpEndpointConfig, IncomingHttpRequest,
|
||||
OutgoingHttpResponse, StatusUpdate,
|
||||
};
|
||||
use near::agent::channel_host::{self, EmittedMessage};
|
||||
|
||||
/// Discord interaction wrapper.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct DiscordInteraction {
|
||||
/// Interaction type (1=Ping, 2=ApplicationCommand, 3=MessageComponent)
|
||||
#[serde(rename = "type")]
|
||||
interaction_type: u8,
|
||||
|
||||
/// Interaction ID
|
||||
id: String,
|
||||
|
||||
/// Application ID
|
||||
application_id: String,
|
||||
|
||||
/// Guild ID (if in server)
|
||||
#[allow(dead_code)] // Part of API payload, currently unused
|
||||
guild_id: Option<String>,
|
||||
|
||||
/// Channel ID
|
||||
channel_id: Option<String>,
|
||||
|
||||
/// Member info (if in server)
|
||||
member: Option<DiscordMember>,
|
||||
|
||||
/// User info (if DM)
|
||||
user: Option<DiscordUser>,
|
||||
|
||||
/// Command data (for slash commands)
|
||||
data: Option<DiscordCommandData>,
|
||||
|
||||
/// Message (for component interactions)
|
||||
message: Option<DiscordMessage>,
|
||||
|
||||
/// Token for responding
|
||||
token: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
struct DiscordMember {
|
||||
user: DiscordUser,
|
||||
#[allow(dead_code)] // Part of API payload, currently unused
|
||||
nick: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
struct DiscordUser {
|
||||
id: String,
|
||||
username: String,
|
||||
global_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
struct DiscordCommandData {
|
||||
#[allow(dead_code)] // Part of API payload, currently unused
|
||||
id: String,
|
||||
name: String,
|
||||
options: Option<Vec<DiscordCommandOption>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
struct DiscordCommandOption {
|
||||
name: String,
|
||||
value: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
struct DiscordMessage {
|
||||
#[allow(dead_code)] // Part of API payload, currently unused
|
||||
id: String,
|
||||
content: String,
|
||||
channel_id: String,
|
||||
#[allow(dead_code)] // Part of API payload, currently unused
|
||||
author: DiscordUser,
|
||||
}
|
||||
|
||||
/// Metadata stored with emitted messages for response routing.
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct DiscordMessageMetadata {
|
||||
/// Discord channel ID
|
||||
channel_id: String,
|
||||
|
||||
/// Interaction ID for followups
|
||||
interaction_id: String,
|
||||
|
||||
/// Interaction token for responding
|
||||
token: String,
|
||||
|
||||
/// Application ID
|
||||
application_id: String,
|
||||
|
||||
/// Thread ID (for forum threads)
|
||||
thread_id: Option<String>,
|
||||
}
|
||||
|
||||
struct DiscordChannel;
|
||||
|
||||
impl Guest for DiscordChannel {
|
||||
fn on_start(_config_json: String) -> Result<ChannelConfig, String> {
|
||||
channel_host::log(channel_host::LogLevel::Info, "Discord channel starting");
|
||||
|
||||
Ok(ChannelConfig {
|
||||
display_name: "Discord".to_string(),
|
||||
http_endpoints: vec![HttpEndpointConfig {
|
||||
path: "/webhook/discord".to_string(),
|
||||
methods: vec!["POST".to_string()],
|
||||
require_secret: true,
|
||||
}],
|
||||
poll: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn on_http_request(req: IncomingHttpRequest) -> OutgoingHttpResponse {
|
||||
let body_str = match std::str::from_utf8(&req.body) {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
return json_response(400, serde_json::json!({"error": "Invalid UTF-8 body"}));
|
||||
}
|
||||
};
|
||||
|
||||
let interaction: DiscordInteraction = match serde_json::from_str(body_str) {
|
||||
Ok(i) => i,
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to parse Discord interaction: {}", e),
|
||||
);
|
||||
return json_response(400, serde_json::json!({"error": "Invalid interaction"}));
|
||||
}
|
||||
};
|
||||
|
||||
match interaction.interaction_type {
|
||||
// Ping - Discord verification
|
||||
1 => {
|
||||
channel_host::log(channel_host::LogLevel::Info, "Responding to Discord ping");
|
||||
json_response(200, serde_json::json!({"type": 1}))
|
||||
}
|
||||
|
||||
// Application Command (slash command)
|
||||
2 => {
|
||||
handle_slash_command(&interaction);
|
||||
json_response(
|
||||
200,
|
||||
serde_json::json!({
|
||||
"type": 5,
|
||||
"data": {
|
||||
"content": "🤔 Thinking..."
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
// Message Component (buttons, selects)
|
||||
3 => {
|
||||
if let Some(ref message) = interaction.message {
|
||||
handle_message_component(&interaction, message);
|
||||
}
|
||||
json_response(200, serde_json::json!({"type": 6}))
|
||||
}
|
||||
|
||||
_ => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Warn,
|
||||
&format!(
|
||||
"Unknown Discord interaction type: {}",
|
||||
interaction.interaction_type
|
||||
),
|
||||
);
|
||||
json_response(200, serde_json::json!({"type": 6}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn on_poll() {}
|
||||
|
||||
fn on_respond(response: AgentResponse) -> Result<(), String> {
|
||||
let metadata: DiscordMessageMetadata = serde_json::from_str(&response.metadata_json)
|
||||
.map_err(|e| format!("Failed to parse metadata: {}", e))?;
|
||||
|
||||
// Use webhook endpoint for followup
|
||||
let url = format!(
|
||||
"https://discord.com/api/v10/webhooks/{}/{}",
|
||||
metadata.application_id, metadata.token
|
||||
);
|
||||
|
||||
// Truncate content to 2000 characters to comply with Discord limits
|
||||
let content = truncate_message(&response.content);
|
||||
|
||||
let mut payload = serde_json::json!({
|
||||
"content": content,
|
||||
});
|
||||
|
||||
// Check for embeds in metadata
|
||||
if let Ok(meta_json) = serde_json::from_str::<serde_json::Value>(&response.metadata_json) {
|
||||
if let Some(embeds) = meta_json.get("embeds") {
|
||||
payload["embeds"] = embeds.clone();
|
||||
}
|
||||
}
|
||||
|
||||
let payload_bytes =
|
||||
serde_json::to_vec(&payload).map_err(|e| format!("Failed to serialize: {}", e))?;
|
||||
|
||||
let headers = serde_json::json!({
|
||||
"Content-Type": "application/json"
|
||||
});
|
||||
|
||||
let result = channel_host::http_request(
|
||||
"POST",
|
||||
&url,
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(http_response) => {
|
||||
if http_response.status >= 200 && http_response.status < 300 {
|
||||
channel_host::log(channel_host::LogLevel::Debug, "Posted followup to Discord");
|
||||
Ok(())
|
||||
} else {
|
||||
let body_str = String::from_utf8_lossy(&http_response.body);
|
||||
Err(format!(
|
||||
"Discord API error: {} - {}",
|
||||
http_response.status, body_str
|
||||
))
|
||||
}
|
||||
}
|
||||
Err(e) => Err(format!("HTTP request failed: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
fn on_status(_update: StatusUpdate) {}
|
||||
|
||||
fn on_shutdown() {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Info,
|
||||
"Discord channel shutting down",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_slash_command(interaction: &DiscordInteraction) {
|
||||
let user = interaction
|
||||
.member
|
||||
.as_ref()
|
||||
.map(|m| &m.user)
|
||||
.or(interaction.user.as_ref());
|
||||
let user_id = user.map(|u| u.id.clone()).unwrap_or_default();
|
||||
let user_name = user
|
||||
.map(|u| {
|
||||
u.global_name
|
||||
.as_ref()
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or(&u.username)
|
||||
.clone()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let channel_id = interaction.channel_id.clone().unwrap_or_default();
|
||||
|
||||
let command_name = interaction
|
||||
.data
|
||||
.as_ref()
|
||||
.map(|d| d.name.clone())
|
||||
.unwrap_or_default();
|
||||
let options = interaction.data.as_ref().and_then(|d| d.options.clone());
|
||||
|
||||
let content = if let Some(opts) = options {
|
||||
let opt_str = opts
|
||||
.iter()
|
||||
.map(|o| format!("{}: {}", o.name, o.value))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
format!("/{} {}", command_name, opt_str)
|
||||
} else {
|
||||
format!("/{}", command_name)
|
||||
};
|
||||
|
||||
let metadata = DiscordMessageMetadata {
|
||||
channel_id: channel_id.clone(),
|
||||
interaction_id: interaction.id.clone(),
|
||||
token: interaction.token.clone(),
|
||||
application_id: interaction.application_id.clone(),
|
||||
thread_id: None,
|
||||
};
|
||||
|
||||
let metadata_json = match serde_json::to_string(&metadata) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to serialize metadata: {}", e),
|
||||
);
|
||||
// Attempt to notify user of internal error
|
||||
let url = format!(
|
||||
"https://discord.com/api/v10/webhooks/{}/{}",
|
||||
interaction.application_id, interaction.token
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"content": "❌ Internal Error: Failed to process command metadata.",
|
||||
"flags": 64 // Ephemeral
|
||||
});
|
||||
let _ = channel_host::http_request(
|
||||
"POST",
|
||||
&url,
|
||||
&serde_json::json!({"Content-Type": "application/json"}).to_string(),
|
||||
Some(&serde_json::to_vec(&payload).unwrap_or_default()),
|
||||
None,
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
channel_host::emit_message(&EmittedMessage {
|
||||
user_id,
|
||||
user_name: Some(user_name),
|
||||
content,
|
||||
thread_id: None,
|
||||
metadata_json,
|
||||
});
|
||||
}
|
||||
|
||||
fn handle_message_component(interaction: &DiscordInteraction, message: &DiscordMessage) {
|
||||
// Check member first (for server contexts), then user (for DMs)
|
||||
let user = interaction
|
||||
.member
|
||||
.as_ref()
|
||||
.map(|m| &m.user)
|
||||
.or(interaction.user.as_ref());
|
||||
let user_id = user.map(|u| u.id.clone()).unwrap_or_default();
|
||||
let user_name = user
|
||||
.map(|u| {
|
||||
u.global_name
|
||||
.as_ref()
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or(&u.username)
|
||||
.clone()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let channel_id = message.channel_id.clone();
|
||||
|
||||
let metadata = DiscordMessageMetadata {
|
||||
channel_id: channel_id.clone(),
|
||||
interaction_id: interaction.id.clone(),
|
||||
token: interaction.token.clone(),
|
||||
application_id: interaction.application_id.clone(),
|
||||
thread_id: None,
|
||||
};
|
||||
|
||||
let metadata_json = match serde_json::to_string(&metadata) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to serialize metadata: {}", e),
|
||||
);
|
||||
return; // Don't emit message if metadata can't be serialized
|
||||
}
|
||||
};
|
||||
|
||||
channel_host::emit_message(&EmittedMessage {
|
||||
user_id,
|
||||
user_name: Some(user_name),
|
||||
content: format!("[Button clicked] {}", message.content),
|
||||
thread_id: None,
|
||||
metadata_json,
|
||||
});
|
||||
}
|
||||
|
||||
fn json_response(status: u16, value: serde_json::Value) -> OutgoingHttpResponse {
|
||||
let body = serde_json::to_vec(&value).unwrap_or_default();
|
||||
let headers = serde_json::json!({"Content-Type": "application/json"});
|
||||
|
||||
OutgoingHttpResponse {
|
||||
status,
|
||||
headers_json: headers.to_string(),
|
||||
body,
|
||||
}
|
||||
}
|
||||
|
||||
export!(DiscordChannel);
|
||||
|
||||
fn truncate_message(content: &str) -> String {
|
||||
if content.len() <= 2000 {
|
||||
content.to_string()
|
||||
} else {
|
||||
let max_bytes = 1990;
|
||||
let cutoff = content
|
||||
.char_indices()
|
||||
.map(|(i, c)| i + c.len_utf8())
|
||||
.take_while(|&end| end <= max_bytes)
|
||||
.last()
|
||||
.unwrap_or(0);
|
||||
let mut truncated = content[..cutoff].to_string();
|
||||
truncated.push_str("\n... (truncated)");
|
||||
truncated
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_truncate_message() {
|
||||
let short = "Hello world";
|
||||
assert_eq!(truncate_message(short), short);
|
||||
|
||||
let long = "a".repeat(2005);
|
||||
let truncated = truncate_message(&long);
|
||||
assert_eq!(truncated.len(), 2006); // 1990 + 16 chars suffix
|
||||
assert!(truncated.ends_with("\n... (truncated)"));
|
||||
|
||||
// Test with multibyte characters (Euro sign is 3 bytes)
|
||||
// 1000 chars * 3 bytes = 3000 bytes
|
||||
let multi = "€".repeat(1000);
|
||||
let truncated_multi = truncate_message(&multi);
|
||||
|
||||
// 1990 bytes limit. 1990 / 3 = 663 with remainder 1.
|
||||
// Should truncate at 663 chars (1989 bytes).
|
||||
// Suffix is 16 bytes. Total: 1989 + 16 = 2005 bytes.
|
||||
assert!(truncated_multi.len() <= 2006);
|
||||
assert!(truncated_multi.len() >= 2006 - 4); // Allow for max utf8 char width variance
|
||||
assert!(truncated_multi.ends_with("\n... (truncated)"));
|
||||
|
||||
let content_part = &truncated_multi[..truncated_multi.len() - 16];
|
||||
assert!(content_part.chars().all(|c| c == '€'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_metadata_serialization() {
|
||||
let metadata = DiscordMessageMetadata {
|
||||
channel_id: "123".into(),
|
||||
interaction_id: "456".into(),
|
||||
token: "abc".into(),
|
||||
application_id: "789".into(),
|
||||
thread_id: None,
|
||||
};
|
||||
let json = serde_json::to_string(&metadata).unwrap();
|
||||
let parsed: DiscordMessageMetadata = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.channel_id, "123");
|
||||
assert_eq!(parsed.interaction_id, "456");
|
||||
}
|
||||
}
|
||||
@@ -338,7 +338,13 @@ fn emit_message(
|
||||
team_id,
|
||||
};
|
||||
|
||||
let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
|
||||
let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|e| {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to serialize Slack metadata: {}", e),
|
||||
);
|
||||
"{}".to_string()
|
||||
});
|
||||
|
||||
// Strip @ mentions of the bot from the text for cleaner messages
|
||||
let cleaned_text = strip_bot_mention(&text);
|
||||
@@ -366,7 +372,13 @@ fn strip_bot_mention(text: &str) -> String {
|
||||
|
||||
/// Create a JSON HTTP response.
|
||||
fn json_response(status: u16, value: serde_json::Value) -> OutgoingHttpResponse {
|
||||
let body = serde_json::to_vec(&value).unwrap_or_default();
|
||||
let body = serde_json::to_vec(&value).unwrap_or_else(|e| {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to serialize JSON response: {}", e),
|
||||
);
|
||||
Vec::new()
|
||||
});
|
||||
let headers = serde_json::json!({"Content-Type": "application/json"});
|
||||
|
||||
OutgoingHttpResponse {
|
||||
|
||||
@@ -285,11 +285,7 @@ impl Guest for TelegramChannel {
|
||||
}
|
||||
|
||||
// Persist dm_policy and allow_from for DM pairing in handle_message
|
||||
let dm_policy = config
|
||||
.dm_policy
|
||||
.as_deref()
|
||||
.unwrap_or("pairing")
|
||||
.to_string();
|
||||
let dm_policy = config.dm_policy.as_deref().unwrap_or("pairing").to_string();
|
||||
let _ = channel_host::workspace_write(DM_POLICY_PATH, &dm_policy);
|
||||
|
||||
let allow_from_json = serde_json::to_string(&config.allow_from.unwrap_or_default())
|
||||
@@ -844,8 +840,8 @@ fn send_pairing_reply(chat_id: i64, code: &str) -> Result<(), String> {
|
||||
"parse_mode": "Markdown",
|
||||
});
|
||||
|
||||
let payload_bytes = serde_json::to_vec(&payload)
|
||||
.map_err(|e| format!("Failed to serialize payload: {}", e))?;
|
||||
let payload_bytes =
|
||||
serde_json::to_vec(&payload).map_err(|e| format!("Failed to serialize payload: {}", e))?;
|
||||
|
||||
let headers = serde_json::json!({
|
||||
"Content-Type": "application/json"
|
||||
@@ -856,6 +852,7 @@ fn send_pairing_reply(chat_id: i64, code: &str) -> Result<(), String> {
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
@@ -914,15 +911,10 @@ fn handle_message(message: TelegramMessage) {
|
||||
let is_private = message.chat.chat_type == "private";
|
||||
|
||||
// Owner validation: when owner_id is set, only that user can message
|
||||
let owner_configured = channel_host::workspace_read(OWNER_ID_PATH)
|
||||
.map(|s| !s.is_empty())
|
||||
.unwrap_or(false);
|
||||
let owner_id_str = channel_host::workspace_read(OWNER_ID_PATH).filter(|s| !s.is_empty());
|
||||
|
||||
if owner_configured {
|
||||
if let Ok(owner_id) = channel_host::workspace_read(OWNER_ID_PATH)
|
||||
.unwrap()
|
||||
.parse::<i64>()
|
||||
{
|
||||
if let Some(ref id_str) = owner_id_str {
|
||||
if let Ok(owner_id) = id_str.parse::<i64>() {
|
||||
if from.id != owner_id {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
@@ -936,8 +928,8 @@ fn handle_message(message: TelegramMessage) {
|
||||
}
|
||||
} else if is_private {
|
||||
// No owner_id: apply dm_policy for private chats
|
||||
let dm_policy = channel_host::workspace_read(DM_POLICY_PATH)
|
||||
.unwrap_or_else(|| "pairing".to_string());
|
||||
let dm_policy =
|
||||
channel_host::workspace_read(DM_POLICY_PATH).unwrap_or_else(|| "pairing".to_string());
|
||||
|
||||
if dm_policy != "open" {
|
||||
// Build effective allow list: config allow_from + pairing store
|
||||
@@ -1000,8 +992,7 @@ fn handle_message(message: TelegramMessage) {
|
||||
|
||||
if !respond_to_all {
|
||||
let has_command = content.starts_with('/');
|
||||
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH)
|
||||
.unwrap_or_default();
|
||||
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH).unwrap_or_default();
|
||||
let has_bot_mention = if bot_username.is_empty() {
|
||||
content.contains('@')
|
||||
} else {
|
||||
|
||||
@@ -254,10 +254,19 @@ struct WhatsAppChannel;
|
||||
|
||||
impl Guest for WhatsAppChannel {
|
||||
fn on_start(config_json: String) -> Result<ChannelConfig, String> {
|
||||
let config: WhatsAppConfig = serde_json::from_str(&config_json).unwrap_or(WhatsAppConfig {
|
||||
api_version: default_api_version(),
|
||||
reply_to_message: default_reply_to_message(),
|
||||
});
|
||||
let config: WhatsAppConfig = match serde_json::from_str(&config_json) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Warn,
|
||||
&format!("Failed to parse WhatsApp config, using defaults: {}", e),
|
||||
);
|
||||
WhatsAppConfig {
|
||||
api_version: default_api_version(),
|
||||
reply_to_message: default_reply_to_message(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Info,
|
||||
@@ -267,6 +276,9 @@ impl Guest for WhatsAppChannel {
|
||||
),
|
||||
);
|
||||
|
||||
// Persist api_version in workspace so on_respond() can read it
|
||||
let _ = channel_host::workspace_write("channels/whatsapp/api_version", &config.api_version);
|
||||
|
||||
// WhatsApp Cloud API is webhook-only, no polling available
|
||||
Ok(ChannelConfig {
|
||||
display_name: "WhatsApp".to_string(),
|
||||
@@ -327,11 +339,16 @@ impl Guest for WhatsAppChannel {
|
||||
let metadata: WhatsAppMessageMetadata = serde_json::from_str(&response.metadata_json)
|
||||
.map_err(|e| format!("Failed to parse metadata: {}", e))?;
|
||||
|
||||
// Read api_version from workspace (set during on_start), fallback to default
|
||||
let api_version = channel_host::workspace_read("channels/whatsapp/api_version")
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or_else(|| "v18.0".to_string());
|
||||
|
||||
// Build WhatsApp API URL with token placeholder
|
||||
// Host will replace {WHATSAPP_ACCESS_TOKEN} with actual token in Authorization header
|
||||
let api_url = format!(
|
||||
"https://graph.facebook.com/v18.0/{}/messages",
|
||||
metadata.phone_number_id
|
||||
"https://graph.facebook.com/{}/{}/messages",
|
||||
api_version, metadata.phone_number_id
|
||||
);
|
||||
|
||||
// Build sendMessage payload
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
[Unit]
|
||||
Description=Cloud SQL Auth Proxy
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
DynamicUser=yes
|
||||
ExecStart=/usr/local/bin/cloud-sql-proxy ironclaw-prod:us-central1:ironclaw-db --port=5432
|
||||
Restart=always
|
||||
RestartSec=5
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,27 @@
|
||||
# WARNING: Replace all CHANGE_ME values before deploying.
|
||||
# Do not use placeholder passwords in production.
|
||||
DATABASE_URL=postgres://ironclaw:CHANGE_ME@localhost:5432/ironclaw
|
||||
|
||||
# NEAR AI
|
||||
NEARAI_SESSION_TOKEN=CHANGE_ME
|
||||
NEARAI_MODEL=claude-3-5-sonnet-20241022
|
||||
NEARAI_BASE_URL=https://cloud-api.near.ai
|
||||
NEARAI_AUTH_URL=https://private.near.ai
|
||||
NEARAI_API_MODE=chat_completions
|
||||
|
||||
# Agent
|
||||
AGENT_NAME=ironclaw
|
||||
CLI_ENABLED=false
|
||||
|
||||
# Web Gateway
|
||||
GATEWAY_ENABLED=true
|
||||
# 0.0.0.0 binds to all interfaces (required for Docker --network=host).
|
||||
# Use 127.0.0.1 if running outside Docker or for local-only access.
|
||||
GATEWAY_HOST=0.0.0.0
|
||||
GATEWAY_PORT=3000
|
||||
GATEWAY_AUTH_TOKEN=CHANGE_ME
|
||||
|
||||
# Disabled for initial deploy
|
||||
SANDBOX_ENABLED=false
|
||||
HEARTBEAT_ENABLED=false
|
||||
EMBEDDING_ENABLED=false
|
||||
@@ -0,0 +1,20 @@
|
||||
[Unit]
|
||||
Description=IronClaw AI Assistant
|
||||
After=cloud-sql-proxy.service docker.service
|
||||
Requires=cloud-sql-proxy.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStartPre=/usr/bin/docker pull us-central1-docker.pkg.dev/ironclaw-prod/ironclaw/agent:latest
|
||||
ExecStart=/usr/bin/docker run --rm \
|
||||
--name ironclaw \
|
||||
--env-file /opt/ironclaw/.env \
|
||||
--network=host \
|
||||
us-central1-docker.pkg.dev/ironclaw-prod/ironclaw/agent:latest \
|
||||
--no-onboard
|
||||
ExecStop=/usr/bin/docker stop ironclaw
|
||||
Restart=always
|
||||
RestartSec=10
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
Executable
+68
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env bash
|
||||
# VM bootstrap script for IronClaw on GCP Compute Engine.
|
||||
#
|
||||
# Run on a fresh Debian 12 VM after SSH:
|
||||
# sudo bash setup.sh
|
||||
#
|
||||
# Prerequisites:
|
||||
# - VM has the ironclaw-vm service account attached
|
||||
# - Cloud SQL Auth Proxy accessible via IAM
|
||||
# - Artifact Registry image pushed
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Must run as root
|
||||
if [ "$(id -u)" -ne 0 ]; then
|
||||
echo "ERROR: This script must be run as root (sudo bash setup.sh)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "==> Installing Docker"
|
||||
apt-get update
|
||||
apt-get install -y docker.io
|
||||
systemctl enable docker
|
||||
systemctl start docker
|
||||
|
||||
echo "==> Installing Cloud SQL Auth Proxy"
|
||||
curl -fsSL -o /usr/local/bin/cloud-sql-proxy \
|
||||
https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.14.3/cloud-sql-proxy.linux.amd64
|
||||
chmod +x /usr/local/bin/cloud-sql-proxy
|
||||
|
||||
echo "==> Installing systemd services"
|
||||
cp /tmp/deploy/cloud-sql-proxy.service /etc/systemd/system/
|
||||
cp /tmp/deploy/ironclaw.service /etc/systemd/system/
|
||||
systemctl daemon-reload
|
||||
|
||||
echo "==> Starting Cloud SQL Auth Proxy"
|
||||
systemctl enable cloud-sql-proxy
|
||||
systemctl start cloud-sql-proxy
|
||||
|
||||
echo "==> Configuring Docker registry auth"
|
||||
# The VM service account provides Artifact Registry access
|
||||
gcloud auth configure-docker us-central1-docker.pkg.dev --quiet
|
||||
|
||||
echo "==> Creating config directory"
|
||||
# Owned by root, readable only by root. Docker reads --env-file as root
|
||||
# before dropping to uid 1000 (ironclaw) inside the container.
|
||||
mkdir -p /opt/ironclaw
|
||||
chmod 700 /opt/ironclaw
|
||||
|
||||
if [ ! -f /opt/ironclaw/.env ]; then
|
||||
echo "WARNING: /opt/ironclaw/.env does not exist."
|
||||
echo "Create it with your configuration before starting IronClaw."
|
||||
echo "See deploy/env.example for the required variables."
|
||||
echo ""
|
||||
echo "Then run: systemctl enable ironclaw && systemctl start ironclaw"
|
||||
else
|
||||
chmod 600 /opt/ironclaw/.env
|
||||
echo "==> Starting IronClaw"
|
||||
systemctl enable ironclaw
|
||||
systemctl start ironclaw
|
||||
fi
|
||||
|
||||
echo "==> Setup complete"
|
||||
echo ""
|
||||
echo "Verify with:"
|
||||
echo " systemctl status cloud-sql-proxy"
|
||||
echo " systemctl status ironclaw"
|
||||
echo " docker logs ironclaw"
|
||||
@@ -0,0 +1,20 @@
|
||||
# Local development only — do NOT use these credentials in production.
|
||||
services:
|
||||
postgres:
|
||||
image: pgvector/pgvector:pg16
|
||||
ports:
|
||||
- "5432:5432"
|
||||
environment:
|
||||
POSTGRES_DB: ironclaw
|
||||
POSTGRES_USER: ironclaw
|
||||
POSTGRES_PASSWORD: ironclaw # dev-only, change for any non-local deployment
|
||||
volumes:
|
||||
- pgdata:/var/lib/postgresql/data
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U ironclaw"]
|
||||
interval: 5s
|
||||
timeout: 3s
|
||||
retries: 5
|
||||
|
||||
volumes:
|
||||
pgdata:
|
||||
@@ -81,7 +81,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
let session = create_session_manager(SessionConfig {
|
||||
auth_base_url: config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let llm = create_llm_provider(&config.llm, session)?;
|
||||
|
||||
+359
-142
@@ -19,16 +19,17 @@ use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusU
|
||||
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::context::JobContext;
|
||||
use crate::db::Database;
|
||||
use crate::error::Error;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// Collapse a tool output string into a single-line preview for display.
|
||||
fn truncate_for_preview(output: &str, max_chars: usize) -> String {
|
||||
pub(crate) fn truncate_for_preview(output: &str, max_chars: usize) -> String {
|
||||
let collapsed: String = output
|
||||
.chars()
|
||||
.take(max_chars + 50)
|
||||
@@ -37,8 +38,14 @@ fn truncate_for_preview(output: &str, max_chars: usize) -> String {
|
||||
.split_whitespace()
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
if collapsed.len() > max_chars {
|
||||
format!("{}...", &collapsed[..max_chars])
|
||||
// char_indices gives us byte offsets at char boundaries, so the slice is always valid UTF-8.
|
||||
if collapsed.chars().count() > max_chars {
|
||||
let byte_offset = collapsed
|
||||
.char_indices()
|
||||
.nth(max_chars)
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(collapsed.len());
|
||||
format!("{}...", &collapsed[..byte_offset])
|
||||
} else {
|
||||
collapsed
|
||||
}
|
||||
@@ -59,12 +66,16 @@ enum AgenticLoopResult {
|
||||
///
|
||||
/// Bundles the shared components to reduce argument count.
|
||||
pub struct AgentDeps {
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
/// Cheap/fast LLM for lightweight tasks (heartbeat, routing, evaluation).
|
||||
/// Falls back to the main `llm` if None.
|
||||
pub cheap_llm: Option<Arc<dyn LlmProvider>>,
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub workspace: Option<Arc<Workspace>>,
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
pub hooks: Arc<HookRegistry>,
|
||||
}
|
||||
|
||||
/// The main agent that coordinates all components.
|
||||
@@ -107,6 +118,7 @@ impl Agent {
|
||||
deps.safety.clone(),
|
||||
deps.tools.clone(),
|
||||
deps.store.clone(),
|
||||
deps.hooks.clone(),
|
||||
));
|
||||
|
||||
Self {
|
||||
@@ -124,7 +136,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Convenience accessors
|
||||
fn store(&self) -> Option<&Arc<Store>> {
|
||||
fn store(&self) -> Option<&Arc<dyn Database>> {
|
||||
self.deps.store.as_ref()
|
||||
}
|
||||
|
||||
@@ -132,6 +144,11 @@ impl Agent {
|
||||
&self.deps.llm
|
||||
}
|
||||
|
||||
/// Get the cheap/fast LLM provider, falling back to the main one.
|
||||
fn cheap_llm(&self) -> &Arc<dyn LlmProvider> {
|
||||
self.deps.cheap_llm.as_ref().unwrap_or(&self.deps.llm)
|
||||
}
|
||||
|
||||
fn safety(&self) -> &Arc<SafetyLayer> {
|
||||
&self.deps.safety
|
||||
}
|
||||
@@ -144,6 +161,10 @@ impl Agent {
|
||||
self.deps.workspace.as_ref()
|
||||
}
|
||||
|
||||
fn hooks(&self) -> &Arc<HookRegistry> {
|
||||
&self.deps.hooks
|
||||
}
|
||||
|
||||
/// Run the agent main loop.
|
||||
pub async fn run(self) -> Result<(), Error> {
|
||||
// Start channels
|
||||
@@ -295,7 +316,7 @@ impl Agent {
|
||||
Some(spawn_heartbeat(
|
||||
config,
|
||||
workspace.clone(),
|
||||
self.llm().clone(),
|
||||
self.cheap_llm().clone(),
|
||||
Some(notify_tx),
|
||||
))
|
||||
} else {
|
||||
@@ -411,10 +432,32 @@ impl Agent {
|
||||
|
||||
match self.handle_message(&message).await {
|
||||
Ok(Some(response)) if !response.is_empty() => {
|
||||
let _ = self
|
||||
.channels
|
||||
.respond(&message, OutgoingResponse::text(response))
|
||||
.await;
|
||||
// Hook: BeforeOutbound — allow hooks to modify or suppress outbound
|
||||
let event = crate::hooks::HookEvent::Outbound {
|
||||
user_id: message.user_id.clone(),
|
||||
channel: message.channel.clone(),
|
||||
content: response.clone(),
|
||||
thread_id: message.thread_id.clone(),
|
||||
};
|
||||
match self.hooks().run(&event).await {
|
||||
Err(err) => {
|
||||
tracing::warn!("BeforeOutbound hook blocked response: {}", err);
|
||||
}
|
||||
Ok(crate::hooks::HookOutcome::Continue {
|
||||
modified: Some(new_content),
|
||||
}) => {
|
||||
let _ = self
|
||||
.channels
|
||||
.respond(&message, OutgoingResponse::text(new_content))
|
||||
.await;
|
||||
}
|
||||
_ => {
|
||||
let _ = self
|
||||
.channels
|
||||
.respond(&message, OutgoingResponse::text(response))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Some(_)) => {
|
||||
// Empty response, nothing to send (e.g. approval handled via send_status)
|
||||
@@ -460,7 +503,33 @@ impl Agent {
|
||||
|
||||
async fn handle_message(&self, message: &IncomingMessage) -> Result<Option<String>, Error> {
|
||||
// Parse submission type first
|
||||
let submission = SubmissionParser::parse(&message.content);
|
||||
let mut submission = SubmissionParser::parse(&message.content);
|
||||
|
||||
// Hook: BeforeInbound — allow hooks to modify or reject user input
|
||||
if let Submission::UserInput { ref content } = submission {
|
||||
let event = crate::hooks::HookEvent::Inbound {
|
||||
user_id: message.user_id.clone(),
|
||||
channel: message.channel.clone(),
|
||||
content: content.clone(),
|
||||
thread_id: message.thread_id.clone(),
|
||||
};
|
||||
match self.hooks().run(&event).await {
|
||||
Err(crate::hooks::HookError::Rejected { reason }) => {
|
||||
return Ok(Some(format!("[Message rejected: {}]", reason)));
|
||||
}
|
||||
Err(err) => {
|
||||
return Ok(Some(format!("[Message blocked by hook policy: {}]", err)));
|
||||
}
|
||||
Ok(crate::hooks::HookOutcome::Continue {
|
||||
modified: Some(new_content),
|
||||
}) => {
|
||||
submission = Submission::UserInput {
|
||||
content: new_content,
|
||||
};
|
||||
}
|
||||
_ => {} // Continue, fail-open errors already logged in registry
|
||||
}
|
||||
}
|
||||
|
||||
// Hydrate thread from DB if it's a historical thread not in memory
|
||||
if let Some(ref external_thread_id) = message.thread_id {
|
||||
@@ -654,19 +723,17 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Restore response chain from conversation metadata
|
||||
if let Some(store) = self.store() {
|
||||
if let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await {
|
||||
if let Some(rid) = metadata
|
||||
.get("last_response_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
{
|
||||
thread.last_response_id = Some(rid.clone());
|
||||
self.llm()
|
||||
.seed_response_chain(&thread_uuid.to_string(), rid);
|
||||
tracing::debug!("Restored response chain for thread {}", thread_uuid);
|
||||
}
|
||||
}
|
||||
if let Some(store) = self.store()
|
||||
&& let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await
|
||||
&& let Some(rid) = metadata
|
||||
.get("last_response_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
{
|
||||
thread.last_response_id = Some(rid.clone());
|
||||
self.llm()
|
||||
.seed_response_chain(&thread_uuid.to_string(), rid);
|
||||
tracing::debug!("Restored response chain for thread {}", thread_uuid);
|
||||
}
|
||||
|
||||
// Insert into session and register with session manager
|
||||
@@ -871,6 +938,27 @@ impl Agent {
|
||||
// Complete, fail, or request approval
|
||||
match result {
|
||||
Ok(AgenticLoopResult::Response(response)) => {
|
||||
// Hook: TransformResponse — allow hooks to modify or reject the final response
|
||||
let response = {
|
||||
let event = crate::hooks::HookEvent::ResponseTransform {
|
||||
user_id: message.user_id.clone(),
|
||||
thread_id: thread_id.to_string(),
|
||||
response: response.clone(),
|
||||
};
|
||||
match self.hooks().run(&event).await {
|
||||
Err(crate::hooks::HookError::Rejected { reason }) => {
|
||||
format!("[Response filtered: {}]", reason)
|
||||
}
|
||||
Err(err) => {
|
||||
format!("[Response blocked by hook policy: {}]", err)
|
||||
}
|
||||
Ok(crate::hooks::HookOutcome::Continue {
|
||||
modified: Some(new_response),
|
||||
}) => new_response,
|
||||
_ => response, // fail-open: use original
|
||||
}
|
||||
};
|
||||
|
||||
thread.complete_turn(&response);
|
||||
self.persist_response_chain(thread);
|
||||
let _ = self
|
||||
@@ -954,13 +1042,12 @@ impl Agent {
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(ref resp) = response {
|
||||
if let Err(e) = store
|
||||
if let Some(ref resp) = response
|
||||
&& let Err(e) = store
|
||||
.add_conversation_message(thread_id, "assistant", resp)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist assistant message: {}", e);
|
||||
}
|
||||
{
|
||||
tracing::warn!("Failed to persist assistant message: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -1058,14 +1145,14 @@ impl Agent {
|
||||
// Check if interrupted
|
||||
{
|
||||
let sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get(&thread_id) {
|
||||
if thread.state == ThreadState::Interrupted {
|
||||
return Err(crate::error::JobError::ContextError {
|
||||
id: thread_id,
|
||||
reason: "Interrupted".to_string(),
|
||||
}
|
||||
.into());
|
||||
if let Some(thread) = sess.threads.get(&thread_id)
|
||||
&& thread.state == ThreadState::Interrupted
|
||||
{
|
||||
return Err(crate::error::JobError::ContextError {
|
||||
id: thread_id,
|
||||
reason: "Interrupted".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1082,9 +1169,16 @@ impl Agent {
|
||||
m
|
||||
});
|
||||
|
||||
let result = reasoning.respond_with_tools(&context).await?;
|
||||
let output = reasoning.respond_with_tools(&context).await?;
|
||||
|
||||
match result {
|
||||
// Track token usage for budget enforcement
|
||||
tracing::debug!(
|
||||
"LLM call used {} input + {} output tokens",
|
||||
output.usage.input_tokens,
|
||||
output.usage.output_tokens
|
||||
);
|
||||
|
||||
match output.result {
|
||||
RespondResult::Text(text) => {
|
||||
// If no tools have been executed yet, prompt the LLM to use tools
|
||||
// This handles the case where the model explains what it will do
|
||||
@@ -1133,39 +1227,90 @@ impl Agent {
|
||||
// Record tool calls in the thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
if let Some(turn) = thread.last_turn_mut() {
|
||||
for tc in &tool_calls {
|
||||
turn.record_tool_call(&tc.name, tc.arguments.clone());
|
||||
}
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
for tc in &tool_calls {
|
||||
turn.record_tool_call(&tc.name, tc.arguments.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute each tool (with approval checking)
|
||||
for tc in tool_calls {
|
||||
// Execute each tool (with approval checking and hook interception)
|
||||
for mut tc in tool_calls {
|
||||
// Check if tool requires approval
|
||||
if let Some(tool) = self.tools().get(&tc.name).await {
|
||||
if tool.requires_approval() {
|
||||
// Check if auto-approved for this session
|
||||
let is_auto_approved = {
|
||||
let sess = session.lock().await;
|
||||
sess.is_tool_auto_approved(&tc.name)
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
// Check if auto-approved for this session
|
||||
let mut is_auto_approved = {
|
||||
let sess = session.lock().await;
|
||||
sess.is_tool_auto_approved(&tc.name)
|
||||
};
|
||||
|
||||
// Let the tool inspect the specific parameters and
|
||||
// override auto-approval (e.g. destructive shell commands).
|
||||
if is_auto_approved && tool.requires_approval_for(&tc.arguments) {
|
||||
tracing::info!(
|
||||
tool = %tc.name,
|
||||
"Tool requires explicit approval for these parameters despite auto-approve"
|
||||
);
|
||||
is_auto_approved = false;
|
||||
}
|
||||
|
||||
if !is_auto_approved {
|
||||
// Need approval - store pending request and return
|
||||
let pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
};
|
||||
|
||||
if !is_auto_approved {
|
||||
// Need approval - store pending request and return
|
||||
let pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
};
|
||||
return Ok(AgenticLoopResult::NeedApproval { pending });
|
||||
}
|
||||
}
|
||||
|
||||
return Ok(AgenticLoopResult::NeedApproval { pending });
|
||||
// Hook: BeforeToolCall — allow hooks to modify or reject tool calls
|
||||
{
|
||||
let event = crate::hooks::HookEvent::ToolCall {
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
user_id: message.user_id.clone(),
|
||||
context: "chat".to_string(),
|
||||
};
|
||||
match self.hooks().run(&event).await {
|
||||
Err(crate::hooks::HookError::Rejected { reason }) => {
|
||||
context_messages.push(ChatMessage::tool_result(
|
||||
&tc.id,
|
||||
&tc.name,
|
||||
format!("Tool call rejected by hook: {}", reason),
|
||||
));
|
||||
continue;
|
||||
}
|
||||
Err(err) => {
|
||||
context_messages.push(ChatMessage::tool_result(
|
||||
&tc.id,
|
||||
&tc.name,
|
||||
format!("Tool call blocked by hook policy: {}", err),
|
||||
));
|
||||
continue;
|
||||
}
|
||||
Ok(crate::hooks::HookOutcome::Continue {
|
||||
modified: Some(new_params),
|
||||
}) => match serde_json::from_str(&new_params) {
|
||||
Ok(parsed) => tc.arguments = parsed,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
tool = %tc.name,
|
||||
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
},
|
||||
_ => {} // Continue, fail-open errors already logged
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1196,34 +1341,34 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Ok(ref output) = tool_result {
|
||||
if !output.is_empty() {
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolResult {
|
||||
name: tc.name.clone(),
|
||||
preview: truncate_for_preview(output, 200),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Ok(ref output) = tool_result
|
||||
&& !output.is_empty()
|
||||
{
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolResult {
|
||||
name: tc.name.clone(),
|
||||
preview: output.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Record result in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
if let Some(turn) = thread.last_turn_mut() {
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1434,6 +1579,9 @@ impl Agent {
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
// Lock session first, then undo manager -- consistent with process_user_input
|
||||
// to avoid potential deadlocks.
|
||||
let mut sess = session.lock().await;
|
||||
let undo_mgr = self.session_manager.get_undo_manager(thread_id).await;
|
||||
let mut mgr = undo_mgr.lock().await;
|
||||
|
||||
@@ -1441,7 +1589,6 @@ impl Agent {
|
||||
return Ok(SubmissionResult::ok_with_message("Nothing to undo."));
|
||||
}
|
||||
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess
|
||||
.threads
|
||||
.get_mut(&thread_id)
|
||||
@@ -1452,12 +1599,10 @@ impl Agent {
|
||||
let current_turn = thread.turn_number();
|
||||
|
||||
if let Some(checkpoint) = mgr.undo(current_turn, current_messages) {
|
||||
// Extract values before consuming the reference
|
||||
let turn_number = checkpoint.turn_number;
|
||||
let messages = checkpoint.messages.clone();
|
||||
let undo_count = mgr.undo_count();
|
||||
// Restore thread from checkpoint
|
||||
thread.restore_from_messages(messages);
|
||||
thread.restore_from_messages(checkpoint.messages);
|
||||
Ok(SubmissionResult::ok_with_message(format!(
|
||||
"Undone to turn {}. {} undo(s) remaining.",
|
||||
turn_number, undo_count
|
||||
@@ -1472,6 +1617,9 @@ impl Agent {
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
// Lock session first, then undo manager -- consistent with process_user_input
|
||||
// to avoid potential deadlocks.
|
||||
let mut sess = session.lock().await;
|
||||
let undo_mgr = self.session_manager.get_undo_manager(thread_id).await;
|
||||
let mut mgr = undo_mgr.lock().await;
|
||||
|
||||
@@ -1479,12 +1627,15 @@ impl Agent {
|
||||
return Ok(SubmissionResult::ok_with_message("Nothing to redo."));
|
||||
}
|
||||
|
||||
if let Some(checkpoint) = mgr.redo() {
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess
|
||||
.threads
|
||||
.get_mut(&thread_id)
|
||||
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
|
||||
let thread = sess
|
||||
.threads
|
||||
.get_mut(&thread_id)
|
||||
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
|
||||
|
||||
let current_messages = thread.messages();
|
||||
let current_turn = thread.turn_number();
|
||||
|
||||
if let Some(checkpoint) = mgr.redo(current_turn, current_messages) {
|
||||
thread.restore_from_messages(checkpoint.messages);
|
||||
Ok(SubmissionResult::ok_with_message(format!(
|
||||
"Redone to turn {}.",
|
||||
@@ -1606,17 +1757,17 @@ impl Agent {
|
||||
};
|
||||
|
||||
// Verify request ID if provided
|
||||
if let Some(req_id) = request_id {
|
||||
if 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 let Some(req_id) = request_id
|
||||
&& 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 {
|
||||
@@ -1670,20 +1821,20 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Ok(ref output) = tool_result {
|
||||
if !output.is_empty() {
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolResult {
|
||||
name: pending.tool_name.clone(),
|
||||
preview: truncate_for_preview(output, 200),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Ok(ref output) = tool_result
|
||||
&& !output.is_empty()
|
||||
{
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolResult {
|
||||
name: pending.tool_name.clone(),
|
||||
preview: output.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Build context including the tool result
|
||||
@@ -1692,15 +1843,15 @@ impl Agent {
|
||||
// Record result in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
if let Some(turn) = thread.last_turn_mut() {
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2060,15 +2211,15 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Persist new job to database (fire-and-forget)
|
||||
if let Some(store) = self.store() {
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
|
||||
let store = store.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store.save_job(&ctx).await {
|
||||
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
|
||||
}
|
||||
});
|
||||
}
|
||||
if let Some(store) = self.store()
|
||||
&& let Ok(ctx) = self.context_manager.get_context(job_id).await
|
||||
{
|
||||
let store = store.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store.save_job(&ctx).await {
|
||||
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Schedule for execution
|
||||
@@ -2148,10 +2299,10 @@ impl Agent {
|
||||
|
||||
let mut output = String::from("Jobs:\n");
|
||||
for job_id in jobs {
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
|
||||
if ctx.user_id == user_id {
|
||||
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
|
||||
}
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await
|
||||
&& ctx.user_id == user_id
|
||||
{
|
||||
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2602,4 +2753,70 @@ mod tests {
|
||||
|
||||
assert!(detect_auth_awaiting("tool_activate", &result).is_none());
|
||||
}
|
||||
|
||||
// --- truncate_for_preview tests ---
|
||||
|
||||
use super::truncate_for_preview;
|
||||
|
||||
#[test]
|
||||
fn test_truncate_short_input() {
|
||||
assert_eq!(truncate_for_preview("hello", 10), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_empty_input() {
|
||||
assert_eq!(truncate_for_preview("", 10), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_exact_length() {
|
||||
assert_eq!(truncate_for_preview("hello", 5), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_over_limit() {
|
||||
let result = truncate_for_preview("hello world, this is long", 10);
|
||||
assert!(result.ends_with("..."));
|
||||
// "hello worl" = 10 chars + "..."
|
||||
assert_eq!(result, "hello worl...");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_collapses_newlines() {
|
||||
let result = truncate_for_preview("line1\nline2\nline3", 100);
|
||||
assert!(!result.contains('\n'));
|
||||
assert_eq!(result, "line1 line2 line3");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_collapses_whitespace() {
|
||||
let result = truncate_for_preview("hello world", 100);
|
||||
assert_eq!(result, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_multibyte_utf8() {
|
||||
// Each emoji is 4 bytes. Truncating at char boundary must not panic.
|
||||
let input = "😀😁😂🤣😃😄😅😆😉😊";
|
||||
let result = truncate_for_preview(input, 5);
|
||||
assert!(result.ends_with("..."));
|
||||
// First 5 chars = 5 emoji
|
||||
assert_eq!(result, "😀😁😂🤣😃...");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_cjk_characters() {
|
||||
// CJK chars are 3 bytes each in UTF-8.
|
||||
let input = "你好世界测试数据很长的字符串";
|
||||
let result = truncate_for_preview(input, 4);
|
||||
assert_eq!(result, "你好世界...");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_mixed_multibyte_and_ascii() {
|
||||
let input = "hello 世界 foo";
|
||||
let result = truncate_for_preview(input, 8);
|
||||
// 'h','e','l','l','o',' ','世','界' = 8 chars
|
||||
assert_eq!(result, "hello 世界...");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ pub mod task;
|
||||
pub mod undo;
|
||||
pub mod worker;
|
||||
|
||||
pub(crate) use agent_loop::truncate_for_preview;
|
||||
pub use agent_loop::{Agent, AgentDeps};
|
||||
pub use compaction::{CompactionResult, ContextCompactor};
|
||||
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
|
||||
|
||||
+17
-22
@@ -11,6 +11,7 @@
|
||||
//! Full-job routines are delegated to the existing `Scheduler`.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::Utc;
|
||||
@@ -23,20 +24,20 @@ use crate::agent::routine::{
|
||||
};
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::history::Store;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// The routine execution engine.
|
||||
pub struct RoutineEngine {
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
/// Sender for notifications (routed to channel manager).
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
/// Currently running routine count (across all routines).
|
||||
running_count: Arc<RwLock<usize>>,
|
||||
running_count: Arc<AtomicUsize>,
|
||||
/// Compiled event regex cache: routine_id -> compiled regex.
|
||||
event_cache: Arc<RwLock<Vec<(Uuid, Routine, Regex)>>>,
|
||||
}
|
||||
@@ -44,7 +45,7 @@ pub struct RoutineEngine {
|
||||
impl RoutineEngine {
|
||||
pub fn new(
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
@@ -55,7 +56,7 @@ impl RoutineEngine {
|
||||
llm,
|
||||
workspace,
|
||||
notify_tx,
|
||||
running_count: Arc::new(RwLock::new(0)),
|
||||
running_count: Arc::new(AtomicUsize::new(0)),
|
||||
event_cache: Arc::new(RwLock::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
@@ -102,10 +103,9 @@ impl RoutineEngine {
|
||||
if let Trigger::Event {
|
||||
channel: Some(ch), ..
|
||||
} = &routine.trigger
|
||||
&& ch != &message.channel
|
||||
{
|
||||
if ch != &message.channel {
|
||||
continue;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Regex match
|
||||
@@ -126,7 +126,7 @@ impl RoutineEngine {
|
||||
}
|
||||
|
||||
// Global capacity check
|
||||
if *self.running_count.read().await >= self.config.max_concurrent_routines {
|
||||
if self.running_count.load(Ordering::Relaxed) >= self.config.max_concurrent_routines {
|
||||
tracing::warn!(routine = %routine.name, "Skipped: global max concurrent reached");
|
||||
continue;
|
||||
}
|
||||
@@ -150,7 +150,7 @@ impl RoutineEngine {
|
||||
};
|
||||
|
||||
for routine in routines {
|
||||
if *self.running_count.read().await >= self.config.max_concurrent_routines {
|
||||
if self.running_count.load(Ordering::Relaxed) >= self.config.max_concurrent_routines {
|
||||
tracing::warn!("Global max concurrent routines reached, skipping remaining");
|
||||
break;
|
||||
}
|
||||
@@ -293,21 +293,18 @@ impl RoutineEngine {
|
||||
|
||||
/// Shared context passed to the execution function.
|
||||
struct EngineContext {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
running_count: Arc<RwLock<usize>>,
|
||||
running_count: Arc<AtomicUsize>,
|
||||
max_lightweight_tokens: u32,
|
||||
}
|
||||
|
||||
/// Execute a routine run. Handles both lightweight and full_job modes.
|
||||
async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) {
|
||||
// Increment running count
|
||||
{
|
||||
let mut count = ctx.running_count.write().await;
|
||||
*count += 1;
|
||||
}
|
||||
// Increment running count (atomic: survives panics in the execution below)
|
||||
ctx.running_count.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
let result = match &routine.action {
|
||||
RoutineAction::Lightweight {
|
||||
@@ -327,10 +324,7 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun)
|
||||
};
|
||||
|
||||
// Decrement running count
|
||||
{
|
||||
let mut count = ctx.running_count.write().await;
|
||||
*count = count.saturating_sub(1);
|
||||
}
|
||||
ctx.running_count.fetch_sub(1, Ordering::Relaxed);
|
||||
|
||||
// Process result
|
||||
let (status, summary, tokens) = match result {
|
||||
@@ -568,7 +562,8 @@ fn truncate(s: &str, max: usize) -> String {
|
||||
if s.len() <= max {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}...", &s[..max])
|
||||
let end = crate::util::floor_char_boundary(s, max);
|
||||
format!("{}...", &s[..end])
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+61
-56
@@ -12,8 +12,9 @@ use crate::agent::task::{Task, TaskContext, TaskOutput};
|
||||
use crate::agent::worker::{Worker, WorkerDeps};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::context::{ContextManager, JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::{Error, JobError};
|
||||
use crate::history::Store;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
@@ -48,7 +49,8 @@ pub struct Scheduler {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
hooks: Arc<HookRegistry>,
|
||||
/// Running jobs (main LLM-driven jobs).
|
||||
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
|
||||
/// Running sub-tasks (tool executions, background tasks).
|
||||
@@ -63,7 +65,8 @@ impl Scheduler {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
hooks: Arc<HookRegistry>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
@@ -72,6 +75,7 @@ impl Scheduler {
|
||||
safety,
|
||||
tools,
|
||||
store,
|
||||
hooks,
|
||||
jobs: Arc::new(RwLock::new(HashMap::new())),
|
||||
subtasks: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
@@ -79,63 +83,64 @@ impl Scheduler {
|
||||
|
||||
/// Schedule a job for execution.
|
||||
pub async fn schedule(&self, job_id: Uuid) -> Result<(), JobError> {
|
||||
// Check if already scheduled
|
||||
if self.jobs.read().await.contains_key(&job_id) {
|
||||
return Ok(());
|
||||
}
|
||||
// Hold write lock for the entire check-insert sequence to prevent
|
||||
// TOCTOU races where two concurrent calls both pass the checks.
|
||||
{
|
||||
let mut jobs = self.jobs.write().await;
|
||||
|
||||
// Check capacity
|
||||
let current_count = self.jobs.read().await.len();
|
||||
if current_count >= self.config.max_parallel_jobs {
|
||||
return Err(JobError::MaxJobsExceeded {
|
||||
max: self.config.max_parallel_jobs,
|
||||
});
|
||||
}
|
||||
|
||||
// Transition job to in_progress
|
||||
self.context_manager
|
||||
.update_context(job_id, |ctx| {
|
||||
ctx.transition_to(
|
||||
JobState::InProgress,
|
||||
Some("Scheduled for execution".to_string()),
|
||||
)
|
||||
})
|
||||
.await?
|
||||
.map_err(|s| JobError::ContextError {
|
||||
id: job_id,
|
||||
reason: s,
|
||||
})?;
|
||||
|
||||
// Create worker channel
|
||||
let (tx, rx) = mpsc::channel(16);
|
||||
|
||||
// Create worker with shared dependencies
|
||||
let deps = WorkerDeps {
|
||||
context_manager: self.context_manager.clone(),
|
||||
llm: self.llm.clone(),
|
||||
safety: self.safety.clone(),
|
||||
tools: self.tools.clone(),
|
||||
store: self.store.clone(),
|
||||
timeout: self.config.job_timeout,
|
||||
use_planning: self.config.use_planning,
|
||||
};
|
||||
let worker = Worker::new(job_id, deps);
|
||||
|
||||
// Spawn worker task
|
||||
let handle = tokio::spawn(async move {
|
||||
if let Err(e) = worker.run(rx).await {
|
||||
tracing::error!("Worker for job {} failed: {}", job_id, e);
|
||||
if jobs.contains_key(&job_id) {
|
||||
return Ok(());
|
||||
}
|
||||
});
|
||||
|
||||
// Start the worker
|
||||
let _ = tx.send(WorkerMessage::Start).await;
|
||||
if jobs.len() >= self.config.max_parallel_jobs {
|
||||
return Err(JobError::MaxJobsExceeded {
|
||||
max: self.config.max_parallel_jobs,
|
||||
});
|
||||
}
|
||||
|
||||
// Store the scheduled job
|
||||
self.jobs
|
||||
.write()
|
||||
.await
|
||||
.insert(job_id, ScheduledJob { handle, tx });
|
||||
// Transition job to in_progress
|
||||
self.context_manager
|
||||
.update_context(job_id, |ctx| {
|
||||
ctx.transition_to(
|
||||
JobState::InProgress,
|
||||
Some("Scheduled for execution".to_string()),
|
||||
)
|
||||
})
|
||||
.await?
|
||||
.map_err(|s| JobError::ContextError {
|
||||
id: job_id,
|
||||
reason: s,
|
||||
})?;
|
||||
|
||||
// Create worker channel
|
||||
let (tx, rx) = mpsc::channel(16);
|
||||
|
||||
// Create worker with shared dependencies
|
||||
let deps = WorkerDeps {
|
||||
context_manager: self.context_manager.clone(),
|
||||
llm: self.llm.clone(),
|
||||
safety: self.safety.clone(),
|
||||
tools: self.tools.clone(),
|
||||
store: self.store.clone(),
|
||||
hooks: self.hooks.clone(),
|
||||
timeout: self.config.job_timeout,
|
||||
use_planning: self.config.use_planning,
|
||||
};
|
||||
let worker = Worker::new(job_id, deps);
|
||||
|
||||
// Spawn worker task
|
||||
let handle = tokio::spawn(async move {
|
||||
if let Err(e) = worker.run(rx).await {
|
||||
tracing::error!("Worker for job {} failed: {}", job_id, e);
|
||||
}
|
||||
});
|
||||
|
||||
// Start the worker
|
||||
let _ = tx.send(WorkerMessage::Start).await;
|
||||
|
||||
// Insert while still holding the write lock
|
||||
jobs.insert(job_id, ScheduledJob { handle, tx });
|
||||
}
|
||||
|
||||
// Cleanup task for this job to avoid capacity leaks
|
||||
let jobs = Arc::clone(&self.jobs);
|
||||
|
||||
+21
-21
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::RepairError;
|
||||
use crate::history::Store;
|
||||
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
||||
|
||||
/// A job that has been detected as stuck.
|
||||
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
|
||||
#[allow(dead_code)] // Will be used for time-based stuck detection
|
||||
stuck_threshold: Duration,
|
||||
max_repair_attempts: u32,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
builder: Option<Arc<dyn SoftwareBuilder>>,
|
||||
#[allow(dead_code)] // Will be used for tool hot-reload after repair
|
||||
tools: Option<Arc<ToolRegistry>>,
|
||||
@@ -94,7 +94,7 @@ impl DefaultSelfRepair {
|
||||
|
||||
/// Add a Store for tool failure tracking.
|
||||
#[allow(dead_code)] // Public API for configuring repair with persistence
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
@@ -119,25 +119,25 @@ impl SelfRepair for DefaultSelfRepair {
|
||||
let mut stuck_jobs = Vec::new();
|
||||
|
||||
for job_id in stuck_ids {
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
|
||||
if ctx.state == JobState::Stuck {
|
||||
let stuck_duration = ctx
|
||||
.started_at
|
||||
.map(|start| {
|
||||
let now = Utc::now();
|
||||
let duration = now.signed_duration_since(start);
|
||||
Duration::from_secs(duration.num_seconds().max(0) as u64)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await
|
||||
&& ctx.state == JobState::Stuck
|
||||
{
|
||||
let stuck_duration = ctx
|
||||
.started_at
|
||||
.map(|start| {
|
||||
let now = Utc::now();
|
||||
let duration = now.signed_duration_since(start);
|
||||
Duration::from_secs(duration.num_seconds().max(0) as u64)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
stuck_jobs.push(StuckJob {
|
||||
job_id,
|
||||
last_activity: ctx.started_at.unwrap_or(ctx.created_at),
|
||||
stuck_duration,
|
||||
last_error: None,
|
||||
repair_attempts: ctx.repair_attempts,
|
||||
});
|
||||
}
|
||||
stuck_jobs.push(StuckJob {
|
||||
job_id,
|
||||
last_activity: ctx.started_at.unwrap_or(ctx.created_at),
|
||||
stuck_duration,
|
||||
last_error: None,
|
||||
repair_attempts: ctx.repair_attempts,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -346,11 +346,11 @@ impl Thread {
|
||||
let mut turn = Turn::new(turn_number, &msg.content);
|
||||
|
||||
// Check if next is assistant response
|
||||
if let Some(next) = iter.peek() {
|
||||
if next.role == crate::llm::Role::Assistant {
|
||||
let response = iter.next().expect("peeked");
|
||||
turn.complete(&response.content);
|
||||
}
|
||||
if let Some(next) = iter.peek()
|
||||
&& next.role == crate::llm::Role::Assistant
|
||||
{
|
||||
let response = iter.next().expect("peeked");
|
||||
turn.complete(&response.content);
|
||||
}
|
||||
|
||||
self.turns.push(turn);
|
||||
|
||||
@@ -11,6 +11,7 @@ use uuid::Uuid;
|
||||
|
||||
use crate::agent::session::Session;
|
||||
use crate::agent::undo::UndoManager;
|
||||
use crate::hooks::HookRegistry;
|
||||
|
||||
/// Key for mapping external thread IDs to internal ones.
|
||||
#[derive(Clone, Hash, Eq, PartialEq)]
|
||||
@@ -25,6 +26,7 @@ pub struct SessionManager {
|
||||
sessions: RwLock<HashMap<String, Arc<Mutex<Session>>>>,
|
||||
thread_map: RwLock<HashMap<ThreadKey, Uuid>>,
|
||||
undo_managers: RwLock<HashMap<Uuid, Arc<Mutex<UndoManager>>>>,
|
||||
hooks: Option<Arc<HookRegistry>>,
|
||||
}
|
||||
|
||||
impl SessionManager {
|
||||
@@ -34,9 +36,16 @@ impl SessionManager {
|
||||
sessions: RwLock::new(HashMap::new()),
|
||||
thread_map: RwLock::new(HashMap::new()),
|
||||
undo_managers: RwLock::new(HashMap::new()),
|
||||
hooks: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a hook registry for session lifecycle events.
|
||||
pub fn with_hooks(mut self, hooks: Arc<HookRegistry>) -> Self {
|
||||
self.hooks = Some(hooks);
|
||||
self
|
||||
}
|
||||
|
||||
/// Get or create a session for a user.
|
||||
pub async fn get_or_create_session(&self, user_id: &str) -> Arc<Mutex<Session>> {
|
||||
// Fast path: check if session exists
|
||||
@@ -54,8 +63,28 @@ impl SessionManager {
|
||||
return Arc::clone(session);
|
||||
}
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new(user_id)));
|
||||
let new_session = Session::new(user_id);
|
||||
let session_id = new_session.id.to_string();
|
||||
let session = Arc::new(Mutex::new(new_session));
|
||||
sessions.insert(user_id.to_string(), Arc::clone(&session));
|
||||
|
||||
// Fire OnSessionStart hook (fire-and-forget)
|
||||
if let Some(ref hooks) = self.hooks {
|
||||
let hooks = hooks.clone();
|
||||
let uid = user_id.to_string();
|
||||
let sid = session_id;
|
||||
tokio::spawn(async move {
|
||||
use crate::hooks::HookEvent;
|
||||
let event = HookEvent::SessionStart {
|
||||
user_id: uid,
|
||||
session_id: sid,
|
||||
};
|
||||
if let Err(e) = hooks.run(&event).await {
|
||||
tracing::warn!("OnSessionStart hook error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
session
|
||||
}
|
||||
|
||||
@@ -173,8 +202,8 @@ impl SessionManager {
|
||||
pub async fn prune_stale_sessions(&self, max_idle: std::time::Duration) -> usize {
|
||||
let cutoff = chrono::Utc::now() - chrono::TimeDelta::seconds(max_idle.as_secs() as i64);
|
||||
|
||||
// Find stale session user_ids
|
||||
let stale_users: Vec<String> = {
|
||||
// Find stale sessions (user_id + session_id)
|
||||
let stale_sessions: Vec<(String, String)> = {
|
||||
let sessions = self.sessions.read().await;
|
||||
sessions
|
||||
.iter()
|
||||
@@ -182,7 +211,7 @@ impl SessionManager {
|
||||
// Try to lock; skip if contended (someone is actively using it)
|
||||
let sess = session.try_lock().ok()?;
|
||||
if sess.last_active_at < cutoff {
|
||||
Some(user_id.clone())
|
||||
Some((user_id.clone(), sess.id.to_string()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
@@ -190,6 +219,11 @@ impl SessionManager {
|
||||
.collect()
|
||||
};
|
||||
|
||||
let stale_users: Vec<String> = stale_sessions
|
||||
.iter()
|
||||
.map(|(user_id, _)| user_id.clone())
|
||||
.collect();
|
||||
|
||||
if stale_users.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
@@ -199,14 +233,33 @@ impl SessionManager {
|
||||
{
|
||||
let sessions = self.sessions.read().await;
|
||||
for user_id in &stale_users {
|
||||
if let Some(session) = sessions.get(user_id) {
|
||||
if let Ok(sess) = session.try_lock() {
|
||||
stale_thread_ids.extend(sess.threads.keys());
|
||||
}
|
||||
if let Some(session) = sessions.get(user_id)
|
||||
&& let Ok(sess) = session.try_lock()
|
||||
{
|
||||
stale_thread_ids.extend(sess.threads.keys());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fire OnSessionEnd hooks for stale sessions (fire-and-forget)
|
||||
if let Some(ref hooks) = self.hooks {
|
||||
for (user_id, session_id) in &stale_sessions {
|
||||
let hooks = hooks.clone();
|
||||
let uid = user_id.clone();
|
||||
let sid = session_id.clone();
|
||||
tokio::spawn(async move {
|
||||
use crate::hooks::HookEvent;
|
||||
let event = HookEvent::SessionEnd {
|
||||
user_id: uid,
|
||||
session_id: sid,
|
||||
};
|
||||
if let Err(e) = hooks.run(&event).await {
|
||||
tracing::warn!("OnSessionEnd hook error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Remove sessions
|
||||
let count = {
|
||||
let mut sessions = self.sessions.write().await;
|
||||
|
||||
+13
-14
@@ -93,27 +93,26 @@ impl SubmissionParser {
|
||||
// /thread <uuid> - switch thread
|
||||
if let Some(rest) = lower.strip_prefix("/thread ") {
|
||||
let rest = rest.trim();
|
||||
if rest != "new" {
|
||||
if let Ok(id) = Uuid::parse_str(rest) {
|
||||
return Submission::SwitchThread { thread_id: id };
|
||||
}
|
||||
if rest != "new"
|
||||
&& let Ok(id) = Uuid::parse_str(rest)
|
||||
{
|
||||
return Submission::SwitchThread { thread_id: id };
|
||||
}
|
||||
}
|
||||
|
||||
// /resume <uuid> - resume from checkpoint
|
||||
if let Some(rest) = lower.strip_prefix("/resume ") {
|
||||
if let Ok(id) = Uuid::parse_str(rest.trim()) {
|
||||
return Submission::Resume { checkpoint_id: id };
|
||||
}
|
||||
if let Some(rest) = lower.strip_prefix("/resume ")
|
||||
&& let Ok(id) = Uuid::parse_str(rest.trim())
|
||||
{
|
||||
return Submission::Resume { checkpoint_id: id };
|
||||
}
|
||||
|
||||
// Try structured JSON approval (from web gateway's /api/chat/approval endpoint)
|
||||
if trimmed.starts_with('{') {
|
||||
if let Ok(submission) = serde_json::from_str::<Submission>(trimmed) {
|
||||
if matches!(submission, Submission::ExecApproval { .. }) {
|
||||
return submission;
|
||||
}
|
||||
}
|
||||
if trimmed.starts_with('{')
|
||||
&& let Ok(submission) = serde_json::from_str::<Submission>(trimmed)
|
||||
&& matches!(submission, Submission::ExecApproval { .. })
|
||||
{
|
||||
return submission;
|
||||
}
|
||||
|
||||
// Approval responses (simple yes/no/always for pending approvals)
|
||||
|
||||
+136
-16
@@ -43,6 +43,10 @@ impl Checkpoint {
|
||||
}
|
||||
|
||||
/// Manager for undo/redo functionality.
|
||||
///
|
||||
/// Each undo/redo operation pops from one stack and pushes the current state
|
||||
/// onto the other, so `undo_count() + redo_count()` stays constant across
|
||||
/// undo/redo cycles (only `checkpoint()` and `clear()` change the total).
|
||||
pub struct UndoManager {
|
||||
/// Stack of past checkpoints (for undo).
|
||||
undo_stack: VecDeque<Checkpoint>,
|
||||
@@ -68,6 +72,14 @@ impl UndoManager {
|
||||
self
|
||||
}
|
||||
|
||||
/// Push a checkpoint onto the undo stack, trimming oldest entries if over limit.
|
||||
fn push_undo(&mut self, checkpoint: Checkpoint) {
|
||||
self.undo_stack.push_back(checkpoint);
|
||||
while self.undo_stack.len() > self.max_checkpoints {
|
||||
self.undo_stack.pop_front();
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a checkpoint at the current state.
|
||||
///
|
||||
/// This clears the redo stack since we're creating a new history branch.
|
||||
@@ -80,24 +92,23 @@ impl UndoManager {
|
||||
// Clear redo stack (new branch of history)
|
||||
self.redo_stack.clear();
|
||||
|
||||
// Create and push checkpoint
|
||||
let checkpoint = Checkpoint::new(turn_number, messages, description);
|
||||
self.undo_stack.push_back(checkpoint);
|
||||
|
||||
// Trim if over limit
|
||||
while self.undo_stack.len() > self.max_checkpoints {
|
||||
self.undo_stack.pop_front();
|
||||
}
|
||||
self.push_undo(checkpoint);
|
||||
}
|
||||
|
||||
/// Undo: pop the last checkpoint and return it.
|
||||
///
|
||||
/// The current state should be saved to redo stack before calling this.
|
||||
/// Saves the current state to the redo stack and pops the most recent
|
||||
/// checkpoint from the undo stack so that repeated undos walk backwards
|
||||
/// through history.
|
||||
///
|
||||
/// Takes ownership of `current_messages`; callers must clone first if
|
||||
/// they need to retain a copy.
|
||||
pub fn undo(
|
||||
&mut self,
|
||||
current_turn: usize,
|
||||
current_messages: Vec<ChatMessage>,
|
||||
) -> Option<&Checkpoint> {
|
||||
) -> Option<Checkpoint> {
|
||||
if self.undo_stack.is_empty() {
|
||||
return None;
|
||||
}
|
||||
@@ -110,9 +121,8 @@ impl UndoManager {
|
||||
);
|
||||
self.redo_stack.push(current);
|
||||
|
||||
// Return the most recent checkpoint without removing it
|
||||
// (we keep it so multiple undos can work)
|
||||
self.undo_stack.back()
|
||||
// Pop and return the most recent checkpoint
|
||||
self.undo_stack.pop_back()
|
||||
}
|
||||
|
||||
/// Pop the last checkpoint from the undo stack.
|
||||
@@ -121,7 +131,29 @@ impl UndoManager {
|
||||
}
|
||||
|
||||
/// Redo: restore a previously undone state.
|
||||
pub fn redo(&mut self) -> Option<Checkpoint> {
|
||||
///
|
||||
/// Saves the current state to the undo stack and pops the most recent
|
||||
/// checkpoint from the redo stack.
|
||||
///
|
||||
/// Takes ownership of `current_messages`; callers must clone first if
|
||||
/// they need to retain a copy.
|
||||
pub fn redo(
|
||||
&mut self,
|
||||
current_turn: usize,
|
||||
current_messages: Vec<ChatMessage>,
|
||||
) -> Option<Checkpoint> {
|
||||
if self.redo_stack.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Save current state to undo stack
|
||||
let current = Checkpoint::new(
|
||||
current_turn,
|
||||
current_messages,
|
||||
format!("Turn {}", current_turn),
|
||||
);
|
||||
self.push_undo(current);
|
||||
|
||||
self.redo_stack.pop()
|
||||
}
|
||||
|
||||
@@ -214,14 +246,16 @@ mod tests {
|
||||
assert!(manager.can_undo());
|
||||
assert!(!manager.can_redo());
|
||||
|
||||
// Undo
|
||||
// Undo - returns owned Checkpoint now
|
||||
let current = vec![ChatMessage::user("Hello"), ChatMessage::assistant("Hi")];
|
||||
let checkpoint = manager.undo(2, current);
|
||||
assert!(checkpoint.is_some());
|
||||
let checkpoint = checkpoint.unwrap();
|
||||
assert_eq!(checkpoint.turn_number, 1);
|
||||
assert!(manager.can_redo());
|
||||
|
||||
// Redo
|
||||
let restored = manager.redo();
|
||||
// Redo - now requires current state parameters
|
||||
let restored = manager.redo(checkpoint.turn_number, checkpoint.messages);
|
||||
assert!(restored.is_some());
|
||||
}
|
||||
|
||||
@@ -249,4 +283,90 @@ mod tests {
|
||||
assert!(restored.is_some());
|
||||
assert_eq!(manager.undo_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_repeated_undo_advances_through_stack() {
|
||||
let mut manager = UndoManager::new();
|
||||
|
||||
// Create 3 checkpoints at turns 0, 1, 2
|
||||
manager.checkpoint(0, vec![], "Turn 0");
|
||||
manager.checkpoint(1, vec![ChatMessage::user("msg1")], "Turn 1");
|
||||
manager.checkpoint(2, vec![ChatMessage::user("msg2")], "Turn 2");
|
||||
assert_eq!(manager.undo_count(), 3);
|
||||
|
||||
// First undo: should return turn 2 checkpoint, stack shrinks to 2
|
||||
let cp1 = manager
|
||||
.undo(3, vec![ChatMessage::user("msg3")])
|
||||
.expect("first undo should succeed");
|
||||
assert_eq!(cp1.turn_number, 2);
|
||||
assert_eq!(manager.undo_count(), 2);
|
||||
|
||||
// Second undo: should return turn 1 checkpoint (different!), stack shrinks to 1
|
||||
let cp2 = manager
|
||||
.undo(cp1.turn_number, cp1.messages)
|
||||
.expect("second undo should succeed");
|
||||
assert_eq!(cp2.turn_number, 1);
|
||||
assert_eq!(manager.undo_count(), 1);
|
||||
|
||||
// Verify we walked backwards through distinct checkpoints
|
||||
assert_ne!(cp1.turn_number, cp2.turn_number);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_undo_redo_cycle_preserves_state() {
|
||||
let mut manager = UndoManager::new();
|
||||
|
||||
let msgs_t0: Vec<ChatMessage> = vec![];
|
||||
let msgs_t1 = vec![ChatMessage::user("hello")];
|
||||
let msgs_t2 = vec![ChatMessage::user("hello"), ChatMessage::assistant("hi")];
|
||||
|
||||
manager.checkpoint(0, msgs_t0, "Turn 0");
|
||||
manager.checkpoint(1, msgs_t1, "Turn 1");
|
||||
|
||||
// Undo from turn 2 -> get turn 1 checkpoint
|
||||
let cp_undo1 = manager
|
||||
.undo(2, msgs_t2.clone())
|
||||
.expect("undo should succeed");
|
||||
assert_eq!(cp_undo1.turn_number, 1);
|
||||
|
||||
// Redo from turn 1 -> get turn 2 state back
|
||||
let cp_redo = manager
|
||||
.redo(cp_undo1.turn_number, cp_undo1.messages)
|
||||
.expect("redo should succeed");
|
||||
assert_eq!(cp_redo.turn_number, 2);
|
||||
assert_eq!(cp_redo.messages.len(), 2);
|
||||
|
||||
// Undo again from turn 2 -> should go back to turn 1 again
|
||||
let cp_undo2 = manager
|
||||
.undo(cp_redo.turn_number, cp_redo.messages)
|
||||
.expect("second undo should succeed");
|
||||
assert_eq!(cp_undo2.turn_number, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_undo_redo_stack_sizes_consistent() {
|
||||
let mut manager = UndoManager::new();
|
||||
|
||||
manager.checkpoint(0, vec![], "Turn 0");
|
||||
manager.checkpoint(1, vec![ChatMessage::user("a")], "Turn 1");
|
||||
manager.checkpoint(2, vec![ChatMessage::user("b")], "Turn 2");
|
||||
|
||||
// Start: undo=3, redo=0, total=3
|
||||
let total = manager.undo_count() + manager.redo_count();
|
||||
assert_eq!(total, 3);
|
||||
|
||||
// After undo: total should still be 3 (one moved from undo to redo,
|
||||
// plus the current state pushed to redo)
|
||||
// Actually: undo pops one (3->2), pushes current to redo (0->1), total=3
|
||||
let cp = manager.undo(3, vec![]).unwrap();
|
||||
assert_eq!(manager.undo_count() + manager.redo_count(), 3);
|
||||
|
||||
// After redo: redo pops one (1->0), pushes current to undo (2->3), total=3
|
||||
let cp2 = manager.redo(cp.turn_number, cp.messages).unwrap();
|
||||
assert_eq!(manager.undo_count() + manager.redo_count(), 3);
|
||||
|
||||
// After another undo: same invariant
|
||||
let _cp3 = manager.undo(cp2.turn_number, cp2.messages).unwrap();
|
||||
assert_eq!(manager.undo_count() + manager.redo_count(), 3);
|
||||
}
|
||||
}
|
||||
|
||||
+170
-72
@@ -10,8 +10,9 @@ use uuid::Uuid;
|
||||
use crate::agent::scheduler::WorkerMessage;
|
||||
use crate::agent::task::TaskOutput;
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::Error;
|
||||
use crate::history::Store;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::{
|
||||
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
|
||||
};
|
||||
@@ -28,7 +29,8 @@ pub struct WorkerDeps {
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub hooks: Arc<HookRegistry>,
|
||||
pub timeout: Duration,
|
||||
pub use_planning: bool,
|
||||
}
|
||||
@@ -67,7 +69,7 @@ impl Worker {
|
||||
&self.deps.tools
|
||||
}
|
||||
|
||||
fn store(&self) -> Option<&Arc<Store>> {
|
||||
fn store(&self) -> Option<&Arc<dyn Database>> {
|
||||
self.deps.store.as_ref()
|
||||
}
|
||||
|
||||
@@ -227,11 +229,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
// Check for cancellation
|
||||
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await {
|
||||
if ctx.state == JobState::Cancelled {
|
||||
tracing::info!("Worker for job {} detected cancellation", self.job_id);
|
||||
return Ok(());
|
||||
}
|
||||
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await
|
||||
&& ctx.state == JobState::Cancelled
|
||||
{
|
||||
tracing::info!("Worker for job {} detected cancellation", self.job_id);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
iteration += 1;
|
||||
@@ -248,16 +250,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
|
||||
if selections.is_empty() {
|
||||
// No tools from select_tools, ask LLM directly (may still return tool calls)
|
||||
let respond_result = reasoning.respond_with_tools(reason_ctx).await?;
|
||||
let respond_output = reasoning.respond_with_tools(reason_ctx).await?;
|
||||
|
||||
match respond_result {
|
||||
match respond_output.result {
|
||||
RespondResult::Text(response) => {
|
||||
// Check for completion keywords
|
||||
let response_lower = response.to_lowercase();
|
||||
if response_lower.contains("complete")
|
||||
|| response_lower.contains("finished")
|
||||
|| response_lower.contains("done")
|
||||
{
|
||||
// Check for explicit completion phrases. Use word-boundary
|
||||
// aware checks to avoid false positives like "incomplete",
|
||||
// "not done", or "unfinished". Only the LLM's own response
|
||||
// (not tool output) can trigger this.
|
||||
if crate::util::llm_signals_completion(&response) {
|
||||
self.mark_completed().await?;
|
||||
return Ok(());
|
||||
}
|
||||
@@ -300,6 +301,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
parameters: tc.arguments.clone(),
|
||||
reasoning: String::new(),
|
||||
alternatives: vec![],
|
||||
tool_call_id: tc.id.clone(),
|
||||
};
|
||||
|
||||
self.process_tool_result(reason_ctx, &selection, result)
|
||||
@@ -352,23 +354,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
.map(|selection| {
|
||||
let tool_name = selection.tool_name.clone();
|
||||
let params = selection.parameters.clone();
|
||||
let tools = self.tools().clone();
|
||||
let context_manager = self.context_manager().clone();
|
||||
let safety = self.safety().clone();
|
||||
let deps = self.deps.clone();
|
||||
let job_id = self.job_id;
|
||||
let store = self.deps.store.clone();
|
||||
|
||||
async move {
|
||||
let result = Self::execute_tool_inner(
|
||||
tools,
|
||||
context_manager,
|
||||
safety,
|
||||
store,
|
||||
job_id,
|
||||
&tool_name,
|
||||
¶ms,
|
||||
)
|
||||
.await;
|
||||
let result = Self::execute_tool_inner(&deps, job_id, &tool_name, ¶ms).await;
|
||||
ToolExecResult { result }
|
||||
}
|
||||
})
|
||||
@@ -379,20 +369,18 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
|
||||
/// Inner tool execution logic that can be called from both single and parallel paths.
|
||||
async fn execute_tool_inner(
|
||||
tools: Arc<ToolRegistry>,
|
||||
context_manager: Arc<ContextManager>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
store: Option<Arc<Store>>,
|
||||
deps: &WorkerDeps,
|
||||
job_id: Uuid,
|
||||
tool_name: &str,
|
||||
params: &serde_json::Value,
|
||||
) -> Result<String, Error> {
|
||||
let tool = tools
|
||||
.get(tool_name)
|
||||
.await
|
||||
.ok_or_else(|| crate::error::ToolError::NotFound {
|
||||
name: tool_name.to_string(),
|
||||
})?;
|
||||
let tool =
|
||||
deps.tools
|
||||
.get(tool_name)
|
||||
.await
|
||||
.ok_or_else(|| crate::error::ToolError::NotFound {
|
||||
name: tool_name.to_string(),
|
||||
})?;
|
||||
|
||||
// Tools requiring approval are blocked in autonomous jobs
|
||||
if tool.requires_approval() {
|
||||
@@ -402,8 +390,46 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
.into());
|
||||
}
|
||||
|
||||
// Get job context for the tool
|
||||
let job_ctx = context_manager.get_context(job_id).await?;
|
||||
// Fetch job context early so we have the real user_id for hooks
|
||||
let job_ctx = deps.context_manager.get_context(job_id).await?;
|
||||
|
||||
// Run BeforeToolCall hook
|
||||
let params = {
|
||||
use crate::hooks::{HookError, HookEvent, HookOutcome};
|
||||
let event = HookEvent::ToolCall {
|
||||
tool_name: tool_name.to_string(),
|
||||
parameters: params.clone(),
|
||||
user_id: job_ctx.user_id.clone(),
|
||||
context: format!("job:{}", job_id),
|
||||
};
|
||||
match deps.hooks.run(&event).await {
|
||||
Err(HookError::Rejected { reason }) => {
|
||||
return Err(crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: format!("Blocked by hook: {}", reason),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: format!("Blocked by hook failure mode: {}", err),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
Ok(HookOutcome::Continue {
|
||||
modified: Some(new_params),
|
||||
}) => serde_json::from_str(&new_params).unwrap_or_else(|e| {
|
||||
tracing::warn!(
|
||||
tool = %tool_name,
|
||||
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
||||
e
|
||||
);
|
||||
params.clone()
|
||||
}),
|
||||
_ => params.clone(),
|
||||
}
|
||||
};
|
||||
if job_ctx.state == JobState::Cancelled {
|
||||
return Err(crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
@@ -413,7 +439,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
// Validate tool parameters
|
||||
let validation = safety.validator().validate_tool_params(params);
|
||||
let validation = deps.safety.validator().validate_tool_params(¶ms);
|
||||
if !validation.is_valid {
|
||||
let details = validation
|
||||
.errors
|
||||
@@ -478,8 +504,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
Ok(Ok(output)) => {
|
||||
let output_str = serde_json::to_string_pretty(&output.result)
|
||||
.ok()
|
||||
.map(|s| safety.sanitize_tool_output(tool_name, &s).content);
|
||||
context_manager
|
||||
.map(|s| deps.safety.sanitize_tool_output(tool_name, &s).content);
|
||||
deps.context_manager
|
||||
.update_memory(job_id, |mem| {
|
||||
let rec = mem.create_action(tool_name, params.clone()).succeed(
|
||||
output_str.clone(),
|
||||
@@ -492,7 +518,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
.await
|
||||
.ok()
|
||||
}
|
||||
Ok(Err(e)) => context_manager
|
||||
Ok(Err(e)) => deps
|
||||
.context_manager
|
||||
.update_memory(job_id, |mem| {
|
||||
let rec = mem
|
||||
.create_action(tool_name, params.clone())
|
||||
@@ -502,7 +529,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
})
|
||||
.await
|
||||
.ok(),
|
||||
Err(_) => context_manager
|
||||
Err(_) => deps
|
||||
.context_manager
|
||||
.update_memory(job_id, |mem| {
|
||||
let rec = mem
|
||||
.create_action(tool_name, params.clone())
|
||||
@@ -515,7 +543,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
};
|
||||
|
||||
// Persist action to database (fire-and-forget)
|
||||
if let (Some(action), Some(store)) = (action, store) {
|
||||
if let (Some(action), Some(store)) = (action, deps.store.clone()) {
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store.save_action(job_id, &action).await {
|
||||
tracing::warn!("Failed to persist action for job {}: {}", job_id, e);
|
||||
@@ -566,17 +594,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
);
|
||||
|
||||
reason_ctx.messages.push(ChatMessage::tool_result(
|
||||
"tool_call_id",
|
||||
&selection.tool_call_id,
|
||||
&selection.tool_name,
|
||||
wrapped,
|
||||
));
|
||||
|
||||
// Check if job is complete
|
||||
if output.contains("TASK_COMPLETE") || output.contains("JOB_DONE") {
|
||||
self.mark_completed().await?;
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
// Tool output never drives job completion. A malicious tool could
|
||||
// emit "TASK_COMPLETE" to force premature completion. Only the LLM's
|
||||
// own structured response (in execution_loop) can mark a job done.
|
||||
Ok(false)
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -601,7 +626,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
reason_ctx.messages.push(ChatMessage::tool_result(
|
||||
"tool_call_id",
|
||||
&selection.tool_call_id,
|
||||
&selection.tool_name,
|
||||
format!("Error: {}", e),
|
||||
));
|
||||
@@ -651,12 +676,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
.execute_tool(&action.tool_name, &action.parameters)
|
||||
.await;
|
||||
|
||||
// Create a synthetic ToolSelection for process_tool_result
|
||||
// Create a synthetic ToolSelection for process_tool_result.
|
||||
// Plan actions don't originate from an LLM tool_call response so
|
||||
// there is no real tool_call_id; generate a unique one.
|
||||
let selection = ToolSelection {
|
||||
tool_name: action.tool_name.clone(),
|
||||
parameters: action.parameters.clone(),
|
||||
reasoning: action.reasoning.clone(),
|
||||
alternatives: vec![],
|
||||
tool_call_id: format!("plan_{}_{}", self.job_id, i),
|
||||
};
|
||||
|
||||
// Process the result
|
||||
@@ -680,11 +708,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
let response = reasoning.respond(reason_ctx).await?;
|
||||
reason_ctx.messages.push(ChatMessage::assistant(&response));
|
||||
|
||||
let response_lower = response.to_lowercase();
|
||||
if response_lower.contains("complete")
|
||||
|| response_lower.contains("finished")
|
||||
|| response_lower.contains("done")
|
||||
{
|
||||
if crate::util::llm_signals_completion(&response) {
|
||||
self.mark_completed().await?;
|
||||
} else {
|
||||
// Job not complete, could re-plan or fall back to direct selection
|
||||
@@ -705,16 +729,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
tool_name: &str,
|
||||
params: &serde_json::Value,
|
||||
) -> Result<String, Error> {
|
||||
Self::execute_tool_inner(
|
||||
self.tools().clone(),
|
||||
self.context_manager().clone(),
|
||||
self.safety().clone(),
|
||||
self.deps.store.clone(),
|
||||
self.job_id,
|
||||
tool_name,
|
||||
params,
|
||||
)
|
||||
.await
|
||||
Self::execute_tool_inner(&self.deps, self.job_id, tool_name, params).await
|
||||
}
|
||||
|
||||
async fn mark_completed(&self) -> Result<(), Error> {
|
||||
@@ -779,3 +794,86 @@ impl From<TaskOutput> for Result<String, Error> {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::llm::ToolSelection;
|
||||
use crate::util::llm_signals_completion;
|
||||
|
||||
#[test]
|
||||
fn test_tool_selection_preserves_call_id() {
|
||||
let selection = ToolSelection {
|
||||
tool_name: "memory_search".to_string(),
|
||||
parameters: serde_json::json!({"query": "test"}),
|
||||
reasoning: "Need to search memory".to_string(),
|
||||
alternatives: vec![],
|
||||
tool_call_id: "call_abc123".to_string(),
|
||||
};
|
||||
|
||||
assert_eq!(selection.tool_call_id, "call_abc123");
|
||||
assert_ne!(
|
||||
selection.tool_call_id, "tool_call_id",
|
||||
"tool_call_id must not be the hardcoded placeholder string"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completion_positive_signals() {
|
||||
assert!(llm_signals_completion("The job is complete."));
|
||||
assert!(llm_signals_completion(
|
||||
"I have completed the task successfully."
|
||||
));
|
||||
assert!(llm_signals_completion("The task is done."));
|
||||
assert!(llm_signals_completion("The task is finished."));
|
||||
assert!(llm_signals_completion(
|
||||
"All steps are complete and verified."
|
||||
));
|
||||
assert!(llm_signals_completion(
|
||||
"I've done all the work. The work is done."
|
||||
));
|
||||
assert!(llm_signals_completion(
|
||||
"Successfully completed the migration."
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completion_negative_signals_block_false_positives() {
|
||||
// These contain completion keywords but also negation, should NOT trigger.
|
||||
assert!(!llm_signals_completion("The task is not complete yet."));
|
||||
assert!(!llm_signals_completion("This is not done."));
|
||||
assert!(!llm_signals_completion("The work is incomplete."));
|
||||
assert!(!llm_signals_completion(
|
||||
"The migration is not yet finished."
|
||||
));
|
||||
assert!(!llm_signals_completion("The job isn't done yet."));
|
||||
assert!(!llm_signals_completion("This remains unfinished."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completion_does_not_match_bare_substrings() {
|
||||
// Bare words embedded in other text should NOT trigger completion.
|
||||
assert!(!llm_signals_completion(
|
||||
"I need to complete more work first."
|
||||
));
|
||||
assert!(!llm_signals_completion(
|
||||
"Let me finish the remaining steps."
|
||||
));
|
||||
assert!(!llm_signals_completion(
|
||||
"I'm done analyzing, now let me fix it."
|
||||
));
|
||||
assert!(!llm_signals_completion(
|
||||
"I completed step 1 but step 2 remains."
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completion_tool_output_injection() {
|
||||
// A malicious tool output echoed by the LLM should not trigger
|
||||
// completion unless it forms a genuine completion phrase.
|
||||
assert!(!llm_signals_completion("TASK_COMPLETE"));
|
||||
assert!(!llm_signals_completion("JOB_DONE"));
|
||||
assert!(!llm_signals_completion(
|
||||
"The tool returned: TASK_COMPLETE signal"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Boot screen displayed after all initialization completes.
|
||||
//!
|
||||
//! Shows a polished ANSI-styled status panel summarizing the agent's runtime
|
||||
//! state: model, database, tool count, enabled features, active channels,
|
||||
//! and the gateway URL.
|
||||
|
||||
/// All displayable fields for the boot screen.
|
||||
pub struct BootInfo {
|
||||
pub version: String,
|
||||
pub agent_name: String,
|
||||
pub llm_backend: String,
|
||||
pub llm_model: String,
|
||||
pub cheap_model: Option<String>,
|
||||
pub db_backend: String,
|
||||
pub db_connected: bool,
|
||||
pub tool_count: usize,
|
||||
pub gateway_url: Option<String>,
|
||||
pub embeddings_enabled: bool,
|
||||
pub embeddings_provider: Option<String>,
|
||||
pub heartbeat_enabled: bool,
|
||||
pub heartbeat_interval_secs: u64,
|
||||
pub sandbox_enabled: bool,
|
||||
pub claude_code_enabled: bool,
|
||||
pub routines_enabled: bool,
|
||||
pub channels: Vec<String>,
|
||||
}
|
||||
|
||||
/// Print the boot screen to stdout.
|
||||
pub fn print_boot_screen(info: &BootInfo) {
|
||||
// ANSI codes matching existing REPL palette
|
||||
let bold = "\x1b[1m";
|
||||
let cyan = "\x1b[36m";
|
||||
let dim = "\x1b[90m";
|
||||
let yellow_underline = "\x1b[33;4m";
|
||||
let reset = "\x1b[0m";
|
||||
|
||||
let border = format!(" {dim}{}{reset}", "\u{2576}".repeat(58));
|
||||
|
||||
println!();
|
||||
println!("{border}");
|
||||
println!();
|
||||
println!(" {bold}{}{reset} v{}", info.agent_name, info.version);
|
||||
println!();
|
||||
|
||||
// Model line
|
||||
let model_display = if let Some(ref cheap) = info.cheap_model {
|
||||
format!(
|
||||
"{cyan}{}{reset} {dim}cheap{reset} {cyan}{}{reset}",
|
||||
info.llm_model, cheap
|
||||
)
|
||||
} else {
|
||||
format!("{cyan}{}{reset}", info.llm_model)
|
||||
};
|
||||
println!(
|
||||
" {dim}model{reset} {model_display} {dim}via {}{reset}",
|
||||
info.llm_backend
|
||||
);
|
||||
|
||||
// Database line
|
||||
let db_status = if info.db_connected {
|
||||
"connected"
|
||||
} else {
|
||||
"none"
|
||||
};
|
||||
println!(
|
||||
" {dim}database{reset} {cyan}{}{reset} {dim}({db_status}){reset}",
|
||||
info.db_backend
|
||||
);
|
||||
|
||||
// Tools line
|
||||
println!(
|
||||
" {dim}tools{reset} {cyan}{}{reset} {dim}registered{reset}",
|
||||
info.tool_count
|
||||
);
|
||||
|
||||
// Features line
|
||||
let mut features = Vec::new();
|
||||
if info.embeddings_enabled {
|
||||
if let Some(ref provider) = info.embeddings_provider {
|
||||
features.push(format!("embeddings ({provider})"));
|
||||
} else {
|
||||
features.push("embeddings".to_string());
|
||||
}
|
||||
}
|
||||
if info.heartbeat_enabled {
|
||||
let mins = info.heartbeat_interval_secs / 60;
|
||||
features.push(format!("heartbeat ({mins}m)"));
|
||||
}
|
||||
if info.sandbox_enabled {
|
||||
features.push("sandbox".to_string());
|
||||
}
|
||||
if info.claude_code_enabled {
|
||||
features.push("claude-code".to_string());
|
||||
}
|
||||
if info.routines_enabled {
|
||||
features.push("routines".to_string());
|
||||
}
|
||||
if !features.is_empty() {
|
||||
println!(
|
||||
" {dim}features{reset} {cyan}{}{reset}",
|
||||
features.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
// Channels line
|
||||
if !info.channels.is_empty() {
|
||||
println!(
|
||||
" {dim}channels{reset} {cyan}{}{reset}",
|
||||
info.channels.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
// Gateway URL (highlighted)
|
||||
if let Some(ref url) = info.gateway_url {
|
||||
println!();
|
||||
println!(" {dim}gateway{reset} {yellow_underline}{url}{reset}");
|
||||
}
|
||||
|
||||
println!();
|
||||
println!("{border}");
|
||||
println!();
|
||||
println!(" /help for commands, /quit to exit");
|
||||
println!();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_print_boot_screen_full() {
|
||||
let info = BootInfo {
|
||||
version: "0.2.0".to_string(),
|
||||
agent_name: "ironclaw".to_string(),
|
||||
llm_backend: "nearai".to_string(),
|
||||
llm_model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
cheap_model: Some("gpt-4o-mini".to_string()),
|
||||
db_backend: "libsql".to_string(),
|
||||
db_connected: true,
|
||||
tool_count: 24,
|
||||
gateway_url: Some("http://127.0.0.1:3001/?token=abc123".to_string()),
|
||||
embeddings_enabled: true,
|
||||
embeddings_provider: Some("openai".to_string()),
|
||||
heartbeat_enabled: true,
|
||||
heartbeat_interval_secs: 1800,
|
||||
sandbox_enabled: true,
|
||||
claude_code_enabled: false,
|
||||
routines_enabled: true,
|
||||
channels: vec![
|
||||
"repl".to_string(),
|
||||
"gateway".to_string(),
|
||||
"telegram".to_string(),
|
||||
],
|
||||
};
|
||||
// Should not panic
|
||||
print_boot_screen(&info);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_print_boot_screen_minimal() {
|
||||
let info = BootInfo {
|
||||
version: "0.2.0".to_string(),
|
||||
agent_name: "ironclaw".to_string(),
|
||||
llm_backend: "nearai".to_string(),
|
||||
llm_model: "gpt-4o".to_string(),
|
||||
cheap_model: None,
|
||||
db_backend: "none".to_string(),
|
||||
db_connected: false,
|
||||
tool_count: 5,
|
||||
gateway_url: None,
|
||||
embeddings_enabled: false,
|
||||
embeddings_provider: None,
|
||||
heartbeat_enabled: false,
|
||||
heartbeat_interval_secs: 0,
|
||||
sandbox_enabled: false,
|
||||
claude_code_enabled: false,
|
||||
routines_enabled: false,
|
||||
channels: vec![],
|
||||
};
|
||||
// Should not panic
|
||||
print_boot_screen(&info);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_print_boot_screen_no_features() {
|
||||
let info = BootInfo {
|
||||
version: "0.1.0".to_string(),
|
||||
agent_name: "test".to_string(),
|
||||
llm_backend: "openai".to_string(),
|
||||
llm_model: "gpt-4o".to_string(),
|
||||
cheap_model: None,
|
||||
db_backend: "postgres".to_string(),
|
||||
db_connected: true,
|
||||
tool_count: 10,
|
||||
gateway_url: None,
|
||||
embeddings_enabled: false,
|
||||
embeddings_provider: None,
|
||||
heartbeat_enabled: false,
|
||||
heartbeat_interval_secs: 0,
|
||||
sandbox_enabled: false,
|
||||
claude_code_enabled: false,
|
||||
routines_enabled: false,
|
||||
channels: vec!["repl".to_string()],
|
||||
};
|
||||
// Should not panic
|
||||
print_boot_screen(&info);
|
||||
}
|
||||
}
|
||||
+306
-167
@@ -1,147 +1,145 @@
|
||||
//! Bootstrap configuration for IronClaw.
|
||||
//! Bootstrap helpers for IronClaw.
|
||||
//!
|
||||
//! These are the only settings that MUST live on disk because they're needed
|
||||
//! before the database connection is established. Everything else lives in the
|
||||
//! `settings` table in PostgreSQL.
|
||||
//! The only setting that truly needs disk persistence before the database is
|
||||
//! available is `DATABASE_URL` (chicken-and-egg: can't connect to DB without
|
||||
//! it). Everything else is auto-detected or read from env vars.
|
||||
//!
|
||||
//! File: `~/.ironclaw/bootstrap.json`
|
||||
//! File: `~/.ironclaw/.env` (standard dotenvy format)
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::settings::KeySource;
|
||||
|
||||
/// Minimal config needed to connect to the database and decrypt secrets.
|
||||
///
|
||||
/// This is the only JSON file IronClaw reads from disk at startup.
|
||||
/// All other configuration lives in the `settings` table in PostgreSQL.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BootstrapConfig {
|
||||
/// Database connection URL (postgres://...).
|
||||
#[serde(default)]
|
||||
pub database_url: Option<String>,
|
||||
|
||||
/// Database connection pool size.
|
||||
#[serde(default)]
|
||||
pub database_pool_size: Option<usize>,
|
||||
|
||||
/// Source for the secrets master key.
|
||||
#[serde(default)]
|
||||
pub secrets_master_key_source: KeySource,
|
||||
|
||||
/// Whether onboarding wizard has been completed.
|
||||
#[serde(default)]
|
||||
pub onboard_completed: bool,
|
||||
/// Path to the IronClaw-specific `.env` file: `~/.ironclaw/.env`.
|
||||
pub fn ironclaw_env_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join(".env")
|
||||
}
|
||||
|
||||
impl Default for BootstrapConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
database_url: None,
|
||||
database_pool_size: None,
|
||||
secrets_master_key_source: KeySource::None,
|
||||
onboard_completed: false,
|
||||
}
|
||||
/// Load env vars from `~/.ironclaw/.env` (in addition to the standard `.env`).
|
||||
///
|
||||
/// Call this **after** `dotenvy::dotenv()` so that the standard `./.env`
|
||||
/// takes priority over `~/.ironclaw/.env`. dotenvy never overwrites
|
||||
/// existing env vars, so the effective priority is:
|
||||
///
|
||||
/// explicit env vars > `./.env` > `~/.ironclaw/.env`
|
||||
///
|
||||
/// If `~/.ironclaw/.env` doesn't exist but the legacy `bootstrap.json` does,
|
||||
/// extracts `DATABASE_URL` from it and writes the `.env` file (one-time
|
||||
/// upgrade from the old config format).
|
||||
pub fn load_ironclaw_env() {
|
||||
let path = ironclaw_env_path();
|
||||
|
||||
if !path.exists() {
|
||||
// One-time upgrade: extract DATABASE_URL from legacy bootstrap.json
|
||||
migrate_bootstrap_json_to_env(&path);
|
||||
}
|
||||
|
||||
if path.exists() {
|
||||
let _ = dotenvy::from_path(&path);
|
||||
}
|
||||
}
|
||||
|
||||
impl BootstrapConfig {
|
||||
/// Default bootstrap file path: `~/.ironclaw/bootstrap.json`.
|
||||
pub fn default_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("bootstrap.json")
|
||||
/// If `bootstrap.json` exists, pull `database_url` out of it and write `.env`.
|
||||
fn migrate_bootstrap_json_to_env(env_path: &std::path::Path) {
|
||||
let ironclaw_dir = env_path
|
||||
.parent()
|
||||
.unwrap_or_else(|| std::path::Path::new("."));
|
||||
let bootstrap_path = ironclaw_dir.join("bootstrap.json");
|
||||
|
||||
if !bootstrap_path.exists() {
|
||||
return;
|
||||
}
|
||||
|
||||
/// Legacy settings.json path (for migration detection).
|
||||
pub fn legacy_settings_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("settings.json")
|
||||
}
|
||||
let content = match std::fs::read_to_string(&bootstrap_path) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
/// Load from the default path, falling back to legacy settings.json,
|
||||
/// then to defaults if neither exists.
|
||||
pub fn load() -> Self {
|
||||
let bootstrap_path = Self::default_path();
|
||||
if bootstrap_path.exists() {
|
||||
return Self::load_from(&bootstrap_path);
|
||||
// Minimal parse: just grab database_url from the JSON
|
||||
let parsed: serde_json::Value = match serde_json::from_str(&content) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
if let Some(url) = parsed.get("database_url").and_then(|v| v.as_str()) {
|
||||
if let Some(parent) = env_path.parent()
|
||||
&& let Err(e) = std::fs::create_dir_all(parent)
|
||||
{
|
||||
eprintln!("Warning: failed to create {}: {}", parent.display(), e);
|
||||
return;
|
||||
}
|
||||
|
||||
// Fall back to legacy settings.json (extract just the 4 bootstrap fields)
|
||||
let legacy_path = Self::legacy_settings_path();
|
||||
if legacy_path.exists() {
|
||||
return Self::load_from_legacy(&legacy_path);
|
||||
if let Err(e) = std::fs::write(env_path, format!("DATABASE_URL=\"{}\"\n", url)) {
|
||||
eprintln!("Warning: failed to migrate bootstrap.json to .env: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Load from a specific path.
|
||||
pub fn load_from(path: &PathBuf) -> Self {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
|
||||
Err(_) => Self::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract bootstrap fields from a legacy settings.json.
|
||||
fn load_from_legacy(path: &PathBuf) -> Self {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(data) => {
|
||||
// The legacy Settings struct is a superset; serde will ignore extra fields.
|
||||
serde_json::from_str(&data).unwrap_or_default()
|
||||
}
|
||||
Err(_) => Self::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Save to the default path.
|
||||
pub fn save(&self) -> std::io::Result<()> {
|
||||
self.save_to(&Self::default_path())
|
||||
}
|
||||
|
||||
/// Save to a specific path.
|
||||
pub fn save_to(&self, path: &PathBuf) -> std::io::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let json = serde_json::to_string_pretty(self)
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
|
||||
std::fs::write(path, json)
|
||||
rename_to_migrated(&bootstrap_path);
|
||||
eprintln!(
|
||||
"Migrated DATABASE_URL from bootstrap.json to {}",
|
||||
env_path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// One-time migration from disk config files to the database settings table.
|
||||
/// Write database bootstrap vars to `~/.ironclaw/.env`.
|
||||
///
|
||||
/// On first boot after upgrade, checks if:
|
||||
/// 1. `~/.ironclaw/settings.json` exists
|
||||
/// 2. The DB settings table is empty for this user
|
||||
/// These settings form the chicken-and-egg layer: they must be available
|
||||
/// from the filesystem (env vars) BEFORE any database connection, because
|
||||
/// they determine which database to connect to. Everything else is stored
|
||||
/// in the database itself.
|
||||
///
|
||||
/// If both conditions hold, migrates settings, MCP servers, and session data
|
||||
/// to the database, writes `bootstrap.json`, and renames old files to `.migrated`.
|
||||
/// Creates the parent directory if it doesn't exist.
|
||||
/// Values are double-quoted so that `#` (common in URL-encoded passwords)
|
||||
/// and other shell-special characters are preserved by dotenvy.
|
||||
pub fn save_bootstrap_env(vars: &[(&str, &str)]) -> std::io::Result<()> {
|
||||
let path = ironclaw_env_path();
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let mut content = String::new();
|
||||
for (key, value) in vars {
|
||||
content.push_str(&format!("{}=\"{}\"\n", key, value));
|
||||
}
|
||||
std::fs::write(&path, content)
|
||||
}
|
||||
|
||||
/// Write `DATABASE_URL` to `~/.ironclaw/.env`.
|
||||
///
|
||||
/// Convenience wrapper around `save_bootstrap_env` for single-value migration
|
||||
/// paths. Prefer `save_bootstrap_env` for new code.
|
||||
pub fn save_database_url(url: &str) -> std::io::Result<()> {
|
||||
save_bootstrap_env(&[("DATABASE_URL", url)])
|
||||
}
|
||||
|
||||
/// One-time migration of legacy `~/.ironclaw/settings.json` into the database.
|
||||
///
|
||||
/// Only runs when a `settings.json` exists on disk AND the DB has no settings
|
||||
/// yet. After the wizard writes directly to the DB, this path is only hit by
|
||||
/// users upgrading from the old disk-only configuration.
|
||||
///
|
||||
/// After syncing, renames `settings.json` to `.migrated` so it won't trigger again.
|
||||
pub async fn migrate_disk_to_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
) -> Result<(), MigrationError> {
|
||||
let legacy_settings_path = BootstrapConfig::legacy_settings_path();
|
||||
let ironclaw_dir = dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw");
|
||||
let legacy_settings_path = ironclaw_dir.join("settings.json");
|
||||
|
||||
if !legacy_settings_path.exists() {
|
||||
tracing::debug!("No legacy settings.json found, skipping disk-to-DB migration");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Only migrate if DB is empty for this user
|
||||
// If DB already has settings, this is not a first boot, the wizard already
|
||||
// wrote directly to the DB. Just clean up the stale file.
|
||||
let has_settings = store.has_settings(user_id).await.map_err(|e| {
|
||||
MigrationError::Database(format!("Failed to check existing settings: {}", e))
|
||||
})?;
|
||||
if has_settings {
|
||||
tracing::debug!(
|
||||
"DB already has settings for user '{}', skipping migration",
|
||||
user_id
|
||||
);
|
||||
tracing::info!("DB already has settings, renaming stale settings.json");
|
||||
rename_to_migrated(&legacy_settings_path);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -160,22 +158,14 @@ pub async fn migrate_disk_to_db(
|
||||
tracing::info!("Migrated {} settings to database", db_map.len());
|
||||
}
|
||||
|
||||
// 2. Write bootstrap.json with the 4 essential fields
|
||||
let bootstrap = BootstrapConfig {
|
||||
database_url: settings.database_url.clone(),
|
||||
database_pool_size: settings.database_pool_size,
|
||||
secrets_master_key_source: settings.secrets_master_key_source,
|
||||
onboard_completed: settings.onboard_completed,
|
||||
};
|
||||
bootstrap
|
||||
.save()
|
||||
.map_err(|e| MigrationError::Io(format!("Failed to write bootstrap.json: {}", e)))?;
|
||||
tracing::info!("Wrote bootstrap.json");
|
||||
// 2. Write DATABASE_URL to ~/.ironclaw/.env
|
||||
if let Some(ref url) = settings.database_url {
|
||||
save_database_url(url)
|
||||
.map_err(|e| MigrationError::Io(format!("Failed to write .env: {}", e)))?;
|
||||
tracing::info!("Wrote DATABASE_URL to {}", ironclaw_env_path().display());
|
||||
}
|
||||
|
||||
// 3. Migrate mcp-servers.json if it exists
|
||||
let ironclaw_dir = dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw");
|
||||
let mcp_path = ironclaw_dir.join("mcp-servers.json");
|
||||
if mcp_path.exists() {
|
||||
match std::fs::read_to_string(&mcp_path) {
|
||||
@@ -211,7 +201,7 @@ pub async fn migrate_disk_to_db(
|
||||
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
|
||||
Ok(value) => {
|
||||
store
|
||||
.set_setting(user_id, "nearai.session", &value)
|
||||
.set_setting(user_id, "nearai.session_token", &value)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
MigrationError::Database(format!(
|
||||
@@ -236,12 +226,19 @@ pub async fn migrate_disk_to_db(
|
||||
// 5. Rename settings.json to .migrated (don't delete, safety net)
|
||||
rename_to_migrated(&legacy_settings_path);
|
||||
|
||||
// 6. Clean up old bootstrap.json if it exists (superseded by .env)
|
||||
let old_bootstrap = ironclaw_dir.join("bootstrap.json");
|
||||
if old_bootstrap.exists() {
|
||||
rename_to_migrated(&old_bootstrap);
|
||||
tracing::info!("Renamed old bootstrap.json to .migrated");
|
||||
}
|
||||
|
||||
tracing::info!("Disk-to-DB migration complete");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rename a file to `<name>.migrated` as a safety net.
|
||||
fn rename_to_migrated(path: &PathBuf) {
|
||||
fn rename_to_migrated(path: &std::path::Path) {
|
||||
let mut migrated = path.as_os_str().to_owned();
|
||||
migrated.push(".migrated");
|
||||
if let Err(e) = std::fs::rename(path, &migrated) {
|
||||
@@ -264,62 +261,204 @@ mod tests {
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_save_load() {
|
||||
fn test_save_and_load_database_url() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("bootstrap.json");
|
||||
let env_path = dir.path().join(".env");
|
||||
|
||||
let config = BootstrapConfig {
|
||||
database_url: Some("postgres://localhost/test".to_string()),
|
||||
database_pool_size: Some(5),
|
||||
secrets_master_key_source: KeySource::Keychain,
|
||||
onboard_completed: true,
|
||||
};
|
||||
// Write in the quoted format that save_database_url uses
|
||||
let url = "postgres://localhost:5432/ironclaw_test";
|
||||
std::fs::write(&env_path, format!("DATABASE_URL=\"{}\"\n", url)).unwrap();
|
||||
|
||||
config.save_to(&path).unwrap();
|
||||
|
||||
let loaded = BootstrapConfig::load_from(&path);
|
||||
// Verify the content is a valid dotenv line (quoted)
|
||||
let content = std::fs::read_to_string(&env_path).unwrap();
|
||||
assert_eq!(
|
||||
loaded.database_url,
|
||||
Some("postgres://localhost/test".to_string())
|
||||
content,
|
||||
"DATABASE_URL=\"postgres://localhost:5432/ironclaw_test\"\n"
|
||||
);
|
||||
assert_eq!(loaded.database_pool_size, Some(5));
|
||||
assert_eq!(loaded.secrets_master_key_source, KeySource::Keychain);
|
||||
assert!(loaded.onboard_completed);
|
||||
|
||||
// Verify dotenvy can parse it (strips quotes automatically)
|
||||
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
assert_eq!(parsed[0].0, "DATABASE_URL");
|
||||
assert_eq!(parsed[0].1, url);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_from_legacy_settings() {
|
||||
fn test_save_database_url_with_hash_in_password() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("settings.json");
|
||||
let env_path = dir.path().join(".env");
|
||||
|
||||
// Write a legacy settings.json with many extra fields
|
||||
let legacy = serde_json::json!({
|
||||
"database_url": "postgres://localhost/ironclaw",
|
||||
"database_pool_size": 10,
|
||||
// URLs with # in the password are common (URL-encoded special chars).
|
||||
// Without quoting, dotenvy treats # as a comment delimiter.
|
||||
let url = "postgres://user:p%23ss@localhost:5432/ironclaw";
|
||||
std::fs::write(&env_path, format!("DATABASE_URL=\"{}\"\n", url)).unwrap();
|
||||
|
||||
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
assert_eq!(parsed[0].0, "DATABASE_URL");
|
||||
assert_eq!(parsed[0].1, url);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_database_url_creates_parent_dirs() {
|
||||
let dir = tempdir().unwrap();
|
||||
let nested = dir.path().join("deep").join("nested");
|
||||
let env_path = nested.join(".env");
|
||||
|
||||
// Parent doesn't exist yet
|
||||
assert!(!nested.exists());
|
||||
|
||||
// The global function uses a fixed path, so we test the logic directly
|
||||
std::fs::create_dir_all(&nested).unwrap();
|
||||
std::fs::write(&env_path, "DATABASE_URL=postgres://test\n").unwrap();
|
||||
|
||||
assert!(env_path.exists());
|
||||
let content = std::fs::read_to_string(&env_path).unwrap();
|
||||
assert!(content.contains("DATABASE_URL=postgres://test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ironclaw_env_path() {
|
||||
let path = ironclaw_env_path();
|
||||
assert!(path.ends_with(".ironclaw/.env"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_migrate_bootstrap_json_to_env() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join(".env");
|
||||
let bootstrap_path = dir.path().join("bootstrap.json");
|
||||
|
||||
// Write a legacy bootstrap.json
|
||||
let bootstrap_json = serde_json::json!({
|
||||
"database_url": "postgres://localhost/ironclaw_upgrade",
|
||||
"database_pool_size": 5,
|
||||
"secrets_master_key_source": "keychain",
|
||||
"onboard_completed": true,
|
||||
"selected_model": "claude-3-5-sonnet",
|
||||
"agent": { "name": "testbot", "max_parallel_jobs": 3 },
|
||||
"heartbeat": { "enabled": true }
|
||||
"onboard_completed": true
|
||||
});
|
||||
std::fs::write(&path, serde_json::to_string_pretty(&legacy).unwrap()).unwrap();
|
||||
std::fs::write(
|
||||
&bootstrap_path,
|
||||
serde_json::to_string_pretty(&bootstrap_json).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = BootstrapConfig::load_from_legacy(&path);
|
||||
assert!(!env_path.exists());
|
||||
assert!(bootstrap_path.exists());
|
||||
|
||||
// Run the migration
|
||||
migrate_bootstrap_json_to_env(&env_path);
|
||||
|
||||
// .env should now exist with DATABASE_URL
|
||||
assert!(env_path.exists());
|
||||
let content = std::fs::read_to_string(&env_path).unwrap();
|
||||
assert_eq!(
|
||||
config.database_url,
|
||||
Some("postgres://localhost/ironclaw".to_string())
|
||||
content,
|
||||
"DATABASE_URL=\"postgres://localhost/ironclaw_upgrade\"\n"
|
||||
);
|
||||
assert_eq!(config.database_pool_size, Some(10));
|
||||
assert_eq!(config.secrets_master_key_source, KeySource::Keychain);
|
||||
assert!(config.onboard_completed);
|
||||
|
||||
// bootstrap.json should be renamed to .migrated
|
||||
assert!(!bootstrap_path.exists());
|
||||
assert!(dir.path().join("bootstrap.json.migrated").exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_defaults() {
|
||||
let config = BootstrapConfig::default();
|
||||
assert!(config.database_url.is_none());
|
||||
assert!(config.database_pool_size.is_none());
|
||||
assert_eq!(config.secrets_master_key_source, KeySource::None);
|
||||
assert!(!config.onboard_completed);
|
||||
fn test_migrate_bootstrap_json_no_database_url() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join(".env");
|
||||
let bootstrap_path = dir.path().join("bootstrap.json");
|
||||
|
||||
// bootstrap.json with no database_url
|
||||
let bootstrap_json = serde_json::json!({
|
||||
"onboard_completed": false
|
||||
});
|
||||
std::fs::write(
|
||||
&bootstrap_path,
|
||||
serde_json::to_string_pretty(&bootstrap_json).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
migrate_bootstrap_json_to_env(&env_path);
|
||||
|
||||
// .env should NOT be created
|
||||
assert!(!env_path.exists());
|
||||
// bootstrap.json should remain (no migration happened)
|
||||
assert!(bootstrap_path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_migrate_bootstrap_json_missing() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join(".env");
|
||||
|
||||
// No bootstrap.json at all
|
||||
migrate_bootstrap_json_to_env(&env_path);
|
||||
|
||||
// Nothing should happen
|
||||
assert!(!env_path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_bootstrap_env_multiple_vars() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join("nested").join(".env");
|
||||
|
||||
std::fs::create_dir_all(env_path.parent().unwrap()).unwrap();
|
||||
|
||||
let vars = [
|
||||
("DATABASE_BACKEND", "libsql"),
|
||||
("LIBSQL_PATH", "/home/user/.ironclaw/ironclaw.db"),
|
||||
];
|
||||
|
||||
// Write manually to the temp path (save_bootstrap_env uses the global path)
|
||||
let mut content = String::new();
|
||||
for (key, value) in &vars {
|
||||
content.push_str(&format!("{}=\"{}\"\n", key, value));
|
||||
}
|
||||
std::fs::write(&env_path, &content).unwrap();
|
||||
|
||||
// Verify dotenvy can parse all entries
|
||||
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
assert_eq!(parsed.len(), 2);
|
||||
assert_eq!(
|
||||
parsed[0],
|
||||
("DATABASE_BACKEND".to_string(), "libsql".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
parsed[1],
|
||||
(
|
||||
"LIBSQL_PATH".to_string(),
|
||||
"/home/user/.ironclaw/ironclaw.db".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_bootstrap_env_overwrites_previous() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join(".env");
|
||||
|
||||
// Write initial content
|
||||
std::fs::write(&env_path, "DATABASE_URL=\"postgres://old\"\n").unwrap();
|
||||
|
||||
// Overwrite with new vars (simulating save_bootstrap_env behavior)
|
||||
let content = "DATABASE_BACKEND=\"libsql\"\nLIBSQL_PATH=\"/new/path.db\"\n";
|
||||
std::fs::write(&env_path, content).unwrap();
|
||||
|
||||
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
// Old DATABASE_URL should be gone
|
||||
assert_eq!(parsed.len(), 2);
|
||||
assert!(parsed.iter().all(|(k, _)| k != "DATABASE_URL"));
|
||||
}
|
||||
}
|
||||
|
||||
+31
-9
@@ -33,9 +33,16 @@ use termimad::MadSkin;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use crate::agent::truncate_for_preview;
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use crate::error::ChannelError;
|
||||
|
||||
/// Max characters for tool result previews in the terminal.
|
||||
const CLI_TOOL_RESULT_MAX: usize = 200;
|
||||
|
||||
/// Max characters for thinking/status messages in the terminal.
|
||||
const CLI_STATUS_MAX: usize = 200;
|
||||
|
||||
/// Slash commands available in the REPL.
|
||||
const SLASH_COMMANDS: &[&str] = &[
|
||||
"/help",
|
||||
@@ -177,6 +184,8 @@ pub struct ReplChannel {
|
||||
debug_mode: Arc<AtomicBool>,
|
||||
/// Whether we're currently streaming (chunks have been printed without a trailing newline).
|
||||
is_streaming: Arc<AtomicBool>,
|
||||
/// When true, the one-liner startup banner is suppressed (boot screen shown instead).
|
||||
suppress_banner: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ReplChannel {
|
||||
@@ -186,6 +195,7 @@ impl ReplChannel {
|
||||
single_message: None,
|
||||
debug_mode: Arc::new(AtomicBool::new(false)),
|
||||
is_streaming: Arc::new(AtomicBool::new(false)),
|
||||
suppress_banner: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,9 +205,15 @@ impl ReplChannel {
|
||||
single_message: Some(message),
|
||||
debug_mode: Arc::new(AtomicBool::new(false)),
|
||||
is_streaming: Arc::new(AtomicBool::new(false)),
|
||||
suppress_banner: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Suppress the one-liner startup banner (boot screen will be shown instead).
|
||||
pub fn suppress_banner(&self) {
|
||||
self.suppress_banner.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn is_debug(&self) -> bool {
|
||||
self.debug_mode.load(Ordering::Relaxed)
|
||||
}
|
||||
@@ -257,11 +273,12 @@ impl Channel for ReplChannel {
|
||||
let (tx, rx) = mpsc::channel(32);
|
||||
let single_message = self.single_message.clone();
|
||||
let debug_mode = Arc::clone(&self.debug_mode);
|
||||
let suppress_banner = Arc::clone(&self.suppress_banner);
|
||||
|
||||
std::thread::spawn(move || {
|
||||
// Single message mode: send it and return
|
||||
if let Some(msg) = single_message {
|
||||
let incoming = IncomingMessage::new("repl", "user", &msg);
|
||||
let incoming = IncomingMessage::new("repl", "default", &msg);
|
||||
let _ = tx.blocking_send(incoming);
|
||||
return;
|
||||
}
|
||||
@@ -291,8 +308,10 @@ impl Channel for ReplChannel {
|
||||
}
|
||||
let _ = rl.load_history(&hist_path);
|
||||
|
||||
println!("\x1b[1mIronClaw\x1b[0m /help for commands, /quit to exit");
|
||||
println!();
|
||||
if !suppress_banner.load(Ordering::Relaxed) {
|
||||
println!("\x1b[1mIronClaw\x1b[0m /help for commands, /quit to exit");
|
||||
println!();
|
||||
}
|
||||
|
||||
loop {
|
||||
let prompt = if debug_mode.load(Ordering::Relaxed) {
|
||||
@@ -329,21 +348,21 @@ impl Channel for ReplChannel {
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let msg = IncomingMessage::new("repl", "user", line);
|
||||
let msg = IncomingMessage::new("repl", "default", line);
|
||||
if tx.blocking_send(msg).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(ReadlineError::Interrupted) => {
|
||||
// Ctrl+C: send /interrupt
|
||||
let msg = IncomingMessage::new("repl", "user", "/interrupt");
|
||||
let msg = IncomingMessage::new("repl", "default", "/interrupt");
|
||||
if tx.blocking_send(msg).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(ReadlineError::Eof) => {
|
||||
// Ctrl+D: send /quit so the agent loop runs graceful shutdown
|
||||
let msg = IncomingMessage::new("repl", "user", "/quit");
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||
let _ = tx.blocking_send(msg);
|
||||
break;
|
||||
}
|
||||
@@ -400,7 +419,8 @@ impl Channel for ReplChannel {
|
||||
|
||||
match status {
|
||||
StatusUpdate::Thinking(msg) => {
|
||||
eprintln!(" \x1b[90m\u{25CB} {msg}\x1b[0m");
|
||||
let display = truncate_for_preview(&msg, CLI_STATUS_MAX);
|
||||
eprintln!(" \x1b[90m\u{25CB} {display}\x1b[0m");
|
||||
}
|
||||
StatusUpdate::ToolStarted { name } => {
|
||||
eprintln!(" \x1b[33m\u{25CB} {name}\x1b[0m");
|
||||
@@ -413,7 +433,8 @@ impl Channel for ReplChannel {
|
||||
}
|
||||
}
|
||||
StatusUpdate::ToolResult { name: _, preview } => {
|
||||
eprintln!(" \x1b[90m{preview}\x1b[0m");
|
||||
let display = truncate_for_preview(&preview, CLI_TOOL_RESULT_MAX);
|
||||
eprintln!(" \x1b[90m{display}\x1b[0m");
|
||||
}
|
||||
StatusUpdate::StreamChunk(chunk) => {
|
||||
// Print separator on the false-to-true transition
|
||||
@@ -438,7 +459,8 @@ impl Channel for ReplChannel {
|
||||
}
|
||||
StatusUpdate::Status(msg) => {
|
||||
if debug || msg.contains("approval") || msg.contains("Approval") {
|
||||
eprintln!(" \x1b[90m{msg}\x1b[0m");
|
||||
let display = truncate_for_preview(&msg, CLI_STATUS_MAX);
|
||||
eprintln!(" \x1b[90m{display}\x1b[0m");
|
||||
}
|
||||
}
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
|
||||
+135
-38
@@ -76,6 +76,9 @@ struct ChannelStoreData {
|
||||
credentials: HashMap<String, String>,
|
||||
/// Pairing store for DM pairing (guest access control).
|
||||
pairing_store: Arc<PairingStore>,
|
||||
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
||||
/// Reused across multiple `http_request` calls within one execution.
|
||||
http_runtime: Option<tokio::runtime::Runtime>,
|
||||
}
|
||||
|
||||
impl ChannelStoreData {
|
||||
@@ -96,6 +99,7 @@ impl ChannelStoreData {
|
||||
table: ResourceTable::new(),
|
||||
credentials,
|
||||
pairing_store,
|
||||
http_runtime: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -134,13 +138,13 @@ impl ChannelStoreData {
|
||||
if result.contains('{') && result.contains('}') {
|
||||
// Only warn if it looks like an unresolved placeholder (not JSON braces)
|
||||
let brace_pattern = regex::Regex::new(r"\{[A-Z_]+\}").ok();
|
||||
if let Some(re) = brace_pattern {
|
||||
if re.is_match(&result) {
|
||||
tracing::warn!(
|
||||
context = %context,
|
||||
"String may contain unresolved credential placeholders"
|
||||
);
|
||||
}
|
||||
if let Some(re) = brace_pattern
|
||||
&& re.is_match(&result)
|
||||
{
|
||||
tracing::warn!(
|
||||
context = %context,
|
||||
"String may contain unresolved credential placeholders"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -273,10 +277,35 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
.scan_http_request(&url, &header_vec, body.as_deref())
|
||||
.map_err(|e| format!("Potential secret leak blocked: {}", e))?;
|
||||
|
||||
// Make the HTTP request using blocking I/O
|
||||
// We're already in a spawn_blocking context, so we can use block_on
|
||||
let result = tokio::runtime::Handle::current().block_on(async {
|
||||
let client = reqwest::Client::new();
|
||||
// Get the max response size from capabilities (default 10MB).
|
||||
let max_response_bytes = self
|
||||
.host_state
|
||||
.capabilities()
|
||||
.tool_capabilities
|
||||
.http
|
||||
.as_ref()
|
||||
.map(|h| h.max_response_bytes)
|
||||
.unwrap_or(10 * 1024 * 1024);
|
||||
|
||||
// Make the HTTP request using a dedicated single-threaded runtime.
|
||||
// We're inside spawn_blocking, so we can't rely on the main runtime's
|
||||
// I/O driver (it may be busy with WASM compilation or other startup work).
|
||||
// A dedicated runtime gives us our own I/O driver and avoids contention.
|
||||
// The runtime is lazily created and reused across calls within one execution.
|
||||
if self.http_runtime.is_none() {
|
||||
self.http_runtime = Some(
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to create HTTP runtime: {e}"))?,
|
||||
);
|
||||
}
|
||||
let rt = self.http_runtime.as_ref().expect("just initialized");
|
||||
let result = rt.block_on(async {
|
||||
let client = reqwest::Client::builder()
|
||||
.connect_timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to build HTTP client: {e}"))?;
|
||||
|
||||
let mut request = match method.to_uppercase().as_str() {
|
||||
"GET" => client.get(&url),
|
||||
@@ -298,9 +327,9 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
request = request.body(body_bytes);
|
||||
}
|
||||
|
||||
// Send request with caller-specified timeout (default 30s).
|
||||
// Cap at callback_timeout to prevent outliving the host wrapper.
|
||||
let timeout = std::time::Duration::from_millis(timeout_ms.unwrap_or(30_000) as u64);
|
||||
// Send request with caller-specified timeout (default 30s, max 5min).
|
||||
let timeout_ms = timeout_ms.unwrap_or(30_000).min(300_000) as u64;
|
||||
let timeout = std::time::Duration::from_millis(timeout_ms);
|
||||
let response = request.timeout(timeout).send().await.map_err(|e| {
|
||||
// Walk the full error chain so we get the actual root cause
|
||||
// (DNS, TLS, connection refused, etc.) instead of just
|
||||
@@ -325,11 +354,29 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
})
|
||||
.collect();
|
||||
let headers_json = serde_json::to_string(&response_headers).unwrap_or_default();
|
||||
|
||||
// Enforce max response body size to prevent memory exhaustion.
|
||||
let max_response = max_response_bytes;
|
||||
if let Some(cl) = response.content_length()
|
||||
&& cl as usize > max_response
|
||||
{
|
||||
return Err(format!(
|
||||
"Response body too large: {} bytes exceeds limit of {} bytes",
|
||||
cl, max_response
|
||||
));
|
||||
}
|
||||
let body = response
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to read response body: {}", e))?
|
||||
.to_vec();
|
||||
.map_err(|e| format!("Failed to read response body: {}", e))?;
|
||||
if body.len() > max_response {
|
||||
return Err(format!(
|
||||
"Response body too large: {} bytes exceeds limit of {} bytes",
|
||||
body.len(),
|
||||
max_response
|
||||
));
|
||||
}
|
||||
let body = body.to_vec();
|
||||
|
||||
tracing::info!(
|
||||
status = status,
|
||||
@@ -767,7 +814,21 @@ impl WasmChannel {
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok((config, _host_state))) => {
|
||||
Ok(Ok((config, mut host_state))) => {
|
||||
// Surface WASM guest logs (errors/warnings from webhook setup, etc.)
|
||||
for entry in host_state.take_logs() {
|
||||
match entry.level {
|
||||
crate::tools::wasm::LogLevel::Error => {
|
||||
tracing::error!(channel = %self.name, "{}", entry.message);
|
||||
}
|
||||
crate::tools::wasm::LogLevel::Warn => {
|
||||
tracing::warn!(channel = %self.name, "{}", entry.message);
|
||||
}
|
||||
_ => {
|
||||
tracing::debug!(channel = %self.name, "{}", entry.message);
|
||||
}
|
||||
}
|
||||
}
|
||||
tracing::info!(
|
||||
channel = %self.name,
|
||||
display_name = %config.display_name,
|
||||
@@ -1467,8 +1528,8 @@ impl WasmChannel {
|
||||
match result {
|
||||
Ok(emitted_messages) => {
|
||||
// Process any emitted messages
|
||||
if !emitted_messages.is_empty() {
|
||||
if let Err(e) = Self::dispatch_emitted_messages(
|
||||
if !emitted_messages.is_empty()
|
||||
&& let Err(e) = Self::dispatch_emitted_messages(
|
||||
&channel_name,
|
||||
emitted_messages,
|
||||
&message_tx,
|
||||
@@ -1480,7 +1541,6 @@ impl WasmChannel {
|
||||
"Failed to dispatch emitted messages from poll"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
@@ -1710,22 +1770,22 @@ impl Channel for WasmChannel {
|
||||
*self.endpoints.write().await = endpoints;
|
||||
|
||||
// Start polling if configured
|
||||
if let Some(poll_config) = &config.poll {
|
||||
if poll_config.enabled {
|
||||
let interval = self
|
||||
.capabilities
|
||||
.validate_poll_interval(poll_config.interval_ms)
|
||||
.map_err(|e| ChannelError::StartupFailed {
|
||||
name: self.name.clone(),
|
||||
reason: e,
|
||||
})?;
|
||||
if let Some(poll_config) = &config.poll
|
||||
&& poll_config.enabled
|
||||
{
|
||||
let interval = self
|
||||
.capabilities
|
||||
.validate_poll_interval(poll_config.interval_ms)
|
||||
.map_err(|e| ChannelError::StartupFailed {
|
||||
name: self.name.clone(),
|
||||
reason: e,
|
||||
})?;
|
||||
|
||||
// Create shutdown channel for polling and store the sender to keep it alive
|
||||
let (poll_shutdown_tx, poll_shutdown_rx) = oneshot::channel();
|
||||
*self.poll_shutdown_tx.write().await = Some(poll_shutdown_tx);
|
||||
// Create shutdown channel for polling and store the sender to keep it alive
|
||||
let (poll_shutdown_tx, poll_shutdown_rx) = oneshot::channel();
|
||||
*self.poll_shutdown_tx.write().await = Some(poll_shutdown_tx);
|
||||
|
||||
self.start_polling(Duration::from_millis(interval as u64), poll_shutdown_rx);
|
||||
}
|
||||
self.start_polling(Duration::from_millis(interval as u64), poll_shutdown_rx);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
@@ -2588,15 +2648,52 @@ mod tests {
|
||||
assert_eq!(store.redact_credentials(input), input);
|
||||
}
|
||||
|
||||
/// Verify that the block_on-inside-spawn_blocking pattern used by the WASM
|
||||
/// channel HTTP host function doesn't deadlock or panic.
|
||||
/// Verify that WASM HTTP host functions work using a dedicated
|
||||
/// current-thread runtime inside spawn_blocking.
|
||||
#[tokio::test]
|
||||
async fn test_block_on_inside_spawn_blocking_does_not_deadlock() {
|
||||
async fn test_dedicated_runtime_inside_spawn_blocking() {
|
||||
let result = tokio::task::spawn_blocking(|| {
|
||||
tokio::runtime::Handle::current().block_on(async { 42 })
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("failed to build runtime");
|
||||
rt.block_on(async { 42 })
|
||||
})
|
||||
.await
|
||||
.expect("spawn_blocking panicked");
|
||||
assert_eq!(result, 42);
|
||||
}
|
||||
|
||||
/// Verify a real HTTP request works using the dedicated-runtime pattern.
|
||||
/// This catches DNS, TLS, and I/O driver issues that trivial tests miss.
|
||||
#[tokio::test]
|
||||
#[ignore] // requires network
|
||||
async fn test_dedicated_runtime_real_http() {
|
||||
let result = tokio::task::spawn_blocking(|| {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("failed to build runtime");
|
||||
rt.block_on(async {
|
||||
let client = reqwest::Client::builder()
|
||||
.connect_timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.expect("failed to build client");
|
||||
let resp = client
|
||||
.get("https://api.telegram.org/bot000/getMe")
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
match resp {
|
||||
Ok(r) => r.status().as_u16(),
|
||||
Err(e) if e.is_timeout() => panic!("request timed out: {e}"),
|
||||
Err(e) => panic!("unexpected error: {e}"),
|
||||
}
|
||||
})
|
||||
})
|
||||
.await
|
||||
.expect("spawn_blocking panicked");
|
||||
// 404 because "000" is not a valid bot token
|
||||
assert_eq!(result, 404);
|
||||
}
|
||||
}
|
||||
|
||||
+13
-14
@@ -6,6 +6,7 @@ use axum::{
|
||||
middleware::Next,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
/// Shared auth state injected via axum middleware state.
|
||||
#[derive(Clone)]
|
||||
@@ -23,24 +24,22 @@ pub async fn auth_middleware(
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
// Try Authorization header first
|
||||
if let Some(auth_header) = headers.get("authorization") {
|
||||
if let Ok(value) = auth_header.to_str() {
|
||||
if let Some(token) = value.strip_prefix("Bearer ") {
|
||||
if token == auth.token {
|
||||
return next.run(request).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Try Authorization header first (constant-time comparison)
|
||||
if let Some(auth_header) = headers.get("authorization")
|
||||
&& let Ok(value) = auth_header.to_str()
|
||||
&& let Some(token) = value.strip_prefix("Bearer ")
|
||||
&& bool::from(token.as_bytes().ct_eq(auth.token.as_bytes()))
|
||||
{
|
||||
return next.run(request).await;
|
||||
}
|
||||
|
||||
// Fall back to query parameter (for SSE EventSource)
|
||||
// Fall back to query parameter for SSE EventSource (constant-time comparison)
|
||||
if let Some(query) = request.uri().query() {
|
||||
for pair in query.split('&') {
|
||||
if let Some(token) = pair.strip_prefix("token=") {
|
||||
if token == auth.token {
|
||||
return next.run(request).await;
|
||||
}
|
||||
if let Some(token) = pair.strip_prefix("token=")
|
||||
&& bool::from(token.as_bytes().ct_eq(auth.token.as_bytes()))
|
||||
{
|
||||
return next.run(request).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,8 @@ use tokio::sync::broadcast;
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing_subscriber::Layer;
|
||||
|
||||
use crate::safety::LeakDetector;
|
||||
|
||||
/// Maximum number of recent log entries kept for late-joining SSE subscribers.
|
||||
const HISTORY_CAP: usize = 500;
|
||||
|
||||
@@ -46,6 +48,8 @@ pub struct LogEntry {
|
||||
pub struct LogBroadcaster {
|
||||
tx: broadcast::Sender<LogEntry>,
|
||||
recent: Mutex<VecDeque<LogEntry>>,
|
||||
/// Scrubs secrets from log messages before broadcasting to SSE clients.
|
||||
leak_detector: LeakDetector,
|
||||
}
|
||||
|
||||
impl LogBroadcaster {
|
||||
@@ -54,10 +58,19 @@ impl LogBroadcaster {
|
||||
Self {
|
||||
tx,
|
||||
recent: Mutex::new(VecDeque::with_capacity(HISTORY_CAP)),
|
||||
leak_detector: LeakDetector::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn send(&self, entry: LogEntry) {
|
||||
pub fn send(&self, mut entry: LogEntry) {
|
||||
// Scrub secrets from the message before it reaches any subscriber.
|
||||
// This is defense-in-depth: even if code elsewhere accidentally logs
|
||||
// a secret, it won't be broadcast to SSE clients.
|
||||
entry.message = self
|
||||
.leak_detector
|
||||
.scan_and_clean(&entry.message)
|
||||
.unwrap_or_else(|_| "[log message redacted: contained blocked secret]".to_string());
|
||||
|
||||
// Stash in ring buffer (for late joiners)
|
||||
if let Ok(mut buf) = self.recent.lock() {
|
||||
if buf.len() >= HISTORY_CAP {
|
||||
@@ -145,6 +158,9 @@ impl Visit for MessageVisitor {
|
||||
///
|
||||
/// Only forwards DEBUG and above. Attach to the tracing subscriber
|
||||
/// alongside the existing fmt layer.
|
||||
///
|
||||
/// Log messages are scrubbed through `LeakDetector` in `LogBroadcaster::send()`
|
||||
/// (the single funnel point for all log output, including late-joiner history).
|
||||
pub struct WebLogLayer {
|
||||
broadcaster: Arc<LogBroadcaster>,
|
||||
}
|
||||
@@ -178,6 +194,7 @@ impl<S: tracing::Subscriber> Layer<S> for WebLogLayer {
|
||||
timestamp: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
|
||||
};
|
||||
|
||||
// LeakDetector scrubbing happens inside broadcaster.send()
|
||||
self.broadcaster.send(entry);
|
||||
}
|
||||
}
|
||||
@@ -313,4 +330,29 @@ mod tests {
|
||||
let v = MessageVisitor::new();
|
||||
assert_eq!(v.finish(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_broadcaster_has_leak_detector() {
|
||||
let broadcaster = LogBroadcaster::new();
|
||||
// Verify the leak detector is initialized with default patterns
|
||||
assert!(broadcaster.leak_detector.pattern_count() > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_leak_detector_scrubs_api_key_in_log() {
|
||||
let detector = crate::safety::LeakDetector::new();
|
||||
let msg = "Connecting with token sk-proj-test1234567890abcdefghij";
|
||||
let result = detector.scan_and_clean(msg);
|
||||
// Should be blocked (OpenAI key pattern)
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_leak_detector_passes_clean_log() {
|
||||
let detector = crate::safety::LeakDetector::new();
|
||||
let msg = "Request completed status=200 url=https://api.example.com/data";
|
||||
let result = detector.scan_and_clean(msg);
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap(), msg);
|
||||
}
|
||||
}
|
||||
|
||||
+13
-2
@@ -16,6 +16,7 @@
|
||||
|
||||
pub mod auth;
|
||||
pub mod log_layer;
|
||||
pub mod openai_compat;
|
||||
pub mod server;
|
||||
pub mod sse;
|
||||
pub mod types;
|
||||
@@ -31,9 +32,9 @@ use tokio_stream::wrappers::ReceiverStream;
|
||||
use crate::agent::SessionManager;
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use crate::config::GatewayConfig;
|
||||
use crate::db::Database;
|
||||
use crate::error::ChannelError;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
@@ -81,6 +82,8 @@ impl GatewayChannel {
|
||||
user_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
chat_rate_limiter: server::RateLimiter::new(30, 60),
|
||||
});
|
||||
|
||||
Self {
|
||||
@@ -106,6 +109,8 @@ impl GatewayChannel {
|
||||
user_id: self.state.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: self.state.ws_tracker.clone(),
|
||||
llm_provider: self.state.llm_provider.clone(),
|
||||
chat_rate_limiter: server::RateLimiter::new(30, 60),
|
||||
};
|
||||
mutate(&mut new_state);
|
||||
self.state = Arc::new(new_state);
|
||||
@@ -142,7 +147,7 @@ impl GatewayChannel {
|
||||
}
|
||||
|
||||
/// Inject the database store for sandbox job persistence.
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
self.rebuild_state(|s| s.store = Some(store));
|
||||
self
|
||||
}
|
||||
@@ -169,6 +174,12 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the LLM provider for OpenAI-compatible API proxy.
|
||||
pub fn with_llm_provider(mut self, llm: Arc<dyn crate::llm::LlmProvider>) -> Self {
|
||||
self.rebuild_state(|s| s.llm_provider = Some(llm));
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the auth token (for printing to console on startup).
|
||||
pub fn auth_token(&self) -> &str {
|
||||
&self.auth_token
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+367
-153
@@ -5,10 +5,11 @@
|
||||
use std::convert::Infallible;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, Query, State, WebSocketUpgrade},
|
||||
extract::{DefaultBodyLimit, Path, Query, State, WebSocketUpgrade},
|
||||
http::{StatusCode, header},
|
||||
middleware,
|
||||
response::{
|
||||
@@ -20,6 +21,7 @@ use axum::{
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio_stream::StreamExt;
|
||||
use tower_http::cors::{AllowHeaders, CorsLayer};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::SessionManager;
|
||||
@@ -28,8 +30,8 @@ use crate::channels::web::auth::{AuthState, auth_middleware};
|
||||
use crate::channels::web::log_layer::LogBroadcaster;
|
||||
use crate::channels::web::sse::SseManager;
|
||||
use crate::channels::web::types::*;
|
||||
use crate::db::Database;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
@@ -44,6 +46,69 @@ pub type PromptQueue = Arc<
|
||||
>,
|
||||
>;
|
||||
|
||||
/// Simple sliding-window rate limiter.
|
||||
///
|
||||
/// Tracks the number of requests in the current window. Resets when the window expires.
|
||||
/// Not per-IP (since this is a single-user gateway with auth), but prevents flooding.
|
||||
pub struct RateLimiter {
|
||||
/// Requests remaining in the current window.
|
||||
remaining: AtomicU64,
|
||||
/// Epoch second when the current window started.
|
||||
window_start: AtomicU64,
|
||||
/// Maximum requests per window.
|
||||
max_requests: u64,
|
||||
/// Window duration in seconds.
|
||||
window_secs: u64,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
pub fn new(max_requests: u64, window_secs: u64) -> Self {
|
||||
Self {
|
||||
remaining: AtomicU64::new(max_requests),
|
||||
window_start: AtomicU64::new(
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
),
|
||||
max_requests,
|
||||
window_secs,
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to consume one request. Returns `true` if allowed, `false` if rate limited.
|
||||
pub fn check(&self) -> bool {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
let window = self.window_start.load(Ordering::Relaxed);
|
||||
if now.saturating_sub(window) >= self.window_secs {
|
||||
// Window expired, reset
|
||||
self.window_start.store(now, Ordering::Relaxed);
|
||||
self.remaining
|
||||
.store(self.max_requests - 1, Ordering::Relaxed);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Try to decrement remaining
|
||||
loop {
|
||||
let current = self.remaining.load(Ordering::Relaxed);
|
||||
if current == 0 {
|
||||
return false;
|
||||
}
|
||||
if self
|
||||
.remaining
|
||||
.compare_exchange_weak(current, current - 1, Ordering::Relaxed, Ordering::Relaxed)
|
||||
.is_ok()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared state for all gateway handlers.
|
||||
pub struct GatewayState {
|
||||
/// Channel to send messages to the agent loop.
|
||||
@@ -61,7 +126,7 @@ pub struct GatewayState {
|
||||
/// Tool registry for listing registered tools.
|
||||
pub tool_registry: Option<Arc<ToolRegistry>>,
|
||||
/// Database store for sandbox job persistence.
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
/// Container job manager for sandbox operations.
|
||||
pub job_manager: Option<Arc<ContainerJobManager>>,
|
||||
/// Prompt queue for Claude Code follow-up prompts.
|
||||
@@ -72,6 +137,10 @@ pub struct GatewayState {
|
||||
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
|
||||
/// WebSocket connection tracker.
|
||||
pub ws_tracker: Option<Arc<crate::channels::web::ws::WsConnectionTracker>>,
|
||||
/// LLM provider for OpenAI-compatible API proxy.
|
||||
pub llm_provider: Option<Arc<dyn crate::llm::LlmProvider>>,
|
||||
/// Rate limiter for chat endpoints (30 messages per 60 seconds).
|
||||
pub chat_rate_limiter: RateLimiter,
|
||||
}
|
||||
|
||||
/// Start the gateway HTTP server.
|
||||
@@ -168,7 +237,16 @@ pub async fn start_server(
|
||||
)
|
||||
// Gateway control plane
|
||||
.route("/api/gateway/status", get(gateway_status_handler))
|
||||
.route_layer(middleware::from_fn_with_state(auth_state, auth_middleware));
|
||||
// OpenAI-compatible API
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
post(super::openai_compat::chat_completions_handler),
|
||||
)
|
||||
.route("/v1/models", get(super::openai_compat::models_handler))
|
||||
.route_layer(middleware::from_fn_with_state(
|
||||
auth_state.clone(),
|
||||
auth_middleware,
|
||||
));
|
||||
|
||||
// Static file routes (no auth, served from embedded strings)
|
||||
let statics = Router::new()
|
||||
@@ -176,19 +254,46 @@ pub async fn start_server(
|
||||
.route("/style.css", get(css_handler))
|
||||
.route("/app.js", get(js_handler));
|
||||
|
||||
// Project file serving (no auth, local browsing of sandbox outputs).
|
||||
// The trailing-slash route serves index.html; the bare route redirects so
|
||||
// relative paths in the HTML (e.g. href="style.css") resolve correctly.
|
||||
// Project file serving (behind auth to prevent unauthorized file access).
|
||||
let projects = Router::new()
|
||||
.route("/projects/{project_id}", get(project_redirect_handler))
|
||||
.route("/projects/{project_id}/", get(project_index_handler))
|
||||
.route("/projects/{project_id}/{*path}", get(project_file_handler));
|
||||
.route("/projects/{project_id}/{*path}", get(project_file_handler))
|
||||
.route_layer(middleware::from_fn_with_state(
|
||||
auth_state.clone(),
|
||||
auth_middleware,
|
||||
));
|
||||
|
||||
// CORS: restrict to same-origin by default. Only localhost/127.0.0.1
|
||||
// origins are allowed, since the gateway is a local-first service.
|
||||
let cors = CorsLayer::new()
|
||||
.allow_origin([
|
||||
format!("http://{}:{}", addr.ip(), addr.port())
|
||||
.parse()
|
||||
.expect("valid origin"),
|
||||
format!("http://localhost:{}", addr.port())
|
||||
.parse()
|
||||
.expect("valid origin"),
|
||||
])
|
||||
.allow_methods([
|
||||
axum::http::Method::GET,
|
||||
axum::http::Method::POST,
|
||||
axum::http::Method::PUT,
|
||||
axum::http::Method::DELETE,
|
||||
])
|
||||
.allow_headers(AllowHeaders::list([
|
||||
header::CONTENT_TYPE,
|
||||
header::AUTHORIZATION,
|
||||
]))
|
||||
.allow_credentials(true);
|
||||
|
||||
let app = Router::new()
|
||||
.merge(public)
|
||||
.merge(statics)
|
||||
.merge(projects)
|
||||
.merge(protected)
|
||||
.layer(cors)
|
||||
.layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body
|
||||
.with_state(state.clone());
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
@@ -244,6 +349,13 @@ async fn chat_send_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Json(req): Json<SendMessageRequest>,
|
||||
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
||||
if !state.chat_rate_limiter.check() {
|
||||
return Err((
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
"Rate limit exceeded. Try again shortly.".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content);
|
||||
|
||||
if let Some(ref thread_id) = req.thread_id {
|
||||
@@ -413,24 +525,58 @@ pub async fn clear_auth_mode(state: &GatewayState) {
|
||||
if let Some(ref sm) = state.session_manager {
|
||||
let session = sm.get_or_create_session(&state.user_id).await;
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread_id) = sess.active_thread {
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
if let Some(thread_id) = sess.active_thread
|
||||
&& let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
{
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn chat_events_handler(State(state): State<Arc<GatewayState>>) -> impl IntoResponse {
|
||||
// subscribe() returns Sse<impl Stream + 'static + use<>> so no lifetime issues
|
||||
state.sse.subscribe()
|
||||
async fn chat_events_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
state.sse.subscribe().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Too many connections".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn chat_ws_handler(
|
||||
headers: axum::http::HeaderMap,
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> impl IntoResponse {
|
||||
ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state))
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
// Validate Origin header to prevent cross-site WebSocket hijacking.
|
||||
// Require the header outright; browsers always send it for WS upgrades,
|
||||
// so a missing Origin means a non-browser client trying to bypass the check.
|
||||
let origin = headers
|
||||
.get("origin")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
(
|
||||
StatusCode::FORBIDDEN,
|
||||
"WebSocket Origin header required".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Extract the host from the origin and compare exactly, so that
|
||||
// crafted origins like "http://localhost.evil.com" are rejected.
|
||||
// Origin format is "scheme://host[:port]".
|
||||
let host = origin
|
||||
.strip_prefix("http://")
|
||||
.or_else(|| origin.strip_prefix("https://"))
|
||||
.and_then(|rest| rest.split(':').next()?.split('/').next())
|
||||
.unwrap_or("");
|
||||
|
||||
let is_local = matches!(host, "localhost" | "127.0.0.1" | "[::1]");
|
||||
if !is_local {
|
||||
return Err((
|
||||
StatusCode::FORBIDDEN,
|
||||
"WebSocket origin not allowed".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -477,57 +623,72 @@ async fn chat_history_handler(
|
||||
.ok_or((StatusCode::NOT_FOUND, "No active thread".to_string()))?
|
||||
};
|
||||
|
||||
// For paginated requests (before cursor set), always go to DB
|
||||
if before_cursor.is_some() {
|
||||
if let Some(ref store) = state.store {
|
||||
let (messages, has_more) = store
|
||||
.list_conversation_messages_paginated(thread_id, before_cursor, limit as i64)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
|
||||
let turns = build_turns_from_db_messages(&messages);
|
||||
return Ok(Json(HistoryResponse {
|
||||
thread_id,
|
||||
turns,
|
||||
has_more,
|
||||
oldest_timestamp,
|
||||
}));
|
||||
// Verify the thread belongs to the authenticated user before returning any data.
|
||||
// In-memory threads are already scoped by user via session_manager, but DB
|
||||
// lookups could expose another user's conversation if the UUID is guessed.
|
||||
if query.thread_id.is_some()
|
||||
&& let Some(ref store) = state.store
|
||||
{
|
||||
let owned = store
|
||||
.conversation_belongs_to_user(thread_id, &state.user_id)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
if !owned && !sess.threads.contains_key(&thread_id) {
|
||||
return Err((StatusCode::NOT_FOUND, "Thread not found".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// Try in-memory first (freshest data for active threads)
|
||||
if let Some(thread) = sess.threads.get(&thread_id) {
|
||||
if !thread.turns.is_empty() {
|
||||
let turns: Vec<TurnInfo> = thread
|
||||
.turns
|
||||
.iter()
|
||||
.map(|t| TurnInfo {
|
||||
turn_number: t.turn_number,
|
||||
user_input: t.user_input.clone(),
|
||||
response: t.response.clone(),
|
||||
state: format!("{:?}", t.state),
|
||||
started_at: t.started_at.to_rfc3339(),
|
||||
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
tool_calls: t
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| ToolCallInfo {
|
||||
name: tc.name.clone(),
|
||||
has_result: tc.result.is_some(),
|
||||
has_error: tc.error.is_some(),
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect();
|
||||
// For paginated requests (before cursor set), always go to DB
|
||||
if before_cursor.is_some()
|
||||
&& let Some(ref store) = state.store
|
||||
{
|
||||
let (messages, has_more) = store
|
||||
.list_conversation_messages_paginated(thread_id, before_cursor, limit as i64)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
return Ok(Json(HistoryResponse {
|
||||
thread_id,
|
||||
turns,
|
||||
has_more: false,
|
||||
oldest_timestamp: None,
|
||||
}));
|
||||
}
|
||||
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
|
||||
let turns = build_turns_from_db_messages(&messages);
|
||||
return Ok(Json(HistoryResponse {
|
||||
thread_id,
|
||||
turns,
|
||||
has_more,
|
||||
oldest_timestamp,
|
||||
}));
|
||||
}
|
||||
|
||||
// Try in-memory first (freshest data for active threads)
|
||||
if let Some(thread) = sess.threads.get(&thread_id)
|
||||
&& !thread.turns.is_empty()
|
||||
{
|
||||
let turns: Vec<TurnInfo> = thread
|
||||
.turns
|
||||
.iter()
|
||||
.map(|t| TurnInfo {
|
||||
turn_number: t.turn_number,
|
||||
user_input: t.user_input.clone(),
|
||||
response: t.response.clone(),
|
||||
state: format!("{:?}", t.state),
|
||||
started_at: t.started_at.to_rfc3339(),
|
||||
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
tool_calls: t
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| ToolCallInfo {
|
||||
name: tc.name.clone(),
|
||||
has_result: tc.result.is_some(),
|
||||
has_error: tc.error.is_some(),
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
return Ok(Json(HistoryResponse {
|
||||
thread_id,
|
||||
turns,
|
||||
has_more: false,
|
||||
oldest_timestamp: None,
|
||||
}));
|
||||
}
|
||||
|
||||
// Fall back to DB for historical threads not in memory (paginated)
|
||||
@@ -577,12 +738,12 @@ fn build_turns_from_db_messages(messages: &[crate::history::ConversationMessage]
|
||||
};
|
||||
|
||||
// Check if next message is an assistant response
|
||||
if let Some(next) = iter.peek() {
|
||||
if next.role == "assistant" {
|
||||
let assistant_msg = iter.next().expect("peeked");
|
||||
turn.response = Some(assistant_msg.content.clone());
|
||||
turn.completed_at = Some(assistant_msg.created_at.to_rfc3339());
|
||||
}
|
||||
if let Some(next) = iter.peek()
|
||||
&& next.role == "assistant"
|
||||
{
|
||||
let assistant_msg = iter.next().expect("peeked");
|
||||
turn.response = Some(assistant_msg.content.clone());
|
||||
turn.completed_at = Some(assistant_msg.created_at.to_rfc3339());
|
||||
}
|
||||
|
||||
// Incomplete turn (user message without response)
|
||||
@@ -901,14 +1062,16 @@ async fn jobs_list_handler(
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
// Fetch sandbox jobs from the DB.
|
||||
// Fetch sandbox jobs scoped to the authenticated user.
|
||||
let sandbox_jobs = store
|
||||
.list_sandbox_jobs()
|
||||
.list_sandbox_jobs_for_user(&state.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
// Scope jobs to the authenticated user.
|
||||
let mut jobs: Vec<JobInfo> = sandbox_jobs
|
||||
.iter()
|
||||
.filter(|j| j.user_id == state.user_id)
|
||||
.map(|j| {
|
||||
let ui_state = match j.status.as_str() {
|
||||
"creating" => "pending",
|
||||
@@ -941,7 +1104,7 @@ async fn jobs_summary_handler(
|
||||
))?;
|
||||
|
||||
let s = store
|
||||
.sandbox_job_summary()
|
||||
.sandbox_job_summary_for_user(&state.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
@@ -962,63 +1125,66 @@ async fn jobs_detail_handler(
|
||||
let job_id = Uuid::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||
|
||||
// Try sandbox job from DB first.
|
||||
if let Some(ref store) = state.store {
|
||||
if let Ok(Some(job)) = store.get_sandbox_job(job_id).await {
|
||||
let browse_id = std::path::Path::new(&job.project_dir)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|| job.id.to_string());
|
||||
|
||||
let ui_state = match job.status.as_str() {
|
||||
"creating" => "pending",
|
||||
"running" => "in_progress",
|
||||
s => s,
|
||||
};
|
||||
|
||||
let elapsed_secs = job.started_at.map(|start| {
|
||||
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
|
||||
(end - start).num_seconds().max(0) as u64
|
||||
});
|
||||
|
||||
// Synthesize transitions from timestamps.
|
||||
let mut transitions = Vec::new();
|
||||
if let Some(started) = job.started_at {
|
||||
transitions.push(TransitionInfo {
|
||||
from: "creating".to_string(),
|
||||
to: "running".to_string(),
|
||||
timestamp: started.to_rfc3339(),
|
||||
reason: None,
|
||||
});
|
||||
}
|
||||
if let Some(completed) = job.completed_at {
|
||||
transitions.push(TransitionInfo {
|
||||
from: "running".to_string(),
|
||||
to: job.status.clone(),
|
||||
timestamp: completed.to_rfc3339(),
|
||||
reason: job.failure_reason.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
return Ok(Json(JobDetailResponse {
|
||||
id: job.id,
|
||||
title: job.task.clone(),
|
||||
description: String::new(),
|
||||
state: ui_state.to_string(),
|
||||
user_id: job.user_id.clone(),
|
||||
created_at: job.created_at.to_rfc3339(),
|
||||
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
|
||||
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
elapsed_secs,
|
||||
project_dir: Some(job.project_dir.clone()),
|
||||
browse_url: Some(format!("/projects/{}/", browse_id)),
|
||||
job_mode: {
|
||||
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
|
||||
mode.filter(|m| m != "worker")
|
||||
},
|
||||
transitions,
|
||||
}));
|
||||
// Try sandbox job from DB first, scoped to the authenticated user.
|
||||
if let Some(ref store) = state.store
|
||||
&& let Ok(Some(job)) = store.get_sandbox_job(job_id).await
|
||||
{
|
||||
if job.user_id != state.user_id {
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
let browse_id = std::path::Path::new(&job.project_dir)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|| job.id.to_string());
|
||||
|
||||
let ui_state = match job.status.as_str() {
|
||||
"creating" => "pending",
|
||||
"running" => "in_progress",
|
||||
s => s,
|
||||
};
|
||||
|
||||
let elapsed_secs = job.started_at.map(|start| {
|
||||
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
|
||||
(end - start).num_seconds().max(0) as u64
|
||||
});
|
||||
|
||||
// Synthesize transitions from timestamps.
|
||||
let mut transitions = Vec::new();
|
||||
if let Some(started) = job.started_at {
|
||||
transitions.push(TransitionInfo {
|
||||
from: "creating".to_string(),
|
||||
to: "running".to_string(),
|
||||
timestamp: started.to_rfc3339(),
|
||||
reason: None,
|
||||
});
|
||||
}
|
||||
if let Some(completed) = job.completed_at {
|
||||
transitions.push(TransitionInfo {
|
||||
from: "running".to_string(),
|
||||
to: job.status.clone(),
|
||||
timestamp: completed.to_rfc3339(),
|
||||
reason: job.failure_reason.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
return Ok(Json(JobDetailResponse {
|
||||
id: job.id,
|
||||
title: job.task.clone(),
|
||||
description: String::new(),
|
||||
state: ui_state.to_string(),
|
||||
user_id: job.user_id.clone(),
|
||||
created_at: job.created_at.to_rfc3339(),
|
||||
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
|
||||
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
elapsed_secs,
|
||||
project_dir: Some(job.project_dir.clone()),
|
||||
browse_url: Some(format!("/projects/{}/", browse_id)),
|
||||
job_mode: {
|
||||
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
|
||||
mode.filter(|m| m != "worker")
|
||||
},
|
||||
transitions,
|
||||
}));
|
||||
}
|
||||
|
||||
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
||||
@@ -1031,31 +1197,36 @@ async fn jobs_cancel_handler(
|
||||
let job_id = Uuid::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||
|
||||
// Try sandbox job cancellation.
|
||||
if let Some(ref store) = state.store {
|
||||
if let Ok(Some(job)) = store.get_sandbox_job(job_id).await {
|
||||
if job.status == "running" || job.status == "creating" {
|
||||
// Stop the container if we have a job manager.
|
||||
if let Some(ref jm) = state.job_manager {
|
||||
let _ = jm.stop_job(job_id).await;
|
||||
}
|
||||
store
|
||||
.update_sandbox_job_status(
|
||||
job_id,
|
||||
"failed",
|
||||
Some(false),
|
||||
Some("Cancelled by user"),
|
||||
None,
|
||||
Some(chrono::Utc::now()),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
}
|
||||
return Ok(Json(serde_json::json!({
|
||||
"status": "cancelled",
|
||||
"job_id": job_id,
|
||||
})));
|
||||
// Try sandbox job cancellation, scoped to the authenticated user.
|
||||
if let Some(ref store) = state.store
|
||||
&& let Ok(Some(job)) = store.get_sandbox_job(job_id).await
|
||||
{
|
||||
if job.user_id != state.user_id {
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
if job.status == "running" || job.status == "creating" {
|
||||
// Stop the container if we have a job manager.
|
||||
if let Some(ref jm) = state.job_manager
|
||||
&& let Err(e) = jm.stop_job(job_id).await
|
||||
{
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container during cancellation");
|
||||
}
|
||||
store
|
||||
.update_sandbox_job_status(
|
||||
job_id,
|
||||
"failed",
|
||||
Some(false),
|
||||
Some("Cancelled by user"),
|
||||
None,
|
||||
Some(chrono::Utc::now()),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
}
|
||||
return Ok(Json(serde_json::json!({
|
||||
"status": "cancelled",
|
||||
"job_id": job_id,
|
||||
})));
|
||||
}
|
||||
|
||||
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
|
||||
@@ -1083,6 +1254,11 @@ async fn jobs_restart_handler(
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
||||
|
||||
// Scope to the authenticated user.
|
||||
if old_job.user_id != state.user_id {
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
|
||||
if old_job.status != "interrupted" && old_job.status != "failed" {
|
||||
return Err((
|
||||
StatusCode::CONFLICT,
|
||||
@@ -1157,6 +1333,16 @@ async fn jobs_prompt_handler(
|
||||
.parse()
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||
|
||||
// Verify user owns this job.
|
||||
if let Some(ref store) = state.store
|
||||
&& !store
|
||||
.sandbox_job_belongs_to_user(job_id, &state.user_id)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
|
||||
let content = body
|
||||
.get("content")
|
||||
.and_then(|v| v.as_str())
|
||||
@@ -1195,6 +1381,15 @@ async fn jobs_events_handler(
|
||||
.parse()
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
|
||||
|
||||
// Verify user owns this job.
|
||||
if !store
|
||||
.sandbox_job_belongs_to_user(job_id, &state.user_id)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
|
||||
let events = store
|
||||
.list_job_events(job_id)
|
||||
.await
|
||||
@@ -1244,6 +1439,11 @@ async fn job_files_list_handler(
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
||||
|
||||
// Verify user owns this job.
|
||||
if job.user_id != state.user_id {
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
|
||||
let base = std::path::PathBuf::from(&job.project_dir);
|
||||
let rel_path = query.path.as_deref().unwrap_or("");
|
||||
let target = base.join(rel_path);
|
||||
@@ -1307,6 +1507,11 @@ async fn job_files_read_handler(
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
|
||||
|
||||
// Verify user owns this job.
|
||||
if job.user_id != state.user_id {
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
|
||||
let path = query.path.as_deref().ok_or((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"path parameter required".to_string(),
|
||||
@@ -1525,6 +1730,15 @@ async fn project_file_handler(
|
||||
/// Shared logic: resolve the file inside `~/.ironclaw/projects/{project_id}/`,
|
||||
/// guard against path traversal, and stream the content with the right MIME type.
|
||||
async fn serve_project_file(project_id: &str, path: &str) -> axum::response::Response {
|
||||
// Reject project_id values that could escape the projects directory.
|
||||
if project_id.contains('/')
|
||||
|| project_id.contains('\\')
|
||||
|| project_id.contains("..")
|
||||
|| project_id.is_empty()
|
||||
{
|
||||
return (StatusCode::BAD_REQUEST, "Invalid project ID").into_response();
|
||||
}
|
||||
|
||||
let base = dirs::home_dir()
|
||||
.unwrap_or_else(|| std::path::PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
|
||||
+59
-12
@@ -13,10 +13,15 @@ use tokio_stream::wrappers::BroadcastStream;
|
||||
|
||||
use crate::channels::web::types::SseEvent;
|
||||
|
||||
/// Maximum number of concurrent SSE/WebSocket connections.
|
||||
/// Prevents resource exhaustion from connection flooding.
|
||||
const MAX_CONNECTIONS: u64 = 100;
|
||||
|
||||
/// Manages SSE broadcast to all connected browser tabs.
|
||||
pub struct SseManager {
|
||||
tx: broadcast::Sender<SseEvent>,
|
||||
connection_count: Arc<AtomicU64>,
|
||||
max_connections: u64,
|
||||
}
|
||||
|
||||
impl SseManager {
|
||||
@@ -27,6 +32,7 @@ impl SseManager {
|
||||
Self {
|
||||
tx,
|
||||
connection_count: Arc::new(AtomicU64::new(0)),
|
||||
max_connections: MAX_CONNECTIONS,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -45,25 +51,50 @@ impl SseManager {
|
||||
///
|
||||
/// Returns a stream of `SseEvent` values and increments/decrements the
|
||||
/// connection counter on creation/drop, just like `subscribe()` does for SSE.
|
||||
pub fn subscribe_raw(&self) -> impl Stream<Item = SseEvent> + Send + 'static + use<> {
|
||||
///
|
||||
/// Returns `None` if the maximum connection limit has been reached.
|
||||
pub fn subscribe_raw(&self) -> Option<impl Stream<Item = SseEvent> + Send + 'static + use<>> {
|
||||
// Atomically increment only if below the limit. This prevents
|
||||
// concurrent callers from overshooting max_connections.
|
||||
let counter = Arc::clone(&self.connection_count);
|
||||
counter.fetch_add(1, Ordering::Relaxed);
|
||||
let max = self.max_connections;
|
||||
counter
|
||||
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
|
||||
if current < max {
|
||||
Some(current + 1)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.ok()?;
|
||||
let rx = self.tx.subscribe();
|
||||
|
||||
let stream = BroadcastStream::new(rx).filter_map(|result| result.ok());
|
||||
|
||||
CountedStream {
|
||||
Some(CountedStream {
|
||||
inner: stream,
|
||||
counter,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new SSE stream for a client connection.
|
||||
///
|
||||
/// Returns `None` if the maximum connection limit has been reached.
|
||||
pub fn subscribe(
|
||||
&self,
|
||||
) -> Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>> {
|
||||
) -> Option<Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>>> {
|
||||
// Atomically increment only if below the limit.
|
||||
let counter = Arc::clone(&self.connection_count);
|
||||
counter.fetch_add(1, Ordering::Relaxed);
|
||||
let max = self.max_connections;
|
||||
counter
|
||||
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
|
||||
if current < max {
|
||||
Some(current + 1)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.ok()?;
|
||||
let rx = self.tx.subscribe();
|
||||
|
||||
let stream = BroadcastStream::new(rx)
|
||||
@@ -99,8 +130,10 @@ impl SseManager {
|
||||
counter,
|
||||
};
|
||||
|
||||
Sse::new(counted_stream)
|
||||
.keep_alive(KeepAlive::new().interval(Duration::from_secs(30)).text(""))
|
||||
Some(
|
||||
Sse::new(counted_stream)
|
||||
.keep_alive(KeepAlive::new().interval(Duration::from_secs(30)).text("")),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -175,7 +208,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_subscribe_raw_receives_events() {
|
||||
let manager = SseManager::new();
|
||||
let mut stream = Box::pin(manager.subscribe_raw());
|
||||
let mut stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
||||
|
||||
assert_eq!(manager.connection_count(), 1);
|
||||
|
||||
@@ -195,7 +228,7 @@ mod tests {
|
||||
async fn test_subscribe_raw_decrements_on_drop() {
|
||||
let manager = SseManager::new();
|
||||
{
|
||||
let _stream = Box::pin(manager.subscribe_raw());
|
||||
let _stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
||||
assert_eq!(manager.connection_count(), 1);
|
||||
}
|
||||
// Stream dropped, counter should decrement
|
||||
@@ -205,8 +238,8 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_subscribe_raw_multiple_subscribers() {
|
||||
let manager = SseManager::new();
|
||||
let mut s1 = Box::pin(manager.subscribe_raw());
|
||||
let mut s2 = Box::pin(manager.subscribe_raw());
|
||||
let mut s1 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
||||
let mut s2 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
|
||||
assert_eq!(manager.connection_count(), 2);
|
||||
|
||||
manager.broadcast(SseEvent::Heartbeat);
|
||||
@@ -221,4 +254,18 @@ mod tests {
|
||||
drop(s2);
|
||||
assert_eq!(manager.connection_count(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_subscribe_raw_rejects_over_limit() {
|
||||
let mut manager = SseManager::new();
|
||||
manager.max_connections = 2; // Low limit for testing
|
||||
|
||||
let _s1 = Box::pin(manager.subscribe_raw().expect("first should succeed"));
|
||||
let _s2 = Box::pin(manager.subscribe_raw().expect("second should succeed"));
|
||||
assert_eq!(manager.connection_count(), 2);
|
||||
|
||||
// Third should be rejected
|
||||
assert!(manager.subscribe_raw().is_none());
|
||||
assert!(manager.subscribe().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ let loadingOlder = false;
|
||||
let jobEvents = new Map(); // job_id -> Array of events
|
||||
let jobListRefreshTimer = null;
|
||||
const JOB_EVENTS_CAP = 500;
|
||||
const MEMORY_SEARCH_QUERY_MAX_LENGTH = 100;
|
||||
|
||||
// --- Auth ---
|
||||
|
||||
@@ -281,6 +282,8 @@ function sendApprovalAction(requestId, action) {
|
||||
function renderMarkdown(text) {
|
||||
if (typeof marked !== 'undefined') {
|
||||
let html = marked.parse(text);
|
||||
// Sanitize HTML output to prevent XSS from tool output or LLM responses.
|
||||
html = sanitizeRenderedHtml(html);
|
||||
// Inject copy buttons into <pre> blocks
|
||||
html = html.replace(/<pre>/g, '<pre class="code-block-wrapper"><button class="copy-btn" onclick="copyCodeBlock(this)">Copy</button>');
|
||||
return html;
|
||||
@@ -288,6 +291,28 @@ function renderMarkdown(text) {
|
||||
return escapeHtml(text);
|
||||
}
|
||||
|
||||
// Strip dangerous HTML elements and attributes from rendered markdown.
|
||||
// This prevents XSS from tool output or prompt injection in LLM responses.
|
||||
function sanitizeRenderedHtml(html) {
|
||||
html = html.replace(/<script\b[^<]*(?:(?!<\/script>)<[^<]*)*<\/script>/gi, '');
|
||||
html = html.replace(/<iframe\b[^>]*>[\s\S]*?<\/iframe>/gi, '');
|
||||
html = html.replace(/<object\b[^>]*>[\s\S]*?<\/object>/gi, '');
|
||||
html = html.replace(/<embed\b[^>]*\/?>/gi, '');
|
||||
html = html.replace(/<form\b[^>]*>[\s\S]*?<\/form>/gi, '');
|
||||
html = html.replace(/<style\b[^>]*>[\s\S]*?<\/style>/gi, '');
|
||||
html = html.replace(/<link\b[^>]*\/?>/gi, '');
|
||||
html = html.replace(/<base\b[^>]*\/?>/gi, '');
|
||||
html = html.replace(/<meta\b[^>]*\/?>/gi, '');
|
||||
// Remove event handler attributes (onclick, onerror, onload, etc.)
|
||||
html = html.replace(/\s+on\w+\s*=\s*"[^"]*"/gi, '');
|
||||
html = html.replace(/\s+on\w+\s*=\s*'[^']*'/gi, '');
|
||||
html = html.replace(/\s+on\w+\s*=\s*[^\s>]+/gi, '');
|
||||
// Remove javascript: and data: URLs in href/src attributes
|
||||
html = html.replace(/(href|src|action)\s*=\s*["']?\s*javascript\s*:/gi, '$1="');
|
||||
html = html.replace(/(href|src|action)\s*=\s*["']?\s*data\s*:/gi, '$1="');
|
||||
return html;
|
||||
}
|
||||
|
||||
function copyCodeBlock(btn) {
|
||||
const pre = btn.parentElement;
|
||||
const code = pre.querySelector('code');
|
||||
@@ -977,9 +1002,12 @@ function buildBreadcrumb(path) {
|
||||
}
|
||||
|
||||
function searchMemory(query) {
|
||||
const normalizedQuery = normalizeSearchQuery(query);
|
||||
if (!normalizedQuery) return;
|
||||
|
||||
apiFetch('/api/memory/search', {
|
||||
method: 'POST',
|
||||
body: { query, limit: 20 },
|
||||
body: { query: normalizedQuery, limit: 20 },
|
||||
}).then((data) => {
|
||||
const tree = document.getElementById('memory-tree');
|
||||
tree.innerHTML = '';
|
||||
@@ -990,18 +1018,23 @@ function searchMemory(query) {
|
||||
for (const result of data.results) {
|
||||
const item = document.createElement('div');
|
||||
item.className = 'search-result';
|
||||
const snippet = snippetAround(result.content, query, 120);
|
||||
const snippet = snippetAround(result.content, normalizedQuery, 120);
|
||||
item.innerHTML = '<div class="path">' + escapeHtml(result.path) + '</div>'
|
||||
+ '<div class="snippet">' + highlightQuery(snippet, query) + '</div>';
|
||||
+ '<div class="snippet">' + highlightQuery(snippet, normalizedQuery) + '</div>';
|
||||
item.addEventListener('click', () => readMemoryFile(result.path));
|
||||
tree.appendChild(item);
|
||||
}
|
||||
}).catch(() => {});
|
||||
}
|
||||
|
||||
function normalizeSearchQuery(query) {
|
||||
return (typeof query === 'string' ? query : '').slice(0, MEMORY_SEARCH_QUERY_MAX_LENGTH);
|
||||
}
|
||||
|
||||
function snippetAround(text, query, len) {
|
||||
const normalizedQuery = normalizeSearchQuery(query);
|
||||
const lower = text.toLowerCase();
|
||||
const idx = lower.indexOf(query.toLowerCase());
|
||||
const idx = lower.indexOf(normalizedQuery.toLowerCase());
|
||||
if (idx < 0) return text.substring(0, len);
|
||||
const start = Math.max(0, idx - Math.floor(len / 2));
|
||||
const end = Math.min(text.length, start + len);
|
||||
@@ -1014,11 +1047,11 @@ function snippetAround(text, query, len) {
|
||||
function highlightQuery(text, query) {
|
||||
if (!query) return escapeHtml(text);
|
||||
const escaped = escapeHtml(text);
|
||||
const queryEscaped = query.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
|
||||
const normalizedQuery = normalizeSearchQuery(query);
|
||||
const queryEscaped = normalizedQuery.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
|
||||
const re = new RegExp('(' + queryEscaped + ')', 'gi');
|
||||
return escaped.replace(re, '<mark>$1</mark>');
|
||||
}
|
||||
|
||||
// --- Logs ---
|
||||
|
||||
const LOG_MAX_ENTRIES = 2000;
|
||||
|
||||
@@ -5,7 +5,11 @@
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>IronClaw</title>
|
||||
<link rel="stylesheet" href="/style.css">
|
||||
<script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
|
||||
<script
|
||||
src="https://cdn.jsdelivr.net/npm/[email protected]/lib/marked.umd.min.js"
|
||||
integrity="sha384-pN9zSKOnTZwXRtYZAu0PBPEgR2B7DOC1aeLxQ33oJ0oy5iN1we6gm57xldM2irDG"
|
||||
crossorigin="anonymous"
|
||||
></script>
|
||||
</head>
|
||||
<body>
|
||||
<!-- Auth Screen -->
|
||||
|
||||
+13
-2
@@ -71,8 +71,17 @@ pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
|
||||
}
|
||||
let tracker_for_drop = state.ws_tracker.clone();
|
||||
|
||||
// Subscribe to broadcast events (same source as SSE)
|
||||
let mut event_stream = Box::pin(state.sse.subscribe_raw());
|
||||
// Subscribe to broadcast events (same source as SSE).
|
||||
// Reject if we've hit the connection limit.
|
||||
let Some(raw_stream) = state.sse.subscribe_raw() else {
|
||||
tracing::warn!("WebSocket rejected: too many connections");
|
||||
// Decrement the WS tracker we already incremented above.
|
||||
if let Some(ref tracker) = tracker_for_drop {
|
||||
tracker.decrement();
|
||||
}
|
||||
return;
|
||||
};
|
||||
let mut event_stream = Box::pin(raw_stream);
|
||||
|
||||
// Channel for the sender task to receive messages from both
|
||||
// the broadcast stream and any direct sends (like Pong)
|
||||
@@ -476,6 +485,8 @@ mod tests {
|
||||
user_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
chat_rate_limiter: crate::channels::web::server::RateLimiter::new(30, 60),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+50
-79
@@ -1,7 +1,9 @@
|
||||
//! Configuration management CLI commands.
|
||||
//!
|
||||
//! Commands for viewing and modifying settings.
|
||||
//! Settings are stored in PostgreSQL (env > DB > default).
|
||||
//! Settings are stored in the database (env > DB > default).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
@@ -46,11 +48,9 @@ pub enum ConfigCommand {
|
||||
/// Connects to the database to read/write settings. Falls back to disk
|
||||
/// if the database is not available.
|
||||
pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
// Try to connect to the DB for settings access
|
||||
let store = match connect_store().await {
|
||||
Ok(s) => Some(s),
|
||||
let db: Option<Arc<dyn crate::db::Database>> = match connect_db().await {
|
||||
Ok(d) => Some(d),
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: Could not connect to database ({}), using disk fallback",
|
||||
@@ -60,41 +60,42 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
}
|
||||
};
|
||||
|
||||
let db_ref = db.as_deref();
|
||||
match cmd {
|
||||
ConfigCommand::List { filter } => list_settings(store.as_ref(), filter).await,
|
||||
ConfigCommand::Get { path } => get_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Set { path, value } => set_setting(store.as_ref(), &path, &value).await,
|
||||
ConfigCommand::Reset { path } => reset_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Path => show_path(store.is_some()),
|
||||
ConfigCommand::List { filter } => list_settings(db_ref, filter).await,
|
||||
ConfigCommand::Get { path } => get_setting(db_ref, &path).await,
|
||||
ConfigCommand::Set { path, value } => set_setting(db_ref, &path, &value).await,
|
||||
ConfigCommand::Reset { path } => reset_setting(db_ref, &path).await,
|
||||
ConfigCommand::Path => show_path(db_ref.is_some()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bootstrap a DB connection for config commands.
|
||||
async fn connect_store() -> anyhow::Result<crate::history::Store> {
|
||||
/// Bootstrap a DB connection for config commands (backend-agnostic).
|
||||
async fn connect_db() -> anyhow::Result<Arc<dyn crate::db::Database>> {
|
||||
let config = crate::config::Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(store)
|
||||
crate::db::connect_from_config(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))
|
||||
}
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Load settings: DB if available, else disk.
|
||||
async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
|
||||
async fn load_settings(store: Option<&dyn crate::db::Database>) -> Settings {
|
||||
if let Some(store) = store {
|
||||
match store.get_all_settings(DEFAULT_USER_ID).await {
|
||||
Ok(map) if !map.is_empty() => return Settings::from_db_map(&map),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Settings::load()
|
||||
Settings::default()
|
||||
}
|
||||
|
||||
/// List all settings.
|
||||
async fn list_settings(
|
||||
store: Option<&crate::history::Store>,
|
||||
store: Option<&dyn crate::db::Database>,
|
||||
filter: Option<String>,
|
||||
) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
@@ -107,10 +108,10 @@ async fn list_settings(
|
||||
println!();
|
||||
|
||||
for (key, value) in all {
|
||||
if let Some(ref f) = filter {
|
||||
if !key.starts_with(f) {
|
||||
continue;
|
||||
}
|
||||
if let Some(ref f) = filter
|
||||
&& !key.starts_with(f)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let display_value = if value.len() > 60 {
|
||||
@@ -126,7 +127,7 @@ async fn list_settings(
|
||||
}
|
||||
|
||||
/// Get a specific setting.
|
||||
async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
async fn get_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
|
||||
match settings.get(path) {
|
||||
@@ -142,7 +143,7 @@ async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyho
|
||||
|
||||
/// Set a setting value.
|
||||
async fn set_setting(
|
||||
store: Option<&crate::history::Store>,
|
||||
store: Option<&dyn crate::db::Database>,
|
||||
path: &str,
|
||||
value: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
@@ -152,42 +153,36 @@ async fn set_setting(
|
||||
.set(path, value)
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
// Save to DB if available, otherwise disk
|
||||
if let Some(store) = store {
|
||||
let json_value = match serde_json::from_str::<serde_json::Value>(value) {
|
||||
Ok(v) => v,
|
||||
Err(_) => serde_json::Value::String(value.to_string()),
|
||||
};
|
||||
store
|
||||
.set_setting(DEFAULT_USER_ID, path, &json_value)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save to database: {}", e))?;
|
||||
} else {
|
||||
settings.save()?;
|
||||
}
|
||||
let store = store.ok_or_else(|| {
|
||||
anyhow::anyhow!("Database connection required to save settings. Check DATABASE_URL.")
|
||||
})?;
|
||||
let json_value = match serde_json::from_str::<serde_json::Value>(value) {
|
||||
Ok(v) => v,
|
||||
Err(_) => serde_json::Value::String(value.to_string()),
|
||||
};
|
||||
store
|
||||
.set_setting(DEFAULT_USER_ID, path, &json_value)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save to database: {}", e))?;
|
||||
|
||||
println!("Set {} = {}", path, value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reset a setting to default.
|
||||
async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
async fn reset_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
|
||||
let default = Settings::default();
|
||||
let default_value = default
|
||||
.get(path)
|
||||
.ok_or_else(|| anyhow::anyhow!("Unknown setting: {}", path))?;
|
||||
|
||||
// Delete from DB (falling back to default) or reset on disk
|
||||
if let Some(store) = store {
|
||||
store
|
||||
.delete_setting(DEFAULT_USER_ID, path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to delete setting from database: {}", e))?;
|
||||
} else {
|
||||
let mut settings = Settings::load();
|
||||
settings.reset(path).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
settings.save()?;
|
||||
}
|
||||
let store = store.ok_or_else(|| {
|
||||
anyhow::anyhow!("Database connection required to reset settings. Check DATABASE_URL.")
|
||||
})?;
|
||||
store
|
||||
.delete_setting(DEFAULT_USER_ID, path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to delete setting from database: {}", e))?;
|
||||
|
||||
println!("Reset {} to default: {}", path, default_value);
|
||||
Ok(())
|
||||
@@ -196,38 +191,14 @@ async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> any
|
||||
/// Show the settings storage info.
|
||||
fn show_path(has_db: bool) -> anyhow::Result<()> {
|
||||
if has_db {
|
||||
println!("Settings stored in: PostgreSQL (settings table)");
|
||||
println!(
|
||||
"Bootstrap config: {}",
|
||||
crate::bootstrap::BootstrapConfig::default_path().display()
|
||||
);
|
||||
println!("Settings stored in: database (settings table)");
|
||||
} else {
|
||||
let path = Settings::default_path();
|
||||
println!("Settings stored in: {} (disk fallback)", path.display());
|
||||
|
||||
if path.exists() {
|
||||
let metadata = std::fs::metadata(&path)?;
|
||||
println!(" Size: {} bytes", metadata.len());
|
||||
if let Ok(modified) = metadata.modified() {
|
||||
use std::time::SystemTime;
|
||||
let duration = SystemTime::now()
|
||||
.duration_since(modified)
|
||||
.unwrap_or_default();
|
||||
let secs = duration.as_secs();
|
||||
if secs < 60 {
|
||||
println!(" Modified: {} seconds ago", secs);
|
||||
} else if secs < 3600 {
|
||||
println!(" Modified: {} minutes ago", secs / 60);
|
||||
} else if secs < 86400 {
|
||||
println!(" Modified: {} hours ago", secs / 3600);
|
||||
} else {
|
||||
println!(" Modified: {} days ago", secs / 86400);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!(" (does not exist, using defaults)");
|
||||
}
|
||||
println!("Settings stored in: PostgreSQL (not connected, using defaults)");
|
||||
}
|
||||
println!(
|
||||
"Env config: {}",
|
||||
crate::bootstrap::ironclaw_env_path().display()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+82
-35
@@ -8,8 +8,10 @@ use std::sync::Arc;
|
||||
use clap::Subcommand;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{SecretsCrypto, SecretsStore};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
|
||||
auth::{authorize_mcp_server, is_authenticated},
|
||||
@@ -172,10 +174,10 @@ async fn add_server(
|
||||
config.validate()?;
|
||||
|
||||
// Save (DB if available, else disk)
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
servers.upsert(config);
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Added MCP server '{}'", name);
|
||||
@@ -193,12 +195,12 @@ async fn add_server(
|
||||
|
||||
/// Remove an MCP server.
|
||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
if !servers.remove(&name) {
|
||||
anyhow::bail!("Server '{}' not found", name);
|
||||
}
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Removed MCP server '{}'", name);
|
||||
@@ -209,8 +211,8 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
|
||||
/// List configured MCP servers.
|
||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
|
||||
if servers.servers.is_empty() {
|
||||
println!();
|
||||
@@ -268,8 +270,8 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
/// Authenticate with an MCP server.
|
||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
@@ -341,8 +343,8 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
/// Test connection to an MCP server.
|
||||
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
@@ -437,8 +439,8 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
|
||||
/// Toggle server enabled/disabled state.
|
||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
|
||||
let server = servers
|
||||
.get_mut(&name)
|
||||
@@ -453,7 +455,7 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
};
|
||||
|
||||
server.enabled = new_state;
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
let status = if new_state { "enabled" } else { "disabled" };
|
||||
println!();
|
||||
@@ -465,18 +467,16 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Try to connect to the database store for DB-backed config.
|
||||
async fn connect_store() -> Option<Store> {
|
||||
/// Try to connect to the database (backend-agnostic).
|
||||
async fn connect_db() -> Option<Arc<dyn Database>> {
|
||||
let config = Config::from_env().await.ok()?;
|
||||
let store = Store::new(&config.database).await.ok()?;
|
||||
store.run_migrations().await.ok()?;
|
||||
Some(store)
|
||||
crate::db::connect_from_config(&config.database).await.ok()
|
||||
}
|
||||
|
||||
/// Load MCP servers (DB if available, else disk).
|
||||
async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::load_mcp_servers_from_db(store, DEFAULT_USER_ID).await
|
||||
async fn load_servers(db: Option<&dyn Database>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(db) = db {
|
||||
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await
|
||||
} else {
|
||||
config::load_mcp_servers().await
|
||||
}
|
||||
@@ -484,11 +484,11 @@ async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::C
|
||||
|
||||
/// Save MCP servers (DB if available, else disk).
|
||||
async fn save_servers(
|
||||
store: Option<&Store>,
|
||||
db: Option<&dyn Database>,
|
||||
servers: &McpServersFile,
|
||||
) -> Result<(), config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::save_mcp_servers_to_db(store, DEFAULT_USER_ID, servers).await
|
||||
if let Some(db) = db {
|
||||
config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, servers).await
|
||||
} else {
|
||||
config::save_mcp_servers(servers).await
|
||||
}
|
||||
@@ -504,14 +504,61 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
|
||||
)
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
return Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+26
-1
@@ -9,6 +9,30 @@ use clap::Subcommand;
|
||||
|
||||
use crate::workspace::{EmbeddingProvider, SearchConfig, Workspace};
|
||||
|
||||
/// Run a memory command using the Database trait (works with any backend).
|
||||
pub async fn run_memory_command_with_db(
|
||||
cmd: MemoryCommand,
|
||||
db: std::sync::Arc<dyn crate::db::Database>,
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut workspace = Workspace::new_with_db("default", db);
|
||||
if let Some(emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb);
|
||||
}
|
||||
|
||||
match cmd {
|
||||
MemoryCommand::Search { query, limit } => search(&workspace, &query, limit).await,
|
||||
MemoryCommand::Read { path } => read(&workspace, &path).await,
|
||||
MemoryCommand::Write {
|
||||
path,
|
||||
content,
|
||||
append,
|
||||
} => write(&workspace, &path, content, append).await,
|
||||
MemoryCommand::Tree { path, depth } => tree(&workspace, &path, depth).await,
|
||||
MemoryCommand::Status => status(&workspace).await,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum MemoryCommand {
|
||||
/// Search workspace memory (hybrid full-text + semantic)
|
||||
@@ -55,7 +79,8 @@ pub enum MemoryCommand {
|
||||
Status,
|
||||
}
|
||||
|
||||
/// Run a memory command.
|
||||
/// Run a memory command (PostgreSQL backend).
|
||||
#[cfg(feature = "postgres")]
|
||||
pub async fn run_memory_command(
|
||||
cmd: MemoryCommand,
|
||||
pool: deadpool_postgres::Pool,
|
||||
|
||||
+5
-1
@@ -12,13 +12,17 @@
|
||||
mod config;
|
||||
mod mcp;
|
||||
pub mod memory;
|
||||
pub mod oauth_defaults;
|
||||
mod pairing;
|
||||
pub mod status;
|
||||
mod tool;
|
||||
|
||||
pub use config::{ConfigCommand, run_config_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use memory::{MemoryCommand, run_memory_command};
|
||||
pub use memory::MemoryCommand;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use memory::run_memory_command;
|
||||
pub use memory::run_memory_command_with_db;
|
||||
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
|
||||
pub use status::run_status_command;
|
||||
pub use tool::{ToolCommand, run_tool_command};
|
||||
|
||||
@@ -0,0 +1,342 @@
|
||||
//! Shared OAuth infrastructure: built-in credentials, callback server, landing pages.
|
||||
//!
|
||||
//! Every OAuth flow in the codebase (WASM tool auth, MCP server auth, NEAR AI login)
|
||||
//! uses the same callback port, landing page, and listener logic from this module.
|
||||
//!
|
||||
//! # Built-in Credentials
|
||||
//!
|
||||
//! Many CLI tools (gcloud, rclone, gdrive) ship with default OAuth credentials
|
||||
//! so users don't need to register their own OAuth app. Google explicitly
|
||||
//! documents that client_secret for "Desktop App" / "Installed App" types
|
||||
//! is NOT actually secret.
|
||||
//!
|
||||
//! Default credentials are hardcoded below. They can be overridden at:
|
||||
//!
|
||||
//! - **Compile time**: Set IRONCLAW_GOOGLE_CLIENT_ID / IRONCLAW_GOOGLE_CLIENT_SECRET
|
||||
//! env vars before building to replace the hardcoded defaults.
|
||||
//! - **Runtime**: Users can set GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET
|
||||
//! env vars, which take priority over built-in defaults.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
// ── Built-in credentials ────────────────────────────────────────────────
|
||||
|
||||
pub struct OAuthCredentials {
|
||||
pub client_id: &'static str,
|
||||
pub client_secret: &'static str,
|
||||
}
|
||||
|
||||
/// Google OAuth "Desktop App" credentials, shared across all Google tools.
|
||||
/// Compile-time env vars override the hardcoded defaults below.
|
||||
const GOOGLE_CLIENT_ID: &str = match option_env!("IRONCLAW_GOOGLE_CLIENT_ID") {
|
||||
Some(v) => v,
|
||||
None => "564604149681-efo25d43rs85v0tibdepsmdv5dsrhhr0.apps.googleusercontent.com",
|
||||
};
|
||||
const GOOGLE_CLIENT_SECRET: &str = match option_env!("IRONCLAW_GOOGLE_CLIENT_SECRET") {
|
||||
Some(v) => v,
|
||||
None => "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",
|
||||
};
|
||||
|
||||
/// Returns built-in OAuth credentials for a provider, keyed by secret_name.
|
||||
///
|
||||
/// The secret_name comes from the tool's capabilities.json `auth.secret_name` field.
|
||||
/// Returns `None` if no built-in credentials are configured for that provider.
|
||||
pub fn builtin_credentials(secret_name: &str) -> Option<OAuthCredentials> {
|
||||
match secret_name {
|
||||
"google_oauth_token" => Some(OAuthCredentials {
|
||||
client_id: GOOGLE_CLIENT_ID,
|
||||
client_secret: GOOGLE_CLIENT_SECRET,
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Shared callback server ──────────────────────────────────────────────
|
||||
|
||||
/// Fixed port for all OAuth callbacks.
|
||||
///
|
||||
/// Every redirect URI registered with providers must use this port:
|
||||
/// `http://localhost:9876/callback` (or `/auth/callback` for NEAR AI).
|
||||
pub const OAUTH_CALLBACK_PORT: u16 = 9876;
|
||||
|
||||
/// Error from the OAuth callback listener.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum OAuthCallbackError {
|
||||
#[error("Port {0} is in use (another auth flow running?): {1}")]
|
||||
PortInUse(u16, String),
|
||||
|
||||
#[error("Authorization denied by user")]
|
||||
Denied,
|
||||
|
||||
#[error("Timed out waiting for authorization")]
|
||||
Timeout,
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(String),
|
||||
}
|
||||
|
||||
/// Bind the OAuth callback listener on the fixed port.
|
||||
///
|
||||
/// Binds to IPv4 `127.0.0.1` first because callback URLs use `127.0.0.1`
|
||||
/// explicitly (e.g., NEAR AI redirects to `http://127.0.0.1:9876/auth/callback`).
|
||||
/// Falls back to IPv6 `[::1]` only if IPv4 binding fails for a reason other
|
||||
/// than `AddrInUse`. If the port is already occupied, fails immediately.
|
||||
pub async fn bind_callback_listener() -> Result<TcpListener, OAuthCallbackError> {
|
||||
let ipv4_addr = format!("127.0.0.1:{}", OAUTH_CALLBACK_PORT);
|
||||
match TcpListener::bind(&ipv4_addr).await {
|
||||
Ok(listener) => return Ok(listener),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::AddrInUse => {
|
||||
return Err(OAuthCallbackError::PortInUse(
|
||||
OAUTH_CALLBACK_PORT,
|
||||
e.to_string(),
|
||||
));
|
||||
}
|
||||
Err(_) => {
|
||||
// IPv4 not available, fall back to IPv6
|
||||
}
|
||||
}
|
||||
TcpListener::bind(format!("[::1]:{}", OAUTH_CALLBACK_PORT))
|
||||
.await
|
||||
.map_err(|e| {
|
||||
if e.kind() == std::io::ErrorKind::AddrInUse {
|
||||
OAuthCallbackError::PortInUse(OAUTH_CALLBACK_PORT, e.to_string())
|
||||
} else {
|
||||
OAuthCallbackError::Io(e.to_string())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Wait for an OAuth callback and extract a query parameter value.
|
||||
///
|
||||
/// Listens for a GET request matching `path_prefix` (e.g., "/callback" or "/auth/callback"),
|
||||
/// extracts the value of `param_name` (e.g., "code" or "token"), and shows a branded
|
||||
/// landing page using `display_name` (e.g., "Google", "Notion", "NEAR AI").
|
||||
///
|
||||
/// Times out after 5 minutes.
|
||||
pub async fn wait_for_callback(
|
||||
listener: TcpListener,
|
||||
path_prefix: &str,
|
||||
param_name: &str,
|
||||
display_name: &str,
|
||||
) -> Result<String, OAuthCallbackError> {
|
||||
let path_prefix = path_prefix.to_string();
|
||||
let param_name = param_name.to_string();
|
||||
let display_name = display_name.to_string();
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(300), async move {
|
||||
loop {
|
||||
let (mut socket, _) = listener
|
||||
.accept()
|
||||
.await
|
||||
.map_err(|e| OAuthCallbackError::Io(e.to_string()))?;
|
||||
|
||||
let mut reader = BufReader::new(&mut socket);
|
||||
let mut request_line = String::new();
|
||||
reader
|
||||
.read_line(&mut request_line)
|
||||
.await
|
||||
.map_err(|e| OAuthCallbackError::Io(e.to_string()))?;
|
||||
|
||||
if let Some(path) = request_line.split_whitespace().nth(1)
|
||||
&& path.starts_with(&path_prefix)
|
||||
&& let Some(query) = path.split('?').nth(1)
|
||||
{
|
||||
// Check for error first
|
||||
if query.contains("error=") {
|
||||
let html = landing_html(&display_name, false);
|
||||
let response = format!(
|
||||
"HTTP/1.1 400 Bad Request\r\n\
|
||||
Content-Type: text/html; charset=utf-8\r\n\
|
||||
Connection: close\r\n\
|
||||
\r\n\
|
||||
{}",
|
||||
html
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
return Err(OAuthCallbackError::Denied);
|
||||
}
|
||||
|
||||
// Look for the target parameter
|
||||
for param in query.split('&') {
|
||||
let parts: Vec<&str> = param.splitn(2, '=').collect();
|
||||
if parts.len() == 2 && parts[0] == param_name {
|
||||
let value = urlencoding::decode(parts[1])
|
||||
.unwrap_or_else(|_| parts[1].into())
|
||||
.into_owned();
|
||||
|
||||
let html = landing_html(&display_name, true);
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\n\
|
||||
Content-Type: text/html; charset=utf-8\r\n\
|
||||
Connection: close\r\n\
|
||||
\r\n\
|
||||
{}",
|
||||
html
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
let _ = socket.shutdown().await;
|
||||
|
||||
return Ok(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Not the callback we're looking for
|
||||
let response = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| OAuthCallbackError::Timeout)?
|
||||
}
|
||||
|
||||
/// Escape a string for safe interpolation into HTML content.
|
||||
fn html_escape(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
for c in s.chars() {
|
||||
match c {
|
||||
'&' => out.push_str("&"),
|
||||
'<' => out.push_str("<"),
|
||||
'>' => out.push_str(">"),
|
||||
'"' => out.push_str("""),
|
||||
'\'' => out.push_str("'"),
|
||||
_ => out.push(c),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// HTML landing page shown in the browser after an OAuth redirect.
|
||||
pub fn landing_html(provider_name: &str, success: bool) -> String {
|
||||
let safe_name = html_escape(provider_name);
|
||||
let (icon, heading, subtitle, accent) = if success {
|
||||
(
|
||||
r##"<div style="width:64px;height:64px;border-radius:50%;background:#22c55e;display:flex;align-items:center;justify-content:center;margin:0 auto 24px">
|
||||
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="#fff" stroke-width="3" stroke-linecap="round" stroke-linejoin="round"><polyline points="20 6 9 17 4 12"/></svg>
|
||||
</div>"##,
|
||||
format!("{} Connected", safe_name),
|
||||
"You can close this window and return to your terminal.",
|
||||
"#22c55e",
|
||||
)
|
||||
} else {
|
||||
(
|
||||
r##"<div style="width:64px;height:64px;border-radius:50%;background:#ef4444;display:flex;align-items:center;justify-content:center;margin:0 auto 24px">
|
||||
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="#fff" stroke-width="3" stroke-linecap="round" stroke-linejoin="round"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg>
|
||||
</div>"##,
|
||||
"Authorization Failed".to_string(),
|
||||
"The request was denied. You can close this window and try again.",
|
||||
"#ef4444",
|
||||
)
|
||||
};
|
||||
|
||||
format!(
|
||||
r#"<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width,initial-scale=1">
|
||||
<title>IronClaw - {heading}</title>
|
||||
<style>
|
||||
* {{ margin:0; padding:0; box-sizing:border-box }}
|
||||
body {{
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
|
||||
background: #0a0a0a;
|
||||
color: #e5e5e5;
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
min-height: 100vh;
|
||||
}}
|
||||
.card {{
|
||||
text-align: center;
|
||||
padding: 48px 40px;
|
||||
max-width: 420px;
|
||||
border: 1px solid #262626;
|
||||
border-radius: 16px;
|
||||
background: #141414;
|
||||
}}
|
||||
h1 {{
|
||||
font-size: 22px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 8px;
|
||||
color: #fafafa;
|
||||
}}
|
||||
p {{
|
||||
font-size: 14px;
|
||||
color: #a3a3a3;
|
||||
line-height: 1.5;
|
||||
}}
|
||||
.accent {{ color: {accent}; }}
|
||||
.brand {{
|
||||
margin-top: 32px;
|
||||
font-size: 12px;
|
||||
color: #525252;
|
||||
letter-spacing: 0.5px;
|
||||
text-transform: uppercase;
|
||||
}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="card">
|
||||
{icon}
|
||||
<h1>{heading}</h1>
|
||||
<p>{subtitle}</p>
|
||||
<div class="brand">IronClaw</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>"#,
|
||||
heading = heading,
|
||||
icon = icon,
|
||||
subtitle = subtitle,
|
||||
accent = accent,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::cli::oauth_defaults::{builtin_credentials, landing_html};
|
||||
|
||||
#[test]
|
||||
fn test_unknown_provider_returns_none() {
|
||||
assert!(builtin_credentials("unknown_token").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_google_returns_based_on_compile_env() {
|
||||
let creds = builtin_credentials("google_oauth_token");
|
||||
assert!(creds.is_some());
|
||||
let creds = creds.unwrap();
|
||||
assert!(!creds.client_id.is_empty());
|
||||
assert!(!creds.client_secret.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_landing_html_success_contains_key_elements() {
|
||||
let html = landing_html("Google", true);
|
||||
assert!(html.contains("Google Connected"));
|
||||
assert!(html.contains("charset"));
|
||||
assert!(html.contains("IronClaw"));
|
||||
assert!(html.contains("#22c55e")); // green accent
|
||||
assert!(!html.contains("Failed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_landing_html_escapes_provider_name() {
|
||||
let html = landing_html("<script>alert(1)</script>", true);
|
||||
assert!(!html.contains("<script>"));
|
||||
assert!(html.contains("<script>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_landing_html_error_contains_key_elements() {
|
||||
let html = landing_html("Notion", false);
|
||||
assert!(html.contains("Authorization Failed"));
|
||||
assert!(html.contains("charset"));
|
||||
assert!(html.contains("IronClaw"));
|
||||
assert!(html.contains("#ef4444")); // red accent
|
||||
assert!(!html.contains("Connected"));
|
||||
}
|
||||
}
|
||||
+51
-24
@@ -9,7 +9,7 @@ use crate::settings::Settings;
|
||||
|
||||
/// Run the status command, printing system health info.
|
||||
pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
let settings = Settings::load();
|
||||
let settings = Settings::default();
|
||||
|
||||
println!("IronClaw Status");
|
||||
println!("===============\n");
|
||||
@@ -22,16 +22,36 @@ pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
);
|
||||
|
||||
// Database
|
||||
let db_url_set = settings.database_url.is_some() || std::env::var("DATABASE_URL").is_ok();
|
||||
print!(" Database: ");
|
||||
if db_url_set {
|
||||
// Try to connect
|
||||
match check_database().await {
|
||||
Ok(()) => println!("connected"),
|
||||
Err(e) => println!("error ({})", e),
|
||||
let db_backend = std::env::var("DATABASE_BACKEND")
|
||||
.ok()
|
||||
.unwrap_or_else(|| "postgres".to_string());
|
||||
match db_backend.as_str() {
|
||||
"libsql" | "turso" | "sqlite" => {
|
||||
let path = std::env::var("LIBSQL_PATH")
|
||||
.map(std::path::PathBuf::from)
|
||||
.unwrap_or_else(|_| crate::config::default_libsql_path());
|
||||
if path.exists() {
|
||||
let turso = if std::env::var("LIBSQL_URL").is_ok() {
|
||||
" + Turso sync"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
println!("libSQL ({}{})", path.display(), turso);
|
||||
} else {
|
||||
println!("libSQL (file missing: {})", path.display());
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if std::env::var("DATABASE_URL").is_ok() {
|
||||
match check_database().await {
|
||||
Ok(()) => println!("connected (PostgreSQL)"),
|
||||
Err(e) => println!("error ({})", e),
|
||||
}
|
||||
} else {
|
||||
println!("not configured");
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!("not configured");
|
||||
}
|
||||
|
||||
// Session / Auth
|
||||
@@ -43,15 +63,17 @@ pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
println!("not found (run `ironclaw onboard`)");
|
||||
}
|
||||
|
||||
// Secrets
|
||||
// Secrets (auto-detect from env only; skip keychain probe to avoid
|
||||
// triggering macOS system password dialogs on a simple status check)
|
||||
print!(" Secrets: ");
|
||||
let secrets_configured = settings.secrets_master_key_source != crate::settings::KeySource::None
|
||||
|| std::env::var("SECRETS_MASTER_KEY").is_ok()
|
||||
|| crate::secrets::keychain::has_master_key().await;
|
||||
if secrets_configured {
|
||||
println!("configured ({:?})", settings.secrets_master_key_source);
|
||||
if std::env::var("SECRETS_MASTER_KEY").is_ok() {
|
||||
println!("configured (env)");
|
||||
} else {
|
||||
println!("not configured");
|
||||
// We don't probe the keychain here because get_generic_password()
|
||||
// triggers macOS unlock+authorization dialogs, which is bad UX for
|
||||
// a read-only status command. If onboarding completed with keychain
|
||||
// storage, the key is there; we just can't cheaply verify it.
|
||||
println!("env not set (keychain may be configured)");
|
||||
}
|
||||
|
||||
// Embeddings
|
||||
@@ -129,19 +151,18 @@ pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
Err(_) => println!("none configured"),
|
||||
}
|
||||
|
||||
// Settings path
|
||||
println!("\n Settings: {}", Settings::default_path().display());
|
||||
// Config path
|
||||
println!(
|
||||
"\n Config: {}",
|
||||
crate::bootstrap::ironclaw_env_path().display()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn check_database() -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
let settings = Settings::load();
|
||||
let url = std::env::var("DATABASE_URL")
|
||||
.ok()
|
||||
.or(settings.database_url)
|
||||
.ok_or_else(|| anyhow::anyhow!("no URL"))?;
|
||||
let url = std::env::var("DATABASE_URL").map_err(|_| anyhow::anyhow!("DATABASE_URL not set"))?;
|
||||
|
||||
let config: deadpool_postgres::Config = deadpool_postgres::Config {
|
||||
url: Some(url),
|
||||
@@ -167,6 +188,12 @@ async fn check_database() -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
async fn check_database() -> anyhow::Result<()> {
|
||||
// For non-postgres backends, just report configured
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn count_wasm_files(dir: &std::path::Path) -> usize {
|
||||
std::fs::read_dir(dir)
|
||||
.map(|entries| {
|
||||
|
||||
+235
-163
@@ -11,8 +11,11 @@ use clap::Subcommand;
|
||||
use tokio::fs;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
#[allow(unused_imports)]
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{CreateSecretParams, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
|
||||
|
||||
/// Default tools directory.
|
||||
@@ -420,11 +423,11 @@ async fn extract_crate_name(cargo_toml: &Path) -> anyhow::Result<String> {
|
||||
// Simple TOML parsing for [package] name
|
||||
for line in content.lines() {
|
||||
let line = line.trim();
|
||||
if line.starts_with("name") {
|
||||
if let Some((_, value)) = line.split_once('=') {
|
||||
let name = value.trim().trim_matches('"').trim_matches('\'');
|
||||
return Ok(name.to_string());
|
||||
}
|
||||
if line.starts_with("name")
|
||||
&& let Some((_, value)) = line.split_once('=')
|
||||
{
|
||||
let name = value.trim().trim_matches('"').trim_matches('\'');
|
||||
return Ok(name.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -488,10 +491,10 @@ async fn list_tools(dir: Option<PathBuf>, verbose: bool) -> anyhow::Result<()> {
|
||||
|
||||
if has_caps {
|
||||
let caps_path = path.with_extension("capabilities.json");
|
||||
if let Ok(content) = fs::read_to_string(&caps_path).await {
|
||||
if let Ok(caps) = CapabilitiesFile::from_json(&content) {
|
||||
print_capabilities_summary(&caps);
|
||||
}
|
||||
if let Ok(content) = fs::read_to_string(&caps_path).await
|
||||
&& let Ok(caps) = CapabilitiesFile::from_json(&content)
|
||||
{
|
||||
print_capabilities_summary(&caps);
|
||||
}
|
||||
}
|
||||
println!();
|
||||
@@ -604,16 +607,16 @@ fn print_capabilities_summary(caps: &CapabilitiesFile) {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref secrets) = caps.secrets {
|
||||
if !secrets.allowed_names.is_empty() {
|
||||
parts.push(format!("secrets: {}", secrets.allowed_names.len()));
|
||||
}
|
||||
if let Some(ref secrets) = caps.secrets
|
||||
&& !secrets.allowed_names.is_empty()
|
||||
{
|
||||
parts.push(format!("secrets: {}", secrets.allowed_names.len()));
|
||||
}
|
||||
|
||||
if let Some(ref ws) = caps.workspace {
|
||||
if !ws.allowed_prefixes.is_empty() {
|
||||
parts.push("workspace: read".to_string());
|
||||
}
|
||||
if let Some(ref ws) = caps.workspace
|
||||
&& !ws.allowed_prefixes.is_empty()
|
||||
{
|
||||
parts.push("workspace: read".to_string());
|
||||
}
|
||||
|
||||
if !parts.is_empty() {
|
||||
@@ -650,30 +653,30 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref secrets) = caps.secrets {
|
||||
if !secrets.allowed_names.is_empty() {
|
||||
println!(" Secrets (existence check only):");
|
||||
for name in &secrets.allowed_names {
|
||||
println!(" {}", name);
|
||||
}
|
||||
if let Some(ref secrets) = caps.secrets
|
||||
&& !secrets.allowed_names.is_empty()
|
||||
{
|
||||
println!(" Secrets (existence check only):");
|
||||
for name in &secrets.allowed_names {
|
||||
println!(" {}", name);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref tool_invoke) = caps.tool_invoke {
|
||||
if !tool_invoke.aliases.is_empty() {
|
||||
println!(" Tool aliases:");
|
||||
for (alias, real_name) in &tool_invoke.aliases {
|
||||
println!(" {} -> {}", alias, real_name);
|
||||
}
|
||||
if let Some(ref tool_invoke) = caps.tool_invoke
|
||||
&& !tool_invoke.aliases.is_empty()
|
||||
{
|
||||
println!(" Tool aliases:");
|
||||
for (alias, real_name) in &tool_invoke.aliases {
|
||||
println!(" {} -> {}", alias, real_name);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref ws) = caps.workspace {
|
||||
if !ws.allowed_prefixes.is_empty() {
|
||||
println!(" Workspace read prefixes:");
|
||||
for prefix in &ws.allowed_prefixes {
|
||||
println!(" {}", prefix);
|
||||
}
|
||||
if let Some(ref ws) = caps.workspace
|
||||
&& !ws.allowed_prefixes.is_empty()
|
||||
{
|
||||
println!(" Workspace read prefixes:");
|
||||
for prefix in &ws.allowed_prefixes {
|
||||
println!(" {}", prefix);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -722,11 +725,58 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
)
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
|
||||
|
||||
let secrets_store: Arc<dyn SecretsStore + Send + Sync> = {
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)))
|
||||
}
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
))
|
||||
}
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Check if already configured
|
||||
let already_configured = secrets_store
|
||||
@@ -752,51 +802,103 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
}
|
||||
|
||||
// Check for environment variable
|
||||
if let Some(ref env_var) = auth.env_var {
|
||||
if let Ok(token) = std::env::var(env_var) {
|
||||
if !token.is_empty() {
|
||||
println!(" Found {} in environment.", env_var);
|
||||
println!();
|
||||
if let Some(ref env_var) = auth.env_var
|
||||
&& let Ok(token) = std::env::var(env_var)
|
||||
&& !token.is_empty()
|
||||
{
|
||||
println!(" Found {} in environment.", env_var);
|
||||
println!();
|
||||
|
||||
// Validate if endpoint is provided
|
||||
if let Some(ref validation) = auth.validation_endpoint {
|
||||
print!(" Validating token...");
|
||||
std::io::stdout().flush()?;
|
||||
// Validate if endpoint is provided
|
||||
if let Some(ref validation) = auth.validation_endpoint {
|
||||
print!(" Validating token...");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
match validate_token(&token, validation, &auth.secret_name).await {
|
||||
Ok(()) => {
|
||||
println!(" ✓");
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" ✗");
|
||||
println!(" Validation failed: {}", e);
|
||||
println!();
|
||||
println!(" Falling back to manual entry...");
|
||||
return auth_tool_manual(&secrets_store, &user_id, &auth).await;
|
||||
}
|
||||
}
|
||||
match validate_token(&token, validation, &auth.secret_name).await {
|
||||
Ok(()) => {
|
||||
println!(" ✓");
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" ✗");
|
||||
println!(" Validation failed: {}", e);
|
||||
println!();
|
||||
println!(" Falling back to manual entry...");
|
||||
return auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await;
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(&secrets_store, &user_id, &auth, &token).await?;
|
||||
print_success(display_name);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(secrets_store.as_ref(), &user_id, &auth, &token, None, None).await?;
|
||||
print_success(display_name);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Check for OAuth configuration
|
||||
if let Some(ref oauth) = auth.oauth {
|
||||
return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await;
|
||||
// For providers with shared tokens (e.g., all Google tools share google_oauth_token),
|
||||
// combine scopes from all installed tools so one auth covers everything.
|
||||
let combined = combine_provider_scopes(&tools_dir, &auth.secret_name, oauth).await;
|
||||
if combined.scopes.len() > oauth.scopes.len() {
|
||||
let extra = combined.scopes.len() - oauth.scopes.len();
|
||||
println!(
|
||||
" Including scopes from {} other installed tool(s) sharing this credential.",
|
||||
extra
|
||||
);
|
||||
println!();
|
||||
}
|
||||
return auth_tool_oauth(secrets_store.as_ref(), &user_id, &auth, &combined).await;
|
||||
}
|
||||
|
||||
// Fall back to manual entry
|
||||
auth_tool_manual(&secrets_store, &user_id, &auth).await
|
||||
auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await
|
||||
}
|
||||
|
||||
/// Scan the tools directory for all capabilities files sharing the same secret_name
|
||||
/// and combine their OAuth scopes. This way, authing any Google tool requests scopes
|
||||
/// for ALL installed Google tools, so one login covers everything.
|
||||
async fn combine_provider_scopes(
|
||||
tools_dir: &Path,
|
||||
secret_name: &str,
|
||||
base_oauth: &crate::tools::wasm::OAuthConfigSchema,
|
||||
) -> crate::tools::wasm::OAuthConfigSchema {
|
||||
let mut all_scopes: std::collections::HashSet<String> =
|
||||
base_oauth.scopes.iter().cloned().collect();
|
||||
|
||||
if let Ok(mut entries) = tokio::fs::read_dir(tools_dir).await {
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("json") {
|
||||
continue;
|
||||
}
|
||||
let name = path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or_default();
|
||||
if !name.ends_with(".capabilities.json") {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Ok(content) = tokio::fs::read_to_string(&path).await
|
||||
&& let Ok(caps) = CapabilitiesFile::from_json(&content)
|
||||
&& let Some(auth) = &caps.auth
|
||||
&& auth.secret_name == secret_name
|
||||
&& let Some(oauth) = &auth.oauth
|
||||
{
|
||||
all_scopes.extend(oauth.scopes.iter().cloned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut combined = base_oauth.clone();
|
||||
combined.scopes = all_scopes.into_iter().collect();
|
||||
combined.scopes.sort(); // deterministic ordering
|
||||
combined
|
||||
}
|
||||
|
||||
/// OAuth browser-based login flow.
|
||||
async fn auth_tool_oauth(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
oauth: &crate::tools::wasm::OAuthConfigSchema,
|
||||
@@ -804,12 +906,14 @@ async fn auth_tool_oauth(
|
||||
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use rand::RngCore;
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
|
||||
|
||||
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
|
||||
|
||||
// Get client_id from config or env
|
||||
// Get client_id: capabilities file > runtime env var > built-in defaults
|
||||
let builtin = oauth_defaults::builtin_credentials(&auth.secret_name);
|
||||
|
||||
let client_id = oauth
|
||||
.client_id
|
||||
.clone()
|
||||
@@ -819,41 +923,32 @@ async fn auth_tool_oauth(
|
||||
.as_ref()
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
})
|
||||
.or_else(|| builtin.as_ref().map(|c| c.client_id.to_string()))
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"OAuth client_id not configured.\n\
|
||||
Set it in the capabilities file or via environment variable."
|
||||
Set {} env var, or build with IRONCLAW_GOOGLE_CLIENT_ID.",
|
||||
oauth.client_id_env.as_deref().unwrap_or("the client_id")
|
||||
)
|
||||
})?;
|
||||
|
||||
// Get client_secret if provided
|
||||
let client_secret = oauth.client_secret.clone().or_else(|| {
|
||||
oauth
|
||||
.client_secret_env
|
||||
.as_ref()
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
});
|
||||
// Get client_secret: capabilities file > runtime env var > built-in defaults
|
||||
let client_secret = oauth
|
||||
.client_secret
|
||||
.clone()
|
||||
.or_else(|| {
|
||||
oauth
|
||||
.client_secret_env
|
||||
.as_ref()
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
})
|
||||
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
|
||||
|
||||
println!(" Starting OAuth authentication...");
|
||||
println!();
|
||||
|
||||
// Find an available port for the callback
|
||||
let mut listener = None;
|
||||
let mut port = 0;
|
||||
|
||||
for p in 9876..=9886 {
|
||||
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
|
||||
Ok(l) => {
|
||||
listener = Some(l);
|
||||
port = p;
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let listener = listener.ok_or_else(|| anyhow::anyhow!("Could not find available port"))?;
|
||||
let redirect_uri = format!("http://localhost:{}/callback", port);
|
||||
let listener = oauth_defaults::bind_callback_listener().await?;
|
||||
let redirect_uri = format!("http://localhost:{}/callback", OAUTH_CALLBACK_PORT);
|
||||
|
||||
// Generate PKCE verifier and challenge
|
||||
let (code_verifier, code_challenge) = if oauth.use_pkce {
|
||||
@@ -912,65 +1007,8 @@ async fn auth_tool_oauth(
|
||||
|
||||
println!(" Waiting for authorization...");
|
||||
|
||||
// Wait for callback with timeout
|
||||
let timeout = std::time::Duration::from_secs(300);
|
||||
let code = tokio::time::timeout(timeout, async {
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await?;
|
||||
|
||||
let mut reader = BufReader::new(&mut socket);
|
||||
let mut request_line = String::new();
|
||||
reader.read_line(&mut request_line).await?;
|
||||
|
||||
// Parse GET /callback?code=xxx HTTP/1.1
|
||||
if let Some(path) = request_line.split_whitespace().nth(1) {
|
||||
if path.starts_with("/callback") {
|
||||
if let Some(query) = path.split('?').nth(1) {
|
||||
for param in query.split('&') {
|
||||
let parts: Vec<&str> = param.splitn(2, '=').collect();
|
||||
if parts.len() == 2 && parts[0] == "code" {
|
||||
let code = urlencoding::decode(parts[1])
|
||||
.unwrap_or_else(|_| parts[1].into())
|
||||
.into_owned();
|
||||
|
||||
// Send success response
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\n\
|
||||
Content-Type: text/html\r\n\
|
||||
\r\n\
|
||||
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
|
||||
display: flex; justify-content: center; align-items: center; \
|
||||
height: 100vh; margin: 0; background: #191919; color: white;\">\
|
||||
<div style=\"text-align: center;\">\
|
||||
<h1>✓ {} Connected!</h1>\
|
||||
<p>You can close this window.</p>\
|
||||
</div></body></html>",
|
||||
display_name
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
let _ = socket.shutdown().await;
|
||||
|
||||
return Ok::<_, anyhow::Error>(code);
|
||||
}
|
||||
}
|
||||
|
||||
// Check for error
|
||||
if query.contains("error=") {
|
||||
let response =
|
||||
"HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
return Err(anyhow::anyhow!("Authorization denied by user"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let response = "HTTP/1.1 404 Not Found\r\n\r\n";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Timed out waiting for authorization"))??;
|
||||
let code =
|
||||
oauth_defaults::wait_for_callback(listener, "/callback", "code", display_name).await?;
|
||||
|
||||
println!();
|
||||
println!(" Exchanging code for token...");
|
||||
@@ -1021,8 +1059,19 @@ async fn auth_tool_oauth(
|
||||
)
|
||||
})?;
|
||||
|
||||
// Save the token
|
||||
save_token(store, user_id, auth, access_token).await?;
|
||||
let refresh_token = token_data.get("refresh_token").and_then(|v| v.as_str());
|
||||
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
|
||||
|
||||
// Save the token (with refresh token and expiry if provided)
|
||||
save_token(
|
||||
store,
|
||||
user_id,
|
||||
auth,
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_in,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Extract any additional info for display
|
||||
let workspace_name = token_data
|
||||
@@ -1044,7 +1093,7 @@ async fn auth_tool_oauth(
|
||||
|
||||
/// Manual token entry flow.
|
||||
async fn auth_tool_manual(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
) -> anyhow::Result<()> {
|
||||
@@ -1124,8 +1173,8 @@ async fn auth_tool_manual(
|
||||
}
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(store, user_id, auth, &token).await?;
|
||||
// Save the token (manual path: no refresh token or expiry)
|
||||
save_token(store, user_id, auth, &token, None, None).await?;
|
||||
print_success(display_name);
|
||||
Ok(())
|
||||
}
|
||||
@@ -1216,11 +1265,16 @@ async fn validate_token(
|
||||
}
|
||||
|
||||
/// Save token to secrets store.
|
||||
///
|
||||
/// Optionally stores a refresh token (as `{secret_name}_refresh_token`) and
|
||||
/// sets `expires_at` on the access token so the runtime can auto-refresh.
|
||||
async fn save_token(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
token: &str,
|
||||
refresh_token: Option<&str>,
|
||||
expires_in: Option<u64>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut params = CreateSecretParams::new(&auth.secret_name, token);
|
||||
|
||||
@@ -1228,11 +1282,29 @@ async fn save_token(
|
||||
params = params.with_provider(provider);
|
||||
}
|
||||
|
||||
if let Some(secs) = expires_in {
|
||||
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64);
|
||||
params = params.with_expiry(expires_at);
|
||||
}
|
||||
|
||||
store
|
||||
.create(user_id, params)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?;
|
||||
|
||||
// Store refresh token separately (no expiry, it's long-lived)
|
||||
if let Some(rt) = refresh_token {
|
||||
let refresh_name = format!("{}_refresh_token", auth.secret_name);
|
||||
let mut refresh_params = CreateSecretParams::new(&refresh_name, rt);
|
||||
if let Some(ref provider) = auth.provider {
|
||||
refresh_params = refresh_params.with_provider(provider);
|
||||
}
|
||||
store
|
||||
.create(user_id, refresh_params)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save refresh token: {}", e))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
+296
-76
@@ -1,11 +1,13 @@
|
||||
//! Configuration for IronClaw.
|
||||
//!
|
||||
//! Settings are loaded with priority: env var > database > default.
|
||||
//! The database replaces the old `settings.json` file for all settings
|
||||
//! except the 4 bootstrap fields (database_url, pool_size, secrets key
|
||||
//! source, onboard_completed) which live in `~/.ironclaw/bootstrap.json`.
|
||||
//! `DATABASE_URL` lives in `~/.ironclaw/.env` (loaded via dotenvy early
|
||||
//! in startup). Everything else comes from env vars, the DB settings
|
||||
//! table, or auto-detection.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
@@ -13,6 +15,13 @@ use secrecy::{ExposeSecret, SecretString};
|
||||
use crate::error::ConfigError;
|
||||
use crate::settings::Settings;
|
||||
|
||||
/// Thread-safe overlay for injected env vars (secrets loaded from DB).
|
||||
///
|
||||
/// Used by `inject_llm_keys_from_secrets()` to make API keys available to
|
||||
/// `optional_env()` without unsafe `set_var` calls. `optional_env()` checks
|
||||
/// real env vars first, then falls back to this overlay.
|
||||
static INJECTED_VARS: OnceLock<HashMap<String, String>> = OnceLock::new();
|
||||
|
||||
/// Main configuration for the agent.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
@@ -38,11 +47,11 @@ impl Config {
|
||||
/// Priority: env var > DB settings > default.
|
||||
/// This is the primary way to load config after DB is connected.
|
||||
pub async fn from_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
bootstrap: &crate::bootstrap::BootstrapConfig,
|
||||
) -> Result<Self, ConfigError> {
|
||||
let _ = dotenvy::dotenv();
|
||||
crate::bootstrap::load_ironclaw_env();
|
||||
|
||||
// Load all settings from DB into a Settings struct
|
||||
let db_settings = match store.get_all_settings(user_id).await {
|
||||
@@ -53,7 +62,7 @@ impl Config {
|
||||
}
|
||||
};
|
||||
|
||||
Self::build(bootstrap, &db_settings).await
|
||||
Self::build(&db_settings).await
|
||||
}
|
||||
|
||||
/// Load configuration from environment variables only (no database).
|
||||
@@ -61,20 +70,20 @@ impl Config {
|
||||
/// Used during early startup before the database is connected,
|
||||
/// and by CLI commands that don't have DB access.
|
||||
/// Falls back to legacy `settings.json` on disk if present.
|
||||
///
|
||||
/// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env`
|
||||
/// (lower priority) via dotenvy, which never overwrites existing vars.
|
||||
pub async fn from_env() -> Result<Self, ConfigError> {
|
||||
let _ = dotenvy::dotenv();
|
||||
let bootstrap = crate::bootstrap::BootstrapConfig::load();
|
||||
crate::bootstrap::load_ironclaw_env();
|
||||
let settings = Settings::load();
|
||||
Self::build(&bootstrap, &settings).await
|
||||
Self::build(&settings).await
|
||||
}
|
||||
|
||||
/// Build config from bootstrap + settings (shared by from_env and from_db).
|
||||
async fn build(
|
||||
bootstrap: &crate::bootstrap::BootstrapConfig,
|
||||
settings: &Settings,
|
||||
) -> Result<Self, ConfigError> {
|
||||
/// Build config from settings (shared by from_env and from_db).
|
||||
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
database: DatabaseConfig::resolve(bootstrap)?,
|
||||
database: DatabaseConfig::resolve()?,
|
||||
llm: LlmConfig::resolve(settings)?,
|
||||
embeddings: EmbeddingsConfig::resolve(settings)?,
|
||||
tunnel: TunnelConfig::resolve(settings)?,
|
||||
@@ -82,7 +91,7 @@ impl Config {
|
||||
agent: AgentConfig::resolve(settings)?,
|
||||
safety: SafetyConfig::resolve()?,
|
||||
wasm: WasmConfig::resolve()?,
|
||||
secrets: SecretsConfig::resolve(bootstrap).await?,
|
||||
secrets: SecretsConfig::resolve().await?,
|
||||
builder: BuilderModeConfig::resolve()?,
|
||||
heartbeat: HeartbeatConfig::resolve(settings)?,
|
||||
routines: RoutineConfig::resolve()?,
|
||||
@@ -107,13 +116,13 @@ impl TunnelConfig {
|
||||
let public_url = optional_env("TUNNEL_URL")?
|
||||
.or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty()));
|
||||
|
||||
if let Some(ref url) = public_url {
|
||||
if !url.starts_with("https://") {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "TUNNEL_URL".to_string(),
|
||||
message: "must start with https:// (webhooks require HTTPS)".to_string(),
|
||||
});
|
||||
}
|
||||
if let Some(ref url) = public_url
|
||||
&& !url.starts_with("https://")
|
||||
{
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "TUNNEL_URL".to_string(),
|
||||
message: "must start with https:// (webhooks require HTTPS)".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self { public_url })
|
||||
@@ -134,35 +143,113 @@ impl TunnelConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Which database backend to use.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum DatabaseBackend {
|
||||
/// PostgreSQL via deadpool-postgres (default).
|
||||
#[default]
|
||||
Postgres,
|
||||
/// libSQL/Turso embedded database.
|
||||
LibSql,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for DatabaseBackend {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Postgres => write!(f, "postgres"),
|
||||
Self::LibSql => write!(f, "libsql"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for DatabaseBackend {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"postgres" | "postgresql" | "pg" => Ok(Self::Postgres),
|
||||
"libsql" | "turso" | "sqlite" => Ok(Self::LibSql),
|
||||
_ => Err(format!(
|
||||
"invalid database backend '{}', expected 'postgres' or 'libsql'",
|
||||
s
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Database configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DatabaseConfig {
|
||||
/// Which backend to use (default: Postgres).
|
||||
pub backend: DatabaseBackend,
|
||||
|
||||
// -- PostgreSQL fields --
|
||||
pub url: SecretString,
|
||||
pub pool_size: usize,
|
||||
|
||||
// -- libSQL fields --
|
||||
/// Path to local libSQL database file (default: ~/.ironclaw/ironclaw.db).
|
||||
pub libsql_path: Option<PathBuf>,
|
||||
/// Turso cloud URL for remote sync (optional).
|
||||
pub libsql_url: Option<String>,
|
||||
/// Turso auth token (required when libsql_url is set).
|
||||
pub libsql_auth_token: Option<SecretString>,
|
||||
}
|
||||
|
||||
impl DatabaseConfig {
|
||||
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
let backend: DatabaseBackend = if let Some(b) = optional_env("DATABASE_BACKEND")? {
|
||||
b.parse().map_err(|e| ConfigError::InvalidValue {
|
||||
key: "DATABASE_BACKEND".to_string(),
|
||||
message: e,
|
||||
})?
|
||||
} else {
|
||||
DatabaseBackend::default()
|
||||
};
|
||||
|
||||
// PostgreSQL URL is required only when using the postgres backend.
|
||||
// For libsql backend, default to an empty placeholder.
|
||||
// DATABASE_URL is loaded from ~/.ironclaw/.env via dotenvy early in startup.
|
||||
let url = optional_env("DATABASE_URL")?
|
||||
.or_else(|| bootstrap.database_url.clone())
|
||||
.or_else(|| {
|
||||
if backend == DatabaseBackend::LibSql {
|
||||
Some("unused://libsql".to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "database_url".to_string(),
|
||||
hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(),
|
||||
})?;
|
||||
|
||||
let pool_size = optional_env("DATABASE_POOL_SIZE")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "DATABASE_POOL_SIZE".to_string(),
|
||||
message: format!("must be a positive integer: {e}"),
|
||||
})?
|
||||
.or(bootstrap.database_pool_size)
|
||||
.unwrap_or(10);
|
||||
let pool_size = parse_optional_env("DATABASE_POOL_SIZE", 10)?;
|
||||
|
||||
let libsql_path = optional_env("LIBSQL_PATH")?.map(PathBuf::from).or_else(|| {
|
||||
if backend == DatabaseBackend::LibSql {
|
||||
Some(default_libsql_path())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
let libsql_url = optional_env("LIBSQL_URL")?;
|
||||
let libsql_auth_token = optional_env("LIBSQL_AUTH_TOKEN")?.map(SecretString::from);
|
||||
|
||||
if libsql_url.is_some() && libsql_auth_token.is_none() {
|
||||
return Err(ConfigError::MissingRequired {
|
||||
key: "LIBSQL_AUTH_TOKEN".to_string(),
|
||||
hint: "LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
backend,
|
||||
url: SecretString::from(url),
|
||||
pool_size,
|
||||
libsql_path,
|
||||
libsql_url,
|
||||
libsql_auth_token,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -172,6 +259,14 @@ impl DatabaseConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Default libSQL database path (~/.ironclaw/ironclaw.db).
|
||||
pub fn default_libsql_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("ironclaw.db")
|
||||
}
|
||||
|
||||
/// Which LLM backend to use.
|
||||
///
|
||||
/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem.
|
||||
@@ -302,6 +397,9 @@ impl std::str::FromStr for NearAiApiMode {
|
||||
pub struct NearAiConfig {
|
||||
/// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o")
|
||||
pub model: String,
|
||||
/// Cheap/fast model for lightweight tasks (heartbeat, routing, evaluation).
|
||||
/// Falls back to the main model if not set.
|
||||
pub cheap_model: Option<String>,
|
||||
/// Base URL for the NEAR AI API (default: https://api.near.ai)
|
||||
pub base_url: String,
|
||||
/// Base URL for auth/refresh endpoints (default: https://private.near.ai)
|
||||
@@ -312,16 +410,44 @@ pub struct NearAiConfig {
|
||||
pub api_mode: NearAiApiMode,
|
||||
/// API key for cloud-api (required for chat_completions mode)
|
||||
pub api_key: Option<SecretString>,
|
||||
/// Optional fallback model for failover (default: None).
|
||||
/// When set, a secondary provider is created with this model and wrapped
|
||||
/// in a `FailoverProvider` so transient errors on the primary model
|
||||
/// automatically fall through to the fallback.
|
||||
pub fallback_model: Option<String>,
|
||||
/// Maximum number of retries for transient errors (default: 3).
|
||||
/// With the default of 3, the provider makes up to 4 total attempts
|
||||
/// (1 initial + 3 retries) before giving up.
|
||||
pub max_retries: u32,
|
||||
/// Cooldown duration in seconds for the failover provider (default: 300).
|
||||
/// When a provider accumulates enough consecutive failures it is skipped
|
||||
/// for this many seconds.
|
||||
pub failover_cooldown_secs: u64,
|
||||
/// Number of consecutive retryable failures before a provider enters
|
||||
/// cooldown (default: 3).
|
||||
pub failover_cooldown_threshold: u32,
|
||||
}
|
||||
|
||||
impl LlmConfig {
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
// Determine backend (default: NearAi)
|
||||
// Determine backend: env var > settings > default (NearAi)
|
||||
let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? {
|
||||
b.parse().map_err(|e| ConfigError::InvalidValue {
|
||||
key: "LLM_BACKEND".to_string(),
|
||||
message: e,
|
||||
})?
|
||||
} else if let Some(ref b) = settings.llm_backend {
|
||||
match b.parse() {
|
||||
Ok(backend) => backend,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Invalid llm_backend '{}' in settings: {}. Using default NearAi.",
|
||||
b,
|
||||
e
|
||||
);
|
||||
LlmBackend::NearAi
|
||||
}
|
||||
}
|
||||
} else {
|
||||
LlmBackend::NearAi
|
||||
};
|
||||
@@ -347,6 +473,7 @@ impl LlmConfig {
|
||||
"fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic"
|
||||
.to_string()
|
||||
}),
|
||||
cheap_model: optional_env("NEARAI_CHEAP_MODEL")?,
|
||||
base_url: optional_env("NEARAI_BASE_URL")?
|
||||
.unwrap_or_else(|| "https://cloud-api.near.ai".to_string()),
|
||||
auth_base_url: optional_env("NEARAI_AUTH_URL")?
|
||||
@@ -356,6 +483,10 @@ impl LlmConfig {
|
||||
.unwrap_or_else(default_session_path),
|
||||
api_mode,
|
||||
api_key: nearai_api_key,
|
||||
fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?,
|
||||
max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?,
|
||||
failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?,
|
||||
failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?,
|
||||
};
|
||||
|
||||
// Resolve provider-specific configs based on backend
|
||||
@@ -388,6 +519,7 @@ impl LlmConfig {
|
||||
|
||||
let ollama = if backend == LlmBackend::Ollama {
|
||||
let base_url = optional_env("OLLAMA_BASE_URL")?
|
||||
.or_else(|| settings.ollama_base_url.clone())
|
||||
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string());
|
||||
Some(OllamaConfig { base_url, model })
|
||||
@@ -396,8 +528,9 @@ impl LlmConfig {
|
||||
};
|
||||
|
||||
let openai_compatible = if backend == LlmBackend::OpenAiCompatible {
|
||||
let base_url =
|
||||
optional_env("LLM_BASE_URL")?.ok_or_else(|| ConfigError::MissingRequired {
|
||||
let base_url = optional_env("LLM_BASE_URL")?
|
||||
.or_else(|| settings.openai_compatible_base_url.clone())
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "LLM_BASE_URL".to_string(),
|
||||
hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(),
|
||||
})?;
|
||||
@@ -767,52 +900,41 @@ impl std::fmt::Debug for SecretsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-wide cache for the keychain master key.
|
||||
///
|
||||
/// Avoids re-prompting the OS keychain on every `SecretsConfig::resolve()` call
|
||||
/// (e.g. `Config::from_env()` then `Config::from_db()`). Thread-safe alternative
|
||||
/// to caching in a process env var.
|
||||
impl SecretsConfig {
|
||||
async fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
|
||||
/// Auto-detect secrets master key from env var, then OS keychain.
|
||||
///
|
||||
/// Sequential probe: SECRETS_MASTER_KEY env var first, then OS keychain.
|
||||
/// No saved "source" needed; just try each source in order.
|
||||
async fn resolve() -> Result<Self, ConfigError> {
|
||||
use crate::settings::KeySource;
|
||||
|
||||
let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? {
|
||||
(Some(SecretString::from(env_key)), KeySource::Env)
|
||||
} else {
|
||||
match bootstrap.secrets_master_key_source {
|
||||
KeySource::Keychain => {
|
||||
// Try to load from OS keychain (async on Linux)
|
||||
match crate::secrets::keychain::get_master_key().await {
|
||||
Ok(key_bytes) => {
|
||||
let key_hex: String =
|
||||
key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
(Some(SecretString::from(key_hex)), KeySource::Keychain)
|
||||
}
|
||||
Err(_) => {
|
||||
// Keychain configured but key not found
|
||||
// This might happen if keychain was cleared
|
||||
tracing::warn!(
|
||||
"Secrets configured for keychain but key not found. \
|
||||
Run 'ironclaw onboard' to reconfigure."
|
||||
);
|
||||
(None, KeySource::None)
|
||||
}
|
||||
}
|
||||
// Probe the OS keychain; if a key is stored, use it
|
||||
match crate::secrets::keychain::get_master_key().await {
|
||||
Ok(key_bytes) => {
|
||||
let key_hex: String = key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
(Some(SecretString::from(key_hex)), KeySource::Keychain)
|
||||
}
|
||||
KeySource::Env => {
|
||||
tracing::warn!(
|
||||
"Secrets configured for env var but SECRETS_MASTER_KEY not set."
|
||||
);
|
||||
(None, KeySource::None)
|
||||
}
|
||||
KeySource::None => (None, KeySource::None),
|
||||
Err(_) => (None, KeySource::None),
|
||||
}
|
||||
};
|
||||
|
||||
let enabled = master_key.is_some();
|
||||
|
||||
if let Some(ref key) = master_key {
|
||||
if key.expose_secret().len() < 32 {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "SECRETS_MASTER_KEY".to_string(),
|
||||
message: "must be at least 32 bytes for AES-256-GCM".to_string(),
|
||||
});
|
||||
}
|
||||
if let Some(ref key) = master_key
|
||||
&& key.expose_secret().len() < 32
|
||||
{
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "SECRETS_MASTER_KEY".to_string(),
|
||||
message: "must be at least 32 bytes for AES-256-GCM".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
@@ -1174,6 +1296,36 @@ pub struct ClaudeCodeConfig {
|
||||
pub max_turns: u32,
|
||||
/// Memory limit in MB for Claude Code containers (heavier than workers).
|
||||
pub memory_limit_mb: u64,
|
||||
/// Allowed tool patterns for Claude Code permission settings.
|
||||
///
|
||||
/// Written to `/workspace/.claude/settings.json` before spawning the CLI.
|
||||
/// Provides defense-in-depth: only explicitly listed tools are auto-approved.
|
||||
/// Any new/unknown tools would require interactive approval (which times out
|
||||
/// in the non-interactive container, failing safely).
|
||||
///
|
||||
/// Patterns follow Claude Code syntax: `"Bash(*)"`, `"Read"`, `"Edit(*)"`, etc.
|
||||
pub allowed_tools: Vec<String>,
|
||||
}
|
||||
|
||||
/// Default allowed tools for Claude Code inside containers.
|
||||
///
|
||||
/// These cover all standard Claude Code tools needed for autonomous operation.
|
||||
/// The Docker container provides the primary security boundary; this allowlist
|
||||
/// provides defense-in-depth by preventing any future unknown tools from being
|
||||
/// silently auto-approved.
|
||||
fn default_claude_code_allowed_tools() -> Vec<String> {
|
||||
[
|
||||
"Bash(*)",
|
||||
"Read",
|
||||
"Edit(*)",
|
||||
"Glob",
|
||||
"Grep",
|
||||
"WebFetch(*)",
|
||||
"Task(*)",
|
||||
]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect()
|
||||
}
|
||||
|
||||
impl Default for ClaudeCodeConfig {
|
||||
@@ -1186,11 +1338,24 @@ impl Default for ClaudeCodeConfig {
|
||||
model: "sonnet".to_string(),
|
||||
max_turns: 50,
|
||||
memory_limit_mb: 4096,
|
||||
allowed_tools: default_claude_code_allowed_tools(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaudeCodeConfig {
|
||||
/// Load from environment variables only (used inside containers where
|
||||
/// there is no database or full config).
|
||||
pub fn from_env() -> Self {
|
||||
match Self::resolve() {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to resolve ClaudeCodeConfig: {e}, using defaults");
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
let defaults = Self::default();
|
||||
Ok(Self {
|
||||
@@ -1211,21 +1376,76 @@ impl ClaudeCodeConfig {
|
||||
"CLAUDE_CODE_MEMORY_LIMIT_MB",
|
||||
defaults.memory_limit_mb,
|
||||
)?,
|
||||
allowed_tools: optional_env("CLAUDE_CODE_ALLOWED_TOOLS")?
|
||||
.map(|s| {
|
||||
s.split(',')
|
||||
.map(|t| t.trim().to_string())
|
||||
.filter(|t| !t.is_empty())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or(defaults.allowed_tools),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Load API keys from the encrypted secrets store into a thread-safe overlay.
|
||||
///
|
||||
/// This bridges the gap between secrets stored during onboarding and the
|
||||
/// env-var-first resolution in `LlmConfig::resolve()`. Keys in the overlay
|
||||
/// are read by `optional_env()` before falling back to `std::env::var()`,
|
||||
/// so explicit env vars always win.
|
||||
pub async fn inject_llm_keys_from_secrets(
|
||||
secrets: &dyn crate::secrets::SecretsStore,
|
||||
user_id: &str,
|
||||
) {
|
||||
let mappings = [
|
||||
("llm_openai_api_key", "OPENAI_API_KEY"),
|
||||
("llm_anthropic_api_key", "ANTHROPIC_API_KEY"),
|
||||
("llm_compatible_api_key", "LLM_API_KEY"),
|
||||
];
|
||||
|
||||
let mut injected = HashMap::new();
|
||||
|
||||
for (secret_name, env_var) in mappings {
|
||||
match std::env::var(env_var) {
|
||||
Ok(val) if !val.is_empty() => continue,
|
||||
_ => {}
|
||||
}
|
||||
match secrets.get_decrypted(user_id, secret_name).await {
|
||||
Ok(decrypted) => {
|
||||
injected.insert(env_var.to_string(), decrypted.expose().to_string());
|
||||
tracing::debug!("Loaded secret '{}' for env var '{}'", secret_name, env_var);
|
||||
}
|
||||
Err(_) => {
|
||||
// Secret doesn't exist, that's fine
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = INJECTED_VARS.set(injected);
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
|
||||
// Check real env vars first (always win over injected secrets)
|
||||
match std::env::var(key) {
|
||||
Ok(val) if val.is_empty() => Ok(None),
|
||||
Ok(val) => Ok(Some(val)),
|
||||
Err(std::env::VarError::NotPresent) => Ok(None),
|
||||
Err(e) => Err(ConfigError::ParseError(format!(
|
||||
"failed to read {key}: {e}"
|
||||
))),
|
||||
Ok(val) if val.is_empty() => {}
|
||||
Ok(val) => return Ok(Some(val)),
|
||||
Err(std::env::VarError::NotPresent) => {}
|
||||
Err(e) => {
|
||||
return Err(ConfigError::ParseError(format!(
|
||||
"failed to read {key}: {e}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to thread-safe overlay (secrets injected from DB)
|
||||
if let Some(val) = INJECTED_VARS.get().and_then(|map| map.get(key)) {
|
||||
return Ok(Some(val.clone()));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn parse_optional_env<T>(key: &str, default: T) -> Result<T, ConfigError>
|
||||
|
||||
@@ -45,20 +45,21 @@ impl ContextManager {
|
||||
title: impl Into<String>,
|
||||
description: impl Into<String>,
|
||||
) -> Result<Uuid, JobError> {
|
||||
let contexts = self.contexts.read().await;
|
||||
// Hold write lock for the entire check-insert to prevent TOCTOU races
|
||||
// where two concurrent calls both pass the active_count check.
|
||||
let mut contexts = self.contexts.write().await;
|
||||
let active_count = contexts.values().filter(|c| c.state.is_active()).count();
|
||||
|
||||
if active_count >= self.max_jobs {
|
||||
return Err(JobError::MaxJobsExceeded { max: self.max_jobs });
|
||||
}
|
||||
drop(contexts);
|
||||
|
||||
let context = JobContext::with_user(user_id, title, description);
|
||||
let job_id = context.job_id;
|
||||
contexts.insert(job_id, context);
|
||||
drop(contexts);
|
||||
|
||||
let memory = Memory::new(job_id);
|
||||
|
||||
self.contexts.write().await.insert(job_id, context);
|
||||
self.memories.write().await.insert(job_id, memory);
|
||||
|
||||
Ok(job_id)
|
||||
|
||||
@@ -119,6 +119,10 @@ pub struct JobContext {
|
||||
pub estimated_duration: Option<Duration>,
|
||||
/// Actual cost so far.
|
||||
pub actual_cost: Decimal,
|
||||
/// Total tokens consumed by LLM calls in this job.
|
||||
pub total_tokens_used: u64,
|
||||
/// Maximum tokens allowed per job (0 = unlimited).
|
||||
pub max_tokens: u64,
|
||||
/// When the job was created.
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// When the job was started.
|
||||
@@ -159,6 +163,8 @@ impl JobContext {
|
||||
estimated_cost: None,
|
||||
estimated_duration: None,
|
||||
actual_cost: Decimal::ZERO,
|
||||
total_tokens_used: 0,
|
||||
max_tokens: 0,
|
||||
created_at: Utc::now(),
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
@@ -189,6 +195,14 @@ impl JobContext {
|
||||
};
|
||||
|
||||
self.transitions.push(transition);
|
||||
|
||||
// Cap transition history to prevent unbounded memory growth
|
||||
const MAX_TRANSITIONS: usize = 200;
|
||||
if self.transitions.len() > MAX_TRANSITIONS {
|
||||
let drain_count = self.transitions.len() - MAX_TRANSITIONS;
|
||||
self.transitions.drain(..drain_count);
|
||||
}
|
||||
|
||||
self.state = new_state;
|
||||
|
||||
// Update timestamps
|
||||
@@ -210,6 +224,29 @@ impl JobContext {
|
||||
self.actual_cost += cost;
|
||||
}
|
||||
|
||||
/// Record token usage from an LLM call. Returns an error string if the
|
||||
/// token budget has been exceeded after this addition.
|
||||
pub fn add_tokens(&mut self, tokens: u64) -> Result<(), String> {
|
||||
self.total_tokens_used += tokens;
|
||||
if self.max_tokens > 0 && self.total_tokens_used > self.max_tokens {
|
||||
Err(format!(
|
||||
"Token budget exceeded: used {} of {} allowed tokens",
|
||||
self.total_tokens_used, self.max_tokens
|
||||
))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Check whether the monetary budget has been exceeded.
|
||||
pub fn budget_exceeded(&self) -> bool {
|
||||
if let Some(ref budget) = self.budget {
|
||||
self.actual_cost > *budget
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the duration since the job started.
|
||||
pub fn elapsed(&self) -> Option<Duration> {
|
||||
self.started_at.map(|start| {
|
||||
@@ -274,6 +311,57 @@ mod tests {
|
||||
assert_eq!(ctx.state, JobState::Completed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transition_history_capped() {
|
||||
let mut ctx = JobContext::new("Test", "Transition cap test");
|
||||
// Cycle through Pending -> InProgress -> Stuck -> InProgress -> Stuck ...
|
||||
ctx.transition_to(JobState::InProgress, None).unwrap();
|
||||
for i in 0..250 {
|
||||
ctx.mark_stuck(format!("stuck {}", i)).unwrap();
|
||||
ctx.attempt_recovery().unwrap();
|
||||
}
|
||||
// 1 initial + 250*2 = 501 transitions, should be capped at 200
|
||||
assert!(
|
||||
ctx.transitions.len() <= 200,
|
||||
"transitions should be capped at 200, got {}",
|
||||
ctx.transitions.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_tokens_enforces_budget() {
|
||||
let mut ctx = JobContext::new("Test", "Budget test");
|
||||
ctx.max_tokens = 1000;
|
||||
assert!(ctx.add_tokens(500).is_ok());
|
||||
assert_eq!(ctx.total_tokens_used, 500);
|
||||
assert!(ctx.add_tokens(600).is_err());
|
||||
assert_eq!(ctx.total_tokens_used, 1100); // tokens still recorded
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_tokens_unlimited() {
|
||||
let mut ctx = JobContext::new("Test", "No budget");
|
||||
// max_tokens = 0 means unlimited
|
||||
assert!(ctx.add_tokens(1_000_000).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_exceeded() {
|
||||
let mut ctx = JobContext::new("Test", "Money test");
|
||||
ctx.budget = Some(Decimal::new(100, 0)); // $100
|
||||
assert!(!ctx.budget_exceeded());
|
||||
ctx.add_cost(Decimal::new(50, 0));
|
||||
assert!(!ctx.budget_exceeded());
|
||||
ctx.add_cost(Decimal::new(60, 0));
|
||||
assert!(ctx.budget_exceeded());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_exceeded_none() {
|
||||
let ctx = JobContext::new("Test", "No budget");
|
||||
assert!(!ctx.budget_exceeded()); // No budget = never exceeded
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stuck_recovery() {
|
||||
let mut ctx = JobContext::new("Test", "Test job");
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,549 @@
|
||||
//! SQLite-dialect migrations for the libSQL/Turso backend.
|
||||
//!
|
||||
//! Consolidates all PostgreSQL migrations (V1-V8) into a single SQLite-compatible
|
||||
//! schema. Run once on database creation; idempotent via `IF NOT EXISTS`.
|
||||
|
||||
/// Consolidated schema for libSQL.
|
||||
///
|
||||
/// Translates PostgreSQL types and features:
|
||||
/// - `UUID` -> `TEXT` (store as hex string)
|
||||
/// - `TIMESTAMPTZ` -> `TEXT` (ISO-8601)
|
||||
/// - `JSONB` -> `TEXT` (JSON encoded)
|
||||
/// - `BYTEA` -> `BLOB`
|
||||
/// - `NUMERIC` -> `TEXT` (preserve precision for rust_decimal)
|
||||
/// - `TEXT[]` -> `TEXT` (JSON array)
|
||||
/// - `VECTOR(1536)` -> `F32_BLOB(1536)` (libsql native)
|
||||
/// - `TSVECTOR` -> FTS5 virtual table
|
||||
/// - `BIGSERIAL` -> `INTEGER PRIMARY KEY AUTOINCREMENT`
|
||||
/// - PL/pgSQL functions -> SQLite triggers
|
||||
pub const SCHEMA: &str = r#"
|
||||
|
||||
-- ==================== Migration tracking ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS _migrations (
|
||||
version INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Conversations ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id TEXT PRIMARY KEY,
|
||||
channel TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
thread_id TEXT,
|
||||
started_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
last_activity TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
metadata TEXT NOT NULL DEFAULT '{}'
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_channel ON conversations(channel);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_last_activity ON conversations(last_activity);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_messages (
|
||||
id TEXT PRIMARY KEY,
|
||||
conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_messages_conversation
|
||||
ON conversation_messages(conversation_id);
|
||||
|
||||
-- ==================== Agent Jobs ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
marketplace_job_id TEXT,
|
||||
conversation_id TEXT REFERENCES conversations(id),
|
||||
title TEXT NOT NULL,
|
||||
description TEXT NOT NULL,
|
||||
category TEXT,
|
||||
status TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL DEFAULT 'default',
|
||||
project_dir TEXT,
|
||||
job_mode TEXT NOT NULL DEFAULT 'worker',
|
||||
budget_amount TEXT,
|
||||
budget_token TEXT,
|
||||
bid_amount TEXT,
|
||||
estimated_cost TEXT,
|
||||
estimated_time_secs INTEGER,
|
||||
estimated_value TEXT,
|
||||
actual_cost TEXT,
|
||||
actual_time_secs INTEGER,
|
||||
success INTEGER,
|
||||
failure_reason TEXT,
|
||||
stuck_since TEXT,
|
||||
repair_attempts INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
started_at TEXT,
|
||||
completed_at TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_status ON agent_jobs(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_marketplace ON agent_jobs(marketplace_job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_conversation ON agent_jobs(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_source ON agent_jobs(source);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_user ON agent_jobs(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_created ON agent_jobs(created_at DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS job_actions (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
sequence_num INTEGER NOT NULL,
|
||||
tool_name TEXT NOT NULL,
|
||||
input TEXT NOT NULL,
|
||||
output_raw TEXT,
|
||||
output_sanitized TEXT,
|
||||
sanitization_warnings TEXT,
|
||||
cost TEXT,
|
||||
duration_ms INTEGER,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE(job_id, sequence_num)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_job_actions_job_id ON job_actions(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_job_actions_tool ON job_actions(tool_name);
|
||||
|
||||
-- ==================== Dynamic Tools ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS dynamic_tools (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
description TEXT NOT NULL,
|
||||
parameters_schema TEXT NOT NULL,
|
||||
code TEXT NOT NULL,
|
||||
sandbox_config TEXT NOT NULL,
|
||||
created_by_job_id TEXT REFERENCES agent_jobs(id),
|
||||
success_count INTEGER NOT NULL DEFAULT 0,
|
||||
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_status ON dynamic_tools(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_name ON dynamic_tools(name);
|
||||
|
||||
-- ==================== LLM Calls ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS llm_calls (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
conversation_id TEXT REFERENCES conversations(id),
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
input_tokens INTEGER NOT NULL,
|
||||
output_tokens INTEGER NOT NULL,
|
||||
cost TEXT NOT NULL,
|
||||
purpose TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_job ON llm_calls(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_conversation ON llm_calls(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_provider ON llm_calls(provider);
|
||||
|
||||
-- ==================== Estimation ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS estimation_snapshots (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
category TEXT NOT NULL,
|
||||
tool_names TEXT NOT NULL DEFAULT '[]',
|
||||
estimated_cost TEXT NOT NULL,
|
||||
actual_cost TEXT,
|
||||
estimated_time_secs INTEGER NOT NULL,
|
||||
actual_time_secs INTEGER,
|
||||
estimated_value TEXT NOT NULL,
|
||||
actual_value TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_estimation_category ON estimation_snapshots(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_estimation_job ON estimation_snapshots(job_id);
|
||||
|
||||
-- ==================== Self Repair ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS repair_attempts (
|
||||
id TEXT PRIMARY KEY,
|
||||
target_type TEXT NOT NULL,
|
||||
target_id TEXT NOT NULL,
|
||||
diagnosis TEXT NOT NULL,
|
||||
action_taken TEXT NOT NULL,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_repair_attempts_target ON repair_attempts(target_type, target_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_repair_attempts_created ON repair_attempts(created_at);
|
||||
|
||||
-- ==================== Workspace: Memory Documents ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memory_documents (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
agent_id TEXT,
|
||||
path TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
UNIQUE (user_id, agent_id, path)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_user ON memory_documents(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_path ON memory_documents(user_id, path);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_updated ON memory_documents(updated_at DESC);
|
||||
|
||||
-- Trigger to auto-update updated_at on memory_documents
|
||||
CREATE TRIGGER IF NOT EXISTS update_memory_documents_updated_at
|
||||
AFTER UPDATE ON memory_documents
|
||||
FOR EACH ROW
|
||||
WHEN NEW.updated_at = OLD.updated_at
|
||||
BEGIN
|
||||
UPDATE memory_documents SET updated_at = datetime('now') WHERE id = NEW.id;
|
||||
END;
|
||||
|
||||
-- ==================== Workspace: Memory Chunks ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memory_chunks (
|
||||
_rowid INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
id TEXT NOT NULL UNIQUE,
|
||||
document_id TEXT NOT NULL REFERENCES memory_documents(id) ON DELETE CASCADE,
|
||||
chunk_index INTEGER NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
embedding F32_BLOB(1536),
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (document_id, chunk_index)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_chunks_document ON memory_chunks(document_id);
|
||||
|
||||
-- Vector index for semantic search (libSQL native)
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_chunks_embedding
|
||||
ON memory_chunks (libsql_vector_idx(embedding));
|
||||
|
||||
-- FTS5 virtual table for full-text search
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS memory_chunks_fts USING fts5(
|
||||
content,
|
||||
content='memory_chunks',
|
||||
content_rowid='_rowid'
|
||||
);
|
||||
|
||||
-- Triggers to keep FTS5 in sync with memory_chunks
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_insert AFTER INSERT ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_delete AFTER DELETE ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
|
||||
VALUES ('delete', old._rowid, old.content);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_update AFTER UPDATE ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
|
||||
VALUES ('delete', old._rowid, old.content);
|
||||
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
|
||||
END;
|
||||
|
||||
-- ==================== Workspace: Heartbeat State ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS heartbeat_state (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
agent_id TEXT,
|
||||
last_run TEXT,
|
||||
next_run TEXT,
|
||||
interval_seconds INTEGER NOT NULL DEFAULT 1800,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
||||
last_checks TEXT NOT NULL DEFAULT '{}',
|
||||
UNIQUE (user_id, agent_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_heartbeat_user ON heartbeat_state(user_id);
|
||||
|
||||
-- ==================== Secrets ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS secrets (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
encrypted_value BLOB NOT NULL,
|
||||
key_salt BLOB NOT NULL,
|
||||
provider TEXT,
|
||||
expires_at TEXT,
|
||||
last_used_at TEXT,
|
||||
usage_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_user ON secrets(user_id);
|
||||
|
||||
-- ==================== WASM Tools ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS wasm_tools (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
version TEXT NOT NULL DEFAULT '1.0.0',
|
||||
description TEXT NOT NULL,
|
||||
wasm_binary BLOB NOT NULL,
|
||||
binary_hash BLOB NOT NULL,
|
||||
parameters_schema TEXT NOT NULL,
|
||||
source_url TEXT,
|
||||
trust_level TEXT NOT NULL DEFAULT 'user',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name, version)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_user ON wasm_tools(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_name ON wasm_tools(user_id, name);
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_status ON wasm_tools(status);
|
||||
|
||||
-- ==================== Tool Capabilities ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_capabilities (
|
||||
id TEXT PRIMARY KEY,
|
||||
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||
http_allowlist TEXT NOT NULL DEFAULT '[]',
|
||||
allowed_secrets TEXT NOT NULL DEFAULT '[]',
|
||||
tool_aliases TEXT NOT NULL DEFAULT '{}',
|
||||
requests_per_minute INTEGER NOT NULL DEFAULT 60,
|
||||
requests_per_hour INTEGER NOT NULL DEFAULT 1000,
|
||||
max_request_body_bytes INTEGER NOT NULL DEFAULT 1048576,
|
||||
max_response_body_bytes INTEGER NOT NULL DEFAULT 10485760,
|
||||
workspace_read_prefixes TEXT NOT NULL DEFAULT '[]',
|
||||
http_timeout_secs INTEGER NOT NULL DEFAULT 30,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (wasm_tool_id)
|
||||
);
|
||||
|
||||
-- ==================== Leak Detection Patterns ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS leak_detection_patterns (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
pattern TEXT NOT NULL,
|
||||
severity TEXT NOT NULL DEFAULT 'high',
|
||||
action TEXT NOT NULL DEFAULT 'block',
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Rate Limit State ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_rate_limit_state (
|
||||
id TEXT PRIMARY KEY,
|
||||
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||
user_id TEXT NOT NULL,
|
||||
minute_window_start TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
minute_count INTEGER NOT NULL DEFAULT 0,
|
||||
hour_window_start TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
hour_count INTEGER NOT NULL DEFAULT 0,
|
||||
UNIQUE (wasm_tool_id, user_id)
|
||||
);
|
||||
|
||||
-- ==================== Secret Usage Audit Log ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS secret_usage_log (
|
||||
id TEXT PRIMARY KEY,
|
||||
secret_id TEXT NOT NULL REFERENCES secrets(id) ON DELETE CASCADE,
|
||||
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
target_host TEXT NOT NULL,
|
||||
target_path TEXT,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_user ON secret_usage_log(user_id);
|
||||
|
||||
-- ==================== Leak Detection Events ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS leak_detection_events (
|
||||
id TEXT PRIMARY KEY,
|
||||
pattern_id TEXT REFERENCES leak_detection_patterns(id) ON DELETE SET NULL,
|
||||
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
action_taken TEXT NOT NULL,
|
||||
context_preview TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Tool Failures ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_failures (
|
||||
id TEXT PRIMARY KEY,
|
||||
tool_name TEXT NOT NULL UNIQUE,
|
||||
error_message TEXT,
|
||||
error_count INTEGER DEFAULT 1,
|
||||
first_failure TEXT DEFAULT (datetime('now')),
|
||||
last_failure TEXT DEFAULT (datetime('now')),
|
||||
last_build_result TEXT,
|
||||
repaired_at TEXT,
|
||||
repair_attempts INTEGER DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_name ON tool_failures(tool_name);
|
||||
|
||||
-- ==================== Job Events ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS job_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id),
|
||||
event_type TEXT NOT NULL,
|
||||
data TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_job_events_job ON job_events(job_id, id);
|
||||
|
||||
-- ==================== Routines ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routines (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
user_id TEXT NOT NULL,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
trigger_type TEXT NOT NULL,
|
||||
trigger_config TEXT NOT NULL,
|
||||
action_type TEXT NOT NULL,
|
||||
action_config TEXT NOT NULL,
|
||||
cooldown_secs INTEGER NOT NULL DEFAULT 300,
|
||||
max_concurrent INTEGER NOT NULL DEFAULT 1,
|
||||
dedup_window_secs INTEGER,
|
||||
notify_channel TEXT,
|
||||
notify_user TEXT NOT NULL DEFAULT 'default',
|
||||
notify_on_success INTEGER NOT NULL DEFAULT 0,
|
||||
notify_on_failure INTEGER NOT NULL DEFAULT 1,
|
||||
notify_on_attention INTEGER NOT NULL DEFAULT 1,
|
||||
state TEXT NOT NULL DEFAULT '{}',
|
||||
last_run_at TEXT,
|
||||
next_fire_at TEXT,
|
||||
run_count INTEGER NOT NULL DEFAULT 0,
|
||||
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_user ON routines(user_id);
|
||||
|
||||
-- ==================== Routine Runs ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routine_runs (
|
||||
id TEXT PRIMARY KEY,
|
||||
routine_id TEXT NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
|
||||
trigger_type TEXT NOT NULL,
|
||||
trigger_detail TEXT,
|
||||
started_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
completed_at TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'running',
|
||||
result_summary TEXT,
|
||||
tokens_used INTEGER,
|
||||
job_id TEXT REFERENCES agent_jobs(id),
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_routine_runs_routine ON routine_runs(routine_id);
|
||||
|
||||
-- ==================== Settings ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
user_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
PRIMARY KEY (user_id, key)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_settings_user ON settings(user_id);
|
||||
|
||||
-- ==================== Missing indexes (parity with PostgreSQL) ====================
|
||||
|
||||
-- agent_jobs
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_stuck ON agent_jobs(stuck_since);
|
||||
|
||||
-- secrets
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_provider ON secrets(provider);
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_expires ON secrets(expires_at);
|
||||
|
||||
-- wasm_tools
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_trust ON wasm_tools(trust_level);
|
||||
|
||||
-- tool_capabilities
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_capabilities_tool ON tool_capabilities(wasm_tool_id);
|
||||
|
||||
-- leak_detection_patterns
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_patterns_enabled ON leak_detection_patterns(enabled);
|
||||
|
||||
-- tool_rate_limit_state
|
||||
CREATE INDEX IF NOT EXISTS idx_rate_limit_tool ON tool_rate_limit_state(wasm_tool_id);
|
||||
|
||||
-- secret_usage_log
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_secret ON secret_usage_log(secret_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_tool ON secret_usage_log(wasm_tool_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_created ON secret_usage_log(created_at DESC);
|
||||
|
||||
-- leak_detection_events
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_pattern ON leak_detection_events(pattern_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_tool ON leak_detection_events(wasm_tool_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_user ON leak_detection_events(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_created ON leak_detection_events(created_at DESC);
|
||||
|
||||
-- tool_failures
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_count ON tool_failures(error_count DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_unrepaired ON tool_failures(tool_name);
|
||||
|
||||
-- routines
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_next_fire ON routines(next_fire_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_event_triggers ON routines(user_id);
|
||||
|
||||
-- routine_runs
|
||||
CREATE INDEX IF NOT EXISTS idx_routine_runs_status ON routine_runs(status);
|
||||
|
||||
-- heartbeat_state
|
||||
CREATE INDEX IF NOT EXISTS idx_heartbeat_next_run ON heartbeat_state(next_run);
|
||||
|
||||
-- ==================== Seed data ====================
|
||||
|
||||
-- Pre-populate leak detection patterns (matches PostgreSQL V2 migration).
|
||||
INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, action, enabled, created_at) VALUES
|
||||
('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440005', 'github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440006', 'github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440007', 'stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440008', 'nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440009', 'bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000a', 'pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000b', 'ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000c', 'google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000d', 'slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000e', 'discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000f', 'twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440010', 'sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440012', 'high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn', 1, datetime('now'));
|
||||
|
||||
"#;
|
||||
+538
@@ -0,0 +1,538 @@
|
||||
//! Database abstraction layer.
|
||||
//!
|
||||
//! Provides a backend-agnostic `Database` trait that unifies all persistence
|
||||
//! operations. Two implementations exist behind feature flags:
|
||||
//!
|
||||
//! - `postgres` (default): Uses `deadpool-postgres` + `tokio-postgres`
|
||||
//! - `libsql`: Uses libSQL (Turso's SQLite fork) for embedded/edge deployment
|
||||
//!
|
||||
//! The existing `Store`, `Repository`, `SecretsStore`, and `WasmToolStore`
|
||||
//! types become thin wrappers that delegate to `Arc<dyn Database>`.
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
pub mod postgres;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
pub mod libsql_backend;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
pub mod libsql_migrations;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use rust_decimal::Decimal;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::BrokenTool;
|
||||
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
use crate::error::DatabaseError;
|
||||
use crate::error::WorkspaceError;
|
||||
use crate::history::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, SettingRow,
|
||||
};
|
||||
use crate::workspace::{MemoryChunk, MemoryDocument, WorkspaceEntry};
|
||||
use crate::workspace::{SearchConfig, SearchResult};
|
||||
|
||||
/// Create a database backend from configuration, run migrations, and return it.
|
||||
///
|
||||
/// This is the shared helper for CLI commands and other call sites that need
|
||||
/// a simple `Arc<dyn Database>` without retaining backend-specific handles
|
||||
/// (e.g., `pg_pool` or `libsql_conn` for the secrets store). The main agent
|
||||
/// startup in `main.rs` uses its own initialization block because it also
|
||||
/// captures those backend-specific handles.
|
||||
pub async fn connect_from_config(
|
||||
config: &crate::config::DatabaseConfig,
|
||||
) -> Result<Arc<dyn Database>, DatabaseError> {
|
||||
match config.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
crate::config::DatabaseBackend::LibSql => {
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config.libsql_path.as_deref().unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.libsql_url {
|
||||
let token = config.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
DatabaseError::Pool(
|
||||
"LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set".to_string(),
|
||||
)
|
||||
})?;
|
||||
libsql_backend::LibSqlBackend::new_remote_replica(
|
||||
db_path,
|
||||
url,
|
||||
token.expose_secret(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
} else {
|
||||
libsql_backend::LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
};
|
||||
backend.run_migrations().await?;
|
||||
Ok(Arc::new(backend))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
pg.run_migrations().await?;
|
||||
Ok(Arc::new(pg))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => Err(DatabaseError::Pool(
|
||||
"No database backend available. Enable 'postgres' or 'libsql' feature.".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Backend-agnostic database trait.
|
||||
///
|
||||
/// Combines all persistence operations from Store, Repository, and related
|
||||
/// stores into a single trait that can be implemented for different backends.
|
||||
#[async_trait]
|
||||
pub trait Database: Send + Sync {
|
||||
/// Run schema migrations for this backend.
|
||||
async fn run_migrations(&self) -> Result<(), DatabaseError>;
|
||||
|
||||
// ==================== Conversations ====================
|
||||
|
||||
/// Create a new conversation.
|
||||
async fn create_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Update conversation last activity.
|
||||
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Add a message to a conversation.
|
||||
async fn add_conversation_message(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Ensure a conversation row exists (upsert).
|
||||
async fn ensure_conversation(
|
||||
&self,
|
||||
id: Uuid,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// List conversations with a title preview.
|
||||
async fn list_conversations_with_preview(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ConversationSummary>, DatabaseError>;
|
||||
|
||||
/// Get or create the singleton assistant conversation.
|
||||
async fn get_or_create_assistant_conversation(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Create a conversation with specific metadata.
|
||||
async fn create_conversation_with_metadata(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Load messages with cursor-based pagination.
|
||||
async fn list_conversation_messages_paginated(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
before: Option<DateTime<Utc>>,
|
||||
limit: i64,
|
||||
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError>;
|
||||
|
||||
/// Merge a single key into conversation metadata.
|
||||
async fn update_conversation_metadata_field(
|
||||
&self,
|
||||
id: Uuid,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Read conversation metadata.
|
||||
async fn get_conversation_metadata(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError>;
|
||||
|
||||
/// Load all messages for a conversation.
|
||||
async fn list_conversation_messages(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Vec<ConversationMessage>, DatabaseError>;
|
||||
|
||||
/// Check if a conversation belongs to a specific user.
|
||||
async fn conversation_belongs_to_user(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError>;
|
||||
|
||||
// ==================== Jobs ====================
|
||||
|
||||
/// Save a job context.
|
||||
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get a job by ID.
|
||||
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError>;
|
||||
|
||||
/// Update job status.
|
||||
async fn update_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: JobState,
|
||||
failure_reason: Option<&str>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Mark job as stuck.
|
||||
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get stuck jobs.
|
||||
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError>;
|
||||
|
||||
// ==================== Actions ====================
|
||||
|
||||
/// Save a job action.
|
||||
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get actions for a job.
|
||||
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError>;
|
||||
|
||||
// ==================== LLM Calls ====================
|
||||
|
||||
/// Record an LLM call.
|
||||
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
// ==================== Estimation Snapshots ====================
|
||||
|
||||
/// Save an estimation snapshot.
|
||||
async fn save_estimation_snapshot(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
category: &str,
|
||||
tool_names: &[String],
|
||||
estimated_cost: Decimal,
|
||||
estimated_time_secs: i32,
|
||||
estimated_value: Decimal,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Update estimation snapshot with actual values.
|
||||
async fn update_estimation_actuals(
|
||||
&self,
|
||||
id: Uuid,
|
||||
actual_cost: Decimal,
|
||||
actual_time_secs: i32,
|
||||
actual_value: Option<Decimal>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
// ==================== Sandbox Jobs ====================
|
||||
|
||||
/// Insert a new sandbox job.
|
||||
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get a sandbox job by ID.
|
||||
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError>;
|
||||
|
||||
/// List all sandbox jobs, most recent first.
|
||||
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError>;
|
||||
|
||||
/// Update sandbox job status.
|
||||
async fn update_sandbox_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: &str,
|
||||
success: Option<bool>,
|
||||
message: Option<&str>,
|
||||
started_at: Option<DateTime<Utc>>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Mark stale sandbox jobs as interrupted.
|
||||
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError>;
|
||||
|
||||
/// Get sandbox job summary.
|
||||
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError>;
|
||||
|
||||
/// List sandbox jobs for a specific user, most recent first.
|
||||
async fn list_sandbox_jobs_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<SandboxJobRecord>, DatabaseError>;
|
||||
|
||||
/// Get sandbox job summary for a specific user.
|
||||
async fn sandbox_job_summary_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<SandboxJobSummary, DatabaseError>;
|
||||
|
||||
/// Check if a sandbox job belongs to a specific user.
|
||||
async fn sandbox_job_belongs_to_user(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError>;
|
||||
|
||||
/// Update sandbox job mode.
|
||||
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get sandbox job mode.
|
||||
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError>;
|
||||
|
||||
// ==================== Job Events ====================
|
||||
|
||||
/// Persist a job event.
|
||||
async fn save_job_event(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
event_type: &str,
|
||||
data: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Load all job events.
|
||||
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError>;
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
/// Create a new routine.
|
||||
async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get a routine by ID.
|
||||
async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError>;
|
||||
|
||||
/// Get a routine by user_id and name.
|
||||
async fn get_routine_by_name(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<Option<Routine>, DatabaseError>;
|
||||
|
||||
/// List routines for a user.
|
||||
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError>;
|
||||
|
||||
/// List all enabled event routines.
|
||||
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError>;
|
||||
|
||||
/// List due cron routines.
|
||||
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError>;
|
||||
|
||||
/// Update a routine.
|
||||
async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Update runtime state after a routine fires.
|
||||
async fn update_routine_runtime(
|
||||
&self,
|
||||
id: Uuid,
|
||||
last_run_at: DateTime<Utc>,
|
||||
next_fire_at: Option<DateTime<Utc>>,
|
||||
run_count: u64,
|
||||
consecutive_failures: u32,
|
||||
state: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Delete a routine.
|
||||
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError>;
|
||||
|
||||
// ==================== Routine Runs ====================
|
||||
|
||||
/// Record a routine run starting.
|
||||
async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Complete a routine run.
|
||||
async fn complete_routine_run(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: RunStatus,
|
||||
result_summary: Option<&str>,
|
||||
tokens_used: Option<i32>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// List recent runs for a routine.
|
||||
async fn list_routine_runs(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
limit: i64,
|
||||
) -> Result<Vec<RoutineRun>, DatabaseError>;
|
||||
|
||||
/// Count currently running runs for a routine.
|
||||
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError>;
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
/// Record a tool failure (upsert).
|
||||
async fn record_tool_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
error_message: &str,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Get broken tools exceeding threshold.
|
||||
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError>;
|
||||
|
||||
/// Mark a tool as repaired.
|
||||
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Increment repair attempts.
|
||||
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError>;
|
||||
|
||||
// ==================== Settings ====================
|
||||
|
||||
/// Get a single setting.
|
||||
async fn get_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError>;
|
||||
|
||||
/// Get a single setting with metadata.
|
||||
async fn get_setting_full(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<SettingRow>, DatabaseError>;
|
||||
|
||||
/// Set a single setting (upsert).
|
||||
async fn set_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Delete a single setting.
|
||||
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError>;
|
||||
|
||||
/// List all settings for a user.
|
||||
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError>;
|
||||
|
||||
/// Get all settings as a flat map.
|
||||
async fn get_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<HashMap<String, serde_json::Value>, DatabaseError>;
|
||||
|
||||
/// Bulk-write settings atomically.
|
||||
async fn set_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
settings: &HashMap<String, serde_json::Value>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Check if settings exist for a user.
|
||||
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError>;
|
||||
|
||||
// ==================== Workspace: Documents ====================
|
||||
|
||||
/// Get a document by path.
|
||||
async fn get_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError>;
|
||||
|
||||
/// Get a document by ID.
|
||||
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError>;
|
||||
|
||||
/// Get or create a document by path.
|
||||
async fn get_or_create_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError>;
|
||||
|
||||
/// Update a document's content.
|
||||
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError>;
|
||||
|
||||
/// Delete a document by path.
|
||||
async fn delete_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<(), WorkspaceError>;
|
||||
|
||||
/// List files and directories in a directory path.
|
||||
async fn list_directory(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
directory: &str,
|
||||
) -> Result<Vec<WorkspaceEntry>, WorkspaceError>;
|
||||
|
||||
/// List all file paths in the workspace.
|
||||
async fn list_all_paths(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<String>, WorkspaceError>;
|
||||
|
||||
/// List all documents for a user.
|
||||
async fn list_documents(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<MemoryDocument>, WorkspaceError>;
|
||||
|
||||
// ==================== Workspace: Chunks ====================
|
||||
|
||||
/// Delete all chunks for a document.
|
||||
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError>;
|
||||
|
||||
/// Insert a chunk.
|
||||
async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError>;
|
||||
|
||||
/// Update a chunk's embedding.
|
||||
async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError>;
|
||||
|
||||
/// Get chunks without embeddings for backfilling.
|
||||
async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError>;
|
||||
|
||||
// ==================== Workspace: Search ====================
|
||||
|
||||
/// Perform hybrid search combining FTS and vector similarity.
|
||||
async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError>;
|
||||
}
|
||||
@@ -0,0 +1,627 @@
|
||||
//! PostgreSQL backend for the Database trait.
|
||||
//!
|
||||
//! Delegates to the existing `Store` (history) and `Repository` (workspace)
|
||||
//! implementations, avoiding SQL duplication.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use deadpool_postgres::Pool;
|
||||
use rust_decimal::Decimal;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::BrokenTool;
|
||||
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||
use crate::config::DatabaseConfig;
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::{DatabaseError, WorkspaceError};
|
||||
use crate::history::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, SettingRow, Store,
|
||||
};
|
||||
use crate::workspace::{
|
||||
MemoryChunk, MemoryDocument, Repository, SearchConfig, SearchResult, WorkspaceEntry,
|
||||
};
|
||||
|
||||
/// PostgreSQL database backend.
|
||||
///
|
||||
/// Wraps the existing `Store` (for history/conversations/jobs/routines/settings)
|
||||
/// and `Repository` (for workspace documents/chunks/search) to implement the
|
||||
/// unified `Database` trait.
|
||||
pub struct PgBackend {
|
||||
store: Store,
|
||||
repo: Repository,
|
||||
}
|
||||
|
||||
impl PgBackend {
|
||||
/// Create a new PostgreSQL backend from configuration.
|
||||
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
|
||||
let store = Store::new(config).await?;
|
||||
let repo = Repository::new(store.pool());
|
||||
Ok(Self { store, repo })
|
||||
}
|
||||
|
||||
/// Get a clone of the connection pool.
|
||||
///
|
||||
/// Useful for sharing with components that still need raw pool access.
|
||||
pub fn pool(&self) -> Pool {
|
||||
self.store.pool()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Database for PgBackend {
|
||||
async fn run_migrations(&self) -> Result<(), DatabaseError> {
|
||||
self.store.run_migrations().await
|
||||
}
|
||||
|
||||
// ==================== Conversations ====================
|
||||
|
||||
async fn create_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.create_conversation(channel, user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.store.touch_conversation(id).await
|
||||
}
|
||||
|
||||
async fn add_conversation_message(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.add_conversation_message(conversation_id, role, content)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn ensure_conversation(
|
||||
&self,
|
||||
id: Uuid,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.ensure_conversation(id, channel, user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_conversations_with_preview(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||
self.store
|
||||
.list_conversations_with_preview(user_id, channel, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_or_create_assistant_conversation(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.get_or_create_assistant_conversation(user_id, channel)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_conversation_with_metadata(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.create_conversation_with_metadata(channel, user_id, metadata)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_conversation_messages_paginated(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
before: Option<DateTime<Utc>>,
|
||||
limit: i64,
|
||||
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
|
||||
self.store
|
||||
.list_conversation_messages_paginated(conversation_id, before, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_conversation_metadata_field(
|
||||
&self,
|
||||
id: Uuid,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_conversation_metadata_field(id, key, value)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_conversation_metadata(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
self.store.get_conversation_metadata(id).await
|
||||
}
|
||||
|
||||
async fn list_conversation_messages(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Vec<ConversationMessage>, DatabaseError> {
|
||||
self.store.list_conversation_messages(conversation_id).await
|
||||
}
|
||||
|
||||
async fn conversation_belongs_to_user(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.store
|
||||
.conversation_belongs_to_user(conversation_id, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Jobs ====================
|
||||
|
||||
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
|
||||
self.store.save_job(ctx).await
|
||||
}
|
||||
|
||||
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||
self.store.get_job(id).await
|
||||
}
|
||||
|
||||
async fn update_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: JobState,
|
||||
failure_reason: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_job_status(id, status, failure_reason)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.store.mark_job_stuck(id).await
|
||||
}
|
||||
|
||||
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
|
||||
self.store.get_stuck_jobs().await
|
||||
}
|
||||
|
||||
// ==================== Actions ====================
|
||||
|
||||
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> {
|
||||
self.store.save_action(job_id, action).await
|
||||
}
|
||||
|
||||
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
|
||||
self.store.get_job_actions(job_id).await
|
||||
}
|
||||
|
||||
// ==================== LLM Calls ====================
|
||||
|
||||
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
|
||||
self.store.record_llm_call(record).await
|
||||
}
|
||||
|
||||
// ==================== Estimation Snapshots ====================
|
||||
|
||||
async fn save_estimation_snapshot(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
category: &str,
|
||||
tool_names: &[String],
|
||||
estimated_cost: Decimal,
|
||||
estimated_time_secs: i32,
|
||||
estimated_value: Decimal,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.save_estimation_snapshot(
|
||||
job_id,
|
||||
category,
|
||||
tool_names,
|
||||
estimated_cost,
|
||||
estimated_time_secs,
|
||||
estimated_value,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_estimation_actuals(
|
||||
&self,
|
||||
id: Uuid,
|
||||
actual_cost: Decimal,
|
||||
actual_time_secs: i32,
|
||||
actual_value: Option<Decimal>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Sandbox Jobs ====================
|
||||
|
||||
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||
self.store.save_sandbox_job(job).await
|
||||
}
|
||||
|
||||
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.get_sandbox_job(id).await
|
||||
}
|
||||
|
||||
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.list_sandbox_jobs().await
|
||||
}
|
||||
|
||||
async fn update_sandbox_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: &str,
|
||||
success: Option<bool>,
|
||||
message: Option<&str>,
|
||||
started_at: Option<DateTime<Utc>>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
|
||||
self.store.cleanup_stale_sandbox_jobs().await
|
||||
}
|
||||
|
||||
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
self.store.sandbox_job_summary().await
|
||||
}
|
||||
|
||||
async fn list_sandbox_jobs_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.list_sandbox_jobs_for_user(user_id).await
|
||||
}
|
||||
|
||||
async fn sandbox_job_summary_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
self.store.sandbox_job_summary_for_user(user_id).await
|
||||
}
|
||||
|
||||
async fn sandbox_job_belongs_to_user(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.store
|
||||
.sandbox_job_belongs_to_user(job_id, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
|
||||
self.store.update_sandbox_job_mode(id, mode).await
|
||||
}
|
||||
|
||||
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
|
||||
self.store.get_sandbox_job_mode(id).await
|
||||
}
|
||||
|
||||
// ==================== Job Events ====================
|
||||
|
||||
async fn save_job_event(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
event_type: &str,
|
||||
data: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.save_job_event(job_id, event_type, data).await
|
||||
}
|
||||
|
||||
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError> {
|
||||
self.store.list_job_events(job_id).await
|
||||
}
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
self.store.create_routine(routine).await
|
||||
}
|
||||
|
||||
async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.store.get_routine(id).await
|
||||
}
|
||||
|
||||
async fn get_routine_by_name(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.store.get_routine_by_name(user_id, name).await
|
||||
}
|
||||
|
||||
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_routines(user_id).await
|
||||
}
|
||||
|
||||
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_event_routines().await
|
||||
}
|
||||
|
||||
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_due_cron_routines().await
|
||||
}
|
||||
|
||||
async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
self.store.update_routine(routine).await
|
||||
}
|
||||
|
||||
async fn update_routine_runtime(
|
||||
&self,
|
||||
id: Uuid,
|
||||
last_run_at: DateTime<Utc>,
|
||||
next_fire_at: Option<DateTime<Utc>>,
|
||||
run_count: u64,
|
||||
consecutive_failures: u32,
|
||||
state: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_routine_runtime(
|
||||
id,
|
||||
last_run_at,
|
||||
next_fire_at,
|
||||
run_count,
|
||||
consecutive_failures,
|
||||
state,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
|
||||
self.store.delete_routine(id).await
|
||||
}
|
||||
|
||||
// ==================== Routine Runs ====================
|
||||
|
||||
async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> {
|
||||
self.store.create_routine_run(run).await
|
||||
}
|
||||
|
||||
async fn complete_routine_run(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: RunStatus,
|
||||
result_summary: Option<&str>,
|
||||
tokens_used: Option<i32>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.complete_routine_run(id, status, result_summary, tokens_used)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_routine_runs(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
limit: i64,
|
||||
) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
self.store.list_routine_runs(routine_id, limit).await
|
||||
}
|
||||
|
||||
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
|
||||
self.store.count_running_routine_runs(routine_id).await
|
||||
}
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
async fn record_tool_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
error_message: &str,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.record_tool_failure(tool_name, error_message)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
|
||||
self.store.get_broken_tools(threshold).await
|
||||
}
|
||||
|
||||
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.store.mark_tool_repaired(tool_name).await
|
||||
}
|
||||
|
||||
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.store.increment_repair_attempts(tool_name).await
|
||||
}
|
||||
|
||||
// ==================== Settings ====================
|
||||
|
||||
async fn get_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
self.store.get_setting(user_id, key).await
|
||||
}
|
||||
|
||||
async fn get_setting_full(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<SettingRow>, DatabaseError> {
|
||||
self.store.get_setting_full(user_id, key).await
|
||||
}
|
||||
|
||||
async fn set_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.set_setting(user_id, key, value).await
|
||||
}
|
||||
|
||||
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.delete_setting(user_id, key).await
|
||||
}
|
||||
|
||||
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError> {
|
||||
self.store.list_settings(user_id).await
|
||||
}
|
||||
|
||||
async fn get_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
|
||||
self.store.get_all_settings(user_id).await
|
||||
}
|
||||
|
||||
async fn set_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
settings: &HashMap<String, serde_json::Value>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.set_all_settings(user_id, settings).await
|
||||
}
|
||||
|
||||
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.has_settings(user_id).await
|
||||
}
|
||||
|
||||
// ==================== Workspace: Documents ====================
|
||||
|
||||
async fn get_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
.get_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo.get_document_by_id(id).await
|
||||
}
|
||||
|
||||
async fn get_or_create_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
.get_or_create_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
|
||||
self.repo.update_document(id, content).await
|
||||
}
|
||||
|
||||
async fn delete_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<(), WorkspaceError> {
|
||||
self.repo
|
||||
.delete_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_directory(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
directory: &str,
|
||||
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
|
||||
self.repo.list_directory(user_id, agent_id, directory).await
|
||||
}
|
||||
|
||||
async fn list_all_paths(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<String>, WorkspaceError> {
|
||||
self.repo.list_all_paths(user_id, agent_id).await
|
||||
}
|
||||
|
||||
async fn list_documents(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
|
||||
self.repo.list_documents(user_id, agent_id).await
|
||||
}
|
||||
|
||||
// ==================== Workspace: Chunks ====================
|
||||
|
||||
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
self.repo.delete_chunks(document_id).await
|
||||
}
|
||||
|
||||
async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError> {
|
||||
self.repo
|
||||
.insert_chunk(document_id, chunk_index, content, embedding)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError> {
|
||||
self.repo.update_chunk_embedding(chunk_id, embedding).await
|
||||
}
|
||||
|
||||
async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
|
||||
self.repo
|
||||
.get_chunks_without_embeddings(user_id, agent_id, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Workspace: Search ====================
|
||||
|
||||
async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
self.repo
|
||||
.hybrid_search(user_id, agent_id, query, embedding, config)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -40,6 +40,9 @@ pub enum Error {
|
||||
#[error("Workspace error: {0}")]
|
||||
Workspace(#[from] WorkspaceError),
|
||||
|
||||
#[error("Hook error: {0}")]
|
||||
Hook(#[from] crate::hooks::HookError),
|
||||
|
||||
#[error("Orchestrator error: {0}")]
|
||||
Orchestrator(#[from] OrchestratorError),
|
||||
|
||||
@@ -87,14 +90,21 @@ pub enum DatabaseError {
|
||||
#[error("Serialization error: {0}")]
|
||||
Serialization(String),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("PostgreSQL error: {0}")]
|
||||
Postgres(#[from] tokio_postgres::Error),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("Pool build error: {0}")]
|
||||
PoolBuild(#[from] deadpool_postgres::BuildError),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("Pool runtime error: {0}")]
|
||||
PoolRuntime(#[from] deadpool_postgres::PoolError),
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[error("LibSQL error: {0}")]
|
||||
LibSql(#[from] libsql::Error),
|
||||
}
|
||||
|
||||
/// Channel-related errors.
|
||||
|
||||
@@ -20,10 +20,6 @@ impl CostEstimator {
|
||||
|
||||
// Default tool costs (in USD or equivalent)
|
||||
tool_costs.insert("http".to_string(), dec!(0.0001)); // API call
|
||||
tool_costs.insert("marketplace".to_string(), dec!(0.01)); // Gas costs
|
||||
tool_costs.insert("ecommerce".to_string(), dec!(0.001)); // API call
|
||||
tool_costs.insert("taskrabbit".to_string(), dec!(0.0)); // Cost comes from task itself
|
||||
tool_costs.insert("restaurant".to_string(), dec!(0.001)); // API call
|
||||
tool_costs.insert("echo".to_string(), dec!(0.0)); // Free
|
||||
tool_costs.insert("time".to_string(), dec!(0.0)); // Free
|
||||
tool_costs.insert("json".to_string(), dec!(0.0)); // Free
|
||||
@@ -74,7 +70,7 @@ mod tests {
|
||||
let estimator = CostEstimator::new();
|
||||
|
||||
assert_eq!(estimator.estimate_tool("echo"), dec!(0.0));
|
||||
assert_eq!(estimator.estimate_tool("marketplace"), dec!(0.01));
|
||||
assert_eq!(estimator.estimate_tool("http"), dec!(0.0001));
|
||||
assert!(estimator.estimate_tool("unknown") > dec!(0.0));
|
||||
}
|
||||
|
||||
|
||||
@@ -16,10 +16,6 @@ impl TimeEstimator {
|
||||
|
||||
// Default tool durations
|
||||
tool_durations.insert("http".to_string(), Duration::from_secs(5));
|
||||
tool_durations.insert("marketplace".to_string(), Duration::from_secs(10));
|
||||
tool_durations.insert("ecommerce".to_string(), Duration::from_secs(8));
|
||||
tool_durations.insert("taskrabbit".to_string(), Duration::from_secs(30)); // Just API, not task itself
|
||||
tool_durations.insert("restaurant".to_string(), Duration::from_secs(5));
|
||||
tool_durations.insert("echo".to_string(), Duration::from_millis(10));
|
||||
tool_durations.insert("time".to_string(), Duration::from_millis(1));
|
||||
tool_durations.insert("json".to_string(), Duration::from_millis(5));
|
||||
|
||||
@@ -144,12 +144,11 @@ impl SuccessEvaluator for RuleBasedEvaluator {
|
||||
|
||||
// Check for critical errors
|
||||
for action in actions.iter().filter(|a| !a.success) {
|
||||
if let Some(ref error) = action.error {
|
||||
if error.to_lowercase().contains("critical")
|
||||
|| error.to_lowercase().contains("fatal")
|
||||
{
|
||||
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
|
||||
}
|
||||
if let Some(ref error) = action.error
|
||||
&& (error.to_lowercase().contains("critical")
|
||||
|| error.to_lowercase().contains("fatal"))
|
||||
{
|
||||
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+64
-27
@@ -57,7 +57,7 @@ pub struct ExtensionManager {
|
||||
_tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
/// Optional database store for DB-backed MCP config.
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
}
|
||||
|
||||
impl ExtensionManager {
|
||||
@@ -71,7 +71,7 @@ impl ExtensionManager {
|
||||
wasm_channels_dir: PathBuf,
|
||||
tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
registry: ExtensionRegistry::new(),
|
||||
@@ -351,7 +351,7 @@ impl ExtensionManager {
|
||||
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
||||
{
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store, &self.user_id).await
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), &self.user_id).await
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
}
|
||||
@@ -375,7 +375,8 @@ impl ExtensionManager {
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
config.validate()?;
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::add_mcp_server_db(store, &self.user_id, config).await
|
||||
crate::tools::mcp::config::add_mcp_server_db(store.as_ref(), &self.user_id, config)
|
||||
.await
|
||||
} else {
|
||||
crate::tools::mcp::config::add_mcp_server(config).await
|
||||
}
|
||||
@@ -386,7 +387,8 @@ impl ExtensionManager {
|
||||
name: &str,
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::remove_mcp_server_db(store, &self.user_id, name).await
|
||||
crate::tools::mcp::config::remove_mcp_server_db(store.as_ref(), &self.user_id, name)
|
||||
.await
|
||||
} else {
|
||||
crate::tools::mcp::config::remove_mcp_server(name).await
|
||||
}
|
||||
@@ -461,7 +463,16 @@ impl ExtensionManager {
|
||||
name: &str,
|
||||
url: &str,
|
||||
) -> Result<InstallResult, ExtensionError> {
|
||||
// Download the WASM binary
|
||||
// Require HTTPS to prevent downgrade attacks
|
||||
if !url.starts_with("https://") {
|
||||
return Err(ExtensionError::InstallFailed(
|
||||
"Only HTTPS URLs are allowed for extension downloads".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 50 MB cap to prevent disk-fill DoS
|
||||
const MAX_WASM_SIZE: usize = 50 * 1024 * 1024;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(60))
|
||||
.build()
|
||||
@@ -480,11 +491,36 @@ impl ExtensionManager {
|
||||
)));
|
||||
}
|
||||
|
||||
// Check Content-Length header before downloading the full body
|
||||
if let Some(len) = response.content_length()
|
||||
&& len as usize > MAX_WASM_SIZE
|
||||
{
|
||||
return Err(ExtensionError::InstallFailed(format!(
|
||||
"WASM binary too large ({} bytes, max {} bytes)",
|
||||
len, MAX_WASM_SIZE
|
||||
)));
|
||||
}
|
||||
|
||||
let bytes = response
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| ExtensionError::DownloadFailed(e.to_string()))?;
|
||||
|
||||
if bytes.len() > MAX_WASM_SIZE {
|
||||
return Err(ExtensionError::InstallFailed(format!(
|
||||
"WASM binary too large ({} bytes, max {} bytes)",
|
||||
bytes.len(),
|
||||
MAX_WASM_SIZE
|
||||
)));
|
||||
}
|
||||
|
||||
// Basic WASM magic number check (\0asm)
|
||||
if bytes.len() < 4 || &bytes[..4] != b"\0asm" {
|
||||
return Err(ExtensionError::InstallFailed(
|
||||
"Downloaded file is not a valid WASM binary (bad magic number)".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Ensure tools directory exists
|
||||
tokio::fs::create_dir_all(&self.wasm_tools_dir)
|
||||
.await
|
||||
@@ -497,9 +533,10 @@ impl ExtensionManager {
|
||||
.map_err(|e| ExtensionError::InstallFailed(e.to_string()))?;
|
||||
|
||||
tracing::info!(
|
||||
"Installed WASM tool '{}' ({} bytes) to {}",
|
||||
"Installed WASM tool '{}' ({} bytes) from {} to {}",
|
||||
name,
|
||||
bytes.len(),
|
||||
url,
|
||||
wasm_path.display()
|
||||
);
|
||||
|
||||
@@ -731,27 +768,27 @@ impl ExtensionManager {
|
||||
};
|
||||
|
||||
// Check env var first
|
||||
if let Some(ref env_var) = auth.env_var {
|
||||
if let Ok(value) = std::env::var(env_var) {
|
||||
// Store the env var value as a secret
|
||||
let params = CreateSecretParams::new(&auth.secret_name, &value)
|
||||
.with_provider(name.to_string());
|
||||
self.secrets
|
||||
.create(&self.user_id, params)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
if let Some(ref env_var) = auth.env_var
|
||||
&& let Ok(value) = std::env::var(env_var)
|
||||
{
|
||||
// Store the env var value as a secret
|
||||
let params =
|
||||
CreateSecretParams::new(&auth.secret_name, &value).with_provider(name.to_string());
|
||||
self.secrets
|
||||
.create(&self.user_id, params)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
|
||||
return Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::WasmTool,
|
||||
auth_url: None,
|
||||
callback_type: None,
|
||||
instructions: None,
|
||||
setup_url: None,
|
||||
awaiting_token: false,
|
||||
status: "authenticated".to_string(),
|
||||
});
|
||||
}
|
||||
return Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::WasmTool,
|
||||
auth_url: None,
|
||||
callback_type: None,
|
||||
instructions: None,
|
||||
setup_url: None,
|
||||
awaiting_token: false,
|
||||
status: "authenticated".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check if already authenticated
|
||||
|
||||
+5
-1
@@ -5,11 +5,15 @@
|
||||
//! - Learning from past executions
|
||||
//! - Analytics and metrics
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
mod analytics;
|
||||
mod store;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use analytics::{JobStats, ToolStats};
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use store::Store;
|
||||
pub use store::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, Store,
|
||||
SandboxJobSummary, SettingRow,
|
||||
};
|
||||
|
||||
@@ -1,13 +1,18 @@
|
||||
//! PostgreSQL store for persisting agent data.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::{Config, Pool, Runtime};
|
||||
use rust_decimal::Decimal;
|
||||
#[cfg(feature = "postgres")]
|
||||
use tokio_postgres::NoTls;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::config::DatabaseConfig;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::error::DatabaseError;
|
||||
|
||||
/// Record for an LLM call to be persisted.
|
||||
@@ -24,11 +29,18 @@ pub struct LlmCallRecord<'a> {
|
||||
}
|
||||
|
||||
/// Database store for the agent.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct Store {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Wrap an existing pool (useful when the caller already has a connection).
|
||||
pub fn from_pool(pool: Pool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// Create a new store and connect to the database.
|
||||
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
|
||||
let mut cfg = Config::new();
|
||||
@@ -144,7 +156,12 @@ impl Store {
|
||||
actual_cost, repair_attempts, created_at, started_at, completed_at
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
|
||||
ON CONFLICT (id) DO UPDATE SET
|
||||
title = EXCLUDED.title,
|
||||
description = EXCLUDED.description,
|
||||
category = EXCLUDED.category,
|
||||
status = EXCLUDED.status,
|
||||
estimated_cost = EXCLUDED.estimated_cost,
|
||||
estimated_time_secs = EXCLUDED.estimated_time_secs,
|
||||
actual_cost = EXCLUDED.actual_cost,
|
||||
repair_attempts = EXCLUDED.repair_attempts,
|
||||
started_at = EXCLUDED.started_at,
|
||||
@@ -220,6 +237,8 @@ impl Store {
|
||||
completed_at: row.get("completed_at"),
|
||||
transitions: Vec::new(), // Not loaded from DB for now
|
||||
metadata: serde_json::Value::Null,
|
||||
total_tokens_used: 0,
|
||||
max_tokens: 0,
|
||||
}))
|
||||
}
|
||||
None => Ok(None),
|
||||
@@ -464,6 +483,7 @@ pub struct SandboxJobSummary {
|
||||
pub interrupted: usize,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Insert a new sandbox job into `agent_jobs`.
|
||||
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||
@@ -565,6 +585,90 @@ impl Store {
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// List sandbox jobs for a specific user, most recent first.
|
||||
pub async fn list_sandbox_jobs_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, title, status, user_id, project_dir,
|
||||
success, failure_reason, created_at, started_at, completed_at
|
||||
FROM agent_jobs WHERE source = 'sandbox' AND user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
&[&user_id],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|r| SandboxJobRecord {
|
||||
id: r.get("id"),
|
||||
task: r.get("title"),
|
||||
status: r.get("status"),
|
||||
user_id: r.get("user_id"),
|
||||
project_dir: r
|
||||
.get::<_, Option<String>>("project_dir")
|
||||
.unwrap_or_default(),
|
||||
success: r.get("success"),
|
||||
failure_reason: r.get("failure_reason"),
|
||||
created_at: r.get("created_at"),
|
||||
started_at: r.get("started_at"),
|
||||
completed_at: r.get("completed_at"),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Get a summary of sandbox job counts by status for a specific user.
|
||||
pub async fn sandbox_job_summary_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = conn
|
||||
.query(
|
||||
"SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' AND user_id = $1 GROUP BY status",
|
||||
&[&user_id],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut summary = SandboxJobSummary::default();
|
||||
for row in &rows {
|
||||
let status: String = row.get("status");
|
||||
let count: i64 = row.get("cnt");
|
||||
let c = count as usize;
|
||||
summary.total += c;
|
||||
match status.as_str() {
|
||||
"creating" => summary.creating += c,
|
||||
"running" => summary.running += c,
|
||||
"completed" => summary.completed += c,
|
||||
"failed" => summary.failed += c,
|
||||
"interrupted" => summary.interrupted += c,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
/// Check if a sandbox job belongs to a specific user.
|
||||
pub async fn sandbox_job_belongs_to_user(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt(
|
||||
"SELECT 1 FROM agent_jobs WHERE id = $1 AND user_id = $2 AND source = 'sandbox'",
|
||||
&[&job_id, &user_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// Update sandbox job status and optional timestamps/result.
|
||||
pub async fn update_sandbox_job_status(
|
||||
&self,
|
||||
@@ -656,6 +760,7 @@ pub struct JobEventRecord {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Persist a job event (fire-and-forget from orchestrator handler).
|
||||
pub async fn save_job_event(
|
||||
@@ -728,10 +833,12 @@ impl Store {
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
|
||||
};
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Create a new routine.
|
||||
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
@@ -1032,6 +1139,7 @@ impl Store {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
|
||||
let trigger_type: String = row.get("trigger_type");
|
||||
let trigger_config: serde_json::Value = row.get("trigger_config");
|
||||
@@ -1076,6 +1184,7 @@ fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_routine_run(row: &tokio_postgres::Row) -> Result<RoutineRun, DatabaseError> {
|
||||
let status_str: String = row.get("status");
|
||||
let status: RunStatus = status_str
|
||||
@@ -1121,6 +1230,7 @@ pub struct ConversationMessage {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Ensure a conversation row exists for a given UUID.
|
||||
///
|
||||
@@ -1258,6 +1368,22 @@ impl Store {
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// Check whether a conversation belongs to the given user.
|
||||
pub async fn conversation_belongs_to_user(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt(
|
||||
"SELECT 1 FROM conversations WHERE id = $1 AND user_id = $2",
|
||||
&[&conversation_id, &user_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// Load messages for a conversation with cursor-based pagination.
|
||||
///
|
||||
/// Returns `(messages_oldest_first, has_more)`.
|
||||
@@ -1375,6 +1501,7 @@ impl Store {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn parse_job_state(s: &str) -> JobState {
|
||||
match s {
|
||||
"pending" => JobState::Pending,
|
||||
@@ -1391,8 +1518,10 @@ fn parse_job_state(s: &str) -> JobState {
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::agent::BrokenTool;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Record a tool failure (upsert: increment count if exists).
|
||||
pub async fn record_tool_failure(
|
||||
@@ -1486,6 +1615,7 @@ pub struct SettingRow {
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Get a single setting by key.
|
||||
pub async fn get_setting(
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
//! Core hook types and traits.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
/// Points in the agent lifecycle where hooks can be attached.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum HookPoint {
|
||||
/// Before processing an inbound user message.
|
||||
BeforeInbound,
|
||||
/// Before executing a tool call.
|
||||
BeforeToolCall,
|
||||
/// Before sending an outbound response.
|
||||
BeforeOutbound,
|
||||
/// When a new session starts.
|
||||
OnSessionStart,
|
||||
/// When a session ends (pruned or expired).
|
||||
OnSessionEnd,
|
||||
/// Transform the final response before completing a turn.
|
||||
TransformResponse,
|
||||
}
|
||||
|
||||
/// Contextual data carried with each hook invocation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum HookEvent {
|
||||
/// An inbound user message about to be processed.
|
||||
Inbound {
|
||||
user_id: String,
|
||||
channel: String,
|
||||
content: String,
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
/// A tool call about to be executed.
|
||||
ToolCall {
|
||||
tool_name: String,
|
||||
parameters: serde_json::Value,
|
||||
user_id: String,
|
||||
/// "chat" for interactive, or a job ID string for autonomous jobs.
|
||||
context: String,
|
||||
},
|
||||
/// An outbound response about to be sent.
|
||||
Outbound {
|
||||
user_id: String,
|
||||
channel: String,
|
||||
content: String,
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
/// A new session was created.
|
||||
SessionStart { user_id: String, session_id: String },
|
||||
/// A session was ended (pruned).
|
||||
SessionEnd { user_id: String, session_id: String },
|
||||
/// The final response is being transformed before completing a turn.
|
||||
ResponseTransform {
|
||||
user_id: String,
|
||||
thread_id: String,
|
||||
response: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl HookEvent {
|
||||
/// Returns the [`HookPoint`] this event corresponds to.
|
||||
pub fn hook_point(&self) -> HookPoint {
|
||||
match self {
|
||||
HookEvent::Inbound { .. } => HookPoint::BeforeInbound,
|
||||
HookEvent::ToolCall { .. } => HookPoint::BeforeToolCall,
|
||||
HookEvent::Outbound { .. } => HookPoint::BeforeOutbound,
|
||||
HookEvent::SessionStart { .. } => HookPoint::OnSessionStart,
|
||||
HookEvent::SessionEnd { .. } => HookPoint::OnSessionEnd,
|
||||
HookEvent::ResponseTransform { .. } => HookPoint::TransformResponse,
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply a modification string to the event's primary content field.
|
||||
pub fn apply_modification(&mut self, modified: &str) {
|
||||
match self {
|
||||
HookEvent::Inbound { content, .. } | HookEvent::Outbound { content, .. } => {
|
||||
*content = modified.to_string();
|
||||
}
|
||||
HookEvent::ToolCall { parameters, .. } => match serde_json::from_str(modified) {
|
||||
Ok(parsed) => *parameters = parsed,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
},
|
||||
HookEvent::ResponseTransform { response, .. } => {
|
||||
*response = modified.to_string();
|
||||
}
|
||||
HookEvent::SessionStart { .. } | HookEvent::SessionEnd { .. } => {
|
||||
// Session events don't have modifiable content
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The result of executing a hook.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum HookOutcome {
|
||||
/// Continue processing, optionally with modified content.
|
||||
Continue {
|
||||
/// If `Some`, replace the event's primary content with this value.
|
||||
modified: Option<String>,
|
||||
},
|
||||
/// Reject the event entirely.
|
||||
Reject {
|
||||
/// Human-readable reason for the rejection.
|
||||
reason: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl HookOutcome {
|
||||
/// Shorthand for `Continue { modified: None }`.
|
||||
pub fn ok() -> Self {
|
||||
HookOutcome::Continue { modified: None }
|
||||
}
|
||||
|
||||
/// Shorthand for `Continue { modified: Some(value) }`.
|
||||
pub fn modify(value: String) -> Self {
|
||||
HookOutcome::Continue {
|
||||
modified: Some(value),
|
||||
}
|
||||
}
|
||||
|
||||
/// Shorthand for `Reject { reason }`.
|
||||
pub fn reject(reason: impl Into<String>) -> Self {
|
||||
HookOutcome::Reject {
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// How to handle hook execution failures.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum HookFailureMode {
|
||||
/// On error/timeout, continue processing as if the hook returned `ok()`.
|
||||
FailOpen,
|
||||
/// On error/timeout, reject the event.
|
||||
FailClosed,
|
||||
}
|
||||
|
||||
/// Hook execution errors.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum HookError {
|
||||
#[error("Hook execution failed: {reason}")]
|
||||
ExecutionFailed { reason: String },
|
||||
|
||||
#[error("Hook timed out after {timeout:?}")]
|
||||
Timeout { timeout: Duration },
|
||||
|
||||
#[error("Hook rejected: {reason}")]
|
||||
Rejected { reason: String },
|
||||
}
|
||||
|
||||
/// Context passed to hooks alongside the event.
|
||||
pub struct HookContext {
|
||||
/// Arbitrary metadata hooks can use.
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
impl Default for HookContext {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
metadata: serde_json::Value::Null,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for implementing lifecycle hooks.
|
||||
///
|
||||
/// Hooks intercept and can modify agent operations at well-defined points.
|
||||
#[async_trait]
|
||||
pub trait Hook: Send + Sync {
|
||||
/// A unique name for this hook.
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// The lifecycle points this hook should be called at.
|
||||
fn hook_points(&self) -> &[HookPoint];
|
||||
|
||||
/// How to handle failures in this hook.
|
||||
///
|
||||
/// Default: `FailOpen` (continue on error).
|
||||
fn failure_mode(&self) -> HookFailureMode {
|
||||
HookFailureMode::FailOpen
|
||||
}
|
||||
|
||||
/// Maximum time this hook is allowed to run.
|
||||
///
|
||||
/// Default: 5 seconds.
|
||||
fn timeout(&self) -> Duration {
|
||||
Duration::from_secs(5)
|
||||
}
|
||||
|
||||
/// Execute the hook.
|
||||
async fn execute(&self, event: &HookEvent, ctx: &HookContext)
|
||||
-> Result<HookOutcome, HookError>;
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Lifecycle hooks for intercepting and transforming agent operations.
|
||||
//!
|
||||
//! The hook system provides 6 well-defined interception points:
|
||||
//!
|
||||
//! - **BeforeInbound** — Before processing an inbound user message
|
||||
//! - **BeforeToolCall** — Before executing a tool call
|
||||
//! - **BeforeOutbound** — Before sending an outbound response
|
||||
//! - **OnSessionStart** — When a new session starts
|
||||
//! - **OnSessionEnd** — When a session ends
|
||||
//! - **TransformResponse** — Transform the final response before completing a turn
|
||||
//!
|
||||
//! Hooks are executed in priority order (lower number = higher priority).
|
||||
//! Each hook can pass through, modify content, or reject the event.
|
||||
|
||||
pub mod hook;
|
||||
pub mod registry;
|
||||
|
||||
pub use hook::{Hook, HookContext, HookError, HookEvent, HookFailureMode, HookOutcome, HookPoint};
|
||||
pub use registry::HookRegistry;
|
||||
@@ -0,0 +1,555 @@
|
||||
//! Hook registry for managing and executing lifecycle hooks.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::hooks::hook::{Hook, HookContext, HookError, HookEvent, HookFailureMode, HookOutcome};
|
||||
|
||||
/// A registered hook with its priority.
|
||||
struct HookEntry {
|
||||
hook: Arc<dyn Hook>,
|
||||
priority: u32,
|
||||
}
|
||||
|
||||
/// Registry that manages hooks and executes them at lifecycle points.
|
||||
///
|
||||
/// Hooks are executed in priority order (lower number = higher priority).
|
||||
/// A `Reject` outcome stops the chain immediately.
|
||||
/// A `Modify` outcome chains through subsequent hooks.
|
||||
pub struct HookRegistry {
|
||||
hooks: RwLock<Vec<HookEntry>>,
|
||||
}
|
||||
|
||||
impl HookRegistry {
|
||||
/// Create an empty registry.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
hooks: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a hook with default priority (100).
|
||||
pub async fn register(&self, hook: Arc<dyn Hook>) {
|
||||
self.register_with_priority(hook, 100).await;
|
||||
}
|
||||
|
||||
/// Register a hook with a specific priority.
|
||||
///
|
||||
/// Lower priority number = runs first.
|
||||
pub async fn register_with_priority(&self, hook: Arc<dyn Hook>, priority: u32) {
|
||||
let mut hooks = self.hooks.write().await;
|
||||
hooks.push(HookEntry { hook, priority });
|
||||
hooks.sort_by_key(|e| e.priority);
|
||||
}
|
||||
|
||||
/// Unregister a hook by name. Returns `true` if it was found and removed.
|
||||
pub async fn unregister(&self, name: &str) -> bool {
|
||||
let mut hooks = self.hooks.write().await;
|
||||
let before = hooks.len();
|
||||
hooks.retain(|e| e.hook.name() != name);
|
||||
hooks.len() < before
|
||||
}
|
||||
|
||||
/// List all registered hook names (in priority order).
|
||||
pub async fn list(&self) -> Vec<String> {
|
||||
let hooks = self.hooks.read().await;
|
||||
hooks.iter().map(|e| e.hook.name().to_string()).collect()
|
||||
}
|
||||
|
||||
/// Run all hooks matching the event's hook point.
|
||||
///
|
||||
/// - Hooks run in priority order (lowest first).
|
||||
/// - `Reject` stops the chain immediately.
|
||||
/// - `Modify` chains the modification through subsequent hooks.
|
||||
/// - Timeout/error handling respects each hook's `failure_mode`.
|
||||
pub async fn run(&self, event: &HookEvent) -> Result<HookOutcome, HookError> {
|
||||
let point = event.hook_point();
|
||||
let ctx = HookContext::default();
|
||||
|
||||
// Clone matching hooks and drop the read guard before executing.
|
||||
// Each hook can run up to its timeout, so holding the guard would
|
||||
// block concurrent register/unregister/run calls.
|
||||
let matching: Vec<Arc<dyn Hook>> = {
|
||||
let hooks = self.hooks.read().await;
|
||||
hooks
|
||||
.iter()
|
||||
.filter(|e| e.hook.hook_points().contains(&point))
|
||||
.map(|e| e.hook.clone())
|
||||
.collect()
|
||||
};
|
||||
|
||||
if matching.is_empty() {
|
||||
return Ok(HookOutcome::ok());
|
||||
}
|
||||
|
||||
let mut current_event = event.clone();
|
||||
|
||||
for hook in &matching {
|
||||
let timeout = hook.timeout();
|
||||
|
||||
let result = tokio::time::timeout(timeout, hook.execute(¤t_event, &ctx)).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(HookOutcome::Reject { reason })) => {
|
||||
tracing::debug!(hook = hook.name(), "Hook rejected: {}", reason);
|
||||
return Err(HookError::Rejected { reason });
|
||||
}
|
||||
Ok(Ok(HookOutcome::Continue {
|
||||
modified: Some(value),
|
||||
})) => {
|
||||
tracing::debug!(hook = hook.name(), "Hook modified content");
|
||||
current_event.apply_modification(&value);
|
||||
}
|
||||
Ok(Ok(HookOutcome::Continue { modified: None })) => {
|
||||
// No-op, continue chain
|
||||
}
|
||||
Ok(Err(err)) => match hook.failure_mode() {
|
||||
HookFailureMode::FailOpen => {
|
||||
tracing::warn!(hook = hook.name(), "Hook failed (fail-open): {}", err);
|
||||
}
|
||||
HookFailureMode::FailClosed => {
|
||||
tracing::warn!(hook = hook.name(), "Hook failed (fail-closed): {}", err);
|
||||
return Err(HookError::ExecutionFailed {
|
||||
reason: format!("Hook '{}' failed: {}", hook.name(), err),
|
||||
});
|
||||
}
|
||||
},
|
||||
Err(_elapsed) => match hook.failure_mode() {
|
||||
HookFailureMode::FailOpen => {
|
||||
tracing::warn!(
|
||||
hook = hook.name(),
|
||||
"Hook timed out (fail-open) after {:?}",
|
||||
timeout
|
||||
);
|
||||
}
|
||||
HookFailureMode::FailClosed => {
|
||||
tracing::warn!(
|
||||
hook = hook.name(),
|
||||
"Hook timed out (fail-closed) after {:?}",
|
||||
timeout
|
||||
);
|
||||
return Err(HookError::Timeout { timeout });
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Determine final outcome by comparing with original event
|
||||
let modified = extract_content(¤t_event);
|
||||
let original = extract_content(event);
|
||||
|
||||
if modified != original {
|
||||
Ok(HookOutcome::modify(modified))
|
||||
} else {
|
||||
Ok(HookOutcome::ok())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HookRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the primary content string from a hook event.
|
||||
fn extract_content(event: &HookEvent) -> String {
|
||||
match event {
|
||||
HookEvent::Inbound { content, .. } | HookEvent::Outbound { content, .. } => content.clone(),
|
||||
HookEvent::ToolCall { parameters, .. } => {
|
||||
serde_json::to_string(parameters).unwrap_or_default()
|
||||
}
|
||||
HookEvent::ResponseTransform { response, .. } => response.clone(),
|
||||
HookEvent::SessionStart { session_id, .. } | HookEvent::SessionEnd { session_id, .. } => {
|
||||
session_id.clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::hooks::hook::{HookFailureMode, HookPoint};
|
||||
use async_trait::async_trait;
|
||||
use std::time::Duration;
|
||||
|
||||
/// A test hook that always returns ok.
|
||||
struct PassthroughHook {
|
||||
name: String,
|
||||
points: Vec<HookPoint>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook for PassthroughHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn hook_points(&self) -> &[HookPoint] {
|
||||
&self.points
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_event: &HookEvent,
|
||||
_ctx: &HookContext,
|
||||
) -> Result<HookOutcome, HookError> {
|
||||
Ok(HookOutcome::ok())
|
||||
}
|
||||
}
|
||||
|
||||
/// A hook that modifies content by appending a suffix.
|
||||
struct ModifyHook {
|
||||
name: String,
|
||||
suffix: String,
|
||||
points: Vec<HookPoint>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook for ModifyHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn hook_points(&self) -> &[HookPoint] {
|
||||
&self.points
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
event: &HookEvent,
|
||||
_ctx: &HookContext,
|
||||
) -> Result<HookOutcome, HookError> {
|
||||
let content = extract_content(event);
|
||||
Ok(HookOutcome::modify(format!("{}{}", content, self.suffix)))
|
||||
}
|
||||
}
|
||||
|
||||
/// A hook that always rejects.
|
||||
struct RejectHook {
|
||||
name: String,
|
||||
reason: String,
|
||||
points: Vec<HookPoint>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook for RejectHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn hook_points(&self) -> &[HookPoint] {
|
||||
&self.points
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_event: &HookEvent,
|
||||
_ctx: &HookContext,
|
||||
) -> Result<HookOutcome, HookError> {
|
||||
Ok(HookOutcome::reject(&self.reason))
|
||||
}
|
||||
}
|
||||
|
||||
/// A hook that always errors.
|
||||
struct ErrorHook {
|
||||
name: String,
|
||||
points: Vec<HookPoint>,
|
||||
failure_mode: HookFailureMode,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook for ErrorHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn hook_points(&self) -> &[HookPoint] {
|
||||
&self.points
|
||||
}
|
||||
fn failure_mode(&self) -> HookFailureMode {
|
||||
self.failure_mode
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_event: &HookEvent,
|
||||
_ctx: &HookContext,
|
||||
) -> Result<HookOutcome, HookError> {
|
||||
Err(HookError::ExecutionFailed {
|
||||
reason: "test error".into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A hook that sleeps longer than its timeout.
|
||||
struct SlowHook {
|
||||
name: String,
|
||||
points: Vec<HookPoint>,
|
||||
failure_mode: HookFailureMode,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook for SlowHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn hook_points(&self) -> &[HookPoint] {
|
||||
&self.points
|
||||
}
|
||||
fn failure_mode(&self) -> HookFailureMode {
|
||||
self.failure_mode
|
||||
}
|
||||
fn timeout(&self) -> Duration {
|
||||
Duration::from_millis(50)
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_event: &HookEvent,
|
||||
_ctx: &HookContext,
|
||||
) -> Result<HookOutcome, HookError> {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
Ok(HookOutcome::ok())
|
||||
}
|
||||
}
|
||||
|
||||
fn test_event() -> HookEvent {
|
||||
HookEvent::Inbound {
|
||||
user_id: "user-1".into(),
|
||||
channel: "test".into(),
|
||||
content: "hello".into(),
|
||||
thread_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_registry_returns_ok() {
|
||||
let registry = HookRegistry::new();
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_ok());
|
||||
assert!(matches!(
|
||||
result.unwrap(),
|
||||
HookOutcome::Continue { modified: None }
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_and_list() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(PassthroughHook {
|
||||
name: "hook-a".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}))
|
||||
.await;
|
||||
registry
|
||||
.register(Arc::new(PassthroughHook {
|
||||
name: "hook-b".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}))
|
||||
.await;
|
||||
|
||||
let names = registry.list().await;
|
||||
assert_eq!(names, vec!["hook-a", "hook-b"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_priority_ordering() {
|
||||
let registry = HookRegistry::new();
|
||||
|
||||
// Register in reverse priority order
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "low-prio".into(),
|
||||
suffix: "-LOW".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
200,
|
||||
)
|
||||
.await;
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "high-prio".into(),
|
||||
suffix: "-HIGH".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Should run in priority order: high-prio first, then low-prio
|
||||
let names = registry.list().await;
|
||||
assert_eq!(names[0], "high-prio");
|
||||
assert_eq!(names[1], "low-prio");
|
||||
|
||||
let result = registry.run(&test_event()).await.unwrap();
|
||||
match result {
|
||||
HookOutcome::Continue { modified: Some(m) } => {
|
||||
// "hello" -> "hello-HIGH" -> "hello-HIGH-LOW"
|
||||
assert_eq!(m, "hello-HIGH-LOW");
|
||||
}
|
||||
other => panic!("Expected modification chain, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reject_stops_chain() {
|
||||
let registry = HookRegistry::new();
|
||||
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(RejectHook {
|
||||
name: "blocker".into(),
|
||||
reason: "blocked".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "modifier".into(),
|
||||
suffix: "-MODIFIED".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
20,
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
HookError::Rejected { reason } => assert_eq!(reason, "blocked"),
|
||||
other => panic!("Expected Rejected, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_modification_chaining() {
|
||||
let registry = HookRegistry::new();
|
||||
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "first".into(),
|
||||
suffix: "-A".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "second".into(),
|
||||
suffix: "-B".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
20,
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await.unwrap();
|
||||
match result {
|
||||
HookOutcome::Continue { modified: Some(m) } => {
|
||||
assert_eq!(m, "hello-A-B");
|
||||
}
|
||||
other => panic!("Expected chained modification, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fail_open_on_error() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(ErrorHook {
|
||||
name: "err-open".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
failure_mode: HookFailureMode::FailOpen,
|
||||
}))
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fail_closed_on_error() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(ErrorHook {
|
||||
name: "err-closed".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
failure_mode: HookFailureMode::FailClosed,
|
||||
}))
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_err());
|
||||
assert!(matches!(
|
||||
result.unwrap_err(),
|
||||
HookError::ExecutionFailed { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fail_open_on_timeout() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(SlowHook {
|
||||
name: "slow-open".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
failure_mode: HookFailureMode::FailOpen,
|
||||
}))
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fail_closed_on_timeout() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(SlowHook {
|
||||
name: "slow-closed".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
failure_mode: HookFailureMode::FailClosed,
|
||||
}))
|
||||
.await;
|
||||
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_err());
|
||||
assert!(matches!(result.unwrap_err(), HookError::Timeout { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unregister() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(PassthroughHook {
|
||||
name: "removable".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}))
|
||||
.await;
|
||||
|
||||
assert_eq!(registry.list().await.len(), 1);
|
||||
assert!(registry.unregister("removable").await);
|
||||
assert_eq!(registry.list().await.len(), 0);
|
||||
|
||||
// Unregistering non-existent returns false
|
||||
assert!(!registry.unregister("nonexistent").await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hooks_only_match_their_points() {
|
||||
let registry = HookRegistry::new();
|
||||
registry
|
||||
.register(Arc::new(RejectHook {
|
||||
name: "outbound-only".into(),
|
||||
reason: "blocked".into(),
|
||||
points: vec![HookPoint::BeforeOutbound],
|
||||
}))
|
||||
.await;
|
||||
|
||||
// Inbound event should not be affected by outbound-only hook
|
||||
let result = registry.run(&test_event()).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
@@ -39,16 +39,19 @@
|
||||
//! - **Continuous learning** - Improve estimates from historical data
|
||||
|
||||
pub mod agent;
|
||||
pub mod boot_screen;
|
||||
pub mod bootstrap;
|
||||
pub mod channels;
|
||||
pub mod cli;
|
||||
pub mod config;
|
||||
pub mod context;
|
||||
pub mod db;
|
||||
pub mod error;
|
||||
pub mod estimation;
|
||||
pub mod evaluation;
|
||||
pub mod extensions;
|
||||
pub mod history;
|
||||
pub mod hooks;
|
||||
pub mod llm;
|
||||
pub mod orchestrator;
|
||||
pub mod pairing;
|
||||
@@ -58,6 +61,8 @@ pub mod secrets;
|
||||
pub mod settings;
|
||||
pub mod setup;
|
||||
pub mod tools;
|
||||
pub mod tracing_fmt;
|
||||
pub mod util;
|
||||
pub mod worker;
|
||||
pub mod workspace;
|
||||
|
||||
|
||||
+1017
File diff suppressed because it is too large
Load Diff
+131
-13
@@ -8,20 +8,26 @@
|
||||
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
|
||||
|
||||
mod costs;
|
||||
pub mod failover;
|
||||
mod nearai;
|
||||
mod nearai_chat;
|
||||
mod provider;
|
||||
mod reasoning;
|
||||
mod retry;
|
||||
mod rig_adapter;
|
||||
pub mod session;
|
||||
|
||||
pub use failover::{CooldownConfig, FailoverProvider};
|
||||
pub use nearai::{ModelInfo, NearAiProvider};
|
||||
pub use nearai_chat::NearAiChatProvider;
|
||||
pub use provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
|
||||
};
|
||||
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
|
||||
pub use reasoning::{
|
||||
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, TokenUsage,
|
||||
ToolSelection,
|
||||
};
|
||||
pub use rig_adapter::RigAdapter;
|
||||
pub use session::{SessionConfig, SessionManager, create_session_manager};
|
||||
|
||||
@@ -30,7 +36,7 @@ use std::sync::Arc;
|
||||
use rig::client::CompletionClient;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode};
|
||||
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode, NearAiConfig};
|
||||
use crate::error::LlmError;
|
||||
|
||||
/// Create an LLM provider based on configuration.
|
||||
@@ -43,7 +49,7 @@ pub fn create_llm_provider(
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
match config.backend {
|
||||
LlmBackend::NearAi => create_nearai_provider(config, session),
|
||||
LlmBackend::NearAi => create_llm_provider_with_config(&config.nearai, session),
|
||||
LlmBackend::OpenAi => create_openai_provider(config),
|
||||
LlmBackend::Anthropic => create_anthropic_provider(config),
|
||||
LlmBackend::Ollama => create_ollama_provider(config),
|
||||
@@ -51,21 +57,28 @@ pub fn create_llm_provider(
|
||||
}
|
||||
}
|
||||
|
||||
fn create_nearai_provider(
|
||||
config: &LlmConfig,
|
||||
/// Create an LLM provider from a `NearAiConfig` directly.
|
||||
///
|
||||
/// This is useful when constructing additional providers for failover,
|
||||
/// where only the model name differs from the primary config.
|
||||
pub fn create_llm_provider_with_config(
|
||||
config: &NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
match config.nearai.api_mode {
|
||||
match config.api_mode {
|
||||
NearAiApiMode::Responses => {
|
||||
tracing::info!("Using NEAR AI Responses API (chat-api) with session auth");
|
||||
Ok(Arc::new(NearAiProvider::new(
|
||||
config.nearai.clone(),
|
||||
session,
|
||||
)))
|
||||
tracing::info!(
|
||||
model = %config.model,
|
||||
"Using Responses API (chat-api) with session auth"
|
||||
);
|
||||
Ok(Arc::new(NearAiProvider::new(config.clone(), session)))
|
||||
}
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
tracing::info!("Using NEAR AI Chat Completions API (cloud-api) with API key auth");
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.nearai.clone())?))
|
||||
tracing::info!(
|
||||
model = %config.model,
|
||||
"Using Chat Completions API (cloud-api) with API key auth"
|
||||
);
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.clone())?))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -170,3 +183,108 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
|
||||
);
|
||||
Ok(Arc::new(RigAdapter::new(model, &compat.model)))
|
||||
}
|
||||
|
||||
/// Create a cheap/fast LLM provider for lightweight tasks (heartbeat, routing, evaluation).
|
||||
///
|
||||
/// Uses `NEARAI_CHEAP_MODEL` if set, otherwise falls back to the main provider.
|
||||
/// Currently only supports NEAR AI backends (Responses and ChatCompletions modes).
|
||||
pub fn create_cheap_llm_provider(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Option<Arc<dyn LlmProvider>>, LlmError> {
|
||||
let Some(ref cheap_model) = config.nearai.cheap_model else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if config.backend != LlmBackend::NearAi {
|
||||
tracing::warn!(
|
||||
"NEARAI_CHEAP_MODEL is set but LLM_BACKEND is {:?}, not NearAi. \
|
||||
Cheap model setting will be ignored.",
|
||||
config.backend
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut cheap_config = config.nearai.clone();
|
||||
cheap_config.model = cheap_model.clone();
|
||||
|
||||
tracing::info!("Cheap LLM provider: {}", cheap_model);
|
||||
|
||||
match cheap_config.api_mode {
|
||||
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)))),
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::{LlmBackend, NearAiApiMode, NearAiConfig};
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn test_nearai_config() -> NearAiConfig {
|
||||
NearAiConfig {
|
||||
model: "test-model".to_string(),
|
||||
cheap_model: None,
|
||||
base_url: "https://api.near.ai".to_string(),
|
||||
auth_base_url: "https://private.near.ai".to_string(),
|
||||
session_path: PathBuf::from("/tmp/test-session.json"),
|
||||
api_mode: NearAiApiMode::Responses,
|
||||
api_key: None,
|
||||
fallback_model: None,
|
||||
max_retries: 3,
|
||||
failover_cooldown_secs: 300,
|
||||
failover_cooldown_threshold: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_llm_config() -> LlmConfig {
|
||||
LlmConfig {
|
||||
backend: LlmBackend::NearAi,
|
||||
nearai: test_nearai_config(),
|
||||
openai: None,
|
||||
anthropic: None,
|
||||
ollama: None,
|
||||
openai_compatible: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_cheap_llm_provider_returns_none_when_not_configured() {
|
||||
let config = test_llm_config();
|
||||
let session = Arc::new(SessionManager::new(SessionConfig::default()));
|
||||
|
||||
let result = create_cheap_llm_provider(&config, session);
|
||||
assert!(result.is_ok());
|
||||
assert!(result.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_cheap_llm_provider_creates_provider_when_configured() {
|
||||
let mut config = test_llm_config();
|
||||
config.nearai.cheap_model = Some("cheap-test-model".to_string());
|
||||
|
||||
let session = Arc::new(SessionManager::new(SessionConfig::default()));
|
||||
let result = create_cheap_llm_provider(&config, session);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let provider = result.unwrap();
|
||||
assert!(provider.is_some());
|
||||
assert_eq!(provider.unwrap().model_name(), "cheap-test-model");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_cheap_llm_provider_ignored_for_non_nearai_backend() {
|
||||
let mut config = test_llm_config();
|
||||
config.backend = LlmBackend::OpenAi;
|
||||
config.nearai.cheap_model = Some("cheap-test-model".to_string());
|
||||
|
||||
let session = Arc::new(SessionManager::new(SessionConfig::default()));
|
||||
let result = create_cheap_llm_provider(&config, session);
|
||||
|
||||
assert!(result.is_ok());
|
||||
assert!(result.unwrap().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+135
-93
@@ -19,6 +19,7 @@ use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
||||
use crate::llm::session::SessionManager;
|
||||
|
||||
/// Information about an available model from NEAR AI API.
|
||||
@@ -209,20 +210,20 @@ impl NearAiProvider {
|
||||
data: Option<Vec<ModelEntry>>,
|
||||
}
|
||||
|
||||
if let Ok(resp) = serde_json::from_str::<ModelsResponse>(&response_text) {
|
||||
if let Some(entries) = resp.models.or(resp.data) {
|
||||
let models: Vec<ModelInfo> = entries
|
||||
.into_iter()
|
||||
.filter_map(|e| {
|
||||
e.get_name().map(|name| ModelInfo {
|
||||
name,
|
||||
provider: None,
|
||||
})
|
||||
if let Ok(resp) = serde_json::from_str::<ModelsResponse>(&response_text)
|
||||
&& let Some(entries) = resp.models.or(resp.data)
|
||||
{
|
||||
let models: Vec<ModelInfo> = entries
|
||||
.into_iter()
|
||||
.filter_map(|e| {
|
||||
e.get_name().map(|name| ModelInfo {
|
||||
name,
|
||||
provider: None,
|
||||
})
|
||||
.collect();
|
||||
if !models.is_empty() {
|
||||
return Ok(models);
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
if !models.is_empty() {
|
||||
return Ok(models);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -270,88 +271,139 @@ impl NearAiProvider {
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner request implementation without retry logic.
|
||||
/// Inner request implementation with retry logic for transient errors.
|
||||
///
|
||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
||||
async fn send_request_inner<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
path: &str,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
let url = self.api_url(path);
|
||||
let token = self.session.get_token().await?;
|
||||
let max_retries = self.config.max_retries;
|
||||
|
||||
tracing::debug!("Sending request to NEAR AI: {}", url);
|
||||
tracing::debug!("Request body: {:?}", body);
|
||||
for attempt in 0..=max_retries {
|
||||
let token = self.session.get_token().await?;
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("NEAR AI request failed: {}", e);
|
||||
e
|
||||
})?;
|
||||
tracing::debug!(
|
||||
"Sending request to NEAR AI: {} (attempt {})",
|
||||
url,
|
||||
attempt + 1
|
||||
);
|
||||
tracing::debug!("Request body: {:?}", body);
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await;
|
||||
|
||||
tracing::debug!("NEAR AI response status: {}", status);
|
||||
tracing::debug!("NEAR AI response body: {}", response_text);
|
||||
let response = match response {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::error!("NEAR AI request failed: {}", e);
|
||||
// Network errors (timeout, connection refused) are transient
|
||||
if attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI request error (attempt {}/{}), retrying in {:?}: {}",
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
e,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
if !status.is_success() {
|
||||
// Check for session expiration (401 with specific message patterns)
|
||||
if status.as_u16() == 401 {
|
||||
let is_session_expired = response_text.to_lowercase().contains("session")
|
||||
&& (response_text.to_lowercase().contains("expired")
|
||||
|| response_text.to_lowercase().contains("invalid"));
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
|
||||
if is_session_expired {
|
||||
return Err(LlmError::SessionExpired {
|
||||
tracing::debug!("NEAR AI response status: {}", status);
|
||||
tracing::debug!("NEAR AI response body: {}", response_text);
|
||||
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
// Check for session expiration (401 with specific message patterns)
|
||||
if status_code == 401 {
|
||||
let lower = response_text.to_lowercase();
|
||||
let is_session_expired = lower.contains("session")
|
||||
&& (lower.contains("expired") || lower.contains("invalid"));
|
||||
|
||||
if is_session_expired {
|
||||
return Err(LlmError::SessionExpired {
|
||||
provider: "nearai".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Generic 401 -- not retryable
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Generic 401 without session expiration indication
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai".to_string(),
|
||||
});
|
||||
}
|
||||
// Check if this is a transient error worth retrying
|
||||
if is_retryable_status(status_code) && attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
||||
status_code,
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Try to parse as JSON error
|
||||
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
||||
if status.as_u16() == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
// Non-retryable error or exhausted retries
|
||||
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
retry_after: None,
|
||||
reason: error.error,
|
||||
});
|
||||
}
|
||||
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: error.error,
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
}
|
||||
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
// Success -- parse the response
|
||||
return match serde_json::from_str::<R>(&response_text) {
|
||||
Ok(parsed) => Ok(parsed),
|
||||
Err(e) => {
|
||||
tracing::debug!("Response is not expected JSON format: {}", e);
|
||||
tracing::debug!("Will try alternative parsing in caller");
|
||||
Err(LlmError::InvalidResponse {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
||||
})
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Try to parse as our expected type
|
||||
match serde_json::from_str::<R>(&response_text) {
|
||||
Ok(parsed) => Ok(parsed),
|
||||
Err(e) => {
|
||||
tracing::debug!("Response is not expected JSON format: {}", e);
|
||||
tracing::debug!("Will try alternative parsing in caller");
|
||||
Err(LlmError::InvalidResponse {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
||||
})
|
||||
}
|
||||
}
|
||||
// This is unreachable because the loop always returns, but the compiler
|
||||
// cannot prove that. Return a generic error as a safety net.
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -456,7 +508,7 @@ impl LlmProvider for NearAiProvider {
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
tracing::debug!("NEAR AI response: {:?}", response);
|
||||
tracing::debug!("NEAR AI response: output_items={}", response.output.len());
|
||||
|
||||
// Extract text from response output
|
||||
// Try multiple formats since API response shape may vary
|
||||
@@ -464,11 +516,6 @@ impl LlmProvider for NearAiProvider {
|
||||
.output
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
tracing::debug!(
|
||||
"Processing output item: type={}, text={:?}",
|
||||
item.item_type,
|
||||
item.text
|
||||
);
|
||||
if item.item_type == "message" {
|
||||
// First check for direct text field on item
|
||||
if let Some(ref text) = item.text {
|
||||
@@ -479,11 +526,6 @@ impl LlmProvider for NearAiProvider {
|
||||
contents
|
||||
.iter()
|
||||
.filter_map(|c| {
|
||||
tracing::debug!(
|
||||
"Content item: type={}, text={:?}",
|
||||
c.content_type,
|
||||
c.text
|
||||
);
|
||||
// Accept various content types that might contain text
|
||||
match c.content_type.as_str() {
|
||||
"output_text" | "text" => c.text.clone(),
|
||||
@@ -694,21 +736,21 @@ impl LlmProvider for NearAiProvider {
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if item.item_type == "function_call" {
|
||||
if let (Some(name), Some(call_id)) = (&item.name, &item.call_id) {
|
||||
// Parse arguments JSON string into Value
|
||||
let arguments = item
|
||||
.arguments
|
||||
.as_ref()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
} else if item.item_type == "function_call"
|
||||
&& let (Some(name), Some(call_id)) = (&item.name, &item.call_id)
|
||||
{
|
||||
// Parse arguments JSON string into Value
|
||||
let arguments = item
|
||||
.arguments
|
||||
.as_ref()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
|
||||
tool_calls.push(ToolCall {
|
||||
id: call_id.clone(),
|
||||
name: name.clone(),
|
||||
arguments,
|
||||
});
|
||||
}
|
||||
tool_calls.push(ToolCall {
|
||||
id: call_id.clone(),
|
||||
name: name.clone(),
|
||||
arguments,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+270
-37
@@ -16,6 +16,7 @@ use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
||||
|
||||
/// NEAR AI Chat Completions API provider.
|
||||
pub struct NearAiChatProvider {
|
||||
@@ -62,63 +63,116 @@ impl NearAiChatProvider {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Send a request to the chat completions API.
|
||||
/// Send a request to the chat completions API with retry on transient errors.
|
||||
///
|
||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
||||
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
let url = self.api_url("chat/completions");
|
||||
let max_retries = self.config.max_retries;
|
||||
|
||||
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
|
||||
for attempt in 0..=max_retries {
|
||||
tracing::debug!(
|
||||
"Sending request to NEAR AI Chat: {} (attempt {})",
|
||||
url,
|
||||
attempt + 1,
|
||||
);
|
||||
|
||||
// Log the request body for debugging tool call issues
|
||||
if let Ok(json) = serde_json::to_string(body) {
|
||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||
}
|
||||
if tracing::enabled!(tracing::Level::DEBUG)
|
||||
&& let Ok(json) = serde_json::to_string(body)
|
||||
{
|
||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("NEAR AI Chat request failed: {}", e);
|
||||
LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: e.to_string(),
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let response = match response {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::error!("NEAR AI Chat request failed: {}", e);
|
||||
if attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI Chat request error (attempt {}/{}), retrying in {:?}: {}",
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
e,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: e.to_string(),
|
||||
});
|
||||
}
|
||||
})?;
|
||||
};
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
|
||||
tracing::debug!("NEAR AI Chat response status: {}", status);
|
||||
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
||||
tracing::debug!("NEAR AI Chat response status: {}", status);
|
||||
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
||||
|
||||
if !status.is_success() {
|
||||
if status.as_u16() == 401 {
|
||||
return Err(LlmError::AuthFailed {
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
// Auth errors are not retryable
|
||||
if status_code == 401 {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Transient errors: retry with backoff
|
||||
if is_retryable_status(status_code) && attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI Chat returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
||||
status_code,
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Non-retryable or exhausted retries
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai_chat".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
}
|
||||
if status.as_u16() == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai_chat".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
|
||||
// Success — parse the response
|
||||
return serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
|
||||
});
|
||||
}
|
||||
|
||||
serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
|
||||
// Safety net: unreachable because the loop always returns
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -222,6 +276,12 @@ impl LlmProvider for NearAiChatProvider {
|
||||
let messages: Vec<ChatCompletionMessage> =
|
||||
req.messages.into_iter().map(|m| m.into()).collect();
|
||||
|
||||
// NEAR AI cloud-api does not support multi-turn tool calling (rejects
|
||||
// any request containing role:"tool" messages with HTTP 400). Rewrite
|
||||
// tool-call / tool-result pairs into plain text so the conversation
|
||||
// history is preserved without using unsupported message roles.
|
||||
let messages = flatten_tool_messages(messages);
|
||||
|
||||
let tools: Vec<ChatCompletionTool> = req
|
||||
.tools
|
||||
.into_iter()
|
||||
@@ -367,6 +427,64 @@ struct ChatCompletionMessage {
|
||||
tool_calls: Option<Vec<ChatCompletionToolCall>>,
|
||||
}
|
||||
|
||||
/// Rewrite tool-call / tool-result messages into plain assistant/user text.
|
||||
///
|
||||
/// NEAR AI cloud-api does not support the OpenAI multi-turn tool-calling
|
||||
/// protocol (`role: "tool"` messages). This function converts:
|
||||
/// - Assistant messages with `tool_calls` → assistant text describing the calls
|
||||
/// - Tool result messages (`role: "tool"`) → user messages with the result
|
||||
///
|
||||
/// Non-tool messages pass through unchanged.
|
||||
fn flatten_tool_messages(messages: Vec<ChatCompletionMessage>) -> Vec<ChatCompletionMessage> {
|
||||
let has_tool_msgs = messages.iter().any(|m| m.role == "tool");
|
||||
if !has_tool_msgs {
|
||||
return messages;
|
||||
}
|
||||
|
||||
tracing::debug!("Flattening tool messages for NEAR AI compatibility");
|
||||
|
||||
messages
|
||||
.into_iter()
|
||||
.map(|msg| {
|
||||
if let (true, Some(calls)) = (msg.role == "assistant", &msg.tool_calls) {
|
||||
// Convert assistant tool_calls into descriptive text
|
||||
let mut parts: Vec<String> = Vec::new();
|
||||
if let Some(ref text) = msg.content
|
||||
&& !text.is_empty()
|
||||
{
|
||||
parts.push(text.clone());
|
||||
}
|
||||
for tc in calls {
|
||||
parts.push(format!(
|
||||
"[Called tool `{}` with arguments: {}]",
|
||||
tc.function.name, tc.function.arguments
|
||||
));
|
||||
}
|
||||
ChatCompletionMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: Some(parts.join("\n")),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
}
|
||||
} else if msg.role == "tool" {
|
||||
// Convert tool result into a user message
|
||||
let tool_name = msg.name.as_deref().unwrap_or("unknown");
|
||||
let result = msg.content.as_deref().unwrap_or("");
|
||||
ChatCompletionMessage {
|
||||
role: "user".to_string(),
|
||||
content: Some(format!("[Tool `{}` returned: {}]", tool_name, result)),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
}
|
||||
} else {
|
||||
msg
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
impl From<ChatMessage> for ChatCompletionMessage {
|
||||
fn from(msg: ChatMessage) -> Self {
|
||||
let role = match msg.role {
|
||||
@@ -544,4 +662,119 @@ mod tests {
|
||||
serde_json::from_str(&calls[0].function.arguments).expect("valid JSON string");
|
||||
assert_eq!(parsed["key"], "value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flatten_no_tool_messages_passthrough() {
|
||||
let messages = vec![
|
||||
ChatCompletionMessage {
|
||||
role: "system".to_string(),
|
||||
content: Some("You are helpful.".to_string()),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
ChatCompletionMessage {
|
||||
role: "user".to_string(),
|
||||
content: Some("Hello".to_string()),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
];
|
||||
let result = flatten_tool_messages(messages);
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0].role, "system");
|
||||
assert_eq!(result[1].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flatten_tool_call_and_result() {
|
||||
let messages = vec![
|
||||
ChatCompletionMessage {
|
||||
role: "user".to_string(),
|
||||
content: Some("test".to_string()),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
ChatCompletionMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: None,
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: Some(vec![ChatCompletionToolCall {
|
||||
id: "call_1".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: ChatCompletionToolCallFunction {
|
||||
name: "echo".to_string(),
|
||||
arguments: r#"{"message":"hi"}"#.to_string(),
|
||||
},
|
||||
}]),
|
||||
},
|
||||
ChatCompletionMessage {
|
||||
role: "tool".to_string(),
|
||||
content: Some("hi".to_string()),
|
||||
tool_call_id: Some("call_1".to_string()),
|
||||
name: Some("echo".to_string()),
|
||||
tool_calls: None,
|
||||
},
|
||||
];
|
||||
|
||||
let result = flatten_tool_messages(messages);
|
||||
assert_eq!(result.len(), 3);
|
||||
|
||||
// Assistant tool_calls → plain assistant text
|
||||
assert_eq!(result[1].role, "assistant");
|
||||
assert!(result[1].tool_calls.is_none());
|
||||
assert!(
|
||||
result[1]
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("[Called tool `echo`")
|
||||
);
|
||||
|
||||
// Tool result → user message
|
||||
assert_eq!(result[2].role, "user");
|
||||
assert!(result[2].tool_call_id.is_none());
|
||||
assert!(
|
||||
result[2]
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("[Tool `echo` returned: hi]")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flatten_preserves_assistant_text_with_tool_calls() {
|
||||
let messages = vec![
|
||||
ChatCompletionMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: Some("Let me check that.".to_string()),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: Some(vec![ChatCompletionToolCall {
|
||||
id: "call_1".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: ChatCompletionToolCallFunction {
|
||||
name: "search".to_string(),
|
||||
arguments: r#"{"q":"test"}"#.to_string(),
|
||||
},
|
||||
}]),
|
||||
},
|
||||
ChatCompletionMessage {
|
||||
role: "tool".to_string(),
|
||||
content: Some("found it".to_string()),
|
||||
tool_call_id: Some("call_1".to_string()),
|
||||
name: Some("search".to_string()),
|
||||
tool_calls: None,
|
||||
},
|
||||
];
|
||||
|
||||
let result = flatten_tool_messages(messages);
|
||||
let text = result[0].content.as_ref().unwrap();
|
||||
assert!(text.starts_with("Let me check that."));
|
||||
assert!(text.contains("[Called tool `search`"));
|
||||
}
|
||||
}
|
||||
|
||||
+77
-32
@@ -113,6 +113,25 @@ pub struct ToolSelection {
|
||||
pub reasoning: String,
|
||||
/// Alternative tools considered.
|
||||
pub alternatives: Vec<String>,
|
||||
/// The tool call ID from the LLM response.
|
||||
///
|
||||
/// OpenAI-compatible providers assign each tool call a unique ID that must
|
||||
/// be echoed back in the corresponding tool result message. Without this,
|
||||
/// the provider cannot match results to their originating calls.
|
||||
pub tool_call_id: String,
|
||||
}
|
||||
|
||||
/// Token usage from a single LLM call.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct TokenUsage {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
}
|
||||
|
||||
impl TokenUsage {
|
||||
pub fn total(&self) -> u32 {
|
||||
self.input_tokens + self.output_tokens
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a response with potential tool calls.
|
||||
@@ -131,6 +150,13 @@ pub enum RespondResult {
|
||||
},
|
||||
}
|
||||
|
||||
/// A `RespondResult` bundled with the token usage from the LLM call that produced it.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RespondOutput {
|
||||
pub result: RespondResult,
|
||||
pub usage: TokenUsage,
|
||||
}
|
||||
|
||||
/// Reasoning engine for the agent.
|
||||
pub struct Reasoning {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
@@ -224,6 +250,7 @@ impl Reasoning {
|
||||
parameters: tool_call.arguments,
|
||||
reasoning: reasoning.clone(),
|
||||
alternatives: vec![],
|
||||
tool_call_id: tool_call.id,
|
||||
})
|
||||
.collect();
|
||||
|
||||
@@ -284,7 +311,8 @@ Respond in JSON format:
|
||||
/// tool calls as text for simple cases. Use `respond_with_tools()` when you
|
||||
/// need to actually execute tool calls in an agentic loop.
|
||||
pub async fn respond(&self, context: &ReasoningContext) -> Result<String, LlmError> {
|
||||
match self.respond_with_tools(context).await? {
|
||||
let output = self.respond_with_tools(context).await?;
|
||||
match output.result {
|
||||
RespondResult::Text(text) => Ok(text),
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls: calls, ..
|
||||
@@ -299,15 +327,14 @@ Respond in JSON format:
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a response that may include tool calls.
|
||||
/// Generate a response that may include tool calls, with token usage tracking.
|
||||
///
|
||||
/// Returns `RespondResult::ToolCalls` if the model wants to call tools,
|
||||
/// allowing the caller to execute them and continue the conversation.
|
||||
/// Returns `RespondResult::Text` when the model has a final text response.
|
||||
/// Returns `RespondOutput` containing the result and token usage from the LLM call.
|
||||
/// The caller should use `usage` to track cost/budget against the job.
|
||||
pub async fn respond_with_tools(
|
||||
&self,
|
||||
context: &ReasoningContext,
|
||||
) -> Result<RespondResult, LlmError> {
|
||||
) -> Result<RespondOutput, LlmError> {
|
||||
let system_prompt = self.build_conversation_prompt(context);
|
||||
|
||||
let mut messages = vec![ChatMessage::system(system_prompt)];
|
||||
@@ -322,12 +349,19 @@ Respond in JSON format:
|
||||
request.metadata = context.metadata.clone();
|
||||
|
||||
let response = self.llm.complete_with_tools(request).await?;
|
||||
let usage = TokenUsage {
|
||||
input_tokens: response.input_tokens,
|
||||
output_tokens: response.output_tokens,
|
||||
};
|
||||
|
||||
// If there were tool calls, return them for execution
|
||||
if !response.tool_calls.is_empty() {
|
||||
return Ok(RespondResult::ToolCalls {
|
||||
tool_calls: response.tool_calls,
|
||||
content: response.content,
|
||||
return Ok(RespondOutput {
|
||||
result: RespondResult::ToolCalls {
|
||||
tool_calls: response.tool_calls,
|
||||
content: response.content,
|
||||
},
|
||||
usage,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -341,17 +375,23 @@ Respond in JSON format:
|
||||
let recovered = recover_tool_calls_from_content(&content, &context.available_tools);
|
||||
if !recovered.is_empty() {
|
||||
let cleaned = clean_response(&content);
|
||||
return Ok(RespondResult::ToolCalls {
|
||||
tool_calls: recovered,
|
||||
content: if cleaned.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(cleaned)
|
||||
return Ok(RespondOutput {
|
||||
result: RespondResult::ToolCalls {
|
||||
tool_calls: recovered,
|
||||
content: if cleaned.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(cleaned)
|
||||
},
|
||||
},
|
||||
usage,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(RespondResult::Text(clean_response(&content)))
|
||||
Ok(RespondOutput {
|
||||
result: RespondResult::Text(clean_response(&content)),
|
||||
usage,
|
||||
})
|
||||
} else {
|
||||
// No tools, use simple completion
|
||||
let mut request = CompletionRequest::new(messages)
|
||||
@@ -360,7 +400,13 @@ Respond in JSON format:
|
||||
request.metadata = context.metadata.clone();
|
||||
|
||||
let response = self.llm.complete(request).await?;
|
||||
Ok(RespondResult::Text(clean_response(&response.content)))
|
||||
Ok(RespondOutput {
|
||||
result: RespondResult::Text(clean_response(&response.content)),
|
||||
usage: TokenUsage {
|
||||
input_tokens: response.input_tokens,
|
||||
output_tokens: response.output_tokens,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -542,21 +588,20 @@ fn recover_tool_calls_from_content(
|
||||
}
|
||||
|
||||
// Try JSON first: {"name":"x","arguments":{}}
|
||||
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner) {
|
||||
if let Some(name) = parsed.get("name").and_then(|v| v.as_str()) {
|
||||
if tool_names.contains(name) {
|
||||
let arguments = parsed
|
||||
.get("arguments")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
calls.push(ToolCall {
|
||||
id: format!("recovered_{}", calls.len()),
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
}
|
||||
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner)
|
||||
&& let Some(name) = parsed.get("name").and_then(|v| v.as_str())
|
||||
&& tool_names.contains(name)
|
||||
{
|
||||
let arguments = parsed
|
||||
.get("arguments")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
calls.push(ToolCall {
|
||||
id: format!("recovered_{}", calls.len()),
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
// Bare tool name (e.g. "<tool_call>tool_list</tool_call>")
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
//! Shared retry helpers for LLM providers.
|
||||
//!
|
||||
//! Provides exponential backoff with jitter and retryable status classification
|
||||
//! used by both `NearAiProvider` and `NearAiChatProvider`.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
/// Returns `true` if the HTTP status code is transient and worth retrying.
|
||||
pub(crate) fn is_retryable_status(status: u16) -> bool {
|
||||
matches!(status, 429 | 500 | 502 | 503 | 504)
|
||||
}
|
||||
|
||||
/// Calculate exponential backoff delay with random jitter.
|
||||
///
|
||||
/// Base delay is 1 second, doubled each attempt, with +/-25% jitter.
|
||||
/// - attempt 0: ~1s (0.75s - 1.25s)
|
||||
/// - attempt 1: ~2s (1.5s - 2.5s)
|
||||
/// - attempt 2: ~4s (3.0s - 5.0s)
|
||||
pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration {
|
||||
let base_ms: u64 = 1000u64.saturating_mul(2u64.saturating_pow(attempt));
|
||||
let jitter_range = base_ms / 4; // 25%
|
||||
let jitter = if jitter_range > 0 {
|
||||
let offset = rand::thread_rng().gen_range(0..=jitter_range * 2);
|
||||
offset as i64 - jitter_range as i64
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let delay_ms = (base_ms as i64 + jitter).max(100) as u64;
|
||||
Duration::from_millis(delay_ms)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_is_retryable_status() {
|
||||
// Transient errors should be retryable
|
||||
assert!(is_retryable_status(429));
|
||||
assert!(is_retryable_status(500));
|
||||
assert!(is_retryable_status(502));
|
||||
assert!(is_retryable_status(503));
|
||||
assert!(is_retryable_status(504));
|
||||
|
||||
// Client errors should not be retryable
|
||||
assert!(!is_retryable_status(400));
|
||||
assert!(!is_retryable_status(401));
|
||||
assert!(!is_retryable_status(403));
|
||||
assert!(!is_retryable_status(404));
|
||||
assert!(!is_retryable_status(422));
|
||||
|
||||
// Success codes should not be retryable
|
||||
assert!(!is_retryable_status(200));
|
||||
assert!(!is_retryable_status(201));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_retry_backoff_delay_exponential_growth() {
|
||||
// Run multiple samples to verify the range, accounting for jitter
|
||||
for _ in 0..20 {
|
||||
let d0 = retry_backoff_delay(0);
|
||||
let d1 = retry_backoff_delay(1);
|
||||
let d2 = retry_backoff_delay(2);
|
||||
|
||||
// Attempt 0: base 1000ms, jitter +/-250ms -> [750, 1250]
|
||||
assert!(d0.as_millis() >= 750, "attempt 0 too low: {:?}", d0);
|
||||
assert!(d0.as_millis() <= 1250, "attempt 0 too high: {:?}", d0);
|
||||
|
||||
// Attempt 1: base 2000ms, jitter +/-500ms -> [1500, 2500]
|
||||
assert!(d1.as_millis() >= 1500, "attempt 1 too low: {:?}", d1);
|
||||
assert!(d1.as_millis() <= 2500, "attempt 1 too high: {:?}", d1);
|
||||
|
||||
// Attempt 2: base 4000ms, jitter +/-1000ms -> [3000, 5000]
|
||||
assert!(d2.as_millis() >= 3000, "attempt 2 too low: {:?}", d2);
|
||||
assert!(d2.as_millis() <= 5000, "attempt 2 too high: {:?}", d2);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_retry_backoff_delay_minimum() {
|
||||
// Even at attempt 0, delay should be at least 100ms (the minimum floor)
|
||||
for _ in 0..20 {
|
||||
let delay = retry_backoff_delay(0);
|
||||
assert!(delay.as_millis() >= 100);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_retry_backoff_delay_no_overflow() {
|
||||
// Very high attempt numbers should not panic from overflow
|
||||
let delay = retry_backoff_delay(30);
|
||||
assert!(delay.as_millis() >= 100);
|
||||
}
|
||||
}
|
||||
+75
-188
@@ -31,8 +31,6 @@ pub struct SessionConfig {
|
||||
pub auth_base_url: String,
|
||||
/// Path to session file (e.g., ~/.ironclaw/session.json).
|
||||
pub session_path: PathBuf,
|
||||
/// Port range for OAuth callback server.
|
||||
pub callback_port_range: (u16, u16),
|
||||
}
|
||||
|
||||
impl Default for SessionConfig {
|
||||
@@ -40,7 +38,6 @@ impl Default for SessionConfig {
|
||||
Self {
|
||||
auth_base_url: "https://private.near.ai".to_string(),
|
||||
session_path: default_session_path(),
|
||||
callback_port_range: (9876, 9886),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -62,7 +59,7 @@ pub struct SessionManager {
|
||||
/// Prevents thundering herd during concurrent 401s.
|
||||
renewal_lock: Mutex<()>,
|
||||
/// Optional database store for persisting session to the settings table.
|
||||
store: RwLock<Option<Arc<crate::history::Store>>>,
|
||||
store: RwLock<Option<Arc<dyn crate::db::Database>>>,
|
||||
/// User ID for DB settings (default: "default").
|
||||
user_id: RwLock<String>,
|
||||
}
|
||||
@@ -83,16 +80,16 @@ impl SessionManager {
|
||||
};
|
||||
|
||||
// Try to load existing session synchronously during construction
|
||||
if let Ok(data) = std::fs::read_to_string(&manager.config.session_path) {
|
||||
if let Ok(session) = serde_json::from_str::<SessionData>(&data) {
|
||||
// We can't await here, so we use try_write
|
||||
if let Ok(mut guard) = manager.token.try_write() {
|
||||
*guard = Some(SecretString::from(session.session_token));
|
||||
tracing::info!(
|
||||
"Loaded session token from {}",
|
||||
manager.config.session_path.display()
|
||||
);
|
||||
}
|
||||
if let Ok(data) = std::fs::read_to_string(&manager.config.session_path)
|
||||
&& let Ok(session) = serde_json::from_str::<SessionData>(&data)
|
||||
{
|
||||
// We can't await here, so we use try_write
|
||||
if let Ok(mut guard) = manager.token.try_write() {
|
||||
*guard = Some(SecretString::from(session.session_token));
|
||||
tracing::info!(
|
||||
"Loaded session token from {}",
|
||||
manager.config.session_path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,7 +122,7 @@ impl SessionManager {
|
||||
/// When a store is attached, session tokens are saved to the `settings`
|
||||
/// table (key: `nearai.session_token`) in addition to the disk file.
|
||||
/// On load, DB is preferred over disk.
|
||||
pub async fn attach_store(&self, store: Arc<crate::history::Store>, user_id: &str) {
|
||||
pub async fn attach_store(&self, store: Arc<dyn crate::db::Database>, user_id: &str) {
|
||||
*self.store.write().await = Some(store);
|
||||
*self.user_id.write().await = user_id.to_string();
|
||||
|
||||
@@ -222,38 +219,21 @@ impl SessionManager {
|
||||
|
||||
/// Start the OAuth login flow.
|
||||
///
|
||||
/// 1. Find an available port for the callback server
|
||||
/// 1. Bind the fixed callback port
|
||||
/// 2. Print the auth URL and attempt to open browser
|
||||
/// 3. Wait for OAuth callback with session token
|
||||
/// 4. Save and return the token
|
||||
async fn initiate_login(&self) -> Result<(), LlmError> {
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
|
||||
|
||||
// Find an available port
|
||||
let mut listener = None;
|
||||
let mut port = 0;
|
||||
let listener = oauth_defaults::bind_callback_listener()
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
for p in self.config.callback_port_range.0..=self.config.callback_port_range.1 {
|
||||
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
|
||||
Ok(l) => {
|
||||
listener = Some(l);
|
||||
port = p;
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let listener = listener.ok_or_else(|| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!(
|
||||
"Could not find available port in range {}-{}",
|
||||
self.config.callback_port_range.0, self.config.callback_port_range.1
|
||||
),
|
||||
})?;
|
||||
|
||||
let callback_url = format!("http://127.0.0.1:{}", port);
|
||||
let callback_url = format!("http://127.0.0.1:{}", OAUTH_CALLBACK_PORT);
|
||||
|
||||
// Show auth provider menu
|
||||
println!();
|
||||
@@ -333,138 +313,16 @@ impl SessionManager {
|
||||
println!();
|
||||
println!("Waiting for authentication...");
|
||||
|
||||
// Wait for callback with timeout
|
||||
// The API redirects to: {frontend_callback}/auth/callback?token=X&session_id=X&expires_at=X&is_new_user=X
|
||||
let timeout = std::time::Duration::from_secs(300); // 5 minutes
|
||||
let selected_provider = auth_provider.to_string();
|
||||
let (session_token, auth_provider) = tokio::time::timeout(timeout, async move {
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await.map_err(|e| {
|
||||
LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to accept connection: {}", e),
|
||||
}
|
||||
// The NEAR AI API redirects to: {frontend_callback}/auth/callback?token=X&...
|
||||
let session_token =
|
||||
oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI")
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let mut reader = BufReader::new(&mut socket);
|
||||
let mut request_line = String::new();
|
||||
reader.read_line(&mut request_line).await.map_err(|e| {
|
||||
LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to read request: {}", e),
|
||||
}
|
||||
})?;
|
||||
|
||||
// Parse GET /auth/callback?token=xxx&session_id=xxx&expires_at=xxx&is_new_user=xxx HTTP/1.1
|
||||
if let Some(path) = request_line.split_whitespace().nth(1) {
|
||||
if path.starts_with("/auth/callback") {
|
||||
// Parse query parameters
|
||||
if let Some(query) = path.split('?').nth(1) {
|
||||
let mut token = None;
|
||||
|
||||
for param in query.split('&') {
|
||||
let parts: Vec<&str> = param.splitn(2, '=').collect();
|
||||
if parts.len() == 2 && parts[0] == "token" {
|
||||
token = Some(
|
||||
urlencoding::decode(parts[1])
|
||||
.unwrap_or_else(|_| parts[1].into())
|
||||
.into_owned(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(token) = token {
|
||||
// Send success response with nice styling
|
||||
let response = concat!(
|
||||
"HTTP/1.1 200 OK\r\n",
|
||||
"Content-Type: text/html; charset=utf-8\r\n",
|
||||
"Connection: close\r\n",
|
||||
"\r\n",
|
||||
"<!DOCTYPE html>\n",
|
||||
"<html>\n",
|
||||
"<head>\n",
|
||||
" <meta charset=\"utf-8\">\n",
|
||||
" <title>NEAR AI - Authentication Successful</title>\n",
|
||||
" <style>\n",
|
||||
" * { margin: 0; padding: 0; box-sizing: border-box; }\n",
|
||||
" body {\n",
|
||||
" font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;\n",
|
||||
" background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);\n",
|
||||
" min-height: 100vh;\n",
|
||||
" display: flex;\n",
|
||||
" align-items: center;\n",
|
||||
" justify-content: center;\n",
|
||||
" color: #fff;\n",
|
||||
" }\n",
|
||||
" .container {\n",
|
||||
" text-align: center;\n",
|
||||
" padding: 3rem;\n",
|
||||
" background: rgba(255,255,255,0.05);\n",
|
||||
" border-radius: 16px;\n",
|
||||
" backdrop-filter: blur(10px);\n",
|
||||
" border: 1px solid rgba(255,255,255,0.1);\n",
|
||||
" max-width: 400px;\n",
|
||||
" }\n",
|
||||
" .checkmark {\n",
|
||||
" width: 80px;\n",
|
||||
" height: 80px;\n",
|
||||
" background: linear-gradient(135deg, #00d9a5 0%, #00b386 100%);\n",
|
||||
" border-radius: 50%;\n",
|
||||
" display: flex;\n",
|
||||
" align-items: center;\n",
|
||||
" justify-content: center;\n",
|
||||
" margin: 0 auto 1.5rem;\n",
|
||||
" font-size: 40px;\n",
|
||||
" }\n",
|
||||
" h1 {\n",
|
||||
" font-size: 1.5rem;\n",
|
||||
" font-weight: 600;\n",
|
||||
" margin-bottom: 0.75rem;\n",
|
||||
" }\n",
|
||||
" p {\n",
|
||||
" color: rgba(255,255,255,0.7);\n",
|
||||
" font-size: 0.95rem;\n",
|
||||
" line-height: 1.5;\n",
|
||||
" }\n",
|
||||
" .brand {\n",
|
||||
" margin-top: 2rem;\n",
|
||||
" padding-top: 1.5rem;\n",
|
||||
" border-top: 1px solid rgba(255,255,255,0.1);\n",
|
||||
" font-size: 0.8rem;\n",
|
||||
" color: rgba(255,255,255,0.4);\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
"</head>\n",
|
||||
"<body>\n",
|
||||
" <div class=\"container\">\n",
|
||||
" <div class=\"checkmark\">✓</div>\n",
|
||||
" <h1>Authentication Successful</h1>\n",
|
||||
" <p>You can close this window and return to the terminal.</p>\n",
|
||||
" <div class=\"brand\">NEAR AI Agent</div>\n",
|
||||
" </div>\n",
|
||||
"</body>\n",
|
||||
"</html>"
|
||||
);
|
||||
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
let _ = socket.shutdown().await;
|
||||
|
||||
return Ok::<_, LlmError>((token, Some(selected_provider.clone())));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Not the callback we're looking for, send 404
|
||||
let response = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "Authentication timed out after 5 minutes".to_string(),
|
||||
})??;
|
||||
let auth_provider = Some(auth_provider.to_string());
|
||||
|
||||
// Save the token
|
||||
self.save_session(&session_token, auth_provider.as_deref())
|
||||
@@ -520,6 +378,25 @@ impl SessionManager {
|
||||
))
|
||||
})?;
|
||||
|
||||
// Restrictive permissions: session file contains a secret token
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let perms = std::fs::Permissions::from_mode(0o600);
|
||||
tokio::fs::set_permissions(&self.config.session_path, perms)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
LlmError::Io(std::io::Error::new(
|
||||
e.kind(),
|
||||
format!(
|
||||
"Failed to set permissions on {}: {}",
|
||||
self.config.session_path.display(),
|
||||
e
|
||||
),
|
||||
))
|
||||
})?;
|
||||
}
|
||||
|
||||
tracing::debug!("Session saved to {}", self.config.session_path.display());
|
||||
|
||||
// Also save to DB if a store is attached
|
||||
@@ -551,17 +428,30 @@ impl SessionManager {
|
||||
})?;
|
||||
|
||||
let user_id = self.user_id.read().await.clone();
|
||||
let value = store
|
||||
let value = if let Some(value) = store
|
||||
.get_setting(&user_id, "nearai.session_token")
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("DB query failed: {}", e),
|
||||
})?
|
||||
.ok_or_else(|| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No session in DB".to_string(),
|
||||
})?;
|
||||
})? {
|
||||
value
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"nearai.session_token missing; falling back to legacy nearai.session for backwards compatibility"
|
||||
);
|
||||
store
|
||||
.get_setting(&user_id, "nearai.session")
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("DB query failed: {}", e),
|
||||
})?
|
||||
.ok_or(LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No session in DB".to_string(),
|
||||
})?
|
||||
};
|
||||
|
||||
let session: SessionData =
|
||||
serde_json::from_value(value).map_err(|e| LlmError::SessionRenewalFailed {
|
||||
@@ -623,15 +513,14 @@ pub async fn create_session_manager(config: SessionConfig) -> Arc<SessionManager
|
||||
let manager = SessionManager::new_async(config).await;
|
||||
|
||||
// Check for legacy env var and migrate if present and no file token
|
||||
if !manager.has_token().await {
|
||||
if let Ok(token) = std::env::var("NEARAI_SESSION_TOKEN") {
|
||||
if !token.is_empty() {
|
||||
tracing::info!("Migrating session token from NEARAI_SESSION_TOKEN env var to file");
|
||||
manager.set_token(SecretString::from(token.clone())).await;
|
||||
if let Err(e) = manager.save_session(&token, None).await {
|
||||
tracing::warn!("Failed to save migrated session: {}", e);
|
||||
}
|
||||
}
|
||||
if !manager.has_token().await
|
||||
&& let Ok(token) = std::env::var("NEARAI_SESSION_TOKEN")
|
||||
&& !token.is_empty()
|
||||
{
|
||||
tracing::info!("Migrating session token from NEARAI_SESSION_TOKEN env var to file");
|
||||
manager.set_token(SecretString::from(token.clone())).await;
|
||||
if let Err(e) = manager.save_session(&token, None).await {
|
||||
tracing::warn!("Failed to save migrated session: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -652,7 +541,6 @@ mod tests {
|
||||
let config = SessionConfig {
|
||||
auth_base_url: "https://example.com".to_string(),
|
||||
session_path: session_path.clone(),
|
||||
callback_port_range: (9900, 9910),
|
||||
};
|
||||
|
||||
let manager = SessionManager::new_async(config.clone()).await;
|
||||
@@ -693,7 +581,6 @@ mod tests {
|
||||
let config = SessionConfig {
|
||||
auth_base_url: "https://example.com".to_string(),
|
||||
session_path: dir.path().join("nonexistent.json"),
|
||||
callback_port_range: (9900, 9910),
|
||||
};
|
||||
|
||||
let manager = SessionManager::new_async(config).await;
|
||||
|
||||
+362
-131
@@ -17,22 +17,23 @@ use ironclaw::{
|
||||
web::log_layer::{LogBroadcaster, WebLogLayer},
|
||||
},
|
||||
cli::{
|
||||
Cli, Command, run_mcp_command, run_memory_command, run_pairing_command, run_status_command,
|
||||
run_tool_command,
|
||||
Cli, Command, run_mcp_command, run_pairing_command, run_status_command, run_tool_command,
|
||||
},
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
extensions::ExtensionManager,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
hooks::HookRegistry,
|
||||
llm::{
|
||||
CooldownConfig, FailoverProvider, LlmProvider, SessionConfig, create_cheap_llm_provider,
|
||||
create_llm_provider, create_llm_provider_with_config, create_session_manager,
|
||||
},
|
||||
orchestrator::{
|
||||
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
||||
api::OrchestratorState,
|
||||
},
|
||||
pairing::PairingStore,
|
||||
safety::SafetyLayer,
|
||||
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
|
||||
setup::{SetupConfig, SetupWizard},
|
||||
secrets::SecretsStore,
|
||||
tools::{
|
||||
ToolRegistry,
|
||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
||||
@@ -41,6 +42,13 @@ use ironclaw::{
|
||||
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
|
||||
};
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
use ironclaw::secrets::LibSqlSecretsStore;
|
||||
#[cfg(feature = "postgres")]
|
||||
use ironclaw::secrets::PostgresSecretsStore;
|
||||
use ironclaw::secrets::SecretsCrypto;
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
use ironclaw::setup::{SetupConfig, SetupWizard};
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let cli = Cli::parse();
|
||||
@@ -85,18 +93,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
.init();
|
||||
|
||||
// Memory commands need database (and optionally embeddings)
|
||||
let _ = dotenvy::dotenv();
|
||||
let config = Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = ironclaw::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
// Set up embeddings if available
|
||||
let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
|
||||
auth_base_url: config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
|
||||
@@ -130,7 +134,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await;
|
||||
// Create a Database-trait-backed workspace for the memory command
|
||||
let db: Arc<dyn ironclaw::db::Database> =
|
||||
ironclaw::db::connect_from_config(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
return ironclaw::cli::run_memory_command_with_db(mem_cmd.clone(), db, embeddings)
|
||||
.await;
|
||||
}
|
||||
Some(Command::Pairing(pairing_cmd)) => {
|
||||
tracing_subscriber::fmt()
|
||||
@@ -142,7 +153,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e));
|
||||
}
|
||||
Some(Command::Status) => {
|
||||
let _ = dotenvy::dotenv();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
|
||||
@@ -216,6 +226,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
max_turns: *max_turns,
|
||||
model: model.clone(),
|
||||
timeout: std::time::Duration::from_secs(1800),
|
||||
allowed_tools: Vec::new(),
|
||||
};
|
||||
|
||||
let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config)
|
||||
@@ -232,15 +243,25 @@ async fn main() -> anyhow::Result<()> {
|
||||
skip_auth,
|
||||
channels_only,
|
||||
}) => {
|
||||
// Load .env before running onboarding wizard
|
||||
// Load .env files before running onboarding wizard.
|
||||
// Standard ./.env first (higher priority), then ~/.ironclaw/.env.
|
||||
let _ = dotenvy::dotenv();
|
||||
ironclaw::bootstrap::load_ironclaw_env();
|
||||
|
||||
let config = SetupConfig {
|
||||
skip_auth: *skip_auth,
|
||||
channels_only: *channels_only,
|
||||
};
|
||||
let mut wizard = SetupWizard::with_config(config);
|
||||
wizard.run().await?;
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
{
|
||||
let config = SetupConfig {
|
||||
skip_auth: *skip_auth,
|
||||
channels_only: *channels_only,
|
||||
};
|
||||
let mut wizard = SetupWizard::with_config(config);
|
||||
wizard.run().await?;
|
||||
}
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = (skip_auth, channels_only);
|
||||
eprintln!("Onboarding wizard requires the 'postgres' or 'libsql' feature.");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
None | Some(Command::Run) => {
|
||||
@@ -248,22 +269,23 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// Load .env if present
|
||||
// Load .env files early so DATABASE_URL (and any other vars) are
|
||||
// available to all subsequent env-based config resolution.
|
||||
// Standard ./.env first (higher priority), then ~/.ironclaw/.env.
|
||||
let _ = dotenvy::dotenv();
|
||||
ironclaw::bootstrap::load_ironclaw_env();
|
||||
|
||||
// Enhanced first-run detection
|
||||
if !cli.no_onboard {
|
||||
if let Some(reason) = check_onboard_needed().await {
|
||||
println!("Onboarding needed: {}", reason);
|
||||
println!();
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.run().await?;
|
||||
}
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
if !cli.no_onboard
|
||||
&& let Some(reason) = check_onboard_needed()
|
||||
{
|
||||
println!("Onboarding needed: {}", reason);
|
||||
println!();
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.run().await?;
|
||||
}
|
||||
|
||||
// Load bootstrap config (4 fields that must live on disk)
|
||||
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
|
||||
|
||||
// Load initial config from env + disk (before DB is available)
|
||||
let mut config = match Config::from_env().await {
|
||||
Ok(c) => c,
|
||||
@@ -283,18 +305,20 @@ async fn main() -> anyhow::Result<()> {
|
||||
let session_config = SessionConfig {
|
||||
auth_base_url: config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
let session = create_session_manager(session_config).await;
|
||||
|
||||
// Ensure we're authenticated before proceeding (only needed for NEAR AI backend)
|
||||
if config.llm.backend == ironclaw::config::LlmBackend::NearAi {
|
||||
// Session-based auth is only needed for NEAR AI backend without an API key.
|
||||
// ChatCompletions mode with an API key skips session auth entirely.
|
||||
if config.llm.backend == ironclaw::config::LlmBackend::NearAi
|
||||
&& config.llm.nearai.api_key.is_none()
|
||||
{
|
||||
session.ensure_authenticated().await?;
|
||||
}
|
||||
|
||||
// Initialize tracing
|
||||
let env_filter = EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=debug"));
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=warn"));
|
||||
|
||||
// Create log broadcaster before tracing init so the WebLogLayer can capture all events.
|
||||
// This gets wired to the gateway's /api/logs/events SSE endpoint later.
|
||||
@@ -302,7 +326,11 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
.with(tracing_subscriber::fmt::layer().with_target(false))
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_target(false)
|
||||
.with_writer(ironclaw::tracing_fmt::TruncatingStderr::default()),
|
||||
)
|
||||
.with(WebLogLayer::new(Arc::clone(&log_broadcaster)))
|
||||
.init();
|
||||
|
||||
@@ -310,7 +338,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
let repl_channel = if let Some(ref msg) = cli.message {
|
||||
Some(ReplChannel::with_message(msg.clone()))
|
||||
} else if config.channels.cli.enabled {
|
||||
Some(ReplChannel::new())
|
||||
let repl = ReplChannel::new();
|
||||
// Suppress the one-liner banner; boot screen will be shown instead.
|
||||
repl.suppress_banner();
|
||||
Some(repl)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -319,23 +350,86 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Loaded configuration for agent: {}", config.agent.name);
|
||||
tracing::info!("LLM backend: {}", config.llm.backend);
|
||||
|
||||
// Initialize database store (optional for testing)
|
||||
let store = if cli.no_db {
|
||||
// Initialize database backend.
|
||||
//
|
||||
// Creates an `Arc<dyn Database>` that all consumers share.
|
||||
// Backend is selected by the `DATABASE_BACKEND` env var / config.
|
||||
//
|
||||
// NOTE: For simpler call sites (CLI commands, Memory handler) use the shared
|
||||
// helper `ironclaw::db::connect_from_config()`. This block is kept inline
|
||||
// because it also captures backend-specific handles (`pg_pool`, `libsql_db`)
|
||||
// needed by the secrets store.
|
||||
#[cfg(feature = "postgres")]
|
||||
let mut pg_pool: Option<deadpool_postgres::Pool> = None;
|
||||
#[cfg(feature = "libsql")]
|
||||
let mut libsql_db: Option<std::sync::Arc<libsql::Database>> = None;
|
||||
|
||||
let db: Option<Arc<dyn ironclaw::db::Database>> = if cli.no_db {
|
||||
tracing::warn!("Running without database connection");
|
||||
None
|
||||
} else {
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
tracing::info!("Database connected and migrations applied");
|
||||
match config.database.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
ironclaw::config::DatabaseBackend::LibSql => {
|
||||
use ironclaw::db::Database as _;
|
||||
use ironclaw::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = ironclaw::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()).await?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path).await?
|
||||
};
|
||||
backend.run_migrations().await?;
|
||||
tracing::info!("libSQL database connected and migrations applied");
|
||||
|
||||
// Capture the Database handle for SecretsStore (connection-per-op)
|
||||
libsql_db = Some(backend.shared_db());
|
||||
|
||||
Some(Arc::new(backend) as Arc<dyn ironclaw::db::Database>)
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
use ironclaw::db::Database as _;
|
||||
let pg = ironclaw::db::postgres::PgBackend::new(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
pg.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
tracing::info!("PostgreSQL database connected and migrations applied");
|
||||
|
||||
pg_pool = Some(pg.pool());
|
||||
Some(Arc::new(pg) as Arc<dyn ironclaw::db::Database>)
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => {
|
||||
anyhow::bail!(
|
||||
"No database backend available. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Post-init operations using the database
|
||||
if let Some(ref db) = db {
|
||||
// One-time migration: move disk config files into the DB settings table.
|
||||
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(&store, "default").await {
|
||||
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(db.as_ref(), "default").await {
|
||||
tracing::warn!("Disk-to-DB settings migration failed: {}", e);
|
||||
}
|
||||
|
||||
// Reload config from DB now that we have a connection.
|
||||
// Priority: env var > DB setting > default.
|
||||
match Config::from_db(&store, "default", &bootstrap).await {
|
||||
match Config::from_db(db.as_ref(), "default").await {
|
||||
Ok(db_config) => {
|
||||
config = db_config;
|
||||
tracing::info!("Configuration reloaded from database");
|
||||
@@ -348,23 +442,120 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
let store = Arc::new(store);
|
||||
|
||||
// Attach store to session manager so tokens save to DB too
|
||||
session.attach_store(Arc::clone(&store), "default").await;
|
||||
// Attach DB to session manager so tokens save to DB too
|
||||
session.attach_store(Arc::clone(db), "default").await;
|
||||
|
||||
// Mark any jobs left in "running" or "creating" state as "interrupted".
|
||||
if let Err(e) = store.cleanup_stale_sandbox_jobs().await {
|
||||
if let Err(e) = db.cleanup_stale_sandbox_jobs().await {
|
||||
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
Some(store)
|
||||
};
|
||||
// Create secrets store early: needed for injecting LLM API keys from encrypted
|
||||
// storage before creating the LLM provider, and later for MCP auth + WASM channels.
|
||||
//
|
||||
// When both `postgres` and `libsql` features are compiled, the runtime-selected
|
||||
// backend determines which store is created: whichever DB init branch ran will
|
||||
// have set its handle (pg_pool or libsql_db), and the or_else chain picks it up.
|
||||
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
|
||||
if let Some(master_key) = config.secrets.master_key() {
|
||||
match SecretsCrypto::new(master_key.clone()) {
|
||||
Ok(crypto) => {
|
||||
let crypto = Arc::new(crypto);
|
||||
let store: Option<Arc<dyn SecretsStore + Send + Sync>> = None;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
let store = store.or_else(|| {
|
||||
libsql_db.take().map(|db| {
|
||||
Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto)))
|
||||
as Arc<dyn SecretsStore + Send + Sync>
|
||||
})
|
||||
});
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
let store = store.or_else(|| {
|
||||
pg_pool.as_ref().map(|pool| {
|
||||
Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto)))
|
||||
as Arc<dyn SecretsStore + Send + Sync>
|
||||
})
|
||||
});
|
||||
|
||||
store
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize secrets crypto: {}", e);
|
||||
#[cfg(feature = "libsql")]
|
||||
let _ = libsql_db.take();
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#[cfg(feature = "libsql")]
|
||||
let _ = libsql_db.take();
|
||||
None
|
||||
};
|
||||
|
||||
// Inject LLM API keys from the encrypted secrets store into a thread-safe
|
||||
// overlay so that optional_env() (used by LlmConfig::resolve()) picks them
|
||||
// up. Then re-resolve LlmConfig with the newly available keys (backend may
|
||||
// have been set during onboarding but the API key is in the secrets store).
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
ironclaw::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await;
|
||||
|
||||
// Re-resolve LlmConfig now that secrets overlay has been populated
|
||||
if let Some(ref db_ref) = db {
|
||||
match Config::from_db(db_ref.as_ref(), "default").await {
|
||||
Ok(refreshed) => {
|
||||
config = refreshed;
|
||||
tracing::debug!("LlmConfig re-resolved after secret injection");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to re-resolve config after secret injection: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize LLM provider (clone session so we can reuse it for embeddings)
|
||||
let llm = create_llm_provider(&config.llm, session.clone())?;
|
||||
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
||||
|
||||
// Wrap in failover if a fallback model is configured
|
||||
let llm: Arc<dyn LlmProvider> =
|
||||
if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
|
||||
if fallback_model == &config.llm.nearai.model {
|
||||
tracing::warn!(
|
||||
"fallback_model is the same as primary model, failover may not be effective"
|
||||
);
|
||||
}
|
||||
let mut fallback_config = config.llm.nearai.clone();
|
||||
fallback_config.model = fallback_model.clone();
|
||||
let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
|
||||
tracing::info!(
|
||||
primary = %llm.model_name(),
|
||||
fallback = %fallback.model_name(),
|
||||
"LLM failover enabled"
|
||||
);
|
||||
let cooldown_config = CooldownConfig {
|
||||
cooldown_duration: std::time::Duration::from_secs(
|
||||
config.llm.nearai.failover_cooldown_secs,
|
||||
),
|
||||
failure_threshold: config.llm.nearai.failover_cooldown_threshold,
|
||||
};
|
||||
Arc::new(FailoverProvider::with_cooldown(
|
||||
vec![llm, fallback],
|
||||
cooldown_config,
|
||||
)?)
|
||||
} else {
|
||||
llm
|
||||
};
|
||||
|
||||
// Initialize cheap LLM provider for lightweight tasks (heartbeat, evaluation)
|
||||
let cheap_llm = create_cheap_llm_provider(&config.llm, session.clone())?;
|
||||
if let Some(ref cheap) = cheap_llm {
|
||||
tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name());
|
||||
}
|
||||
|
||||
// Initialize safety layer
|
||||
let safety = Arc::new(SafetyLayer::new(&config.safety));
|
||||
tracing::info!("Safety layer initialized");
|
||||
@@ -414,8 +605,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
// Register memory tools if database is available
|
||||
if let Some(ref store) = store {
|
||||
let mut workspace = Workspace::new("default", store.pool());
|
||||
if let Some(ref db) = db {
|
||||
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
|
||||
if let Some(ref emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb.clone());
|
||||
}
|
||||
@@ -438,23 +629,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Builder mode enabled");
|
||||
}
|
||||
|
||||
// Create secrets store if master key is configured (needed for MCP auth and WASM channels)
|
||||
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
|
||||
if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) {
|
||||
match SecretsCrypto::new(master_key.clone()) {
|
||||
Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
))),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize secrets crypto: {}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
|
||||
// Create WASM tool runtime (sync, just builds the wasmtime engine)
|
||||
@@ -475,7 +649,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe.
|
||||
let wasm_tools_future = async {
|
||||
if let Some(ref runtime) = wasm_tool_runtime {
|
||||
let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
|
||||
let mut loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
loader = loader.with_secrets_store(Arc::clone(secrets));
|
||||
}
|
||||
|
||||
// Load installed tools from ~/.ironclaw/tools/
|
||||
match loader.load_from_dir(&config.wasm.tools_dir).await {
|
||||
@@ -515,8 +692,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
let mcp_servers_future = async {
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
let servers_result = if let Some(ref s) = store {
|
||||
load_mcp_servers_from_db(s, "default").await
|
||||
let servers_result = if let Some(ref d) = db {
|
||||
load_mcp_servers_from_db(d.as_ref(), "default").await
|
||||
} else {
|
||||
ironclaw::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
@@ -629,7 +806,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.channels.wasm_channels_dir.clone(),
|
||||
config.tunnel.public_url.clone(),
|
||||
"default".to_string(),
|
||||
store.clone(),
|
||||
db.clone(),
|
||||
));
|
||||
tools.register_extension_tools(Arc::clone(&manager));
|
||||
tracing::info!("Extension manager initialized with in-chat discovery tools");
|
||||
@@ -681,6 +858,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
claude_code_model: config.claude_code.model.clone(),
|
||||
claude_code_max_turns: config.claude_code.max_turns,
|
||||
claude_code_memory_limit_mb: config.claude_code.memory_limit_mb,
|
||||
claude_code_allowed_tools: config.claude_code.allowed_tools.clone(),
|
||||
};
|
||||
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
|
||||
|
||||
@@ -691,7 +869,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
token_store,
|
||||
job_event_tx: job_event_tx.clone(),
|
||||
prompt_queue: Arc::clone(&prompt_queue),
|
||||
store: store.clone(),
|
||||
store: db.clone(),
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -720,12 +898,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
// Initialize channel manager
|
||||
let mut channels = ChannelManager::new();
|
||||
let mut channel_names: Vec<String> = Vec::new();
|
||||
|
||||
if let Some(repl) = repl_channel {
|
||||
channels.add(Box::new(repl));
|
||||
if cli.message.is_some() {
|
||||
tracing::info!("Single message mode");
|
||||
} else {
|
||||
channel_names.push("repl".to_string());
|
||||
tracing::info!("REPL mode enabled");
|
||||
}
|
||||
}
|
||||
@@ -797,13 +977,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
// Inject owner_id for Telegram so the bot only responds
|
||||
// to the bound user account.
|
||||
if channel_name == "telegram" {
|
||||
if let Some(owner_id) = config.channels.telegram_owner_id {
|
||||
config_updates.insert(
|
||||
"owner_id".to_string(),
|
||||
serde_json::json!(owner_id),
|
||||
);
|
||||
}
|
||||
if channel_name == "telegram"
|
||||
&& let Some(owner_id) = config.channels.telegram_owner_id
|
||||
{
|
||||
config_updates.insert(
|
||||
"owner_id".to_string(),
|
||||
serde_json::json!(owner_id),
|
||||
);
|
||||
}
|
||||
|
||||
if !config_updates.is_empty() {
|
||||
@@ -861,6 +1041,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
channel_names.push(channel_name.clone());
|
||||
channels.add(Box::new(SharedWasmChannel::new(channel_arc)));
|
||||
}
|
||||
|
||||
@@ -894,23 +1075,24 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Extract its routes for the unified server; the channel itself just
|
||||
// provides the mpsc stream.
|
||||
let mut webhook_server_addr: Option<std::net::SocketAddr> = None;
|
||||
if !cli.cli_only {
|
||||
if let Some(ref http_config) = config.channels.http {
|
||||
let http_channel = HttpChannel::new(http_config.clone());
|
||||
webhook_routes.push(http_channel.routes());
|
||||
let (host, port) = http_channel.addr();
|
||||
webhook_server_addr = Some(
|
||||
format!("{}:{}", host, port)
|
||||
.parse()
|
||||
.expect("HttpConfig host:port must be a valid SocketAddr"),
|
||||
);
|
||||
channels.add(Box::new(http_channel));
|
||||
tracing::info!(
|
||||
"HTTP channel enabled on {}:{}",
|
||||
http_config.host,
|
||||
http_config.port
|
||||
);
|
||||
}
|
||||
if !cli.cli_only
|
||||
&& let Some(ref http_config) = config.channels.http
|
||||
{
|
||||
let http_channel = HttpChannel::new(http_config.clone());
|
||||
webhook_routes.push(http_channel.routes());
|
||||
let (host, port) = http_channel.addr();
|
||||
webhook_server_addr = Some(
|
||||
format!("{}:{}", host, port)
|
||||
.parse()
|
||||
.expect("HttpConfig host:port must be a valid SocketAddr"),
|
||||
);
|
||||
channel_names.push("http".to_string());
|
||||
channels.add(Box::new(http_channel));
|
||||
tracing::info!(
|
||||
"HTTP channel enabled on {}:{}",
|
||||
http_config.host,
|
||||
http_config.port
|
||||
);
|
||||
}
|
||||
|
||||
// Start the unified webhook server if any routes were registered.
|
||||
@@ -928,13 +1110,15 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
// Create workspace for agent (shared with memory tools)
|
||||
let workspace = store.as_ref().map(|s| {
|
||||
let mut ws = Workspace::new("default", s.pool());
|
||||
let workspace = if let Some(ref db_ref) = db {
|
||||
let mut ws = Workspace::new_with_db("default", Arc::clone(db_ref));
|
||||
if let Some(ref emb) = embeddings {
|
||||
ws = ws.with_embeddings(emb.clone());
|
||||
}
|
||||
Arc::new(ws)
|
||||
});
|
||||
Some(Arc::new(ws))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Seed workspace with core identity files on first boot
|
||||
if let Some(ref ws) = workspace {
|
||||
@@ -965,17 +1149,21 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Create context manager (shared between job tools and agent)
|
||||
let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs));
|
||||
|
||||
// Create hook registry
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
|
||||
// Create session manager (shared between agent and web gateway)
|
||||
let session_manager = Arc::new(SessionManager::new());
|
||||
let session_manager = Arc::new(SessionManager::new().with_hooks(hooks.clone()));
|
||||
|
||||
// Register job tools (sandbox deps auto-injected when container_job_manager is available)
|
||||
tools.register_job_tools(
|
||||
Arc::clone(&context_manager),
|
||||
container_job_manager.clone(),
|
||||
store.clone(),
|
||||
db.clone(),
|
||||
);
|
||||
|
||||
// Add web gateway channel if configured
|
||||
let mut gateway_url: Option<String> = None;
|
||||
if let Some(ref gw_config) = config.channels.gateway {
|
||||
let mut gw = GatewayChannel::new(gw_config.clone());
|
||||
if let Some(ref ws) = workspace {
|
||||
@@ -987,8 +1175,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
if let Some(ref ext_mgr) = extension_manager {
|
||||
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
|
||||
}
|
||||
if let Some(ref s) = store {
|
||||
gw = gw.with_store(Arc::clone(s));
|
||||
if let Some(ref d) = db {
|
||||
gw = gw.with_store(Arc::clone(d));
|
||||
}
|
||||
if let Some(ref jm) = container_job_manager {
|
||||
gw = gw.with_job_manager(Arc::clone(jm));
|
||||
@@ -1008,29 +1196,39 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
gateway_url = Some(format!(
|
||||
"http://{}:{}/?token={}",
|
||||
gw_config.host,
|
||||
gw_config.port,
|
||||
gw.auth_token()
|
||||
));
|
||||
|
||||
tracing::info!(
|
||||
"Web gateway enabled on {}:{}",
|
||||
gw_config.host,
|
||||
gw_config.port
|
||||
);
|
||||
tracing::info!(
|
||||
"Web UI: http://{}:{}/?token={}",
|
||||
gw_config.host,
|
||||
gw_config.port,
|
||||
gw.auth_token()
|
||||
);
|
||||
tracing::info!("Web UI: http://{}:{}/", gw_config.host, gw_config.port);
|
||||
|
||||
channel_names.push("gateway".to_string());
|
||||
channels.add(Box::new(gw));
|
||||
}
|
||||
|
||||
// Capture boot screen info before moving Arcs into AgentDeps.
|
||||
let boot_tool_count = tools.count();
|
||||
let boot_llm_model = llm.model_name().to_string();
|
||||
let boot_cheap_model = cheap_llm.as_ref().map(|c| c.model_name().to_string());
|
||||
|
||||
// Create and run the agent
|
||||
let deps = AgentDeps {
|
||||
store,
|
||||
store: db,
|
||||
llm,
|
||||
cheap_llm,
|
||||
safety,
|
||||
tools,
|
||||
workspace,
|
||||
extension_manager,
|
||||
hooks,
|
||||
};
|
||||
let agent = Agent::new(
|
||||
config.agent.clone(),
|
||||
@@ -1044,6 +1242,38 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
tracing::info!("Agent initialized, starting main loop...");
|
||||
|
||||
// Print boot screen for interactive CLI mode (not single-message mode).
|
||||
if config.channels.cli.enabled && cli.message.is_none() {
|
||||
let boot_info = ironclaw::boot_screen::BootInfo {
|
||||
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
agent_name: config.agent.name.clone(),
|
||||
llm_backend: config.llm.backend.to_string(),
|
||||
llm_model: boot_llm_model,
|
||||
cheap_model: boot_cheap_model,
|
||||
db_backend: if cli.no_db {
|
||||
"none".to_string()
|
||||
} else {
|
||||
config.database.backend.to_string()
|
||||
},
|
||||
db_connected: !cli.no_db,
|
||||
tool_count: boot_tool_count,
|
||||
gateway_url,
|
||||
embeddings_enabled: config.embeddings.enabled,
|
||||
embeddings_provider: if config.embeddings.enabled {
|
||||
Some(config.embeddings.provider.clone())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
heartbeat_enabled: config.heartbeat.enabled,
|
||||
heartbeat_interval_secs: config.heartbeat.interval_secs,
|
||||
sandbox_enabled: config.sandbox.enabled,
|
||||
claude_code_enabled: config.claude_code.enabled,
|
||||
routines_enabled: config.routines.enabled,
|
||||
channels: channel_names,
|
||||
};
|
||||
ironclaw::boot_screen::print_boot_screen(&boot_info);
|
||||
}
|
||||
|
||||
// Run the agent (blocks until shutdown)
|
||||
agent.run().await?;
|
||||
|
||||
@@ -1059,27 +1289,28 @@ async fn main() -> anyhow::Result<()> {
|
||||
/// Check if onboarding is needed and return the reason.
|
||||
///
|
||||
/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise.
|
||||
async fn check_onboard_needed() -> Option<&'static str> {
|
||||
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
|
||||
/// Called after `load_ironclaw_env()`, so DATABASE_URL from `~/.ironclaw/.env`
|
||||
/// is already in the environment.
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
fn check_onboard_needed() -> Option<&'static str> {
|
||||
let has_db = std::env::var("DATABASE_URL").is_ok()
|
||||
|| std::env::var("LIBSQL_PATH").is_ok()
|
||||
|| ironclaw::config::default_libsql_path().exists();
|
||||
|
||||
// Database not configured (and not in env)
|
||||
if bootstrap.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
|
||||
if !has_db {
|
||||
return Some("Database not configured");
|
||||
}
|
||||
|
||||
// Secrets not configured (and not in env)
|
||||
if bootstrap.secrets_master_key_source == ironclaw::settings::KeySource::None
|
||||
&& std::env::var("SECRETS_MASTER_KEY").is_err()
|
||||
&& !ironclaw::secrets::keychain::has_master_key().await
|
||||
{
|
||||
// Only require secrets setup if user hasn't explicitly disabled it
|
||||
// For now, we don't require it for first run
|
||||
}
|
||||
|
||||
// First run (onboarding never completed and no session)
|
||||
let session_path = ironclaw::llm::session::default_session_path();
|
||||
if !bootstrap.onboard_completed && !session_path.exists() {
|
||||
return Some("First run");
|
||||
// First run (onboarding never completed and no session).
|
||||
// Reads NEARAI_API_KEY env var directly because this function runs
|
||||
// before Config is loaded -- Config::from_env() may fail without a
|
||||
// database URL, which is what triggers onboarding in the first place.
|
||||
if std::env::var("NEARAI_API_KEY").is_err() {
|
||||
let settings = ironclaw::settings::Settings::load();
|
||||
let session_path = ironclaw::llm::session::default_session_path();
|
||||
if !settings.onboard_completed && !session_path.exists() {
|
||||
return Some("First run");
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
|
||||
+14
-14
@@ -15,7 +15,7 @@ use tokio::sync::{Mutex, broadcast};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::types::SseEvent;
|
||||
use crate::history::Store;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
@@ -43,7 +43,7 @@ pub struct OrchestratorState {
|
||||
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
|
||||
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
|
||||
/// Database handle for persisting job events.
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
}
|
||||
|
||||
/// The orchestrator's internal API server.
|
||||
@@ -202,7 +202,7 @@ async fn report_complete(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(report): Json<CompletionReport>,
|
||||
) -> Result<StatusCode, StatusCode> {
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
if report.success {
|
||||
tracing::info!(
|
||||
job_id = %job_id,
|
||||
@@ -223,7 +223,7 @@ async fn report_complete(
|
||||
};
|
||||
let _ = state.job_manager.complete_job(job_id, result).await;
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
Ok(Json(serde_json::json!({"status": "ok"})))
|
||||
}
|
||||
|
||||
// -- Sandbox job event handlers --
|
||||
@@ -339,16 +339,16 @@ async fn get_prompt_handler(
|
||||
Path(job_id): Path<Uuid>,
|
||||
) -> Result<(StatusCode, Json<serde_json::Value>), StatusCode> {
|
||||
let mut queue = state.prompt_queue.lock().await;
|
||||
if let Some(prompts) = queue.get_mut(&job_id) {
|
||||
if let Some(prompt) = prompts.pop_front() {
|
||||
return Ok((
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"content": prompt.content,
|
||||
"done": prompt.done,
|
||||
})),
|
||||
));
|
||||
}
|
||||
if let Some(prompts) = queue.get_mut(&job_id)
|
||||
&& let Some(prompt) = prompts.pop_front()
|
||||
{
|
||||
return Ok((
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"content": prompt.content,
|
||||
"done": prompt.done,
|
||||
})),
|
||||
));
|
||||
}
|
||||
|
||||
// Return 204 with an empty body. The Json wrapper requires some value
|
||||
|
||||
@@ -14,6 +14,7 @@ use axum::http::StatusCode;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use rand::Rng;
|
||||
use subtle::ConstantTimeEq;
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -38,13 +39,13 @@ impl TokenStore {
|
||||
token
|
||||
}
|
||||
|
||||
/// Validate a token for a specific job.
|
||||
/// Validate a token for a specific job (constant-time comparison).
|
||||
pub async fn validate(&self, job_id: Uuid, token: &str) -> bool {
|
||||
self.tokens
|
||||
.read()
|
||||
.await
|
||||
.get(&job_id)
|
||||
.map(|stored| stored == token)
|
||||
.map(|stored| stored.as_bytes().ct_eq(token.as_bytes()).into())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
|
||||
@@ -58,6 +58,8 @@ pub struct ContainerJobConfig {
|
||||
pub claude_code_max_turns: u32,
|
||||
/// Memory limit in MB for Claude Code containers (heavier than workers).
|
||||
pub claude_code_memory_limit_mb: u64,
|
||||
/// Allowed tool patterns for Claude Code (passed as CLAUDE_CODE_ALLOWED_TOOLS env var).
|
||||
pub claude_code_allowed_tools: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for ContainerJobConfig {
|
||||
@@ -71,6 +73,7 @@ impl Default for ContainerJobConfig {
|
||||
claude_code_model: "sonnet".to_string(),
|
||||
claude_code_max_turns: 50,
|
||||
claude_code_memory_limit_mb: 4096,
|
||||
claude_code_allowed_tools: crate::config::ClaudeCodeConfig::default().allowed_tools,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -161,6 +164,29 @@ impl ContainerJobManager {
|
||||
};
|
||||
self.containers.write().await.insert(job_id, handle);
|
||||
|
||||
// Run the actual container creation. On any failure, revoke the token
|
||||
// and remove the handle so we don't leak resources.
|
||||
match self
|
||||
.create_job_inner(job_id, &token, project_dir, mode)
|
||||
.await
|
||||
{
|
||||
Ok(()) => Ok(token),
|
||||
Err(e) => {
|
||||
self.token_store.revoke(job_id).await;
|
||||
self.containers.write().await.remove(&job_id);
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner implementation of container creation (separated for cleanup).
|
||||
async fn create_job_inner(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
token: &str,
|
||||
project_dir: Option<PathBuf>,
|
||||
mode: JobMode,
|
||||
) -> Result<(), OrchestratorError> {
|
||||
// Connect to Docker
|
||||
let docker = connect_docker()
|
||||
.await
|
||||
@@ -203,27 +229,34 @@ impl ContainerJobManager {
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("projects");
|
||||
if let Ok(canonical_base) = projects_base.canonicalize() {
|
||||
if !canonical.starts_with(&canonical_base) {
|
||||
return Err(OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: format!(
|
||||
"project directory {} is outside allowed base {}",
|
||||
canonical.display(),
|
||||
canonical_base.display()
|
||||
),
|
||||
});
|
||||
}
|
||||
if let Ok(canonical_base) = projects_base.canonicalize()
|
||||
&& !canonical.starts_with(&canonical_base)
|
||||
{
|
||||
return Err(OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: format!(
|
||||
"project directory {} is outside allowed base {}",
|
||||
canonical.display(),
|
||||
canonical_base.display()
|
||||
),
|
||||
});
|
||||
}
|
||||
binds.push(format!("{}:/workspace:rw", canonical.display()));
|
||||
env_vec.push("IRONCLAW_WORKSPACE=/workspace".to_string());
|
||||
}
|
||||
|
||||
// Claude Code mode: mount host ~/.claude read-only for auth
|
||||
// Claude Code mode: mount host ~/.claude read-only for auth,
|
||||
// and pass the tool allowlist so the bridge can write settings.json.
|
||||
if mode == JobMode::ClaudeCode {
|
||||
if let Some(ref claude_dir) = self.config.claude_config_dir {
|
||||
binds.push(format!("{}:/home/sandbox/.claude:ro", claude_dir.display()));
|
||||
}
|
||||
if !self.config.claude_code_allowed_tools.is_empty() {
|
||||
env_vec.push(format!(
|
||||
"CLAUDE_CODE_ALLOWED_TOOLS={}",
|
||||
self.config.claude_code_allowed_tools.join(",")
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Memory limit: Claude Code gets more memory
|
||||
@@ -243,11 +276,7 @@ impl ContainerJobManager {
|
||||
network_mode: Some("bridge".to_string()),
|
||||
extra_hosts: Some(vec!["host.docker.internal:host-gateway".to_string()]),
|
||||
cap_drop: Some(vec!["ALL".to_string()]),
|
||||
cap_add: Some(vec![
|
||||
"CHOWN".to_string(),
|
||||
"SETUID".to_string(),
|
||||
"SETGID".to_string(),
|
||||
]),
|
||||
cap_add: Some(vec!["CHOWN".to_string()]),
|
||||
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
|
||||
tmpfs: Some(
|
||||
[("/tmp".to_string(), "size=512M".to_string())]
|
||||
@@ -328,7 +357,7 @@ impl ContainerJobManager {
|
||||
"Created and started worker container"
|
||||
);
|
||||
|
||||
Ok(token)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop a running container job.
|
||||
@@ -355,15 +384,18 @@ impl ContainerJobManager {
|
||||
})?;
|
||||
|
||||
// Stop the container (10 second grace period)
|
||||
let _ = docker
|
||||
if let Err(e) = docker
|
||||
.stop_container(
|
||||
&container_id,
|
||||
Some(bollard::container::StopContainerOptions { t: 10 }),
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container (may already be stopped)");
|
||||
}
|
||||
|
||||
// Remove the container
|
||||
let _ = docker
|
||||
if let Err(e) = docker
|
||||
.remove_container(
|
||||
&container_id,
|
||||
Some(bollard::container::RemoveContainerOptions {
|
||||
@@ -371,7 +403,10 @@ impl ContainerJobManager {
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to remove container (may require manual cleanup)");
|
||||
}
|
||||
|
||||
// Update state
|
||||
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
|
||||
@@ -407,16 +442,21 @@ impl ContainerJobManager {
|
||||
let containers = self.containers.read().await;
|
||||
containers.get(&job_id).map(|h| h.container_id.clone())
|
||||
};
|
||||
if let Some(cid) = container_id {
|
||||
if !cid.is_empty() {
|
||||
if let Ok(docker) = connect_docker().await {
|
||||
let _ = docker
|
||||
if let Some(cid) = container_id
|
||||
&& !cid.is_empty()
|
||||
{
|
||||
match connect_docker().await {
|
||||
Ok(docker) => {
|
||||
if let Err(e) = docker
|
||||
.stop_container(
|
||||
&cid,
|
||||
Some(bollard::container::StopContainerOptions { t: 5 }),
|
||||
)
|
||||
.await;
|
||||
let _ = docker
|
||||
.await
|
||||
{
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop completed container");
|
||||
}
|
||||
if let Err(e) = docker
|
||||
.remove_container(
|
||||
&cid,
|
||||
Some(bollard::container::RemoveContainerOptions {
|
||||
@@ -424,7 +464,13 @@ impl ContainerJobManager {
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to remove completed container");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(job_id = %job_id, error = %e, "Failed to connect to Docker for container cleanup");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+11
-3
@@ -320,19 +320,27 @@ impl PairingStore {
|
||||
fn record_failed_approve(&self, channel: &str) -> Result<(), PairingStoreError> {
|
||||
let path = approve_attempts_path(&self.base_dir, channel)?;
|
||||
fs::create_dir_all(path.parent().unwrap())?;
|
||||
|
||||
// Open (or create) and lock before reading so concurrent callers
|
||||
// don't clobber each other's writes.
|
||||
let file = fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.truncate(false)
|
||||
.open(&path)?;
|
||||
file.lock_exclusive()?;
|
||||
let content = fs::read_to_string(&path).unwrap_or_default();
|
||||
let mut data: ApproveAttemptsFile = serde_json::from_str(&content).unwrap_or_default();
|
||||
|
||||
let mut data: ApproveAttemptsFile = fs::read_to_string(&path)
|
||||
.ok()
|
||||
.and_then(|c| serde_json::from_str(&c).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
let now = now_secs();
|
||||
data.failed_at.push(now);
|
||||
let cutoff = now.saturating_sub(PAIRING_APPROVE_RATE_WINDOW_SECS);
|
||||
data.failed_at.retain(|&t| t >= cutoff);
|
||||
|
||||
let json = serde_json::to_string_pretty(&data)?;
|
||||
fs::write(&path, json)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
|
||||
@@ -147,10 +147,10 @@ impl LeakDetector {
|
||||
// Build prefix matcher for patterns that start with a known prefix
|
||||
let mut prefixes = Vec::new();
|
||||
for (idx, pattern) in patterns.iter().enumerate() {
|
||||
if let Some(prefix) = extract_literal_prefix(pattern.regex.as_str()) {
|
||||
if prefix.len() >= 3 {
|
||||
prefixes.push((prefix, idx));
|
||||
}
|
||||
if let Some(prefix) = extract_literal_prefix(pattern.regex.as_str())
|
||||
&& prefix.len() >= 3
|
||||
{
|
||||
prefixes.push((prefix, idx));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -306,12 +306,11 @@ impl LeakDetector {
|
||||
})?;
|
||||
}
|
||||
|
||||
// Scan body if present and valid UTF-8
|
||||
// Scan body if present. Use lossy UTF-8 conversion so a leading
|
||||
// non-UTF8 byte can't be used to skip scanning entirely.
|
||||
if let Some(body_bytes) = body {
|
||||
if let Ok(body_str) = std::str::from_utf8(body_bytes) {
|
||||
self.scan_and_clean(body_str)?;
|
||||
}
|
||||
// Binary bodies are not scanned (could add hex pattern detection later)
|
||||
let body_str = String::from_utf8_lossy(body_bytes);
|
||||
self.scan_and_clean(&body_str)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -705,4 +704,17 @@ mod tests {
|
||||
let result = detector.scan_http_request("https://api.example.com/webhook", &[], Some(body));
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_http_request_blocks_secret_in_binary_body() {
|
||||
let detector = LeakDetector::new();
|
||||
|
||||
// Attacker prepends a non-UTF8 byte to bypass strict from_utf8 check.
|
||||
// The lossy conversion should still detect the secret.
|
||||
let mut body = vec![0xFF]; // invalid UTF-8 leading byte
|
||||
body.extend_from_slice(b"sk-proj-test1234567890abcdefghij");
|
||||
|
||||
let result = detector.scan_http_request("https://api.example.com/exfil", &[], Some(&body));
|
||||
assert!(result.is_err(), "binary body should still be scanned");
|
||||
}
|
||||
}
|
||||
|
||||
+21
-5
@@ -98,15 +98,15 @@ impl SafetyLayer {
|
||||
was_modified: true,
|
||||
};
|
||||
}
|
||||
if violations
|
||||
let force_sanitize = violations
|
||||
.iter()
|
||||
.any(|rule| rule.action == crate::safety::PolicyAction::Sanitize)
|
||||
{
|
||||
.any(|rule| rule.action == crate::safety::PolicyAction::Sanitize);
|
||||
if force_sanitize {
|
||||
was_modified = true;
|
||||
}
|
||||
|
||||
// Run sanitization if enabled
|
||||
if self.config.injection_check_enabled {
|
||||
// Run sanitization once: if injection_check is enabled OR policy requires it
|
||||
if self.config.injection_check_enabled || force_sanitize {
|
||||
let mut sanitized = self.sanitizer.sanitize(&content);
|
||||
sanitized.was_modified = sanitized.was_modified || was_modified;
|
||||
sanitized
|
||||
@@ -190,4 +190,20 @@ mod tests {
|
||||
assert!(wrapped.contains("sanitized=\"true\""));
|
||||
assert!(wrapped.contains("Hello <world>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_action_forces_sanitization_when_injection_check_disabled() {
|
||||
let config = SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
};
|
||||
let safety = SafetyLayer::new(&config);
|
||||
|
||||
// Content with an injection-like pattern that a policy might flag
|
||||
let output = safety.sanitize_tool_output("test", "normal text");
|
||||
// With injection_check disabled and no policy violations, content
|
||||
// should pass through unmodified
|
||||
assert_eq!(output.content, "normal text");
|
||||
assert!(!output.was_modified);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -279,11 +279,7 @@ impl ContainerRunner {
|
||||
network_mode: Some("bridge".to_string()),
|
||||
// Security: drop all capabilities and add back only what's needed
|
||||
cap_drop: Some(vec!["ALL".to_string()]),
|
||||
cap_add: Some(vec![
|
||||
"CHOWN".to_string(),
|
||||
"SETUID".to_string(),
|
||||
"SETGID".to_string(),
|
||||
]),
|
||||
cap_add: Some(vec!["CHOWN".to_string()]),
|
||||
// Prevent privilege escalation
|
||||
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
|
||||
// Read-only root filesystem (workspace is still writable if policy allows)
|
||||
@@ -498,10 +494,10 @@ impl ContainerRunner {
|
||||
/// 3. `~/.docker/run/docker.sock` (Docker Desktop on macOS)
|
||||
pub async fn connect_docker() -> Result<Docker> {
|
||||
// First try bollard defaults (checks DOCKER_HOST, then /var/run/docker.sock)
|
||||
if let Ok(docker) = Docker::connect_with_local_defaults() {
|
||||
if docker.ping().await.is_ok() {
|
||||
return Ok(docker);
|
||||
}
|
||||
if let Ok(docker) = Docker::connect_with_local_defaults()
|
||||
&& docker.ping().await.is_ok()
|
||||
{
|
||||
return Ok(docker);
|
||||
}
|
||||
|
||||
// Try Docker Desktop socket (macOS)
|
||||
@@ -511,10 +507,9 @@ pub async fn connect_docker() -> Result<Docker> {
|
||||
let sock_str = desktop_sock.to_string_lossy();
|
||||
if let Ok(docker) =
|
||||
Docker::connect_with_socket(&sock_str, 120, bollard::API_DEFAULT_VERSION)
|
||||
&& docker.ping().await.is_ok()
|
||||
{
|
||||
if docker.ping().await.is_ok() {
|
||||
return Ok(docker);
|
||||
}
|
||||
return Ok(docker);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -259,11 +259,11 @@ async fn handle_connect(
|
||||
|
||||
let decision = state.decider.decide(&network_req).await;
|
||||
|
||||
if !decision.is_allowed() {
|
||||
if let NetworkDecision::Deny { reason } = decision {
|
||||
tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason);
|
||||
return error_response(StatusCode::FORBIDDEN, reason);
|
||||
}
|
||||
if !decision.is_allowed()
|
||||
&& let NetworkDecision::Deny { reason } = decision
|
||||
{
|
||||
tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason);
|
||||
return error_response(StatusCode::FORBIDDEN, reason);
|
||||
}
|
||||
|
||||
tracing::debug!("Proxy: allowing CONNECT to {}", host);
|
||||
@@ -294,10 +294,10 @@ async fn forward_request(
|
||||
|
||||
// Copy headers (except hop-by-hop headers)
|
||||
for (name, value) in req.headers() {
|
||||
if !is_hop_by_hop_header(name.as_str()) {
|
||||
if let Ok(v) = value.to_str() {
|
||||
builder = builder.header(name.as_str(), v);
|
||||
}
|
||||
if !is_hop_by_hop_header(name.as_str())
|
||||
&& let Ok(v) = value.to_str()
|
||||
{
|
||||
builder = builder.header(name.as_str(), v);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -109,12 +109,11 @@ impl NetworkPolicyDecider for DefaultPolicyDecider {
|
||||
async fn decide(&self, request: &NetworkRequest) -> NetworkDecision {
|
||||
// First check if the domain is allowed
|
||||
let validation = self.allowlist.is_allowed(&request.host);
|
||||
if !validation.is_allowed() {
|
||||
if let crate::sandbox::proxy::allowlist::DomainValidationResult::Denied(reason) =
|
||||
if !validation.is_allowed()
|
||||
&& let crate::sandbox::proxy::allowlist::DomainValidationResult::Denied(reason) =
|
||||
validation
|
||||
{
|
||||
return NetworkDecision::Deny { reason };
|
||||
}
|
||||
{
|
||||
return NetworkDecision::Deny { reason };
|
||||
}
|
||||
|
||||
// Check if we need to inject credentials
|
||||
|
||||
@@ -261,7 +261,7 @@ pub use platform::{delete_master_key, get_master_key, has_master_key, store_mast
|
||||
|
||||
/// Parse a hex string to bytes.
|
||||
fn hex_to_bytes(hex: &str) -> Result<Vec<u8>, SecretError> {
|
||||
if hex.len() % 2 != 0 {
|
||||
if !hex.len().is_multiple_of(2) {
|
||||
return Err(SecretError::KeychainError(
|
||||
"Invalid hex string length".to_string(),
|
||||
));
|
||||
|
||||
+5
-1
@@ -64,7 +64,11 @@ mod store;
|
||||
mod types;
|
||||
|
||||
pub use crypto::SecretsCrypto;
|
||||
pub use store::{PostgresSecretsStore, SecretsStore};
|
||||
#[cfg(feature = "libsql")]
|
||||
pub use store::LibSqlSecretsStore;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use store::PostgresSecretsStore;
|
||||
pub use store::SecretsStore;
|
||||
pub use types::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, DecryptedSecret, Secret,
|
||||
SecretError, SecretRef,
|
||||
|
||||
+382
-14
@@ -10,6 +10,7 @@ use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::Pool;
|
||||
use secrecy::ExposeSecret;
|
||||
use uuid::Uuid;
|
||||
@@ -61,11 +62,13 @@ pub trait SecretsStore: Send + Sync {
|
||||
}
|
||||
|
||||
/// PostgreSQL implementation of SecretsStore.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct PostgresSecretsStore {
|
||||
pool: Pool,
|
||||
crypto: Arc<SecretsCrypto>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl PostgresSecretsStore {
|
||||
/// Create a new store with the given database pool and crypto instance.
|
||||
pub fn new(pool: Pool, crypto: Arc<SecretsCrypto>) -> Self {
|
||||
@@ -73,6 +76,7 @@ impl PostgresSecretsStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[async_trait]
|
||||
impl SecretsStore for PostgresSecretsStore {
|
||||
async fn create(
|
||||
@@ -149,10 +153,10 @@ impl SecretsStore for PostgresSecretsStore {
|
||||
let secret = row_to_secret(&r);
|
||||
|
||||
// Check expiration
|
||||
if let Some(expires_at) = secret.expires_at {
|
||||
if expires_at < Utc::now() {
|
||||
return Err(SecretError::Expired);
|
||||
}
|
||||
if let Some(expires_at) = secret.expires_at
|
||||
&& expires_at < Utc::now()
|
||||
{
|
||||
return Err(SecretError::Expired);
|
||||
}
|
||||
|
||||
Ok(secret)
|
||||
@@ -272,10 +276,10 @@ impl SecretsStore for PostgresSecretsStore {
|
||||
}
|
||||
|
||||
// Simple glob: * matches any suffix
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
if secret_name.starts_with(prefix) {
|
||||
return Ok(true);
|
||||
}
|
||||
if let Some(prefix) = pattern.strip_suffix('*')
|
||||
&& secret_name.starts_with(prefix)
|
||||
{
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -283,6 +287,7 @@ impl SecretsStore for PostgresSecretsStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
|
||||
Secret {
|
||||
id: row.get("id"),
|
||||
@@ -299,6 +304,332 @@ fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== libSQL implementation ====================
|
||||
|
||||
/// libSQL/Turso implementation of SecretsStore.
|
||||
///
|
||||
/// Holds an `Arc<Database>` handle and creates a fresh connection per operation,
|
||||
/// matching the connection-per-request pattern used by the main `LibSqlBackend`.
|
||||
#[cfg(feature = "libsql")]
|
||||
pub struct LibSqlSecretsStore {
|
||||
db: Arc<libsql::Database>,
|
||||
crypto: Arc<SecretsCrypto>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
impl LibSqlSecretsStore {
|
||||
/// Create a new store with the given shared libsql database handle and crypto instance.
|
||||
pub fn new(db: Arc<libsql::Database>, crypto: Arc<SecretsCrypto>) -> Self {
|
||||
Self { db, crypto }
|
||||
}
|
||||
|
||||
fn connect(&self) -> Result<libsql::Connection, SecretError> {
|
||||
self.db
|
||||
.connect()
|
||||
.map_err(|e| SecretError::Database(format!("Connection failed: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[async_trait]
|
||||
impl SecretsStore for LibSqlSecretsStore {
|
||||
async fn create(
|
||||
&self,
|
||||
user_id: &str,
|
||||
params: CreateSecretParams,
|
||||
) -> Result<Secret, SecretError> {
|
||||
let plaintext = params.value.expose_secret().as_bytes();
|
||||
let (encrypted_value, key_salt) = self.crypto.encrypt(plaintext)?;
|
||||
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
let now_str = now.to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let expires_at_str = params
|
||||
.expires_at
|
||||
.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
|
||||
|
||||
// Start transaction for atomic upsert + read-back
|
||||
let conn = self.connect()?;
|
||||
let tx = conn
|
||||
.transaction()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
tx.execute(
|
||||
r#"
|
||||
INSERT INTO secrets (id, user_id, name, encrypted_value, key_salt, provider, expires_at, created_at, updated_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8)
|
||||
ON CONFLICT (user_id, name) DO UPDATE SET
|
||||
encrypted_value = excluded.encrypted_value,
|
||||
key_salt = excluded.key_salt,
|
||||
provider = excluded.provider,
|
||||
expires_at = excluded.expires_at,
|
||||
updated_at = ?8
|
||||
"#,
|
||||
libsql::params![
|
||||
id.to_string(),
|
||||
user_id,
|
||||
params.name.as_str(),
|
||||
libsql::Value::Blob(encrypted_value.clone()),
|
||||
libsql::Value::Blob(key_salt.clone()),
|
||||
libsql_opt_text(params.provider.as_deref()),
|
||||
libsql_opt_text(expires_at_str.as_deref()),
|
||||
now_str.as_str(),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
// Read back the row (may have been upserted)
|
||||
let mut rows = tx
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||
last_used_at, usage_count, created_at, updated_at
|
||||
FROM secrets
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
"#,
|
||||
libsql::params![user_id, params.name.as_str()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
let row = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
.ok_or_else(|| SecretError::Database("Insert succeeded but row not found".into()))?;
|
||||
|
||||
let secret = libsql_row_to_secret(&row)?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(secret)
|
||||
}
|
||||
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||
last_used_at, usage_count, created_at, updated_at
|
||||
FROM secrets
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
"#,
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let secret = libsql_row_to_secret(&row)?;
|
||||
|
||||
if let Some(expires_at) = secret.expires_at
|
||||
&& expires_at < Utc::now()
|
||||
{
|
||||
return Err(SecretError::Expired);
|
||||
}
|
||||
|
||||
Ok(secret)
|
||||
}
|
||||
None => Err(SecretError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_decrypted(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<DecryptedSecret, SecretError> {
|
||||
let secret = self.get(user_id, name).await?;
|
||||
self.crypto
|
||||
.decrypt(&secret.encrypted_value, &secret.key_salt)
|
||||
}
|
||||
|
||||
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2",
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
.is_some())
|
||||
}
|
||||
|
||||
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT name, provider FROM secrets WHERE user_id = ?1 ORDER BY name",
|
||||
libsql::params![user_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
let mut refs = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
{
|
||||
refs.push(SecretRef {
|
||||
name: row.get::<String>(0).unwrap_or_default(),
|
||||
provider: row.get::<String>(1).ok(),
|
||||
});
|
||||
}
|
||||
Ok(refs)
|
||||
}
|
||||
|
||||
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let affected = conn
|
||||
.execute(
|
||||
"DELETE FROM secrets WHERE user_id = ?1 AND name = ?2",
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> {
|
||||
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let conn = self.connect()?;
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
UPDATE secrets
|
||||
SET last_used_at = ?1, usage_count = usage_count + 1
|
||||
WHERE id = ?2
|
||||
"#,
|
||||
libsql::params![now.as_str(), secret_id.to_string()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_accessible(
|
||||
&self,
|
||||
user_id: &str,
|
||||
secret_name: &str,
|
||||
allowed_secrets: &[String],
|
||||
) -> Result<bool, SecretError> {
|
||||
if !self.exists(user_id, secret_name).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
for pattern in allowed_secrets {
|
||||
if pattern == secret_name {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
if let Some(prefix) = pattern.strip_suffix('*')
|
||||
&& secret_name.starts_with(prefix)
|
||||
{
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_opt_text(s: Option<&str>) -> libsql::Value {
|
||||
match s {
|
||||
Some(s) => libsql::Value::Text(s.to_string()),
|
||||
None => libsql::Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_parse_timestamp(s: &str) -> Result<chrono::DateTime<Utc>, SecretError> {
|
||||
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
|
||||
return Ok(dt.with_timezone(&Utc));
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
Err(SecretError::Database(format!(
|
||||
"unparseable timestamp: {:?}",
|
||||
s
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_row_to_secret(row: &libsql::Row) -> Result<Secret, SecretError> {
|
||||
let id_str: String = row
|
||||
.get(0)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let user_id: String = row
|
||||
.get(1)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let name: String = row
|
||||
.get(2)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let encrypted_value: Vec<u8> = row
|
||||
.get(3)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let key_salt: Vec<u8> = row
|
||||
.get(4)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let provider: Option<String> = row.get::<String>(5).ok().filter(|s| !s.is_empty());
|
||||
let expires_at = row
|
||||
.get::<String>(6)
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.and_then(|s| libsql_parse_timestamp(&s).ok());
|
||||
let last_used_at = row
|
||||
.get::<String>(7)
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.and_then(|s| libsql_parse_timestamp(&s).ok());
|
||||
let usage_count: i64 = row.get::<i64>(8).unwrap_or(0);
|
||||
let created_at_str: String = row
|
||||
.get(9)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let updated_at_str: String = row
|
||||
.get(10)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(Secret {
|
||||
id: id_str
|
||||
.parse()
|
||||
.map_err(|e: uuid::Error| SecretError::Database(e.to_string()))?,
|
||||
user_id,
|
||||
name,
|
||||
encrypted_value,
|
||||
key_salt,
|
||||
provider,
|
||||
expires_at,
|
||||
last_used_at,
|
||||
usage_count,
|
||||
created_at: libsql_parse_timestamp(&created_at_str)?,
|
||||
updated_at: libsql_parse_timestamp(&updated_at_str)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// In-memory implementation for testing.
|
||||
#[cfg(test)]
|
||||
pub mod testing {
|
||||
@@ -364,12 +695,21 @@ pub mod testing {
|
||||
}
|
||||
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
|
||||
self.secrets
|
||||
let secret = self
|
||||
.secrets
|
||||
.read()
|
||||
.await
|
||||
.get(&(user_id.to_string(), name.to_string()))
|
||||
.cloned()
|
||||
.ok_or_else(|| SecretError::NotFound(name.to_string()))
|
||||
.ok_or_else(|| SecretError::NotFound(name.to_string()))?;
|
||||
|
||||
if let Some(expires_at) = secret.expires_at
|
||||
&& expires_at < Utc::now()
|
||||
{
|
||||
return Err(SecretError::Expired);
|
||||
}
|
||||
|
||||
Ok(secret)
|
||||
}
|
||||
|
||||
async fn get_decrypted(
|
||||
@@ -430,10 +770,10 @@ pub mod testing {
|
||||
if pattern == secret_name {
|
||||
return Ok(true);
|
||||
}
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
if secret_name.starts_with(prefix) {
|
||||
return Ok(true);
|
||||
}
|
||||
if let Some(prefix) = pattern.strip_suffix('*')
|
||||
&& secret_name.starts_with(prefix)
|
||||
{
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
@@ -558,6 +898,34 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_expired_secret_returns_error() {
|
||||
let store = test_store();
|
||||
let expires_at = chrono::Utc::now() - chrono::Duration::hours(1);
|
||||
let params = CreateSecretParams::new("expired_key", "value").with_expiry(expires_at);
|
||||
|
||||
store.create("user1", params).await.unwrap();
|
||||
|
||||
let result = store.get("user1", "expired_key").await;
|
||||
assert!(result.is_err());
|
||||
assert!(matches!(
|
||||
result.unwrap_err(),
|
||||
crate::secrets::SecretError::Expired
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_non_expired_secret_succeeds() {
|
||||
let store = test_store();
|
||||
let expires_at = chrono::Utc::now() + chrono::Duration::hours(1);
|
||||
let params = CreateSecretParams::new("fresh_key", "value").with_expiry(expires_at);
|
||||
|
||||
store.create("user1", params).await.unwrap();
|
||||
|
||||
let result = store.get("user1", "fresh_key").await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_user_isolation() {
|
||||
let store = test_store();
|
||||
|
||||
+87
-83
@@ -15,6 +15,10 @@ pub struct Settings {
|
||||
pub onboard_completed: bool,
|
||||
|
||||
// === Step 1: Database ===
|
||||
/// Database backend: "postgres" or "libsql".
|
||||
#[serde(default)]
|
||||
pub database_backend: Option<String>,
|
||||
|
||||
/// Database connection URL (postgres://...).
|
||||
#[serde(default)]
|
||||
pub database_url: Option<String>,
|
||||
@@ -23,13 +27,31 @@ pub struct Settings {
|
||||
#[serde(default)]
|
||||
pub database_pool_size: Option<usize>,
|
||||
|
||||
/// Path to local libSQL database file.
|
||||
#[serde(default)]
|
||||
pub libsql_path: Option<String>,
|
||||
|
||||
/// Turso cloud URL for remote replica sync.
|
||||
#[serde(default)]
|
||||
pub libsql_url: Option<String>,
|
||||
|
||||
// === Step 2: Security ===
|
||||
/// Source for the secrets master key.
|
||||
#[serde(default)]
|
||||
pub secrets_master_key_source: KeySource,
|
||||
|
||||
// === Step 3: NEAR AI Auth ===
|
||||
// Session stored separately in session.json
|
||||
// === Step 3: Inference Provider ===
|
||||
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible".
|
||||
#[serde(default)]
|
||||
pub llm_backend: Option<String>,
|
||||
|
||||
/// Ollama base URL (when llm_backend = "ollama").
|
||||
#[serde(default)]
|
||||
pub ollama_base_url: Option<String>,
|
||||
|
||||
/// OpenAI-compatible endpoint base URL (when llm_backend = "openai_compatible").
|
||||
#[serde(default)]
|
||||
pub openai_compatible_base_url: Option<String>,
|
||||
|
||||
// === Step 4: Model Selection ===
|
||||
/// Currently selected model.
|
||||
@@ -487,20 +509,16 @@ impl Default for BuilderSettings {
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
/// Get the default settings file path (~/.ironclaw/settings.json).
|
||||
pub fn default_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("settings.json")
|
||||
}
|
||||
|
||||
/// Reconstruct Settings from a flat key-value map (as stored in the DB).
|
||||
///
|
||||
/// Each key is a dotted path (e.g., "agent.name"), value is a JSONB value.
|
||||
/// Missing keys get their default value.
|
||||
pub fn from_db_map(map: &std::collections::HashMap<String, serde_json::Value>) -> Self {
|
||||
// Start with defaults, then overlay each DB setting
|
||||
// Start with defaults, then overlay each DB setting.
|
||||
//
|
||||
// The settings table stores both Settings struct fields and app-specific
|
||||
// data (e.g. nearai.session_token). Skip keys that don't correspond to
|
||||
// a known Settings path.
|
||||
let mut settings = Self::default();
|
||||
|
||||
for (key, value) in map {
|
||||
@@ -509,17 +527,23 @@ impl Settings {
|
||||
serde_json::Value::String(s) => s.clone(),
|
||||
serde_json::Value::Bool(b) => b.to_string(),
|
||||
serde_json::Value::Number(n) => n.to_string(),
|
||||
serde_json::Value::Null => "null".to_string(),
|
||||
serde_json::Value::Null => continue, // null means default, skip
|
||||
other => other.to_string(),
|
||||
};
|
||||
|
||||
if let Err(e) = settings.set(key, &value_str) {
|
||||
tracing::warn!(
|
||||
"Failed to apply DB setting '{}' = '{}': {}",
|
||||
key,
|
||||
value_str,
|
||||
e
|
||||
);
|
||||
match settings.set(key, &value_str) {
|
||||
Ok(()) => {}
|
||||
// The settings table stores both Settings fields and app-specific
|
||||
// data (e.g. nearai.session_token). Silently skip unknown paths.
|
||||
Err(e) if e.starts_with("Path not found") => {}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to apply DB setting '{}' = '{}': {}",
|
||||
key,
|
||||
value_str,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -540,50 +564,27 @@ impl Settings {
|
||||
map
|
||||
}
|
||||
|
||||
/// Get the default settings file path (~/.ironclaw/settings.json).
|
||||
pub fn default_path() -> std::path::PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| std::path::PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("settings.json")
|
||||
}
|
||||
|
||||
/// Load settings from disk, returning default if not found.
|
||||
pub fn load() -> Self {
|
||||
Self::load_from(&Self::default_path())
|
||||
}
|
||||
|
||||
/// Load settings from a specific path.
|
||||
pub fn load_from(path: &PathBuf) -> Self {
|
||||
/// Load settings from a specific path (used by bootstrap legacy migration).
|
||||
pub fn load_from(path: &std::path::Path) -> Self {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
|
||||
Err(_) => Self::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Save settings to disk.
|
||||
pub fn save(&self) -> std::io::Result<()> {
|
||||
self.save_to(&Self::default_path())
|
||||
}
|
||||
|
||||
/// Save settings to a specific path.
|
||||
pub fn save_to(&self, path: &PathBuf) -> std::io::Result<()> {
|
||||
// Ensure parent directory exists
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
let json = serde_json::to_string_pretty(self)
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
|
||||
|
||||
std::fs::write(path, json)
|
||||
}
|
||||
|
||||
/// Get the selected model, falling back to the provided default.
|
||||
pub fn model_or(&self, default: &str) -> String {
|
||||
self.selected_model
|
||||
.clone()
|
||||
.unwrap_or_else(|| default.to_string())
|
||||
}
|
||||
|
||||
/// Set the selected model and save.
|
||||
pub fn set_model(&mut self, model: &str) -> std::io::Result<()> {
|
||||
self.selected_model = Some(model.to_string());
|
||||
self.save()
|
||||
}
|
||||
|
||||
/// Get a setting value by dotted path (e.g., "agent.max_parallel_jobs").
|
||||
pub fn get(&self, path: &str) -> Option<String> {
|
||||
let json = serde_json::to_value(self).ok()?;
|
||||
@@ -768,42 +769,22 @@ fn collect_settings(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn test_settings_save_load() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("settings.json");
|
||||
|
||||
fn test_db_map_round_trip() {
|
||||
let settings = Settings {
|
||||
selected_model: Some("claude-3-5-sonnet-20241022".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
settings.save_to(&path).unwrap();
|
||||
|
||||
let loaded = Settings::load_from(&path);
|
||||
let map = settings.to_db_map();
|
||||
let restored = Settings::from_db_map(&map);
|
||||
assert_eq!(
|
||||
loaded.selected_model,
|
||||
restored.selected_model,
|
||||
Some("claude-3-5-sonnet-20241022".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_or_default() {
|
||||
let settings = Settings::default();
|
||||
assert_eq!(
|
||||
settings.model_or("default-model"),
|
||||
"default-model".to_string()
|
||||
);
|
||||
|
||||
let settings = Settings {
|
||||
selected_model: Some("my-model".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(settings.model_or("default-model"), "my-model".to_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_setting() {
|
||||
let settings = Settings::default();
|
||||
@@ -874,16 +855,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_telegram_owner_id_round_trip() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("settings.json");
|
||||
|
||||
fn test_telegram_owner_id_db_round_trip() {
|
||||
let mut settings = Settings::default();
|
||||
settings.channels.telegram_owner_id = Some(123456789);
|
||||
settings.save_to(&path).unwrap();
|
||||
|
||||
let loaded = Settings::load_from(&path);
|
||||
assert_eq!(loaded.channels.telegram_owner_id, Some(123456789));
|
||||
let map = settings.to_db_map();
|
||||
let restored = Settings::from_db_map(&map);
|
||||
assert_eq!(restored.channels.telegram_owner_id, Some(123456789));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -900,4 +878,30 @@ mod tests {
|
||||
.unwrap();
|
||||
assert_eq!(settings.channels.telegram_owner_id, Some(987654321));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_llm_backend_round_trip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("settings.json");
|
||||
|
||||
let settings = Settings {
|
||||
llm_backend: Some("anthropic".to_string()),
|
||||
ollama_base_url: Some("http://localhost:11434".to_string()),
|
||||
openai_compatible_base_url: Some("http://my-vllm:8000/v1".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let json = serde_json::to_string_pretty(&settings).unwrap();
|
||||
std::fs::write(&path, json).unwrap();
|
||||
|
||||
let loaded = Settings::load_from(&path);
|
||||
assert_eq!(loaded.llm_backend, Some("anthropic".to_string()));
|
||||
assert_eq!(
|
||||
loaded.ollama_base_url,
|
||||
Some("http://localhost:11434".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
loaded.openai_compatible_base_url,
|
||||
Some("http://my-vllm:8000/v1".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,539 @@
|
||||
# Setup / Onboarding Specification
|
||||
|
||||
This document is the authoritative specification for IronClaw's onboarding
|
||||
wizard. Any code change to `src/setup/` **must** keep this document in sync.
|
||||
If a future contributor or coding agent modifies setup behavior, update this
|
||||
file first, then adjust the code to match.
|
||||
|
||||
---
|
||||
|
||||
## Entry Points
|
||||
|
||||
```
|
||||
ironclaw onboard [--skip-auth] [--channels-only]
|
||||
```
|
||||
|
||||
Explicit invocation. Loads `.env` files, runs the wizard, exits.
|
||||
|
||||
```
|
||||
ironclaw (first run, no database configured)
|
||||
```
|
||||
|
||||
Auto-detection via `check_onboard_needed()` in `main.rs`. Triggers when
|
||||
none of these are true:
|
||||
- `DATABASE_URL` env var is set
|
||||
- `LIBSQL_PATH` env var is set
|
||||
- `~/.ironclaw/ironclaw.db` exists on disk
|
||||
|
||||
The `--no-onboard` CLI flag suppresses auto-detection.
|
||||
|
||||
---
|
||||
|
||||
## Startup Sequence (main.rs)
|
||||
|
||||
```
|
||||
1. Parse CLI args
|
||||
2. If Command::Onboard → load .env, run wizard, exit
|
||||
3. If Command::Run or no command:
|
||||
a. Load .env files (dotenvy::dotenv() then load_ironclaw_env())
|
||||
b. check_onboard_needed() → run wizard if needed
|
||||
c. Config::from_env() → build config from env vars
|
||||
d. Create SessionManager → load session token
|
||||
e. ensure_authenticated() → validate session (NEAR AI only)
|
||||
f. ... rest of agent startup
|
||||
```
|
||||
|
||||
**Critical ordering:** `.env` files must be loaded (step 3a) before
|
||||
`Config::from_env()` (step 3c) because bootstrap vars like
|
||||
`DATABASE_BACKEND` live in `~/.ironclaw/.env`.
|
||||
|
||||
---
|
||||
|
||||
## The 7-Step Wizard
|
||||
|
||||
### Overview
|
||||
|
||||
```
|
||||
Step 1: Database Connection
|
||||
Step 2: Security (master key)
|
||||
Step 3: Inference Provider ← skipped if --skip-auth
|
||||
Step 4: Model Selection
|
||||
Step 5: Embeddings
|
||||
Step 6: Channel Configuration
|
||||
Step 7: Background Tasks (heartbeat)
|
||||
↓
|
||||
save_and_summarize()
|
||||
```
|
||||
|
||||
`--channels-only` mode runs only Step 6, skipping everything else.
|
||||
|
||||
---
|
||||
|
||||
### Step 1: Database Connection
|
||||
|
||||
**Module:** `wizard.rs` → `step_database()`
|
||||
|
||||
**Goal:** Select backend, establish connection, run migrations.
|
||||
|
||||
**Decision tree:**
|
||||
|
||||
```
|
||||
Both features compiled?
|
||||
├─ Yes → DATABASE_BACKEND env var set?
|
||||
│ ├─ Yes → use that backend
|
||||
│ └─ No → interactive selection (PostgreSQL vs libSQL)
|
||||
├─ Only postgres feature → step_database_postgres()
|
||||
└─ Only libsql feature → step_database_libsql()
|
||||
```
|
||||
|
||||
**PostgreSQL path** (`step_database_postgres`):
|
||||
1. Check `DATABASE_URL` from env or settings
|
||||
2. Test connection (creates `deadpool_postgres::Pool`)
|
||||
3. Optionally run refinery migrations
|
||||
4. Store pool in `self.db_pool`
|
||||
|
||||
**libSQL path** (`step_database_libsql`):
|
||||
1. Offer local path (default: `~/.ironclaw/ironclaw.db`)
|
||||
2. Optional Turso cloud sync (URL + auth token)
|
||||
3. Test connection (creates `LibSqlBackend`)
|
||||
4. Always run migrations (idempotent CREATE IF NOT EXISTS)
|
||||
5. Store backend in `self.db_backend`
|
||||
|
||||
**Invariant:** After Step 1, exactly one of `self.db_pool` or
|
||||
`self.db_backend` is `Some`. This is required for settings persistence
|
||||
in `save_and_summarize()`.
|
||||
|
||||
---
|
||||
|
||||
### Step 2: Security (Master Key)
|
||||
|
||||
**Module:** `wizard.rs` → `step_security()`
|
||||
|
||||
**Goal:** Configure encryption for API tokens and secrets.
|
||||
|
||||
**Decision tree:**
|
||||
|
||||
```
|
||||
SECRETS_MASTER_KEY env var set?
|
||||
├─ Yes → use env var, done
|
||||
└─ No → try get_master_key() from OS keychain
|
||||
├─ Ok(bytes) → cache in self.secrets_crypto, ask "use existing?"
|
||||
│ ├─ Yes → done (keychain)
|
||||
│ └─ No → clear cache, fall through to options
|
||||
└─ Err → fall through to options
|
||||
├─ OS Keychain: generate + store + build SecretsCrypto
|
||||
├─ Env variable: generate + print export command
|
||||
└─ Skip: disable secrets features
|
||||
```
|
||||
|
||||
**CRITICAL CAVEAT: macOS Keychain Dialogs**
|
||||
|
||||
On macOS, `security_framework::get_generic_password()` can trigger TWO
|
||||
system dialogs:
|
||||
1. "Enter your password to unlock the keychain" (keychain locked)
|
||||
2. "Allow ironclaw to access this keychain item" (per-app authorization)
|
||||
|
||||
This is OS-level behavior we cannot prevent. To minimize pain:
|
||||
|
||||
- **Use `get_master_key()` not `has_master_key()`** in step 2. Both call
|
||||
the same underlying API, but `get_master_key()` returns the key bytes
|
||||
so we can cache them. `has_master_key()` throws them away, forcing a
|
||||
second keychain access later.
|
||||
|
||||
- **Build `SecretsCrypto` eagerly.** When the keychain key is retrieved,
|
||||
immediately construct `SecretsCrypto` and store in `self.secrets_crypto`.
|
||||
Later calls to `init_secrets_context()` check this field first, avoiding
|
||||
redundant keychain probes.
|
||||
|
||||
- **Never probe the keychain in read-only commands** (e.g., `ironclaw status`).
|
||||
The status command reports "env not set (keychain may be configured)"
|
||||
rather than triggering system dialogs.
|
||||
|
||||
**Invariant:** After Step 2, `self.secrets_crypto` is `Some` if the user
|
||||
chose Keychain or generated a new key. It may be `None` if the user chose
|
||||
env-var mode or skipped secrets.
|
||||
|
||||
---
|
||||
|
||||
### Step 3: Inference Provider
|
||||
|
||||
**Module:** `wizard.rs` → `step_inference_provider()`
|
||||
|
||||
**Goal:** Choose LLM backend and authenticate.
|
||||
|
||||
**Providers:**
|
||||
|
||||
| Provider | Auth Method | Secret Name | Env Var |
|
||||
|----------|-------------|-------------|---------|
|
||||
| NEAR AI | Browser OAuth | (session token) | `NEARAI_SESSION_TOKEN` |
|
||||
| Anthropic | API key | `anthropic_api_key` | `ANTHROPIC_API_KEY` |
|
||||
| OpenAI | API key | `openai_api_key` | `OPENAI_API_KEY` |
|
||||
| Ollama | None | - | - |
|
||||
| OpenAI-compatible | Optional API key | `llm_compatible_api_key` | `LLM_API_KEY` |
|
||||
|
||||
**API-key providers** (`setup_api_key_provider`):
|
||||
1. Check env var → if set, ask to reuse, persist to secrets store
|
||||
2. Otherwise prompt for key entry via `secret_input()`
|
||||
3. Store encrypted in secrets via `init_secrets_context()`
|
||||
4. **Cache key in `self.llm_api_key`** for model fetching in Step 4
|
||||
|
||||
**NEAR AI** (`setup_nearai`):
|
||||
- Calls `session_manager.ensure_authenticated()` which opens browser
|
||||
- Session token saved to `~/.ironclaw/session.json`
|
||||
|
||||
**`self.llm_api_key` caching:** The wizard caches the API key as
|
||||
`Option<SecretString>` so that Step 4 (model fetching) and Step 5
|
||||
(embeddings) can use it without re-reading from the secrets store or
|
||||
mutating environment variables.
|
||||
|
||||
---
|
||||
|
||||
### Step 4: Model Selection
|
||||
|
||||
**Module:** `wizard.rs` → `step_model_selection()`
|
||||
|
||||
**Goal:** Choose which model to use.
|
||||
|
||||
**Flow:**
|
||||
1. If model already set → offer to keep it
|
||||
2. Fetch models from provider API (5-second timeout)
|
||||
3. On timeout or error → use static fallback list
|
||||
4. Present list + "Custom model ID" escape hatch
|
||||
5. Store in `self.settings.selected_model`
|
||||
|
||||
**Model fetchers pass the cached API key explicitly:**
|
||||
```rust
|
||||
let cached = self.llm_api_key.as_ref().map(|k| k.expose_secret().to_string());
|
||||
let models = fetch_anthropic_models(cached.as_deref()).await;
|
||||
```
|
||||
|
||||
This avoids mutating environment variables. The fetcher checks the explicit
|
||||
key first, then falls back to the standard env var.
|
||||
|
||||
---
|
||||
|
||||
### Step 5: Embeddings
|
||||
|
||||
**Module:** `wizard.rs` → `step_embeddings()`
|
||||
|
||||
**Goal:** Configure semantic search for workspace memory.
|
||||
|
||||
**Flow:**
|
||||
1. Ask "Enable semantic search?" (default: yes)
|
||||
2. Detect available providers:
|
||||
- NEAR AI: if backend is `nearai` OR valid session exists
|
||||
- OpenAI: if `OPENAI_API_KEY` in env OR (backend is `openai` AND cached key)
|
||||
3. If both available → let user choose
|
||||
4. If only one → use it
|
||||
5. If neither → disable embeddings
|
||||
|
||||
**Default model:** `text-embedding-3-small` (for both providers)
|
||||
|
||||
---
|
||||
|
||||
### Step 6: Channel Configuration
|
||||
|
||||
**Module:** `wizard.rs` → `step_channels()`, delegating to `channels.rs`
|
||||
|
||||
**Goal:** Enable input channels (TUI, HTTP, Telegram, etc.).
|
||||
|
||||
**Sub-steps:**
|
||||
|
||||
```
|
||||
6a. Tunnel setup (if webhook channels needed)
|
||||
6b. Discover WASM channels from ~/.ironclaw/channels/
|
||||
6c. Multi-select: CLI/TUI, HTTP, discovered channels, bundled channels
|
||||
6d. Install missing bundled channels (copy WASM binaries)
|
||||
6e. Initialize SecretsContext (for token storage)
|
||||
6f. Setup HTTP webhook (if selected)
|
||||
6g. Setup each WASM channel (secrets, owner binding)
|
||||
```
|
||||
|
||||
**Tunnel setup** (`setup_tunnel`):
|
||||
- Options: ngrok, Cloudflare Tunnel, localtunnel, custom URL
|
||||
- Validates HTTPS requirement
|
||||
- Stored in `self.settings.tunnel.public_url`
|
||||
|
||||
**WASM channel setup** (`setup_wasm_channel`):
|
||||
- Reads `capabilities.json` for `setup.required_secrets`
|
||||
- For each secret: check existing, prompt or auto-generate, validate regex
|
||||
- Save each secret via `SecretsContext`
|
||||
|
||||
**Telegram special case** (`setup_telegram`):
|
||||
- Validates bot token via Telegram `getMe` API
|
||||
- Owner binding: polls `getUpdates` for 120s to capture sender's user ID
|
||||
- Optional webhook secret generation
|
||||
|
||||
**SecretsContext creation** (`init_secrets_context`):
|
||||
1. Check `self.secrets_crypto` (set in Step 2) → use if available
|
||||
2. Else try `SECRETS_MASTER_KEY` env var
|
||||
3. Else try `get_master_key()` from keychain (only in `channels_only` mode)
|
||||
4. Create backend-appropriate secrets store (respects selected database backend)
|
||||
|
||||
---
|
||||
|
||||
### Step 7: Heartbeat
|
||||
|
||||
**Module:** `wizard.rs` → `step_heartbeat()`
|
||||
|
||||
**Goal:** Configure periodic background execution.
|
||||
|
||||
**Flow:**
|
||||
1. Ask "Enable heartbeat?" (default: no)
|
||||
2. If yes: interval in minutes (default: 30), notification channel
|
||||
3. Store in `self.settings.heartbeat`
|
||||
|
||||
---
|
||||
|
||||
## Settings Persistence
|
||||
|
||||
### Two-Layer Architecture
|
||||
|
||||
Settings are persisted in two places:
|
||||
|
||||
**Layer 1: `~/.ironclaw/.env`** (bootstrap vars)
|
||||
|
||||
Contains only the settings needed BEFORE database connection. Written by
|
||||
`save_bootstrap_env()` in `bootstrap.rs`.
|
||||
|
||||
```env
|
||||
DATABASE_BACKEND="libsql"
|
||||
LIBSQL_PATH="/Users/name/.ironclaw/ironclaw.db"
|
||||
```
|
||||
|
||||
Or for PostgreSQL:
|
||||
```env
|
||||
DATABASE_BACKEND="postgres"
|
||||
DATABASE_URL="postgres://user:pass@localhost/ironclaw"
|
||||
```
|
||||
|
||||
**Why separate?** Chicken-and-egg: you need `DATABASE_BACKEND` to know
|
||||
which database to connect to, so it can't be stored in the database.
|
||||
|
||||
**Layer 2: Database settings table** (everything else)
|
||||
|
||||
All other settings are stored as key-value pairs in the `settings` table,
|
||||
keyed by `(user_id, key)`. Written by `set_all_settings()`.
|
||||
|
||||
Settings are serialized via `Settings::to_db_map()` as dotted paths:
|
||||
```
|
||||
database_backend = "libsql"
|
||||
llm_backend = "nearai"
|
||||
selected_model = "anthropic/claude-sonnet-4-5"
|
||||
embeddings.enabled = "true"
|
||||
embeddings.provider = "nearai"
|
||||
channels.http_enabled = "true"
|
||||
heartbeat.enabled = "true"
|
||||
heartbeat.interval_secs = "300"
|
||||
```
|
||||
|
||||
### save_and_summarize()
|
||||
|
||||
Final step of the wizard:
|
||||
|
||||
```
|
||||
1. Mark onboard_completed = true
|
||||
2. Write ALL settings to database (try postgres pool, then libSQL backend)
|
||||
3. Write bootstrap vars to ~/.ironclaw/.env:
|
||||
- DATABASE_BACKEND (always)
|
||||
- DATABASE_URL (if postgres)
|
||||
- LIBSQL_PATH (if libsql)
|
||||
- LIBSQL_URL (if turso sync)
|
||||
4. Print configuration summary
|
||||
```
|
||||
|
||||
**Invariant:** Both Layer 1 and Layer 2 must be written. If the database
|
||||
write fails, the wizard returns an error and the `.env` file is not written.
|
||||
|
||||
### Legacy Migration
|
||||
|
||||
`bootstrap.rs` handles one-time upgrades from older config formats:
|
||||
- `bootstrap.json` → extracts `DATABASE_URL`, writes `.env`, renames to `.migrated`
|
||||
- `settings.json` → migrated to database via `migrate_disk_to_db()`
|
||||
|
||||
---
|
||||
|
||||
## Settings Struct
|
||||
|
||||
**Module:** `settings.rs`
|
||||
|
||||
```rust
|
||||
pub struct Settings {
|
||||
// Meta
|
||||
pub onboard_completed: bool,
|
||||
|
||||
// Step 1: Database
|
||||
pub database_backend: Option<String>, // "postgres" | "libsql"
|
||||
pub database_url: Option<String>,
|
||||
pub libsql_path: Option<String>,
|
||||
pub libsql_url: Option<String>,
|
||||
|
||||
// Step 2: Security
|
||||
pub secrets_master_key_source: KeySource, // Keychain | Env | None
|
||||
|
||||
// Step 3: Inference
|
||||
pub llm_backend: Option<String>, // "nearai" | "anthropic" | "openai" | "ollama" | "openai_compatible"
|
||||
pub ollama_base_url: Option<String>,
|
||||
pub openai_compatible_base_url: Option<String>,
|
||||
|
||||
// Step 4: Model
|
||||
pub selected_model: Option<String>,
|
||||
|
||||
// Step 5: Embeddings
|
||||
pub embeddings: EmbeddingsSettings, // enabled, provider, model
|
||||
|
||||
// Step 6: Channels
|
||||
pub tunnel: TunnelSettings, // provider, public_url
|
||||
pub channels: ChannelSettings, // http config, telegram owner, etc.
|
||||
|
||||
// Step 7: Heartbeat
|
||||
pub heartbeat: HeartbeatSettings, // enabled, interval, notify
|
||||
|
||||
// Advanced (not in wizard, set via `ironclaw config set`)
|
||||
pub agent: AgentSettings,
|
||||
pub wasm: WasmSettings,
|
||||
pub sandbox: SandboxSettings,
|
||||
pub safety: SafetySettings,
|
||||
pub builder: BuilderSettings,
|
||||
}
|
||||
```
|
||||
|
||||
**KeySource enum:** `Keychain | Env | None`
|
||||
|
||||
---
|
||||
|
||||
## Secrets Flow
|
||||
|
||||
### SecretsContext
|
||||
|
||||
Thin wrapper for setup-time secret operations:
|
||||
|
||||
```rust
|
||||
pub struct SecretsContext {
|
||||
store: Arc<dyn SecretsStore>,
|
||||
user_id: String,
|
||||
}
|
||||
```
|
||||
|
||||
Created by `init_secrets_context()` which:
|
||||
1. Gets `SecretsCrypto` from `self.secrets_crypto` or loads from keychain/env
|
||||
2. Creates the appropriate backend store:
|
||||
- If both features compiled: respects `self.settings.database_backend`
|
||||
- Tries selected backend first, falls back to the other
|
||||
3. Returns `SecretsContext` wrapping the store
|
||||
|
||||
### Secret Storage
|
||||
|
||||
Secrets are encrypted with AES-256-GCM using the master key, then stored
|
||||
in the database `secrets` table. The wizard writes secrets like:
|
||||
|
||||
```
|
||||
telegram_bot_token → encrypted bot token
|
||||
telegram_webhook_secret → encrypted webhook HMAC secret
|
||||
anthropic_api_key → encrypted API key
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Prompt Utilities
|
||||
|
||||
**Module:** `prompts.rs`
|
||||
|
||||
| Function | Description |
|
||||
|----------|-------------|
|
||||
| `select_one(label, options)` | Numbered single-choice menu |
|
||||
| `select_many(label, options, defaults)` | Checkbox multi-select (raw terminal mode) |
|
||||
| `input(label)` | Single line text input |
|
||||
| `optional_input(label, hint)` | Text input that can be empty |
|
||||
| `secret_input(label)` | Hidden input (shows `*` per char), returns `SecretString` |
|
||||
| `confirm(label, default)` | `[Y/n]` or `[y/N]` prompt |
|
||||
| `print_header(text)` | Bold section header with underline |
|
||||
| `print_step(n, total, text)` | `[1/7] Step Name` |
|
||||
| `print_success(text)` | Green checkmark prefix |
|
||||
| `print_error(text)` | Red X prefix |
|
||||
| `print_info(text)` | Blue info prefix |
|
||||
|
||||
`select_many` uses `crossterm` raw mode for arrow key navigation.
|
||||
Must properly restore terminal state on all exit paths.
|
||||
|
||||
---
|
||||
|
||||
## Platform Caveats
|
||||
|
||||
### macOS Keychain
|
||||
|
||||
- `get_generic_password()` triggers system dialogs (unlock + authorize)
|
||||
- Two dialogs per call is normal, not a bug
|
||||
- Cache the result after first access to avoid repeat prompts
|
||||
- Never probe keychain in read-only commands (`status`, `--help`)
|
||||
- Service name: `"ironclaw"`, account: `"master_key"`
|
||||
|
||||
### Linux Secret Service
|
||||
|
||||
- Uses GNOME Keyring or KWallet via `secret-service` crate
|
||||
- May need `gnome-keyring` daemon running
|
||||
- Collection unlock may prompt for password
|
||||
|
||||
### URL Passwords
|
||||
|
||||
- `#` is common in URL-encoded passwords (`%23` decoded)
|
||||
- `.env` values must be double-quoted to preserve `#`
|
||||
- Display masked: `postgres://user:****@host/db`
|
||||
|
||||
### Telegram API
|
||||
|
||||
- Bot token format: `123456:ABC-DEF...`
|
||||
- Token goes in URL path: `https://api.telegram.org/bot{TOKEN}/method`
|
||||
- Webhook secret header: `X-Telegram-Bot-Api-Secret-Token`
|
||||
- Owner binding polls `getUpdates` (must delete webhook first)
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
Tests live in `mod tests {}` at the bottom of each file.
|
||||
|
||||
**What to test when modifying setup:**
|
||||
|
||||
- Settings round-trip: `to_db_map()` then `from_db_map()` preserves values
|
||||
- Bootstrap `.env`: dotenvy can parse what `save_bootstrap_env()` writes
|
||||
- Model fetchers: static fallback works when API is unreachable
|
||||
- Channel discovery: handles missing dir, invalid JSON, deduplication
|
||||
- Prompt functions: not tested (interactive I/O), but ensure error paths
|
||||
don't panic
|
||||
|
||||
**Run setup tests:**
|
||||
```bash
|
||||
cargo test --lib -- setup
|
||||
cargo test --lib -- bootstrap
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Modification Checklist
|
||||
|
||||
When changing the onboarding flow:
|
||||
|
||||
1. Update this README first with the intended behavior change
|
||||
2. If adding a new wizard step:
|
||||
- Add to the step enum in `run()`, adjust `total_steps`
|
||||
- Add corresponding settings fields to `Settings`
|
||||
- Add `to_db_map` / `from_db_map` serialization
|
||||
- If the setting is needed before DB connection, add to `save_bootstrap_env()`
|
||||
3. If adding a new provider or channel:
|
||||
- Add to the selection menu in the appropriate step
|
||||
- Add authentication flow (API key or OAuth)
|
||||
- Add model fetcher with static fallback + 5s timeout
|
||||
4. If touching keychain:
|
||||
- Cache the result, never call `get_master_key()` twice
|
||||
- Test on macOS (dialog behavior differs from Linux)
|
||||
5. If touching secrets:
|
||||
- Ensure `init_secrets_context()` respects the selected database backend
|
||||
- Test with both postgres and libsql features
|
||||
6. Run the full shipping checklist:
|
||||
```bash
|
||||
cargo fmt
|
||||
cargo clippy --all --benches --tests --examples --all-features -- -D warnings
|
||||
cargo test --lib -- setup bootstrap
|
||||
```
|
||||
7. Test a fresh onboarding: `rm -rf ~/.ironclaw && cargo run`
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user