mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
53
Commits
Generated
+2
@@ -3482,6 +3482,7 @@ dependencies = [
|
||||
"wasmparser 0.220.1",
|
||||
"wasmtime",
|
||||
"wasmtime-wasi",
|
||||
"webpki-roots 0.26.11",
|
||||
"zbus",
|
||||
"zip",
|
||||
]
|
||||
@@ -6991,6 +6992,7 @@ dependencies = [
|
||||
"futures-util",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"iri-string",
|
||||
"pin-project-lite",
|
||||
"tower 0.5.3",
|
||||
|
||||
+4
-1
@@ -57,6 +57,7 @@ refinery = { version = "0.8", features = ["tokio-postgres"], optional = true }
|
||||
tokio-postgres-rustls = { version = "0.13", optional = true }
|
||||
rustls = { version = "0.23", optional = true, default-features = false }
|
||||
rustls-native-certs = { version = "0.8", optional = true }
|
||||
webpki-roots = { version = "0.26", optional = true }
|
||||
|
||||
# Database - libSQL/Turso (optional embedded database)
|
||||
libsql = { version = "0.6", optional = true, default-features = false, features = ["core", "replication", "remote", "tls"] }
|
||||
@@ -95,7 +96,7 @@ termimad = "0.34"
|
||||
# Channel integrations
|
||||
axum = { version = "0.8", features = ["ws"] }
|
||||
tower = "0.5"
|
||||
tower-http = { version = "0.6", features = ["trace", "cors", "set-header"] }
|
||||
tower-http = { version = "0.6", features = ["trace", "cors", "set-header", "catch-panic"] }
|
||||
|
||||
# Cron scheduling for routines
|
||||
cron = "0.13"
|
||||
@@ -219,6 +220,7 @@ postgres = [
|
||||
"dep:tokio-postgres-rustls",
|
||||
"dep:rustls",
|
||||
"dep:rustls-native-certs",
|
||||
"dep:webpki-roots",
|
||||
"dep:postgres-types",
|
||||
"dep:refinery",
|
||||
"dep:pgvector",
|
||||
@@ -230,6 +232,7 @@ libsql = ["dep:libsql"]
|
||||
integration = []
|
||||
html-to-markdown = ["dep:html-to-markdown-rs", "dep:readabilityrs"]
|
||||
bedrock = ["dep:aws-config", "dep:aws-sdk-bedrockruntime", "dep:aws-smithy-types"]
|
||||
demo = []
|
||||
import = ["dep:json5", "libsql"]
|
||||
|
||||
[[test]]
|
||||
|
||||
+39
-9
@@ -1,49 +1,78 @@
|
||||
# Multi-stage Dockerfile for the IronClaw agent (cloud deployment).
|
||||
#
|
||||
# Uses cargo-chef for dependency caching — only rebuilds deps when
|
||||
# Cargo.toml/Cargo.lock change, not on every source edit.
|
||||
#
|
||||
# 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
|
||||
# Stage 1: Install cargo-chef
|
||||
FROM rust:1.92-slim-bookworm AS chef
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
pkg-config libssl-dev cmake gcc g++ \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& rustup target add wasm32-wasip2 \
|
||||
&& cargo install wasm-tools
|
||||
&& cargo install cargo-chef wasm-tools
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy manifests first for layer caching
|
||||
# Stage 2: Generate the dependency recipe (changes only when Cargo.toml/lock change)
|
||||
FROM chef AS planner
|
||||
|
||||
COPY Cargo.toml Cargo.lock ./
|
||||
COPY crates/ crates/
|
||||
|
||||
# Copy source, build script, tests, and supporting directories
|
||||
COPY build.rs build.rs
|
||||
COPY src/ src/
|
||||
COPY tests/ tests/
|
||||
COPY benches/ benches/
|
||||
COPY migrations/ migrations/
|
||||
COPY registry/ registry/
|
||||
COPY channels-src/ channels-src/
|
||||
COPY wit/ wit/
|
||||
COPY providers.json providers.json
|
||||
# [[bench]] entries in Cargo.toml require bench sources to exist for cargo to parse the manifest
|
||||
|
||||
RUN cargo chef prepare --recipe-path recipe.json
|
||||
|
||||
# Stage 3: Build dependencies (cached unless Cargo.toml/lock change)
|
||||
FROM chef AS deps
|
||||
|
||||
COPY --from=planner /app/recipe.json recipe.json
|
||||
RUN cargo chef cook --release --recipe-path recipe.json
|
||||
|
||||
# Stage 4: Build the actual binary (only recompiles ironclaw source)
|
||||
FROM deps AS builder
|
||||
|
||||
COPY Cargo.toml Cargo.lock ./
|
||||
COPY crates/ crates/
|
||||
COPY build.rs build.rs
|
||||
COPY src/ src/
|
||||
COPY tests/ tests/
|
||||
COPY benches/ benches/
|
||||
COPY migrations/ migrations/
|
||||
COPY registry/ registry/
|
||||
COPY channels-src/ channels-src/
|
||||
COPY wit/ wit/
|
||||
COPY providers.json providers.json
|
||||
|
||||
RUN cargo build --release --bin ironclaw
|
||||
COPY skills/ skills/
|
||||
|
||||
# Stage 2: Runtime
|
||||
RUN cargo build --release --features demo --bin ironclaw
|
||||
|
||||
# Stage 5: Runtime
|
||||
FROM debian:bookworm-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates libssl3 \
|
||||
&& update-ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=builder /app/target/release/ironclaw /usr/local/bin/ironclaw
|
||||
COPY --from=builder /app/migrations /app/migrations
|
||||
COPY --from=builder /app/skills /app/skills
|
||||
|
||||
# Non-root user
|
||||
RUN useradd -m -u 1000 -s /bin/bash ironclaw
|
||||
@@ -52,5 +81,6 @@ USER ironclaw
|
||||
EXPOSE 3000
|
||||
|
||||
ENV RUST_LOG=ironclaw=info
|
||||
ENV SKILLS_DIR=/app/skills
|
||||
|
||||
ENTRYPOINT ["ironclaw"]
|
||||
|
||||
@@ -0,0 +1,571 @@
|
||||
# User Management API
|
||||
|
||||
DB-backed user management for multi-tenant IronClaw deployments. Covers admin user CRUD, per-user secrets provisioning, self-service profile, API token management, and usage reporting.
|
||||
|
||||
## Authentication
|
||||
|
||||
All endpoints require `Authorization: Bearer <token>`. Tokens are either:
|
||||
- **Env-var tokens** — configured via `GATEWAY_AUTH_TOKEN` (single-user) at startup
|
||||
- **DB-backed tokens** — created via `POST /api/tokens` or `POST /api/admin/users`
|
||||
|
||||
DB tokens are SHA-256 hashed at rest; plaintext is returned exactly once at creation time.
|
||||
|
||||
Auth is cached in a bounded LRU (1024 entries, 60s TTL). Suspending a user or revoking a token may take up to 60s to take effect.
|
||||
|
||||
## Roles
|
||||
|
||||
| Role | Scope |
|
||||
|------|-------|
|
||||
| `admin` | Full access to all endpoints |
|
||||
| `member` | Self-service profile + own token management only |
|
||||
|
||||
Endpoints marked **Admin** return `403 Forbidden` for `member` role.
|
||||
|
||||
---
|
||||
|
||||
## Admin: Users
|
||||
|
||||
### POST /api/admin/users
|
||||
|
||||
Create a new user. Returns the user record and a one-time plaintext API token.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"display_name": "Alice Smith",
|
||||
"email": "[email protected]",
|
||||
"role": "member"
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Default | Notes |
|
||||
|-------|------|----------|---------|-------|
|
||||
| `display_name` | string | yes | | |
|
||||
| `email` | string | no | `null` | Must be unique if provided |
|
||||
| `role` | string | no | `"member"` | `"admin"` or `"member"` |
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"email": "[email protected]",
|
||||
"display_name": "Alice Smith",
|
||||
"status": "active",
|
||||
"role": "member",
|
||||
"token": "a1b2c3d4e5f6...64-char hex...",
|
||||
"created_at": "2026-03-25T12:00:00+00:00",
|
||||
"created_by": "admin-user-id"
|
||||
}
|
||||
```
|
||||
|
||||
The `token` field is the plaintext API token. It is shown **only once** — store it securely.
|
||||
|
||||
**Errors:** `400` (missing display_name, invalid role), `403` (not admin), `503` (no database)
|
||||
|
||||
---
|
||||
|
||||
### GET /api/admin/users
|
||||
|
||||
List all users.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"users": [
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"email": "[email protected]",
|
||||
"display_name": "Alice Smith",
|
||||
"status": "active",
|
||||
"role": "member",
|
||||
"created_at": "2026-03-25T12:00:00+00:00",
|
||||
"updated_at": "2026-03-25T12:00:00+00:00",
|
||||
"last_login_at": "2026-03-25T14:30:00+00:00",
|
||||
"created_by": "admin-user-id"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### GET /api/admin/users/{id}
|
||||
|
||||
Get a single user by ID.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"email": "[email protected]",
|
||||
"display_name": "Alice Smith",
|
||||
"status": "active",
|
||||
"role": "member",
|
||||
"created_at": "2026-03-25T12:00:00+00:00",
|
||||
"updated_at": "2026-03-25T12:00:00+00:00",
|
||||
"last_login_at": "2026-03-25T14:30:00+00:00",
|
||||
"created_by": "admin-user-id",
|
||||
"metadata": {}
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `404` (user not found), `403` (not admin)
|
||||
|
||||
---
|
||||
|
||||
### PATCH /api/admin/users/{id}
|
||||
|
||||
Update a user's display name and/or metadata. Omitted fields are left unchanged.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"display_name": "Alice Johnson",
|
||||
"metadata": {"department": "engineering"}
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Notes |
|
||||
|-------|------|----------|-------|
|
||||
| `display_name` | string | no | |
|
||||
| `metadata` | object | no | Replaces entire metadata object (merge patch) |
|
||||
|
||||
**Response:** `200 OK` — returns the full updated user record (same shape as GET detail, without `last_login_at`/`created_by`).
|
||||
|
||||
**Errors:** `404` (user not found), `403` (not admin)
|
||||
|
||||
---
|
||||
|
||||
### POST /api/admin/users/{id}/suspend
|
||||
|
||||
Suspend a user. Suspended users cannot authenticate (DB auth checks user status).
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"status": "suspended"
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `404` (user not found), `403` (not admin)
|
||||
|
||||
---
|
||||
|
||||
### POST /api/admin/users/{id}/activate
|
||||
|
||||
Re-activate a suspended user.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"status": "active"
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `404` (user not found), `403` (not admin)
|
||||
|
||||
---
|
||||
|
||||
### DELETE /api/admin/users/{id}
|
||||
|
||||
Permanently delete a user and all associated data (tokens, jobs, conversations, memory, routines, settings, secrets).
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"deleted": true
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `404` (user not found), `403` (not admin)
|
||||
|
||||
**Cascade:** Deletes from `api_tokens`, `agent_jobs`, `conversations`, `memory_documents`, `routines`, `secrets`, `settings`, `wasm_tools`, and related tables. On PostgreSQL this uses FK cascades; on libSQL it uses explicit deletes.
|
||||
|
||||
---
|
||||
|
||||
## Admin: Per-User Secrets
|
||||
|
||||
Provision secrets on behalf of individual users. The primary use case is an application backend (acting as admin) that configures per-user credentials so each user's IronClaw agent can call back to external services.
|
||||
|
||||
Secrets are encrypted at rest with AES-256-GCM using a per-secret HKDF-derived key. Plaintext values are **never returned** by any endpoint — they can only be used by the agent's tool system at runtime.
|
||||
|
||||
### PUT /api/admin/users/{user_id}/secrets/{name}
|
||||
|
||||
Create or update a secret for the specified user. If a secret with the same name already exists, it is overwritten.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Path parameters:**
|
||||
|
||||
| Param | Type | Notes |
|
||||
|-------|------|-------|
|
||||
| `user_id` | string | The user's ID |
|
||||
| `name` | string | Secret name (normalized to lowercase) |
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"value": "sk-live-abc123...",
|
||||
"provider": "my-app-backend",
|
||||
"expires_in_days": 90
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Notes |
|
||||
|-------|------|----------|-------|
|
||||
| `value` | string | yes | The secret value (encrypted at rest, never returned) |
|
||||
| `provider` | string | no | Tag for grouping (e.g. `"stripe"`, `"my-app"`) |
|
||||
| `expires_in_days` | integer | no | Auto-expire after N days; `null` = never |
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"user_id": "550e8400-...",
|
||||
"name": "my_app_callback_token",
|
||||
"status": "created"
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `400` (missing value), `403` (not admin), `503` (secrets store not available)
|
||||
|
||||
**Example — application backend provisioning a callback token:**
|
||||
|
||||
```bash
|
||||
# Admin creates a user
|
||||
curl -X POST https://ironclaw.example.com/api/admin/users \
|
||||
-H "Authorization: Bearer $ADMIN_TOKEN" \
|
||||
-d '{"display_name": "Alice", "role": "member"}'
|
||||
# Response includes: {"id": "alice-uuid", "token": "alice-bearer-token", ...}
|
||||
|
||||
# Admin provisions a per-user callback secret
|
||||
curl -X PUT https://ironclaw.example.com/api/admin/users/alice-uuid/secrets/app_callback_token \
|
||||
-H "Authorization: Bearer $ADMIN_TOKEN" \
|
||||
-d '{"value": "per-user-jwt-for-alice", "provider": "my-app"}'
|
||||
|
||||
# Now Alice's IronClaw agent can use the "app_callback_token" secret
|
||||
# when calling tools that need to authenticate back to the app backend.
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### GET /api/admin/users/{user_id}/secrets
|
||||
|
||||
List a user's secrets. Returns names and providers only — **never values or hashes**.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"user_id": "550e8400-...",
|
||||
"secrets": [
|
||||
{"name": "app_callback_token", "provider": "my-app"},
|
||||
{"name": "openai_api_key", "provider": "openai"}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### DELETE /api/admin/users/{user_id}/secrets/{name}
|
||||
|
||||
Delete a specific secret for a user.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"user_id": "550e8400-...",
|
||||
"name": "app_callback_token",
|
||||
"deleted": true
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `404` (secret not found), `403` (not admin), `503` (secrets store not available)
|
||||
|
||||
---
|
||||
|
||||
## Admin: Usage
|
||||
|
||||
### GET /api/admin/usage
|
||||
|
||||
Per-user LLM usage statistics aggregated from `llm_calls` via `agent_jobs.user_id`.
|
||||
|
||||
**Auth:** Admin
|
||||
|
||||
**Query parameters:**
|
||||
|
||||
| Param | Type | Default | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| `user_id` | string | all users | Filter to a single user |
|
||||
| `period` | string | `"day"` | `"day"` (24h), `"week"` (7d), or `"month"` (30d) |
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"period": "week",
|
||||
"since": "2026-03-18T12:00:00+00:00",
|
||||
"usage": [
|
||||
{
|
||||
"user_id": "alice-id",
|
||||
"model": "claude-sonnet-4-5-20250514",
|
||||
"call_count": 42,
|
||||
"input_tokens": 150000,
|
||||
"output_tokens": 30000,
|
||||
"total_cost": "1.23"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Self-Service: Profile
|
||||
|
||||
### GET /api/profile
|
||||
|
||||
Get the authenticated user's own profile.
|
||||
|
||||
**Auth:** Any authenticated user
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"email": "[email protected]",
|
||||
"display_name": "Alice Smith",
|
||||
"status": "active",
|
||||
"role": "member",
|
||||
"created_at": "2026-03-25T12:00:00+00:00",
|
||||
"last_login_at": "2026-03-25T14:30:00+00:00"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### PATCH /api/profile
|
||||
|
||||
Update the authenticated user's own display name and/or metadata.
|
||||
|
||||
**Auth:** Any authenticated user
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"display_name": "Alice Johnson",
|
||||
"metadata": {"theme": "dark"}
|
||||
}
|
||||
```
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "550e8400-...",
|
||||
"display_name": "Alice Johnson",
|
||||
"updated": true
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Self-Service: Tokens
|
||||
|
||||
### POST /api/tokens
|
||||
|
||||
Create a new API token for the authenticated user. Admins can optionally create tokens for other users by including `user_id`.
|
||||
|
||||
**Auth:** Any authenticated user
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "CI pipeline",
|
||||
"expires_in_days": 90,
|
||||
"user_id": "other-user-id"
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Notes |
|
||||
|-------|------|----------|-------|
|
||||
| `name` | string | yes | Human-readable label |
|
||||
| `expires_in_days` | integer | no | `null` = never expires |
|
||||
| `user_id` | string | no | Admin-only; create token for another user |
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"token": "a1b2c3d4...64-char hex...",
|
||||
"id": "token-uuid",
|
||||
"name": "CI pipeline",
|
||||
"token_prefix": "a1b2c3d4",
|
||||
"expires_at": "2026-06-23T12:00:00+00:00",
|
||||
"created_at": "2026-03-25T12:00:00+00:00"
|
||||
}
|
||||
```
|
||||
|
||||
The `token` field is shown **only once**.
|
||||
|
||||
---
|
||||
|
||||
### GET /api/tokens
|
||||
|
||||
List the authenticated user's tokens. Token hashes are never returned.
|
||||
|
||||
**Auth:** Any authenticated user
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"tokens": [
|
||||
{
|
||||
"id": "token-uuid",
|
||||
"name": "CI pipeline",
|
||||
"token_prefix": "a1b2c3d4",
|
||||
"expires_at": "2026-06-23T12:00:00+00:00",
|
||||
"last_used_at": "2026-03-25T14:00:00+00:00",
|
||||
"created_at": "2026-03-25T12:00:00+00:00",
|
||||
"revoked_at": null
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### DELETE /api/tokens/{id}
|
||||
|
||||
Revoke one of the authenticated user's tokens. Users can only revoke their own tokens.
|
||||
|
||||
**Auth:** Any authenticated user
|
||||
|
||||
**Path:** `id` — UUID of the token to revoke
|
||||
|
||||
**Response:** `200 OK`
|
||||
|
||||
```json
|
||||
{
|
||||
"status": "revoked",
|
||||
"id": "token-uuid"
|
||||
}
|
||||
```
|
||||
|
||||
**Errors:** `400` (invalid UUID), `404` (token not found or belongs to another user)
|
||||
|
||||
---
|
||||
|
||||
## Error Format
|
||||
|
||||
All error responses return a plain text body with the error message and the corresponding HTTP status code:
|
||||
|
||||
| Code | Meaning |
|
||||
|------|---------|
|
||||
| `400` | Bad request (missing fields, invalid input) |
|
||||
| `401` | Missing or invalid bearer token |
|
||||
| `403` | Authenticated but insufficient role (member accessing admin endpoint) |
|
||||
| `404` | Resource not found |
|
||||
| `503` | Database or secrets store not available |
|
||||
| `500` | Internal server error |
|
||||
|
||||
---
|
||||
|
||||
## Security Model
|
||||
|
||||
### Secrets Encryption
|
||||
|
||||
- **Algorithm:** AES-256-GCM with per-secret HKDF-SHA256 derived keys
|
||||
- **Master key:** 32+ bytes, resolved from `SECRETS_MASTER_KEY` env var or OS keychain
|
||||
- **Storage format:** `nonce (12B) || ciphertext || tag (16B)` in `encrypted_value` column
|
||||
- **Per-secret salt:** 32 random bytes stored alongside the ciphertext
|
||||
- **Zero-exposure:** Plaintext never appears in logs, debug output, API responses, or LLM conversations
|
||||
|
||||
### Auth Cache
|
||||
|
||||
- Bounded LRU cache (1024 entries max)
|
||||
- 60-second TTL per entry
|
||||
- Suspending a user or revoking a token takes up to 60s to propagate
|
||||
|
||||
---
|
||||
|
||||
## Database Schema
|
||||
|
||||
### users
|
||||
|
||||
| Column | Type (PG / libSQL) | Notes |
|
||||
|--------|--------------------|-------|
|
||||
| `id` | `UUID` / `TEXT` | Primary key, UUID v4 |
|
||||
| `email` | `TEXT UNIQUE` | Nullable |
|
||||
| `display_name` | `TEXT NOT NULL` | |
|
||||
| `status` | `TEXT NOT NULL` | `"active"` or `"suspended"` |
|
||||
| `role` | `TEXT NOT NULL` | `"admin"` or `"member"` |
|
||||
| `created_at` | `TIMESTAMPTZ` / `TEXT` | |
|
||||
| `updated_at` | `TIMESTAMPTZ` / `TEXT` | |
|
||||
| `last_login_at` | `TIMESTAMPTZ` / `TEXT` | Nullable |
|
||||
| `created_by` | `TEXT` | Nullable, references `users.id` |
|
||||
| `metadata` | `JSONB` / `TEXT` | Default `{}` |
|
||||
|
||||
### api_tokens
|
||||
|
||||
| Column | Type (PG / libSQL) | Notes |
|
||||
|--------|--------------------|-------|
|
||||
| `id` | `UUID` / `TEXT` | Primary key |
|
||||
| `user_id` | `TEXT NOT NULL` | FK to `users.id` (PG cascades; libSQL explicit cleanup) |
|
||||
| `token_hash` | `BYTEA` / `BLOB` | SHA-256 of hex-encoded plaintext |
|
||||
| `token_prefix` | `TEXT NOT NULL` | First 8 chars for identification |
|
||||
| `name` | `TEXT NOT NULL` | Human-readable label |
|
||||
| `expires_at` | `TIMESTAMPTZ` / `TEXT` | Nullable |
|
||||
| `last_used_at` | `TIMESTAMPTZ` / `TEXT` | Nullable |
|
||||
| `created_at` | `TIMESTAMPTZ` / `TEXT` | |
|
||||
| `revoked_at` | `TIMESTAMPTZ` / `TEXT` | Nullable; set on revocation |
|
||||
|
||||
### secrets
|
||||
|
||||
| Column | Type (PG / libSQL) | Notes |
|
||||
|--------|--------------------|-------|
|
||||
| `id` | `UUID` / `TEXT` | Primary key |
|
||||
| `user_id` | `TEXT NOT NULL` | Scoped to user |
|
||||
| `name` | `TEXT NOT NULL` | Unique per user (lowercase normalized) |
|
||||
| `encrypted_value` | `BYTEA` / `BLOB` | AES-256-GCM (nonce + ciphertext + tag) |
|
||||
| `key_salt` | `BYTEA` / `BLOB` | Per-secret HKDF salt |
|
||||
| `provider` | `TEXT` | Optional grouping tag |
|
||||
| `expires_at` | `TIMESTAMPTZ` / `TEXT` | Nullable |
|
||||
| `last_used_at` | `TIMESTAMPTZ` / `TEXT` | Audit: last injection time |
|
||||
| `usage_count` | `BIGINT` / `INTEGER` | Audit: total injections |
|
||||
| `created_at` | `TIMESTAMPTZ` / `TEXT` | |
|
||||
| `updated_at` | `TIMESTAMPTZ` / `TEXT` | |
|
||||
@@ -0,0 +1,31 @@
|
||||
-- User management tables for multi-tenant deployments.
|
||||
--
|
||||
-- Replaces the static GATEWAY_USER_TOKENS env var with DB-backed
|
||||
-- user registration, API token management, and invitation flow.
|
||||
|
||||
CREATE TABLE users (
|
||||
id TEXT PRIMARY KEY, -- matches existing user_id pattern (string, not UUID)
|
||||
email TEXT UNIQUE, -- nullable for token-only users
|
||||
display_name TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active', -- active | suspended | deactivated
|
||||
role TEXT NOT NULL DEFAULT 'member', -- admin | member
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
last_login_at TIMESTAMPTZ,
|
||||
created_by TEXT REFERENCES users(id), -- who invited this user (nullable for bootstrap)
|
||||
metadata JSONB NOT NULL DEFAULT '{}' -- extensible profile data
|
||||
);
|
||||
|
||||
CREATE TABLE api_tokens (
|
||||
id UUID PRIMARY KEY,
|
||||
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
token_hash BYTEA NOT NULL, -- SHA-256 hash (never store plaintext)
|
||||
token_prefix TEXT NOT NULL, -- first 8 hex chars for display
|
||||
name TEXT NOT NULL, -- human label ("my-laptop", "ci-bot")
|
||||
expires_at TIMESTAMPTZ, -- nullable = never expires
|
||||
last_used_at TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
revoked_at TIMESTAMPTZ -- soft-revoke: set this instead of deleting
|
||||
);
|
||||
CREATE INDEX idx_api_tokens_user ON api_tokens(user_id);
|
||||
CREATE INDEX idx_api_tokens_hash ON api_tokens(token_hash);
|
||||
@@ -19,8 +19,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"sha256": "52def36121a93cf0b06dd6d594dc9528745fae56ded7d632998700cc595be762",
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.21.0/channel-feishu-0.1.2-wasm32-wasip2.tar.gz"
|
||||
"sha256": "a66ff0dafb67d2216d8161bb7e96e724a94acb0ab993b85d2782d30412f8fe94",
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/channel-feishu-0.1.3-wasm32-wasip2.tar.gz"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -19,8 +19,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-github-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "92c530b3ad172e2372d819744b5233f1d8f65768e26eb5a6c213eba3ce1de758"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-github-0.2.2-wasm32-wasip2.tar.gz",
|
||||
"sha256": "70b55af593193d8fa495c0f702ea23284d83a624124f8a5f7564916ec5032c3f"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/gmail-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "ee9574e02e92bc1d481f1310eb88afd99ee52bf6971074ab33bd76bf99b34b1d"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-gmail-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "79025b40ee70ce1120acc4320bae50da095d7afb0ef67bd56d99b064b72ea779"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-calendar-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "2fa47150ea222e787c122182ad6f4dfa2ffaf5fe490d05e8de887a76445f8d2d"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-calendar-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "86bcc075010b08f5ab2f98f504cec1c6c9e0ca144857d185cbecf72a11f504bf"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-docs-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "40e134a1c1564f832ca861c3396895d4e33ec67b99313fc1f97baf8d971423a9"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-docs-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "39d476029764949498a53a6a223f9952b5f4df151be7b8b19bf3fe4d401a57cd"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-drive-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "002a341a1d58125563a7c69561b26fbc2629b04ea723cade744102bdc0fbb71f"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-drive-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "6e9a700fab93865c852af718666af64c5b534ad6a419fb4b736e07740188f494"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-sheets-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "8aa2c9d52f033edea3a6c2311b0ec694ccb6d0a54ef07e94d72bf8be1ce8009a"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-sheets-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "1f8c381799a916be83263cac9d497d52946e21b1b588592a3a42ca94a73b7051"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -17,8 +17,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-slides-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "e931a97d4fd0b0b938e464dc7c7f2be6ea6b4d1508f5ea3cd931d44db23f05f5"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-google-slides-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "e2528be5da02f1b8cfc8ee9b0cdd849516c53d412e2f75c6175b3bded7f512cb"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -21,8 +21,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-llm-context-0.1.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "d9ced2b1226b879135891e0ee40e072c7c95412e1b2462925a23853e1f92497e"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-llm-context-0.1.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "9b19e2fd05dbbbe3c8bd55309a91db09124e8415eb0f767828b6e10b55771e63"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -17,8 +17,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-slack-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "ccfb0415d7a04f9497726c712d15216de36e86f498b849101283c017f5ab4efb"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-slack-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "927519e5b7734beeb022d3b8bbd152e0e6b9f67c9452a8ad47809d3c4221a137"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-telegram-0.2.0-wasm32-wasip2.tar.gz",
|
||||
"sha256": "c17065ca41fae5f2a7c43b36144686718cd310a2f22442313bb1aa82bbad0ae4"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-telegram-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "1e57d0755fc9c7b3ec013d079f30168898b484a6919f9edd105f0cd80131c1cd"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -18,8 +18,8 @@
|
||||
},
|
||||
"artifacts": {
|
||||
"wasm32-wasip2": {
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/v0.19.0/tool-web-search-0.2.1-wasm32-wasip2.tar.gz",
|
||||
"sha256": "bad275ca4ec314adea5241d6b92c44ccf9cebcbca8e30ba2493cc0bcb4b57218"
|
||||
"url": "https://github.com/nearai/ironclaw/releases/download/ironclaw-v0.22.0/tool-web-search-0.2.2-wasm32-wasip2.tar.gz",
|
||||
"sha256": "47382b50c1ea7525b20d59dc02fab04e336d018665826c2f24710bdf460779ae"
|
||||
}
|
||||
},
|
||||
"auth_summary": {
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
---
|
||||
name: abound-remittance
|
||||
version: 0.1.0
|
||||
description: Smart remittance assistant for Abound — helps users send money to India with intelligent forex timing and transfer management.
|
||||
activation:
|
||||
keywords:
|
||||
- send money
|
||||
- transfer
|
||||
- remittance
|
||||
- exchange rate
|
||||
- forex
|
||||
- INR
|
||||
- India
|
||||
- wire
|
||||
- schedule trade
|
||||
- trade tomorrow
|
||||
- convert currency
|
||||
- send dollars
|
||||
- rupees
|
||||
- beneficiary
|
||||
- funding source
|
||||
- payment
|
||||
- how much
|
||||
- rate today
|
||||
- best time
|
||||
- family maintenance
|
||||
patterns:
|
||||
- "send \\$?\\d+"
|
||||
- "schedule.*(trade|transfer|send|wire)"
|
||||
- "how much.*(INR|rupees|India)"
|
||||
- "best time to (send|transfer|convert)"
|
||||
- "(rate|forex).*(good|bad|high|low|today|now)"
|
||||
- "transfer.*tomorrow|tomorrow.*transfer"
|
||||
tags:
|
||||
- fintech
|
||||
- remittance
|
||||
- forex
|
||||
max_context_tokens: 2500
|
||||
---
|
||||
|
||||
# Abound Remittance Assistant
|
||||
|
||||
You are a smart remittance assistant for Abound, helping users send money from USD to INR (India) with intelligent timing advice.
|
||||
|
||||
## Available Tools
|
||||
|
||||
You have these Abound-specific tools:
|
||||
|
||||
- **abound_get_account_info** — Get the user's account: limits, recipients, funding sources, payment reasons
|
||||
- **abound_get_exchange_rate** — Get current USD/INR exchange rate (current + effective after fees)
|
||||
- **abound_get_forex_score** — Get a 0-100 forex timing score with a signal (convert_now / split_transfer / wait)
|
||||
- **abound_send_wire** — Execute a wire transfer (requires: funding_source_id, beneficiary_ref_id, amount, payment_reason_key)
|
||||
- **abound_create_notification** — Send a notification to the user's Abound app
|
||||
|
||||
You also have **routine_create** for scheduling future/recurring transfers.
|
||||
|
||||
## Workflow: "Send $X" or "Transfer money"
|
||||
|
||||
1. **Always check the rate first.** Call `abound_get_exchange_rate` to get the current rate.
|
||||
2. **Check the forex score.** Call `abound_get_forex_score` to assess timing.
|
||||
3. **Get account info.** Call `abound_get_account_info` to know the user's limits, recipients, and funding sources.
|
||||
4. **Advise based on the score:**
|
||||
- **Score >= 60 (convert_now):** Tell the user it's a good time. Show the rate, the INR equivalent of their amount, and recommend proceeding.
|
||||
- **Score 40-59 (split_transfer):** Suggest splitting — send half now at the current rate, schedule the rest for later when the rate may improve.
|
||||
- **Score < 40 (wait):** Unless the transfer is urgent, recommend waiting. Explain why (rate below average, unfavorable season).
|
||||
5. **Execute if user confirms.** Use `abound_send_wire` with the correct funding_source_id, beneficiary_ref_id, amount, and payment_reason_key from the account info.
|
||||
6. **Notify.** After a successful wire, call `abound_create_notification` with relevant metadata.
|
||||
|
||||
## Workflow: "Schedule a trade" or "Send tomorrow morning"
|
||||
|
||||
1. Gather the same info (rate, score, account).
|
||||
2. Use **routine_create** to schedule the transfer:
|
||||
- For "tomorrow morning": use cron `"0 9 * * *"` with the user's timezone, set to fire once
|
||||
- For "every week": use cron `"0 9 * * MON"` (or the user's preferred day)
|
||||
- The routine prompt should instruct the agent to check the rate and execute the wire
|
||||
3. Confirm the schedule with the user, showing when it will fire.
|
||||
|
||||
## Presentation Rules
|
||||
|
||||
- Always show amounts in **both USD and INR**: "$1,000 (~INR 85,420 at today's rate of 85.42)"
|
||||
- Show the **effective rate** (after fees), not just the market rate
|
||||
- When showing the forex score, explain it simply: "The forex timing score is 72/100 — this is a good time to send."
|
||||
- If the user's amount exceeds their limit ($5,000), tell them and suggest splitting into multiple transfers
|
||||
- Always mention the **estimated delivery time** (1-3 business days) after a wire
|
||||
|
||||
## Payment Reasons
|
||||
|
||||
When asking about the purpose, offer these options:
|
||||
- Family Maintenance
|
||||
- Gift
|
||||
- Education Support
|
||||
- Medical Support
|
||||
|
||||
If the user doesn't specify, ask which applies.
|
||||
+136
-37
@@ -13,7 +13,7 @@ use futures::StreamExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::context_monitor::ContextMonitor;
|
||||
use crate::agent::heartbeat::spawn_heartbeat;
|
||||
use crate::agent::heartbeat::{spawn_heartbeat, spawn_multi_user_heartbeat};
|
||||
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
|
||||
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
|
||||
use crate::agent::session::ThreadState;
|
||||
@@ -182,6 +182,8 @@ pub struct AgentDeps {
|
||||
/// Resolved LLM backend identifier (e.g., "nearai", "openai", "groq").
|
||||
/// Used by `/model` persistence to determine which env var to update.
|
||||
pub llm_backend: String,
|
||||
/// Per-tenant rate limiting registry (lazily creates rate state per user).
|
||||
pub tenant_rates: Arc<crate::tenant::TenantRateRegistry>,
|
||||
}
|
||||
|
||||
/// The main agent that coordinates all components.
|
||||
@@ -244,7 +246,10 @@ impl Agent {
|
||||
SchedulerDeps {
|
||||
tools: deps.tools.clone(),
|
||||
extension_manager: deps.extension_manager.clone(),
|
||||
store: deps.store.clone(),
|
||||
store: deps
|
||||
.store
|
||||
.as_ref()
|
||||
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db))),
|
||||
hooks: deps.hooks.clone(),
|
||||
},
|
||||
);
|
||||
@@ -325,6 +330,50 @@ impl Agent {
|
||||
&self.deps.cost_guard
|
||||
}
|
||||
|
||||
/// Build a tenant-scoped execution context for the given user.
|
||||
///
|
||||
/// This is the standard entry point for per-user operations. The returned
|
||||
/// [`TenantCtx`] provides a [`TenantScope`] that auto-binds `user_id` on
|
||||
/// every database operation and a per-user rate limiter.
|
||||
pub(super) async fn tenant_ctx(&self, user_id: &str) -> crate::tenant::TenantCtx {
|
||||
let rate = self.deps.tenant_rates.get_or_create(user_id).await;
|
||||
|
||||
let store = self
|
||||
.deps
|
||||
.store
|
||||
.as_ref()
|
||||
.map(|db| crate::tenant::TenantScope::new(user_id, Arc::clone(db)));
|
||||
|
||||
// Reuse the owner workspace if user matches, otherwise create per-user.
|
||||
let workspace = match &self.deps.workspace {
|
||||
Some(ws) if ws.user_id() == user_id => Some(Arc::clone(ws)),
|
||||
_ => self
|
||||
.deps
|
||||
.store
|
||||
.as_ref()
|
||||
.map(|db| Arc::new(Workspace::new_with_db(user_id, Arc::clone(db)))),
|
||||
};
|
||||
|
||||
crate::tenant::TenantCtx::new(
|
||||
user_id,
|
||||
store,
|
||||
workspace,
|
||||
Arc::clone(&self.deps.cost_guard),
|
||||
rate,
|
||||
)
|
||||
}
|
||||
|
||||
/// Get an admin-scoped database accessor for cross-tenant operations.
|
||||
///
|
||||
/// Only for system-level components (heartbeat, routine engine, self-repair,
|
||||
/// scheduler). Handler code should use [`tenant_ctx()`](Self::tenant_ctx) instead.
|
||||
pub(super) fn admin_store(&self) -> Option<crate::tenant::AdminScope> {
|
||||
self.deps
|
||||
.store
|
||||
.as_ref()
|
||||
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db)))
|
||||
}
|
||||
|
||||
pub(super) fn skill_registry(&self) -> Option<&Arc<std::sync::RwLock<SkillRegistry>>> {
|
||||
self.deps.skill_registry.as_ref()
|
||||
}
|
||||
@@ -410,8 +459,8 @@ impl Agent {
|
||||
self.config.stuck_threshold,
|
||||
self.config.max_repair_attempts,
|
||||
);
|
||||
if let Some(ref store) = self.deps.store {
|
||||
self_repair = self_repair.with_store(Arc::clone(store));
|
||||
if let Some(admin) = self.admin_store() {
|
||||
self_repair = self_repair.with_store(admin);
|
||||
}
|
||||
if let Some(ref builder) = self.deps.builder {
|
||||
self_repair = self_repair.with_builder(Arc::clone(builder), Arc::clone(self.tools()));
|
||||
@@ -518,6 +567,7 @@ impl Agent {
|
||||
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
|
||||
config.quiet_hours_start = hb_config.quiet_hours_start;
|
||||
config.quiet_hours_end = hb_config.quiet_hours_end;
|
||||
config.multi_tenant = hb_config.multi_tenant;
|
||||
config.timezone = hb_config
|
||||
.timezone
|
||||
.clone()
|
||||
@@ -547,30 +597,52 @@ impl Agent {
|
||||
.await;
|
||||
let notify_user = heartbeat_notify_user;
|
||||
let channels = self.channels.clone();
|
||||
let is_multi_tenant = hb_config.multi_tenant;
|
||||
tokio::spawn(async move {
|
||||
while let Some(response) = notify_rx.recv().await {
|
||||
// In multi-tenant mode, extract the owning user_id from
|
||||
// the response metadata so notifications reach the
|
||||
// correct user rather than the agent's owner.
|
||||
// This intentionally overrides the configured notify_target
|
||||
// because each user's heartbeat should notify that user.
|
||||
let effective_user = if is_multi_tenant {
|
||||
response
|
||||
.metadata
|
||||
.get("owner_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Try the configured channel first, fall back to
|
||||
// broadcasting on all channels.
|
||||
let targeted_ok = if let Some(ref channel) = notify_channel
|
||||
&& let Some(ref user) = notify_target
|
||||
{
|
||||
channels
|
||||
.broadcast(channel, user, response.clone())
|
||||
.await
|
||||
.is_ok()
|
||||
let targeted_ok = if let Some(ref channel) = notify_channel {
|
||||
let target = effective_user.as_deref().or(notify_target.as_deref());
|
||||
if let Some(user) = target {
|
||||
channels
|
||||
.broadcast(channel, user, response.clone())
|
||||
.await
|
||||
.is_ok()
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if !targeted_ok && let Some(ref user) = notify_user {
|
||||
let results = channels.broadcast_all(user, response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to broadcast heartbeat to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
if !targeted_ok {
|
||||
let fallback = effective_user.as_deref().or(notify_user.as_deref());
|
||||
if let Some(user) = fallback {
|
||||
let results = channels.broadcast_all(user, response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to broadcast heartbeat to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -583,14 +655,29 @@ impl Agent {
|
||||
.map(|h| h.to_workspace_config())
|
||||
.unwrap_or_default();
|
||||
|
||||
Some(spawn_heartbeat(
|
||||
config,
|
||||
hygiene,
|
||||
workspace.clone(),
|
||||
self.cheap_llm().clone(),
|
||||
Some(notify_tx),
|
||||
self.store().map(Arc::clone),
|
||||
))
|
||||
if config.multi_tenant {
|
||||
if let Some(admin) = self.admin_store() {
|
||||
Some(spawn_multi_user_heartbeat(
|
||||
config,
|
||||
hygiene,
|
||||
self.cheap_llm().clone(),
|
||||
Some(notify_tx),
|
||||
admin,
|
||||
))
|
||||
} else {
|
||||
tracing::warn!("Multi-tenant heartbeat requires a database store");
|
||||
None
|
||||
}
|
||||
} else {
|
||||
Some(spawn_heartbeat(
|
||||
config,
|
||||
hygiene,
|
||||
workspace.clone(),
|
||||
self.cheap_llm().clone(),
|
||||
Some(notify_tx),
|
||||
self.admin_store(),
|
||||
))
|
||||
}
|
||||
} else {
|
||||
tracing::warn!("Heartbeat enabled but no workspace available");
|
||||
None
|
||||
@@ -612,7 +699,7 @@ impl Agent {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
rt_config.clone(),
|
||||
Arc::clone(store),
|
||||
crate::tenant::AdminScope::new(Arc::clone(store)),
|
||||
self.llm().clone(),
|
||||
Arc::clone(workspace),
|
||||
notify_tx,
|
||||
@@ -1173,13 +1260,22 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// Build per-tenant execution context once; threaded through all handlers.
|
||||
let tenant = self.tenant_ctx(&message.user_id).await;
|
||||
|
||||
let session_for_empty_exit = Arc::clone(&session);
|
||||
|
||||
// Process based on submission type
|
||||
let result = match submission {
|
||||
Submission::UserInput { content } => {
|
||||
let mut result = self
|
||||
.process_user_input(message, session.clone(), thread_id, &content)
|
||||
.process_user_input(
|
||||
message,
|
||||
tenant.clone(),
|
||||
session.clone(),
|
||||
thread_id,
|
||||
&content,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Drain any messages queued during processing.
|
||||
@@ -1246,7 +1342,13 @@ impl Agent {
|
||||
let mut queued_msg = message.clone();
|
||||
queued_msg.attachments.clear();
|
||||
result = self
|
||||
.process_user_input(&queued_msg, session.clone(), thread_id, &next_content)
|
||||
.process_user_input(
|
||||
&queued_msg,
|
||||
tenant.clone(),
|
||||
session.clone(),
|
||||
thread_id,
|
||||
&next_content,
|
||||
)
|
||||
.await;
|
||||
|
||||
// If processing failed, re-queue the drained content so it
|
||||
@@ -1294,7 +1396,7 @@ impl Agent {
|
||||
};
|
||||
}
|
||||
// Authorization checks (including restart channel check) are enforced in handle_system_command
|
||||
self.handle_system_command(&command, &args, &message.channel)
|
||||
self.handle_system_command(&command, &args, &message.channel, &tenant)
|
||||
.await
|
||||
}
|
||||
Submission::Undo => self.process_undo(session, thread_id).await,
|
||||
@@ -1307,12 +1409,9 @@ impl Agent {
|
||||
Submission::Summarize => self.process_summarize(session, thread_id).await,
|
||||
Submission::Suggest => self.process_suggest(session, thread_id).await,
|
||||
Submission::JobStatus { job_id } => {
|
||||
self.process_job_status(&message.user_id, job_id.as_deref())
|
||||
.await
|
||||
}
|
||||
Submission::JobCancel { job_id } => {
|
||||
self.process_job_cancel(&message.user_id, &job_id).await
|
||||
self.process_job_status(&tenant, job_id.as_deref()).await
|
||||
}
|
||||
Submission::JobCancel { job_id } => self.process_job_cancel(&tenant, &job_id).await,
|
||||
Submission::Quit => return Ok(None),
|
||||
Submission::SwitchThread { thread_id: target } => {
|
||||
self.process_switch_thread(message, target).await
|
||||
|
||||
+125
-1
@@ -10,7 +10,7 @@ use std::borrow::Cow;
|
||||
|
||||
use crate::agent::session::PendingApproval;
|
||||
use crate::error::Error;
|
||||
use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult};
|
||||
use crate::llm::{ChatMessage, FinishReason, Reasoning, ReasoningContext, RespondResult};
|
||||
|
||||
/// Signal from the delegate indicating how the loop should proceed.
|
||||
pub enum LoopSignal {
|
||||
@@ -134,6 +134,9 @@ pub async fn run_agentic_loop(
|
||||
config: &AgenticLoopConfig,
|
||||
) -> Result<LoopOutcome, Error> {
|
||||
let mut consecutive_tool_intent_nudges: u32 = 0;
|
||||
// Accumulates across all iterations (not reset by text responses) so
|
||||
// non-consecutive truncations still escalate to force_text.
|
||||
let mut truncation_count: u32 = 0;
|
||||
|
||||
for iteration in 1..=config.max_iterations {
|
||||
// Check for external signals (stop, cancellation, user messages)
|
||||
@@ -215,7 +218,35 @@ pub async fn run_agentic_loop(
|
||||
tool_calls,
|
||||
content,
|
||||
} => {
|
||||
// If the response was truncated, tool call parameters are likely
|
||||
// incomplete. Discard them and tell the LLM to try a different
|
||||
// approach rather than executing malformed tool calls.
|
||||
if output.finish_reason == FinishReason::Length {
|
||||
truncation_count += 1;
|
||||
let names: Vec<&str> = tool_calls.iter().map(|tc| tc.name.as_str()).collect();
|
||||
tracing::warn!(
|
||||
iteration,
|
||||
tools = ?names,
|
||||
truncation_count,
|
||||
"Discarding truncated tool calls (finish_reason=Length)"
|
||||
);
|
||||
if let Some(ref text) = content {
|
||||
reason_ctx.messages.push(ChatMessage::assistant(text));
|
||||
}
|
||||
reason_ctx
|
||||
.messages
|
||||
.push(ChatMessage::user(crate::llm::TRUNCATED_TOOL_CALL_NOTICE));
|
||||
// After repeated truncations, force text-only mode so the LLM
|
||||
// stops attempting tool calls it can't fit in the output budget.
|
||||
if truncation_count >= 3 {
|
||||
reason_ctx.force_text = true;
|
||||
}
|
||||
delegate.after_iteration(iteration).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
consecutive_tool_intent_nudges = 0;
|
||||
truncation_count = 0;
|
||||
|
||||
if let Some(outcome) = delegate
|
||||
.execute_tool_calls(tool_calls, content, reason_ctx)
|
||||
@@ -271,6 +302,7 @@ mod tests {
|
||||
RespondOutput {
|
||||
result: RespondResult::Text(text.to_string()),
|
||||
usage: zero_usage(),
|
||||
finish_reason: FinishReason::Stop,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -281,6 +313,7 @@ mod tests {
|
||||
content: None,
|
||||
},
|
||||
usage: zero_usage(),
|
||||
finish_reason: FinishReason::ToolUse,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -622,4 +655,95 @@ mod tests {
|
||||
let result = truncate_for_preview("café", 4);
|
||||
assert_eq!(result, "caf...");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_truncated_tool_calls_discarded_on_length() {
|
||||
let truncated_tool_call = ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "memory_write".to_string(),
|
||||
arguments: serde_json::json!({}), // empty — truncated
|
||||
reasoning: None,
|
||||
};
|
||||
let truncated_output = RespondOutput {
|
||||
result: RespondResult::ToolCalls {
|
||||
tool_calls: vec![truncated_tool_call],
|
||||
content: Some("I'll write the report.".to_string()),
|
||||
},
|
||||
usage: zero_usage(),
|
||||
finish_reason: FinishReason::Length, // response was truncated
|
||||
};
|
||||
let delegate = MockDelegate::new(vec![truncated_output, text_output("Summarized it.")]);
|
||||
let reasoning = stub_reasoning();
|
||||
let mut ctx = ReasoningContext::new();
|
||||
let config = AgenticLoopConfig {
|
||||
max_iterations: 5,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Tool calls should NOT have been executed
|
||||
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
|
||||
// The loop should have continued and returned the text response
|
||||
assert!(matches!(outcome, LoopOutcome::Response(ref t) if t == "Summarized it."));
|
||||
// A truncation notice should have been injected into context
|
||||
assert!(
|
||||
ctx.messages
|
||||
.iter()
|
||||
.any(|m| m.role == crate::llm::Role::User && m.content.contains("truncated")),
|
||||
"Should inject truncation notice into context"
|
||||
);
|
||||
// The partial assistant content should have been preserved
|
||||
assert!(
|
||||
ctx.messages
|
||||
.iter()
|
||||
.any(|m| m.role == crate::llm::Role::Assistant
|
||||
&& m.content.contains("write the report")),
|
||||
"Should preserve partial assistant content"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_repeated_truncations_force_text_mode() {
|
||||
let make_truncated = || RespondOutput {
|
||||
result: RespondResult::ToolCalls {
|
||||
tool_calls: vec![ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "memory_write".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
reasoning: None,
|
||||
}],
|
||||
content: None,
|
||||
},
|
||||
usage: zero_usage(),
|
||||
finish_reason: FinishReason::Length,
|
||||
};
|
||||
// Three truncated responses, then a text response
|
||||
let delegate = MockDelegate::new(vec![
|
||||
make_truncated(),
|
||||
make_truncated(),
|
||||
make_truncated(),
|
||||
text_output("Gave up on tool calls."),
|
||||
]);
|
||||
let reasoning = stub_reasoning();
|
||||
let mut ctx = ReasoningContext::new();
|
||||
let config = AgenticLoopConfig {
|
||||
max_iterations: 5,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(outcome, LoopOutcome::Response(_)));
|
||||
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
|
||||
// After 3 truncations, force_text should be set
|
||||
assert!(
|
||||
ctx.force_text,
|
||||
"Should escalate to force_text after repeated truncations"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+88
-53
@@ -33,6 +33,7 @@ impl Agent {
|
||||
&self,
|
||||
intent: MessageIntent,
|
||||
message: &IncomingMessage,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
// Send thinking status for non-trivial operations
|
||||
if let MessageIntent::CreateJob { .. } = &intent {
|
||||
@@ -52,24 +53,18 @@ impl Agent {
|
||||
description,
|
||||
category,
|
||||
} => {
|
||||
self.handle_create_job(&message.user_id, title, description, category)
|
||||
self.handle_create_job(tenant, title, description, category)
|
||||
.await?
|
||||
}
|
||||
MessageIntent::CheckJobStatus { job_id } => {
|
||||
self.handle_check_status(&message.user_id, job_id).await?
|
||||
}
|
||||
MessageIntent::CancelJob { job_id } => {
|
||||
self.handle_cancel_job(&message.user_id, &job_id).await?
|
||||
}
|
||||
MessageIntent::ListJobs { filter } => {
|
||||
self.handle_list_jobs(&message.user_id, filter).await?
|
||||
}
|
||||
MessageIntent::HelpJob { job_id } => {
|
||||
self.handle_help_job(&message.user_id, &job_id).await?
|
||||
self.handle_check_status(tenant, job_id).await?
|
||||
}
|
||||
MessageIntent::CancelJob { job_id } => self.handle_cancel_job(tenant, &job_id).await?,
|
||||
MessageIntent::ListJobs { filter } => self.handle_list_jobs(tenant, filter).await?,
|
||||
MessageIntent::HelpJob { job_id } => self.handle_help_job(tenant, &job_id).await?,
|
||||
MessageIntent::Command { command, args } => {
|
||||
match self
|
||||
.handle_command(&command, &args, &message.channel)
|
||||
.handle_command(&command, &args, &message.channel, tenant)
|
||||
.await?
|
||||
{
|
||||
Some(s) => s,
|
||||
@@ -83,14 +78,14 @@ impl Agent {
|
||||
|
||||
async fn handle_create_job(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
title: String,
|
||||
description: String,
|
||||
category: Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
let job_id = self
|
||||
.scheduler
|
||||
.dispatch_job(user_id, &title, &description, None)
|
||||
.dispatch_job(tenant.user_id(), &title, &description, None)
|
||||
.await?;
|
||||
|
||||
// Set the dedicated category field (not stored in metadata)
|
||||
@@ -113,7 +108,7 @@ impl Agent {
|
||||
|
||||
async fn handle_check_status(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
job_id: Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
match job_id {
|
||||
@@ -122,7 +117,8 @@ impl Agent {
|
||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||
|
||||
// Try DB first for persistent state, fall back to ContextManager.
|
||||
if let Some(store) = self.store()
|
||||
// TenantScope.get_job() auto-filters by ownership — no manual check needed.
|
||||
if let Some(store) = tenant.store()
|
||||
&& let Ok(Some(ctx)) = store.get_job(uuid).await
|
||||
{
|
||||
return Ok(format!(
|
||||
@@ -138,7 +134,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
let ctx = self.context_manager.get_context(uuid).await?;
|
||||
if ctx.user_id != user_id {
|
||||
if ctx.user_id != tenant.user_id() {
|
||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||
}
|
||||
|
||||
@@ -155,7 +151,8 @@ impl Agent {
|
||||
}
|
||||
None => {
|
||||
// Show summary from DB for consistency with Jobs tab.
|
||||
if let Some(store) = self.store() {
|
||||
// TenantScope methods auto-scope to user — no user_id parameter needed.
|
||||
if let Some(store) = tenant.store() {
|
||||
let mut total = 0;
|
||||
let mut in_progress = 0;
|
||||
let mut completed = 0;
|
||||
@@ -183,7 +180,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Fallback to ContextManager if no DB.
|
||||
let summary = self.context_manager.summary_for(user_id).await;
|
||||
let summary = self.context_manager.summary_for(tenant.user_id()).await;
|
||||
Ok(format!(
|
||||
"Jobs summary: Total: {} In Progress: {} Completed: {} Failed: {} Stuck: {}",
|
||||
summary.total,
|
||||
@@ -196,19 +193,24 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
|
||||
async fn handle_cancel_job(
|
||||
&self,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
job_id: &str,
|
||||
) -> Result<String, Error> {
|
||||
let uuid = Uuid::parse_str(job_id)
|
||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||
|
||||
let ctx = self.context_manager.get_context(uuid).await?;
|
||||
if ctx.user_id != user_id {
|
||||
if ctx.user_id != tenant.user_id() {
|
||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||
}
|
||||
|
||||
self.scheduler.stop(uuid).await?;
|
||||
|
||||
// Also update DB so the Jobs tab reflects cancellation immediately.
|
||||
if let Some(store) = self.store()
|
||||
// Use TenantScope — ownership already verified above.
|
||||
if let Some(store) = tenant.store()
|
||||
&& let Err(e) = store
|
||||
.update_job_status(uuid, JobState::Cancelled, Some("Cancelled by user"))
|
||||
.await
|
||||
@@ -221,11 +223,12 @@ impl Agent {
|
||||
|
||||
async fn handle_list_jobs(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
_filter: Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
// List from DB for consistency with Jobs tab.
|
||||
if let Some(store) = self.store() {
|
||||
// TenantScope methods auto-scope to user.
|
||||
if let Some(store) = tenant.store() {
|
||||
let agent_jobs = match store.list_agent_jobs().await {
|
||||
Ok(jobs) => jobs,
|
||||
Err(e) => {
|
||||
@@ -256,7 +259,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Fallback to ContextManager if no DB.
|
||||
let jobs = self.context_manager.all_jobs_for(user_id).await;
|
||||
let jobs = self.context_manager.all_jobs_for(tenant.user_id()).await;
|
||||
if jobs.is_empty() {
|
||||
return Ok("No jobs found.".to_string());
|
||||
}
|
||||
@@ -270,12 +273,16 @@ impl Agent {
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
|
||||
async fn handle_help_job(
|
||||
&self,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
job_id: &str,
|
||||
) -> Result<String, Error> {
|
||||
let uuid = Uuid::parse_str(job_id)
|
||||
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
|
||||
|
||||
let ctx = self.context_manager.get_context(uuid).await?;
|
||||
if ctx.user_id != user_id {
|
||||
if ctx.user_id != tenant.user_id() {
|
||||
return Err(crate::error::JobError::NotFound { id: uuid }.into());
|
||||
}
|
||||
|
||||
@@ -308,11 +315,11 @@ impl Agent {
|
||||
/// Show job status inline — either all jobs (no id) or a specific job.
|
||||
pub(super) async fn process_job_status(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
job_id: Option<&str>,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
match self
|
||||
.handle_check_status(user_id, job_id.map(|s| s.to_string()))
|
||||
.handle_check_status(tenant, job_id.map(|s| s.to_string()))
|
||||
.await
|
||||
{
|
||||
Ok(text) => Ok(SubmissionResult::response(text)),
|
||||
@@ -323,10 +330,10 @@ impl Agent {
|
||||
/// Cancel a job by ID.
|
||||
pub(super) async fn process_job_cancel(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
job_id: &str,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
match self.handle_cancel_job(user_id, job_id).await {
|
||||
match self.handle_cancel_job(tenant, job_id).await {
|
||||
Ok(text) => Ok(SubmissionResult::response(text)),
|
||||
Err(e) => Ok(SubmissionResult::error(format!("Cancel error: {}", e))),
|
||||
}
|
||||
@@ -559,6 +566,7 @@ impl Agent {
|
||||
command: &str,
|
||||
args: &[String],
|
||||
channel: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
match command {
|
||||
"help" => Ok(SubmissionResult::response(concat!(
|
||||
@@ -752,19 +760,32 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
match self.llm().set_model(requested) {
|
||||
Ok(()) => {
|
||||
// Persist the model choice so it survives restarts.
|
||||
self.persist_selected_model(requested).await;
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Switched model to: {}",
|
||||
requested
|
||||
)))
|
||||
if self.config.multi_tenant {
|
||||
// Multi-tenant: only persist to per-user DB settings.
|
||||
// Do NOT call set_model() on the shared provider — that
|
||||
// would change the default for all users. The per-request
|
||||
// model_override in the dispatcher reads from the same
|
||||
// "selected_model" setting and applies it per-user.
|
||||
self.persist_selected_model(tenant, requested).await;
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Model preference set to: {} (per-user)",
|
||||
requested
|
||||
)))
|
||||
} else {
|
||||
match self.llm().set_model(requested) {
|
||||
Ok(()) => {
|
||||
// Persist the model choice so it survives restarts.
|
||||
self.persist_selected_model(tenant, requested).await;
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Switched model to: {}",
|
||||
requested
|
||||
)))
|
||||
}
|
||||
Err(e) => Ok(SubmissionResult::error(format!(
|
||||
"Failed to switch model: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
Err(e) => Ok(SubmissionResult::error(format!(
|
||||
"Failed to switch model: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -906,10 +927,14 @@ impl Agent {
|
||||
command: &str,
|
||||
args: &[String],
|
||||
channel: &str,
|
||||
tenant: &crate::tenant::TenantCtx,
|
||||
) -> Result<Option<String>, Error> {
|
||||
// System commands are now handled directly via Submission::SystemCommand,
|
||||
// but the router may still send us unknown /commands.
|
||||
match self.handle_system_command(command, args, channel).await? {
|
||||
match self
|
||||
.handle_system_command(command, args, channel, tenant)
|
||||
.await?
|
||||
{
|
||||
SubmissionResult::Response { content } => Ok(Some(content)),
|
||||
SubmissionResult::Ok { message } => Ok(message),
|
||||
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
||||
@@ -921,23 +946,33 @@ impl Agent {
|
||||
///
|
||||
/// Best-effort: logs warnings on failure but does not propagate errors,
|
||||
/// since the in-memory model switch already succeeded.
|
||||
async fn persist_selected_model(&self, model: &str) {
|
||||
// 1. Persist to DB if available.
|
||||
if let Some(store) = self.store() {
|
||||
///
|
||||
/// In multi-tenant mode, only the per-user DB setting is written — global
|
||||
/// .env and TOML files are shared across users and must not be mutated.
|
||||
async fn persist_selected_model(&self, tenant: &crate::tenant::TenantCtx, model: &str) {
|
||||
// 1. Persist to DB if available (per-user scoped via TenantScope).
|
||||
if let Some(store) = tenant.store() {
|
||||
let value = serde_json::Value::String(model.to_string());
|
||||
if let Err(e) = store
|
||||
.set_setting(self.owner_id(), "selected_model", &value)
|
||||
.await
|
||||
{
|
||||
if let Err(e) = store.set_setting("selected_model", &value).await {
|
||||
tracing::warn!("Failed to persist model to DB: {}", e);
|
||||
} else {
|
||||
tracing::debug!("Persisted selected_model to DB: {}", model);
|
||||
tracing::debug!(
|
||||
user_id = tenant.user_id(),
|
||||
"Persisted selected_model to DB: {}",
|
||||
model
|
||||
);
|
||||
}
|
||||
} else {
|
||||
tracing::warn!("No database store available — model choice will not persist to DB");
|
||||
}
|
||||
|
||||
// 2. Update .env and TOML config file (sync I/O in spawn_blocking).
|
||||
// 2. In multi-tenant mode, skip .env/TOML writes — these are global
|
||||
// files shared by all users. The per-user DB setting is sufficient.
|
||||
if self.config.multi_tenant {
|
||||
return;
|
||||
}
|
||||
|
||||
// 3. Update .env and TOML config file (sync I/O in spawn_blocking).
|
||||
let model_owned = model.to_string();
|
||||
let backend = self.deps.llm_backend.clone();
|
||||
if let Err(e) = tokio::task::spawn_blocking(move || {
|
||||
|
||||
+236
-3
@@ -21,6 +21,9 @@ pub struct CostGuardConfig {
|
||||
pub max_cost_per_day_cents: Option<u64>,
|
||||
/// Maximum LLM calls per hour. None = unlimited.
|
||||
pub max_actions_per_hour: Option<u64>,
|
||||
/// Maximum spend per user per day in cents. None = unlimited.
|
||||
/// Applied independently per user alongside the global budget.
|
||||
pub max_cost_per_user_per_day_cents: Option<u64>,
|
||||
}
|
||||
|
||||
/// Error returned when a cost limit is exceeded.
|
||||
@@ -30,6 +33,12 @@ pub enum CostLimitExceeded {
|
||||
DailyBudget { spent_cents: u64, limit_cents: u64 },
|
||||
/// Hourly action rate limit reached.
|
||||
HourlyRate { actions: u64, limit: u64 },
|
||||
/// Per-user daily spending cap reached.
|
||||
UserDailyBudget {
|
||||
user_id: String,
|
||||
spent_cents: u64,
|
||||
limit_cents: u64,
|
||||
},
|
||||
}
|
||||
|
||||
impl std::fmt::Display for CostLimitExceeded {
|
||||
@@ -49,6 +58,17 @@ impl std::fmt::Display for CostLimitExceeded {
|
||||
"Hourly action limit exceeded: {} actions of {} allowed per hour",
|
||||
actions, limit
|
||||
),
|
||||
Self::UserDailyBudget {
|
||||
user_id,
|
||||
spent_cents,
|
||||
limit_cents,
|
||||
} => write!(
|
||||
f,
|
||||
"User '{}' daily cost limit exceeded: spent ${:.2} of ${:.2} allowed",
|
||||
user_id,
|
||||
*spent_cents as f64 / 100.0,
|
||||
*limit_cents as f64 / 100.0
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -78,6 +98,9 @@ pub struct CostGuard {
|
||||
|
||||
/// Per-model token usage since startup.
|
||||
model_tokens: Mutex<HashMap<String, ModelTokens>>,
|
||||
|
||||
/// Per-user daily cost tracking. Each entry resets independently at midnight UTC.
|
||||
per_user_daily_cost: Mutex<HashMap<String, DailyCost>>,
|
||||
}
|
||||
|
||||
struct DailyCost {
|
||||
@@ -97,6 +120,7 @@ impl CostGuard {
|
||||
action_window: Mutex::new(VecDeque::new()),
|
||||
budget_exceeded: AtomicBool::new(false),
|
||||
model_tokens: Mutex::new(HashMap::new()),
|
||||
per_user_daily_cost: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,6 +227,11 @@ impl CostGuard {
|
||||
daily.reset_date = today;
|
||||
self.budget_exceeded.store(false, Ordering::Relaxed);
|
||||
tracing::info!("Cost guard: daily counter reset for {}", today);
|
||||
|
||||
// Prune per-user entries from previous days to prevent
|
||||
// unbounded HashMap growth in long-lived deployments.
|
||||
let mut per_user = self.per_user_daily_cost.lock().await;
|
||||
per_user.retain(|_, entry| entry.reset_date == today);
|
||||
}
|
||||
daily.total += cost;
|
||||
|
||||
@@ -248,6 +277,85 @@ impl CostGuard {
|
||||
cost
|
||||
}
|
||||
|
||||
/// Record an LLM call with per-user attribution.
|
||||
///
|
||||
/// Delegates to `record_llm_call` for global tracking, then additionally
|
||||
/// records the cost against the user's daily budget.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn record_llm_call_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
model: &str,
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
cache_read_input_tokens: u32,
|
||||
cache_creation_input_tokens: u32,
|
||||
cache_read_discount: Decimal,
|
||||
cache_write_multiplier: Decimal,
|
||||
cost_per_token: Option<(Decimal, Decimal)>,
|
||||
) -> Decimal {
|
||||
let cost = self
|
||||
.record_llm_call(
|
||||
model,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_creation_input_tokens,
|
||||
cache_read_discount,
|
||||
cache_write_multiplier,
|
||||
cost_per_token,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Track per-user daily cost
|
||||
{
|
||||
let today = chrono::Utc::now().date_naive();
|
||||
let mut per_user = self.per_user_daily_cost.lock().await;
|
||||
let entry = per_user
|
||||
.entry(user_id.to_string())
|
||||
.or_insert_with(|| DailyCost {
|
||||
total: Decimal::ZERO,
|
||||
reset_date: today,
|
||||
});
|
||||
if today != entry.reset_date {
|
||||
entry.total = Decimal::ZERO;
|
||||
entry.reset_date = today;
|
||||
}
|
||||
entry.total += cost;
|
||||
}
|
||||
|
||||
cost
|
||||
}
|
||||
|
||||
/// Check whether the next action is allowed for a specific user.
|
||||
///
|
||||
/// Checks the global limits first (via `check_allowed`), then additionally
|
||||
/// checks the per-user daily budget if configured.
|
||||
pub async fn check_allowed_for_user(&self, user_id: &str) -> Result<(), CostLimitExceeded> {
|
||||
// Check global limits first
|
||||
self.check_allowed().await?;
|
||||
|
||||
// Check per-user daily budget
|
||||
if let Some(limit_cents) = self.config.max_cost_per_user_per_day_cents {
|
||||
let today = chrono::Utc::now().date_naive();
|
||||
let per_user = self.per_user_daily_cost.lock().await;
|
||||
if let Some(entry) = per_user.get(user_id)
|
||||
&& entry.reset_date == today
|
||||
{
|
||||
let spent_cents = to_cents(entry.total);
|
||||
if spent_cents >= limit_cents {
|
||||
return Err(CostLimitExceeded::UserDailyBudget {
|
||||
user_id: user_id.to_string(),
|
||||
spent_cents,
|
||||
limit_cents,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Current daily spend in USD (as Decimal).
|
||||
pub async fn daily_spend(&self) -> Decimal {
|
||||
let daily = self.daily_cost.lock().await;
|
||||
@@ -259,6 +367,16 @@ impl CostGuard {
|
||||
}
|
||||
}
|
||||
|
||||
/// Current daily spend for a specific user in USD (as Decimal).
|
||||
pub async fn daily_spend_for_user(&self, user_id: &str) -> Decimal {
|
||||
let today = chrono::Utc::now().date_naive();
|
||||
let per_user = self.per_user_daily_cost.lock().await;
|
||||
match per_user.get(user_id) {
|
||||
Some(entry) if entry.reset_date == today => entry.total,
|
||||
_ => Decimal::ZERO,
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of actions in the current hourly window.
|
||||
pub async fn actions_this_hour(&self) -> u64 {
|
||||
let mut window = self.action_window.lock().await;
|
||||
@@ -314,7 +432,7 @@ mod tests {
|
||||
async fn test_daily_budget_enforcement() {
|
||||
let guard = CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: Some(1), // $0.01 limit
|
||||
max_actions_per_hour: None,
|
||||
..CostGuardConfig::default()
|
||||
});
|
||||
|
||||
// First call allowed
|
||||
@@ -350,8 +468,8 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_hourly_rate_enforcement() {
|
||||
let guard = CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: Some(3),
|
||||
..CostGuardConfig::default()
|
||||
});
|
||||
|
||||
// First 3 actions allowed
|
||||
@@ -633,8 +751,8 @@ mod tests {
|
||||
// A fresh CostGuard with rate limits should not panic even if
|
||||
// checked_sub returns None (simulating short uptime).
|
||||
let guard = CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: Some(100),
|
||||
..CostGuardConfig::default()
|
||||
});
|
||||
|
||||
// These must not panic regardless of system uptime
|
||||
@@ -656,4 +774,119 @@ mod tests {
|
||||
let result = Instant::now().checked_sub(std::time::Duration::MAX);
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_per_user_daily_budget_enforcement() {
|
||||
let guard = CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
|
||||
});
|
||||
|
||||
// Both users initially allowed
|
||||
assert!(guard.check_allowed_for_user("alice").await.is_ok());
|
||||
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||
|
||||
// Alice makes an expensive call
|
||||
guard
|
||||
.record_llm_call_for_user(
|
||||
"alice",
|
||||
"gpt-4o",
|
||||
10_000,
|
||||
10_000,
|
||||
0,
|
||||
0,
|
||||
Decimal::ONE,
|
||||
Decimal::ONE,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Alice should be blocked, Bob should still be allowed
|
||||
let result = guard.check_allowed_for_user("alice").await;
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
CostLimitExceeded::UserDailyBudget {
|
||||
user_id,
|
||||
limit_cents,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(user_id, "alice");
|
||||
assert_eq!(limit_cents, 1);
|
||||
}
|
||||
other => panic!("Expected UserDailyBudget, got {:?}", other),
|
||||
}
|
||||
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_per_user_daily_spend_tracking() {
|
||||
let guard = CostGuard::new(CostGuardConfig::default());
|
||||
|
||||
assert_eq!(guard.daily_spend_for_user("alice").await, Decimal::ZERO);
|
||||
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
|
||||
|
||||
let cost = guard
|
||||
.record_llm_call_for_user(
|
||||
"alice",
|
||||
"gpt-4o",
|
||||
1000,
|
||||
500,
|
||||
0,
|
||||
0,
|
||||
Decimal::ONE,
|
||||
Decimal::ONE,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(guard.daily_spend_for_user("alice").await, cost);
|
||||
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
|
||||
// Global spend should also be tracked
|
||||
assert_eq!(guard.daily_spend().await, cost);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_per_user_budget_independent_of_global() {
|
||||
let guard = CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: Some(100_000), // $1000 global limit
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
|
||||
});
|
||||
|
||||
// User hits their personal limit
|
||||
guard
|
||||
.record_llm_call_for_user(
|
||||
"alice",
|
||||
"gpt-4o",
|
||||
10_000,
|
||||
10_000,
|
||||
0,
|
||||
0,
|
||||
Decimal::ONE,
|
||||
Decimal::ONE,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Alice blocked by per-user limit, not global
|
||||
assert!(guard.check_allowed_for_user("alice").await.is_err());
|
||||
// Global limit is far from reached
|
||||
assert!(guard.check_allowed().await.is_ok());
|
||||
// Bob is unaffected
|
||||
assert!(guard.check_allowed_for_user("bob").await.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_cost_limit_display() {
|
||||
let limit = CostLimitExceeded::UserDailyBudget {
|
||||
user_id: "alice".to_string(),
|
||||
spent_cents: 150,
|
||||
limit_cents: 100,
|
||||
};
|
||||
let msg = limit.to_string();
|
||||
assert!(msg.contains("alice"));
|
||||
assert!(msg.contains("$1.50"));
|
||||
assert!(msg.contains("$1.00"));
|
||||
}
|
||||
}
|
||||
|
||||
+63
-9
@@ -42,6 +42,7 @@ impl Agent {
|
||||
pub(super) async fn run_agentic_loop(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
tenant: crate::tenant::TenantCtx,
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
initial_messages: Vec<ChatMessage>,
|
||||
@@ -168,6 +169,7 @@ impl Agent {
|
||||
|
||||
let delegate = ChatDelegate {
|
||||
agent: self,
|
||||
tenant,
|
||||
session: session.clone(),
|
||||
thread_id,
|
||||
message,
|
||||
@@ -240,6 +242,7 @@ impl Agent {
|
||||
/// auth intercept, and cost tracking.
|
||||
struct ChatDelegate<'a> {
|
||||
agent: &'a Agent,
|
||||
tenant: crate::tenant::TenantCtx,
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
message: &'a IncomingMessage,
|
||||
@@ -303,6 +306,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
|
||||
// Update context for this iteration
|
||||
reason_ctx.available_tools = tool_defs;
|
||||
// Preserve force_text if already set (e.g. by truncation escalation).
|
||||
let force_text = force_text || reason_ctx.force_text;
|
||||
reason_ctx.system_prompt = Some(if force_text {
|
||||
self.cached_prompt_no_tools.clone()
|
||||
} else {
|
||||
@@ -336,8 +341,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
reason_ctx: &mut ReasoningContext,
|
||||
iteration: usize,
|
||||
) -> Result<crate::llm::RespondOutput, Error> {
|
||||
// Enforce cost guardrails before the LLM call
|
||||
if let Err(limit) = self.agent.cost_guard().check_allowed().await {
|
||||
// Enforce cost guardrails before the LLM call (global + per-user)
|
||||
if let Err(limit) = self.tenant.check_cost_allowed().await {
|
||||
return Err(crate::error::LlmError::InvalidResponse {
|
||||
provider: "agent".to_string(),
|
||||
reason: limit.to_string(),
|
||||
@@ -345,6 +350,21 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
.into());
|
||||
}
|
||||
|
||||
// Apply per-user model override from settings (first iteration only
|
||||
// to avoid repeated DB lookups within the same agentic loop).
|
||||
// Uses "selected_model" — the same key the /model command persists to
|
||||
// via SettingsStore (per-user scoped via TenantScope).
|
||||
if iteration == 0
|
||||
&& let Some(store) = self.tenant.store()
|
||||
&& let Ok(Some(value)) = store.get_setting("selected_model").await
|
||||
&& let Some(model) = value.as_str()
|
||||
{
|
||||
let model = model.trim();
|
||||
if !model.is_empty() {
|
||||
reason_ctx.model_override = Some(model.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let output = match reasoning.respond_with_tools(reason_ctx).await {
|
||||
Ok(output) => output,
|
||||
Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => {
|
||||
@@ -379,13 +399,27 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
// Record cost and track token usage
|
||||
let model_name = self.agent.llm().active_model_name();
|
||||
// Record cost and track token usage (global + per-user).
|
||||
// Use the provider's effective_model_name so cost attribution matches
|
||||
// the model that actually served the request. When the override is
|
||||
// honoured (e.g. NearAI), this returns the override name; when the
|
||||
// provider ignores overrides (e.g. Rig-based), it returns the active
|
||||
// model, keeping attribution accurate in both cases.
|
||||
let model_name = self
|
||||
.agent
|
||||
.llm()
|
||||
.effective_model_name(reason_ctx.model_override.as_deref());
|
||||
let cost_per_token = if reason_ctx.model_override.is_some() {
|
||||
// Override may use different pricing; let CostGuard fall back to
|
||||
// costs::model_cost() for the effective model.
|
||||
None
|
||||
} else {
|
||||
Some(self.agent.llm().cost_per_token())
|
||||
};
|
||||
let read_discount = self.agent.llm().cache_read_discount();
|
||||
let write_multiplier = self.agent.llm().cache_write_multiplier();
|
||||
let call_cost = self
|
||||
.agent
|
||||
.cost_guard()
|
||||
.tenant
|
||||
.record_llm_call(
|
||||
&model_name,
|
||||
output.usage.input_tokens,
|
||||
@@ -394,7 +428,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
output.usage.cache_creation_input_tokens,
|
||||
read_discount,
|
||||
write_multiplier,
|
||||
Some(self.agent.llm().cost_per_token()),
|
||||
cost_per_token,
|
||||
)
|
||||
.await;
|
||||
tracing::debug!(
|
||||
@@ -1305,6 +1339,7 @@ mod tests {
|
||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
Agent::new(
|
||||
@@ -1320,10 +1355,15 @@ mod tests {
|
||||
allow_local_tools: false,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
max_tool_iterations: 50,
|
||||
auto_approve_tools: false,
|
||||
default_timezone: "UTC".to_string(),
|
||||
max_jobs_per_user: None,
|
||||
max_tokens_per_job: 0,
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: None,
|
||||
max_jobs_concurrent_per_user: None,
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -2181,6 +2221,7 @@ mod tests {
|
||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
Agent::new(
|
||||
@@ -2196,10 +2237,15 @@ mod tests {
|
||||
allow_local_tools: false,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
max_tool_iterations,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
max_jobs_per_user: None,
|
||||
max_tokens_per_job: 0,
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: None,
|
||||
max_jobs_concurrent_per_user: None,
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -2234,13 +2280,14 @@ mod tests {
|
||||
|
||||
let message = IncomingMessage::new("test", "test-user", "do something");
|
||||
let initial_messages = vec![ChatMessage::user("do something")];
|
||||
let tenant = agent.tenant_ctx("test-user").await;
|
||||
|
||||
// The dispatcher must terminate within 5 seconds. If there is an
|
||||
// infinite loop bug (e.g., index not advancing on tool failure), the
|
||||
// timeout will fire and the test will fail.
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
|
||||
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -2302,6 +2349,7 @@ mod tests {
|
||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
Agent::new(
|
||||
@@ -2317,10 +2365,15 @@ mod tests {
|
||||
allow_local_tools: false,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
max_tool_iterations: max_iter,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
max_jobs_per_user: None,
|
||||
max_tokens_per_job: 0,
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: None,
|
||||
max_jobs_concurrent_per_user: None,
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -2340,13 +2393,14 @@ mod tests {
|
||||
|
||||
let message = IncomingMessage::new("test", "test-user", "keep calling tools");
|
||||
let initial_messages = vec![ChatMessage::user("keep calling tools")];
|
||||
let tenant = agent.tenant_ctx("test-user").await;
|
||||
|
||||
// Even with an LLM that always wants to call tools, the dispatcher
|
||||
// must terminate within the timeout thanks to force_text at
|
||||
// max_tool_iterations.
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
|
||||
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
+185
-6
@@ -31,8 +31,8 @@ use chrono_tz::Tz;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::channels::OutgoingResponse;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||
use crate::tenant::AdminScope;
|
||||
use crate::workspace::Workspace;
|
||||
use crate::workspace::hygiene::HygieneConfig;
|
||||
|
||||
@@ -57,6 +57,9 @@ pub struct HeartbeatConfig {
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
/// When true, cycle through all users with routines instead of
|
||||
/// running heartbeat for a single user. Requires a database store.
|
||||
pub multi_tenant: bool,
|
||||
}
|
||||
|
||||
impl Default for HeartbeatConfig {
|
||||
@@ -71,6 +74,7 @@ impl Default for HeartbeatConfig {
|
||||
quiet_hours_start: None,
|
||||
quiet_hours_end: None,
|
||||
timezone: None,
|
||||
multi_tenant: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -178,7 +182,7 @@ pub struct HeartbeatRunner {
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
store: Option<AdminScope>,
|
||||
consecutive_failures: u32,
|
||||
}
|
||||
|
||||
@@ -207,8 +211,8 @@ impl HeartbeatRunner {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the database store for persistent heartbeat conversations.
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
/// Set the admin-scoped database store for persistent heartbeat conversations.
|
||||
pub fn with_store(mut self, store: AdminScope) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
@@ -493,7 +497,7 @@ pub fn spawn_heartbeat(
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
store: Option<AdminScope>,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
||||
if let Some(tx) = response_tx {
|
||||
@@ -508,6 +512,181 @@ pub fn spawn_heartbeat(
|
||||
})
|
||||
}
|
||||
|
||||
/// Spawn a multi-user heartbeat runner that cycles through all users who
|
||||
/// have routines (enabled or not). Each tick, it queries the DB for distinct
|
||||
/// user_ids, creates a per-user workspace, and runs a heartbeat check for
|
||||
/// each user concurrently. Per-user failure counts are tracked independently.
|
||||
pub fn spawn_multi_user_heartbeat(
|
||||
config: HeartbeatConfig,
|
||||
hygiene_config: HygieneConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: AdminScope,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
if !config.enabled {
|
||||
tracing::info!("Multi-user heartbeat is disabled");
|
||||
return;
|
||||
}
|
||||
|
||||
let mut tick_interval = if config.fire_at.is_none() {
|
||||
let mut iv = tokio::time::interval(config.interval);
|
||||
iv.tick().await; // skip immediate tick
|
||||
Some(iv)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Track consecutive failures per user so we can disable heartbeat
|
||||
// for persistently-failing users (same semantics as single-user mode).
|
||||
let mut user_failures: std::collections::HashMap<String, u32> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
tracing::info!("Starting multi-user heartbeat loop");
|
||||
|
||||
loop {
|
||||
if let Some(fire_at) = config.fire_at {
|
||||
let sleep_dur = duration_until_next_fire(fire_at, config.resolved_tz());
|
||||
tokio::time::sleep(sleep_dur).await;
|
||||
} else if let Some(ref mut iv) = tick_interval {
|
||||
iv.tick().await;
|
||||
}
|
||||
|
||||
if config.is_quiet_hours() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Get distinct user_ids from routines
|
||||
let user_ids = match store.list_all_routines().await {
|
||||
Ok(routines) => {
|
||||
let mut ids: Vec<String> = routines
|
||||
.iter()
|
||||
.map(|r| r.user_id.clone())
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
ids.sort();
|
||||
ids
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Multi-user heartbeat: failed to list routines: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Run user heartbeats (and hygiene) concurrently so one slow LLM
|
||||
// call doesn't block others. Cap concurrency to avoid flooding the
|
||||
// LLM provider. Hygiene runs inside the same JoinSet so it is
|
||||
// tracked and bounded by the same concurrency cap.
|
||||
const MAX_CONCURRENT_HEARTBEATS: usize = 8;
|
||||
let mut join_set = tokio::task::JoinSet::new();
|
||||
|
||||
for user_id in &user_ids {
|
||||
// Skip users that have exceeded max_failures
|
||||
let failures = user_failures.get(user_id).copied().unwrap_or(0);
|
||||
if failures >= config.max_failures {
|
||||
continue;
|
||||
}
|
||||
|
||||
let workspace = Arc::new(Workspace::new_with_db(user_id, Arc::clone(store.db())));
|
||||
|
||||
// Drain completed tasks to stay within the concurrency cap.
|
||||
while join_set.len() >= MAX_CONCURRENT_HEARTBEATS {
|
||||
if let Some(join_result) = join_set.join_next().await {
|
||||
collect_heartbeat_result(join_result, &mut user_failures, &config);
|
||||
}
|
||||
}
|
||||
|
||||
let uid = user_id.clone();
|
||||
// In multi-tenant mode, clear notify_user_id so that
|
||||
// HeartbeatRunner::send_notification falls back to
|
||||
// workspace.user_id() — each user's heartbeat should persist
|
||||
// and notify that user, not the shared config target.
|
||||
let mut cfg = config.clone();
|
||||
cfg.notify_user_id = None;
|
||||
let hyg = hygiene_config.clone();
|
||||
let llm_clone = llm.clone();
|
||||
let tx = response_tx.clone();
|
||||
let admin = store.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
// Run memory hygiene per user (same as single-user heartbeat)
|
||||
// inside the tracked task so concurrency is bounded.
|
||||
let report = crate::workspace::hygiene::run_if_due(&workspace, &hyg).await;
|
||||
if report.had_work() {
|
||||
tracing::info!(
|
||||
user_id = uid,
|
||||
daily_logs_deleted = report.daily_logs_deleted,
|
||||
conversation_docs_deleted = report.conversation_docs_deleted,
|
||||
"multi-user heartbeat: memory hygiene deleted stale documents"
|
||||
);
|
||||
}
|
||||
|
||||
let mut runner = HeartbeatRunner::new(cfg, hyg, workspace, llm_clone);
|
||||
if let Some(tx) = tx {
|
||||
runner = runner.with_response_channel(tx);
|
||||
}
|
||||
runner = runner.with_store(admin);
|
||||
|
||||
let result = runner.check_heartbeat().await;
|
||||
if let HeartbeatResult::NeedsAttention(msg) = &result {
|
||||
runner.send_notification(msg).await;
|
||||
}
|
||||
(uid, result)
|
||||
});
|
||||
}
|
||||
|
||||
// Collect remaining results and update failure counts
|
||||
while let Some(join_result) = join_set.join_next().await {
|
||||
collect_heartbeat_result(join_result, &mut user_failures, &config);
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Process a single JoinSet result from the multi-user heartbeat loop.
|
||||
fn collect_heartbeat_result(
|
||||
join_result: Result<(String, HeartbeatResult), tokio::task::JoinError>,
|
||||
user_failures: &mut std::collections::HashMap<String, u32>,
|
||||
config: &HeartbeatConfig,
|
||||
) {
|
||||
let (uid, result) = match join_result {
|
||||
Ok(pair) => pair,
|
||||
Err(e) => {
|
||||
tracing::error!("Multi-user heartbeat task panicked: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
match result {
|
||||
HeartbeatResult::Ok => {
|
||||
tracing::trace!(user_id = uid, "Multi-user heartbeat OK");
|
||||
user_failures.remove(&uid);
|
||||
}
|
||||
HeartbeatResult::NeedsAttention(_) => {
|
||||
tracing::info!(user_id = uid, "Multi-user heartbeat needs attention");
|
||||
user_failures.remove(&uid);
|
||||
}
|
||||
HeartbeatResult::Skipped => {}
|
||||
HeartbeatResult::Failed(err) => {
|
||||
let count = user_failures.entry(uid.clone()).or_insert(0);
|
||||
*count += 1;
|
||||
tracing::error!(
|
||||
user_id = uid,
|
||||
consecutive_failures = *count,
|
||||
"Multi-user heartbeat failed: {}",
|
||||
err
|
||||
);
|
||||
if *count >= config.max_failures {
|
||||
tracing::error!(
|
||||
user_id = uid,
|
||||
"Multi-user heartbeat disabled for user after {} consecutive failures",
|
||||
count
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -726,7 +905,7 @@ mod tests {
|
||||
Arc<crate::workspace::Workspace>,
|
||||
Arc<dyn crate::llm::LlmProvider>,
|
||||
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
|
||||
Option<Arc<dyn crate::db::Database>>,
|
||||
Option<AdminScope>,
|
||||
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
|
||||
let _ = _fn_ptr;
|
||||
}
|
||||
|
||||
+3
-1
@@ -36,7 +36,9 @@ 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};
|
||||
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
|
||||
pub use heartbeat::{
|
||||
HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat, spawn_multi_user_heartbeat,
|
||||
};
|
||||
pub use router::{MessageIntent, Router};
|
||||
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
|
||||
pub use routine_engine::{RoutineEngine, SandboxReadiness};
|
||||
|
||||
@@ -28,12 +28,12 @@ use crate::agent::routine::{
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::context::{JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::RoutineError;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::llm::{
|
||||
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
|
||||
};
|
||||
use crate::tenant::AdminScope;
|
||||
use crate::tools::{
|
||||
ToolError, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_message,
|
||||
prepare_tool_params,
|
||||
@@ -99,7 +99,7 @@ pub(crate) fn routine_matches_message(routine: &Routine, message: &IncomingMessa
|
||||
/// The routine execution engine.
|
||||
pub struct RoutineEngine {
|
||||
config: RoutineConfig,
|
||||
store: Arc<dyn Database>,
|
||||
store: AdminScope,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
/// Sender for notifications (routed to channel manager).
|
||||
@@ -128,7 +128,7 @@ impl RoutineEngine {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
config: RoutineConfig,
|
||||
store: Arc<dyn Database>,
|
||||
store: AdminScope,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
@@ -782,12 +782,22 @@ impl RoutineEngine {
|
||||
});
|
||||
}
|
||||
|
||||
// Per-user workspace (same pattern as spawn_fire).
|
||||
let routine_workspace = if routine.user_id == self.workspace.user_id() {
|
||||
self.workspace.clone()
|
||||
} else {
|
||||
Arc::new(Workspace::new_with_db(
|
||||
&routine.user_id,
|
||||
Arc::clone(self.store.db()),
|
||||
))
|
||||
};
|
||||
|
||||
// Execute inline for manual triggers (caller wants to wait)
|
||||
let engine = EngineContext {
|
||||
config: self.config.clone(),
|
||||
store: self.store.clone(),
|
||||
llm: self.llm.clone(),
|
||||
workspace: self.workspace.clone(),
|
||||
workspace: routine_workspace,
|
||||
notify_tx: self.notify_tx.clone(),
|
||||
running_count: self.running_count.clone(),
|
||||
scheduler: self.scheduler.clone(),
|
||||
@@ -910,11 +920,23 @@ impl RoutineEngine {
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
|
||||
// Use per-user workspace so each routine executes in the correct
|
||||
// user's context. Fall back to the engine-wide workspace when the
|
||||
// routine belongs to the same user (avoids unnecessary allocation).
|
||||
let routine_workspace = if routine.user_id == self.workspace.user_id() {
|
||||
self.workspace.clone()
|
||||
} else {
|
||||
Arc::new(Workspace::new_with_db(
|
||||
&routine.user_id,
|
||||
Arc::clone(self.store.db()),
|
||||
))
|
||||
};
|
||||
|
||||
let engine = EngineContext {
|
||||
config: self.config.clone(),
|
||||
store: self.store.clone(),
|
||||
llm: self.llm.clone(),
|
||||
workspace: self.workspace.clone(),
|
||||
workspace: routine_workspace,
|
||||
notify_tx: self.notify_tx.clone(),
|
||||
running_count: self.running_count.clone(),
|
||||
scheduler: self.scheduler.clone(),
|
||||
@@ -967,7 +989,7 @@ impl RoutineEngine {
|
||||
/// an active state (Pending/InProgress/Stuck). Maps the final `JobState` to
|
||||
/// a `RunStatus` for the routine run.
|
||||
struct FullJobWatcher {
|
||||
store: Arc<dyn Database>,
|
||||
store: AdminScope,
|
||||
job_id: Uuid,
|
||||
routine_name: String,
|
||||
}
|
||||
@@ -978,7 +1000,7 @@ impl FullJobWatcher {
|
||||
/// Safety ceiling: 24 hours, derived from POLL_INTERVAL.
|
||||
const MAX_POLLS: u32 = (24 * 60 * 60) / Self::POLL_INTERVAL.as_secs() as u32;
|
||||
|
||||
fn new(store: Arc<dyn Database>, job_id: Uuid, routine_name: String) -> Self {
|
||||
fn new(store: AdminScope, job_id: Uuid, routine_name: String) -> Self {
|
||||
Self {
|
||||
store,
|
||||
job_id,
|
||||
@@ -1050,7 +1072,7 @@ impl FullJobWatcher {
|
||||
/// Shared context passed to the execution function.
|
||||
struct EngineContext {
|
||||
config: RoutineConfig,
|
||||
store: Arc<dyn Database>,
|
||||
store: AdminScope,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
|
||||
+22
-3
@@ -11,12 +11,12 @@ use uuid::Uuid;
|
||||
use crate::agent::task::{Task, TaskContext, TaskOutput};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::context::{ContextManager, JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::{Error, JobError};
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tenant::AdminScope;
|
||||
use crate::tools::{
|
||||
ApprovalContext, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_error,
|
||||
prepare_tool_params,
|
||||
@@ -52,7 +52,7 @@ struct ScheduledSubtask {
|
||||
pub struct SchedulerDeps {
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub store: Option<AdminScope>,
|
||||
pub hooks: Arc<HookRegistry>,
|
||||
}
|
||||
|
||||
@@ -64,7 +64,7 @@ pub struct Scheduler {
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
extension_manager: Option<Arc<ExtensionManager>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
store: Option<AdminScope>,
|
||||
hooks: Arc<HookRegistry>,
|
||||
/// SSE manager for live job event streaming.
|
||||
sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||
@@ -267,6 +267,20 @@ impl Scheduler {
|
||||
});
|
||||
}
|
||||
|
||||
// Per-user concurrency check
|
||||
if let Some(max_per_user) = self.config.max_jobs_per_user
|
||||
&& let Ok(ctx) = self.context_manager.get_context(job_id).await
|
||||
{
|
||||
let user_active = self
|
||||
.context_manager
|
||||
.active_jobs_for(&ctx.user_id)
|
||||
.await
|
||||
.len();
|
||||
if user_active >= max_per_user {
|
||||
return Err(JobError::MaxJobsExceeded { max: max_per_user });
|
||||
}
|
||||
}
|
||||
|
||||
// Transition job to in_progress
|
||||
self.context_manager
|
||||
.update_context(job_id, |ctx| {
|
||||
@@ -780,10 +794,15 @@ mod tests {
|
||||
allow_local_tools: true,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
max_tool_iterations: 10,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
max_jobs_per_user: None,
|
||||
max_tokens_per_job,
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: None,
|
||||
max_jobs_concurrent_per_user: None,
|
||||
};
|
||||
let cm = Arc::new(ContextManager::new(5));
|
||||
let llm: Arc<dyn LlmProvider> = Arc::new(StubLlm);
|
||||
|
||||
@@ -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::tenant::AdminScope;
|
||||
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
||||
|
||||
/// A job that has been detected as stuck.
|
||||
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
|
||||
/// Jobs in `InProgress` longer than this are treated as stuck.
|
||||
stuck_threshold: Duration,
|
||||
max_repair_attempts: u32,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
store: Option<AdminScope>,
|
||||
builder: Option<Arc<dyn SoftwareBuilder>>,
|
||||
tools: Option<Arc<ToolRegistry>>,
|
||||
}
|
||||
@@ -91,8 +91,8 @@ impl DefaultSelfRepair {
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a Store for tool failure tracking.
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
/// Add an admin-scoped store for tool failure tracking.
|
||||
pub fn with_store(mut self, store: AdminScope) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
@@ -806,7 +806,7 @@ mod tests {
|
||||
// Create self-repair with zero threshold (detect immediately),
|
||||
// wired with store, builder, and tools.
|
||||
let repair = DefaultSelfRepair::new(Arc::clone(&cm), Duration::from_secs(0), 3)
|
||||
.with_store(Arc::clone(&db))
|
||||
.with_store(crate::tenant::AdminScope::new(Arc::clone(&db)))
|
||||
.with_builder(
|
||||
Arc::clone(&builder) as Arc<dyn crate::tools::SoftwareBuilder>,
|
||||
tools,
|
||||
|
||||
+10
-3
@@ -175,6 +175,7 @@ impl Agent {
|
||||
pub(super) async fn process_user_input(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
tenant: crate::tenant::TenantCtx,
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
content: &str,
|
||||
@@ -351,7 +352,7 @@ impl Agent {
|
||||
|
||||
if let Some(intent) = self.router.route_command(&temp_message) {
|
||||
// Explicit command like /status, /job, /list - handle directly
|
||||
return self.handle_job_or_command(intent, message).await;
|
||||
return self.handle_job_or_command(intent, message, &tenant).await;
|
||||
}
|
||||
|
||||
// Natural language goes through the agentic loop
|
||||
@@ -462,7 +463,7 @@ impl Agent {
|
||||
|
||||
// Run the agentic tool execution loop
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, turn_messages)
|
||||
.run_agentic_loop(message, tenant, session.clone(), thread_id, turn_messages)
|
||||
.await;
|
||||
|
||||
// Re-acquire lock and check if interrupted
|
||||
@@ -1473,7 +1474,13 @@ impl Agent {
|
||||
|
||||
// Continue the agentic loop (a tool was already executed this turn)
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages)
|
||||
.run_agentic_loop(
|
||||
message,
|
||||
self.tenant_ctx(&message.user_id).await,
|
||||
session.clone(),
|
||||
thread_id,
|
||||
context_messages,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Handle the result
|
||||
|
||||
+3
-7
@@ -336,17 +336,12 @@ impl AppBuilder {
|
||||
ws = ws.with_memory_layers(self.config.workspace.memory_layers.clone());
|
||||
let ws = Arc::new(ws);
|
||||
|
||||
// Detect multi-tenant mode: when GATEWAY_USER_TOKENS is configured,
|
||||
// Detect multi-tenant mode: when the database has registered users,
|
||||
// each authenticated user needs their own workspace scope. Use
|
||||
// WorkspacePool (which implements WorkspaceResolver) to create
|
||||
// per-user workspaces on demand instead of sharing the startup
|
||||
// workspace across all users.
|
||||
let is_multi_tenant = self
|
||||
.config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.is_some_and(|gw| gw.user_tokens.is_some());
|
||||
let is_multi_tenant = db.has_any_users().await.unwrap_or(false);
|
||||
|
||||
if is_multi_tenant {
|
||||
let pool = Arc::new(crate::channels::web::server::WorkspacePool::new(
|
||||
@@ -880,6 +875,7 @@ impl AppBuilder {
|
||||
crate::agent::cost_guard::CostGuardConfig {
|
||||
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
|
||||
max_actions_per_hour: self.config.agent.max_actions_per_hour,
|
||||
max_cost_per_user_per_day_cents: self.config.agent.max_cost_per_user_per_day_cents,
|
||||
},
|
||||
));
|
||||
|
||||
|
||||
@@ -91,6 +91,34 @@ Browser-facing HTTP API and SSE/WebSocket real-time streaming. Axum-based, singl
|
||||
| DELETE | `/api/routines/{id}` | Delete a routine |
|
||||
| GET | `/api/routines/{id}/runs` | List runs for a specific routine |
|
||||
|
||||
### User Management (admin — requires `admin` role, see `docs/USER_MANAGEMENT_API.md`)
|
||||
| Method | Path | Description |
|
||||
|--------|------|-------------|
|
||||
| POST | `/api/admin/users` | Create a new user (returns one-time token) |
|
||||
| GET | `/api/admin/users` | List all users |
|
||||
| GET | `/api/admin/users/{id}` | Get a single user |
|
||||
| PATCH | `/api/admin/users/{id}` | Update user profile/metadata |
|
||||
| DELETE | `/api/admin/users/{id}` | Delete user and all data |
|
||||
| POST | `/api/admin/users/{id}/suspend` | Suspend a user |
|
||||
| POST | `/api/admin/users/{id}/activate` | Re-activate a user |
|
||||
| GET | `/api/admin/usage` | Per-user LLM usage stats |
|
||||
| GET | `/api/admin/users/{id}/secrets` | List a user's secrets (names only) |
|
||||
| PUT | `/api/admin/users/{id}/secrets/{name}` | Create or update a user's secret |
|
||||
| DELETE | `/api/admin/users/{id}/secrets/{name}` | Delete a user's secret |
|
||||
|
||||
### Profile (self-service)
|
||||
| Method | Path | Description |
|
||||
|--------|------|-------------|
|
||||
| GET | `/api/profile` | Get own profile |
|
||||
| PATCH | `/api/profile` | Update own display name/metadata |
|
||||
|
||||
### Tokens (self-service)
|
||||
| Method | Path | Description |
|
||||
|--------|------|-------------|
|
||||
| POST | `/api/tokens` | Create API token (returns plaintext once) |
|
||||
| GET | `/api/tokens` | List own tokens |
|
||||
| DELETE | `/api/tokens/{id}` | Revoke a token |
|
||||
|
||||
### Settings
|
||||
| Method | Path | Description |
|
||||
|--------|------|-------------|
|
||||
|
||||
+201
-20
@@ -5,6 +5,7 @@
|
||||
//! handlers can extract it via `AuthenticatedUser`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::num::NonZeroUsize;
|
||||
|
||||
use axum::{
|
||||
extract::{FromRequestParts, Request, State},
|
||||
@@ -13,18 +14,25 @@ use axum::{
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use subtle::ConstantTimeEq;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::db::Database;
|
||||
|
||||
/// Identity resolved from a bearer token.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UserIdentity {
|
||||
pub user_id: String,
|
||||
/// `admin` or `member`.
|
||||
pub role: String,
|
||||
/// Additional user scopes this identity can read from.
|
||||
pub workspace_read_scopes: Vec<String>,
|
||||
}
|
||||
|
||||
/// Hash a token with SHA-256 for constant-size, timing-safe storage.
|
||||
fn hash_token(token: &str) -> [u8; 32] {
|
||||
pub fn hash_token(token: &str) -> [u8; 32] {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(token.as_bytes());
|
||||
hasher.finalize().into()
|
||||
@@ -56,6 +64,7 @@ impl MultiAuthState {
|
||||
hash,
|
||||
UserIdentity {
|
||||
user_id,
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
)],
|
||||
@@ -64,6 +73,11 @@ impl MultiAuthState {
|
||||
}
|
||||
|
||||
/// Create a multi-user auth state from a map of tokens to identities.
|
||||
///
|
||||
/// **Test-only** — production multi-user auth is DB-backed via
|
||||
/// `DbAuthenticator`. This constructor is kept public (not `#[cfg(test)]`)
|
||||
/// because integration tests in `tests/` compile the crate as a library
|
||||
/// where `cfg(test)` is not set.
|
||||
pub fn multi(tokens: HashMap<String, UserIdentity>) -> Self {
|
||||
let hashed_tokens: Vec<([u8; 32], UserIdentity)> = tokens
|
||||
.into_iter()
|
||||
@@ -108,6 +122,112 @@ impl MultiAuthState {
|
||||
}
|
||||
}
|
||||
|
||||
/// DB-backed token authenticator with a bounded LRU cache.
|
||||
///
|
||||
/// Checks an LRU cache first (TTL 60s), then falls back to a DB query.
|
||||
/// The cache is bounded to `MAX_CACHE_ENTRIES` — when full, the least
|
||||
/// recently used entry is evicted regardless of TTL.
|
||||
///
|
||||
/// Revoking a token or suspending a user has at most 60s of stale
|
||||
/// authentication before the cache entry expires.
|
||||
#[derive(Clone)]
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub struct DbAuthenticator {
|
||||
store: Arc<dyn Database>,
|
||||
/// Bounded LRU cache: token_hash → (identity, inserted_at).
|
||||
cache: Arc<RwLock<lru::LruCache<[u8; 32], (UserIdentity, Instant)>>>,
|
||||
}
|
||||
|
||||
impl DbAuthenticator {
|
||||
/// Cache TTL — how long a successful auth is cached before re-querying the DB.
|
||||
const CACHE_TTL_SECS: u64 = 60;
|
||||
/// Maximum cache entries to prevent unbounded growth.
|
||||
// SAFETY: 1024 is non-zero, so the unwrap in `new()` is infallible.
|
||||
const MAX_CACHE_ENTRIES: NonZeroUsize = match NonZeroUsize::new(1024) {
|
||||
Some(v) => v,
|
||||
None => unreachable!(),
|
||||
};
|
||||
|
||||
pub fn new(store: Arc<dyn Database>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
cache: Arc::new(RwLock::new(lru::LruCache::new(Self::MAX_CACHE_ENTRIES))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Authenticate a token against the database, using cache when possible.
|
||||
///
|
||||
/// Returns `Ok(Some(identity))` on success, `Ok(None)` if the token is
|
||||
/// not found, or `Err(())` if the database is unreachable (so the caller
|
||||
/// can return 503 instead of 401).
|
||||
pub async fn authenticate(&self, candidate: &str) -> Result<Option<UserIdentity>, ()> {
|
||||
let hash = hash_token(candidate);
|
||||
|
||||
// Check cache first (promotes to most-recent on hit)
|
||||
{
|
||||
let mut cache = self.cache.write().await;
|
||||
if let Some((identity, inserted_at)) = cache.get(&hash) {
|
||||
if inserted_at.elapsed().as_secs() < Self::CACHE_TTL_SECS {
|
||||
return Ok(Some(identity.clone()));
|
||||
}
|
||||
// Expired — remove stale entry
|
||||
cache.pop(&hash);
|
||||
}
|
||||
}
|
||||
|
||||
// Cache miss or expired — query DB
|
||||
let (token_record, user_record) = match self.store.authenticate_token(&hash).await {
|
||||
Ok(Some(pair)) => pair,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "DB auth lookup failed, returning 503");
|
||||
return Err(());
|
||||
}
|
||||
};
|
||||
|
||||
let identity = UserIdentity {
|
||||
user_id: user_record.id.clone(),
|
||||
role: user_record.role.clone(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
};
|
||||
|
||||
// Record token usage (best-effort, don't block auth)
|
||||
let store = self.store.clone();
|
||||
let token_id = token_record.id;
|
||||
let user_id = user_record.id;
|
||||
tokio::spawn(async move {
|
||||
let _ = store.record_token_usage(token_id).await;
|
||||
let _ = store.record_login(&user_id).await;
|
||||
});
|
||||
|
||||
// Insert into bounded LRU — if full, least-recently-used entry is evicted
|
||||
{
|
||||
let mut cache = self.cache.write().await;
|
||||
cache.put(hash, (identity.clone(), Instant::now()));
|
||||
}
|
||||
|
||||
Ok(Some(identity))
|
||||
}
|
||||
}
|
||||
|
||||
/// Combined auth state: tries env-var tokens first, then DB-backed tokens.
|
||||
#[derive(Clone)]
|
||||
pub struct CombinedAuthState {
|
||||
/// In-memory tokens from GATEWAY_AUTH_TOKEN.
|
||||
pub env_auth: MultiAuthState,
|
||||
/// DB-backed token authenticator (optional — only when a database is available).
|
||||
pub db_auth: Option<DbAuthenticator>,
|
||||
}
|
||||
|
||||
impl From<MultiAuthState> for CombinedAuthState {
|
||||
fn from(env_auth: MultiAuthState) -> Self {
|
||||
Self {
|
||||
env_auth,
|
||||
db_auth: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Axum extractor that provides the authenticated user identity.
|
||||
///
|
||||
/// Only available on routes behind `auth_middleware`. Extracts the
|
||||
@@ -130,6 +250,31 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// Axum extractor that requires the authenticated user to have the `admin` role.
|
||||
///
|
||||
/// Use instead of `AuthenticatedUser` on endpoints that modify system-wide
|
||||
/// state (user management, model selection, extension/skill installation).
|
||||
pub struct AdminUser(pub UserIdentity);
|
||||
|
||||
impl<S> FromRequestParts<S> for AdminUser
|
||||
where
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Rejection = (StatusCode, &'static str);
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
|
||||
let identity = parts
|
||||
.extensions
|
||||
.get::<UserIdentity>()
|
||||
.cloned()
|
||||
.ok_or((StatusCode::UNAUTHORIZED, "Not authenticated"))?;
|
||||
if identity.role != "admin" {
|
||||
return Err((StatusCode::FORBIDDEN, "Admin role required"));
|
||||
}
|
||||
Ok(AdminUser(identity))
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether query-string token auth is allowed for this request.
|
||||
///
|
||||
/// Only GET requests to streaming endpoints may use `?token=xxx`. This
|
||||
@@ -166,39 +311,65 @@ fn query_token(request: &Request) -> Option<String> {
|
||||
|
||||
/// Auth middleware that validates bearer token from header or query param.
|
||||
///
|
||||
/// SSE connections can't set headers from `EventSource`, so we also accept
|
||||
/// `?token=xxx` as a query parameter, but only on SSE/WS endpoints.
|
||||
/// Tries env-var tokens first (constant-time, in-memory), then falls back
|
||||
/// to DB-backed token lookup if configured. SSE connections can't set
|
||||
/// headers from `EventSource`, so we also accept `?token=xxx` as a query
|
||||
/// parameter, but only on SSE/WS endpoints.
|
||||
///
|
||||
/// On successful authentication, inserts the matching `UserIdentity` into
|
||||
/// request extensions for downstream extraction via `AuthenticatedUser`.
|
||||
pub async fn auth_middleware(
|
||||
State(auth): State<MultiAuthState>,
|
||||
State(auth): State<CombinedAuthState>,
|
||||
headers: HeaderMap,
|
||||
mut request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
// Try Authorization header first.
|
||||
// RFC 6750 Section 2.1: auth-scheme comparison is case-insensitive.
|
||||
// Extract the candidate token from header or query param.
|
||||
let token = extract_token(&headers, &request);
|
||||
|
||||
if let Some(ref tok) = token {
|
||||
// 1. Try env-var tokens first (fast, constant-time, in-memory).
|
||||
if let Some(identity) = auth.env_auth.authenticate(tok) {
|
||||
request.extensions_mut().insert(identity.clone());
|
||||
return next.run(request).await;
|
||||
}
|
||||
|
||||
// 2. Fall back to DB-backed token lookup.
|
||||
if let Some(ref db_auth) = auth.db_auth {
|
||||
match db_auth.authenticate(tok).await {
|
||||
Ok(Some(identity)) => {
|
||||
request.extensions_mut().insert(identity);
|
||||
return next.run(request).await;
|
||||
}
|
||||
Err(()) => {
|
||||
return (StatusCode::SERVICE_UNAVAILABLE, "Database unavailable")
|
||||
.into_response();
|
||||
}
|
||||
Ok(None) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(StatusCode::UNAUTHORIZED, "Invalid or missing auth token").into_response()
|
||||
}
|
||||
|
||||
/// Extract a bearer token from the Authorization header or query parameter.
|
||||
fn extract_token(headers: &HeaderMap, request: &Request) -> Option<String> {
|
||||
// Try Authorization header first (RFC 6750).
|
||||
if let Some(auth_header) = headers.get("authorization")
|
||||
&& let Ok(value) = auth_header.to_str()
|
||||
&& value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("Bearer ")
|
||||
&& let Some(identity) = auth.authenticate(&value[7..])
|
||||
{
|
||||
request.extensions_mut().insert(identity.clone());
|
||||
return next.run(request).await;
|
||||
return Some(value[7..].to_string());
|
||||
}
|
||||
|
||||
// Fall back to query parameter, but only for SSE/WS endpoints.
|
||||
if allows_query_token_auth(&request)
|
||||
&& let Some(token) = query_token(&request)
|
||||
&& let Some(identity) = auth.authenticate(&token)
|
||||
{
|
||||
request.extensions_mut().insert(identity.clone());
|
||||
return next.run(request).await;
|
||||
// Fall back to query parameter for SSE/WS endpoints.
|
||||
if allows_query_token_auth(request) {
|
||||
return query_token(request);
|
||||
}
|
||||
|
||||
(StatusCode::UNAUTHORIZED, "Invalid or missing auth token").into_response()
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -227,6 +398,7 @@ mod tests {
|
||||
"tok-alice".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
@@ -234,6 +406,7 @@ mod tests {
|
||||
"tok-bob".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "bob".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
@@ -274,7 +447,10 @@ mod tests {
|
||||
/// Router with streaming endpoints (query auth allowed) and regular
|
||||
/// endpoints (query auth rejected).
|
||||
fn test_app(token: &str) -> Router {
|
||||
let state = MultiAuthState::single(token.to_string(), "test-user".to_string());
|
||||
let state = CombinedAuthState::from(MultiAuthState::single(
|
||||
token.to_string(),
|
||||
"test-user".to_string(),
|
||||
));
|
||||
Router::new()
|
||||
.route("/api/chat/events", get(dummy_handler))
|
||||
.route("/api/logs/events", get(dummy_handler))
|
||||
@@ -486,7 +662,7 @@ mod tests {
|
||||
|
||||
/// Build a multi-user router where each token maps to a distinct identity.
|
||||
fn multi_user_app(tokens: HashMap<String, UserIdentity>) -> Router {
|
||||
let state = MultiAuthState::multi(tokens);
|
||||
let state = CombinedAuthState::from(MultiAuthState::multi(tokens));
|
||||
Router::new()
|
||||
.route("/api/chat/events", get(identity_handler))
|
||||
.route("/api/chat/send", post(identity_handler))
|
||||
@@ -500,6 +676,7 @@ mod tests {
|
||||
"tok-alice".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["shared".to_string()],
|
||||
},
|
||||
);
|
||||
@@ -507,6 +684,7 @@ mod tests {
|
||||
"tok-bob".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "bob".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["shared".to_string(), "alice".to_string()],
|
||||
},
|
||||
);
|
||||
@@ -643,7 +821,10 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_multi_user_empty_scopes_for_single_user() {
|
||||
// Single-user mode creates identity with empty workspace_read_scopes.
|
||||
let state = MultiAuthState::single("tok-only".to_string(), "solo".to_string());
|
||||
let state = CombinedAuthState::from(MultiAuthState::single(
|
||||
"tok-only".to_string(),
|
||||
"solo".to_string(),
|
||||
));
|
||||
let app = Router::new()
|
||||
.route("/api/scopes", get(scopes_handler))
|
||||
.layer(middleware::from_fn_with_state(state, auth_middleware));
|
||||
|
||||
@@ -5,7 +5,10 @@
|
||||
pub mod jobs;
|
||||
pub mod memory;
|
||||
pub mod routines;
|
||||
pub mod secrets;
|
||||
pub mod skills;
|
||||
pub mod tokens;
|
||||
pub mod users;
|
||||
|
||||
// Modules not yet wired into server.rs router -- suppress dead_code until
|
||||
// they replace their inline counterparts.
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
//! Admin secrets provisioning handlers.
|
||||
//!
|
||||
//! Allows an admin (typically an application backend) to create, list, and
|
||||
//! delete secrets on behalf of individual users so their IronClaw agent can
|
||||
//! call back to external services with per-user credentials.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
};
|
||||
|
||||
use crate::channels::web::auth::AdminUser;
|
||||
use crate::channels::web::server::GatewayState;
|
||||
use crate::secrets::CreateSecretParams;
|
||||
|
||||
/// PUT /api/admin/users/{user_id}/secrets/{name} — create or update a secret.
|
||||
///
|
||||
/// Upserts: if a secret with the same (user_id, name) already exists it is
|
||||
/// overwritten. The plaintext value is encrypted at rest (AES-256-GCM) and
|
||||
/// never returned by any endpoint.
|
||||
pub async fn secrets_put_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_admin): AdminUser,
|
||||
Path((user_id, name)): Path<(String, String)>,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let secrets = state.secrets_store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Secrets store not available".to_string(),
|
||||
))?;
|
||||
|
||||
let value = body
|
||||
.get("value")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Missing required field 'value'".to_string(),
|
||||
))?
|
||||
.to_string();
|
||||
|
||||
let provider = body
|
||||
.get("provider")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
|
||||
let expires_at = body
|
||||
.get("expires_in_days")
|
||||
.and_then(|v| v.as_u64())
|
||||
.map(|d| d.min(36500))
|
||||
.map(|days| chrono::Utc::now() + chrono::Duration::days(days as i64));
|
||||
|
||||
let mut params = CreateSecretParams::new(name.clone(), value);
|
||||
if let Some(p) = provider {
|
||||
params = params.with_provider(p);
|
||||
}
|
||||
if let Some(exp) = expires_at {
|
||||
params = params.with_expiry(exp);
|
||||
}
|
||||
|
||||
secrets
|
||||
.create(&user_id, params)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"user_id": user_id,
|
||||
"name": name.to_lowercase(),
|
||||
"status": "created",
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/admin/users/{user_id}/secrets — list a user's secrets (names only).
|
||||
///
|
||||
/// Never returns secret values or hashes.
|
||||
pub async fn secrets_list_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_admin): AdminUser,
|
||||
Path(user_id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let secrets = state.secrets_store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Secrets store not available".to_string(),
|
||||
))?;
|
||||
|
||||
let refs = secrets
|
||||
.list(&user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let secrets_json: Vec<serde_json::Value> = refs
|
||||
.into_iter()
|
||||
.map(|r| {
|
||||
serde_json::json!({
|
||||
"name": r.name,
|
||||
"provider": r.provider,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"user_id": user_id,
|
||||
"secrets": secrets_json,
|
||||
})))
|
||||
}
|
||||
|
||||
/// DELETE /api/admin/users/{user_id}/secrets/{name} — delete a user's secret.
|
||||
pub async fn secrets_delete_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_admin): AdminUser,
|
||||
Path((user_id, name)): Path<(String, String)>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let secrets = state.secrets_store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Secrets store not available".to_string(),
|
||||
))?;
|
||||
|
||||
let deleted = secrets
|
||||
.delete(&user_id, &name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
if !deleted {
|
||||
return Err((StatusCode::NOT_FOUND, "Secret not found".to_string()));
|
||||
}
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"user_id": user_id,
|
||||
"name": name,
|
||||
"deleted": true,
|
||||
})))
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
//! API token management handlers.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
};
|
||||
use rand::RngCore;
|
||||
use rand::rngs::OsRng;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::auth::AuthenticatedUser;
|
||||
use crate::channels::web::server::GatewayState;
|
||||
|
||||
/// POST /api/tokens — create a new API token (returns plaintext ONCE).
|
||||
pub async fn tokens_create_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let name = body
|
||||
.get("name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Missing required field 'name'".to_string(),
|
||||
))?
|
||||
.to_string();
|
||||
|
||||
let expires_in_days: Option<i64> = body
|
||||
.get("expires_in_days")
|
||||
.and_then(|v| v.as_u64())
|
||||
.map(|d| d.min(36500) as i64);
|
||||
|
||||
let expires_at = expires_in_days.map(|days| chrono::Utc::now() + chrono::Duration::days(days));
|
||||
|
||||
// Generate 32 random bytes for the token.
|
||||
// Hash the hex-encoded plaintext (what the user sends as Bearer token),
|
||||
// NOT the raw bytes — must match hash_token() in auth.rs.
|
||||
let mut token_bytes = [0u8; 32];
|
||||
OsRng.fill_bytes(&mut token_bytes);
|
||||
let plaintext_token = hex::encode(token_bytes);
|
||||
let hash = crate::channels::web::auth::hash_token(&plaintext_token);
|
||||
|
||||
// First 8 chars of the hex token as a prefix for identification.
|
||||
let token_prefix = &plaintext_token[..8];
|
||||
|
||||
// Admin users can create tokens for other users via optional "user_id" field.
|
||||
let target_user = body
|
||||
.get("user_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|_| user.role == "admin")
|
||||
.unwrap_or(&user.user_id);
|
||||
|
||||
// Verify the target user exists to prevent orphan tokens.
|
||||
if target_user != user.user_id {
|
||||
store
|
||||
.get_user(target_user)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("Target user '{target_user}' not found"),
|
||||
))?;
|
||||
}
|
||||
|
||||
let record = store
|
||||
.create_api_token(target_user, &name, &hash, token_prefix, expires_at)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
// Return the plaintext token — this is the ONLY time it is shown.
|
||||
Ok(Json(serde_json::json!({
|
||||
"token": plaintext_token,
|
||||
"id": record.id.to_string(),
|
||||
"name": record.name,
|
||||
"token_prefix": record.token_prefix,
|
||||
"expires_at": record.expires_at.map(|dt| dt.to_rfc3339()),
|
||||
"created_at": record.created_at.to_rfc3339(),
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/tokens — list the current user's tokens (no hashes).
|
||||
pub async fn tokens_list_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let tokens = store
|
||||
.list_api_tokens(&user.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let tokens_json: Vec<serde_json::Value> = tokens
|
||||
.into_iter()
|
||||
.map(|t| {
|
||||
serde_json::json!({
|
||||
"id": t.id.to_string(),
|
||||
"name": t.name,
|
||||
"token_prefix": t.token_prefix,
|
||||
"expires_at": t.expires_at.map(|dt| dt.to_rfc3339()),
|
||||
"last_used_at": t.last_used_at.map(|dt| dt.to_rfc3339()),
|
||||
"created_at": t.created_at.to_rfc3339(),
|
||||
"revoked_at": t.revoked_at.map(|dt| dt.to_rfc3339()),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(serde_json::json!({ "tokens": tokens_json })))
|
||||
}
|
||||
|
||||
/// DELETE /api/tokens/{id} — revoke a token.
|
||||
pub async fn tokens_revoke_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let token_id = Uuid::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid token ID".to_string()))?;
|
||||
|
||||
let revoked = store
|
||||
.revoke_api_token(token_id, &user.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
if !revoked {
|
||||
return Err((StatusCode::NOT_FOUND, "Token not found".to_string()));
|
||||
}
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"status": "revoked",
|
||||
"id": token_id.to_string(),
|
||||
})))
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
//! User management API handlers (admin).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
};
|
||||
use rand::RngCore;
|
||||
use rand::rngs::OsRng;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::auth::{AdminUser, AuthenticatedUser};
|
||||
use crate::channels::web::server::GatewayState;
|
||||
use crate::db::UserRecord;
|
||||
|
||||
/// POST /api/admin/users — create a new user.
|
||||
pub async fn users_create_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(user): AdminUser,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let display_name = body
|
||||
.get("display_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Missing required field 'display_name'".to_string(),
|
||||
))?
|
||||
.to_string();
|
||||
|
||||
let email = body.get("email").and_then(|v| v.as_str()).map(String::from);
|
||||
let role = body
|
||||
.get("role")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("member")
|
||||
.to_string();
|
||||
if role != "admin" && role != "member" {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"role must be 'admin' or 'member'".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let user_id = Uuid::new_v4().to_string();
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let user_record = UserRecord {
|
||||
id: user_id.clone(),
|
||||
email,
|
||||
display_name: display_name.clone(),
|
||||
status: "active".to_string(),
|
||||
role,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
last_login_at: None,
|
||||
created_by: match store.get_user(&user.user_id).await {
|
||||
Ok(Some(_)) => Some(user.user_id.clone()),
|
||||
_ => None,
|
||||
},
|
||||
metadata: serde_json::json!({}),
|
||||
};
|
||||
|
||||
store
|
||||
.create_user(&user_record)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
// Generate a first API token so the new user can authenticate immediately.
|
||||
// Hash the hex-encoded plaintext (what the user sends as Bearer token),
|
||||
// NOT the raw bytes — must match hash_token() in auth.rs.
|
||||
let mut token_bytes = [0u8; 32];
|
||||
OsRng.fill_bytes(&mut token_bytes);
|
||||
let plaintext_token = hex::encode(token_bytes);
|
||||
let token_hash = crate::channels::web::auth::hash_token(&plaintext_token);
|
||||
let token_prefix = &plaintext_token[..8];
|
||||
|
||||
let _token_record = store
|
||||
.create_api_token(&user_id, "initial", &token_hash, token_prefix, None)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": user_record.id,
|
||||
"email": user_record.email,
|
||||
"display_name": user_record.display_name,
|
||||
"status": user_record.status,
|
||||
"role": user_record.role,
|
||||
"token": plaintext_token,
|
||||
"created_at": user_record.created_at.to_rfc3339(),
|
||||
"created_by": user_record.created_by,
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/admin/users — list all users.
|
||||
pub async fn users_list_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let users = store
|
||||
.list_users(None)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let users_json: Vec<serde_json::Value> = users
|
||||
.into_iter()
|
||||
.map(|u| {
|
||||
serde_json::json!({
|
||||
"id": u.id,
|
||||
"email": u.email,
|
||||
"display_name": u.display_name,
|
||||
"status": u.status,
|
||||
"role": u.role,
|
||||
"created_at": u.created_at.to_rfc3339(),
|
||||
"updated_at": u.updated_at.to_rfc3339(),
|
||||
"last_login_at": u.last_login_at.map(|dt| dt.to_rfc3339()),
|
||||
"created_by": u.created_by,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(serde_json::json!({ "users": users_json })))
|
||||
}
|
||||
|
||||
/// GET /api/admin/users/{id} — get a single user.
|
||||
pub async fn users_detail_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let user_record = store
|
||||
.get_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": user_record.id,
|
||||
"email": user_record.email,
|
||||
"display_name": user_record.display_name,
|
||||
"status": user_record.status,
|
||||
"role": user_record.role,
|
||||
"created_at": user_record.created_at.to_rfc3339(),
|
||||
"updated_at": user_record.updated_at.to_rfc3339(),
|
||||
"last_login_at": user_record.last_login_at.map(|dt| dt.to_rfc3339()),
|
||||
"created_by": user_record.created_by,
|
||||
"metadata": user_record.metadata,
|
||||
})))
|
||||
}
|
||||
|
||||
/// PATCH /api/admin/users/{id} — update a user's profile.
|
||||
pub async fn users_update_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
Path(id): Path<String>,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
// Verify the user exists.
|
||||
let existing = store
|
||||
.get_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
let display_name = body
|
||||
.get("display_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&existing.display_name);
|
||||
|
||||
let metadata = body.get("metadata").unwrap_or(&existing.metadata);
|
||||
|
||||
store
|
||||
.update_user_profile(&id, display_name, metadata)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
// Re-fetch the updated record to return consistent data.
|
||||
let updated = store
|
||||
.get_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": updated.id,
|
||||
"email": updated.email,
|
||||
"display_name": updated.display_name,
|
||||
"status": updated.status,
|
||||
"role": updated.role,
|
||||
"created_at": updated.created_at.to_rfc3339(),
|
||||
"updated_at": updated.updated_at.to_rfc3339(),
|
||||
"metadata": updated.metadata,
|
||||
})))
|
||||
}
|
||||
|
||||
/// POST /api/admin/users/{id}/suspend — suspend a user.
|
||||
pub async fn users_suspend_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
// Verify the user exists.
|
||||
store
|
||||
.get_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
store
|
||||
.update_user_status(&id, "suspended")
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": id,
|
||||
"status": "suspended",
|
||||
})))
|
||||
}
|
||||
|
||||
/// POST /api/admin/users/{id}/activate — activate a user.
|
||||
pub async fn users_activate_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
// Verify the user exists.
|
||||
store
|
||||
.get_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
store
|
||||
.update_user_status(&id, "active")
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": id,
|
||||
"status": "active",
|
||||
})))
|
||||
}
|
||||
|
||||
/// DELETE /api/admin/users/{id} — delete a user and all their data.
|
||||
pub async fn users_delete_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let deleted = store
|
||||
.delete_user(&id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
if !deleted {
|
||||
return Err((StatusCode::NOT_FOUND, "User not found".to_string()));
|
||||
}
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": id,
|
||||
"deleted": true,
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/profile — get the authenticated user's own profile.
|
||||
pub async fn profile_get_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let record = store
|
||||
.get_user(&user.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": record.id,
|
||||
"email": record.email,
|
||||
"display_name": record.display_name,
|
||||
"status": record.status,
|
||||
"role": record.role,
|
||||
"created_at": record.created_at.to_rfc3339(),
|
||||
"last_login_at": record.last_login_at.map(|dt| dt.to_rfc3339()),
|
||||
})))
|
||||
}
|
||||
|
||||
/// PATCH /api/profile — update the authenticated user's own profile.
|
||||
pub async fn profile_update_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let current = store
|
||||
.get_user(&user.user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?;
|
||||
|
||||
let display_name = body
|
||||
.get("display_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(¤t.display_name);
|
||||
let metadata = body.get("metadata").unwrap_or(¤t.metadata);
|
||||
|
||||
store
|
||||
.update_user_profile(&user.user_id, display_name, metadata)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"id": user.user_id,
|
||||
"display_name": display_name,
|
||||
"updated": true,
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/admin/usage — per-user LLM usage stats.
|
||||
pub async fn usage_stats_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AdminUser(_user): AdminUser,
|
||||
axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let store = state.store.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
let user_id = params.get("user_id").map(|s| s.as_str());
|
||||
let period = params.get("period").map(|s| s.as_str()).unwrap_or("day");
|
||||
let since = match period {
|
||||
"week" => chrono::Utc::now() - chrono::Duration::days(7),
|
||||
"month" => chrono::Utc::now() - chrono::Duration::days(30),
|
||||
_ => chrono::Utc::now() - chrono::Duration::days(1),
|
||||
};
|
||||
|
||||
let stats = store
|
||||
.user_usage_stats(user_id, since)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let entries: Vec<serde_json::Value> = stats
|
||||
.iter()
|
||||
.map(|s| {
|
||||
serde_json::json!({
|
||||
"user_id": s.user_id,
|
||||
"model": s.model,
|
||||
"call_count": s.call_count,
|
||||
"input_tokens": s.input_tokens,
|
||||
"output_tokens": s.output_tokens,
|
||||
"total_cost": s.total_cost.to_string(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(serde_json::json!({
|
||||
"period": period,
|
||||
"since": since.to_rfc3339(),
|
||||
"usage": entries,
|
||||
})))
|
||||
}
|
||||
@@ -54,10 +54,37 @@ fn validate_webhook_secret(
|
||||
///
|
||||
/// This endpoint is **public** (no gateway auth token required) but protected
|
||||
/// by the per-routine webhook secret sent via the `X-Webhook-Secret` header.
|
||||
///
|
||||
/// **Single-user/backward-compatible**: looks up routines by path across all
|
||||
/// users. For multi-tenant isolation, use the user-scoped endpoint at
|
||||
/// `/api/webhooks/u/{user_id}/{path}` instead.
|
||||
pub async fn webhook_trigger_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Path(path): Path<String>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
fire_webhook_inner(state, &path, None, &headers).await
|
||||
}
|
||||
|
||||
/// Handle incoming webhook POST to `/api/webhooks/u/{user_id}/{path}`.
|
||||
///
|
||||
/// User-scoped variant for multi-tenant deployments. The `user_id` in the URL
|
||||
/// restricts the routine lookup to that user only, preventing cross-user
|
||||
/// webhook triggering even when paths collide.
|
||||
pub async fn webhook_trigger_user_scoped_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Path((user_id, path)): Path<(String, String)>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
fire_webhook_inner(state, &path, Some(&user_id), &headers).await
|
||||
}
|
||||
|
||||
/// Shared webhook logic for both scoped and unscoped endpoints.
|
||||
async fn fire_webhook_inner(
|
||||
state: Arc<GatewayState>,
|
||||
path: &str,
|
||||
user_id: Option<&str>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
// Rate limit check
|
||||
if !state.webhook_rate_limiter.check() {
|
||||
@@ -72,9 +99,9 @@ pub async fn webhook_trigger_handler(
|
||||
"Database not available".to_string(),
|
||||
))?;
|
||||
|
||||
// Targeted query instead of loading all routines
|
||||
// Targeted query — when user_id is provided, restrict to that user's routines
|
||||
let routine = store
|
||||
.get_webhook_routine_by_path(&path)
|
||||
.get_webhook_routine_by_path(path, user_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
|
||||
.ok_or((
|
||||
@@ -99,7 +126,7 @@ pub async fn webhook_trigger_handler(
|
||||
))?
|
||||
};
|
||||
|
||||
let run_id = engine.fire_webhook(routine.id, &path).await.map_err(|e| {
|
||||
let run_id = engine.fire_webhook(routine.id, path).await.map_err(|e| {
|
||||
let status = match &e {
|
||||
crate::error::RoutineError::NotFound { .. } => StatusCode::NOT_FOUND,
|
||||
crate::error::RoutineError::Disabled { .. }
|
||||
|
||||
+28
-65
@@ -18,6 +18,7 @@ pub mod auth;
|
||||
pub(crate) mod handlers;
|
||||
pub mod log_layer;
|
||||
pub mod openai_compat;
|
||||
pub mod responses_api;
|
||||
pub mod server;
|
||||
pub mod sse;
|
||||
pub mod types;
|
||||
@@ -55,7 +56,7 @@ use crate::workspace::Workspace;
|
||||
|
||||
use self::log_layer::{LogBroadcaster, LogLevelHandle};
|
||||
|
||||
use self::auth::MultiAuthState;
|
||||
use self::auth::{CombinedAuthState, DbAuthenticator, MultiAuthState};
|
||||
use self::server::GatewayState;
|
||||
use self::sse::SseManager;
|
||||
use self::types::AppEvent;
|
||||
@@ -64,8 +65,8 @@ use self::types::AppEvent;
|
||||
pub struct GatewayChannel {
|
||||
config: GatewayConfig,
|
||||
state: Arc<GatewayState>,
|
||||
/// Multi-user auth state (replaces bare auth_token).
|
||||
auth: MultiAuthState,
|
||||
/// Combined auth state: env-var tokens + optional DB-backed tokens.
|
||||
auth: CombinedAuthState,
|
||||
}
|
||||
|
||||
impl GatewayChannel {
|
||||
@@ -73,7 +74,7 @@ impl GatewayChannel {
|
||||
///
|
||||
/// If no auth token is configured, generates a random one and prints it.
|
||||
/// Builds a single-user `MultiAuthState` from the config.
|
||||
pub fn new(config: GatewayConfig) -> Self {
|
||||
pub fn new(config: GatewayConfig, owner_id: String) -> Self {
|
||||
let auth_token = config.auth_token.clone().unwrap_or_else(|| {
|
||||
use rand::RngCore;
|
||||
use rand::rngs::OsRng;
|
||||
@@ -82,64 +83,11 @@ impl GatewayChannel {
|
||||
bytes.iter().map(|b| format!("{b:02x}")).collect()
|
||||
});
|
||||
|
||||
let auth = MultiAuthState::single(auth_token, config.user_id.clone());
|
||||
|
||||
let state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(None),
|
||||
sse: Arc::new(SseManager::new()),
|
||||
workspace: None,
|
||||
workspace_pool: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
skill_registry: None,
|
||||
skill_catalog: None,
|
||||
chat_rate_limiter: server::PerUserRateLimiter::new(30, 60),
|
||||
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
||||
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||
registry_entries: Vec::new(),
|
||||
cost_guard: None,
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: server::ActiveConfigSnapshot::default(),
|
||||
});
|
||||
|
||||
Self {
|
||||
config,
|
||||
state,
|
||||
auth,
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebind the single-user auth identity to the durable owner scope while
|
||||
/// preserving the configured gateway sender/routing identity.
|
||||
pub fn with_owner_scope(mut self, owner_id: impl Into<String>) -> Self {
|
||||
let owner_id = owner_id.into();
|
||||
let single_user_token = if self.config.user_tokens.is_none() {
|
||||
self.auth.first_token().map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
let auth = CombinedAuthState {
|
||||
env_auth: MultiAuthState::single(auth_token, owner_id.clone()),
|
||||
db_auth: None,
|
||||
};
|
||||
if let Some(token) = single_user_token {
|
||||
self.auth = MultiAuthState::single(token, owner_id.clone());
|
||||
}
|
||||
self.rebuild_state(|s| s.owner_id = owner_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Create a gateway channel with a pre-built multi-user auth state.
|
||||
pub fn new_multi_auth(config: GatewayConfig, auth: MultiAuthState) -> Self {
|
||||
let state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(None),
|
||||
sse: Arc::new(SseManager::new()),
|
||||
@@ -154,8 +102,7 @@ impl GatewayChannel {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
owner_id,
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -163,12 +110,13 @@ impl GatewayChannel {
|
||||
skill_catalog: None,
|
||||
chat_rate_limiter: server::PerUserRateLimiter::new(30, 60),
|
||||
oauth_rate_limiter: server::RateLimiter::new(10, 60),
|
||||
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||
registry_entries: Vec::new(),
|
||||
cost_guard: None,
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
webhook_rate_limiter: server::RateLimiter::new(10, 60),
|
||||
active_config: server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
Self {
|
||||
@@ -196,7 +144,6 @@ impl GatewayChannel {
|
||||
prompt_queue: self.state.prompt_queue.clone(),
|
||||
scheduler: self.state.scheduler.clone(),
|
||||
owner_id: self.state.owner_id.clone(),
|
||||
default_sender_id: self.state.default_sender_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: self.state.ws_tracker.clone(),
|
||||
llm_provider: self.state.llm_provider.clone(),
|
||||
@@ -210,6 +157,7 @@ impl GatewayChannel {
|
||||
routine_engine: Arc::clone(&self.state.routine_engine),
|
||||
startup_time: self.state.startup_time,
|
||||
active_config: self.state.active_config.clone(),
|
||||
secrets_store: self.state.secrets_store.clone(),
|
||||
};
|
||||
mutate(&mut new_state);
|
||||
self.state = Arc::new(new_state);
|
||||
@@ -257,6 +205,12 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable DB-backed token authentication alongside env-var tokens.
|
||||
pub fn with_db_auth(mut self, store: Arc<dyn Database>) -> Self {
|
||||
self.auth.db_auth = Some(DbAuthenticator::new(store));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the container job manager for sandbox operations.
|
||||
pub fn with_job_manager(mut self, jm: Arc<ContainerJobManager>) -> Self {
|
||||
self.rebuild_state(|s| s.job_manager = Some(jm));
|
||||
@@ -327,6 +281,15 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the secrets store for admin secret provisioning.
|
||||
pub fn with_secrets_store(
|
||||
mut self,
|
||||
store: Arc<dyn crate::secrets::SecretsStore + Send + Sync>,
|
||||
) -> Self {
|
||||
self.rebuild_state(|s| s.secrets_store = Some(store));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the per-user workspace pool for multi-user mode.
|
||||
pub fn with_workspace_pool(mut self, pool: Arc<server::WorkspacePool>) -> Self {
|
||||
self.rebuild_state(|s| s.workspace_pool = Some(pool));
|
||||
@@ -335,7 +298,7 @@ impl GatewayChannel {
|
||||
|
||||
/// Get the first auth token (for printing to console on startup).
|
||||
pub fn auth_token(&self) -> &str {
|
||||
self.auth.first_token().unwrap_or("")
|
||||
self.auth.env_auth.first_token().unwrap_or("")
|
||||
}
|
||||
|
||||
/// Get a reference to the shared gateway state (for the agent to push SSE events).
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+97
-13
@@ -16,7 +16,7 @@ use axum::{
|
||||
IntoResponse,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
},
|
||||
routing::{get, post},
|
||||
routing::{get, post, put},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -31,7 +31,7 @@ use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::channels::IncomingMessage;
|
||||
use crate::channels::relay::DEFAULT_RELAY_NAME;
|
||||
use crate::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
AuthenticatedUser, CombinedAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
use crate::channels::web::handlers::jobs::{
|
||||
job_files_list_handler, job_files_read_handler, jobs_cancel_handler, jobs_detail_handler,
|
||||
@@ -347,8 +347,6 @@ pub struct GatewayState {
|
||||
pub prompt_queue: Option<PromptQueue>,
|
||||
/// Durable owner scope for persistence and unauthenticated callback flows.
|
||||
pub owner_id: String,
|
||||
/// Default sender/routing identity for gateway-originated messages.
|
||||
pub default_sender_id: String,
|
||||
/// Shutdown signal sender.
|
||||
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
|
||||
/// WebSocket connection tracker.
|
||||
@@ -378,6 +376,8 @@ pub struct GatewayState {
|
||||
pub startup_time: std::time::Instant,
|
||||
/// Snapshot of active (resolved) configuration for the frontend.
|
||||
pub active_config: ActiveConfigSnapshot,
|
||||
/// Secrets store for admin secret provisioning.
|
||||
pub secrets_store: Option<Arc<dyn crate::secrets::SecretsStore + Send + Sync>>,
|
||||
}
|
||||
|
||||
/// Start the gateway HTTP server.
|
||||
@@ -386,7 +386,7 @@ pub struct GatewayState {
|
||||
pub async fn start_server(
|
||||
addr: SocketAddr,
|
||||
state: Arc<GatewayState>,
|
||||
auth: MultiAuthState,
|
||||
auth: CombinedAuthState,
|
||||
) -> Result<SocketAddr, crate::error::ChannelError> {
|
||||
let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
|
||||
crate::error::ChannelError::StartupFailed {
|
||||
@@ -414,6 +414,11 @@ pub async fn start_server(
|
||||
.route(
|
||||
"/api/webhooks/{path}",
|
||||
post(crate::channels::web::handlers::webhooks::webhook_trigger_handler),
|
||||
)
|
||||
// User-scoped webhook endpoint for multi-tenant isolation
|
||||
.route(
|
||||
"/api/webhooks/u/{user_id}/{path}",
|
||||
post(crate::channels::web::handlers::webhooks::webhook_trigger_user_scoped_handler),
|
||||
);
|
||||
|
||||
// Protected routes (require auth)
|
||||
@@ -507,6 +512,57 @@ pub async fn start_server(
|
||||
"/api/settings/{key}",
|
||||
axum::routing::delete(settings_delete_handler),
|
||||
)
|
||||
// User management (admin)
|
||||
.route(
|
||||
"/api/admin/users",
|
||||
get(super::handlers::users::users_list_handler)
|
||||
.post(super::handlers::users::users_create_handler),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/users/{id}",
|
||||
get(super::handlers::users::users_detail_handler)
|
||||
.patch(super::handlers::users::users_update_handler)
|
||||
.delete(super::handlers::users::users_delete_handler),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/users/{id}/suspend",
|
||||
post(super::handlers::users::users_suspend_handler),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/users/{id}/activate",
|
||||
post(super::handlers::users::users_activate_handler),
|
||||
)
|
||||
// Admin secrets provisioning (per-user)
|
||||
.route(
|
||||
"/api/admin/users/{user_id}/secrets",
|
||||
get(super::handlers::secrets::secrets_list_handler),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/users/{user_id}/secrets/{name}",
|
||||
put(super::handlers::secrets::secrets_put_handler)
|
||||
.delete(super::handlers::secrets::secrets_delete_handler),
|
||||
)
|
||||
// Usage reporting (admin)
|
||||
.route(
|
||||
"/api/admin/usage",
|
||||
get(super::handlers::users::usage_stats_handler),
|
||||
)
|
||||
// User self-service profile
|
||||
.route(
|
||||
"/api/profile",
|
||||
get(super::handlers::users::profile_get_handler)
|
||||
.patch(super::handlers::users::profile_update_handler),
|
||||
)
|
||||
// Token management
|
||||
.route(
|
||||
"/api/tokens",
|
||||
get(super::handlers::tokens::tokens_list_handler)
|
||||
.post(super::handlers::tokens::tokens_create_handler),
|
||||
)
|
||||
.route(
|
||||
"/api/tokens/{id}",
|
||||
axum::routing::delete(super::handlers::tokens::tokens_revoke_handler),
|
||||
)
|
||||
// Gateway control plane
|
||||
.route("/api/gateway/status", get(gateway_status_handler))
|
||||
// OpenAI-compatible API
|
||||
@@ -515,6 +571,15 @@ pub async fn start_server(
|
||||
post(super::openai_compat::chat_completions_handler),
|
||||
)
|
||||
.route("/v1/models", get(super::openai_compat::models_handler))
|
||||
// OpenAI Responses API (routes through the full agent loop)
|
||||
.route(
|
||||
"/v1/responses",
|
||||
post(super::responses_api::create_response_handler),
|
||||
)
|
||||
.route(
|
||||
"/v1/responses/{id}",
|
||||
get(super::responses_api::get_response_handler),
|
||||
)
|
||||
.route_layer(middleware::from_fn_with_state(
|
||||
auth_state.clone(),
|
||||
auth_middleware,
|
||||
@@ -557,6 +622,7 @@ pub async fn start_server(
|
||||
axum::http::Method::GET,
|
||||
axum::http::Method::POST,
|
||||
axum::http::Method::PUT,
|
||||
axum::http::Method::PATCH,
|
||||
axum::http::Method::DELETE,
|
||||
])
|
||||
.allow_headers(AllowHeaders::list([
|
||||
@@ -571,6 +637,25 @@ pub async fn start_server(
|
||||
.merge(projects)
|
||||
.merge(protected)
|
||||
.layer(DefaultBodyLimit::max(10 * 1024 * 1024)) // 10 MB max request body (image uploads)
|
||||
.layer(tower_http::catch_panic::CatchPanicLayer::custom(
|
||||
|panic_info: Box<dyn std::any::Any + Send + 'static>| {
|
||||
let detail = if let Some(s) = panic_info.downcast_ref::<String>() {
|
||||
s.clone()
|
||||
} else if let Some(s) = panic_info.downcast_ref::<&str>() {
|
||||
(*s).to_string()
|
||||
} else {
|
||||
"unknown panic".to_string()
|
||||
};
|
||||
tracing::error!("Handler panicked: {}", detail);
|
||||
axum::http::Response::builder()
|
||||
.status(axum::http::StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.header("content-type", "text/plain")
|
||||
.body(axum::body::Body::from("Internal Server Error"))
|
||||
.unwrap_or_else(|_| {
|
||||
axum::http::Response::new(axum::body::Body::from("Internal Server Error"))
|
||||
})
|
||||
},
|
||||
))
|
||||
.layer(cors)
|
||||
.layer(SetResponseHeaderLayer::if_not_present(
|
||||
header::X_CONTENT_TYPE_OPTIONS,
|
||||
@@ -1298,9 +1383,6 @@ async fn chat_send_handler(
|
||||
}
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, &req.content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
// Prefer timezone from JSON body, fall back to X-Timezone header
|
||||
let tz = req
|
||||
.timezone
|
||||
@@ -1402,9 +1484,6 @@ async fn chat_approval_handler(
|
||||
})?;
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
|
||||
if let Some(ref thread_id) = req.thread_id {
|
||||
msg = msg.with_thread(thread_id);
|
||||
@@ -2980,7 +3059,6 @@ mod tests {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
@@ -2995,6 +3073,7 @@ mod tests {
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -3061,6 +3140,7 @@ mod tests {
|
||||
// without needing the full auth middleware layer.
|
||||
req.extensions_mut().insert(UserIdentity {
|
||||
user_id: "test".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
});
|
||||
|
||||
@@ -3145,6 +3225,7 @@ mod tests {
|
||||
// without needing the full auth middleware layer.
|
||||
req.extensions_mut().insert(UserIdentity {
|
||||
user_id: "test".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
});
|
||||
|
||||
@@ -3194,7 +3275,10 @@ mod tests {
|
||||
let state = test_gateway_state(None);
|
||||
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let auth = MultiAuthState::single("test-token".to_string(), "test".to_string());
|
||||
let auth = CombinedAuthState::from(crate::channels::web::auth::MultiAuthState::single(
|
||||
"test-token".to_string(),
|
||||
"test".to_string(),
|
||||
));
|
||||
let bound = start_server(addr, state.clone(), auth)
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
@@ -186,6 +186,13 @@ function authenticate() {
|
||||
connectSSE();
|
||||
connectLogSSE();
|
||||
startGatewayStatusPolling();
|
||||
// Hide the Users settings tab for non-admin users.
|
||||
apiFetch('/api/profile').then(function(profile) {
|
||||
if (profile && profile.role !== 'admin') {
|
||||
var usersTab = document.querySelector('[data-settings-subtab="users"]');
|
||||
if (usersTab) usersTab.style.display = 'none';
|
||||
}
|
||||
}).catch(function() {});
|
||||
checkTeeStatus();
|
||||
loadThreads();
|
||||
loadMemoryTree();
|
||||
@@ -4339,6 +4346,142 @@ function formatRelativeTime(isoString) {
|
||||
return future ? I18n.t('time.daysFromNow', { n: days }) : I18n.t('time.daysAgo', { n: days });
|
||||
}
|
||||
|
||||
// --- Users (admin) ---
|
||||
|
||||
function loadUsers() {
|
||||
apiFetch('/api/admin/users').then(function(data) {
|
||||
renderUsersList(data.users || []);
|
||||
}).catch(function(err) {
|
||||
// Non-admin users get 403 — show a message instead of an error
|
||||
var tbody = document.getElementById('users-tbody');
|
||||
var empty = document.getElementById('users-empty');
|
||||
if (tbody) tbody.innerHTML = '';
|
||||
if (empty) {
|
||||
empty.style.display = 'block';
|
||||
empty.textContent = 'Admin access required to manage users.';
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
function renderUsersList(users) {
|
||||
var tbody = document.getElementById('users-tbody');
|
||||
var empty = document.getElementById('users-empty');
|
||||
if (!users || users.length === 0) {
|
||||
tbody.innerHTML = '';
|
||||
empty.style.display = 'block';
|
||||
empty.textContent = 'No users found. Create the first user to get started.';
|
||||
return;
|
||||
}
|
||||
empty.style.display = 'none';
|
||||
tbody.innerHTML = users.map(function(u) {
|
||||
var statusClass = u.status === 'active' ? 'active' : 'failed';
|
||||
var roleLabel = u.role === 'admin' ? '<span class="badge badge-admin">admin</span>' : '<span class="badge">member</span>';
|
||||
var actions = '';
|
||||
if (u.status === 'active') {
|
||||
actions += '<button class="btn-small btn-danger" data-action="suspend-user" data-user-id="' + escapeHtml(u.id) + '">Suspend</button> ';
|
||||
} else {
|
||||
actions += '<button class="btn-small btn-primary" data-action="activate-user" data-user-id="' + escapeHtml(u.id) + '">Activate</button> ';
|
||||
}
|
||||
actions += '<button class="btn-small" data-action="create-token" data-user-id="' + escapeHtml(u.id) + '" data-user-name="' + escapeHtml(u.display_name) + '">+ Token</button>';
|
||||
return '<tr>'
|
||||
+ '<td class="user-id" title="' + escapeHtml(u.id) + '">' + escapeHtml(u.id.substring(0, 8)) + '…</td>'
|
||||
+ '<td>' + escapeHtml(u.display_name) + '</td>'
|
||||
+ '<td>' + escapeHtml(u.email || '—') + '</td>'
|
||||
+ '<td>' + roleLabel + '</td>'
|
||||
+ '<td><span class="status-badge ' + statusClass + '">' + escapeHtml(u.status) + '</span></td>'
|
||||
+ '<td>' + formatRelativeTime(u.created_at) + '</td>'
|
||||
+ '<td>' + actions + '</td>'
|
||||
+ '</tr>';
|
||||
}).join('');
|
||||
}
|
||||
|
||||
function suspendUser(userId) {
|
||||
apiFetch('/api/admin/users/' + userId + '/suspend', { method: 'POST' })
|
||||
.then(function() { loadUsers(); })
|
||||
.catch(function(e) { alert('Failed to suspend user: ' + e.message); });
|
||||
}
|
||||
|
||||
function activateUser(userId) {
|
||||
apiFetch('/api/admin/users/' + userId + '/activate', { method: 'POST' })
|
||||
.then(function() { loadUsers(); })
|
||||
.catch(function(e) { alert('Failed to activate user: ' + e.message); });
|
||||
}
|
||||
|
||||
function createTokenForUser(userId, displayName) {
|
||||
var tokenName = prompt('Token name for ' + displayName + ':', 'api-token');
|
||||
if (!tokenName) return;
|
||||
apiFetch('/api/tokens', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ name: tokenName, user_id: userId }),
|
||||
}).then(function(data) {
|
||||
showTokenBanner(data.token);
|
||||
}).catch(function(e) { alert('Failed to create token: ' + e.message); });
|
||||
}
|
||||
|
||||
function showTokenBanner(tokenValue) {
|
||||
var banner = document.getElementById('users-token-result');
|
||||
if (!banner) return;
|
||||
var loginUrl = window.location.origin + '/?token=' + encodeURIComponent(tokenValue);
|
||||
banner.style.display = 'block';
|
||||
banner.innerHTML = '<strong>User created!</strong> Share this login link — it won\'t be shown again:<br>'
|
||||
+ '<code class="token-display" id="token-copy-value">' + escapeHtml(loginUrl) + '</code>'
|
||||
+ '<button class="btn-small" id="token-copy-link">Copy Link</button>'
|
||||
+ '<br><span style="font-size:0.8em;color:var(--text-muted)">Raw token: ' + escapeHtml(tokenValue) + '</span>';
|
||||
document.getElementById('token-copy-link').addEventListener('click', function() {
|
||||
navigator.clipboard.writeText(loginUrl);
|
||||
this.textContent = 'Copied!';
|
||||
});
|
||||
}
|
||||
|
||||
// Delegated click handler for user action buttons (CSP-safe, no inline onclick)
|
||||
document.getElementById('users-table')?.addEventListener('click', function(e) {
|
||||
var btn = e.target.closest('[data-action]');
|
||||
if (!btn) return;
|
||||
var action = btn.getAttribute('data-action');
|
||||
var userId = btn.getAttribute('data-user-id');
|
||||
var userName = btn.getAttribute('data-user-name');
|
||||
if (action === 'suspend-user') suspendUser(userId);
|
||||
else if (action === 'activate-user') activateUser(userId);
|
||||
else if (action === 'create-token') createTokenForUser(userId, userName || '');
|
||||
});
|
||||
|
||||
// Wire up Users tab create form
|
||||
document.getElementById('users-create-btn')?.addEventListener('click', function() {
|
||||
document.getElementById('users-create-form').style.display = 'flex';
|
||||
document.getElementById('users-token-result').style.display = 'none';
|
||||
document.getElementById('user-display-name').focus();
|
||||
});
|
||||
|
||||
document.getElementById('users-create-cancel')?.addEventListener('click', function() {
|
||||
document.getElementById('users-create-form').style.display = 'none';
|
||||
});
|
||||
|
||||
document.getElementById('users-create-submit')?.addEventListener('click', function() {
|
||||
var displayName = document.getElementById('user-display-name').value.trim();
|
||||
var email = document.getElementById('user-email').value.trim();
|
||||
var role = document.getElementById('user-role').value;
|
||||
if (!displayName) { alert('Display name is required'); return; }
|
||||
|
||||
apiFetch('/api/admin/users', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
display_name: displayName,
|
||||
email: email || undefined,
|
||||
role: role,
|
||||
}),
|
||||
}).then(function(data) {
|
||||
document.getElementById('users-create-form').style.display = 'none';
|
||||
document.getElementById('user-display-name').value = '';
|
||||
document.getElementById('user-email').value = '';
|
||||
if (data.token) {
|
||||
showTokenBanner(data.token);
|
||||
}
|
||||
loadUsers();
|
||||
}).catch(function(e) { alert('Failed to create user: ' + e.message); });
|
||||
});
|
||||
|
||||
// --- Gateway status widget ---
|
||||
|
||||
let gatewayStatusInterval = null;
|
||||
@@ -5028,6 +5171,7 @@ function loadSettingsSubtab(subtab) {
|
||||
else if (subtab === 'extensions') { loadExtensions(); startPairingPoll(); }
|
||||
else if (subtab === 'mcp') loadMcpServers();
|
||||
else if (subtab === 'skills') loadSkills();
|
||||
else if (subtab === 'users') loadUsers();
|
||||
if (subtab !== 'extensions' && subtab !== 'channels') stopPairingPoll();
|
||||
}
|
||||
|
||||
|
||||
@@ -9,8 +9,6 @@ I18n.register('en', {
|
||||
'auth.connect': 'Connect',
|
||||
'auth.errorRequired': 'Token required',
|
||||
'auth.errorInvalid': 'Invalid token',
|
||||
'auth.hint': 'Enter the GATEWAY_AUTH_TOKEN from your .env file',
|
||||
|
||||
// Chat
|
||||
'chat.inputPlaceholder': 'Message or / for commands...',
|
||||
|
||||
@@ -44,7 +42,8 @@ I18n.register('en', {
|
||||
'settings.channels': 'Channels',
|
||||
'settings.networking': 'Networking',
|
||||
'settings.mcp': 'MCP',
|
||||
|
||||
'settings.users': 'Users',
|
||||
|
||||
// Status
|
||||
'status.connected': 'Connected',
|
||||
'status.disconnected': 'Disconnected',
|
||||
|
||||
@@ -9,8 +9,6 @@ I18n.register('zh-CN', {
|
||||
'auth.connect': '连接',
|
||||
'auth.errorRequired': '请输入令牌',
|
||||
'auth.errorInvalid': '令牌无效',
|
||||
'auth.hint': '输入 .env 配置文件中的 GATEWAY_AUTH_TOKEN',
|
||||
|
||||
// 聊天
|
||||
'chat.inputPlaceholder': '输入消息或 / 以使用命令...',
|
||||
|
||||
@@ -44,7 +42,8 @@ I18n.register('zh-CN', {
|
||||
'settings.channels': '频道',
|
||||
'settings.networking': '网络',
|
||||
'settings.mcp': 'MCP',
|
||||
|
||||
'settings.users': '用户管理',
|
||||
|
||||
// 状态
|
||||
'status.connected': '已连接',
|
||||
'status.disconnected': '已断开',
|
||||
|
||||
@@ -41,7 +41,6 @@
|
||||
<button id="auth-connect-btn" data-i18n="auth.connect">Connect</button>
|
||||
</div>
|
||||
<div id="auth-error"></div>
|
||||
<p class="auth-hint" data-i18n="auth.hint">Enter the GATEWAY_AUTH_TOKEN from your .env configuration.</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -293,6 +292,7 @@
|
||||
<button class="settings-subtab" data-settings-subtab="extensions" data-i18n="tab.extensions">Extensions</button>
|
||||
<button class="settings-subtab" data-settings-subtab="mcp" data-i18n="settings.mcp">MCP</button>
|
||||
<button class="settings-subtab" data-settings-subtab="skills" data-i18n="tab.skills">Skills</button>
|
||||
<button class="settings-subtab" data-settings-subtab="users" data-i18n="settings.users">Users</button>
|
||||
<button class="settings-theme-toggle" id="settings-theme-toggle" data-i18n="theme.tooltipSystem" title="Toggle theme">Theme</button>
|
||||
</div>
|
||||
<div class="settings-content">
|
||||
@@ -390,6 +390,29 @@
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="settings-subpanel" id="settings-users">
|
||||
<div class="users-container">
|
||||
<div class="users-header">
|
||||
<h3>User Management</h3>
|
||||
<button id="users-create-btn" class="btn-primary">+ New User</button>
|
||||
</div>
|
||||
<div id="users-create-form" style="display:none" class="users-form" autocomplete="off">
|
||||
<input type="text" id="user-display-name" placeholder="Display name" autocomplete="off" />
|
||||
<input type="text" id="user-email" placeholder="Email (optional)" autocomplete="off" />
|
||||
<select id="user-role"><option value="member">Member</option><option value="admin">Admin</option></select>
|
||||
<button id="users-create-submit" class="btn-primary">Create</button>
|
||||
<button id="users-create-cancel" class="btn-secondary">Cancel</button>
|
||||
</div>
|
||||
<div id="users-token-result" style="display:none" class="users-token-banner"></div>
|
||||
<table class="routines-table" id="users-table">
|
||||
<thead><tr>
|
||||
<th>ID</th><th>Display Name</th><th>Email</th><th>Role</th><th>Status</th><th>Created</th><th>Actions</th>
|
||||
</tr></thead>
|
||||
<tbody id="users-tbody"></tbody>
|
||||
</table>
|
||||
<div id="users-empty" class="empty-state" style="display:none">No users found. Create the first user to get started.</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -5429,3 +5429,22 @@ body.theme-transition *:not(svg):not(path):not(line):not(circle):not(rect) {
|
||||
--text-muted: #a1a1aa;
|
||||
}
|
||||
}
|
||||
|
||||
/* --- Users Tab --- */
|
||||
.users-container { padding: 1rem; }
|
||||
.users-header { display: flex; align-items: center; justify-content: space-between; margin-bottom: 1rem; }
|
||||
.users-header h3 { margin: 0; font-size: 1.1rem; }
|
||||
.users-form { display: flex; gap: 0.5rem; align-items: center; margin-bottom: 1rem; flex-wrap: wrap; }
|
||||
.users-form input, .users-form select { padding: 0.4rem 0.6rem; border-radius: 6px; border: 1px solid var(--border); background: var(--bg-secondary); color: var(--text-primary); font-size: 0.85rem; }
|
||||
.users-token-banner { background: var(--bg-tertiary); border: 1px solid var(--accent); border-radius: 8px; padding: 0.75rem 1rem; margin-bottom: 1rem; font-size: 0.85rem; }
|
||||
.token-display { display: inline-block; padding: 0.3rem 0.6rem; background: var(--bg-primary); border-radius: 4px; font-family: var(--font-mono); word-break: break-all; margin: 0.4rem 0; user-select: all; }
|
||||
.user-id { font-family: var(--font-mono); font-size: 0.8rem; color: var(--text-muted); }
|
||||
.badge { display: inline-block; padding: 0.15rem 0.5rem; border-radius: 10px; font-size: 0.75rem; background: var(--bg-tertiary); color: var(--text-secondary); }
|
||||
.badge-admin { background: var(--accent); color: #fff; }
|
||||
.btn-small { padding: 0.25rem 0.5rem; font-size: 0.75rem; border-radius: 4px; border: 1px solid var(--border); background: var(--bg-secondary); color: var(--text-primary); cursor: pointer; }
|
||||
.btn-small:hover { background: var(--bg-tertiary); }
|
||||
.btn-danger { border-color: #ef4444; color: #ef4444; }
|
||||
.btn-danger:hover { background: #ef4444; color: #fff; }
|
||||
.btn-primary { background: var(--accent); color: #fff; border: none; padding: 0.4rem 0.8rem; border-radius: 6px; cursor: pointer; font-size: 0.85rem; }
|
||||
.btn-primary:hover { opacity: 0.9; }
|
||||
.btn-secondary { background: var(--bg-tertiary); color: var(--text-primary); border: 1px solid var(--border); padding: 0.4rem 0.8rem; border-radius: 6px; cursor: pointer; font-size: 0.85rem; }
|
||||
|
||||
@@ -77,7 +77,6 @@ impl TestGatewayBuilder {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
owner_id: self.user_id.clone(),
|
||||
default_sender_id: self.user_id,
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: self.llm_provider,
|
||||
@@ -92,6 +91,7 @@ impl TestGatewayBuilder {
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: crate::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ impl TestGatewayBuilder {
|
||||
let addr: SocketAddr = "127.0.0.1:0"
|
||||
.parse()
|
||||
.expect("hard-coded address must parse"); // safety: constant literal
|
||||
let bound = start_server(addr, state.clone(), auth).await?;
|
||||
let bound = start_server(addr, state.clone(), auth.into()).await?;
|
||||
Ok((bound, state))
|
||||
}
|
||||
|
||||
@@ -120,7 +120,7 @@ impl TestGatewayBuilder {
|
||||
let addr: SocketAddr = "127.0.0.1:0"
|
||||
.parse()
|
||||
.expect("hard-coded address must parse"); // safety: constant literal
|
||||
let bound = start_server(addr, state.clone(), auth).await?;
|
||||
let bound = start_server(addr, state.clone(), auth.into()).await?;
|
||||
Ok((bound, state))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,6 @@ use axum::routing::{delete, get, post};
|
||||
use tower::ServiceExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::GatewayChannel;
|
||||
use crate::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
@@ -24,7 +23,6 @@ use crate::channels::web::server::{
|
||||
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
||||
};
|
||||
use crate::channels::web::sse::SseManager;
|
||||
use crate::config::GatewayConfig;
|
||||
|
||||
// ── Helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -35,6 +33,7 @@ fn two_user_auth() -> MultiAuthState {
|
||||
"tok-alice".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["shared".to_string()],
|
||||
},
|
||||
);
|
||||
@@ -42,6 +41,7 @@ fn two_user_auth() -> MultiAuthState {
|
||||
"tok-bob".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "bob".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["shared".to_string(), "alice".to_string()],
|
||||
},
|
||||
);
|
||||
@@ -67,7 +67,6 @@ fn build_state(
|
||||
job_manager: None,
|
||||
prompt_queue,
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
@@ -82,43 +81,10 @@ fn build_state(
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn gateway_config() -> GatewayConfig {
|
||||
GatewayConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: Some("gateway-auth".to_string()),
|
||||
user_id: "gateway-sender".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_owner_scope_updates_gateway_owner_scope_in_multi_user_mode() {
|
||||
let mut gateway = GatewayChannel::new(gateway_config());
|
||||
gateway.auth = two_user_auth();
|
||||
gateway.config.user_tokens = Some(HashMap::new());
|
||||
let gateway = gateway.with_owner_scope("owner-scope");
|
||||
|
||||
assert_eq!(gateway.state.owner_id, "owner-scope");
|
||||
assert_eq!(gateway.state.default_sender_id, "gateway-sender");
|
||||
|
||||
let alice = gateway
|
||||
.auth
|
||||
.authenticate("tok-alice")
|
||||
.expect("alice token should remain valid");
|
||||
let bob = gateway
|
||||
.auth
|
||||
.authenticate("tok-bob")
|
||||
.expect("bob token should remain valid");
|
||||
assert_eq!(alice.user_id, "alice");
|
||||
assert_eq!(bob.user_id, "bob");
|
||||
}
|
||||
|
||||
/// Create a libSQL-backed test database in a temporary directory.
|
||||
///
|
||||
/// Returns the database and a `TempDir` guard — the database file is
|
||||
@@ -225,6 +191,7 @@ mod workspace_pool {
|
||||
);
|
||||
let identity = UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
};
|
||||
let ws = pool.get_or_create(&identity).await;
|
||||
@@ -253,6 +220,7 @@ mod workspace_pool {
|
||||
);
|
||||
let identity = UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
};
|
||||
let ws = pool.get_or_create(&identity).await;
|
||||
@@ -276,6 +244,7 @@ mod workspace_pool {
|
||||
);
|
||||
let identity = UserIdentity {
|
||||
user_id: "bob".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["alice".to_string(), "shared".to_string()],
|
||||
};
|
||||
let ws = pool.get_or_create(&identity).await;
|
||||
@@ -302,10 +271,12 @@ mod workspace_pool {
|
||||
);
|
||||
let alice_id = UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
};
|
||||
let bob_id = UserIdentity {
|
||||
user_id: "bob".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
};
|
||||
|
||||
@@ -337,6 +308,7 @@ mod workspace_pool {
|
||||
);
|
||||
let identity = UserIdentity {
|
||||
user_id: "alice".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["token-scope".to_string()],
|
||||
};
|
||||
let ws = pool.get_or_create(&identity).await;
|
||||
@@ -377,7 +349,10 @@ mod jobs_isolation {
|
||||
.route("/api/jobs/{id}/cancel", post(jobs_cancel_handler))
|
||||
.route("/api/jobs/{id}/restart", post(jobs_restart_handler))
|
||||
.route("/api/jobs/{id}/prompt", post(jobs_prompt_handler))
|
||||
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||
.layer(middleware::from_fn_with_state(
|
||||
crate::channels::web::auth::CombinedAuthState::from(auth),
|
||||
auth_middleware,
|
||||
))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -583,7 +558,10 @@ mod routines_isolation {
|
||||
.route("/api/routines/{id}", get(routines_detail_handler))
|
||||
.route("/api/routines/{id}/toggle", post(routines_toggle_handler))
|
||||
.route("/api/routines/{id}", delete(routines_delete_handler))
|
||||
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||
.layer(middleware::from_fn_with_state(
|
||||
crate::channels::web::auth::CombinedAuthState::from(auth),
|
||||
auth_middleware,
|
||||
))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -708,7 +686,10 @@ mod auth_enforcement {
|
||||
.route("/api/logs/level", get(authed_handler).put(authed_handler))
|
||||
// Gateway status
|
||||
.route("/api/gateway/status", get(authed_handler))
|
||||
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||
.layer(middleware::from_fn_with_state(
|
||||
crate::channels::web::auth::CombinedAuthState::from(auth),
|
||||
auth_middleware,
|
||||
))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -831,3 +812,140 @@ mod auth_enforcement {
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Admin Endpoint Role Enforcement Tests
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
mod admin_role_enforcement {
|
||||
use super::*;
|
||||
use crate::channels::web::handlers::users::{
|
||||
users_activate_handler, users_detail_handler, users_list_handler, users_suspend_handler,
|
||||
users_update_handler,
|
||||
};
|
||||
use axum::routing::patch;
|
||||
|
||||
/// Build a router with admin user endpoints behind multi-user auth.
|
||||
/// Uses a member-role token and an admin-role token.
|
||||
fn admin_router() -> Router {
|
||||
let mut tokens = HashMap::new();
|
||||
tokens.insert(
|
||||
"tok-admin".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "admin-user".to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
},
|
||||
);
|
||||
tokens.insert(
|
||||
"tok-member".to_string(),
|
||||
UserIdentity {
|
||||
user_id: "member-user".to_string(),
|
||||
role: "member".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
},
|
||||
);
|
||||
let auth = MultiAuthState::multi(tokens);
|
||||
let state = build_state(None, None);
|
||||
|
||||
Router::new()
|
||||
.route("/api/admin/users", get(users_list_handler))
|
||||
.route("/api/admin/users/{id}", get(users_detail_handler))
|
||||
.route("/api/admin/users/{id}", patch(users_update_handler))
|
||||
.route("/api/admin/users/{id}/suspend", post(users_suspend_handler))
|
||||
.route(
|
||||
"/api/admin/users/{id}/activate",
|
||||
post(users_activate_handler),
|
||||
)
|
||||
.layer(middleware::from_fn_with_state(
|
||||
crate::channels::web::auth::CombinedAuthState::from(auth),
|
||||
auth_middleware,
|
||||
))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
/// Assert a request returns FORBIDDEN for a member token.
|
||||
async fn assert_forbidden_for_member(app: &Router, method: Method, uri: &str) {
|
||||
let req = Request::builder()
|
||||
.method(method)
|
||||
.uri(uri)
|
||||
.header("Authorization", "Bearer tok-member")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(req).await.unwrap();
|
||||
assert_eq!(
|
||||
resp.status(),
|
||||
StatusCode::FORBIDDEN,
|
||||
"expected 403 for member on {}",
|
||||
uri
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_admin_user_endpoints_reject_member_role() {
|
||||
let app = admin_router();
|
||||
|
||||
assert_forbidden_for_member(&app, Method::GET, "/api/admin/users").await;
|
||||
assert_forbidden_for_member(&app, Method::GET, "/api/admin/users/some-id").await;
|
||||
assert_forbidden_for_member(&app, Method::POST, "/api/admin/users/some-id/suspend").await;
|
||||
assert_forbidden_for_member(&app, Method::POST, "/api/admin/users/some-id/activate").await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_admin_user_endpoints_accept_admin_role() {
|
||||
let app = admin_router();
|
||||
|
||||
// Admin token should pass auth (will get 503 since no DB, but not 403).
|
||||
let req = Request::builder()
|
||||
.uri("/api/admin/users")
|
||||
.header("Authorization", "Bearer tok-admin")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(req).await.unwrap();
|
||||
assert_ne!(
|
||||
resp.status(),
|
||||
StatusCode::FORBIDDEN,
|
||||
"admin should not get 403"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// DbAuthenticator Cache Bounded Tests
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
mod db_auth_cache {
|
||||
use super::*;
|
||||
use std::time::Instant;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_bounded_by_max_entries() {
|
||||
// Access the internal cache and verify LRU eviction.
|
||||
// We can't easily test through `authenticate()` since it hits the DB,
|
||||
// so we test the LRU cache directly.
|
||||
let cap = std::num::NonZeroUsize::new(4).unwrap(); // safety: test-only, 4 is non-zero
|
||||
let cache: lru::LruCache<[u8; 32], (UserIdentity, Instant)> = lru::LruCache::new(cap);
|
||||
let cache = Arc::new(tokio::sync::RwLock::new(cache));
|
||||
|
||||
{
|
||||
let mut c = cache.write().await;
|
||||
for i in 0..10u8 {
|
||||
let mut hash = [0u8; 32];
|
||||
hash[0] = i;
|
||||
c.put(
|
||||
hash,
|
||||
(
|
||||
UserIdentity {
|
||||
user_id: format!("user-{i}"),
|
||||
role: "member".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
},
|
||||
Instant::now(),
|
||||
),
|
||||
);
|
||||
}
|
||||
// Cache must be bounded at capacity, not grown to 10.
|
||||
assert_eq!(c.len(), 4, "cache should be bounded to capacity"); // safety: test assertion
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -521,7 +521,6 @@ mod tests {
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -535,6 +534,7 @@ mod tests {
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: crate::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+2
-1
@@ -352,7 +352,8 @@ pub async fn run_routines_cli(
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{e:#}"))?;
|
||||
|
||||
let user_id = std::env::var("GATEWAY_USER_ID").unwrap_or_else(|_| "default".to_string());
|
||||
let user_id =
|
||||
std::env::var("IRONCLAW_OWNER_ID").unwrap_or_else(|_| "default".to_string());
|
||||
run_routines_command(routines_cmd.clone(), db, &user_id).await
|
||||
}
|
||||
|
||||
|
||||
@@ -23,14 +23,26 @@ pub struct AgentConfig {
|
||||
pub max_cost_per_day_cents: Option<u64>,
|
||||
/// Maximum LLM/tool actions per hour. None = unlimited.
|
||||
pub max_actions_per_hour: Option<u64>,
|
||||
/// Maximum daily LLM spend per user in cents. None = unlimited.
|
||||
pub max_cost_per_user_per_day_cents: Option<u64>,
|
||||
/// Maximum tool-call iterations per agentic loop invocation. Default 50.
|
||||
pub max_tool_iterations: usize,
|
||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||
pub auto_approve_tools: bool,
|
||||
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
|
||||
pub default_timezone: String,
|
||||
/// Maximum concurrent jobs per user. None = use global max_parallel_jobs.
|
||||
pub max_jobs_per_user: Option<usize>,
|
||||
/// Maximum tokens per job (0 = unlimited).
|
||||
pub max_tokens_per_job: u64,
|
||||
/// Whether the deployment is multi-tenant (multiple users sharing one
|
||||
/// instance). Detected at runtime after DB initialization, not from config.
|
||||
/// See app.rs startup logic.
|
||||
pub multi_tenant: bool,
|
||||
/// Maximum concurrent LLM calls per user. None = use default (4).
|
||||
pub max_llm_concurrent_per_user: Option<usize>,
|
||||
/// Maximum concurrent jobs per user. None = use default (3).
|
||||
pub max_jobs_concurrent_per_user: Option<usize>,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@@ -49,10 +61,15 @@ impl AgentConfig {
|
||||
allow_local_tools: true,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
max_tool_iterations: 10,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
max_jobs_per_user: None,
|
||||
max_tokens_per_job: 0,
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: None,
|
||||
max_jobs_concurrent_per_user: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -87,6 +104,7 @@ impl AgentConfig {
|
||||
allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?,
|
||||
max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?,
|
||||
max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?,
|
||||
max_cost_per_user_per_day_cents: parse_option_env("MAX_COST_PER_USER_PER_DAY_CENTS")?,
|
||||
max_tool_iterations: parse_optional_env(
|
||||
"AGENT_MAX_TOOL_ITERATIONS",
|
||||
settings.agent.max_tool_iterations,
|
||||
@@ -108,10 +126,16 @@ impl AgentConfig {
|
||||
}
|
||||
tz
|
||||
},
|
||||
max_jobs_per_user: parse_option_env("MAX_JOBS_PER_USER")?,
|
||||
max_tokens_per_job: parse_optional_env(
|
||||
"AGENT_MAX_TOKENS_PER_JOB",
|
||||
settings.agent.max_tokens_per_job,
|
||||
)?,
|
||||
// Multi-tenant mode is detected at runtime after DB initialization,
|
||||
// not from config. See app.rs startup logic.
|
||||
multi_tenant: false,
|
||||
max_llm_concurrent_per_user: parse_option_env("TENANT_MAX_LLM_CONCURRENT")?,
|
||||
max_jobs_concurrent_per_user: parse_option_env("TENANT_MAX_JOBS_CONCURRENT")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+2
-64
@@ -1,13 +1,11 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use secrecy::SecretString;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env};
|
||||
use crate::error::ConfigError;
|
||||
use crate::settings::Settings;
|
||||
use secrecy::SecretString;
|
||||
|
||||
/// Channel configurations.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -45,7 +43,6 @@ pub struct GatewayConfig {
|
||||
pub port: u16,
|
||||
/// Bearer token for authentication. Random hex generated at startup if unset.
|
||||
pub auth_token: Option<String>,
|
||||
pub user_id: String,
|
||||
/// Additional user scopes for workspace reads.
|
||||
///
|
||||
/// When set, the workspace will be able to read (search, read, list) from
|
||||
@@ -54,18 +51,6 @@ pub struct GatewayConfig {
|
||||
pub workspace_read_scopes: Vec<String>,
|
||||
/// Memory layer definitions (JSON in env var, or from external config).
|
||||
pub memory_layers: Vec<crate::workspace::layer::MemoryLayer>,
|
||||
/// Multi-user token map. When set, each token maps to a user identity.
|
||||
/// Parsed from `GATEWAY_USER_TOKENS` (JSON string). When absent, falls back
|
||||
/// to single-user mode via `auth_token` + `user_id`.
|
||||
pub user_tokens: Option<HashMap<String, UserTokenConfig>>,
|
||||
}
|
||||
|
||||
/// Per-user token configuration for multi-user mode.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct UserTokenConfig {
|
||||
pub user_id: String,
|
||||
#[serde(default)]
|
||||
pub workspace_read_scopes: Vec<String>,
|
||||
}
|
||||
|
||||
/// Signal channel configuration (signal-cli daemon HTTP/JSON-RPC).
|
||||
@@ -132,10 +117,6 @@ impl ChannelsConfig {
|
||||
|
||||
let gateway_enabled = parse_bool_env("GATEWAY_ENABLED", cs.gateway_enabled)?;
|
||||
let gateway = if gateway_enabled {
|
||||
let user_id = optional_env("GATEWAY_USER_ID")?
|
||||
.or_else(|| cs.gateway_user_id.clone())
|
||||
.unwrap_or_else(|| owner_id.to_string());
|
||||
|
||||
let memory_layers: Vec<crate::workspace::layer::MemoryLayer> =
|
||||
match optional_env("MEMORY_LAYERS")? {
|
||||
Some(json_str) => {
|
||||
@@ -144,7 +125,7 @@ impl ChannelsConfig {
|
||||
message: format!("must be valid JSON array of layer objects: {e}"),
|
||||
})?
|
||||
}
|
||||
None => crate::workspace::layer::MemoryLayer::default_for_user(&user_id),
|
||||
None => crate::workspace::layer::MemoryLayer::default_for_user(owner_id),
|
||||
};
|
||||
|
||||
// Validate layer names and scopes
|
||||
@@ -196,41 +177,6 @@ impl ChannelsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
let user_tokens: Option<HashMap<String, UserTokenConfig>> =
|
||||
match optional_env("GATEWAY_USER_TOKENS")? {
|
||||
Some(json_str) => {
|
||||
let tokens: HashMap<String, UserTokenConfig> = serde_json::from_str(
|
||||
&json_str,
|
||||
)
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||
message: format!(
|
||||
"must be valid JSON object mapping tokens to user configs: {e}"
|
||||
),
|
||||
})?;
|
||||
if tokens.is_empty() {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||
message:
|
||||
"token map is empty — remove the variable to use single-user mode"
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
for (tok, cfg) in &tokens {
|
||||
if cfg.user_id.trim().is_empty() {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "GATEWAY_USER_TOKENS".to_string(),
|
||||
message: format!(
|
||||
"token '{}...' has an empty user_id",
|
||||
&tok[..tok.len().min(8)]
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
Some(tokens)
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
let workspace_read_scopes: Vec<String> = optional_env("WORKSPACE_READ_SCOPES")?
|
||||
.map(|s| {
|
||||
s.split(',')
|
||||
@@ -258,10 +204,8 @@ impl ChannelsConfig {
|
||||
)?,
|
||||
auth_token: optional_env("GATEWAY_AUTH_TOKEN")?
|
||||
.or_else(|| cs.gateway_auth_token.clone()),
|
||||
user_id,
|
||||
workspace_read_scopes,
|
||||
memory_layers,
|
||||
user_tokens,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
@@ -416,15 +360,12 @@ mod tests {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: Some("tok-abc".to_string()),
|
||||
user_id: "default".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
memory_layers: vec![],
|
||||
user_tokens: None,
|
||||
};
|
||||
assert_eq!(cfg.host, "127.0.0.1");
|
||||
assert_eq!(cfg.port, 3000);
|
||||
assert_eq!(cfg.auth_token.as_deref(), Some("tok-abc"));
|
||||
assert_eq!(cfg.user_id, "default");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -433,10 +374,8 @@ mod tests {
|
||||
host: "0.0.0.0".to_string(),
|
||||
port: 3001,
|
||||
auth_token: None,
|
||||
user_id: "anon".to_string(),
|
||||
workspace_read_scopes: vec![],
|
||||
memory_layers: vec![],
|
||||
user_tokens: None,
|
||||
};
|
||||
assert!(cfg.auth_token.is_none());
|
||||
}
|
||||
@@ -563,7 +502,6 @@ mod tests {
|
||||
assert_eq!(gateway.host, "127.0.0.3");
|
||||
assert_eq!(gateway.port, 9191);
|
||||
assert_eq!(gateway.auth_token.as_deref(), Some("tok"));
|
||||
assert_eq!(gateway.user_id, "owner-scope");
|
||||
|
||||
let signal = cfg.signal.expect("signal config");
|
||||
assert_eq!(signal.account, "+15551234567");
|
||||
|
||||
@@ -21,6 +21,9 @@ pub struct HeartbeatConfig {
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
/// When true, cycle through all users with routines. Set explicitly via
|
||||
/// HEARTBEAT_MULTI_TENANT or detected at runtime after DB initialization.
|
||||
pub multi_tenant: bool,
|
||||
}
|
||||
|
||||
impl Default for HeartbeatConfig {
|
||||
@@ -34,6 +37,7 @@ impl Default for HeartbeatConfig {
|
||||
quiet_hours_start: None,
|
||||
quiet_hours_end: None,
|
||||
timezone: None,
|
||||
multi_tenant: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -101,6 +105,7 @@ impl HeartbeatConfig {
|
||||
}
|
||||
tz
|
||||
},
|
||||
multi_tenant: parse_bool_env("HEARTBEAT_MULTI_TENANT", false)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ mod routines;
|
||||
mod sandbox;
|
||||
mod settings;
|
||||
mod tool_failures;
|
||||
mod users;
|
||||
mod workspace;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
@@ -530,10 +530,24 @@ impl RoutineStore for LibSqlBackend {
|
||||
async fn get_webhook_routine_by_path(
|
||||
&self,
|
||||
path: &str,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
let mut rows = if let Some(uid) = user_id {
|
||||
conn.query(
|
||||
&format!(
|
||||
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
|
||||
AND user_id = ?2 \
|
||||
AND (json_extract(trigger_config, '$.path') = ?1 \
|
||||
OR (json_extract(trigger_config, '$.path') IS NULL AND CAST(id AS TEXT) = ?1))",
|
||||
ROUTINE_COLUMNS
|
||||
),
|
||||
params![path, uid],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
} else {
|
||||
conn.query(
|
||||
&format!(
|
||||
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
|
||||
AND (json_extract(trigger_config, '$.path') = ?1 \
|
||||
@@ -543,7 +557,8 @@ impl RoutineStore for LibSqlBackend {
|
||||
params![path],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
};
|
||||
|
||||
match rows
|
||||
.next()
|
||||
|
||||
@@ -0,0 +1,735 @@
|
||||
//! UserStore implementation for LibSqlBackend.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use libsql::params;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{fmt_opt_ts, fmt_ts, get_opt_text, get_opt_ts, get_text, get_ts, opt_text};
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
use crate::db::{ApiTokenRecord, DatabaseError, UserRecord, UserStore};
|
||||
|
||||
fn row_to_user(row: &libsql::Row) -> Result<UserRecord, DatabaseError> {
|
||||
let metadata_str = get_text(row, 9);
|
||||
let metadata: serde_json::Value = serde_json::from_str(&metadata_str)
|
||||
.map_err(|e| DatabaseError::Serialization(e.to_string()))?;
|
||||
Ok(UserRecord {
|
||||
id: get_text(row, 0),
|
||||
email: get_opt_text(row, 1),
|
||||
display_name: get_text(row, 2),
|
||||
status: get_text(row, 3),
|
||||
role: get_text(row, 4),
|
||||
created_at: get_ts(row, 5),
|
||||
updated_at: get_ts(row, 6),
|
||||
last_login_at: get_opt_ts(row, 7),
|
||||
created_by: get_opt_text(row, 8),
|
||||
metadata,
|
||||
})
|
||||
}
|
||||
|
||||
fn row_to_api_token(row: &libsql::Row) -> Result<ApiTokenRecord, DatabaseError> {
|
||||
let id_str = get_text(row, 0);
|
||||
let id: Uuid = id_str
|
||||
.parse()
|
||||
.map_err(|e| DatabaseError::Serialization(format!("invalid UUID: {e}")))?;
|
||||
Ok(ApiTokenRecord {
|
||||
id,
|
||||
user_id: get_text(row, 1),
|
||||
name: get_text(row, 2),
|
||||
token_prefix: get_text(row, 3),
|
||||
expires_at: get_opt_ts(row, 4),
|
||||
last_used_at: get_opt_ts(row, 5),
|
||||
created_at: get_ts(row, 6),
|
||||
revoked_at: get_opt_ts(row, 7),
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UserStore for LibSqlBackend {
|
||||
async fn create_user(&self, user: &UserRecord) -> Result<(), DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let metadata_json = serde_json::to_string(&user.metadata)
|
||||
.map_err(|e| DatabaseError::Serialization(e.to_string()))?;
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO users (id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
|
||||
"#,
|
||||
params![
|
||||
user.id.as_str(),
|
||||
opt_text(user.email.as_deref()),
|
||||
user.display_name.as_str(),
|
||||
user.status.as_str(),
|
||||
user.role.as_str(),
|
||||
fmt_ts(&user.created_at),
|
||||
fmt_ts(&user.updated_at),
|
||||
fmt_opt_ts(&user.last_login_at),
|
||||
opt_text(user.created_by.as_deref()),
|
||||
metadata_json,
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_user(&self, id: &str) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, email, display_name, status, role, created_at, updated_at,
|
||||
last_login_at, created_by, metadata
|
||||
FROM users WHERE id = ?1
|
||||
"#,
|
||||
params![id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
Some(row) => Ok(Some(row_to_user(&row)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_user_by_email(&self, email: &str) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, email, display_name, status, role, created_at, updated_at,
|
||||
last_login_at, created_by, metadata
|
||||
FROM users WHERE email = ?1
|
||||
"#,
|
||||
params![email],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
Some(row) => Ok(Some(row_to_user(&row)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_users(&self, status: Option<&str>) -> Result<Vec<UserRecord>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut users = Vec::new();
|
||||
|
||||
let mut rows = if let Some(status) = status {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, email, display_name, status, role, created_at, updated_at,
|
||||
last_login_at, created_by, metadata
|
||||
FROM users WHERE status = ?1
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
params![status],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
} else {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, email, display_name, status, role, created_at, updated_at,
|
||||
last_login_at, created_by, metadata
|
||||
FROM users
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
};
|
||||
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
users.push(row_to_user(&row)?);
|
||||
}
|
||||
Ok(users)
|
||||
}
|
||||
|
||||
async fn update_user_status(&self, id: &str, status: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
conn.execute(
|
||||
"UPDATE users SET status = ?2, updated_at = ?3 WHERE id = ?1",
|
||||
params![id, status, now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn update_user_profile(
|
||||
&self,
|
||||
id: &str,
|
||||
display_name: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
let metadata_json = serde_json::to_string(metadata)
|
||||
.map_err(|e| DatabaseError::Serialization(e.to_string()))?;
|
||||
conn.execute(
|
||||
"UPDATE users SET display_name = ?2, metadata = ?3, updated_at = ?4 WHERE id = ?1",
|
||||
params![id, display_name, metadata_json, now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn record_login(&self, id: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
conn.execute(
|
||||
"UPDATE users SET last_login_at = ?2, updated_at = ?2 WHERE id = ?1",
|
||||
params![id, now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_api_token(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
token_hash: &[u8; 32],
|
||||
token_prefix: &str,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
) -> Result<ApiTokenRecord, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO api_tokens (id, user_id, token_hash, token_prefix, name, expires_at, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
|
||||
"#,
|
||||
params![
|
||||
id.to_string(),
|
||||
user_id,
|
||||
libsql::Value::Blob(token_hash.to_vec()),
|
||||
token_prefix,
|
||||
name,
|
||||
fmt_opt_ts(&expires_at),
|
||||
fmt_ts(&now),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
Ok(ApiTokenRecord {
|
||||
id,
|
||||
user_id: user_id.to_string(),
|
||||
name: name.to_string(),
|
||||
token_prefix: token_prefix.to_string(),
|
||||
expires_at,
|
||||
last_used_at: None,
|
||||
created_at: now,
|
||||
revoked_at: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_api_tokens(&self, user_id: &str) -> Result<Vec<ApiTokenRecord>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, token_prefix, expires_at, last_used_at, created_at, revoked_at
|
||||
FROM api_tokens WHERE user_id = ?1
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
params![user_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
let mut tokens = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
tokens.push(row_to_api_token(&row)?);
|
||||
}
|
||||
Ok(tokens)
|
||||
}
|
||||
|
||||
async fn revoke_api_token(&self, token_id: Uuid, user_id: &str) -> Result<bool, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
let rows_affected = conn
|
||||
.execute(
|
||||
r#"
|
||||
UPDATE api_tokens SET revoked_at = ?3
|
||||
WHERE id = ?1 AND user_id = ?2 AND revoked_at IS NULL
|
||||
"#,
|
||||
params![token_id.to_string(), user_id, now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
async fn authenticate_token(
|
||||
&self,
|
||||
token_hash: &[u8; 32],
|
||||
) -> Result<Option<(ApiTokenRecord, UserRecord)>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT
|
||||
t.id, t.user_id, t.name, t.token_prefix, t.expires_at,
|
||||
t.last_used_at, t.created_at, t.revoked_at,
|
||||
u.id, u.email, u.display_name, u.status, u.role, u.created_at,
|
||||
u.updated_at, u.last_login_at, u.created_by, u.metadata
|
||||
FROM api_tokens t
|
||||
JOIN users u ON u.id = t.user_id
|
||||
WHERE t.token_hash = ?1
|
||||
AND t.revoked_at IS NULL
|
||||
AND (t.expires_at IS NULL OR t.expires_at > ?2)
|
||||
AND u.status = 'active'
|
||||
"#,
|
||||
params![libsql::Value::Blob(token_hash.to_vec()), now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let id_str = get_text(&row, 0);
|
||||
let token_id: Uuid = id_str
|
||||
.parse()
|
||||
.map_err(|e| DatabaseError::Serialization(format!("invalid UUID: {e}")))?;
|
||||
let token = ApiTokenRecord {
|
||||
id: token_id,
|
||||
user_id: get_text(&row, 1),
|
||||
name: get_text(&row, 2),
|
||||
token_prefix: get_text(&row, 3),
|
||||
expires_at: get_opt_ts(&row, 4),
|
||||
last_used_at: get_opt_ts(&row, 5),
|
||||
created_at: get_ts(&row, 6),
|
||||
revoked_at: get_opt_ts(&row, 7),
|
||||
};
|
||||
|
||||
let metadata_str = get_text(&row, 17);
|
||||
let metadata: serde_json::Value = serde_json::from_str(&metadata_str)
|
||||
.map_err(|e| DatabaseError::Serialization(e.to_string()))?;
|
||||
|
||||
let user = UserRecord {
|
||||
id: get_text(&row, 8),
|
||||
email: get_opt_text(&row, 9),
|
||||
display_name: get_text(&row, 10),
|
||||
status: get_text(&row, 11),
|
||||
role: get_text(&row, 12),
|
||||
created_at: get_ts(&row, 13),
|
||||
updated_at: get_ts(&row, 14),
|
||||
last_login_at: get_opt_ts(&row, 15),
|
||||
created_by: get_opt_text(&row, 16),
|
||||
metadata,
|
||||
};
|
||||
|
||||
Ok(Some((token, user)))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_token_usage(&self, token_id: Uuid) -> Result<(), DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
conn.execute(
|
||||
"UPDATE api_tokens SET last_used_at = ?2 WHERE id = ?1",
|
||||
params![token_id.to_string(), now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn has_any_users(&self) -> Result<bool, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query("SELECT 1 FROM users LIMIT 1", ())
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
let has_users = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
.is_some();
|
||||
Ok(has_users)
|
||||
}
|
||||
|
||||
async fn delete_user(&self, id: &str) -> Result<bool, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
// Delete from child tables first to avoid FK violations.
|
||||
// agent_jobs cascades to job_actions, llm_calls, estimation_snapshots
|
||||
// conversations cascades to conversation_messages
|
||||
// memory_documents cascades to memory_chunks
|
||||
// routines cascades to routine_runs
|
||||
for table in &[
|
||||
"settings",
|
||||
"heartbeat_state",
|
||||
"tool_rate_limit_state",
|
||||
"secret_usage_log",
|
||||
"leak_detection_events",
|
||||
"secrets",
|
||||
"wasm_tools",
|
||||
"routines",
|
||||
"memory_documents",
|
||||
"conversations",
|
||||
"api_tokens",
|
||||
] {
|
||||
conn.execute(
|
||||
&format!("DELETE FROM {} WHERE user_id = ?1", table),
|
||||
params![id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
}
|
||||
// job_events references agent_jobs(id) without CASCADE — delete via subquery.
|
||||
conn.execute(
|
||||
"DELETE FROM job_events WHERE job_id IN (SELECT id FROM agent_jobs WHERE user_id = ?1)",
|
||||
params![id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
conn.execute("DELETE FROM agent_jobs WHERE user_id = ?1", params![id])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
// Nullify self-referencing created_by before deleting the user
|
||||
conn.execute(
|
||||
"UPDATE users SET created_by = NULL WHERE created_by = ?1",
|
||||
params![id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
let rows = conn
|
||||
.execute("DELETE FROM users WHERE id = ?1", params![id])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(rows > 0)
|
||||
}
|
||||
|
||||
async fn user_usage_stats(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
since: DateTime<Utc>,
|
||||
) -> Result<Vec<crate::db::UserUsageStats>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let since_str = fmt_ts(&since);
|
||||
let mut rows = if let Some(uid) = user_id {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT j.user_id, l.model, COUNT(*) as call_count,
|
||||
COALESCE(SUM(l.input_tokens), 0) as input_tokens,
|
||||
COALESCE(SUM(l.output_tokens), 0) as output_tokens,
|
||||
COALESCE(SUM(l.cost), 0) as total_cost
|
||||
FROM llm_calls l
|
||||
JOIN agent_jobs j ON l.job_id = j.id
|
||||
WHERE l.created_at >= ?1
|
||||
AND j.user_id = ?2
|
||||
GROUP BY j.user_id, l.model
|
||||
ORDER BY total_cost DESC
|
||||
"#,
|
||||
params![since_str, uid],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
} else {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT j.user_id, l.model, COUNT(*) as call_count,
|
||||
COALESCE(SUM(l.input_tokens), 0) as input_tokens,
|
||||
COALESCE(SUM(l.output_tokens), 0) as output_tokens,
|
||||
COALESCE(SUM(l.cost), 0) as total_cost
|
||||
FROM llm_calls l
|
||||
JOIN agent_jobs j ON l.job_id = j.id
|
||||
WHERE l.created_at >= ?1
|
||||
GROUP BY j.user_id, l.model
|
||||
ORDER BY total_cost DESC
|
||||
"#,
|
||||
params![since_str],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
};
|
||||
let mut stats = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
let cost_str = get_text(&row, 5);
|
||||
let total_cost = rust_decimal::Decimal::from_str_exact(&cost_str).unwrap_or_default();
|
||||
stats.push(crate::db::UserUsageStats {
|
||||
user_id: get_text(&row, 0),
|
||||
model: get_text(&row, 1),
|
||||
call_count: row
|
||||
.get::<i64>(2)
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?,
|
||||
input_tokens: row
|
||||
.get::<i64>(3)
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?,
|
||||
output_tokens: row
|
||||
.get::<i64>(4)
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?,
|
||||
total_cost,
|
||||
});
|
||||
}
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
use crate::db::{Database, UserStore};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
fn hash(s: &str) -> [u8; 32] {
|
||||
let mut h = Sha256::new();
|
||||
h.update(s.as_bytes());
|
||||
h.finalize().into()
|
||||
}
|
||||
|
||||
async fn setup() -> (LibSqlBackend, tempfile::TempDir) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir.path().join("test_users.db");
|
||||
let db = LibSqlBackend::new_local(&db_path).await.unwrap();
|
||||
db.run_migrations().await.unwrap();
|
||||
(db, dir) // keep dir alive so the DB file isn't deleted
|
||||
}
|
||||
|
||||
fn test_user(id: &str) -> UserRecord {
|
||||
UserRecord {
|
||||
id: id.to_string(),
|
||||
email: Some(format!("{}@test.com", id)),
|
||||
display_name: id.to_string(),
|
||||
status: "active".to_string(),
|
||||
role: "member".to_string(),
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
last_login_at: None,
|
||||
created_by: None,
|
||||
metadata: serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_has_any_users_empty() {
|
||||
let (db, _dir) = setup().await;
|
||||
assert!(!db.has_any_users().await.unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_and_get_user() {
|
||||
let (db, _dir) = setup().await;
|
||||
let user = test_user("alice");
|
||||
db.create_user(&user).await.unwrap();
|
||||
|
||||
assert!(db.has_any_users().await.unwrap());
|
||||
|
||||
let found = db.get_user("alice").await.unwrap().unwrap();
|
||||
assert_eq!(found.id, "alice");
|
||||
assert_eq!(found.email, Some("[email protected]".to_string()));
|
||||
assert_eq!(found.status, "active");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_user_by_email() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("bob")).await.unwrap();
|
||||
|
||||
let found = db.get_user_by_email("[email protected]").await.unwrap();
|
||||
assert!(found.is_some());
|
||||
assert_eq!(found.unwrap().id, "bob");
|
||||
|
||||
assert!(
|
||||
db.get_user_by_email("[email protected]")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_users_with_status_filter() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
db.create_user(&test_user("bob")).await.unwrap();
|
||||
db.update_user_status("bob", "suspended").await.unwrap();
|
||||
|
||||
let all = db.list_users(None).await.unwrap();
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let active = db.list_users(Some("active")).await.unwrap();
|
||||
assert_eq!(active.len(), 1);
|
||||
assert_eq!(active[0].id, "alice");
|
||||
|
||||
let suspended = db.list_users(Some("suspended")).await.unwrap();
|
||||
assert_eq!(suspended.len(), 1);
|
||||
assert_eq!(suspended[0].id, "bob");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_user_profile() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
|
||||
let meta = serde_json::json!({"role": "admin"});
|
||||
db.update_user_profile("alice", "Alice Smith", &meta)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let user = db.get_user("alice").await.unwrap().unwrap();
|
||||
assert_eq!(user.display_name, "Alice Smith");
|
||||
assert_eq!(user.metadata["role"], "admin");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_lifecycle_create_authenticate_revoke() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
|
||||
// Create token
|
||||
let token_hash = hash("secret-token-123");
|
||||
let record = db
|
||||
.create_api_token("alice", "laptop", &token_hash, "secret-t", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(record.user_id, "alice");
|
||||
assert_eq!(record.name, "laptop");
|
||||
assert_eq!(record.token_prefix, "secret-t");
|
||||
|
||||
// Authenticate
|
||||
let (tok, user) = db.authenticate_token(&token_hash).await.unwrap().unwrap();
|
||||
assert_eq!(tok.id, record.id);
|
||||
assert_eq!(user.id, "alice");
|
||||
|
||||
// List tokens
|
||||
let tokens = db.list_api_tokens("alice").await.unwrap();
|
||||
assert_eq!(tokens.len(), 1);
|
||||
|
||||
// Revoke
|
||||
assert!(db.revoke_api_token(record.id, "alice").await.unwrap());
|
||||
|
||||
// Auth should fail after revoke
|
||||
assert!(db.authenticate_token(&token_hash).await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_auth_fails_for_suspended_user() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
|
||||
let token_hash = hash("token-abc");
|
||||
db.create_api_token("alice", "test", &token_hash, "token-ab", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Auth works while active
|
||||
assert!(db.authenticate_token(&token_hash).await.unwrap().is_some());
|
||||
|
||||
// Suspend user
|
||||
db.update_user_status("alice", "suspended").await.unwrap();
|
||||
|
||||
// Auth should fail
|
||||
assert!(db.authenticate_token(&token_hash).await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_revoke_wrong_user_returns_false() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
db.create_user(&test_user("bob")).await.unwrap();
|
||||
|
||||
let token_hash = hash("alice-token");
|
||||
let record = db
|
||||
.create_api_token("alice", "test", &token_hash, "alice-to", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Bob can't revoke Alice's token
|
||||
assert!(!db.revoke_api_token(record.id, "bob").await.unwrap());
|
||||
|
||||
// Alice can
|
||||
assert!(db.revoke_api_token(record.id, "alice").await.unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_record_login_and_token_usage() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
|
||||
let token_hash = hash("tok");
|
||||
let record = db
|
||||
.create_api_token("alice", "test", &token_hash, "tok", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Record usage
|
||||
db.record_token_usage(record.id).await.unwrap();
|
||||
db.record_login("alice").await.unwrap();
|
||||
|
||||
// Verify timestamps updated
|
||||
let user = db.get_user("alice").await.unwrap().unwrap();
|
||||
assert!(user.last_login_at.is_some());
|
||||
|
||||
let tokens = db.list_api_tokens("alice").await.unwrap();
|
||||
assert!(tokens[0].last_used_at.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_delete_user_removes_api_tokens() {
|
||||
let (db, _dir) = setup().await;
|
||||
db.create_user(&test_user("alice")).await.unwrap();
|
||||
|
||||
let token_hash = hash("alice-tok");
|
||||
db.create_api_token("alice", "primary", &token_hash, "alice-to", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify token exists before deletion.
|
||||
let tokens = db.list_api_tokens("alice").await.unwrap();
|
||||
assert_eq!(tokens.len(), 1);
|
||||
|
||||
// Delete user — should also remove their api_tokens.
|
||||
assert!(db.delete_user("alice").await.unwrap());
|
||||
|
||||
// api_tokens must be gone (not orphaned).
|
||||
let tokens = db.list_api_tokens("alice").await.unwrap();
|
||||
assert!(
|
||||
tokens.is_empty(),
|
||||
"expected api_tokens to be deleted with user, found {}",
|
||||
tokens.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -579,6 +579,36 @@ INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, acti
|
||||
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', '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, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'));
|
||||
|
||||
|
||||
-- ==================== User management (V14) ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id TEXT PRIMARY KEY,
|
||||
email TEXT UNIQUE,
|
||||
display_name TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
role TEXT NOT NULL DEFAULT 'member',
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
last_login_at TEXT,
|
||||
created_by TEXT,
|
||||
metadata TEXT NOT NULL DEFAULT '{}'
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS api_tokens (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
token_hash BLOB NOT NULL,
|
||||
token_prefix TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
expires_at TEXT,
|
||||
last_used_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
revoked_at TEXT
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_user ON api_tokens(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_hash ON api_tokens(token_hash);
|
||||
|
||||
"#;
|
||||
|
||||
/// Incremental migrations applied after the base schema.
|
||||
@@ -723,6 +753,38 @@ CREATE INDEX IF NOT EXISTS idx_routines_event_triggers
|
||||
WHERE enabled = 1 AND trigger_type IN ('event', 'system_event');
|
||||
|
||||
PRAGMA foreign_keys=ON;
|
||||
"#,
|
||||
),
|
||||
(
|
||||
14,
|
||||
"users",
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id TEXT PRIMARY KEY,
|
||||
email TEXT UNIQUE,
|
||||
display_name TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
role TEXT NOT NULL DEFAULT 'member',
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
last_login_at TEXT,
|
||||
created_by TEXT,
|
||||
metadata TEXT NOT NULL DEFAULT '{}'
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS api_tokens (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
token_hash BLOB NOT NULL,
|
||||
token_prefix TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
expires_at TEXT,
|
||||
last_used_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
revoked_at TEXT
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_user ON api_tokens(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_hash ON api_tokens(token_hash);
|
||||
"#,
|
||||
),
|
||||
];
|
||||
|
||||
+114
@@ -309,6 +309,43 @@ async fn validate_postgres(pool: &deadpool_postgres::Pool) -> Result<(), Databas
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ==================== User management record types ====================
|
||||
|
||||
/// A registered user.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UserRecord {
|
||||
/// User identifier (string, matches existing `user_id` throughout the codebase).
|
||||
pub id: String,
|
||||
pub email: Option<String>,
|
||||
pub display_name: String,
|
||||
/// `active`, `suspended`, or `deactivated`.
|
||||
pub status: String,
|
||||
/// `admin` or `member`.
|
||||
pub role: String,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub last_login_at: Option<DateTime<Utc>>,
|
||||
/// Who created/invited this user (nullable for bootstrap users).
|
||||
pub created_by: Option<String>,
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
/// An API token for authenticating requests (hash stored, never plaintext).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ApiTokenRecord {
|
||||
pub id: Uuid,
|
||||
pub user_id: String,
|
||||
/// Human label (e.g. "my-laptop", "ci-bot").
|
||||
pub name: String,
|
||||
/// First 8 hex chars of the plaintext token for display/identification.
|
||||
pub token_prefix: String,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// Soft-revoke timestamp. Non-null means revoked.
|
||||
pub revoked_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
// ==================== Sub-traits ====================
|
||||
//
|
||||
// Each sub-trait groups related persistence methods. The `Database` supertrait
|
||||
@@ -545,6 +582,7 @@ pub trait RoutineStore: Send + Sync {
|
||||
async fn get_webhook_routine_by_path(
|
||||
&self,
|
||||
path: &str,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Option<Routine>, DatabaseError>;
|
||||
|
||||
/// List routine runs that were dispatched as full_job but have not yet
|
||||
@@ -760,6 +798,81 @@ pub trait WorkspaceStore: Send + Sync {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UserStore: Send + Sync {
|
||||
// ---- Users ----
|
||||
|
||||
/// Create a new user record.
|
||||
async fn create_user(&self, user: &UserRecord) -> Result<(), DatabaseError>;
|
||||
/// Get a user by their string id.
|
||||
async fn get_user(&self, id: &str) -> Result<Option<UserRecord>, DatabaseError>;
|
||||
/// Get a user by email address.
|
||||
async fn get_user_by_email(&self, email: &str) -> Result<Option<UserRecord>, DatabaseError>;
|
||||
/// List users, optionally filtered by status.
|
||||
async fn list_users(&self, status: Option<&str>) -> Result<Vec<UserRecord>, DatabaseError>;
|
||||
/// Update a user's status (active/suspended/deactivated).
|
||||
async fn update_user_status(&self, id: &str, status: &str) -> Result<(), DatabaseError>;
|
||||
/// Update a user's display name and metadata.
|
||||
async fn update_user_profile(
|
||||
&self,
|
||||
id: &str,
|
||||
display_name: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError>;
|
||||
/// Record a login timestamp.
|
||||
async fn record_login(&self, id: &str) -> Result<(), DatabaseError>;
|
||||
|
||||
// ---- API Tokens ----
|
||||
|
||||
/// Create a new API token. The `token_hash` is SHA-256 of the plaintext.
|
||||
async fn create_api_token(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
token_hash: &[u8; 32],
|
||||
token_prefix: &str,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
) -> Result<ApiTokenRecord, DatabaseError>;
|
||||
/// List tokens for a user (never includes the hash).
|
||||
async fn list_api_tokens(&self, user_id: &str) -> Result<Vec<ApiTokenRecord>, DatabaseError>;
|
||||
/// Soft-revoke a token. Returns false if the token doesn't exist or doesn't belong to the user.
|
||||
async fn revoke_api_token(&self, token_id: Uuid, user_id: &str) -> Result<bool, DatabaseError>;
|
||||
/// Look up a token by hash, returning the token record and its owning user.
|
||||
/// Only returns active (non-revoked, non-expired) tokens for active users.
|
||||
async fn authenticate_token(
|
||||
&self,
|
||||
token_hash: &[u8; 32],
|
||||
) -> Result<Option<(ApiTokenRecord, UserRecord)>, DatabaseError>;
|
||||
/// Update `last_used_at` for a token.
|
||||
async fn record_token_usage(&self, token_id: Uuid) -> Result<(), DatabaseError>;
|
||||
|
||||
/// Check whether any user records exist (for first-run bootstrap detection).
|
||||
async fn has_any_users(&self) -> Result<bool, DatabaseError>;
|
||||
|
||||
/// Delete a user and all their data across all user-scoped tables.
|
||||
/// Returns false if the user doesn't exist.
|
||||
async fn delete_user(&self, id: &str) -> Result<bool, DatabaseError>;
|
||||
|
||||
/// Get per-user LLM usage stats for a time period.
|
||||
/// Aggregates from llm_calls via agent_jobs.user_id.
|
||||
async fn user_usage_stats(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
since: DateTime<Utc>,
|
||||
) -> Result<Vec<UserUsageStats>, DatabaseError>;
|
||||
}
|
||||
|
||||
/// Per-user LLM usage statistics.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UserUsageStats {
|
||||
pub user_id: String,
|
||||
pub model: String,
|
||||
pub call_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub total_cost: Decimal,
|
||||
}
|
||||
|
||||
/// Backend-agnostic database supertrait.
|
||||
///
|
||||
/// Combines all sub-traits into one. Existing `Arc<dyn Database>` consumers
|
||||
@@ -773,6 +886,7 @@ pub trait Database:
|
||||
+ ToolFailureStore
|
||||
+ SettingsStore
|
||||
+ WorkspaceStore
|
||||
+ UserStore
|
||||
+ Send
|
||||
+ Sync
|
||||
{
|
||||
|
||||
+92
-3
@@ -16,8 +16,8 @@ use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||
use crate::config::DatabaseConfig;
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
use crate::db::{
|
||||
ConversationStore, Database, JobStore, RoutineStore, SandboxStore, SettingsStore,
|
||||
ToolFailureStore, WorkspaceStore,
|
||||
ApiTokenRecord, ConversationStore, Database, JobStore, RoutineStore, SandboxStore,
|
||||
SettingsStore, ToolFailureStore, UserRecord, UserStore, WorkspaceStore,
|
||||
};
|
||||
use crate::error::{DatabaseError, WorkspaceError};
|
||||
use crate::history::{
|
||||
@@ -529,8 +529,9 @@ impl RoutineStore for PgBackend {
|
||||
async fn get_webhook_routine_by_path(
|
||||
&self,
|
||||
path: &str,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.store.get_webhook_routine_by_path(path).await
|
||||
self.store.get_webhook_routine_by_path(path, user_id).await
|
||||
}
|
||||
|
||||
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
@@ -785,3 +786,91 @@ impl WorkspaceStore for PgBackend {
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== UserStore ====================
|
||||
|
||||
#[async_trait]
|
||||
impl UserStore for PgBackend {
|
||||
async fn create_user(&self, user: &UserRecord) -> Result<(), DatabaseError> {
|
||||
self.store.create_user(user).await
|
||||
}
|
||||
|
||||
async fn get_user(&self, id: &str) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
self.store.get_user(id).await
|
||||
}
|
||||
|
||||
async fn get_user_by_email(&self, email: &str) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
self.store.get_user_by_email(email).await
|
||||
}
|
||||
|
||||
async fn list_users(&self, status: Option<&str>) -> Result<Vec<UserRecord>, DatabaseError> {
|
||||
self.store.list_users(status).await
|
||||
}
|
||||
|
||||
async fn update_user_status(&self, id: &str, status: &str) -> Result<(), DatabaseError> {
|
||||
self.store.update_user_status(id, status).await
|
||||
}
|
||||
|
||||
async fn update_user_profile(
|
||||
&self,
|
||||
id: &str,
|
||||
display_name: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_user_profile(id, display_name, metadata)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn record_login(&self, id: &str) -> Result<(), DatabaseError> {
|
||||
self.store.record_login(id).await
|
||||
}
|
||||
|
||||
async fn create_api_token(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
token_hash: &[u8; 32],
|
||||
token_prefix: &str,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
) -> Result<ApiTokenRecord, DatabaseError> {
|
||||
self.store
|
||||
.create_api_token(user_id, name, token_hash, token_prefix, expires_at)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_api_tokens(&self, user_id: &str) -> Result<Vec<ApiTokenRecord>, DatabaseError> {
|
||||
self.store.list_api_tokens(user_id).await
|
||||
}
|
||||
|
||||
async fn revoke_api_token(&self, token_id: Uuid, user_id: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.revoke_api_token(token_id, user_id).await
|
||||
}
|
||||
|
||||
async fn authenticate_token(
|
||||
&self,
|
||||
token_hash: &[u8; 32],
|
||||
) -> Result<Option<(ApiTokenRecord, UserRecord)>, DatabaseError> {
|
||||
self.store.authenticate_token(token_hash).await
|
||||
}
|
||||
|
||||
async fn record_token_usage(&self, token_id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.store.record_token_usage(token_id).await
|
||||
}
|
||||
|
||||
async fn has_any_users(&self) -> Result<bool, DatabaseError> {
|
||||
self.store.has_any_users().await
|
||||
}
|
||||
|
||||
async fn delete_user(&self, id: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.delete_user(id).await
|
||||
}
|
||||
|
||||
async fn user_usage_stats(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
since: DateTime<Utc>,
|
||||
) -> Result<Vec<crate::db::UserUsageStats>, DatabaseError> {
|
||||
self.store.user_usage_stats(user_id, since).await
|
||||
}
|
||||
}
|
||||
|
||||
+18
-7
@@ -2,7 +2,8 @@
|
||||
//!
|
||||
//! Builds a [`deadpool_postgres::Pool`] with the appropriate TLS connector
|
||||
//! based on the configured [`SslMode`]. Uses `rustls` with system root
|
||||
//! certificates — the same TLS stack that `reqwest` already uses for HTTP.
|
||||
//! certificates, falling back to Mozilla's bundled roots via `webpki-roots`
|
||||
//! when the system store is empty (common in minimal container images).
|
||||
|
||||
use deadpool_postgres::{Pool, Runtime};
|
||||
use thiserror::Error;
|
||||
@@ -19,9 +20,15 @@ pub enum CreatePoolError {
|
||||
TlsConfig(#[from] rustls::Error),
|
||||
}
|
||||
|
||||
/// Build a rustls-based TLS connector using the platform's root certificate store.
|
||||
/// Build a rustls-based TLS connector.
|
||||
///
|
||||
/// Tries the platform's native certificate store first. If that yields zero
|
||||
/// certificates (slim container images, missing ca-certificates package),
|
||||
/// falls back to Mozilla's root certificates bundled via `webpki-roots`.
|
||||
fn make_rustls_connector() -> Result<MakeRustlsConnect, rustls::Error> {
|
||||
let mut root_store = rustls::RootCertStore::empty();
|
||||
|
||||
// Try native certs first.
|
||||
let native = rustls_native_certs::load_native_certs();
|
||||
for e in &native.errors {
|
||||
tracing::warn!("error loading system root certs: {e}");
|
||||
@@ -31,11 +38,16 @@ fn make_rustls_connector() -> Result<MakeRustlsConnect, rustls::Error> {
|
||||
tracing::warn!("skipping invalid system root cert: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to bundled Mozilla roots when the system store is empty.
|
||||
if root_store.is_empty() {
|
||||
tracing::error!("no system root certificates found -- TLS connections will fail");
|
||||
tracing::info!(
|
||||
"no system root certificates found, using bundled Mozilla roots"
|
||||
);
|
||||
root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
}
|
||||
// `--all-features` brings in both aws-lc-rs and ring-backed rustls providers.
|
||||
// Pick the same ring provider reqwest already uses so postgres TLS setup stays deterministic.
|
||||
|
||||
// Pick the ring crypto provider (same one reqwest uses).
|
||||
let config = rustls::ClientConfig::builder_with_provider(
|
||||
rustls::crypto::ring::default_provider().into(),
|
||||
)
|
||||
@@ -48,7 +60,7 @@ fn make_rustls_connector() -> Result<MakeRustlsConnect, rustls::Error> {
|
||||
/// Create a [`deadpool_postgres::Pool`] with the appropriate TLS connector.
|
||||
///
|
||||
/// - `Disable` → plain TCP (no TLS)
|
||||
/// - `Prefer` / `Require` → rustls with system root certificates
|
||||
/// - `Prefer` / `Require` → rustls with system or bundled root certificates
|
||||
///
|
||||
/// **Note:** `Prefer` and `Require` currently behave identically — both
|
||||
/// provide a TLS connector and will fail if the server rejects the TLS
|
||||
@@ -81,7 +93,6 @@ mod tests {
|
||||
fn create_pool_disable_mode() {
|
||||
let mut config = deadpool_postgres::Config::new();
|
||||
config.url = Some("postgres://localhost/test".to_string());
|
||||
// Should succeed — pool is created lazily, no actual connection needed.
|
||||
let pool = create_pool(&config, SslMode::Disable);
|
||||
assert!(pool.is_ok());
|
||||
}
|
||||
|
||||
+417
-3
@@ -1162,15 +1162,25 @@ impl Store {
|
||||
pub async fn get_webhook_routine_by_path(
|
||||
&self,
|
||||
path: &str,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt(
|
||||
let row = if let Some(uid) = user_id {
|
||||
conn.query_opt(
|
||||
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
|
||||
AND user_id = $2 \
|
||||
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
|
||||
&[&path, &uid],
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
conn.query_opt(
|
||||
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
|
||||
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
|
||||
&[&path],
|
||||
)
|
||||
.await?;
|
||||
.await?
|
||||
};
|
||||
row.as_ref().map(row_to_routine).transpose()
|
||||
}
|
||||
|
||||
@@ -2269,6 +2279,410 @@ impl Store {
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Users / API Tokens / Invitations ====================
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::db::{ApiTokenRecord, UserRecord};
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Create a new user record.
|
||||
pub async fn create_user(&self, user: &UserRecord) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO users (id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
"#,
|
||||
&[
|
||||
&user.id,
|
||||
&user.email,
|
||||
&user.display_name,
|
||||
&user.status,
|
||||
&user.role,
|
||||
&user.created_at,
|
||||
&user.updated_at,
|
||||
&user.last_login_at,
|
||||
&user.created_by,
|
||||
&user.metadata,
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get a user by their string id.
|
||||
pub async fn get_user(&self, id: &str) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt("SELECT id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata FROM users WHERE id = $1", &[&id])
|
||||
.await?;
|
||||
Ok(row.map(|r| row_to_user(&r)))
|
||||
}
|
||||
|
||||
/// Get a user by email address.
|
||||
pub async fn get_user_by_email(
|
||||
&self,
|
||||
email: &str,
|
||||
) -> Result<Option<UserRecord>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt("SELECT id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata FROM users WHERE email = $1", &[&email])
|
||||
.await?;
|
||||
Ok(row.map(|r| row_to_user(&r)))
|
||||
}
|
||||
|
||||
/// List users, optionally filtered by status.
|
||||
pub async fn list_users(&self, status: Option<&str>) -> Result<Vec<UserRecord>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = match status {
|
||||
Some(s) => {
|
||||
conn.query(
|
||||
"SELECT id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata FROM users WHERE status = $1 ORDER BY created_at DESC",
|
||||
&[&s],
|
||||
)
|
||||
.await?
|
||||
}
|
||||
None => {
|
||||
conn.query("SELECT id, email, display_name, status, role, created_at, updated_at, last_login_at, created_by, metadata FROM users ORDER BY created_at DESC", &[])
|
||||
.await?
|
||||
}
|
||||
};
|
||||
Ok(rows.iter().map(row_to_user).collect())
|
||||
}
|
||||
|
||||
/// Update a user's status.
|
||||
pub async fn update_user_status(&self, id: &str, status: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
conn.execute(
|
||||
"UPDATE users SET status = $1, updated_at = NOW() WHERE id = $2",
|
||||
&[&status, &id],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update a user's display name and metadata.
|
||||
pub async fn update_user_profile(
|
||||
&self,
|
||||
id: &str,
|
||||
display_name: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
conn.execute(
|
||||
"UPDATE users SET display_name = $1, metadata = $2, updated_at = NOW() WHERE id = $3",
|
||||
&[&display_name, metadata, &id],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record a login timestamp for a user.
|
||||
pub async fn record_login(&self, id: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
conn.execute(
|
||||
"UPDATE users SET last_login_at = NOW(), updated_at = NOW() WHERE id = $1",
|
||||
&[&id],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create a new API token.
|
||||
pub async fn create_api_token(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
token_hash: &[u8; 32],
|
||||
token_prefix: &str,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
) -> Result<ApiTokenRecord, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO api_tokens (id, user_id, token_hash, token_prefix, name, expires_at, created_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
"#,
|
||||
&[
|
||||
&id,
|
||||
&user_id,
|
||||
&token_hash.to_vec(),
|
||||
&token_prefix,
|
||||
&name,
|
||||
&expires_at,
|
||||
&now,
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
Ok(ApiTokenRecord {
|
||||
id,
|
||||
user_id: user_id.to_string(),
|
||||
name: name.to_string(),
|
||||
token_prefix: token_prefix.to_string(),
|
||||
expires_at,
|
||||
last_used_at: None,
|
||||
created_at: now,
|
||||
revoked_at: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// List tokens for a user.
|
||||
pub async fn list_api_tokens(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<ApiTokenRecord>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, token_prefix, expires_at, last_used_at, created_at, revoked_at
|
||||
FROM api_tokens
|
||||
WHERE user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
&[&user_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows.iter().map(row_to_api_token).collect())
|
||||
}
|
||||
|
||||
/// Soft-revoke a token. Returns false if the token doesn't exist or doesn't belong to the user.
|
||||
pub async fn revoke_api_token(
|
||||
&self,
|
||||
token_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let count = conn
|
||||
.execute(
|
||||
"UPDATE api_tokens SET revoked_at = NOW() WHERE id = $1 AND user_id = $2 AND revoked_at IS NULL",
|
||||
&[&token_id, &user_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(count > 0)
|
||||
}
|
||||
|
||||
/// Authenticate a token by hash. Returns the token record and its owning user
|
||||
/// if the token is active (non-revoked, non-expired) and the user is active.
|
||||
pub async fn authenticate_token(
|
||||
&self,
|
||||
token_hash: &[u8; 32],
|
||||
) -> Result<Option<(ApiTokenRecord, UserRecord)>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT t.id, t.user_id, t.name, t.token_prefix, t.expires_at, t.last_used_at, t.created_at, t.revoked_at,
|
||||
u.id as u_id, u.email, u.display_name, u.status, u.role, u.created_at as u_created_at, u.updated_at, u.last_login_at, u.created_by, u.metadata
|
||||
FROM api_tokens t
|
||||
JOIN users u ON t.user_id = u.id
|
||||
WHERE t.token_hash = $1
|
||||
AND t.revoked_at IS NULL
|
||||
AND (t.expires_at IS NULL OR t.expires_at > NOW())
|
||||
AND u.status = 'active'
|
||||
"#,
|
||||
&[&token_hash.to_vec()],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.map(|r| {
|
||||
let token = ApiTokenRecord {
|
||||
id: r.get("id"),
|
||||
user_id: r.get("user_id"),
|
||||
name: r.get("name"),
|
||||
token_prefix: r.get("token_prefix"),
|
||||
expires_at: r.get("expires_at"),
|
||||
last_used_at: r.get("last_used_at"),
|
||||
created_at: r.get("created_at"),
|
||||
revoked_at: r.get("revoked_at"),
|
||||
};
|
||||
let user = UserRecord {
|
||||
id: r.get("u_id"),
|
||||
email: r.get("email"),
|
||||
display_name: r.get("display_name"),
|
||||
status: r.get("status"),
|
||||
role: r.get("role"),
|
||||
created_at: r.get("u_created_at"),
|
||||
updated_at: r.get("updated_at"),
|
||||
last_login_at: r.get("last_login_at"),
|
||||
created_by: r.get("created_by"),
|
||||
metadata: r.get("metadata"),
|
||||
};
|
||||
(token, user)
|
||||
}))
|
||||
}
|
||||
|
||||
/// Update `last_used_at` for a token.
|
||||
pub async fn record_token_usage(&self, token_id: Uuid) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
conn.execute(
|
||||
"UPDATE api_tokens SET last_used_at = NOW() WHERE id = $1",
|
||||
&[&token_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check whether any user records exist.
|
||||
pub async fn has_any_users(&self) -> Result<bool, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_one(
|
||||
"SELECT EXISTS(SELECT 1 FROM users LIMIT 1) as has_users",
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.get("has_users"))
|
||||
}
|
||||
|
||||
/// Delete a user and all their data across all user-scoped tables.
|
||||
/// Returns false if the user doesn't exist.
|
||||
pub async fn delete_user(&self, id: &str) -> Result<bool, DatabaseError> {
|
||||
let mut conn = self.conn().await?;
|
||||
let tx = conn
|
||||
.transaction()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
// Delete from child tables first to avoid FK violations.
|
||||
// job_events must come before agent_jobs (FK without CASCADE).
|
||||
// agent_jobs cascades to job_actions, llm_calls, estimation_snapshots.
|
||||
// conversations cascades to conversation_messages.
|
||||
// memory_documents cascades to memory_chunks.
|
||||
// routines cascades to routine_runs.
|
||||
// api_tokens cascade automatically via FK on users.
|
||||
for table in &[
|
||||
"settings",
|
||||
"heartbeat_state",
|
||||
"tool_rate_limit_state",
|
||||
"secret_usage_log",
|
||||
"leak_detection_events",
|
||||
"secrets",
|
||||
"wasm_tools",
|
||||
"routines",
|
||||
"memory_documents",
|
||||
"conversations",
|
||||
] {
|
||||
tx.execute(&format!("DELETE FROM {table} WHERE user_id = $1"), &[&id])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
}
|
||||
// job_events references agent_jobs(id) without CASCADE — delete via subquery.
|
||||
tx.execute(
|
||||
"DELETE FROM job_events WHERE job_id IN (SELECT id FROM agent_jobs WHERE user_id = $1)",
|
||||
&[&id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
tx.execute("DELETE FROM agent_jobs WHERE user_id = $1", &[&id])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
// Nullify self-referencing created_by before deleting the user
|
||||
tx.execute(
|
||||
"UPDATE users SET created_by = NULL WHERE created_by = $1",
|
||||
&[&id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
// api_tokens cascade automatically via FK
|
||||
let result = tx
|
||||
.execute("DELETE FROM users WHERE id = $1", &[&id])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(result > 0)
|
||||
}
|
||||
|
||||
/// Get per-user LLM usage stats for a time period.
|
||||
/// Aggregates from llm_calls via agent_jobs.user_id.
|
||||
pub async fn user_usage_stats(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
since: DateTime<Utc>,
|
||||
) -> Result<Vec<crate::db::UserUsageStats>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = if let Some(uid) = user_id {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT j.user_id, l.model, COUNT(*) as call_count,
|
||||
COALESCE(SUM(l.input_tokens), 0) as input_tokens,
|
||||
COALESCE(SUM(l.output_tokens), 0) as output_tokens,
|
||||
COALESCE(SUM(l.cost), 0) as total_cost
|
||||
FROM llm_calls l
|
||||
JOIN agent_jobs j ON l.job_id = j.id
|
||||
WHERE l.created_at >= $1
|
||||
AND j.user_id = $2
|
||||
GROUP BY j.user_id, l.model
|
||||
ORDER BY total_cost DESC
|
||||
"#,
|
||||
&[&since, &uid],
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT j.user_id, l.model, COUNT(*) as call_count,
|
||||
COALESCE(SUM(l.input_tokens), 0) as input_tokens,
|
||||
COALESCE(SUM(l.output_tokens), 0) as output_tokens,
|
||||
COALESCE(SUM(l.cost), 0) as total_cost
|
||||
FROM llm_calls l
|
||||
JOIN agent_jobs j ON l.job_id = j.id
|
||||
WHERE l.created_at >= $1
|
||||
GROUP BY j.user_id, l.model
|
||||
ORDER BY total_cost DESC
|
||||
"#,
|
||||
&[&since],
|
||||
)
|
||||
.await?
|
||||
};
|
||||
let mut stats = Vec::with_capacity(rows.len());
|
||||
for row in &rows {
|
||||
stats.push(crate::db::UserUsageStats {
|
||||
user_id: row.get("user_id"),
|
||||
model: row.get("model"),
|
||||
call_count: row.get("call_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
total_cost: row.get("total_cost"),
|
||||
});
|
||||
}
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_user(row: &tokio_postgres::Row) -> UserRecord {
|
||||
UserRecord {
|
||||
id: row.get("id"),
|
||||
email: row.get("email"),
|
||||
display_name: row.get("display_name"),
|
||||
status: row.get("status"),
|
||||
role: row.get("role"),
|
||||
created_at: row.get("created_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
last_login_at: row.get("last_login_at"),
|
||||
created_by: row.get("created_by"),
|
||||
metadata: row.get("metadata"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_api_token(row: &tokio_postgres::Row) -> ApiTokenRecord {
|
||||
ApiTokenRecord {
|
||||
id: row.get("id"),
|
||||
user_id: row.get("user_id"),
|
||||
name: row.get("name"),
|
||||
token_prefix: row.get("token_prefix"),
|
||||
expires_at: row.get("expires_at"),
|
||||
last_used_at: row.get("last_used_at"),
|
||||
created_at: row.get("created_at"),
|
||||
revoked_at: row.get("revoked_at"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -69,6 +69,7 @@ pub mod service;
|
||||
pub mod settings;
|
||||
pub mod setup;
|
||||
pub mod skills;
|
||||
pub mod tenant;
|
||||
pub mod timezone;
|
||||
pub mod tools;
|
||||
pub mod tracing_fmt;
|
||||
|
||||
+2
-1
@@ -63,7 +63,8 @@ pub use provider::{
|
||||
};
|
||||
pub use reasoning::{
|
||||
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN,
|
||||
TOOL_INTENT_NUDGE, TokenUsage, ToolSelection, is_silent_reply, llm_signals_tool_intent,
|
||||
TOOL_INTENT_NUDGE, TRUNCATED_TOOL_CALL_NOTICE, TokenUsage, ToolSelection, is_silent_reply,
|
||||
llm_signals_tool_intent,
|
||||
};
|
||||
pub use recording::RecordingLlm;
|
||||
pub use registry::{ProviderDefinition, ProviderProtocol, ProviderRegistry};
|
||||
|
||||
+211
-8
@@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize};
|
||||
use crate::llm::error::LlmError;
|
||||
|
||||
use crate::llm::{
|
||||
ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest,
|
||||
ToolDefinition,
|
||||
ChatMessage, CompletionRequest, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolDefinition,
|
||||
};
|
||||
|
||||
/// Token the agent returns when it has nothing to say (e.g. in group chats).
|
||||
@@ -23,6 +23,13 @@ You said you would perform an action, but you did not include any tool calls.\n\
|
||||
Do NOT describe what you intend to do — actually call the tool now.\n\
|
||||
Use the tool_calls mechanism to invoke the appropriate tool.";
|
||||
|
||||
/// Notice injected when the LLM's response was truncated mid-tool-call,
|
||||
/// causing incomplete parameters. Tells the LLM to try a different approach.
|
||||
pub const TRUNCATED_TOOL_CALL_NOTICE: &str = "\
|
||||
Your previous response was truncated while generating tool call parameters. \
|
||||
The tool calls were discarded. Please try a different approach — \
|
||||
summarize or transform the data instead of echoing it verbatim in a tool call.";
|
||||
|
||||
/// Seed value used as the second argument to `generate_tool_call_id` when
|
||||
/// recovering tool calls from malformed LLM text responses. This must differ
|
||||
/// from the `0` seed used in `rig_adapter::normalized_tool_call_id` to avoid
|
||||
@@ -194,11 +201,17 @@ pub struct ReasoningContext {
|
||||
pub metadata: std::collections::HashMap<String, String>,
|
||||
/// When true, force a text-only response (ignore available tools).
|
||||
/// Used by the agentic loop to guarantee termination near the iteration limit.
|
||||
/// Sticky: once set, never cleared within a loop invocation. Callers must
|
||||
/// create a fresh `ReasoningContext` per `run_agentic_loop()` call.
|
||||
pub force_text: bool,
|
||||
/// Pre-built system prompt. When set, `respond_with_tools` uses this directly
|
||||
/// instead of calling `build_system_prompt_with_tools`. Allows callers to build
|
||||
/// the prompt once and reuse it across iterations.
|
||||
pub system_prompt: Option<String>,
|
||||
/// Per-user model override. When set, completion requests use this model
|
||||
/// instead of the provider's default. Only effective with providers that
|
||||
/// support per-request model overrides (e.g. NearAI).
|
||||
pub model_override: Option<String>,
|
||||
}
|
||||
|
||||
impl ReasoningContext {
|
||||
@@ -212,6 +225,7 @@ impl ReasoningContext {
|
||||
metadata: std::collections::HashMap::new(),
|
||||
force_text: false,
|
||||
system_prompt: None,
|
||||
model_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -344,6 +358,7 @@ pub enum RespondResult {
|
||||
pub struct RespondOutput {
|
||||
pub result: RespondResult,
|
||||
pub usage: TokenUsage,
|
||||
pub finish_reason: FinishReason,
|
||||
}
|
||||
|
||||
/// Reasoning engine for the agent.
|
||||
@@ -525,6 +540,17 @@ impl Reasoning {
|
||||
|
||||
let response = self.llm.complete_with_tools(request).await?;
|
||||
|
||||
// If the response was truncated, tool call parameters are likely incomplete.
|
||||
// Return empty so the caller can fall through to respond_with_tools() which
|
||||
// has a larger output token budget.
|
||||
if response.finish_reason == FinishReason::Length {
|
||||
tracing::warn!(
|
||||
"select_tools response truncated (finish_reason=Length), \
|
||||
discarding potentially incomplete tool selections"
|
||||
);
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
let shared_reasoning = response
|
||||
.content
|
||||
.map(|c| {
|
||||
@@ -671,6 +697,9 @@ Respond in JSON format:
|
||||
.with_temperature(0.7)
|
||||
.with_tool_choice("auto");
|
||||
request.metadata = context.metadata.clone();
|
||||
if let Some(ref model) = context.model_override {
|
||||
request.model = Some(model.clone());
|
||||
}
|
||||
|
||||
let response = self.llm.complete_with_tools(request).await?;
|
||||
let usage = TokenUsage {
|
||||
@@ -714,6 +743,7 @@ Respond in JSON format:
|
||||
content: narrative,
|
||||
},
|
||||
usage,
|
||||
finish_reason: response.finish_reason,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -741,6 +771,7 @@ Respond in JSON format:
|
||||
},
|
||||
},
|
||||
usage,
|
||||
finish_reason: response.finish_reason,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -766,6 +797,7 @@ Respond in JSON format:
|
||||
Ok(RespondOutput {
|
||||
result: RespondResult::Text(final_text),
|
||||
usage,
|
||||
finish_reason: response.finish_reason,
|
||||
})
|
||||
} else {
|
||||
// No tools, use simple completion
|
||||
@@ -773,6 +805,9 @@ Respond in JSON format:
|
||||
.with_max_tokens(4096)
|
||||
.with_temperature(0.7);
|
||||
request.metadata = context.metadata.clone();
|
||||
if let Some(ref model) = context.model_override {
|
||||
request.model = Some(model.clone());
|
||||
}
|
||||
|
||||
let response = self.llm.complete(request).await?;
|
||||
let pre_truncated = truncate_at_tool_tags(&response.content);
|
||||
@@ -794,6 +829,7 @@ Respond in JSON format:
|
||||
cache_read_input_tokens: response.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
||||
},
|
||||
finish_reason: response.finish_reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1334,6 +1370,49 @@ fn is_inside_code(pos: usize, regions: &[CodeRegion]) -> bool {
|
||||
regions.iter().any(|r| pos >= r.start && pos < r.end)
|
||||
}
|
||||
|
||||
/// Check whether a byte range overlaps any code region.
|
||||
fn overlaps_code_region(start: usize, end: usize, regions: &[CodeRegion]) -> bool {
|
||||
regions.iter().any(|r| start < r.end && end > r.start)
|
||||
}
|
||||
|
||||
/// Return the byte bounds of the line containing `pos`, excluding the trailing newline.
|
||||
fn line_bounds(text: &str, pos: usize) -> (usize, usize) {
|
||||
let start = text[..pos].rfind('\n').map_or(0, |idx| idx + 1);
|
||||
let end = text[pos..].find('\n').map_or(text.len(), |idx| pos + idx);
|
||||
(start, end)
|
||||
}
|
||||
|
||||
/// Only recover XML-style tool calls when they are isolated content outside
|
||||
/// markdown code and quote contexts. This avoids converting code examples or
|
||||
/// quoted snippets into executable tool calls.
|
||||
fn is_recoverable_tool_call_segment(
|
||||
text: &str,
|
||||
start: usize,
|
||||
end: usize,
|
||||
code_regions: &[CodeRegion],
|
||||
) -> bool {
|
||||
if overlaps_code_region(start, end, code_regions) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let (first_line_start, first_line_end) = line_bounds(text, start);
|
||||
let first_line = &text[first_line_start..first_line_end];
|
||||
|
||||
if first_line.trim_start().starts_with('>') {
|
||||
return false;
|
||||
}
|
||||
|
||||
let (_, last_line_end) = line_bounds(text, end.saturating_sub(1));
|
||||
let first_line_prefix = &text[first_line_start..start];
|
||||
let last_line_suffix = &text[end..last_line_end];
|
||||
|
||||
if !first_line_prefix.trim().is_empty() || !last_line_suffix.trim().is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
|
||||
///
|
||||
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
|
||||
@@ -1353,6 +1432,7 @@ fn recover_tool_calls_from_content(
|
||||
) -> Vec<ToolCall> {
|
||||
let tool_names: std::collections::HashSet<&str> =
|
||||
available_tools.iter().map(|t| t.name.as_str()).collect();
|
||||
let code_regions = find_code_regions(content);
|
||||
let mut calls = Vec::new();
|
||||
|
||||
for (open, close) in &[
|
||||
@@ -1361,15 +1441,23 @@ fn recover_tool_calls_from_content(
|
||||
("<function_call>", "</function_call>"),
|
||||
("<|function_call|>", "<|/function_call|>"),
|
||||
] {
|
||||
let mut remaining = content;
|
||||
while let Some(start) = remaining.find(open) {
|
||||
let mut search_from = 0;
|
||||
while let Some(offset) = content[search_from..].find(open) {
|
||||
let start = search_from + offset;
|
||||
let inner_start = start + open.len();
|
||||
let after = &remaining[inner_start..];
|
||||
let Some(end) = after.find(close) else {
|
||||
let after = &content[inner_start..];
|
||||
let Some(end_offset) = after.find(close) else {
|
||||
break;
|
||||
};
|
||||
let inner = after[..end].trim();
|
||||
remaining = &after[end + close.len()..];
|
||||
let end = inner_start + end_offset;
|
||||
let segment_end = end + close.len();
|
||||
search_from = segment_end;
|
||||
|
||||
if !is_recoverable_tool_call_segment(content, start, segment_end, &code_regions) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let inner = content[inner_start..end].trim();
|
||||
|
||||
if inner.is_empty() {
|
||||
continue;
|
||||
@@ -2302,6 +2390,40 @@ That's my plan."#;
|
||||
assert_eq!(calls[0].name, "tool_list");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_tool_call_in_fenced_code_block_ignored() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "Here is the XML format:\n\n```xml\n<tool_call>tool_list</tool_call>\n```";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert!(calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_tool_call_in_inline_code_ignored() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "Use `<tool_call>tool_list</tool_call>` to illustrate the syntax.";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert!(calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_tool_call_in_blockquote_ignored() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "The page replied:\n> <tool_call>tool_list</tool_call>";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert!(calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_multiline_json_tool_call_on_own_line() {
|
||||
let tools = make_tools(&["memory_search"]);
|
||||
let content = "Let me check.\n\n<tool_call>\n{\"name\": \"memory_search\", \"arguments\": {\"query\": \"test\"}}\n</tool_call>\n\nDone.";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "memory_search");
|
||||
assert_eq!(calls[0].arguments, serde_json::json!({"query": "test"}));
|
||||
}
|
||||
|
||||
// ---- System prompt building tests (issue #565) ----
|
||||
|
||||
fn make_test_reasoning() -> Reasoning {
|
||||
@@ -3218,4 +3340,85 @@ That's my plan."#;
|
||||
let cleaned = clean_response(&pre_truncated);
|
||||
assert!(cleaned.trim().is_empty());
|
||||
}
|
||||
|
||||
// ---- select_tools truncation guard ----
|
||||
|
||||
/// Mock provider that returns tool calls with a configurable finish_reason.
|
||||
struct TruncatingLlm {
|
||||
finish_reason: crate::llm::FinishReason,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl crate::llm::LlmProvider for TruncatingLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
"truncating-stub"
|
||||
}
|
||||
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
|
||||
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
|
||||
}
|
||||
async fn complete(
|
||||
&self,
|
||||
_request: crate::llm::CompletionRequest,
|
||||
) -> Result<crate::llm::CompletionResponse, crate::llm::error::LlmError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_request: crate::llm::ToolCompletionRequest,
|
||||
) -> Result<crate::llm::ToolCompletionResponse, crate::llm::error::LlmError> {
|
||||
Ok(crate::llm::ToolCompletionResponse {
|
||||
content: Some("I'll write the report.".to_string()),
|
||||
tool_calls: vec![ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "memory_write".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
reasoning: None,
|
||||
}],
|
||||
input_tokens: 5000,
|
||||
output_tokens: 1024,
|
||||
finish_reason: self.finish_reason,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_tools_returns_empty_on_truncation() {
|
||||
let llm = Arc::new(TruncatingLlm {
|
||||
finish_reason: FinishReason::Length,
|
||||
});
|
||||
let reasoning = Reasoning::new(llm);
|
||||
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
|
||||
ctx.available_tools.push(ToolDefinition {
|
||||
name: "memory_write".to_string(),
|
||||
description: "Write to memory".to_string(),
|
||||
parameters: serde_json::json!({"type": "object"}),
|
||||
});
|
||||
|
||||
let selections = reasoning.select_tools(&ctx).await.unwrap();
|
||||
assert!(
|
||||
selections.is_empty(),
|
||||
"Truncated tool selections should be discarded (got {} selections)",
|
||||
selections.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_tools_returns_selections_when_not_truncated() {
|
||||
let llm = Arc::new(TruncatingLlm {
|
||||
finish_reason: FinishReason::ToolUse,
|
||||
});
|
||||
let reasoning = Reasoning::new(llm);
|
||||
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
|
||||
ctx.available_tools.push(ToolDefinition {
|
||||
name: "memory_write".to_string(),
|
||||
description: "Write to memory".to_string(),
|
||||
parameters: serde_json::json!({"type": "object"}),
|
||||
});
|
||||
|
||||
let selections = reasoning.select_tools(&ctx).await.unwrap();
|
||||
assert_eq!(selections.len(), 1);
|
||||
assert_eq!(selections[0].tool_name, "memory_write");
|
||||
}
|
||||
}
|
||||
|
||||
+82
-20
@@ -598,6 +598,33 @@ fn build_rig_request(
|
||||
})
|
||||
}
|
||||
|
||||
/// Inject a per-request model override into the rig request's `additional_params`.
|
||||
///
|
||||
/// Rig-core bakes the model name at construction time inside each provider's
|
||||
/// `CompletionModel` implementation. This helper inserts a top-level `"model"`
|
||||
/// key into `additional_params`, which rig-core flattens into the provider's
|
||||
/// request payload via `#[serde(flatten)]`.
|
||||
///
|
||||
/// Whether the override takes effect depends on the downstream API server's
|
||||
/// handling of duplicate JSON keys (most Python/Go servers use last-key-wins,
|
||||
/// but this is not guaranteed by the JSON spec). The `effective_model_name()`
|
||||
/// trait method should be consulted to determine the model actually used.
|
||||
fn inject_model_override(rig_req: &mut RigRequest, model_override: Option<&str>) {
|
||||
let Some(model) = model_override else {
|
||||
return;
|
||||
};
|
||||
match rig_req.additional_params {
|
||||
Some(ref mut params) => {
|
||||
if let Some(obj) = params.as_object_mut() {
|
||||
obj.insert("model".to_string(), serde_json::json!(model));
|
||||
}
|
||||
}
|
||||
None => {
|
||||
rig_req.additional_params = Some(serde_json::json!({ "model": model }));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<M> LlmProvider for RigAdapter<M>
|
||||
where
|
||||
@@ -632,15 +659,7 @@ where
|
||||
&self,
|
||||
mut request: CompletionRequest,
|
||||
) -> Result<CompletionResponse, LlmError> {
|
||||
if let Some(requested_model) = request.model.as_deref()
|
||||
&& requested_model != self.model_name.as_str()
|
||||
{
|
||||
tracing::warn!(
|
||||
requested_model = requested_model,
|
||||
active_model = %self.model_name,
|
||||
"Per-request model override is not supported for this provider; using configured model"
|
||||
);
|
||||
}
|
||||
let model_override = request.model.take();
|
||||
|
||||
self.strip_unsupported_completion_params(&mut request);
|
||||
|
||||
@@ -648,7 +667,7 @@ where
|
||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
|
||||
let rig_req = build_rig_request(
|
||||
let mut rig_req = build_rig_request(
|
||||
preamble,
|
||||
history,
|
||||
Vec::new(),
|
||||
@@ -658,6 +677,8 @@ where
|
||||
self.cache_retention,
|
||||
)?;
|
||||
|
||||
inject_model_override(&mut rig_req, model_override.as_deref());
|
||||
|
||||
let response =
|
||||
self.model
|
||||
.completion(rig_req)
|
||||
@@ -695,15 +716,7 @@ where
|
||||
&self,
|
||||
mut request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
if let Some(requested_model) = request.model.as_deref()
|
||||
&& requested_model != self.model_name.as_str()
|
||||
{
|
||||
tracing::warn!(
|
||||
requested_model = requested_model,
|
||||
active_model = %self.model_name,
|
||||
"Per-request model override is not supported for this provider; using configured model"
|
||||
);
|
||||
}
|
||||
let model_override = request.model.take();
|
||||
|
||||
self.strip_unsupported_tool_params(&mut request);
|
||||
|
||||
@@ -716,7 +729,7 @@ where
|
||||
let tools = convert_tools(&request.tools);
|
||||
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
||||
|
||||
let rig_req = build_rig_request(
|
||||
let mut rig_req = build_rig_request(
|
||||
preamble,
|
||||
history,
|
||||
tools,
|
||||
@@ -726,6 +739,8 @@ where
|
||||
self.cache_retention,
|
||||
)?;
|
||||
|
||||
inject_model_override(&mut rig_req, model_override.as_deref());
|
||||
|
||||
let response =
|
||||
self.model
|
||||
.completion(rig_req)
|
||||
@@ -1503,4 +1518,51 @@ mod tests {
|
||||
"different raw IDs should produce different hashed IDs"
|
||||
);
|
||||
}
|
||||
|
||||
fn make_rig_request(additional_params: Option<serde_json::Value>) -> RigRequest {
|
||||
RigRequest {
|
||||
preamble: None,
|
||||
chat_history: OneOrMany::one(RigMessage::user("test")),
|
||||
documents: Vec::new(),
|
||||
tools: Vec::new(),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
tool_choice: None,
|
||||
additional_params,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_model_override_creates_params_when_none() {
|
||||
let mut req = make_rig_request(None);
|
||||
inject_model_override(&mut req, Some("test-model"));
|
||||
|
||||
let params = req
|
||||
.additional_params
|
||||
.expect("additional_params should be Some");
|
||||
assert_eq!(params, serde_json::json!({ "model": "test-model" }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_model_override_preserves_existing_params() {
|
||||
let mut req = make_rig_request(Some(serde_json::json!({
|
||||
"cache_control": { "type": "ephemeral" },
|
||||
})));
|
||||
inject_model_override(&mut req, Some("override-model"));
|
||||
|
||||
let params = req.additional_params.expect("should remain Some");
|
||||
let obj = params.as_object().expect("should be object");
|
||||
assert_eq!(
|
||||
obj.get("cache_control"),
|
||||
Some(&serde_json::json!({ "type": "ephemeral" }))
|
||||
);
|
||||
assert_eq!(obj.get("model"), Some(&serde_json::json!("override-model")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_model_override_noop_when_none() {
|
||||
let mut req = make_rig_request(None);
|
||||
inject_model_override(&mut req, None);
|
||||
assert!(req.additional_params.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+54
-27
@@ -591,27 +591,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
let mut gateway_url: Option<String> = None;
|
||||
let mut sse_manager: Option<std::sync::Arc<ironclaw::channels::web::sse::SseManager>> = None;
|
||||
if let Some(ref gw_config) = config.channels.gateway {
|
||||
// Build multi-user auth state if user_tokens is configured, else single-user.
|
||||
let mut gw = if let Some(ref user_tokens) = gw_config.user_tokens {
|
||||
use ironclaw::channels::web::auth::{MultiAuthState, UserIdentity};
|
||||
let tokens = user_tokens
|
||||
.iter()
|
||||
.map(|(token, cfg)| {
|
||||
(
|
||||
token.clone(),
|
||||
UserIdentity {
|
||||
user_id: cfg.user_id.clone(),
|
||||
workspace_read_scopes: cfg.workspace_read_scopes.clone(),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let auth = MultiAuthState::multi(tokens);
|
||||
GatewayChannel::new_multi_auth(gw_config.clone(), auth)
|
||||
} else {
|
||||
GatewayChannel::new(gw_config.clone())
|
||||
};
|
||||
gw = gw.with_owner_scope(config.owner_id.clone());
|
||||
let mut gw = GatewayChannel::new(gw_config.clone(), config.owner_id.clone());
|
||||
gw = gw.with_llm_provider(Arc::clone(&components.llm));
|
||||
if let Some(ref ws) = components.workspace {
|
||||
gw = gw.with_workspace(Arc::clone(ws));
|
||||
@@ -650,6 +630,54 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
}
|
||||
if let Some(ref d) = components.db {
|
||||
gw = gw.with_store(Arc::clone(d));
|
||||
gw = gw.with_db_auth(Arc::clone(d));
|
||||
if let Some(ref ss) = components.secrets_store {
|
||||
gw = gw.with_secrets_store(Arc::clone(ss));
|
||||
}
|
||||
|
||||
// Bootstrap: create the first admin user from single-user config
|
||||
// so the owner appears in the Users admin panel immediately.
|
||||
if let Ok(false) = d.has_any_users().await {
|
||||
let now = chrono::Utc::now();
|
||||
let user = ironclaw::db::UserRecord {
|
||||
id: config.owner_id.clone(),
|
||||
email: None,
|
||||
display_name: config.owner_id.clone(),
|
||||
status: "active".to_string(),
|
||||
role: "admin".to_string(),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
last_login_at: None,
|
||||
created_by: None,
|
||||
metadata: serde_json::json!({"source": "bootstrap"}),
|
||||
};
|
||||
if let Err(e) = d.create_user(&user).await {
|
||||
tracing::warn!("Failed to bootstrap admin user: {}", e);
|
||||
} else {
|
||||
// Also create an API token from the gateway auth token so
|
||||
// DB-backed auth works for the bootstrapped user.
|
||||
let auth_token = gw.auth_token();
|
||||
if !auth_token.is_empty() {
|
||||
use ironclaw::channels::web::auth::hash_token;
|
||||
let hash = hash_token(auth_token);
|
||||
let prefix = if auth_token.len() >= 8 {
|
||||
&auth_token[..8]
|
||||
} else {
|
||||
auth_token
|
||||
};
|
||||
if let Err(e) = d
|
||||
.create_api_token(&config.owner_id, "bootstrap", &hash, prefix, None)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to create bootstrap token: {}", e);
|
||||
}
|
||||
}
|
||||
tracing::info!(
|
||||
user_id = config.owner_id,
|
||||
"Bootstrapped admin user from gateway config"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(ref jm) = container_job_manager {
|
||||
gw = gw.with_job_manager(Arc::clone(jm));
|
||||
@@ -791,12 +819,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
.await;
|
||||
|
||||
// Default user ID for extension operations (single-user mode).
|
||||
let ext_user_id = config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|g| g.user_id.clone())
|
||||
.unwrap_or_else(|| "default".to_string());
|
||||
let ext_user_id = config.owner_id.clone();
|
||||
|
||||
// Wire up channel runtime for hot-activation of WASM channels.
|
||||
if let Some(ref ext_mgr) = components.extension_manager
|
||||
@@ -914,6 +937,10 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
},
|
||||
builder: components.builder,
|
||||
llm_backend: config.llm.backend.clone(),
|
||||
tenant_rates: Arc::new(ironclaw::tenant::TenantRateRegistry::new(
|
||||
config.agent.max_llm_concurrent_per_user.unwrap_or(4),
|
||||
config.agent.max_jobs_concurrent_per_user.unwrap_or(3),
|
||||
)),
|
||||
};
|
||||
|
||||
let channels_for_warnings = Arc::clone(&channels);
|
||||
|
||||
@@ -269,10 +269,6 @@ pub struct ChannelSettings {
|
||||
#[serde(default)]
|
||||
pub gateway_auth_token: Option<String>,
|
||||
|
||||
/// Web gateway user ID.
|
||||
#[serde(default)]
|
||||
pub gateway_user_id: Option<String>,
|
||||
|
||||
/// Whether the CLI channel is enabled.
|
||||
#[serde(default = "default_true")]
|
||||
pub cli_enabled: bool,
|
||||
@@ -342,7 +338,6 @@ impl Default for ChannelSettings {
|
||||
gateway_host: None,
|
||||
gateway_port: None,
|
||||
gateway_auth_token: None,
|
||||
gateway_user_id: None,
|
||||
cli_enabled: true,
|
||||
signal_enabled: false,
|
||||
signal_http_url: None,
|
||||
|
||||
+947
@@ -0,0 +1,947 @@
|
||||
//! Compile-time tenant isolation.
|
||||
//!
|
||||
//! Provides two database access tiers:
|
||||
//!
|
||||
//! - **[`TenantScope`]** (default): All operations are bound to a single user.
|
||||
//! ID-based lookups return `None` if the resource doesn't belong to this user.
|
||||
//! This is the only way handler code should access the database.
|
||||
//!
|
||||
//! - **[`AdminScope`]**: Cross-tenant access for system-level operations
|
||||
//! (heartbeat, routine engine, self-repair). Must be obtained explicitly via
|
||||
//! [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
|
||||
//!
|
||||
//! [`TenantCtx`] bundles a `TenantScope` with workspace, cost guard, and
|
||||
//! per-tenant rate limiting. Constructed once per request at the entry point
|
||||
//! where a `user_id` becomes known.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use rust_decimal::Decimal;
|
||||
use tokio::sync::{Semaphore, SemaphorePermit};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::BrokenTool;
|
||||
use crate::agent::cost_guard::{CostGuard, CostLimitExceeded};
|
||||
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::DatabaseError;
|
||||
use crate::history::{
|
||||
AgentJobRecord, AgentJobSummary, ConversationMessage, ConversationSummary, LlmCallRecord,
|
||||
SandboxJobRecord, SandboxJobSummary, SettingRow,
|
||||
};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TenantScope — scoped database access (default tier)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Scoped database view. All operations are bound to a single user.
|
||||
///
|
||||
/// This is the **only** way handler code should access the database.
|
||||
/// ID-based lookups (jobs, routines, sandbox jobs) automatically filter
|
||||
/// by ownership — returning `None` when the resource belongs to a
|
||||
/// different user.
|
||||
#[derive(Clone)]
|
||||
pub struct TenantScope {
|
||||
user_id: String,
|
||||
inner: Arc<dyn Database>,
|
||||
}
|
||||
|
||||
impl TenantScope {
|
||||
pub fn new(user_id: impl Into<String>, db: Arc<dyn Database>) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
inner: db,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn user_id(&self) -> &str {
|
||||
&self.user_id
|
||||
}
|
||||
|
||||
// === Jobs ===
|
||||
|
||||
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||
self.inner.list_agent_jobs_for_user(&self.user_id).await
|
||||
}
|
||||
|
||||
pub async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError> {
|
||||
self.inner.agent_job_summary_for_user(&self.user_id).await
|
||||
}
|
||||
|
||||
/// Fetch a job by ID, returning `None` if it doesn't belong to this user.
|
||||
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||
match self.inner.get_job(id).await? {
|
||||
Some(ctx) if ctx.user_id == self.user_id => Ok(Some(ctx)),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_agent_job_failure_reason(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError> {
|
||||
// Verify ownership first
|
||||
if self.get_job(id).await?.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
self.inner.get_agent_job_failure_reason(id).await
|
||||
}
|
||||
|
||||
pub async fn update_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: JobState,
|
||||
failure_reason: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
// Verify ownership before mutating
|
||||
if self.get_job(id).await?.is_none() {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "job".to_string(),
|
||||
id: id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner
|
||||
.update_job_status(id, status, failure_reason)
|
||||
.await
|
||||
}
|
||||
|
||||
// === Sandbox jobs ===
|
||||
|
||||
pub async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
self.inner.list_sandbox_jobs_for_user(&self.user_id).await
|
||||
}
|
||||
|
||||
pub async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
self.inner.sandbox_job_summary_for_user(&self.user_id).await
|
||||
}
|
||||
|
||||
/// Fetch a sandbox job by ID, returning `None` if it doesn't belong to this user.
|
||||
pub async fn get_sandbox_job(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||
match self.inner.get_sandbox_job(id).await? {
|
||||
Some(job) if job.user_id == self.user_id => Ok(Some(job)),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn sandbox_job_belongs_to_user(&self, job_id: Uuid) -> Result<bool, DatabaseError> {
|
||||
self.inner
|
||||
.sandbox_job_belongs_to_user(job_id, &self.user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
// === Routines ===
|
||||
|
||||
pub async fn list_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.inner.list_routines(&self.user_id).await
|
||||
}
|
||||
|
||||
pub async fn get_routine_by_name(&self, name: &str) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.inner.get_routine_by_name(&self.user_id, name).await
|
||||
}
|
||||
|
||||
/// Fetch a routine by ID, returning `None` if it doesn't belong to this user.
|
||||
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||
match self.inner.get_routine(id).await? {
|
||||
Some(r) if r.user_id == self.user_id => Ok(Some(r)),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
debug_assert_eq!(
|
||||
routine.user_id, self.user_id,
|
||||
"routine.user_id must match TenantScope user"
|
||||
);
|
||||
self.inner.create_routine(routine).await
|
||||
}
|
||||
|
||||
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
// Verify ownership
|
||||
if self.get_routine(routine.id).await?.is_none() {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "routine".to_string(),
|
||||
id: routine.id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.update_routine(routine).await
|
||||
}
|
||||
|
||||
pub async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
|
||||
// Verify ownership
|
||||
if self.get_routine(id).await?.is_none() {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "routine".to_string(),
|
||||
id: id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.delete_routine(id).await
|
||||
}
|
||||
|
||||
/// List routine runs, verifying the routine belongs to this user.
|
||||
pub async fn list_routine_runs(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
limit: i64,
|
||||
) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
// Verify routine ownership first
|
||||
if self.get_routine(routine_id).await?.is_none() {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "routine".to_string(),
|
||||
id: routine_id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.list_routine_runs(routine_id, limit).await
|
||||
}
|
||||
|
||||
pub async fn get_webhook_routine_by_path(
|
||||
&self,
|
||||
path: &str,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.inner
|
||||
.get_webhook_routine_by_path(path, Some(&self.user_id))
|
||||
.await
|
||||
}
|
||||
|
||||
// === Settings ===
|
||||
|
||||
pub async fn get_setting(&self, key: &str) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
self.inner.get_setting(&self.user_id, key).await
|
||||
}
|
||||
|
||||
pub async fn get_setting_full(&self, key: &str) -> Result<Option<SettingRow>, DatabaseError> {
|
||||
self.inner.get_setting_full(&self.user_id, key).await
|
||||
}
|
||||
|
||||
pub async fn set_setting(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner.set_setting(&self.user_id, key, value).await
|
||||
}
|
||||
|
||||
pub async fn delete_setting(&self, key: &str) -> Result<bool, DatabaseError> {
|
||||
self.inner.delete_setting(&self.user_id, key).await
|
||||
}
|
||||
|
||||
pub async fn list_settings(&self) -> Result<Vec<SettingRow>, DatabaseError> {
|
||||
self.inner.list_settings(&self.user_id).await
|
||||
}
|
||||
|
||||
pub async fn get_all_settings(
|
||||
&self,
|
||||
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
|
||||
self.inner.get_all_settings(&self.user_id).await
|
||||
}
|
||||
|
||||
pub async fn set_all_settings(
|
||||
&self,
|
||||
settings: &HashMap<String, serde_json::Value>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner.set_all_settings(&self.user_id, settings).await
|
||||
}
|
||||
|
||||
pub async fn has_settings(&self) -> Result<bool, DatabaseError> {
|
||||
self.inner.has_settings(&self.user_id).await
|
||||
}
|
||||
|
||||
// === Conversations ===
|
||||
|
||||
pub async fn create_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.create_conversation(channel, &self.user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn ensure_conversation(
|
||||
&self,
|
||||
id: Uuid,
|
||||
channel: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.inner
|
||||
.ensure_conversation(id, channel, &self.user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn list_conversations_with_preview(
|
||||
&self,
|
||||
channel: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||
self.inner
|
||||
.list_conversations_with_preview(&self.user_id, channel, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn list_conversations_all_channels(
|
||||
&self,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||
self.inner
|
||||
.list_conversations_all_channels(&self.user_id, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_or_create_routine_conversation(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
routine_name: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.get_or_create_routine_conversation(routine_id, routine_name, &self.user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_or_create_heartbeat_conversation(&self) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.get_or_create_heartbeat_conversation(&self.user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_or_create_assistant_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.get_or_create_assistant_conversation(&self.user_id, channel)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn conversation_belongs_to_user(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.inner
|
||||
.conversation_belongs_to_user(conversation_id, &self.user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Add a message to a conversation owned by this tenant.
|
||||
///
|
||||
/// Returns `NotFound` if the conversation does not belong to this user.
|
||||
pub async fn add_conversation_message(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(conversation_id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: conversation_id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner
|
||||
.add_conversation_message(conversation_id, role, content)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Touch a conversation timestamp. Returns `NotFound` if not owned by this user.
|
||||
pub async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.touch_conversation(id).await
|
||||
}
|
||||
|
||||
/// List messages in a conversation. Returns `NotFound` if not owned by this user.
|
||||
pub async fn list_conversation_messages(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Vec<ConversationMessage>, DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(conversation_id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: conversation_id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.list_conversation_messages(conversation_id).await
|
||||
}
|
||||
|
||||
/// Paginated message listing. Returns `NotFound` if not owned by this user.
|
||||
pub async fn list_conversation_messages_paginated(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
before: Option<DateTime<Utc>>,
|
||||
limit: i64,
|
||||
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(conversation_id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: conversation_id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner
|
||||
.list_conversation_messages_paginated(conversation_id, before, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn create_conversation_with_metadata(
|
||||
&self,
|
||||
channel: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.create_conversation_with_metadata(channel, &self.user_id, metadata)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Update metadata on a conversation. Returns `NotFound` if not owned by this user.
|
||||
pub async fn update_conversation_metadata_field(
|
||||
&self,
|
||||
id: Uuid,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner
|
||||
.update_conversation_metadata_field(id, key, value)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Get conversation metadata. Returns `NotFound` if not owned by this user.
|
||||
pub async fn get_conversation_metadata(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
if !self.conversation_belongs_to_user(id).await? {
|
||||
return Err(DatabaseError::NotFound {
|
||||
entity: "conversation".to_string(),
|
||||
id: id.to_string(),
|
||||
});
|
||||
}
|
||||
self.inner.get_conversation_metadata(id).await
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AdminScope — explicit cross-tenant access
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Cross-tenant database access for system-level operations.
|
||||
///
|
||||
/// **Not** available through [`TenantCtx`] — must be obtained explicitly via
|
||||
/// [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
|
||||
///
|
||||
/// Used by: heartbeat enumeration, routine engine scheduling, self-repair,
|
||||
/// scheduler job persistence, worker status updates.
|
||||
#[derive(Clone)]
|
||||
pub struct AdminScope {
|
||||
inner: Arc<dyn Database>,
|
||||
}
|
||||
|
||||
impl AdminScope {
|
||||
pub fn new(db: Arc<dyn Database>) -> Self {
|
||||
Self { inner: db }
|
||||
}
|
||||
|
||||
/// Access the raw Database trait object.
|
||||
///
|
||||
/// Prefer using the typed methods on AdminScope instead. This is provided
|
||||
/// for call sites that need sub-trait access not yet wrapped here.
|
||||
pub fn db(&self) -> &Arc<dyn Database> {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
// === Routine engine ===
|
||||
|
||||
pub async fn list_all_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.inner.list_all_routines().await
|
||||
}
|
||||
|
||||
pub async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.inner.list_event_routines().await
|
||||
}
|
||||
|
||||
pub async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.inner.list_due_cron_routines().await
|
||||
}
|
||||
|
||||
pub async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
self.inner.list_dispatched_routine_runs().await
|
||||
}
|
||||
|
||||
pub async fn count_running_routine_runs_batch(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, i64>, DatabaseError> {
|
||||
self.inner
|
||||
.count_running_routine_runs_batch(routine_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn batch_get_last_run_status(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
|
||||
self.inner.batch_get_last_run_status(routine_ids).await
|
||||
}
|
||||
|
||||
pub async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
|
||||
self.inner.count_running_routine_runs(routine_id).await
|
||||
}
|
||||
|
||||
pub async fn update_routine_runtime(
|
||||
&self,
|
||||
id: Uuid,
|
||||
last_run_at: DateTime<Utc>,
|
||||
next_fire_at: Option<DateTime<Utc>>,
|
||||
run_count: u64,
|
||||
consecutive_failures: u32,
|
||||
state: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.update_routine_runtime(
|
||||
id,
|
||||
last_run_at,
|
||||
next_fire_at,
|
||||
run_count,
|
||||
consecutive_failures,
|
||||
state,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> {
|
||||
self.inner.create_routine_run(run).await
|
||||
}
|
||||
|
||||
pub async fn complete_routine_run(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: RunStatus,
|
||||
result_summary: Option<&str>,
|
||||
tokens_used: Option<i32>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.complete_routine_run(id, status, result_summary, tokens_used)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn link_routine_run_to_job(
|
||||
&self,
|
||||
run_id: Uuid,
|
||||
job_id: Uuid,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner.link_routine_run_to_job(run_id, job_id).await
|
||||
}
|
||||
|
||||
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.inner.get_routine(id).await
|
||||
}
|
||||
|
||||
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
self.inner.update_routine(routine).await
|
||||
}
|
||||
|
||||
// === Self-repair ===
|
||||
|
||||
pub async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
|
||||
self.inner.get_stuck_jobs().await
|
||||
}
|
||||
|
||||
pub async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
|
||||
self.inner.get_broken_tools(threshold).await
|
||||
}
|
||||
|
||||
pub async fn record_tool_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
error_message: &str,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.record_tool_failure(tool_name, error_message)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.inner.mark_tool_repaired(tool_name).await
|
||||
}
|
||||
|
||||
pub async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.inner.increment_repair_attempts(tool_name).await
|
||||
}
|
||||
|
||||
// === Sandbox housekeeping ===
|
||||
|
||||
pub async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
|
||||
self.inner.cleanup_stale_sandbox_jobs().await
|
||||
}
|
||||
|
||||
pub async fn get_sandbox_job(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||
self.inner.get_sandbox_job(id).await
|
||||
}
|
||||
|
||||
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||
self.inner.save_sandbox_job(job).await
|
||||
}
|
||||
|
||||
pub async fn update_sandbox_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: &str,
|
||||
success: Option<bool>,
|
||||
message: Option<&str>,
|
||||
started_at: Option<DateTime<Utc>>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
|
||||
self.inner.update_sandbox_job_mode(id, mode).await
|
||||
}
|
||||
|
||||
pub async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
|
||||
self.inner.get_sandbox_job_mode(id).await
|
||||
}
|
||||
|
||||
pub async fn save_job_event(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
event_type: &str,
|
||||
data: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner.save_job_event(job_id, event_type, data).await
|
||||
}
|
||||
|
||||
pub async fn list_job_events(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
limit: Option<i64>,
|
||||
) -> Result<Vec<crate::history::JobEventRecord>, DatabaseError> {
|
||||
self.inner.list_job_events(job_id, limit).await
|
||||
}
|
||||
|
||||
// === Job persistence (scheduler, worker) ===
|
||||
|
||||
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||
self.inner.get_job(id).await
|
||||
}
|
||||
|
||||
pub async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
|
||||
self.inner.save_job(ctx).await
|
||||
}
|
||||
|
||||
pub async fn update_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: JobState,
|
||||
failure_reason: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.update_job_status(id, status, failure_reason)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.inner.mark_job_stuck(id).await
|
||||
}
|
||||
|
||||
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
|
||||
self.inner.list_agent_jobs().await
|
||||
}
|
||||
|
||||
pub async fn get_agent_job_failure_reason(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError> {
|
||||
self.inner.get_agent_job_failure_reason(id).await
|
||||
}
|
||||
|
||||
// === LLM call recording ===
|
||||
|
||||
pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
|
||||
self.inner.record_llm_call(record).await
|
||||
}
|
||||
|
||||
pub async fn save_action(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
action: &ActionRecord,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner.save_action(job_id, action).await
|
||||
}
|
||||
|
||||
pub async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
|
||||
self.inner.get_job_actions(job_id).await
|
||||
}
|
||||
|
||||
// === Estimation ===
|
||||
|
||||
pub async fn save_estimation_snapshot(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
category: &str,
|
||||
tool_names: &[String],
|
||||
estimated_cost: Decimal,
|
||||
estimated_time_secs: i32,
|
||||
estimated_value: Decimal,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.save_estimation_snapshot(
|
||||
job_id,
|
||||
category,
|
||||
tool_names,
|
||||
estimated_cost,
|
||||
estimated_time_secs,
|
||||
estimated_value,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn update_estimation_actuals(
|
||||
&self,
|
||||
id: Uuid,
|
||||
actual_cost: Decimal,
|
||||
actual_time_secs: i32,
|
||||
actual_value: Option<Decimal>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.inner
|
||||
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
|
||||
.await
|
||||
}
|
||||
|
||||
// === Conversations (admin context) ===
|
||||
|
||||
pub async fn add_conversation_message(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.add_conversation_message(conversation_id, role, content)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_or_create_routine_conversation(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
routine_name: &str,
|
||||
user_id: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.get_or_create_routine_conversation(routine_id, routine_name, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_or_create_heartbeat_conversation(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.inner
|
||||
.get_or_create_heartbeat_conversation(user_id)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TenantRateState / TenantRateRegistry — per-user concurrency
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Per-tenant concurrency limits.
|
||||
pub struct TenantRateState {
|
||||
/// Limits concurrent LLM calls for this user.
|
||||
pub llm_semaphore: Arc<Semaphore>,
|
||||
/// Limits concurrent jobs for this user.
|
||||
pub job_semaphore: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl TenantRateState {
|
||||
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
|
||||
Self {
|
||||
llm_semaphore: Arc::new(Semaphore::new(max_llm_concurrent)),
|
||||
job_semaphore: Arc::new(Semaphore::new(max_job_concurrent)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Registry that lazily creates per-tenant rate state.
|
||||
///
|
||||
/// Uses `tokio::sync::RwLock<HashMap>` (consistent with the rest of the
|
||||
/// codebase — no DashMap dependency).
|
||||
pub struct TenantRateRegistry {
|
||||
state: tokio::sync::RwLock<HashMap<String, Arc<TenantRateState>>>,
|
||||
max_llm_concurrent: usize,
|
||||
max_job_concurrent: usize,
|
||||
}
|
||||
|
||||
impl TenantRateRegistry {
|
||||
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
|
||||
Self {
|
||||
state: tokio::sync::RwLock::new(HashMap::new()),
|
||||
max_llm_concurrent,
|
||||
max_job_concurrent,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or lazily create rate state for a user.
|
||||
pub async fn get_or_create(&self, user_id: &str) -> Arc<TenantRateState> {
|
||||
// Fast path: read lock
|
||||
{
|
||||
let map = self.state.read().await;
|
||||
if let Some(s) = map.get(user_id) {
|
||||
return Arc::clone(s);
|
||||
}
|
||||
}
|
||||
// Slow path: write lock with double-check
|
||||
let mut map = self.state.write().await;
|
||||
if let Some(s) = map.get(user_id) {
|
||||
return Arc::clone(s);
|
||||
}
|
||||
let s = Arc::new(TenantRateState::new(
|
||||
self.max_llm_concurrent,
|
||||
self.max_job_concurrent,
|
||||
));
|
||||
map.insert(user_id.to_string(), Arc::clone(&s));
|
||||
s
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TenantCtx — per-request tenant execution context
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Per-request tenant execution context.
|
||||
///
|
||||
/// Bundles a [`TenantScope`] (scoped DB access), workspace, cost guard,
|
||||
/// and per-tenant rate limiting. Constructed once per request via
|
||||
/// [`AgentDeps::tenant_ctx()`](crate::agent::AgentDeps::tenant_ctx).
|
||||
///
|
||||
/// `Clone + Send + Sync` — safe to store on `ChatDelegate` without lifetime issues.
|
||||
#[derive(Clone)]
|
||||
pub struct TenantCtx {
|
||||
user_id: String,
|
||||
store: Option<TenantScope>,
|
||||
workspace: Option<Arc<Workspace>>,
|
||||
cost_guard: Arc<CostGuard>,
|
||||
rate: Arc<TenantRateState>,
|
||||
}
|
||||
|
||||
impl TenantCtx {
|
||||
pub fn new(
|
||||
user_id: impl Into<String>,
|
||||
store: Option<TenantScope>,
|
||||
workspace: Option<Arc<Workspace>>,
|
||||
cost_guard: Arc<CostGuard>,
|
||||
rate: Arc<TenantRateState>,
|
||||
) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
store,
|
||||
workspace,
|
||||
cost_guard,
|
||||
rate,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn user_id(&self) -> &str {
|
||||
&self.user_id
|
||||
}
|
||||
|
||||
pub fn store(&self) -> Option<&TenantScope> {
|
||||
self.store.as_ref()
|
||||
}
|
||||
|
||||
pub fn workspace(&self) -> Option<&Arc<Workspace>> {
|
||||
self.workspace.as_ref()
|
||||
}
|
||||
|
||||
pub fn cost_guard(&self) -> &CostGuard {
|
||||
&self.cost_guard
|
||||
}
|
||||
|
||||
/// Check cost limits for this tenant (global + per-user).
|
||||
pub async fn check_cost_allowed(&self) -> Result<(), CostLimitExceeded> {
|
||||
self.cost_guard.check_allowed_for_user(&self.user_id).await
|
||||
}
|
||||
|
||||
/// Record an LLM call for this tenant.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn record_llm_call(
|
||||
&self,
|
||||
model: &str,
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
cache_read_input_tokens: u32,
|
||||
cache_creation_input_tokens: u32,
|
||||
cache_read_discount: Decimal,
|
||||
cache_write_multiplier: Decimal,
|
||||
cost_per_token: Option<(Decimal, Decimal)>,
|
||||
) -> Decimal {
|
||||
self.cost_guard
|
||||
.record_llm_call_for_user(
|
||||
&self.user_id,
|
||||
model,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_creation_input_tokens,
|
||||
cache_read_discount,
|
||||
cache_write_multiplier,
|
||||
cost_per_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Acquire an LLM concurrency permit for this tenant.
|
||||
pub async fn acquire_llm_permit(&self) -> Result<SemaphorePermit<'_>, crate::error::Error> {
|
||||
self.rate.llm_semaphore.acquire().await.map_err(|_| {
|
||||
crate::error::Error::Config(crate::error::ConfigError::InvalidValue {
|
||||
key: "llm_semaphore".to_string(),
|
||||
message: "semaphore closed".to_string(),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_rate_registry_returns_same_state_for_same_user() {
|
||||
let registry = TenantRateRegistry::new(4, 3);
|
||||
let a1 = registry.get_or_create("alice").await;
|
||||
let a2 = registry.get_or_create("alice").await;
|
||||
assert!(Arc::ptr_eq(&a1, &a2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_rate_registry_different_users_get_different_state() {
|
||||
let registry = TenantRateRegistry::new(4, 3);
|
||||
let alice = registry.get_or_create("alice").await;
|
||||
let bob = registry.get_or_create("bob").await;
|
||||
assert!(!Arc::ptr_eq(&alice, &bob));
|
||||
}
|
||||
}
|
||||
@@ -532,6 +532,7 @@ impl TestHarnessBuilder {
|
||||
let cost_guard = Arc::new(CostGuard::new(CostGuardConfig {
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_cost_per_user_per_day_cents: None,
|
||||
}));
|
||||
|
||||
let channel = if self.stub_channel {
|
||||
@@ -564,6 +565,7 @@ impl TestHarnessBuilder {
|
||||
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: std::sync::Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
TestHarness {
|
||||
|
||||
@@ -0,0 +1,563 @@
|
||||
//! Mock Abound API tools for demo purposes.
|
||||
//!
|
||||
//! These tools simulate Abound's backend API (account info, wire transfers,
|
||||
//! exchange rates, notifications, forex scoring) with realistic mock data.
|
||||
//! They are feature-gated behind `--features demo` and will be replaced by
|
||||
//! real WASM tools once Abound's backend is live.
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{Datelike, Utc};
|
||||
use rand::Rng;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool 1: Get Account Info
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Returns mock Abound account data (limits, recipients, funding sources).
|
||||
pub struct AboundGetAccountInfoTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AboundGetAccountInfoTool {
|
||||
fn name(&self) -> &str {
|
||||
"abound_get_account_info"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Retrieve the authenticated user's Abound account information including \
|
||||
transfer limits, payment reasons, recipients, and funding sources."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let data = json!({
|
||||
"status": "success",
|
||||
"data": {
|
||||
"user_id": "acc_123456",
|
||||
"user_name": "John Doe",
|
||||
"limits": {
|
||||
"ach_limit": {
|
||||
"limit": 5000,
|
||||
"formatted_limit": "$5,000"
|
||||
}
|
||||
},
|
||||
"payment_reasons": [
|
||||
{ "key": "FAMILY_MAINTENANCE", "value": "Family Maintenance" },
|
||||
{ "key": "GIFT", "value": "Gift" },
|
||||
{ "key": "EDUCATION_SUPPORT", "value": "Education Support" },
|
||||
{ "key": "MEDICAL_SUPPORT", "value": "Medical Support" }
|
||||
],
|
||||
"recipients": [
|
||||
{
|
||||
"beneficiary_ref_id": "ben_001",
|
||||
"name": "Rahul Sharma",
|
||||
"mask": "****2222"
|
||||
}
|
||||
],
|
||||
"funding_sources": [
|
||||
{
|
||||
"funding_source_id": "fs_001",
|
||||
"bank_name": "HDFC Bank",
|
||||
"mask": "****2222"
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(data, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool 2: Get Exchange Rate
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Returns mock USD/INR exchange rate with slight randomization.
|
||||
pub struct AboundGetExchangeRateTool;
|
||||
|
||||
/// Generate a mock USD/INR rate with slight jitter around 85.42.
|
||||
fn mock_exchange_rate() -> f64 {
|
||||
let mut rng = rand::thread_rng();
|
||||
let jitter: f64 = rng.gen_range(-0.30..=0.30);
|
||||
((85.42 + jitter) * 100.0).round() / 100.0
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AboundGetExchangeRateTool {
|
||||
fn name(&self) -> &str {
|
||||
"abound_get_exchange_rate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Get the current USD to INR exchange rate including the effective rate \
|
||||
after fees. Use this before initiating any wire transfer."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let rate = mock_exchange_rate();
|
||||
let effective = ((rate - 0.32) * 100.0).round() / 100.0;
|
||||
|
||||
let data = json!({
|
||||
"status": "success",
|
||||
"data": {
|
||||
"from_currency": "USD",
|
||||
"to_currency": "INR",
|
||||
"current_exchange_rate": {
|
||||
"formatted_value": format!("{rate:.2}"),
|
||||
"value": rate
|
||||
},
|
||||
"effective_exchange_rate": {
|
||||
"formatted_value": format!("{effective:.2}"),
|
||||
"value": effective
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(data, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool 3: Send Wire
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Simulates a wire transfer. Requires approval before execution.
|
||||
pub struct AboundSendWireTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AboundSendWireTool {
|
||||
fn name(&self) -> &str {
|
||||
"abound_send_wire"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Submit a wire transfer to send USD to an INR recipient. Requires a \
|
||||
funding source, beneficiary, amount in USD, and payment reason. \
|
||||
The transfer amount must not exceed the user's ACH limit."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"funding_source_id": {
|
||||
"type": "string",
|
||||
"description": "Funding source ID (e.g. 'fs_001')"
|
||||
},
|
||||
"beneficiary_ref_id": {
|
||||
"type": "string",
|
||||
"description": "Beneficiary reference ID (e.g. 'ben_001')"
|
||||
},
|
||||
"amount": {
|
||||
"type": "number",
|
||||
"description": "Amount in USD to send"
|
||||
},
|
||||
"payment_reason_key": {
|
||||
"type": "string",
|
||||
"description": "Payment reason key: FAMILY_MAINTENANCE, GIFT, EDUCATION_SUPPORT, or MEDICAL_SUPPORT"
|
||||
}
|
||||
},
|
||||
"required": ["funding_source_id", "beneficiary_ref_id", "amount", "payment_reason_key"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let _funding_source = require_str(¶ms, "funding_source_id")?;
|
||||
let _beneficiary = require_str(¶ms, "beneficiary_ref_id")?;
|
||||
let _reason = require_str(¶ms, "payment_reason_key")?;
|
||||
|
||||
let amount = params
|
||||
.get("amount")
|
||||
.and_then(|v| v.as_f64())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("missing or invalid 'amount' parameter".to_string())
|
||||
})?;
|
||||
|
||||
// Enforce mock ACH limit
|
||||
if amount > 5000.0 {
|
||||
let data = json!({
|
||||
"status": "error",
|
||||
"error": {
|
||||
"code": "TRANSFER_NOT_ALLOWED",
|
||||
"message": format!(
|
||||
"Transfer amount ${:.2} exceeds your ACH limit of $5,000.00",
|
||||
amount
|
||||
)
|
||||
}
|
||||
});
|
||||
return Ok(ToolOutput::success(data, start.elapsed()));
|
||||
}
|
||||
|
||||
if amount <= 0.0 {
|
||||
return Err(ToolError::InvalidParameters(
|
||||
"amount must be greater than zero".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let txn_id = uuid::Uuid::new_v4();
|
||||
let trk_id = uuid::Uuid::new_v4();
|
||||
|
||||
let data = json!({
|
||||
"status": "success",
|
||||
"data": {
|
||||
"transaction_id": format!("txn_{}", &txn_id.to_string()[..8]),
|
||||
"tracking_id": format!("trk_{}", &trk_id.to_string()[..8]),
|
||||
"amount_usd": amount,
|
||||
"completion_time": {
|
||||
"min_calendar_days": 1,
|
||||
"min_business_days": 1,
|
||||
"max_calendar_days": 3,
|
||||
"max_business_days": 2
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(data, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool 4: Create Notification
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Simulates sending a notification to the Abound system.
|
||||
pub struct AboundCreateNotificationTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AboundCreateNotificationTool {
|
||||
fn name(&self) -> &str {
|
||||
"abound_create_notification"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create a notification in the Abound app (e.g. rate alert, transfer \
|
||||
confirmation, forex scoring signal). Returns 202 accepted."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message_id": {
|
||||
"type": "string",
|
||||
"description": "Unique message identifier"
|
||||
},
|
||||
"action_type": {
|
||||
"type": "string",
|
||||
"description": "Notification type: 'notification' or 'token_refresh'"
|
||||
},
|
||||
"meta_data": {
|
||||
"type": "object",
|
||||
"description": "Additional metadata (e.g. score, rate, signal)"
|
||||
}
|
||||
},
|
||||
"required": ["message_id", "action_type"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let message_id = require_str(¶ms, "message_id")?;
|
||||
|
||||
let data = json!({
|
||||
"status": "accepted",
|
||||
"message": "Notification request accepted for processing",
|
||||
"data": {
|
||||
"message_id": message_id
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(data, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool 5: Forex Score
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Seasonal bias factors by month (1-indexed: Jan=1 .. Dec=12).
|
||||
/// High bias in Oct-Mar (favorable remittance window), low Apr-Sep.
|
||||
const SEASONAL_BIAS: [f64; 13] = [
|
||||
0.0, // placeholder for 0-index
|
||||
0.75, // Jan
|
||||
0.70, // Feb
|
||||
0.65, // Mar
|
||||
0.35, // Apr
|
||||
0.30, // May
|
||||
0.25, // Jun
|
||||
0.25, // Jul
|
||||
0.30, // Aug
|
||||
0.35, // Sep
|
||||
0.70, // Oct
|
||||
0.75, // Nov
|
||||
0.65, // Dec
|
||||
];
|
||||
|
||||
/// Computes a forex timing score for USD/INR, biased toward interesting
|
||||
/// signals (60-80 range) for demo purposes.
|
||||
pub struct AboundGetForexScoreTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AboundGetForexScoreTool {
|
||||
fn name(&self) -> &str {
|
||||
"abound_get_forex_score"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Compute a forex timing score (0-100) for USD/INR transfers. Returns \
|
||||
a score with a signal: 'convert_now' (>=60, good time to send), \
|
||||
'split_transfer' (40-59, send half now), or 'wait' (<40, hold off). \
|
||||
Use this to advise users on optimal transfer timing."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let mut rng = rand::thread_rng();
|
||||
|
||||
// Mock current rate and MA50
|
||||
let rate = mock_exchange_rate();
|
||||
let ma50_jitter: f64 = rng.gen_range(-0.20..=0.20);
|
||||
let ma50 = ((84.50 + ma50_jitter) * 100.0).round() / 100.0;
|
||||
|
||||
// Get seasonal bias for current month
|
||||
let month = Utc::now().month() as usize;
|
||||
let month_bias = SEASONAL_BIAS[month.clamp(1, 12)];
|
||||
|
||||
// Scoring weights
|
||||
let w_ma = 0.7;
|
||||
let w_s = 0.3;
|
||||
|
||||
// MA-based signal: how far current rate is above the 50-day average
|
||||
let ma_signal = (50.0 + ((rate - ma50) / ma50 * 100.0) * 15.0).clamp(0.0, 100.0);
|
||||
|
||||
// Combined score
|
||||
let raw_score = (ma_signal * w_ma + month_bias * 100.0 * w_s) / (w_ma + w_s);
|
||||
|
||||
// Bias toward 60-80 for demo (blend raw score with a favorable base)
|
||||
let demo_base = 68.0;
|
||||
let score = ((raw_score * 0.4 + demo_base * 0.6) as u32).clamp(55, 85);
|
||||
|
||||
let signal = if score >= 60 {
|
||||
"convert_now"
|
||||
} else if score >= 40 {
|
||||
"split_transfer"
|
||||
} else {
|
||||
"wait"
|
||||
};
|
||||
|
||||
let explanation = match signal {
|
||||
"convert_now" => format!(
|
||||
"The current USD/INR rate of {rate:.2} is above the 50-day moving average \
|
||||
of {ma50:.2}, and seasonal trends are favorable. This is a good time to \
|
||||
convert and send money."
|
||||
),
|
||||
"split_transfer" => format!(
|
||||
"The rate of {rate:.2} is near the 50-day average of {ma50:.2}. Consider \
|
||||
splitting your transfer \u{2014} send half now and hold the rest for a \
|
||||
potentially better rate."
|
||||
),
|
||||
_ => format!(
|
||||
"The current rate of {rate:.2} is below the 50-day average of {ma50:.2}. \
|
||||
Unless urgent, consider waiting for a better rate."
|
||||
),
|
||||
};
|
||||
|
||||
let data = json!({
|
||||
"score": score,
|
||||
"signal": signal,
|
||||
"rate": rate,
|
||||
"ma50": ma50,
|
||||
"month_bias": month_bias,
|
||||
"explanation": explanation
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(data, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::context::JobContext;
|
||||
|
||||
fn test_ctx() -> JobContext {
|
||||
JobContext::with_user("test-user", "test", "mock abound tool test")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_info_returns_valid_json() {
|
||||
let tool = AboundGetAccountInfoTool;
|
||||
let result = tool.execute(json!({}), &test_ctx()).await.unwrap();
|
||||
let data = &result.result["data"];
|
||||
assert_eq!(data["user_id"], "acc_123456");
|
||||
assert_eq!(data["limits"]["ach_limit"]["limit"], 5000);
|
||||
assert_eq!(data["recipients"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(data["funding_sources"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exchange_rate_returns_valid_range() {
|
||||
let tool = AboundGetExchangeRateTool;
|
||||
let result = tool.execute(json!({}), &test_ctx()).await.unwrap();
|
||||
let rate = result.result["data"]["current_exchange_rate"]["value"]
|
||||
.as_f64()
|
||||
.unwrap();
|
||||
// Rate should be within jitter range of base 85.42
|
||||
assert!(rate > 85.0 && rate < 85.8, "rate {rate} out of range");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_wire_succeeds_within_limit() {
|
||||
let tool = AboundSendWireTool;
|
||||
let params = json!({
|
||||
"funding_source_id": "fs_001",
|
||||
"beneficiary_ref_id": "ben_001",
|
||||
"amount": 1000.0,
|
||||
"payment_reason_key": "FAMILY_MAINTENANCE"
|
||||
});
|
||||
let result = tool.execute(params, &test_ctx()).await.unwrap();
|
||||
assert_eq!(result.result["status"], "success");
|
||||
assert!(result.result["data"]["transaction_id"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.starts_with("txn_"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_wire_rejects_over_limit() {
|
||||
let tool = AboundSendWireTool;
|
||||
let params = json!({
|
||||
"funding_source_id": "fs_001",
|
||||
"beneficiary_ref_id": "ben_001",
|
||||
"amount": 6000.0,
|
||||
"payment_reason_key": "FAMILY_MAINTENANCE"
|
||||
});
|
||||
let result = tool.execute(params, &test_ctx()).await.unwrap();
|
||||
assert_eq!(result.result["status"], "error");
|
||||
assert_eq!(result.result["error"]["code"], "TRANSFER_NOT_ALLOWED");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_wire_requires_approval() {
|
||||
let tool = AboundSendWireTool;
|
||||
assert!(matches!(
|
||||
tool.requires_approval(&json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_notification_returns_accepted() {
|
||||
let tool = AboundCreateNotificationTool;
|
||||
let params = json!({
|
||||
"message_id": "msg_001",
|
||||
"action_type": "notification",
|
||||
"meta_data": { "score": 72 }
|
||||
});
|
||||
let result = tool.execute(params, &test_ctx()).await.unwrap();
|
||||
assert_eq!(result.result["status"], "accepted");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forex_score_in_demo_range() {
|
||||
let tool = AboundGetForexScoreTool;
|
||||
// Run multiple times to check range stability
|
||||
for _ in 0..20 {
|
||||
let result = tool.execute(json!({}), &test_ctx()).await.unwrap();
|
||||
let score = result.result["score"].as_u64().unwrap();
|
||||
assert!(
|
||||
(55..=85).contains(&score),
|
||||
"score {score} outside demo range [55, 85]"
|
||||
);
|
||||
let signal = result.result["signal"].as_str().unwrap();
|
||||
assert!(
|
||||
signal == "convert_now" || signal == "split_transfer" || signal == "wait",
|
||||
"unexpected signal: {signal}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
//! Built-in tools that come with the agent.
|
||||
|
||||
#[cfg(feature = "demo")]
|
||||
mod abound;
|
||||
mod echo;
|
||||
pub mod extension_tools;
|
||||
mod file;
|
||||
@@ -17,6 +19,11 @@ pub mod skill_tools;
|
||||
mod time;
|
||||
mod tool_info;
|
||||
|
||||
#[cfg(feature = "demo")]
|
||||
pub use abound::{
|
||||
AboundCreateNotificationTool, AboundGetAccountInfoTool, AboundGetExchangeRateTool,
|
||||
AboundGetForexScoreTool, AboundSendWireTool,
|
||||
};
|
||||
pub use echo::EchoTool;
|
||||
pub use extension_tools::{
|
||||
ExtensionInfoTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool,
|
||||
|
||||
@@ -16,6 +16,11 @@ use crate::skills::registry::SkillRegistry;
|
||||
use crate::tools::builder::{
|
||||
BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder, SoftwareBuilder,
|
||||
};
|
||||
#[cfg(feature = "demo")]
|
||||
use crate::tools::builtin::{
|
||||
AboundCreateNotificationTool, AboundGetAccountInfoTool, AboundGetExchangeRateTool,
|
||||
AboundGetForexScoreTool, AboundSendWireTool,
|
||||
};
|
||||
use crate::tools::builtin::{
|
||||
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
||||
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
||||
@@ -248,6 +253,16 @@ impl ToolRegistry {
|
||||
}
|
||||
self.register_sync(Arc::new(http));
|
||||
|
||||
// Abound demo tools (mock APIs — only compiled with --features demo)
|
||||
#[cfg(feature = "demo")]
|
||||
{
|
||||
self.register_sync(Arc::new(AboundGetAccountInfoTool));
|
||||
self.register_sync(Arc::new(AboundGetExchangeRateTool));
|
||||
self.register_sync(Arc::new(AboundSendWireTool));
|
||||
self.register_sync(Arc::new(AboundCreateNotificationTool));
|
||||
self.register_sync(Arc::new(AboundGetForexScoreTool));
|
||||
}
|
||||
|
||||
tracing::debug!("Registered {} built-in tools", self.count());
|
||||
}
|
||||
|
||||
|
||||
@@ -428,10 +428,8 @@ mod tests {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: None,
|
||||
user_id: "test".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
});
|
||||
c
|
||||
}
|
||||
@@ -442,10 +440,8 @@ mod tests {
|
||||
host: host.to_string(),
|
||||
port,
|
||||
auth_token: None,
|
||||
user_id: "test".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
});
|
||||
c
|
||||
}
|
||||
|
||||
+5
-3
@@ -20,7 +20,6 @@ use crate::agent::scheduler::WorkerMessage;
|
||||
use crate::agent::task::TaskOutput;
|
||||
use crate::channels::web::types::ToolDecisionDto;
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::Error;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::{
|
||||
@@ -28,6 +27,7 @@ use crate::llm::{
|
||||
ToolSelection,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tenant::AdminScope;
|
||||
use crate::tools::execute::process_tool_result;
|
||||
use crate::tools::rate_limiter::RateLimitResult;
|
||||
use crate::tools::{
|
||||
@@ -45,7 +45,7 @@ pub struct WorkerDeps {
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub store: Option<AdminScope>,
|
||||
pub hooks: Arc<HookRegistry>,
|
||||
pub timeout: Duration,
|
||||
pub use_planning: bool,
|
||||
@@ -94,7 +94,7 @@ impl Worker {
|
||||
&self.deps.tools
|
||||
}
|
||||
|
||||
fn store(&self) -> Option<&Arc<dyn Database>> {
|
||||
fn store(&self) -> Option<&AdminScope> {
|
||||
self.deps.store.as_ref()
|
||||
}
|
||||
|
||||
@@ -1158,6 +1158,7 @@ impl<'a> JobDelegate<'a> {
|
||||
Ok(crate::llm::RespondOutput {
|
||||
result: RespondResult::Text(String::new()),
|
||||
usage: crate::llm::TokenUsage::default(),
|
||||
finish_reason: crate::llm::FinishReason::Stop,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1283,6 +1284,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
|
||||
content: reasoning_text,
|
||||
},
|
||||
usage: crate::llm::TokenUsage::default(),
|
||||
finish_reason: crate::llm::FinishReason::ToolUse,
|
||||
});
|
||||
}
|
||||
Ok(_) => {} // empty selections, fall through
|
||||
|
||||
@@ -271,7 +271,6 @@ async def ironclaw_server(
|
||||
"GATEWAY_HOST": "127.0.0.1",
|
||||
"GATEWAY_PORT": str(gateway_port),
|
||||
"GATEWAY_AUTH_TOKEN": AUTH_TOKEN,
|
||||
"GATEWAY_USER_ID": "e2e-web-sender",
|
||||
"HTTP_HOST": "127.0.0.1",
|
||||
"HTTP_PORT": str(http_port),
|
||||
"HTTP_WEBHOOK_SECRET": HTTP_WEBHOOK_SECRET,
|
||||
@@ -371,7 +370,6 @@ async def hosted_oauth_refresh_server(
|
||||
"GATEWAY_HOST": "127.0.0.1",
|
||||
"GATEWAY_PORT": str(gateway_port),
|
||||
"GATEWAY_AUTH_TOKEN": AUTH_TOKEN,
|
||||
"GATEWAY_USER_ID": OWNER_SCOPE_ID,
|
||||
"HTTP_HOST": "127.0.0.1",
|
||||
"HTTP_PORT": str(http_port),
|
||||
"HTTP_WEBHOOK_SECRET": HTTP_WEBHOOK_SECRET,
|
||||
@@ -411,7 +409,6 @@ async def hosted_oauth_refresh_server(
|
||||
yield {
|
||||
"base_url": base_url,
|
||||
"db_path": db_path,
|
||||
"gateway_user_id": OWNER_SCOPE_ID,
|
||||
"mock_llm_url": mock_llm_server,
|
||||
}
|
||||
except TimeoutError:
|
||||
@@ -473,7 +470,6 @@ async def http_channel_server_without_secret(
|
||||
"GATEWAY_HOST": "127.0.0.1",
|
||||
"GATEWAY_PORT": str(gateway_port),
|
||||
"GATEWAY_AUTH_TOKEN": AUTH_TOKEN,
|
||||
"GATEWAY_USER_ID": "e2e-tester",
|
||||
"HTTP_HOST": "127.0.0.1",
|
||||
"HTTP_PORT": str(http_port),
|
||||
"CLI_ENABLED": "false",
|
||||
|
||||
@@ -337,14 +337,14 @@ mod tests {
|
||||
SchedulerDeps {
|
||||
tools: registry.clone(),
|
||||
extension_manager: extension_manager.clone(),
|
||||
store: Some(db.clone()),
|
||||
store: Some(ironclaw::tenant::AdminScope::new(db.clone())),
|
||||
hooks: Arc::new(HookRegistry::new()),
|
||||
},
|
||||
));
|
||||
|
||||
Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db,
|
||||
ironclaw::tenant::AdminScope::new(db),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -448,7 +448,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -527,7 +527,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -614,7 +614,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -723,7 +723,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -866,7 +866,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -1049,7 +1049,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
Arc::clone(&db),
|
||||
ironclaw::tenant::AdminScope::new(Arc::clone(&db)),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -1171,7 +1171,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
@@ -1279,7 +1279,7 @@ mod tests {
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
config,
|
||||
db.clone(),
|
||||
ironclaw::tenant::AdminScope::new(db.clone()),
|
||||
llm,
|
||||
ws,
|
||||
notify_tx,
|
||||
|
||||
@@ -201,6 +201,7 @@ mod tests {
|
||||
sandbox_readiness: ironclaw::agent::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
let gateway = Arc::new(TestChannel::new());
|
||||
|
||||
@@ -42,8 +42,6 @@ const ALICE_USER_ID: &str = "alice";
|
||||
const BOB_USER_ID: &str = "bob";
|
||||
const OWNER_TOKEN: &str = "tok-owner-secret";
|
||||
const OWNER_SCOPE_ID: &str = "owner-scope";
|
||||
const GATEWAY_SENDER_ID: &str = "gateway-sender";
|
||||
|
||||
/// Build a MultiAuthState with two users.
|
||||
fn two_user_auth() -> MultiAuthState {
|
||||
let mut tokens = HashMap::new();
|
||||
@@ -51,6 +49,7 @@ fn two_user_auth() -> MultiAuthState {
|
||||
ALICE_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: ALICE_USER_ID.to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
@@ -58,6 +57,7 @@ fn two_user_auth() -> MultiAuthState {
|
||||
BOB_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: BOB_USER_ID.to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: vec!["shared".to_string()],
|
||||
},
|
||||
);
|
||||
@@ -79,7 +79,10 @@ fn user_echo_app(auth: MultiAuthState) -> Router {
|
||||
.route("/api/whoami/scopes", get(echo_user_with_scopes))
|
||||
.route("/api/action", post(echo_user))
|
||||
.route("/api/chat/events", get(echo_user)) // SSE endpoint (allows query token)
|
||||
.layer(middleware::from_fn_with_state(auth, auth_middleware))
|
||||
.layer(middleware::from_fn_with_state(
|
||||
ironclaw::channels::web::auth::CombinedAuthState::from(auth),
|
||||
auth_middleware,
|
||||
))
|
||||
}
|
||||
|
||||
// ===========================================================================
|
||||
@@ -544,7 +547,6 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: "fallback".to_string(),
|
||||
default_sender_id: "fallback".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -558,10 +560,10 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
startup_time: std::time::Instant::now(),
|
||||
webhook_rate_limiter: RateLimiter::new(10, 60),
|
||||
active_config: Default::default(),
|
||||
secrets_store: None,
|
||||
};
|
||||
|
||||
assert_eq!(state.owner_id, "fallback");
|
||||
assert_eq!(state.default_sender_id, "fallback");
|
||||
assert!(state.workspace_pool.is_none());
|
||||
}
|
||||
|
||||
@@ -592,6 +594,7 @@ async fn start_owner_scoped_sender_server() -> (
|
||||
OWNER_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: OWNER_SCOPE_ID.to_string(),
|
||||
role: "admin".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
@@ -599,6 +602,7 @@ async fn start_owner_scoped_sender_server() -> (
|
||||
BOB_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: BOB_USER_ID.to_string(),
|
||||
role: "member".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
@@ -618,7 +622,6 @@ async fn start_owner_scoped_sender_server() -> (
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: OWNER_SCOPE_ID.to_string(),
|
||||
default_sender_id: GATEWAY_SENDER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -632,9 +635,10 @@ async fn start_owner_scoped_sender_server() -> (
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: Default::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let auth = MultiAuthState::multi(tokens);
|
||||
let auth = MultiAuthState::multi(tokens).into();
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound = start_server(addr, state.clone(), auth)
|
||||
.await
|
||||
@@ -769,7 +773,7 @@ async fn full_server_chat_send_rewrites_sender_only_for_owner_scope_rebind() {
|
||||
.expect("Timed out waiting for owner message")
|
||||
.expect("Agent channel closed");
|
||||
assert_eq!(owner_msg.user_id, OWNER_SCOPE_ID);
|
||||
assert_eq!(owner_msg.sender_id, GATEWAY_SENDER_ID);
|
||||
assert_eq!(owner_msg.sender_id, OWNER_SCOPE_ID);
|
||||
assert_eq!(owner_msg.content, "hello from owner");
|
||||
|
||||
let other_resp = client
|
||||
@@ -1003,7 +1007,6 @@ async fn start_multi_user_server_with_db() -> (
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: ALICE_USER_ID.to_string(),
|
||||
default_sender_id: ALICE_USER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -1017,10 +1020,11 @@ async fn start_multi_user_server_with_db() -> (
|
||||
startup_time: std::time::Instant::now(),
|
||||
webhook_rate_limiter: RateLimiter::new(10, 60),
|
||||
active_config: Default::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound = ironclaw::channels::web::server::start_server(addr, state.clone(), auth)
|
||||
let bound = ironclaw::channels::web::server::start_server(addr, state.clone(), auth.into())
|
||||
.await
|
||||
.expect("Failed to start server with DB");
|
||||
|
||||
|
||||
@@ -205,7 +205,6 @@ async fn start_test_server_with_provider(
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(llm_provider),
|
||||
@@ -219,6 +218,7 @@ async fn start_test_server_with_provider(
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: ironclaw::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let auth = ironclaw::channels::web::auth::MultiAuthState::single(
|
||||
@@ -226,7 +226,7 @@ async fn start_test_server_with_provider(
|
||||
"test-user".to_string(),
|
||||
);
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound_addr = start_server(addr, state.clone(), auth)
|
||||
let bound_addr = start_server(addr, state.clone(), auth.into())
|
||||
.await
|
||||
.expect("Failed to start test server");
|
||||
|
||||
@@ -704,7 +704,6 @@ async fn test_no_llm_provider_returns_503() {
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None, // No LLM!
|
||||
@@ -718,6 +717,7 @@ async fn test_no_llm_provider_returns_503() {
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: ironclaw::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let auth = ironclaw::channels::web::auth::MultiAuthState::single(
|
||||
@@ -725,7 +725,7 @@ async fn test_no_llm_provider_returns_503() {
|
||||
"test-user".to_string(),
|
||||
);
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound_addr = start_server(addr, state, auth).await.unwrap();
|
||||
let bound_addr = start_server(addr, state, auth.into()).await.unwrap();
|
||||
|
||||
let url = format!("http://{}/v1/chat/completions", bound_addr);
|
||||
let resp = client()
|
||||
@@ -763,7 +763,7 @@ async fn test_chat_completions_body_too_large() {
|
||||
post(ironclaw::channels::web::openai_compat::chat_completions_handler),
|
||||
)
|
||||
.route_layer(middleware::from_fn_with_state(
|
||||
auth_state,
|
||||
ironclaw::channels::web::auth::CombinedAuthState::from(auth_state),
|
||||
ironclaw::channels::web::auth::auth_middleware,
|
||||
))
|
||||
.layer(DefaultBodyLimit::max(10 * 1024 * 1024))
|
||||
|
||||
@@ -227,7 +227,6 @@ impl GatewayWorkflowHarness {
|
||||
prompt_queue: None,
|
||||
scheduler: Some(scheduler_slot.clone()),
|
||||
owner_id: user_id.clone(),
|
||||
default_sender_id: user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(Arc::clone(&components.llm)),
|
||||
@@ -241,6 +240,7 @@ impl GatewayWorkflowHarness {
|
||||
routine_engine: Arc::clone(&routine_slot),
|
||||
startup_time: Instant::now(),
|
||||
active_config: ironclaw::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let mut agent = Agent::new(
|
||||
@@ -266,6 +266,7 @@ impl GatewayWorkflowHarness {
|
||||
sandbox_readiness: ironclaw::agent::SandboxReadiness::DisabledByConfig,
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
|
||||
},
|
||||
channels,
|
||||
None,
|
||||
@@ -298,7 +299,7 @@ impl GatewayWorkflowHarness {
|
||||
let addr = start_server(
|
||||
"127.0.0.1:0".parse().expect("valid localhost addr"),
|
||||
Arc::clone(&gateway_state),
|
||||
auth,
|
||||
auth.into(),
|
||||
)
|
||||
.await
|
||||
.expect("failed to start gateway server");
|
||||
|
||||
@@ -642,7 +642,7 @@ impl TestRigBuilder {
|
||||
let (notify_tx, _notify_rx) = tokio::sync::mpsc::channel(16);
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
routine_config,
|
||||
Arc::clone(db_arc),
|
||||
ironclaw::tenant::AdminScope::new(Arc::clone(db_arc)),
|
||||
components.llm.clone(),
|
||||
Arc::clone(ws),
|
||||
notify_tx,
|
||||
@@ -762,6 +762,7 @@ impl TestRigBuilder {
|
||||
sandbox_readiness: ironclaw::agent::SandboxReadiness::Available, // tests don't use real Docker
|
||||
builder: None,
|
||||
llm_backend: "nearai".to_string(),
|
||||
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
|
||||
};
|
||||
|
||||
// 7. Create TestChannel and ChannelManager.
|
||||
|
||||
@@ -52,7 +52,6 @@ async fn start_test_server() -> (
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -66,6 +65,7 @@ async fn start_test_server() -> (
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: ironclaw::channels::web::server::ActiveConfigSnapshot::default(),
|
||||
secrets_store: None,
|
||||
});
|
||||
|
||||
let auth = ironclaw::channels::web::auth::MultiAuthState::single(
|
||||
@@ -73,7 +73,7 @@ async fn start_test_server() -> (
|
||||
"test-user".to_string(),
|
||||
);
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound_addr = start_server(addr, state.clone(), auth)
|
||||
let bound_addr = start_server(addr, state.clone(), auth.into())
|
||||
.await
|
||||
.expect("Failed to start test server");
|
||||
|
||||
|
||||
Reference in New Issue
Block a user