gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement tool approval, fix tool definition refresh, and wire embeddings

Browse Source 
This commit addresses three critical issues from code review:

1. Tool approval enforcement: Tools declaring requires_approval() (shell,
   http, file write/patch, build_software) now gate execution. Adds
   PendingApproval struct, session-scoped auto-approved tools set, and
   approval flow with yes/no/always commands.

2. Tool definition refresh: Tool definitions now refresh each iteration
   in both chat and job loops, so newly built tools become visible
   immediately within the same session.

3. Worker tool call handling: Changed respond() to respond_with_tools()
   when select_tools returns empty, properly executing tool calls instead
   of formatting them as text.

Also includes prior work from the plan:
- Wire embeddings provider (OpenAI + NEAR AI) to workspace
- Load workspace system prompt (identity files) into LLM context
- Route heartbeat notifications through channel manager
- Enable auto-context compaction when threshold exceeded
- Refactor to config structs (AgentDeps, WorkerDeps, LlmCallRecord)
- Fix clippy warnings (saturating_sub, too_many_arguments)

Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-03 11:34:10 -08:00
co-authored by Claude Opus 4.5
parent 8af48390a9
commit 2cc9aed364
18 changed files with 1079 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/agent_loop.rs
+431 -62
View File
@@ -10,7 +10,7 @@ use crate::agent::compaction::ContextCompactor;
use crate::agent::context_monitor::ContextMonitor;
use crate::agent::heartbeat::spawn_heartbeat;
use crate::agent::self_repair::DefaultSelfRepair;
use crate::agent::session::{Session, ThreadState};
use crate::agent::session::{PendingApproval, Session, ThreadState};
use crate::agent::session_manager::SessionManager;
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
use crate::agent::{
@@ -27,20 +27,38 @@ use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
use crate::workspace::Workspace;
/// Result of the agentic loop execution.
enum AgenticLoopResult {
/// Completed with a response.
Response(String),
/// A tool requires approval before continuing.
NeedApproval {
/// The pending approval request to store.
pending: PendingApproval,
},
}
/// Core dependencies for the agent.
///
/// Bundles the shared components to reduce argument count.
pub struct AgentDeps {
pub store: Option<Arc<Store>>,
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub workspace: Option<Arc<Workspace>>,
}
/// The main agent that coordinates all components.
pub struct Agent {
config: AgentConfig,
store: Option<Arc<Store>>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
channels: ChannelManager,
deps: AgentDeps,
channels: Arc<ChannelManager>,
context_manager: Arc<ContextManager>,
scheduler: Arc<Scheduler>,
router: Router,
session_manager: Arc<SessionManager>,
context_monitor: ContextMonitor,
workspace: Option<Arc<Workspace>>,
heartbeat_config: Option<HeartbeatConfig>,
}
@@ -48,12 +66,8 @@ impl Agent {
/// Create a new agent.
pub fn new(
config: AgentConfig,
store: Option<Arc<Store>>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
deps: AgentDeps,
channels: ChannelManager,
workspace: Option<Arc<Workspace>>,
heartbeat_config: Option<HeartbeatConfig>,
) -> Self {
let context_manager = Arc::new(ContextManager::new(config.max_parallel_jobs));
@@ -61,29 +75,46 @@ impl Agent {
let scheduler = Arc::new(Scheduler::new(
config.clone(),
context_manager.clone(),
llm.clone(),
safety.clone(),
tools.clone(),
store.clone(),
deps.llm.clone(),
deps.safety.clone(),
deps.tools.clone(),
deps.store.clone(),
));
Self {
config,
store,
llm,
safety,
tools,
channels,
deps,
channels: Arc::new(channels),
context_manager,
scheduler,
router: Router::new(),
session_manager: Arc::new(SessionManager::new()),
context_monitor: ContextMonitor::new(),
workspace,
heartbeat_config,
}
}
// Convenience accessors
fn store(&self) -> Option<&Arc<Store>> {
self.deps.store.as_ref()
}
fn llm(&self) -> &Arc<dyn LlmProvider> {
&self.deps.llm
}
fn safety(&self) -> &Arc<SafetyLayer> {
&self.deps.safety
}
fn tools(&self) -> &Arc<ToolRegistry> {
&self.deps.tools
}
fn workspace(&self) -> Option<&Arc<Workspace>> {
self.deps.workspace.as_ref()
}
/// Run the agent main loop.
pub async fn run(self) -> Result<(), Error> {
// Start channels
@@ -104,30 +135,61 @@ impl Agent {
// Spawn heartbeat if enabled
let heartbeat_handle = if let Some(ref hb_config) = self.heartbeat_config {
if hb_config.enabled {
if let Some(ref workspace) = self.workspace {
if let Some(workspace) = self.workspace() {
let config = AgentHeartbeatConfig::default()
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
// Set up notification channel if configured
// Set up notification channel
let (notify_tx, mut notify_rx) =
tokio::sync::mpsc::channel::<OutgoingResponse>(16);
// Spawn notification forwarder
// We can't clone ChannelManager directly, so we just log the notifications
// The heartbeat system will handle notifications via the response_tx
// Spawn notification forwarder that routes through channel manager
let notify_channel = hb_config.notify_channel.clone();
let notify_user = hb_config.notify_user.clone();
let channels = self.channels.clone();
tokio::spawn(async move {
while let Some(response) = notify_rx.recv().await {
if let (Some(ch), Some(user)) = (&notify_channel, &notify_user) {
// Log the heartbeat notification
// In a full implementation, we'd route this through a shared channel reference
tracing::info!(
"Heartbeat notification for {}/{}: {}",
ch,
user,
&response.content
);
// Route notification to configured channel/user, or broadcast to all
match (&notify_channel, &notify_user) {
(Some(channel), Some(user)) => {
// Send to specific channel and user
if let Err(e) =
channels.broadcast(channel, user, response.clone()).await
{
tracing::warn!(
"Failed to send heartbeat to {}/{}: {}",
channel,
user,
e
);
} else {
tracing::debug!(
"Heartbeat notification sent to {}/{}",
channel,
user
);
}
}
(None, Some(user)) => {
// Broadcast to all channels for this 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
);
}
}
}
_ => {
// No target configured, just log
tracing::info!(
"Heartbeat notification (no target configured): {}",
&response.content
);
}
}
}
});
@@ -139,7 +201,7 @@ impl Agent {
Some(spawn_heartbeat(
config,
workspace.clone(),
self.llm.clone(),
self.llm().clone(),
Some(notify_tx),
))
} else {
@@ -230,11 +292,24 @@ impl Agent {
Submission::Resume { checkpoint_id } => {
self.process_resume(session, thread_id, checkpoint_id).await
}
Submission::ExecApproval { .. } => {
// Not supported in simple chat flow
Ok(SubmissionResult::error(
"Approval flow not supported in this context",
))
Submission::ExecApproval {
request_id,
approved,
always,
} => {
self.process_approval(
message,
session,
thread_id,
Some(request_id),
approved,
always,
)
.await
}
Submission::ApprovalResponse { approved, always } => {
self.process_approval(message, session, thread_id, None, approved, always)
.await
}
};
@@ -244,8 +319,32 @@ impl Agent {
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())),
SubmissionResult::NeedApproval { .. } => {
Ok(Some("Approval required but not supported.".into()))
SubmissionResult::NeedApproval {
request_id,
tool_name,
description,
parameters,
} => {
// Format approval request for user
let params_preview = serde_json::to_string_pretty(&parameters)
.unwrap_or_else(|_| parameters.to_string());
let params_truncated = if params_preview.len() > 200 {
format!("{}...", &params_preview[..200])
} else {
params_preview
};
Ok(Some(format!(
"🔒 Tool requires approval:\n\n\
**Tool:** {}\n\
**Description:** {}\n\
**Parameters:** ```\n{}\n```\n\n\
Reply with:\n\
- `yes` or `approve` to allow this tool\n\
- `always` to always allow this tool in this session\n\
- `no` or `deny` to reject\n\n\
Request ID: {}",
tool_name, description, params_truncated, request_id
)))
}
}
}
@@ -322,9 +421,9 @@ impl Agent {
"Context at {:.1}% capacity, auto-compacting",
self.context_monitor.usage_percent(&messages)
);
let compactor = ContextCompactor::new(self.llm.clone());
let compactor = ContextCompactor::new(self.llm().clone());
if let Err(e) = compactor
.compact(thread, strategy, self.workspace.as_deref())
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
.await
{
tracing::warn!("Auto-compaction failed: {}", e);
@@ -389,9 +488,9 @@ impl Agent {
return Ok(SubmissionResult::Interrupted);
}
// Complete or fail the turn
// Complete, fail, or request approval
match result {
Ok(response) => {
Ok(AgenticLoopResult::Response(response)) => {
thread.complete_turn(&response);
let _ = self
.channels
@@ -399,6 +498,27 @@ impl Agent {
.await;
Ok(SubmissionResult::response(response))
}
Ok(AgenticLoopResult::NeedApproval { pending }) => {
// Store pending approval in thread and update state
let request_id = pending.request_id;
let tool_name = pending.tool_name.clone();
let description = pending.description.clone();
let parameters = pending.parameters.clone();
thread.await_approval(pending);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::Status("Awaiting approval".into()),
)
.await;
Ok(SubmissionResult::NeedApproval {
request_id,
tool_name,
description,
parameters,
})
}
Err(e) => {
thread.fail_turn(e.to_string());
Ok(SubmissionResult::error(e.to_string()))
@@ -407,15 +527,34 @@ impl Agent {
}
/// Run the agentic loop: call LLM, execute tools, repeat until text response.
///
/// Returns `AgenticLoopResult::Response` on completion, or
/// `AgenticLoopResult::NeedApproval` if a tool requires user approval.
async fn run_agentic_loop(
&self,
message: &IncomingMessage,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
initial_messages: Vec<ChatMessage>,
) -> Result<String, Error> {
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let tool_defs = self.tools.tool_definitions().await;
) -> Result<AgenticLoopResult, Error> {
// Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.)
let system_prompt = if let Some(ws) = self.workspace() {
match ws.system_prompt().await {
Ok(prompt) if !prompt.is_empty() => Some(prompt),
Ok(_) => None,
Err(e) => {
tracing::debug!("Could not load workspace system prompt: {}", e);
None
}
}
} else {
None
};
let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
if let Some(prompt) = system_prompt {
reasoning = reasoning.with_system_prompt(prompt);
}
// Build context with messages that we'll mutate during the loop
let mut context_messages = initial_messages;
@@ -425,6 +564,7 @@ impl Agent {
const MAX_TOOL_ITERATIONS: usize = 10;
let mut iteration = 0;
let mut tools_executed = false;
loop {
iteration += 1;
@@ -450,19 +590,38 @@ impl Agent {
}
}
// Refresh tool definitions each iteration so newly built tools become visible
let tool_defs = self.tools().tool_definitions().await;
// Call LLM with current context
let context = ReasoningContext::new()
.with_messages(context_messages.clone())
.with_tools(tool_defs.clone());
.with_tools(tool_defs);
let result = reasoning.respond_with_tools(&context).await?;
match result {
RespondResult::Text(text) => {
// Final response, return it
return Ok(text);
// If no tools have been executed yet, prompt the LLM to use tools
// This handles the case where the model explains what it will do
// instead of actually calling tools
if !tools_executed && iteration < 3 {
tracing::debug!(
"No tools executed yet (iteration {}), prompting for tool use",
iteration
);
context_messages.push(ChatMessage::assistant(&text));
context_messages.push(ChatMessage::user(
"Please proceed and use the available tools to complete this task.",
));
continue;
}
// Tools have been executed or we've tried multiple times, return response
return Ok(AgenticLoopResult::Response(text));
}
RespondResult::ToolCalls(tool_calls) => {
tools_executed = true;
// Execute tools and add results to context
let _ = self
.channels
@@ -487,8 +646,33 @@ impl Agent {
}
}
// Execute each tool
// Execute each tool (with approval checking)
for tc in tool_calls {
// Check if tool requires approval
if let Some(tool) = self.tools().get(&tc.name).await {
if tool.requires_approval() {
// Check if auto-approved for this session
let is_auto_approved = {
let sess = session.lock().await;
sess.is_tool_auto_approved(&tc.name)
};
if !is_auto_approved {
// Need approval - store pending request and return
let pending = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
description: tool.description().to_string(),
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
};
return Ok(AgenticLoopResult::NeedApproval { pending });
}
}
}
let tool_result = self
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
.await;
@@ -514,8 +698,9 @@ impl Agent {
let result_content = match tool_result {
Ok(output) => {
// Sanitize output before showing to LLM
let sanitized = self.safety.sanitize_tool_output(&tc.name, &output);
self.safety.wrap_for_llm(
let sanitized =
self.safety().sanitize_tool_output(&tc.name, &output);
self.safety().wrap_for_llm(
&tc.name,
&sanitized.content,
sanitized.was_modified,
@@ -543,7 +728,7 @@ impl Agent {
job_ctx: &JobContext,
) -> Result<String, Error> {
let tool =
self.tools
self.tools()
.get(tool_name)
.await
.ok_or_else(|| crate::error::ToolError::NotFound {
@@ -718,9 +903,9 @@ impl Agent {
crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 },
);
let compactor = ContextCompactor::new(self.llm.clone());
let compactor = ContextCompactor::new(self.llm().clone());
match compactor
.compact(thread, strategy, self.workspace.as_deref())
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
.await
{
Ok(result) => {
@@ -757,6 +942,190 @@ impl Agent {
Ok(SubmissionResult::ok_with_message("Thread cleared."))
}
/// Process an approval or rejection of a pending tool execution.
async fn process_approval(
&self,
message: &IncomingMessage,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
request_id: Option<Uuid>,
approved: bool,
always: bool,
) -> Result<SubmissionResult, Error> {
// Get thread state and pending approval
let (_thread_state, pending) = {
let mut sess = session.lock().await;
let thread = sess
.threads
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
if thread.state != ThreadState::AwaitingApproval {
return Ok(SubmissionResult::error("No pending approval request."));
}
let pending = thread.take_pending_approval();
(thread.state, pending)
};
let pending = match pending {
Some(p) => p,
None => return Ok(SubmissionResult::error("No pending approval request.")),
};
// Verify request ID if provided
if let Some(req_id) = request_id {
if req_id != pending.request_id {
// Put it back and return error
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.await_approval(pending);
}
return Ok(SubmissionResult::error(
"Request ID mismatch. Use the correct request ID.",
));
}
}
if approved {
// If always, add to auto-approved set
if always {
let mut sess = session.lock().await;
sess.auto_approve_tool(&pending.tool_name);
tracing::info!(
"Auto-approved tool '{}' for session {}",
pending.tool_name,
sess.id
);
}
// Reset thread state to processing
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.state = ThreadState::Processing;
}
}
// Execute the approved tool and continue the loop
let job_ctx =
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
let tool_result = self
.execute_chat_tool(&pending.tool_name, &pending.parameters, &job_ctx)
.await;
// Build context including the tool result
let mut context_messages = pending.context_messages;
// Record result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
if let Some(turn) = thread.last_turn_mut() {
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
}
}
}
}
// Add tool result to context
let result_content = match tool_result {
Ok(output) => {
let sanitized = self
.safety()
.sanitize_tool_output(&pending.tool_name, &output);
self.safety().wrap_for_llm(
&pending.tool_name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
context_messages.push(ChatMessage::tool_result(
&pending.tool_call_id,
&pending.tool_name,
result_content,
));
// Continue the agentic loop
let result = self
.run_agentic_loop(message, session.clone(), thread_id, context_messages)
.await;
// Handle the result
let mut sess = session.lock().await;
let thread = sess
.threads
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
match result {
Ok(AgenticLoopResult::Response(response)) => {
thread.complete_turn(&response);
let _ = self
.channels
.send_status(&message.channel, StatusUpdate::Status("Done".into()))
.await;
Ok(SubmissionResult::response(response))
}
Ok(AgenticLoopResult::NeedApproval {
pending: new_pending,
}) => {
let request_id = new_pending.request_id;
let tool_name = new_pending.tool_name.clone();
let description = new_pending.description.clone();
let parameters = new_pending.parameters.clone();
thread.await_approval(new_pending);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::Status("Awaiting approval".into()),
)
.await;
Ok(SubmissionResult::NeedApproval {
request_id,
tool_name,
description,
parameters,
})
}
Err(e) => {
thread.fail_turn(e.to_string());
Ok(SubmissionResult::error(e.to_string()))
}
}
} else {
// Rejected - clear approval and return to idle
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.clear_pending_approval();
}
}
let _ = self
.channels
.send_status(&message.channel, StatusUpdate::Status("Rejected".into()))
.await;
Ok(SubmissionResult::response(format!(
"Tool '{}' was rejected. The agent will not execute this tool.\n\n\
You can continue the conversation or try a different approach.",
pending.tool_name
)))
}
}
async fn process_new_thread(
&self,
message: &IncomingMessage,
@@ -842,7 +1211,7 @@ impl Agent {
}
// Persist new job to database (fire-and-forget)
if let Some(ref store) = self.store {
if let Some(store) = self.store() {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
let store = store.clone();
tokio::spawn(async move {
@@ -990,7 +1359,7 @@ impl Agent {
))),
"tools" => {
let tools = self.tools.list().await;
let tools = self.tools().list().await;
Ok(Some(format!("Available tools: {}", tools.join(", "))))
}
src/agent/mod.rs
+4 -4
View File
@@ -21,18 +21,18 @@ mod session_manager;
pub mod submission;
pub mod task;
pub mod undo;
mod worker;
pub mod worker;
pub use agent_loop::Agent;
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 router::{MessageIntent, Router};
pub use scheduler::Scheduler;
pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob};
pub use session::{Session, Thread, ThreadState, Turn, TurnState};
pub use session::{PendingApproval, Session, Thread, ThreadState, Turn, TurnState};
pub use session_manager::SessionManager;
pub use submission::{Submission, SubmissionParser, SubmissionResult};
pub use task::{Task, TaskContext, TaskHandler, TaskOutput, TaskStatus};
pub use undo::{Checkpoint, UndoManager};
pub use worker::Worker;
pub use worker::{Worker, WorkerDeps};
src/agent/scheduler.rs
+12 -12
View File
@@ -8,8 +8,8 @@ use tokio::sync::{RwLock, mpsc, oneshot};
use tokio::task::JoinHandle;
use uuid::Uuid;
use crate::agent::Worker;
use crate::agent::task::{Task, TaskContext, TaskOutput};
use crate::agent::worker::{Worker, WorkerDeps};
use crate::config::AgentConfig;
use crate::context::{ContextManager, JobContext, JobState};
use crate::error::{Error, JobError};
@@ -109,17 +109,17 @@ impl Scheduler {
// Create worker channel
let (tx, rx) = mpsc::channel(16);
// Create worker
let worker = Worker::new(
job_id,
self.context_manager.clone(),
self.llm.clone(),
self.safety.clone(),
self.tools.clone(),
self.store.clone(),
self.config.job_timeout,
self.config.use_planning,
);
// Create worker with shared dependencies
let deps = WorkerDeps {
context_manager: self.context_manager.clone(),
llm: self.llm.clone(),
safety: self.safety.clone(),
tools: self.tools.clone(),
store: self.store.clone(),
timeout: self.config.job_timeout,
use_planning: self.config.use_planning,
};
let worker = Worker::new(job_id, deps);
// Spawn worker task
let handle = tokio::spawn(async move {
src/agent/session.rs
+51 -3
View File
@@ -10,7 +10,7 @@
//! - Compaction: Summarize old turns to save context
//! - Resume: Continue from a saved checkpoint
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
@@ -35,6 +35,9 @@ pub struct Session {
pub last_active_at: DateTime<Utc>,
/// Session metadata.
pub metadata: serde_json::Value,
/// Tools that have been auto-approved for this session ("always approve").
#[serde(default)]
pub auto_approved_tools: HashSet<String>,
}
impl Session {
@@ -49,9 +52,20 @@ impl Session {
created_at: now,
last_active_at: now,
metadata: serde_json::Value::Null,
auto_approved_tools: HashSet::new(),
}
}
/// Check if a tool has been auto-approved for this session.
pub fn is_tool_auto_approved(&self, tool_name: &str) -> bool {
self.auto_approved_tools.contains(tool_name)
}
/// Add a tool to the auto-approved set.
pub fn auto_approve_tool(&mut self, tool_name: impl Into<String>) {
self.auto_approved_tools.insert(tool_name.into());
}
/// Create a new thread in this session.
pub fn create_thread(&mut self) -> &mut Thread {
let thread = Thread::new(self.id);
@@ -107,6 +121,23 @@ pub enum ThreadState {
Interrupted,
}
/// Pending tool approval request stored on a thread.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingApproval {
/// Unique request ID.
pub request_id: Uuid,
/// Tool name requiring approval.
pub tool_name: String,
/// Tool parameters.
pub parameters: serde_json::Value,
/// Description of what the tool will do.
pub description: String,
/// Tool call ID from LLM (for proper context continuation).
pub tool_call_id: String,
/// Context messages at the time of the request (to resume from).
pub context_messages: Vec<ChatMessage>,
}
/// A conversation thread within a session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Thread {
@@ -124,6 +155,9 @@ pub struct Thread {
pub updated_at: DateTime<Utc>,
/// Thread metadata (e.g., title, tags).
pub metadata: serde_json::Value,
/// Pending approval request (when state is AwaitingApproval).
#[serde(default)]
pub pending_approval: Option<PendingApproval>,
}
impl Thread {
@@ -138,6 +172,7 @@ impl Thread {
created_at: now,
updated_at: now,
metadata: serde_json::Value::Null,
pending_approval: None,
}
}
@@ -184,9 +219,22 @@ impl Thread {
self.updated_at = Utc::now();
}
/// Mark the thread as awaiting approval.
pub fn await_approval(&mut self) {
/// Mark the thread as awaiting approval with pending request details.
pub fn await_approval(&mut self, pending: PendingApproval) {
self.state = ThreadState::AwaitingApproval;
self.pending_approval = Some(pending);
self.updated_at = Utc::now();
}
/// Take the pending approval (clearing it from the thread).
pub fn take_pending_approval(&mut self) -> Option<PendingApproval> {
self.pending_approval.take()
}
/// Clear pending approval and return to idle state.
pub fn clear_pending_approval(&mut self) {
self.pending_approval = None;
self.state = ThreadState::Idle;
self.updated_at = Utc::now();
}
src/agent/submission.rs
+33 -1
View File
@@ -52,6 +52,30 @@ impl SubmissionParser {
}
}
// Approval responses (simple yes/no/always for pending approvals)
// These are short enough to check explicitly
match lower.as_str() {
"yes" | "y" | "approve" | "ok" => {
return Submission::ApprovalResponse {
approved: true,
always: false,
};
}
"always" | "yes always" | "approve always" => {
return Submission::ApprovalResponse {
approved: true,
always: true,
};
}
"no" | "n" | "deny" | "reject" | "cancel" => {
return Submission::ApprovalResponse {
approved: false,
always: false,
};
}
_ => {}
}
// Default: user input
Submission::UserInput {
content: content.to_string(),
@@ -68,7 +92,7 @@ pub enum Submission {
content: String,
},
/// Response to an execution approval request.
/// Response to an execution approval request (with explicit request ID).
ExecApproval {
/// ID of the approval request being responded to.
request_id: Uuid,
@@ -78,6 +102,14 @@ pub enum Submission {
always: bool,
},
/// Simple approval response (yes/no/always) for the current pending approval.
ApprovalResponse {
/// Whether the execution was approved.
approved: bool,
/// If true, auto-approve this tool for the rest of the session.
always: bool,
},
/// Interrupt the current turn.
Interrupt,
src/agent/worker.rs
+128 -68
View File
@@ -13,22 +13,30 @@ use crate::context::{ContextManager, JobState};
use crate::error::Error;
use crate::history::Store;
use crate::llm::{
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolSelection,
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
};
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
/// Shared dependencies for worker execution.
///
/// This bundles the dependencies that are shared across all workers,
/// reducing the number of arguments to `Worker::new`.
#[derive(Clone)]
pub struct WorkerDeps {
pub context_manager: Arc<ContextManager>,
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub store: Option<Arc<Store>>,
pub timeout: Duration,
pub use_planning: bool,
}
/// Worker that executes a single job.
pub struct Worker {
job_id: Uuid,
context_manager: Arc<ContextManager>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
timeout: Duration,
/// Whether to use planning before tool execution.
use_planning: bool,
deps: WorkerDeps,
}
/// Result of a tool execution with metadata for context building.
@@ -37,32 +45,43 @@ struct ToolExecResult {
}
impl Worker {
/// Create a new worker.
pub fn new(
job_id: Uuid,
context_manager: Arc<ContextManager>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
timeout: Duration,
use_planning: bool,
) -> Self {
Self {
job_id,
context_manager,
llm,
safety,
tools,
store,
timeout,
use_planning,
}
/// Create a new worker for a specific job.
pub fn new(job_id: Uuid, deps: WorkerDeps) -> Self {
Self { job_id, deps }
}
// Convenience accessors to avoid deps.field everywhere
fn context_manager(&self) -> &Arc<ContextManager> {
&self.deps.context_manager
}
fn llm(&self) -> &Arc<dyn LlmProvider> {
&self.deps.llm
}
fn safety(&self) -> &Arc<SafetyLayer> {
&self.deps.safety
}
fn tools(&self) -> &Arc<ToolRegistry> {
&self.deps.tools
}
fn store(&self) -> Option<&Arc<Store>> {
self.deps.store.as_ref()
}
fn timeout(&self) -> Duration {
self.deps.timeout
}
fn use_planning(&self) -> bool {
self.deps.use_planning
}
/// Fire-and-forget persistence of job status.
fn persist_status(&self, status: JobState, reason: Option<String>) {
if let Some(ref store) = self.store {
if let Some(store) = self.store() {
let store = store.clone();
let job_id = self.job_id;
tokio::spawn(async move {
@@ -91,16 +110,13 @@ impl Worker {
}
// Get job context
let job_ctx = self.context_manager.get_context(self.job_id).await?;
let job_ctx = self.context_manager().get_context(self.job_id).await?;
// Create reasoning engine
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
// Build initial reasoning context
let tool_defs = self.tools.tool_definitions().await;
let mut reason_ctx = ReasoningContext::new()
.with_job(&job_ctx.description)
.with_tools(tool_defs);
// Build initial reasoning context (tool definitions refreshed each iteration in execution_loop)
let mut reason_ctx = ReasoningContext::new().with_job(&job_ctx.description);
// Add system message
reason_ctx.messages.push(ChatMessage::system(format!(
@@ -116,7 +132,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
)));
// Main execution loop with timeout
let result = tokio::time::timeout(self.timeout, async {
let result = tokio::time::timeout(self.timeout(), async {
self.execution_loop(&mut rx, &reasoning, &mut reason_ctx)
.await
})
@@ -148,8 +164,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let max_iterations = 50;
let mut iteration = 0;
// Initial tool definitions for planning (will be refreshed in loop)
reason_ctx.available_tools = self.tools().tool_definitions().await;
// Generate plan if planning is enabled
let plan = if self.use_planning {
let plan = if self.use_planning() {
match reasoning.plan(reason_ctx).await {
Ok(p) => {
tracing::info!(
@@ -213,29 +232,61 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
return Ok(());
}
// Refresh tool definitions so newly built tools become visible
reason_ctx.available_tools = self.tools().tool_definitions().await;
// Select next tool(s) to use
let selections = reasoning.select_tools(reason_ctx).await?;
if selections.is_empty() {
// No tools selected, ask LLM for next steps
let response = reasoning.respond(reason_ctx).await?;
// No tools from select_tools, ask LLM directly (may still return tool calls)
let respond_result = reasoning.respond_with_tools(reason_ctx).await?;
if response.to_lowercase().contains("complete")
|| response.to_lowercase().contains("finished")
|| response.to_lowercase().contains("done")
{
self.mark_completed().await?;
return Ok(());
}
match respond_result {
RespondResult::Text(response) => {
// Check for completion keywords
let response_lower = response.to_lowercase();
if response_lower.contains("complete")
|| response_lower.contains("finished")
|| response_lower.contains("done")
{
self.mark_completed().await?;
return Ok(());
}
// Add assistant response to context
reason_ctx.messages.push(ChatMessage::assistant(&response));
// Add assistant response to context
reason_ctx.messages.push(ChatMessage::assistant(&response));
// Give it one more chance to select a tool
if iteration > 3 && iteration % 5 == 0 {
reason_ctx.messages.push(ChatMessage::user(
"Are you stuck? Do you need help completing this job?",
));
// Give it one more chance to select a tool
if iteration > 3 && iteration % 5 == 0 {
reason_ctx.messages.push(ChatMessage::user(
"Are you stuck? Do you need help completing this job?",
));
}
}
RespondResult::ToolCalls(tool_calls) => {
// Model returned tool calls - execute them
tracing::debug!(
"Job {} respond_with_tools returned {} tool calls",
self.job_id,
tool_calls.len()
);
for tc in tool_calls {
let result = self.execute_tool(&tc.name, &tc.arguments).await;
// Create synthetic selection for process_tool_result
let selection = ToolSelection {
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
reasoning: String::new(),
alternatives: vec![],
};
self.process_tool_result(reason_ctx, &selection, result)
.await?;
}
}
}
} else if selections.len() == 1 {
// Single tool: execute directly
@@ -282,10 +333,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
.map(|selection| {
let tool_name = selection.tool_name.clone();
let params = selection.parameters.clone();
let tools = self.tools.clone();
let context_manager = self.context_manager.clone();
let tools = self.tools().clone();
let context_manager = self.context_manager().clone();
let job_id = self.job_id;
let store = self.store.clone();
let store = self.deps.store.clone();
async move {
let result = Self::execute_tool_inner(
@@ -321,6 +372,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
name: tool_name.to_string(),
})?;
// Log warning if tool requires approval (autonomous jobs auto-approve for now)
if tool.requires_approval() {
tracing::warn!(
job_id = %job_id,
tool = %tool_name,
"Executing sensitive tool in autonomous job (auto-approved)"
);
}
// Get job context for the tool
let job_ctx = context_manager.get_context(job_id).await?;
@@ -412,11 +472,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
Ok(output) => {
// Sanitize output
let sanitized = self
.safety
.safety()
.sanitize_tool_output(&selection.tool_name, &output);
// Add to context
let wrapped = self.safety.wrap_for_llm(
let wrapped = self.safety().wrap_for_llm(
&selection.tool_name,
&sanitized.content,
sanitized.was_modified,
@@ -445,7 +505,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
);
// Record failure for self-repair tracking
if let Some(ref store) = self.store {
if let Some(store) = self.store() {
let store = store.clone();
let tool_name = selection.tool_name.clone();
let error_msg = e.to_string();
@@ -563,9 +623,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
params: &serde_json::Value,
) -> Result<String, Error> {
Self::execute_tool_inner(
self.tools.clone(),
self.context_manager.clone(),
self.store.clone(),
self.tools().clone(),
self.context_manager().clone(),
self.deps.store.clone(),
self.job_id,
tool_name,
params,
@@ -574,7 +634,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
async fn mark_completed(&self) -> Result<(), Error> {
self.context_manager
self.context_manager()
.update_context(self.job_id, |ctx| {
ctx.transition_to(
JobState::Completed,
@@ -595,7 +655,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
async fn mark_failed(&self, reason: &str) -> Result<(), Error> {
self.context_manager
self.context_manager()
.update_context(self.job_id, |ctx| {
ctx.transition_to(JobState::Failed, Some(reason.to_string()))
})
@@ -610,7 +670,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
async fn mark_stuck(&self, reason: &str) -> Result<(), Error> {
self.context_manager
self.context_manager()
.update_context(self.job_id, |ctx| ctx.mark_stuck(reason))
.await?
.map_err(|s| crate::error::JobError::ContextError {
src/channels/channel.rs
+14
View File
@@ -146,6 +146,20 @@ pub trait Channel: Send + Sync {
Ok(())
}
/// Send a proactive message without a prior incoming message.
///
/// Used for alerts, heartbeat notifications, and other agent-initiated communication.
/// The user_id helps target a specific user within the channel.
///
/// Default implementation does nothing (for channels that don't support broadcast).
async fn broadcast(
&self,
_user_id: &str,
_response: OutgoingResponse,
) -> Result<(), ChannelError> {
Ok(())
}
/// Check if the channel is healthy.
async fn health_check(&self) -> Result<(), ChannelError>;
src/channels/cli/mod.rs
+16
View File
@@ -128,6 +128,22 @@ impl Channel for TuiChannel {
Ok(())
}
async fn broadcast(
&self,
_user_id: &str,
response: OutgoingResponse,
) -> Result<(), ChannelError> {
// For TUI, broadcasts appear as regular agent responses with a notification indicator
self.event_tx
.send(AppEvent::Response(response.content))
.await
.map_err(|e| ChannelError::SendFailed {
name: "tui".to_string(),
reason: e.to_string(),
})?;
Ok(())
}
async fn health_check(&self) -> Result<(), ChannelError> {
// Channel is healthy if we haven't been closed
if self.event_tx.is_closed() {
src/channels/cli/render.rs
+1 -5
View File
@@ -89,11 +89,7 @@ fn render_messages(frame: &mut Frame, app: &AppState, area: Rect) {
// Calculate scroll - show most recent messages
let visible_height = area.height.saturating_sub(2) as usize; // Account for borders
let total_lines = lines.len();
let scroll_offset = if total_lines > visible_height {
total_lines - visible_height
} else {
0
};
let scroll_offset = total_lines.saturating_sub(visible_height);
let text = Text::from(lines);
let messages = Paragraph::new(text)
src/channels/manager.rs
+39
View File
@@ -97,6 +97,45 @@ impl ChannelManager {
}
}
/// Broadcast a message to a specific user on a specific channel.
///
/// Used for proactive notifications like heartbeat alerts.
pub async fn broadcast(
&self,
channel_name: &str,
user_id: &str,
response: OutgoingResponse,
) -> Result<(), ChannelError> {
let channels = self.channels.read().await;
if let Some(channel) = channels.get(channel_name) {
channel.broadcast(user_id, response).await
} else {
Err(ChannelError::SendFailed {
name: channel_name.to_string(),
reason: "Channel not found".to_string(),
})
}
}
/// Broadcast a message to all channels.
///
/// Sends to the specified user on every registered channel.
pub async fn broadcast_all(
&self,
user_id: &str,
response: OutgoingResponse,
) -> Vec<(String, Result<(), ChannelError>)> {
let channels = self.channels.read().await;
let mut results = Vec::new();
for (name, channel) in channels.iter() {
let result = channel.broadcast(user_id, response.clone()).await;
results.push((name.clone(), result));
}
results
}
/// Check health of all channels.
pub async fn health_check_all(&self) -> HashMap<String, Result<(), ChannelError>> {
let channels = self.channels.read().await;
src/config.rs
+58
View File
@@ -12,6 +12,7 @@ use crate::error::ConfigError;
pub struct Config {
pub database: DatabaseConfig,
pub llm: LlmConfig,
pub embeddings: EmbeddingsConfig,
pub channels: ChannelsConfig,
pub agent: AgentConfig,
pub safety: SafetyConfig,
@@ -30,6 +31,7 @@ impl Config {
Ok(Self {
database: DatabaseConfig::from_env()?,
llm: LlmConfig::from_env()?,
embeddings: EmbeddingsConfig::from_env()?,
channels: ChannelsConfig::from_env()?,
agent: AgentConfig::from_env()?,
safety: SafetyConfig::from_env()?,
@@ -106,6 +108,62 @@ impl LlmConfig {
}
}
/// Embeddings provider configuration.
#[derive(Debug, Clone)]
pub struct EmbeddingsConfig {
/// Whether embeddings are enabled.
pub enabled: bool,
/// Provider to use: "openai" or "nearai"
pub provider: String,
/// OpenAI API key (for OpenAI provider).
pub openai_api_key: Option<SecretString>,
/// Model to use for embeddings.
/// For OpenAI: "text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002"
/// For NEAR AI: Uses the configured session for auth.
pub model: String,
}
impl Default for EmbeddingsConfig {
fn default() -> Self {
Self {
enabled: false,
provider: "openai".to_string(),
openai_api_key: None,
model: "text-embedding-3-small".to_string(),
}
}
}
impl EmbeddingsConfig {
fn from_env() -> Result<Self, ConfigError> {
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
let provider = optional_env("EMBEDDING_PROVIDER")?.unwrap_or_else(|| "openai".to_string());
// Auto-enable if we have an API key
let enabled = optional_env("EMBEDDING_ENABLED")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "EMBEDDING_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(openai_api_key.is_some());
Ok(Self {
enabled,
provider,
openai_api_key,
model: optional_env("EMBEDDING_MODEL")?
.unwrap_or_else(|| "text-embedding-3-small".to_string()),
})
}
/// Get the OpenAI API key if configured.
pub fn openai_api_key(&self) -> Option<&str> {
self.openai_api_key.as_ref().map(|s| s.expose_secret())
}
}
/// Get the default session file path (~/.near-agent/session.json).
fn default_session_path() -> PathBuf {
dirs::home_dir()
src/history/mod.rs
+1 -1
View File
@@ -9,4 +9,4 @@ mod analytics;
mod store;
pub use analytics::{JobStats, ToolStats};
pub use store::Store;
pub use store::{LlmCallRecord, Store};
src/history/store.rs
+22 -19
View File
@@ -9,6 +9,19 @@ use crate::config::DatabaseConfig;
use crate::context::{ActionRecord, JobContext, JobState};
use crate::error::DatabaseError;
/// Record for an LLM call to be persisted.
#[derive(Debug, Clone)]
pub struct LlmCallRecord<'a> {
pub job_id: Option<Uuid>,
pub conversation_id: Option<Uuid>,
pub provider: &'a str,
pub model: &'a str,
pub input_tokens: u32,
pub output_tokens: u32,
pub cost: Decimal,
pub purpose: Option<&'a str>,
}
/// Database store for the agent.
pub struct Store {
pool: Pool,
@@ -335,17 +348,7 @@ impl Store {
// ==================== LLM Calls ====================
/// Record an LLM call.
pub async fn record_llm_call(
&self,
job_id: Option<Uuid>,
conversation_id: Option<Uuid>,
provider: &str,
model: &str,
input_tokens: u32,
output_tokens: u32,
cost: Decimal,
purpose: Option<&str>,
) -> Result<Uuid, DatabaseError> {
pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
let conn = self.conn().await?;
let id = Uuid::new_v4();
@@ -356,14 +359,14 @@ impl Store {
"#,
&[
&id,
&job_id,
&conversation_id,
&provider,
&model,
&(input_tokens as i32),
&(output_tokens as i32),
&cost,
&purpose,
&record.job_id,
&record.conversation_id,
&record.provider,
&record.model,
&(record.input_tokens as i32),
&(record.output_tokens as i32),
&record.cost,
&record.purpose,
],
)
.await?;
src/llm/reasoning.rs
+28 -4
View File
@@ -121,12 +121,29 @@ pub struct Reasoning {
llm: Arc<dyn LlmProvider>,
#[allow(dead_code)] // Will be used for sanitizing tool outputs
safety: Arc<SafetyLayer>,
/// Optional workspace for loading identity/system prompts.
workspace_system_prompt: Option<String>,
}
impl Reasoning {
/// Create a new reasoning engine.
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
Self { llm, safety }
Self {
llm,
safety,
workspace_system_prompt: None,
}
}
/// Set a custom system prompt from workspace identity files.
///
/// This is typically loaded from workspace.system_prompt() which combines
/// AGENTS.md, SOUL.md, USER.md, and IDENTITY.md into a unified prompt.
pub fn with_system_prompt(mut self, prompt: String) -> Self {
if !prompt.is_empty() {
self.workspace_system_prompt = Some(prompt);
}
self
}
/// Generate a plan for completing a goal.
@@ -361,11 +378,18 @@ Respond with a JSON plan in this format:
.map(|t| format!(" - {}: {}", t.name, t.description))
.collect();
format!(
"\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools directly when needed.",
"\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools when they would help accomplish the task.",
tool_list.join("\n")
)
};
// Include workspace identity prompt if available
let identity_section = if let Some(ref identity) = self.workspace_system_prompt {
format!("\n\n---\n\n{}", identity)
} else {
String::new()
};
format!(
r#"You are NEAR AI Agent, an autonomous assistant.
@@ -388,8 +412,8 @@ Here's the solution: [actual response to user]
- For code, use appropriate code blocks with language tags
- Call tools when they would help accomplish the task{}
The user sees ONLY content outside <thinking> tags."#,
tools_section
The user sees ONLY content outside <thinking> tags.{}"#,
tools_section, identity_section
)
}
src/main.rs
+74 -11
View File
@@ -6,7 +6,7 @@ use clap::Parser;
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
use near_agent::{
agent::Agent,
agent::{Agent, AgentDeps},
channels::{ChannelManager, HttpChannel, TuiChannel},
cli::{Cli, Command, run_tool_command},
config::Config,
@@ -17,7 +17,7 @@ use near_agent::{
ToolRegistry,
wasm::{WasmToolLoader, WasmToolRuntime},
},
workspace::Workspace,
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
};
#[tokio::main]
@@ -93,8 +93,8 @@ async fn main() -> anyhow::Result<()> {
Some(Arc::new(store))
};
// Initialize LLM provider
let llm = create_llm_provider(&config.llm, session)?;
// Initialize LLM provider (clone session so we can reuse it for embeddings)
let llm = create_llm_provider(&config.llm, session.clone())?;
tracing::info!("LLM provider initialized: {}", llm.model_name());
// Initialize safety layer
@@ -106,9 +106,52 @@ async fn main() -> anyhow::Result<()> {
tools.register_builtin_tools();
tracing::info!("Registered {} built-in tools", tools.count());
// Create embeddings provider if configured
let embeddings: Option<Arc<dyn EmbeddingProvider>> = if config.embeddings.enabled {
match config.embeddings.provider.as_str() {
"nearai" => {
tracing::info!(
"Embeddings enabled via NEAR AI (model: {})",
config.embeddings.model
);
Some(Arc::new(
NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone())
.with_model(&config.embeddings.model, 1536),
))
}
_ => {
// Default to OpenAI for unknown providers
if let Some(api_key) = config.embeddings.openai_api_key() {
tracing::info!(
"Embeddings enabled via OpenAI (model: {})",
config.embeddings.model
);
Some(Arc::new(OpenAiEmbeddings::with_model(
api_key,
&config.embeddings.model,
match config.embeddings.model.as_str() {
"text-embedding-3-large" => 3072,
_ => 1536, // text-embedding-3-small and ada-002
},
)))
} else {
tracing::warn!("Embeddings configured but OPENAI_API_KEY not set");
None
}
}
}
} else {
tracing::info!("Embeddings disabled (set OPENAI_API_KEY or EMBEDDING_ENABLED=true)");
None
};
// Register memory tools if database is available
if let Some(ref store) = store {
let workspace = Arc::new(Workspace::new("default", store.pool()));
let mut workspace = Workspace::new("default", store.pool());
if let Some(ref emb) = embeddings {
workspace = workspace.with_embeddings(emb.clone());
}
let workspace = Arc::new(workspace);
tools.register_memory_tools(workspace);
}
@@ -181,19 +224,39 @@ async fn main() -> anyhow::Result<()> {
}
// Create workspace for agent (shared with memory tools)
let workspace = store
.as_ref()
.map(|s| Arc::new(Workspace::new("default", s.pool())));
let workspace = store.as_ref().map(|s| {
let mut ws = Workspace::new("default", s.pool());
if let Some(ref emb) = embeddings {
ws = ws.with_embeddings(emb.clone());
}
Arc::new(ws)
});
// Backfill embeddings if we just enabled the provider
if let (Some(ws), Some(_)) = (&workspace, &embeddings) {
match ws.backfill_embeddings().await {
Ok(count) if count > 0 => {
tracing::info!("Backfilled embeddings for {} chunks", count);
}
Ok(_) => {}
Err(e) => {
tracing::warn!("Failed to backfill embeddings: {}", e);
}
}
}
// Create and run the agent
let agent = Agent::new(
config.agent.clone(),
let deps = AgentDeps {
store,
llm,
safety,
tools,
channels,
workspace,
};
let agent = Agent::new(
config.agent.clone(),
deps,
channels,
Some(config.heartbeat.clone()),
);
src/workspace/embeddings.rs
+141
View File
@@ -213,6 +213,147 @@ impl EmbeddingProvider for OpenAiEmbeddings {
}
}
/// NEAR AI embedding provider using the NEAR AI API.
///
/// Uses the same session-based auth as the LLM provider.
pub struct NearAiEmbeddings {
client: reqwest::Client,
base_url: String,
session: std::sync::Arc<crate::llm::SessionManager>,
model: String,
dimension: usize,
}
impl NearAiEmbeddings {
/// Create a new NEAR AI embedding provider.
///
/// Uses the same session manager as the LLM provider for auth.
pub fn new(
base_url: impl Into<String>,
session: std::sync::Arc<crate::llm::SessionManager>,
) -> Self {
Self {
client: reqwest::Client::new(),
base_url: base_url.into(),
session,
model: "text-embedding-3-small".to_string(),
dimension: 1536,
}
}
/// Use a specific model.
pub fn with_model(mut self, model: impl Into<String>, dimension: usize) -> Self {
self.model = model.into();
self.dimension = dimension;
self
}
}
#[derive(Debug, Serialize)]
struct NearAiEmbeddingRequest<'a> {
model: &'a str,
input: &'a [String],
}
#[derive(Debug, Deserialize)]
struct NearAiEmbeddingResponse {
data: Vec<NearAiEmbeddingData>,
}
#[derive(Debug, Deserialize)]
struct NearAiEmbeddingData {
embedding: Vec<f32>,
}
#[async_trait]
impl EmbeddingProvider for NearAiEmbeddings {
fn dimension(&self) -> usize {
self.dimension
}
fn model_name(&self) -> &str {
&self.model
}
fn max_input_length(&self) -> usize {
32_000
}
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
if text.len() > self.max_input_length() {
return Err(EmbeddingError::TextTooLong {
length: text.len(),
max: self.max_input_length(),
});
}
let embeddings = self.embed_batch(&[text.to_string()]).await?;
embeddings
.into_iter()
.next()
.ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string()))
}
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
use secrecy::ExposeSecret;
if texts.is_empty() {
return Ok(Vec::new());
}
let request = NearAiEmbeddingRequest {
model: &self.model,
input: texts,
};
let token = self
.session
.get_token()
.await
.map_err(|_| EmbeddingError::AuthFailed)?;
let url = format!("{}/v1/embeddings", self.base_url);
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", token.expose_secret()))
.json(&request)
.send()
.await?;
let status = response.status();
if status == reqwest::StatusCode::UNAUTHORIZED {
return Err(EmbeddingError::AuthFailed);
}
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
let retry_after = response
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<u64>().ok())
.map(std::time::Duration::from_secs);
return Err(EmbeddingError::RateLimited { retry_after });
}
if !status.is_success() {
let error_text = response.text().await.unwrap_or_default();
return Err(EmbeddingError::HttpError(format!(
"Status {}: {}",
status, error_text
)));
}
let result: NearAiEmbeddingResponse = response.json().await.map_err(|e| {
EmbeddingError::InvalidResponse(format!("Failed to parse response: {}", e))
})?;
Ok(result.data.into_iter().map(|d| d.embedding).collect())
}
}
/// A mock embedding provider for testing.
///
/// Generates deterministic embeddings based on text hash.
src/workspace/mod.rs
+1 -1
View File
@@ -48,7 +48,7 @@ mod search;
pub use chunker::{ChunkConfig, chunk_document};
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
pub use embeddings::{EmbeddingProvider, MockEmbeddings, OpenAiEmbeddings};
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
pub use repository::Repository;
pub use search::{SearchConfig, SearchResult};