gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-09-01 17:19:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Sandbox jobs (#4)

Browse Source 
* Orchestrating jobs and running them in sandboxes

* Fix heartbeat: dynamic max_tokens, empty content guard, notification fallback

- Query /v1/models API for context_length and set max_tokens to half
  (floor 4096) instead of hardcoded 1024; reasoning models like GLM-4.7
  need much larger budgets
- Guard against empty LLM content (reasoning models can burn all tokens
  on chain-of-thought and return content: null)
- Simplify notification routing: try configured channel first, fall back
  to broadcast_all so heartbeat alerts always reach someone
- Add ModelMetadata struct and model_metadata() to LlmProvider trait
- Refactor NearAiChatProvider::list_models into shared fetch_models()
- Add standalone test_heartbeat example for isolated debugging

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

* Add job detail view with drill-down from jobs list

Click a job row to see full details across four sub-tabs:
Overview (metadata grid, description, state transitions timeline),
Actions (expandable tool call cards with input/output JSON),
Thinking (conversation messages styled by role), and
Files (embedded workspace tree browser).

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

* Strip model-internal XML tags from LLM responses, fix Telegram parse_mode 400

Some models (GLM-4.7, etc.) emit <tool_call>tool_list</tool_call> in the
content field instead of using the OpenAI tool_calls array. This XML leaks
through to channels as text, and Telegram's Markdown parser chokes on the
underscores, returning 400 "can't parse entities".

Two fixes:
- Generalize clean_response() to strip <tool_call>, <function_call>,
  <tool_calls>, and pipe-delimited variants (<|tool_call|>) alongside
  the existing <thinking> tag stripping
- Add Telegram send_message helper with parse_mode fallback: try Markdown
  first, retry as plain text on "can't parse entities" 400 errors

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

* Add SystemCommand submission type for thread-state-independent commands

System commands (/help, /model, /version, /tools, /ping, /debug) now
bypass thread-state checks and safety validation via a dedicated
Submission::SystemCommand variant. Previously these flowed through
process_user_input() which blocked them during Processing/AwaitingApproval
/Completed states.

- Add /model [name] for runtime model switching with provider validation
- Add active_model_name()/set_model() to LlmProvider trait with RwLock
  hot-swap in both NEAR AI providers
- Rewrite /help with aligned columns grouped by category
- Expand REPL tab-completion from 10 to 23 slash commands
- Remove REPL-local /help interception (now handled by agent)

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

* Add per-tool execution timeouts, auto-create sandbox project dirs, serve built files

The sandbox e2e pipeline (agent -> container -> built website -> browsable URL)
was broken by three gaps: hardcoded 60s timeouts killed sandbox jobs that need
minutes, no auto-created project directory meant container output vanished, and
no HTTP route to browse the built files.

- Add `execution_timeout()` to the `Tool` trait (default 60s), replace all four
  hardcoded `Duration::from_secs(60)` call sites (agent_loop, worker, scheduler,
  worker/runtime) with the per-tool value
- Override to 660s in `RunInSandboxTool` (10 min polling + 60s buffer)
- Auto-create `~/.ironclaw/projects/{uuid}/` when no `project_dir` is specified,
  so every sandbox job gets a persistent bind mount
- Include `project_dir` and `browse_url` in sandbox tool output JSON
- Add `/projects/{id}` and `/projects/{id}/{path}` static file serving routes
  to the web gateway with path traversal protection and MIME type detection
- Add `mime_guess` dependency for content-type detection

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

* Apply cargo fmt to wizard.rs after merge

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

* Persist sandbox jobs in DB, fix web UI, unify job model

Sandbox container jobs were invisible to the web UI because they lived
only in ContainerJobManager's in-memory HashMap while the API queried
ContextManager. This persists them to the agent_jobs table and fixes
all six front-end bugs (empty job list, broken back button, empty
actions/thinking tabs, wrong files tab, stuck status, no persistence).

Key changes:
- V4 migration adds project_dir and user_id columns to agent_jobs
- Embedded migrations via refinery (no external CLI needed)
- SandboxJobRecord CRUD in Store with fire-and-forget DB writes
- Unified job_id: sandbox tool generates UUID, passes to ContainerJobManager
- Web API queries DB for sandbox jobs, merges with ContextManager direct jobs
- New endpoints: restart, project file list/read with path traversal protection
- Front-end: rebuild DOM on back navigation, sandbox-aware tabs, job cards in
  chat stream, source badges, restart button for failed/interrupted jobs
- Gateway defaults to enabled, prints Web UI URL on startup
- Stale jobs marked "interrupted" on restart for visibility and restartability

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

* Secure in-chat auth: tokens never touch the LLM or chat history

Remove the token parameter from tool_auth so the LLM cannot pass raw
API keys. Add dedicated REST (POST /api/chat/auth-token) and WebSocket
(auth_token) endpoints that route tokens directly to ext_mgr.auth(),
completely bypassing the message pipeline, turns, history, and compaction.

Web UI shows an auth card (password input + OAuth button) when the agent
enters auth mode, submitted via the dedicated endpoint. CLI auth mode
interception is unchanged (already secure).

New StatusUpdate::AuthRequired/AuthCompleted variants propagate through
all channels (SSE, WebSocket, REPL, WASM).

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

* feat: Add Claude Code mode for sandbox jobs

Run Claude Code CLI inside Docker containers as an alternative to the
standard worker mode. The bridge spawns `claude -p` with stream-json
output, posts events to the orchestrator, and supports follow-up
prompts via `--resume`.

Key additions:
- `claude-bridge` CLI subcommand and ClaudeBridgeRuntime
- JobMode enum (Worker vs ClaudeCode) with per-mode container config
- Orchestrator endpoints for Claude events and prompt polling
- SSE event variants for real-time Claude Code streaming to frontend
- Claude Code sub-tab in web UI with terminal-style output and input bar
- Database migration for job_mode column and claude_code_events table
- ClaudeCodeConfig with env var support (CLAUDE_CODE_ENABLED, etc.)
- Mode parameter on run_in_sandbox tool schema

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

* fix: Skip create_job tool when sandbox is enabled to prevent duplicate jobs

When sandbox mode is on, the LLM would call create_job (creating a
pending "direct" entry) then run_in_sandbox (creating a second "sandbox"
entry), producing two jobs in the list for a single user request.

Now register_job_tools() skips create_job when sandbox is enabled since
run_in_sandbox already creates tracked jobs. Also improved the
run_in_sandbox description to guide the LLM to use it directly and to
mention wait=false for async execution.

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

* feat: Web gateway UI quality-of-life improvements

Phase 1: Send button disabled state to prevent double-sends, copy button
on code blocks, confirm() guards on destructive actions, SSE-driven job
list auto-refresh, log filters re-applied on tab switch, jobEvents memory
leak fix (cap at 500, cleanup after 60s).

Phase 2: Toast notification system replacing chat-based system messages,
memory search highlighting with centered snippets, keyboard shortcuts
(Ctrl+1-5 tabs, Ctrl+K focus, Ctrl+N new thread, Escape close/blur),
activity tab toolbar with event type filter and auto-scroll toggle.

Phase 3: Thread sidebar with load/switch/create, thread_id passed with
messages, collapsible to hamburger. Memory inline editing with textarea,
Save/Cancel, POST to /api/memory/write.

Phase 4: Gateway status popover on hover (polls every 30s), extension
install form (name/URL/kind), markdown rendering in memory viewer for
.md files, mobile responsive layout at 768px breakpoint.

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

* feat: Add routines system, remove non-sandbox job mode from web UI

Routines: scheduled & reactive job system with cron and event triggers,
lightweight (single LLM call) and full-job execution modes, guardrails
(cooldown, max concurrent, dedup), and LLM-facing tools for CRUD.

Web UI: remove ContextManager-backed "direct" job mode entirely. Jobs
are now exclusively sandbox-backed (DB + container). Simplify job detail
response, drop dead types (ActionInfo, MessageInfo, MessageToolCallInfo),
fix Browse Files CSS loading (trailing-slash redirect), fix Activity tab
event rendering.

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

* fix: Re-enable chat input on agent completion, auto-auth on tool_activate, recover tool calls from content XML

Three fixes:

1. Chat input stays disabled after agent finishes: the "Done" status
   SSE event now calls enableChatInput() as a safety net when the
   response event is empty or lost. Same for auth_completed and
   cancelAuth().

2. tool_activate never triggers auth: when activation fails due to
   missing authentication, it now auto-initiates the auth flow
   (same pattern as the web API handler). detect_auth_awaiting()
   also matches tool_activate results now.

3. Models like GLM-4.7 emit tool calls as XML tags in content
   (<tool_call>tool_list</tool_call>) instead of using the structured
   tool_calls array. recover_tool_calls_from_content() extracts and
   validates these before falling back to plain text.

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

* feat: Add routines web UI tab, update docs for sandbox-jobs branch

Add full routines management to the web gateway (list, detail, trigger,
toggle, delete) with 7 new API endpoints, response types, and frontend
(HTML, JS, CSS). Update FEATURE_PARITY.md (~23 rows), CLAUDE.md (new
subsystems, config, TODOs), and README.md (architecture diagram,
features, components, fix onboard command).

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

* fix: Bind Telegram bot to owner account during setup

Without owner binding, anyone who discovers the bot can send it messages.
The setup wizard now prompts the user to message their bot, captures their
Telegram user ID via getUpdates, and persists it as telegram_owner_id in
settings. On startup, the owner_id is injected into the WASM channel config
so the existing owner restriction logic drops messages from non-owners.

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

* feat: Move settings from disk to PostgreSQL database

Settings previously lived in three JSON files on disk (settings.json,
mcp-servers.json, session.json). This made them inaccessible from the
web UI and caused redundant disk reads (Settings::load() called 8+
times during startup).

Now all settings live in a `settings` table (user_id + key -> JSONB)
with only 4 bootstrap fields remaining on disk (database_url, pool
size, secrets key source, onboard_completed) since they're needed
before the DB connection exists.

- Add V8 migration for settings table
- Add BootstrapConfig (thin disk file) and Settings DB round-trip
- Add Store CRUD methods for settings (get/set/delete/list/bulk)
- Refactor Config to load from DB (env > DB > default cascade)
- Add SessionManager DB persistence for session tokens
- Add DB-backed MCP server config load/save functions
- Add 6 settings web API endpoints (list/get/set/delete/export/import)
- Add one-time disk-to-DB migration on first boot
- Make CLI config commands async with DB access (disk fallback)

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

* feat: Seed workspace on boot, fix gateway duplicate logs and URL auto-auth

- Add Workspace::seed_if_empty() to create core identity files (README,
  MEMORY, IDENTITY, SOUL, AGENTS, USER, HEARTBEAT) when missing, called
  on every boot without overwriting existing user edits
- Remove duplicate gateway log lines from web/mod.rs (main.rs has the
  useful clickable ?token= URL)
- Auto-authenticate from ?token= URL parameter in the web UI and strip
  the token from the address bar after successful auth

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

* fix: Harden sandbox security (path traversal + orchestrator auth)

Two vulnerabilities fixed:

1. project_dir path traversal: The create_job tool let the LLM specify
   arbitrary host paths for Docker bind mounts. Removed project_dir from
   the tool schema entirely, and added canonicalization + prefix validation
   at both resolve_project_dir() and the job_manager bind mount point.

2. Orchestrator API auth bypass: worker_auth_middleware was defined but
   never applied. Each handler manually called validate_token(), so any
   new endpoint that forgot would be publicly accessible. Applied the
   middleware as route_layer on all /worker/ routes, removed manual auth
   from all 7 handlers. Bind to 127.0.0.1 on macOS/Windows (Linux keeps
   0.0.0.0 since containers reach host via docker bridge, not loopback).

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

* feat: Rework gateway chat with pinned assistant, pagination, and NEAR AI response chaining

Implements the 4-phase plan for overhauling the web gateway chat:

- Phase 1: Pinned "Assistant" thread at top of sidebar, regular threads below
- Phase 2: Cursor-based history pagination with infinite scroll
- Phase 3: NEAR AI previous_response_id chaining (delta-only messages),
  with fallback to full history on chain errors, and DB persistence of
  chain state across restarts
- Phase 4: SSE thread isolation (events filtered by thread_id)

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

* fix: Add per-request HTTP timeout to WASM host, redact credentials in errors

Three fixes for WASM channel reliability:

1. Per-request timeout: Add optional timeout-ms parameter to http-request
   in both channel and tool WIT interfaces. Telegram long-poll now specifies
   35s (outliving the 30s server-side hold), while regular API calls use
   the 30s default. Fixes the triple-30s timeout race that caused polling
   failures.

2. Credential redaction: reqwest::Error includes the full URL (with injected
   bot tokens) in its Display output. Scrub credential values from error
   messages before logging or returning to WASM.

3. Webhook route registration: Remove tunnel URL gate so webhook routes are
   always available when webhook channels exist, not only when TUNNEL_URL
   is configured.

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

* chore: Fix clippy warnings in WASM tools and channels

- slack channel: allow dead_code on signing_secret_name (forward compat field)
- gmail tool: use div_ceil() instead of manual (n+2)/3
- google-calendar tool: extract CreateEventParams/UpdateEventParams structs
  to fix too-many-arguments warnings

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

* Fix approval flow

* fix: Rebuild bundled telegram.wasm with updated WIT interface

The bundled WASM binary must match the host's WIT definition.
Previous binary was compiled against the old 4-arg http-request;
this rebuild includes the new timeout-ms parameter.

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

* refactor: Load WASM channels from disk instead of bundling in binary

Remove include_bytes! embedding of telegram.wasm. Channels are now
loaded from their build output directories (channels-src/<name>/target/)
during onboarding, then from ~/.ironclaw/channels/ at runtime.

- bundled.rs: locate_channel_artifacts() finds WASM + capabilities from
  build output; IRONCLAW_CHANNELS_SRC env var overrides the default path
- available_channel_names(): only lists channels with build artifacts
- bundled_channel_names(): lists all known channels (manifest)
- Setup wizard uses available_channel_names() to offer installable channels
- Add *.wasm to .gitignore, remove tracked telegram.wasm

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

* fix: Persist gateway auth token, fix thread hydration race, polish auth screen

Three web gateway UX fixes:

1. Token persistence: Store auth token in sessionStorage so refreshing
   the page doesn't force re-authentication. Hide the auth screen
   immediately when a saved token exists to prevent flash.

2. Thread hydration: Remove the !msgs.is_empty() bail-out in
   maybe_hydrate_thread so that even brand-new (empty) assistant threads
   get hydrated with their correct DB UUID. Previously resolve_thread
   would mint a fresh UUID, causing messages to land in the wrong
   conversation and duplicate threads to appear.

3. Auth screen: Redesign as a centered card with brand, tagline, labeled
   input, and hint text.

Also adds 34 new tests covering session/thread lifecycle, thread
resolution isolation (user, channel, external ID), hydration edge cases,
serialization round-trips, approval flows, and stale mapping recovery.

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

* fix: Use bindgen! for WASM tool wrapper, add dev tool loading

Three changes:

1. Rewrite src/tools/wasm/wrapper.rs to use wasmtime::component::bindgen!
   instead of manual linker.root().func_wrap(). This fixes the
   "component imports instance 'near:agent/host', but a matching
   implementation was not found in the linker" error. All 6 host functions
   (log, now-millis, workspace-read, http-request, secret-exists,
   tool-invoke) are now properly registered under the near:agent/host
   namespace. Also adds WASI support, credential injection, and leak
   detection for HTTP requests made by WASM tools.

2. Add dev tool loading to src/tools/wasm/loader.rs. During startup, the
   loader now also scans tools-src/*/target/wasm32-wasip2/release/ for
   build artifacts that are newer than installed copies. This means during
   development you just rebuild the WASM and restart the host; no manual
   copy step needed. Set IRONCLAW_TOOLS_SRC to override the source dir.

3. Wire up load_dev_tools() in main.rs alongside the existing
   load_from_dir() call.

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

* feat: Wire main startup and CLI to use DB-backed settings

main.rs now reloads Config from the database after connecting,
attaches the store to the session manager for dual-write tokens,
and loads MCP servers from DB instead of disk. ExtensionManager
and MCP CLI commands use DB when available with disk fallback.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-11 08:31:25 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 202665a55c
commit ced83d5b4d
99 changed files with 17095 additions and 1162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llm/mod.rs
+2 -2
View File
@@ -13,8 +13,8 @@ pub mod session;
pub use nearai::{ModelInfo, NearAiProvider};
pub use nearai_chat::NearAiChatProvider;
pub use provider::{
ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
};
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
pub use session::{SessionConfig, SessionManager, create_session_manager};
src/llm/nearai.rs
+308 -20
View File
@@ -3,6 +3,7 @@
//! This provider uses the NEAR AI chat-api which provides a unified interface
//! to multiple LLM models (OpenAI, Anthropic, etc.) with user authentication.
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
@@ -31,11 +32,23 @@ pub struct ModelInfo {
pub provider: Option<String>,
}
/// Per-thread chaining state: the last response ID and how many input
/// messages were included in that request. This lets subsequent calls send
/// only the delta (new messages since last call).
struct ChainState {
response_id: String,
input_count: usize,
}
/// NEAR AI Chat API provider.
pub struct NearAiProvider {
client: Client,
config: NearAiConfig,
session: Arc<SessionManager>,
active_model: std::sync::RwLock<String>,
/// Per-thread response ID chaining state.
/// Key is thread_id from request metadata.
response_chains: std::sync::RwLock<HashMap<String, ChainState>>,
}
impl NearAiProvider {
@@ -46,13 +59,64 @@ impl NearAiProvider {
.build()
.unwrap_or_else(|_| Client::new());
let active_model = std::sync::RwLock::new(config.model.clone());
Self {
client,
config,
session,
active_model,
response_chains: std::sync::RwLock::new(HashMap::new()),
}
}
/// Seed a response chain for a thread (e.g. when restoring from DB).
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.insert(
thread_id.to_string(),
ChainState {
response_id,
input_count: 0,
},
);
}
/// Get the last response ID for a thread (for persistence).
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
let chains = self
.response_chains
.read()
.expect("response_chains lock poisoned");
chains.get(thread_id).map(|c| c.response_id.clone())
}
/// Store a response chain state after a successful call.
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.insert(
thread_id.to_string(),
ChainState {
response_id,
input_count,
},
);
}
/// Clear the chain for a thread (on error / fallback).
fn clear_chain(&self, thread_id: &str) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.remove(thread_id);
}
fn api_url(&self, path: &str) -> String {
format!(
"{}/v1/{}",
@@ -291,18 +355,34 @@ impl NearAiProvider {
}
}
/// Split messages into system instructions and non-system input messages.
/// Split messages into system instructions and non-system input items.
/// The OpenAI Responses API expects system prompts in an `instructions` field,
/// not as a message with role "system" in the input array.
fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMessage>) {
///
/// When `chaining` is true, tool result messages (role=tool) are converted to
/// `NearAiInputItem::FunctionCallOutput` for the Responses API protocol.
fn split_messages(
messages: Vec<ChatMessage>,
chaining: bool,
) -> (Option<String>, Vec<NearAiInputItem>) {
let mut instructions: Vec<String> = Vec::new();
let mut input: Vec<NearAiMessage> = Vec::new();
let mut input: Vec<NearAiInputItem> = Vec::new();
for msg in messages {
if msg.role == Role::System {
instructions.push(msg.content);
} else if chaining && msg.role == Role::Tool {
if let Some(ref call_id) = msg.tool_call_id {
input.push(NearAiInputItem::FunctionCallOutput {
item_type: "function_call_output".to_string(),
call_id: call_id.clone(),
output: msg.content,
});
} else {
input.push(NearAiInputItem::Message(msg.into()));
}
} else {
input.push(msg.into());
input.push(NearAiInputItem::Message(msg.into()));
}
}
@@ -318,12 +398,14 @@ fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMess
#[async_trait]
impl LlmProvider for NearAiProvider {
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
let (instructions, input) = split_messages(req.messages);
let thread_id = req.metadata.get("thread_id").cloned();
let (instructions, input) = split_messages(req.messages, false);
let request = NearAiRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
instructions,
input,
previous_response_id: None,
temperature: req.temperature,
max_output_tokens: req.max_tokens,
stream: Some(false),
@@ -350,6 +432,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
response_id: None,
});
}
@@ -367,6 +450,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: 0,
output_tokens: 0,
response_id: None,
});
}
Err(e) => return Err(e),
@@ -423,11 +507,17 @@ impl LlmProvider for NearAiProvider {
);
}
// Store response ID for chaining
if let Some(ref tid) = thread_id {
self.store_chain(tid, response.id.clone(), 0);
}
Ok(CompletionResponse {
content: text,
finish_reason: FinishReason::Stop,
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
response_id: Some(response.id),
})
}
@@ -435,7 +525,33 @@ impl LlmProvider for NearAiProvider {
&self,
req: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
let (instructions, input) = split_messages(req.messages);
let thread_id = req.metadata.get("thread_id").cloned();
// Look up chaining state for this thread
let chain_state = thread_id.as_ref().and_then(|tid| {
let chains = self
.response_chains
.read()
.expect("response_chains lock poisoned");
chains
.get(tid)
.map(|c| (c.response_id.clone(), c.input_count))
});
let chaining = chain_state.is_some();
let (previous_response_id, prev_input_count) = chain_state
.map(|(rid, count)| (Some(rid), count))
.unwrap_or((None, 0));
// When chaining, only send new messages (the delta since last call).
// Tool results are converted to function_call_output items.
let (instructions, all_input) = split_messages(req.messages, chaining);
let input = if chaining && all_input.len() > prev_input_count {
all_input[prev_input_count..].to_vec()
} else {
all_input.clone()
};
let total_input_count = all_input.len();
let tools: Vec<NearAiTool> = req
.tools
@@ -449,18 +565,58 @@ impl LlmProvider for NearAiProvider {
.collect();
let request = NearAiRequest {
model: self.config.model.clone(),
instructions,
model: self.active_model_name(),
instructions: if chaining { None } else { instructions.clone() },
input,
previous_response_id: previous_response_id.clone(),
temperature: req.temperature,
max_output_tokens: req.max_tokens,
stream: Some(false),
tools: if tools.is_empty() { None } else { Some(tools) },
tools: if tools.is_empty() {
None
} else {
Some(tools.clone())
},
};
// Try to get structured response, fall back to alternative formats
// Try to get structured response, fall back to alternative formats.
// If chaining fails (bad previous_response_id), retry with full history.
let response: NearAiResponse = match self.send_request("responses", &request).await {
Ok(r) => r,
Err(ref e) if chaining && is_chain_error(e) => {
tracing::warn!(
"Response chaining failed, retrying with full history: {}",
e
);
if let Some(ref tid) = thread_id {
self.clear_chain(tid);
}
let (instructions_full, input_full) = split_messages(
// Rebuild from the original input (non-chaining mode)
{
let mut msgs = Vec::new();
if let Some(ref instr) = instructions {
msgs.push(ChatMessage::system(instr.clone()));
}
for item in &all_input {
msgs.push(item.to_chat_message());
}
msgs
},
false,
);
let retry_request = NearAiRequest {
model: self.active_model_name(),
instructions: instructions_full,
input: input_full,
previous_response_id: None,
temperature: request.temperature,
max_output_tokens: request.max_output_tokens,
stream: Some(false),
tools: request.tools.clone(),
};
self.send_request("responses", &retry_request).await?
}
Err(LlmError::InvalidResponse { reason, .. }) if reason.contains("Raw: ") => {
let raw_text = reason.split("Raw: ").nth(1).unwrap_or("");
@@ -490,6 +646,7 @@ impl LlmProvider for NearAiProvider {
finish_reason,
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
response_id: None,
});
}
@@ -507,6 +664,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: 0,
output_tokens: 0,
response_id: None,
});
}
Err(e) => return Err(e),
@@ -560,12 +718,18 @@ impl LlmProvider for NearAiProvider {
FinishReason::ToolUse
};
// Store response ID for chaining on subsequent calls
if let Some(ref tid) = thread_id {
self.store_chain(tid, response.id.clone(), total_input_count);
}
Ok(ToolCompletionResponse {
content: if text.is_empty() { None } else { Some(text) },
tool_calls,
finish_reason,
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
response_id: Some(response.id),
})
}
@@ -584,6 +748,30 @@ impl LlmProvider for NearAiProvider {
let models = NearAiProvider::list_models(self).await?;
Ok(models.into_iter().map(|m| m.name).collect())
}
fn active_model_name(&self) -> String {
self.active_model
.read()
.expect("active_model lock poisoned")
.clone()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
let mut guard = self
.active_model
.write()
.expect("active_model lock poisoned");
*guard = model.to_string();
Ok(())
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.seed_response_id(thread_id, response_id);
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.get_response_id(thread_id)
}
}
// NEAR AI API types
@@ -597,8 +785,11 @@ struct NearAiRequest {
/// System instructions (replaces sending system role in input)
#[serde(skip_serializing_if = "Option::is_none")]
instructions: Option<String>,
/// Input messages (user/assistant/tool only, NOT system)
input: Vec<NearAiMessage>,
/// Input items: messages and/or function_call_output entries.
input: Vec<NearAiInputItem>,
/// Chain this request to a previous response (avoids resending full context).
#[serde(skip_serializing_if = "Option::is_none")]
previous_response_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -609,7 +800,7 @@ struct NearAiRequest {
tools: Option<Vec<NearAiTool>>,
}
#[derive(Debug, Serialize, Deserialize)]
#[derive(Debug, Serialize, Deserialize, Clone)]
struct NearAiMessage {
role: String,
content: String,
@@ -630,7 +821,68 @@ impl From<ChatMessage> for NearAiMessage {
}
}
#[derive(Debug, Serialize)]
/// Input item for the Responses API. Either a regular message or a
/// function_call_output (for returning tool results when chaining).
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
enum NearAiInputItem {
Message(NearAiMessage),
FunctionCallOutput {
#[serde(rename = "type")]
item_type: String,
call_id: String,
output: String,
},
}
impl NearAiInputItem {
/// Convert back to a ChatMessage (used for fallback retry).
fn to_chat_message(&self) -> ChatMessage {
match self {
NearAiInputItem::Message(msg) => {
let role = match msg.role.as_str() {
"system" => Role::System,
"user" => Role::User,
"assistant" => Role::Assistant,
"tool" => Role::Tool,
_ => Role::User,
};
ChatMessage {
role,
content: msg.content.clone(),
tool_call_id: None,
name: None,
tool_calls: None,
}
}
NearAiInputItem::FunctionCallOutput {
call_id, output, ..
} => ChatMessage {
role: Role::Tool,
content: output.clone(),
tool_call_id: Some(call_id.clone()),
name: None,
tool_calls: None,
},
}
}
}
/// Check if an LLM error is likely caused by an invalid previous_response_id.
fn is_chain_error(err: &LlmError) -> bool {
match err {
LlmError::RequestFailed { reason, .. } => {
let lower = reason.to_lowercase();
lower.contains("previous_response_id")
|| lower.contains("previous response")
|| lower.contains("not found")
|| lower.contains("invalid response id")
}
_ => false,
}
}
#[derive(Debug, Clone, Serialize)]
struct NearAiTool {
#[serde(rename = "type")]
tool_type: String,
@@ -833,14 +1085,17 @@ mod tests {
ChatMessage::user("Hello"),
ChatMessage::assistant("Hi there!"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert_eq!(
instructions,
Some("You are a helpful assistant".to_string())
);
assert_eq!(input.len(), 2);
assert_eq!(input[0].role, "user");
assert_eq!(input[1].role, "assistant");
// Verify the input items are messages
match &input[0] {
NearAiInputItem::Message(m) => assert_eq!(m.role, "user"),
_ => panic!("expected Message"),
}
}
#[test]
@@ -849,7 +1104,7 @@ mod tests {
ChatMessage::user("Hello"),
ChatMessage::assistant("Hi there!"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert!(instructions.is_none());
assert_eq!(input.len(), 2);
}
@@ -861,11 +1116,44 @@ mod tests {
ChatMessage::system("Second instruction"),
ChatMessage::user("Hello"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert_eq!(
instructions,
Some("First instruction\n\nSecond instruction".to_string())
);
assert_eq!(input.len(), 1);
}
#[test]
fn test_split_messages_chaining_converts_tool_results() {
let messages = vec![
ChatMessage::user("Hello"),
ChatMessage::tool_result("call_123", "my_tool", "result data"),
];
let (_, input) = split_messages(messages, true);
assert_eq!(input.len(), 2);
match &input[1] {
NearAiInputItem::FunctionCallOutput {
call_id, output, ..
} => {
assert_eq!(call_id, "call_123");
assert_eq!(output, "result data");
}
_ => panic!("expected FunctionCallOutput"),
}
}
#[test]
fn test_split_messages_no_chaining_keeps_tool_as_message() {
let messages = vec![
ChatMessage::user("Hello"),
ChatMessage::tool_result("call_123", "my_tool", "result data"),
];
let (_, input) = split_messages(messages, false);
assert_eq!(input.len(), 2);
match &input[1] {
NearAiInputItem::Message(m) => assert_eq!(m.role, "tool"),
_ => panic!("expected Message"),
}
}
}
src/llm/nearai_chat.rs
+69 -18
View File
@@ -13,14 +13,15 @@ use serde::{Deserialize, Serialize};
use crate::config::NearAiConfig;
use crate::error::LlmError;
use crate::llm::provider::{
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse,
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
};
/// NEAR AI Chat Completions API provider.
pub struct NearAiChatProvider {
client: Client,
config: NearAiConfig,
active_model: std::sync::RwLock<String>,
}
impl NearAiChatProvider {
@@ -37,7 +38,12 @@ impl NearAiChatProvider {
.build()
.unwrap_or_else(|_| Client::new());
Ok(Self { client, config })
let active_model = std::sync::RwLock::new(config.model.clone());
Ok(Self {
client,
config,
active_model,
})
}
fn api_url(&self, path: &str) -> String {
@@ -65,6 +71,11 @@ impl NearAiChatProvider {
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
// Log the request body for debugging tool call issues
if let Ok(json) = serde_json::to_string(body) {
tracing::debug!("NEAR AI Chat request body: {}", json);
}
let response = self
.client
.post(&url)
@@ -111,8 +122,8 @@ impl NearAiChatProvider {
})
}
/// Fetch available models.
pub async fn list_models(&self) -> Result<Vec<String>, LlmError> {
/// Fetch available models with full metadata from the `/v1/models` endpoint.
async fn fetch_models(&self) -> Result<Vec<ApiModelEntry>, LlmError> {
let url = self.api_url("models");
let response = self
@@ -138,12 +149,7 @@ impl NearAiChatProvider {
#[derive(Deserialize)]
struct ModelsResponse {
data: Vec<ModelEntry>,
}
#[derive(Deserialize)]
struct ModelEntry {
id: String,
data: Vec<ApiModelEntry>,
}
let resp: ModelsResponse =
@@ -152,10 +158,18 @@ impl NearAiChatProvider {
reason: format!("JSON parse error: {}", e),
})?;
Ok(resp.data.into_iter().map(|m| m.id).collect())
Ok(resp.data)
}
}
/// Model entry as returned by the `/v1/models` API.
#[derive(Debug, Deserialize)]
struct ApiModelEntry {
id: String,
#[serde(default)]
context_length: Option<u32>,
}
#[async_trait]
impl LlmProvider for NearAiChatProvider {
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
@@ -163,7 +177,7 @@ impl LlmProvider for NearAiChatProvider {
req.messages.into_iter().map(|m| m.into()).collect();
let request = ChatCompletionRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
@@ -197,6 +211,7 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens: response.usage.prompt_tokens,
output_tokens: response.usage.completion_tokens,
response_id: None,
})
}
@@ -221,7 +236,7 @@ impl LlmProvider for NearAiChatProvider {
.collect();
let request = ChatCompletionRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
@@ -278,6 +293,7 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens: response.usage.prompt_tokens,
output_tokens: response.usage.completion_tokens,
response_id: None,
})
}
@@ -291,7 +307,34 @@ impl LlmProvider for NearAiChatProvider {
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
NearAiChatProvider::list_models(self).await
let models = self.fetch_models().await?;
Ok(models.into_iter().map(|m| m.id).collect())
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
let active = self.active_model_name();
let models = self.fetch_models().await?;
let current = models.iter().find(|m| m.id == active);
Ok(ModelMetadata {
id: active,
context_length: current.and_then(|m| m.context_length),
})
}
fn active_model_name(&self) -> String {
self.active_model
.read()
.expect("active_model lock poisoned")
.clone()
}
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
let mut guard = self
.active_model
.write()
.expect("active_model lock poisoned");
*guard = model.to_string();
Ok(())
}
}
@@ -332,6 +375,7 @@ impl From<ChatMessage> for ChatCompletionMessage {
Role::Assistant => "assistant",
Role::Tool => "tool",
};
let tool_calls = msg.tool_calls.map(|calls| {
calls
.into_iter()
@@ -345,9 +389,16 @@ impl From<ChatMessage> for ChatCompletionMessage {
})
.collect()
});
let content = if role == "assistant" && tool_calls.is_some() && msg.content.is_empty() {
None
} else {
Some(msg.content)
};
Self {
role: role.to_string(),
content: Some(msg.content),
content,
tool_call_id: msg.tool_call_id,
name: msg.name,
tool_calls,
@@ -454,7 +505,7 @@ mod tests {
},
];
let msg = ChatMessage::assistant_with_tool_calls("", tool_calls);
let msg = ChatMessage::assistant_with_tool_calls(None, tool_calls);
let chat_msg: ChatCompletionMessage = msg.into();
assert_eq!(chat_msg.role, "assistant");
@@ -484,7 +535,7 @@ mod tests {
name: "test".to_string(),
arguments: serde_json::json!({"key": "value"}),
};
let msg = ChatMessage::assistant_with_tool_calls("", vec![tc]);
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
let chat_msg: ChatCompletionMessage = msg.into();
let calls = chat_msg.tool_calls.unwrap();
src/llm/provider.rs
+64 -11
View File
@@ -27,9 +27,8 @@ pub struct ChatMessage {
/// Name of the tool for tool results.
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
/// Tool calls requested by the assistant (for conversation replay).
/// OpenAI-compatible APIs require the assistant message to include
/// tool_calls when followed by tool result messages.
/// Tool calls made by the assistant (OpenAI protocol requires these
/// to appear on the assistant message preceding tool result messages).
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
@@ -68,17 +67,14 @@ impl ChatMessage {
}
}
/// Create an assistant message that requested tool calls.
/// Create an assistant message that includes tool calls.
///
/// OpenAI-compatible APIs require the assistant message to carry the
/// `tool_calls` array when followed by tool-result messages.
pub fn assistant_with_tool_calls(
content: impl Into<String>,
tool_calls: Vec<ToolCall>,
) -> Self {
/// Per the OpenAI protocol, an assistant message with tool_calls must
/// precede the corresponding tool result messages in the conversation.
pub fn assistant_with_tool_calls(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
content: content.unwrap_or_default(),
tool_call_id: None,
name: None,
tool_calls: if tool_calls.is_empty() {
@@ -112,6 +108,8 @@ pub struct CompletionRequest {
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub stop_sequences: Option<Vec<String>>,
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl CompletionRequest {
@@ -122,6 +120,7 @@ impl CompletionRequest {
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -145,6 +144,8 @@ pub struct CompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
pub finish_reason: FinishReason,
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
pub response_id: Option<String>,
}
/// Why the completion finished.
@@ -191,6 +192,8 @@ pub struct ToolCompletionRequest {
pub temperature: Option<f32>,
/// How to handle tool use: "auto", "required", or "none".
pub tool_choice: Option<String>,
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl ToolCompletionRequest {
@@ -202,6 +205,7 @@ impl ToolCompletionRequest {
max_tokens: None,
temperature: None,
tool_choice: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -234,6 +238,16 @@ pub struct ToolCompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
pub finish_reason: FinishReason,
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
pub response_id: Option<String>,
}
/// Metadata about a model returned by the provider's API.
#[derive(Debug, Clone)]
pub struct ModelMetadata {
pub id: String,
/// Total context window size in tokens.
pub context_length: Option<u32>,
}
/// Trait for LLM providers.
@@ -260,6 +274,45 @@ pub trait LlmProvider: Send + Sync {
Ok(Vec::new())
}
/// Fetch metadata for the current model (context length, etc.).
/// Default returns the model name with no size info.
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
Ok(ModelMetadata {
id: self.model_name().to_string(),
context_length: None,
})
}
/// Get the currently active model name.
///
/// May differ from `model_name()` if the model was switched at runtime
/// via `set_model()`. Default returns `model_name()`.
fn active_model_name(&self) -> String {
self.model_name().to_string()
}
/// Switch the active model at runtime. Not all providers support this.
fn set_model(&self, _model: &str) -> Result<(), LlmError> {
Err(LlmError::RequestFailed {
provider: "unknown".to_string(),
reason: "Runtime model switching not supported by this provider".to_string(),
})
}
/// Seed a response chain for a thread (e.g. restoring from DB).
///
/// Providers that support response chaining (e.g. NEAR AI `previous_response_id`)
/// store this so subsequent calls send only delta messages.
fn seed_response_chain(&self, _thread_id: &str, _response_id: String) {}
/// Get the last response chain ID for a thread.
///
/// Returns `None` if the provider doesn't support chaining or has no
/// stored state for this thread.
fn get_response_chain_id(&self, _thread_id: &str) -> Option<String> {
None
}
/// Calculate cost for a completion.
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
let (input_cost, output_cost) = self.cost_per_token();
src/llm/reasoning.rs
+353 -29
View File
@@ -21,6 +21,8 @@ pub struct ReasoningContext {
pub job_description: Option<String>,
/// Current state description.
pub current_state: Option<String>,
/// Opaque metadata forwarded to the LLM provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl ReasoningContext {
@@ -31,6 +33,7 @@ impl ReasoningContext {
available_tools: Vec::new(),
job_description: None,
current_state: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -57,6 +60,12 @@ impl ReasoningContext {
self.job_description = Some(description.into());
self
}
/// Set metadata (forwarded to the LLM provider).
pub fn with_metadata(mut self, metadata: std::collections::HashMap<String, String>) -> Self {
self.metadata = metadata;
self
}
}
impl Default for ReasoningContext {
@@ -114,7 +123,12 @@ pub enum RespondResult {
/// A text response (no tools needed).
Text(String),
/// The model wants to call tools. Caller should execute them and call back.
ToolCalls(Vec<ToolCall>),
/// Includes the optional content from the assistant message (some models
/// include explanatory text alongside tool calls).
ToolCalls {
tool_calls: Vec<ToolCall>,
content: Option<String>,
},
}
/// Reasoning engine for the agent.
@@ -192,10 +206,11 @@ impl Reasoning {
return Ok(vec![]);
}
let request =
let mut request =
ToolCompletionRequest::new(context.messages.clone(), context.available_tools.clone())
.with_max_tokens(1024)
.with_tool_choice("auto");
request.metadata = context.metadata.clone();
let response = self.llm.complete_with_tools(request).await?;
@@ -271,7 +286,9 @@ Respond in JSON format:
pub async fn respond(&self, context: &ReasoningContext) -> Result<String, LlmError> {
match self.respond_with_tools(context).await? {
RespondResult::Text(text) => Ok(text),
RespondResult::ToolCalls(calls) => {
RespondResult::ToolCalls {
tool_calls: calls, ..
} => {
// Format tool calls as text (legacy behavior for non-agentic callers)
let tool_info: Vec<String> = calls
.iter()
@@ -298,28 +315,49 @@ Respond in JSON format:
// If we have tools, use tool completion mode
if !context.available_tools.is_empty() {
let request = ToolCompletionRequest::new(messages, context.available_tools.clone())
let mut request = ToolCompletionRequest::new(messages, context.available_tools.clone())
.with_max_tokens(4096)
.with_temperature(0.7)
.with_tool_choice("auto");
request.metadata = context.metadata.clone();
let response = self.llm.complete_with_tools(request).await?;
// If there were tool calls, return them for execution
if !response.tool_calls.is_empty() {
return Ok(RespondResult::ToolCalls(response.tool_calls));
return Ok(RespondResult::ToolCalls {
tool_calls: response.tool_calls,
content: response.content,
});
}
let content = response
.content
.unwrap_or_else(|| "I'm not sure how to respond to that.".to_string());
// Some models (e.g. GLM-4.7) emit tool calls as XML tags in content
// instead of using the structured tool_calls field. Try to recover
// them before giving up and returning plain text.
let recovered = recover_tool_calls_from_content(&content, &context.available_tools);
if !recovered.is_empty() {
let cleaned = clean_response(&content);
return Ok(RespondResult::ToolCalls {
tool_calls: recovered,
content: if cleaned.is_empty() {
None
} else {
Some(cleaned)
},
});
}
Ok(RespondResult::Text(clean_response(&content)))
} else {
// No tools, use simple completion
let request = CompletionRequest::new(messages)
let mut request = CompletionRequest::new(messages)
.with_max_tokens(4096)
.with_temperature(0.7);
request.metadata = context.metadata.clone();
let response = self.llm.complete(request).await?;
Ok(RespondResult::Text(clean_response(&response.content)))
@@ -462,47 +500,178 @@ fn extract_json(text: &str) -> Option<&str> {
}
}
/// Clean up LLM response by stripping thinking tags and reasoning patterns.
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
///
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
/// Try to extract tool calls from content text where the model emitted them
/// as XML tags instead of using the structured tool_calls field.
///
/// Handles these formats:
/// - `<tool_call>tool_name</tool_call>` (bare name)
/// - `<tool_call>{"name":"x","arguments":{}}</tool_call>` (JSON)
/// - `<|tool_call|>...<|/tool_call|>` (pipe-delimited variant)
/// - `<function_call>...</function_call>` (function_call variant)
///
/// Only returns calls whose name matches an available tool.
fn recover_tool_calls_from_content(
content: &str,
available_tools: &[ToolDefinition],
) -> Vec<ToolCall> {
let tool_names: std::collections::HashSet<&str> =
available_tools.iter().map(|t| t.name.as_str()).collect();
let mut calls = Vec::new();
for (open, close) in &[
("<tool_call>", "</tool_call>"),
("<|tool_call|>", "<|/tool_call|>"),
("<function_call>", "</function_call>"),
("<|function_call|>", "<|/function_call|>"),
] {
let mut remaining = content;
while let Some(start) = remaining.find(open) {
let inner_start = start + open.len();
let after = &remaining[inner_start..];
let Some(end) = after.find(close) else {
break;
};
let inner = after[..end].trim();
remaining = &after[end + close.len()..];
if inner.is_empty() {
continue;
}
// Try JSON first: {"name":"x","arguments":{}}
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner) {
if let Some(name) = parsed.get("name").and_then(|v| v.as_str()) {
if tool_names.contains(name) {
let arguments = parsed
.get("arguments")
.cloned()
.unwrap_or(serde_json::Value::Object(Default::default()));
calls.push(ToolCall {
id: format!("recovered_{}", calls.len()),
name: name.to_string(),
arguments,
});
continue;
}
}
}
// Bare tool name (e.g. "<tool_call>tool_list</tool_call>")
let name = inner.trim();
if tool_names.contains(name) {
calls.push(ToolCall {
id: format!("recovered_{}", calls.len()),
name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()),
});
}
}
}
calls
}
/// `<tool_call>tool_list</tool_call>` or `<|tool_call|>` in the content field
/// instead of using the standard OpenAI tool_calls array. We strip all of
/// these before the response reaches channels/users.
fn clean_response(text: &str) -> String {
let text = strip_thinking_tags(text);
let text = strip_internal_tags(text);
strip_reasoning_patterns(&text)
}
/// Strip `<thinking>...</thinking>` blocks from LLM output.
/// Tags that are model-internal and should never reach users.
const INTERNAL_TAGS: &[&str] = &["thinking", "tool_call", "function_call", "tool_calls"];
/// Strip all model-internal XML tags from LLM output.
///
/// Some models (especially Claude with extended thinking) include internal
/// reasoning in thinking tags. We strip these before showing to users.
fn strip_thinking_tags(text: &str) -> String {
/// Handles standard XML tags (`<tag>...</tag>`) and pipe-delimited variants
/// (`<|tag|>...<|/tag|>`) used by some models (e.g. GLM-4.7).
fn strip_internal_tags(text: &str) -> String {
let mut result = text.to_string();
for tag in INTERNAL_TAGS {
result = strip_xml_tag(&result, tag);
result = strip_pipe_tag(&result, tag);
}
// Collapse triple+ newlines left behind by removed blocks
while result.contains("\n\n\n") {
result = result.replace("\n\n\n", "\n\n");
}
result.trim().to_string()
}
/// Strip `<tag>...</tag>` and `<tag ...>...</tag>` blocks from text.
fn strip_xml_tag(text: &str, tag: &str) -> String {
let open_exact = format!("<{}>", tag);
let open_prefix = format!("<{} ", tag); // for <tag attr="...">
let close = format!("</{}>", tag);
let mut result = String::with_capacity(text.len());
let mut remaining = text;
while let Some(start) = remaining.find("<thinking>") {
loop {
// Find the next opening tag (exact or with attributes)
let exact_pos = remaining.find(&open_exact);
let prefix_pos = remaining.find(&open_prefix);
let start = match (exact_pos, prefix_pos) {
(Some(a), Some(b)) => a.min(b),
(Some(a), None) => a,
(None, Some(b)) => b,
(None, None) => break,
};
// Add everything before the tag
result.push_str(&remaining[..start]);
// Find the end of the opening tag (the closing >)
let after_open = &remaining[start..];
let open_end = match after_open.find('>') {
Some(pos) => start + pos + 1,
None => break, // malformed, stop
};
// Find the closing tag
if let Some(end_offset) = remaining[start..].find("</thinking>") {
// Skip past the closing tag (start + offset + tag length)
let end = start + end_offset + "</thinking>".len();
if let Some(close_offset) = remaining[open_end..].find(&close) {
let end = open_end + close_offset + close.len();
remaining = &remaining[end..];
} else {
// No closing tag found, discard everything from here
// (malformed, but handle gracefully by not including the unclosed tag)
// No closing tag, discard from here (malformed)
remaining = "";
break;
}
}
// Add any remaining content after the last thinking block
result.push_str(remaining);
result
}
// Clean up any double newlines left behind
let mut cleaned = result.trim().to_string();
while cleaned.contains("\n\n\n") {
cleaned = cleaned.replace("\n\n\n", "\n\n");
/// Strip `<|tag|>...<|/tag|>` pipe-delimited blocks from text.
///
/// Some models (e.g. certain Chinese LLMs) use this format instead of
/// standard XML tags.
fn strip_pipe_tag(text: &str, tag: &str) -> String {
let open = format!("<|{}|>", tag);
let close = format!("<|/{}|>", tag);
let mut result = String::with_capacity(text.len());
let mut remaining = text;
while let Some(start) = remaining.find(&open) {
result.push_str(&remaining[..start]);
if let Some(close_offset) = remaining[start..].find(&close) {
let end = start + close_offset + close.len();
remaining = &remaining[end..];
} else {
remaining = "";
break;
}
}
cleaned
result.push_str(remaining);
result
}
/// Strip any remaining reasoning that wasn't in proper <thinking> tags.
@@ -574,7 +743,7 @@ That's my plan."#;
#[test]
fn test_strip_thinking_tags_basic() {
let input = "<thinking>Let me think about this...</thinking>Hello, user!";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello, user!");
}
@@ -582,7 +751,7 @@ That's my plan."#;
fn test_strip_thinking_tags_multiple() {
let input =
"<thinking>First thought</thinking>Hello<thinking>Second thought</thinking> world!";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello world!");
}
@@ -594,14 +763,14 @@ I need to consider:
2. How to respond
</thinking>
Here is my response to your question."#;
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Here is my response to your question.");
}
#[test]
fn test_strip_thinking_tags_no_tags() {
let input = "Just a normal response without thinking tags.";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Just a normal response without thinking tags.");
}
@@ -609,10 +778,77 @@ Here is my response to your question."#;
fn test_strip_thinking_tags_unclosed() {
// Malformed: unclosed tag should strip from there to end
let input = "Hello <thinking>this never closes";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello");
}
#[test]
fn test_strip_tool_call_tags() {
// GLM-4.7 emits this garbage instead of using the tool_calls array
let input = "<tool_call>tool_list</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "");
}
#[test]
fn test_strip_tool_call_with_surrounding_text() {
let input = "Here is my answer.\n\n<tool_call>\n{\"name\": \"search\", \"arguments\": {}}\n</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Here is my answer.");
}
#[test]
fn test_strip_multiple_internal_tags() {
let input = "<thinking>Let me think</thinking>Hello!\n<tool_call>some_tool</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Hello!");
}
#[test]
fn test_strip_function_call_tags() {
let input = "Response text<function_call>{\"name\": \"foo\"}</function_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Response text");
}
#[test]
fn test_strip_tool_calls_plural() {
let input = "<tool_calls>[{\"id\": \"1\"}]</tool_calls>Actual response.";
let output = strip_internal_tags(input);
assert_eq!(output, "Actual response.");
}
#[test]
fn test_strip_pipe_delimited_tags() {
let input = "<|tool_call|>{\"name\": \"search\"}<|/tool_call|>Hello!";
let output = strip_internal_tags(input);
assert_eq!(output, "Hello!");
}
#[test]
fn test_strip_pipe_delimited_thinking() {
let input = "<|thinking|>reasoning here<|/thinking|>The answer is 42.";
let output = strip_internal_tags(input);
assert_eq!(output, "The answer is 42.");
}
#[test]
fn test_strip_xml_tag_with_attributes() {
let input = "<tool_call type=\"function\">search()</tool_call>Done.";
let output = strip_internal_tags(input);
assert_eq!(output, "Done.");
}
#[test]
fn test_clean_response_preserves_normal_content() {
let input = "The function tool_call_handler works great. No tags here!";
let output = clean_response(input);
assert_eq!(
output,
"The function tool_call_handler works great. No tags here!"
);
}
#[test]
fn test_strip_reasoning_paragraph_break() {
// Content after paragraph break with "here" marker
@@ -660,4 +896,92 @@ Here is my response to your question."#;
let output = clean_response(input);
assert_eq!(output, "Here's the answer.");
}
// -- recover_tool_calls_from_content tests --
fn make_tools(names: &[&str]) -> Vec<ToolDefinition> {
names
.iter()
.map(|n| ToolDefinition {
name: n.to_string(),
description: String::new(),
parameters: serde_json::json!({}),
})
.collect()
}
#[test]
fn test_recover_bare_tool_name() {
let tools = make_tools(&["tool_list", "tool_auth"]);
let content = "<tool_call>tool_list</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
assert_eq!(calls[0].arguments, serde_json::json!({}));
}
#[test]
fn test_recover_json_tool_call() {
let tools = make_tools(&["memory_search"]);
let content =
r#"<tool_call>{"name": "memory_search", "arguments": {"query": "test"}}</tool_call>"#;
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"}));
}
#[test]
fn test_recover_pipe_delimited() {
let tools = make_tools(&["tool_list"]);
let content = "<|tool_call|>tool_list<|/tool_call|>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
}
#[test]
fn test_recover_unknown_tool_ignored() {
let tools = make_tools(&["tool_list"]);
let content = "<tool_call>nonexistent_tool</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_no_tags() {
let tools = make_tools(&["tool_list"]);
let content = "Just a normal response.";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_multiple_tool_calls() {
let tools = make_tools(&["tool_list", "tool_auth"]);
let content = "<tool_call>tool_list</tool_call>\n<tool_call>tool_auth</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 2);
assert_eq!(calls[0].name, "tool_list");
assert_eq!(calls[1].name, "tool_auth");
}
#[test]
fn test_recover_function_call_variant() {
let tools = make_tools(&["shell"]);
let content =
r#"<function_call>{"name": "shell", "arguments": {"cmd": "ls"}}</function_call>"#;
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "shell");
}
#[test]
fn test_recover_with_surrounding_text() {
let tools = make_tools(&["tool_list"]);
let content = "Let me check.\n\n<tool_call>tool_list</tool_call>\n\nDone.";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
}
}
src/llm/session.rs
+77 -2
View File
@@ -61,6 +61,10 @@ pub struct SessionManager {
token: RwLock<Option<SecretString>>,
/// Prevents thundering herd during concurrent 401s.
renewal_lock: Mutex<()>,
/// Optional database store for persisting session to the settings table.
store: RwLock<Option<Arc<crate::history::Store>>>,
/// User ID for DB settings (default: "default").
user_id: RwLock<String>,
}
impl SessionManager {
@@ -74,6 +78,8 @@ impl SessionManager {
.unwrap_or_else(|_| Client::new()),
token: RwLock::new(None),
renewal_lock: Mutex::new(()),
store: RwLock::new(None),
user_id: RwLock::new("default".to_string()),
};
// Try to load existing session synchronously during construction
@@ -103,6 +109,8 @@ impl SessionManager {
.unwrap_or_else(|_| Client::new()),
token: RwLock::new(None),
renewal_lock: Mutex::new(()),
store: RwLock::new(None),
user_id: RwLock::new("default".to_string()),
};
if let Err(e) = manager.load_session().await {
@@ -112,6 +120,21 @@ impl SessionManager {
manager
}
/// Attach a database store for persisting session tokens.
///
/// When a store is attached, session tokens are saved to the `settings`
/// table (key: `nearai.session_token`) in addition to the disk file.
/// On load, DB is preferred over disk.
pub async fn attach_store(&self, store: Arc<crate::history::Store>, user_id: &str) {
*self.store.write().await = Some(store);
*self.user_id.write().await = user_id.to_string();
// Try to load from DB (may have been saved by a previous run)
if let Err(e) = self.load_session_from_db().await {
tracing::debug!("No session in DB: {}", e);
}
}
/// Get the current session token, returning an error if not authenticated.
pub async fn get_token(&self) -> Result<SecretString, LlmError> {
let guard = self.token.read().await;
@@ -460,7 +483,7 @@ impl SessionManager {
Ok(())
}
/// Save session data to disk.
/// Save session data to disk and (if available) to the database.
async fn save_session(&self, token: &str, auth_provider: Option<&str>) -> Result<(), LlmError> {
let session = SessionData {
session_token: token.to_string(),
@@ -468,7 +491,7 @@ impl SessionManager {
auth_provider: auth_provider.map(String::from),
};
// Ensure parent directory exists
// Save to disk (always, as bootstrap fallback)
if let Some(parent) = self.config.session_path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
LlmError::Io(std::io::Error::new(
@@ -498,6 +521,58 @@ impl SessionManager {
})?;
tracing::debug!("Session saved to {}", self.config.session_path.display());
// Also save to DB if a store is attached
if let Some(ref store) = *self.store.read().await {
let user_id = self.user_id.read().await.clone();
let session_json = serde_json::to_value(&session)
.unwrap_or(serde_json::Value::String(token.to_string()));
if let Err(e) = store
.set_setting(&user_id, "nearai.session_token", &session_json)
.await
{
tracing::warn!("Failed to save session to DB: {}", e);
} else {
tracing::debug!("Session also saved to DB settings");
}
}
Ok(())
}
/// Try to load session from the database.
async fn load_session_from_db(&self) -> Result<(), LlmError> {
let store_guard = self.store.read().await;
let store = store_guard
.as_ref()
.ok_or_else(|| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: "No DB store attached".to_string(),
})?;
let user_id = self.user_id.read().await.clone();
let value = store
.get_setting(&user_id, "nearai.session_token")
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})?
.ok_or_else(|| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: "No session in DB".to_string(),
})?;
let session: SessionData =
serde_json::from_value(value).map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("Failed to parse DB session: {}", e),
})?;
let mut guard = self.token.write().await;
*guard = Some(SecretString::from(session.session_token));
tracing::info!("Loaded session from DB settings");
Ok(())
}