diff --git a/CLAUDE.md b/CLAUDE.md index 724f51d8..32a89fec 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,7 +2,13 @@ ## Project Overview -LLM-powered autonomous agent for the NEAR AI marketplace. Handles multi-channel input (CLI, HTTP, Slack, Telegram), parallel job execution, extensible tools (including MCP), prompt injection defense, and self-repair for stuck jobs. +LLM-powered autonomous agent for the NEAR AI marketplace. Features: +- **Multi-channel input**: Full TUI (Ratatui), HTTP webhook with secret auth (Slack/Telegram stubs) +- **Parallel job execution** with state machine and self-repair for stuck jobs +- **Extensible tools**: Built-in tools, WASM sandbox, MCP client, dynamic builder +- **Persistent memory**: Workspace with hybrid search (FTS + vector via RRF) +- **Prompt injection defense**: Sanitizer, validator, policy rules, leak detection +- **Heartbeat system**: Proactive periodic execution with checklist ## Build & Test @@ -38,39 +44,69 @@ src/ │ ├── scheduler.rs # Parallel job scheduling │ ├── worker.rs # Per-job execution with LLM reasoning │ ├── self_repair.rs # Stuck job detection and recovery -│ └── heartbeat.rs # Proactive periodic execution (OpenClaw-inspired) +│ ├── heartbeat.rs # Proactive periodic execution +│ ├── session.rs # Session/thread/turn model with state machine +│ ├── session_manager.rs # Thread/session lifecycle management +│ ├── compaction.rs # Context window management with turn summarization +│ ├── context_monitor.rs # Memory pressure detection +│ ├── undo.rs # Turn-based undo/redo with checkpoints +│ ├── submission.rs # Submission parsing (undo, redo, compact, clear, etc.) +│ └── task.rs # Sub-task execution framework │ ├── channels/ # Multi-channel input │ ├── channel.rs # Channel trait, IncomingMessage, OutgoingResponse │ ├── manager.rs # ChannelManager merges streams -│ ├── cli.rs # Interactive CLI (stdin) -│ ├── http.rs # HTTP webhook (axum) +│ ├── cli/ # Full TUI with Ratatui +│ │ ├── mod.rs # TuiChannel implementation +│ │ ├── app.rs # Application state +│ │ ├── render.rs # UI rendering +│ │ ├── events.rs # Input handling +│ │ ├── overlay.rs # Approval overlays +│ │ └── composer.rs # Message composition +│ ├── http.rs # HTTP webhook (axum) with secret validation │ ├── slack.rs # Stub │ └── telegram.rs # Stub │ ├── safety/ # Prompt injection defense │ ├── sanitizer.rs # Pattern detection, content escaping │ ├── validator.rs # Input validation (length, encoding, patterns) -│ └── policy.rs # PolicyRule system with severity/actions +│ ├── policy.rs # PolicyRule system with severity/actions +│ └── leak_detector.rs # Secret detection (API keys, tokens, etc.) │ ├── llm/ # LLM integration (NEAR AI only) │ ├── provider.rs # LlmProvider trait, message types │ ├── nearai.rs # NEAR AI chat-api implementation -│ └── reasoning.rs # Planning, tool selection, evaluation +│ ├── reasoning.rs # Planning, tool selection, evaluation +│ └── session.rs # Session token management with auto-renewal │ ├── tools/ # Extensible tool system │ ├── tool.rs # Tool trait, ToolOutput, ToolError │ ├── registry.rs # ToolRegistry for discovery -│ ├── builder.rs # Dynamic tool creation (stub) -│ ├── sandbox.rs # Sandboxed execution (stub) +│ ├── sandbox.rs # Process-based sandbox (stub, superseded by wasm/) │ ├── builtin/ # Built-in tools │ │ ├── echo.rs, time.rs, json.rs, http.rs -│ │ ├── marketplace.rs, ecommerce.rs -│ │ ├── taskrabbit.rs, restaurant.rs -│ │ └── memory.rs # Memory tools (search, write, read) -│ └── mcp/ # Model Context Protocol -│ ├── client.rs # MCP client over HTTP -│ └── protocol.rs # JSON-RPC types +│ │ ├── file.rs # ReadFile, WriteFile, ListDir, ApplyPatch +│ │ ├── shell.rs # Shell command execution +│ │ ├── memory.rs # Memory tools (search, write, read, tree) +│ │ └── marketplace.rs, ecommerce.rs, taskrabbit.rs, restaurant.rs (stubs) +│ ├── builder/ # Dynamic tool building +│ │ ├── core.rs # BuildRequirement, SoftwareType, Language +│ │ ├── templates.rs # Project scaffolding +│ │ ├── testing.rs # Test harness integration +│ │ └── validation.rs # WASM validation +│ ├── mcp/ # Model Context Protocol +│ │ ├── client.rs # MCP client over HTTP +│ │ └── protocol.rs # JSON-RPC types +│ └── wasm/ # Full WASM sandbox (wasmtime) +│ ├── runtime.rs # Module compilation and caching +│ ├── wrapper.rs # Tool trait wrapper for WASM modules +│ ├── host.rs # Host functions (logging, time, workspace) +│ ├── limits.rs # Fuel metering and memory limiting +│ ├── allowlist.rs # Network endpoint allowlisting +│ ├── credential_injector.rs # Safe credential injection +│ ├── loader.rs # WASM tool discovery from filesystem +│ ├── rate_limiter.rs # Per-tool rate limiting +│ └── storage.rs # Linear memory persistence │ ├── workspace/ # Persistent memory system (OpenClaw-inspired) │ ├── mod.rs # Workspace struct, memory operations @@ -92,12 +128,17 @@ src/ │ └── learner.rs # Exponential moving average learning │ ├── evaluation/ # Success evaluation -│ ├── success.rs # SuccessEvaluator trait, RuleBasedEvaluator +│ ├── success.rs # SuccessEvaluator trait, RuleBasedEvaluator, LlmEvaluator │ └── metrics.rs # MetricsCollector, QualityMetrics │ +├── secrets/ # Secrets management +│ ├── crypto.rs # AES-256-GCM encryption +│ ├── store.rs # Secret storage +│ └── types.rs # Credential types +│ └── history/ # Persistence ├── store.rs # PostgreSQL repositories - └── analytics.rs # Aggregation queries for learning + └── analytics.rs # Aggregation queries (JobStats, ToolStats) ``` ## Key Patterns @@ -232,13 +273,23 @@ Key test patterns: ## Current Limitations / TODOs 1. **Slack/Telegram channels** - Stubs only, need implementation -2. **Tool sandboxing** - `sandbox.rs` is a stub, needs WASM integration -3. **Dynamic tool building** - `builder.rs` placeholder, needs LLM code generation -4. **Integration tests** - Need testcontainers setup for PostgreSQL -5. **MCP stdio transport** - Only HTTP transport implemented -6. **Workspace integration** - Memory tools need to be registered and workspace passed to workers -7. **Embedding backfill** - Background job to generate embeddings for chunks missing them -8. **Context compaction** - Auto-trigger memory preservation before context window fills +2. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations +3. **Integration tests** - Need testcontainers setup for PostgreSQL +4. **MCP stdio transport** - Only HTTP transport implemented +5. **Embedding backfill** - Background job to generate embeddings for chunks missing them +6. **Auto-context compaction** - Context monitor exists but doesn't auto-trigger (requires manual `/compact`) + +### Completed + +- ✅ **Workspace integration** - Memory tools registered, workspace passed to Agent and heartbeat +- ✅ **WASM sandboxing** - Full implementation in `tools/wasm/` with fuel metering, memory limits, capabilities +- ✅ **Dynamic tool building** - `tools/builder/` has LlmSoftwareBuilder with iterative build loop +- ✅ **HTTP webhook security** - Secret validation implemented, proper error handling (no panics) + +### Known Clippy Warnings (non-blocking) + +- `too_many_arguments` on Agent::new, Worker::new, Store::record_llm_call (refactor to config structs if desired) +- `implicit_saturating_sub` in CLI render (use `.saturating_sub()` instead of manual check) ## Adding a New Tool diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index f6644808..cf6a6492 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -368,9 +368,12 @@ impl Agent { ) .await; - // Call LLM with thread context + // Call LLM with thread context and available tools let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone()); - let context = ReasoningContext::new().with_messages(turn_messages); + let tool_defs = self.tools.tool_definitions().await; + let context = ReasoningContext::new() + .with_messages(turn_messages) + .with_tools(tool_defs); let llm_result = reasoning.respond(&context).await; // Re-acquire lock and check if interrupted diff --git a/src/agent/heartbeat.rs b/src/agent/heartbeat.rs index b6508ab1..a89cf75e 100644 --- a/src/agent/heartbeat.rs +++ b/src/agent/heartbeat.rs @@ -29,7 +29,6 @@ use std::time::Duration; use tokio::sync::mpsc; use crate::channels::OutgoingResponse; -use crate::error::WorkspaceError; use crate::llm::{ChatMessage, CompletionRequest, LlmProvider}; use crate::workspace::Workspace; @@ -277,21 +276,6 @@ pub fn spawn_heartbeat( }) } -/// Update heartbeat state in the database. -pub async fn update_heartbeat_state( - workspace: &Workspace, - last_run: chrono::DateTime, -) -> Result<(), WorkspaceError> { - // This would update the heartbeat_state table - // For now, we just log - tracing::debug!( - "Heartbeat state updated for user {} at {}", - workspace.user_id(), - last_run - ); - Ok(()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 85ec059f..2e7cd2fd 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -32,14 +32,12 @@ pub enum WorkerMessage { /// Status of a scheduled job. #[derive(Debug)] pub struct ScheduledJob { - pub job_id: Uuid, pub handle: JoinHandle<()>, pub tx: mpsc::Sender, } /// Status of a scheduled sub-task. struct ScheduledSubtask { - task_id: Uuid, handle: JoinHandle>, } @@ -137,7 +135,7 @@ impl Scheduler { self.jobs .write() .await - .insert(job_id, ScheduledJob { job_id, handle, tx }); + .insert(job_id, ScheduledJob { handle, tx }); tracing::info!("Scheduled job {} for execution", job_id); Ok(()) @@ -203,7 +201,6 @@ impl Scheduler { self.subtasks.write().await.insert( task_id, ScheduledSubtask { - task_id, handle: tokio::spawn(async move { // Wrap the handle to get its result match handle.await { diff --git a/src/agent/self_repair.rs b/src/agent/self_repair.rs index 5a62dca1..4d514405 100644 --- a/src/agent/self_repair.rs +++ b/src/agent/self_repair.rs @@ -66,10 +66,12 @@ pub trait SelfRepair: Send + Sync { /// Default self-repair implementation. pub struct DefaultSelfRepair { context_manager: Arc, + #[allow(dead_code)] // Will be used for time-based stuck detection stuck_threshold: Duration, max_repair_attempts: u32, store: Option>, builder: Option>, + #[allow(dead_code)] // Will be used for tool hot-reload after repair tools: Option>, } @@ -91,12 +93,14 @@ impl DefaultSelfRepair { } /// Add a Store for tool failure tracking. + #[allow(dead_code)] // Public API for configuring repair with persistence pub fn with_store(mut self, store: Arc) -> Self { self.store = Some(store); self } /// Add a Builder and ToolRegistry for automatic tool repair. + #[allow(dead_code)] // Public API for enabling automatic tool repair pub fn with_builder( mut self, builder: Arc, diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 4b5f8e44..94dbd7f6 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -33,9 +33,7 @@ pub struct Worker { /// Result of a tool execution with metadata for context building. struct ToolExecResult { - tool_name: String, result: Result, - duration: Duration, } impl Worker { @@ -290,7 +288,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# let store = self.store.clone(); async move { - let start = std::time::Instant::now(); let result = Self::execute_tool_inner( tools, context_manager, @@ -300,11 +297,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# ¶ms, ) .await; - ToolExecResult { - tool_name, - result, - duration: start.elapsed(), - } + ToolExecResult { result } } }) .collect(); diff --git a/src/channels/http.rs b/src/channels/http.rs index 35c3ff51..121b8eda 100644 --- a/src/channels/http.rs +++ b/src/channels/http.rs @@ -11,6 +11,7 @@ use axum::{ response::IntoResponse, routing::{get, post}, }; +use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; use tokio::sync::{RwLock, mpsc, oneshot}; use tokio_stream::wrappers::ReceiverStream; @@ -33,17 +34,25 @@ struct HttpChannelState { pending_responses: RwLock>>, /// Server shutdown signal. shutdown_tx: RwLock>>, + /// Expected webhook secret for authentication (if configured). + webhook_secret: Option, } impl HttpChannel { /// Create a new HTTP channel. pub fn new(config: HttpConfig) -> Self { + let webhook_secret = config + .webhook_secret + .as_ref() + .map(|s| s.expose_secret().to_string()); + Self { config, state: Arc::new(HttpChannelState { tx: RwLock::new(None), pending_responses: RwLock::new(std::collections::HashMap::new()), shutdown_tx: RwLock::new(None), + webhook_secret, }), } } @@ -90,8 +99,35 @@ async fn health_handler() -> impl IntoResponse { async fn webhook_handler( State(state): State>, Json(req): Json, -) -> impl IntoResponse { - // TODO: Validate secret if configured +) -> (StatusCode, Json) { + // Validate secret if configured + if let Some(ref expected_secret) = state.webhook_secret { + match &req.secret { + Some(provided) if provided == expected_secret => { + // Secret matches, continue + } + Some(_) => { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Invalid webhook secret".to_string()), + }), + ); + } + None => { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Webhook secret required".to_string()), + }), + ); + } + } + } let msg = IncomingMessage::new("http", &req.user_id, &req.content).with_metadata(serde_json::json!({ @@ -110,7 +146,7 @@ async fn process_message( state: Arc, msg: IncomingMessage, wait_for_response: bool, -) -> impl IntoResponse { +) -> (StatusCode, Json) { let msg_id = msg.id; // Set up response channel if waiting @@ -182,6 +218,26 @@ impl Channel for HttpChannel { let host = self.config.host.clone(); let port = self.config.port; + // Parse address before spawning so we can return errors + let addr: SocketAddr = + format!("{}:{}", host, port) + .parse() + .map_err(|e| ChannelError::StartupFailed { + name: "http".to_string(), + reason: format!("Invalid address '{}:{}': {}", host, port, e), + })?; + + // Bind listener before spawning so we can return errors + let listener = + tokio::net::TcpListener::bind(addr) + .await + .map_err(|e| ChannelError::StartupFailed { + name: "http".to_string(), + reason: format!("Failed to bind to {}: {}", addr, e), + })?; + + tracing::info!("HTTP channel listening on {}", addr); + // Create router let app = Router::new() .route("/health", get(health_handler)) @@ -192,23 +248,17 @@ impl Channel for HttpChannel { let (shutdown_tx, shutdown_rx) = oneshot::channel(); *self.state.shutdown_tx.write().await = Some(shutdown_tx); - // Spawn server + // Spawn server (listener is already bound, serve errors are logged) tokio::spawn(async move { - let addr: SocketAddr = format!("{}:{}", host, port) - .parse() - .expect("Invalid address"); - - tracing::info!("HTTP channel listening on {}", addr); - - let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); - - axum::serve(listener, app) + if let Err(e) = axum::serve(listener, app) .with_graceful_shutdown(async { let _ = shutdown_rx.await; tracing::info!("HTTP channel shutting down"); }) .await - .unwrap(); + { + tracing::error!("HTTP server error: {}", e); + } }); Ok(Box::pin(ReceiverStream::new(rx))) diff --git a/src/evaluation/success.rs b/src/evaluation/success.rs index 61b6029a..41e55bb8 100644 --- a/src/evaluation/success.rs +++ b/src/evaluation/success.rs @@ -82,12 +82,14 @@ impl RuleBasedEvaluator { } /// Set minimum action success rate. + #[allow(dead_code)] // Public API for configuring evaluation threshold pub fn with_min_success_rate(mut self, rate: f64) -> Self { self.min_action_success_rate = rate; self } /// Set maximum failures. + #[allow(dead_code)] // Public API for configuring failure tolerance pub fn with_max_failures(mut self, max: u32) -> Self { self.max_failures = max; self @@ -204,6 +206,7 @@ pub struct LlmEvaluator { impl LlmEvaluator { /// Create a new LLM-based evaluator. + #[allow(dead_code)] // Public API for LLM-based evaluation pub fn new(llm: Arc) -> Self { Self { llm } } diff --git a/src/history/analytics.rs b/src/history/analytics.rs index 24891bf6..4ee18c6a 100644 --- a/src/history/analytics.rs +++ b/src/history/analytics.rs @@ -1,22 +1,12 @@ //! Analytics and aggregation for learning. +//! +//! Analytics methods are implemented directly on [`Store`] for convenience. use rust_decimal::Decimal; use crate::error::DatabaseError; use crate::history::Store; -/// Analytics queries for the store. -pub struct Analytics<'a> { - store: &'a Store, -} - -impl<'a> Analytics<'a> { - /// Create analytics wrapper for a store. - pub fn new(store: &'a Store) -> Self { - Self { store } - } -} - /// Statistics about jobs. #[derive(Debug, Default)] pub struct JobStats { diff --git a/src/history/mod.rs b/src/history/mod.rs index 83616361..ea321821 100644 --- a/src/history/mod.rs +++ b/src/history/mod.rs @@ -8,5 +8,5 @@ mod analytics; mod store; -pub use analytics::{Analytics, JobStats, ToolStats}; +pub use analytics::{JobStats, ToolStats}; pub use store::Store; diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 6b486198..64487e5f 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -181,6 +181,12 @@ impl ToolCompletionRequest { self } + /// Set temperature. + pub fn with_temperature(mut self, temperature: f32) -> Self { + self.temperature = Some(temperature); + self + } + /// Set tool choice mode. pub fn with_tool_choice(mut self, choice: impl Into) -> Self { self.tool_choice = Some(choice.into()); diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index dda52a34..46035594 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -108,6 +108,7 @@ pub struct ToolSelection { /// Reasoning engine for the agent. pub struct Reasoning { llm: Arc, + #[allow(dead_code)] // Will be used for sanitizing tool outputs safety: Arc, } @@ -233,20 +234,48 @@ Respond in JSON format: } /// Generate a response to a user message. + /// + /// If tools are available in the context, uses tool completion mode. pub async fn respond(&self, context: &ReasoningContext) -> Result { - let system_prompt = self.build_conversation_prompt(); + let system_prompt = self.build_conversation_prompt(context); let mut messages = vec![ChatMessage::system(system_prompt)]; messages.extend(context.messages.clone()); - let request = CompletionRequest::new(messages) - .with_max_tokens(2048) - .with_temperature(0.7); + // If we have tools, use tool completion mode + if !context.available_tools.is_empty() { + let request = ToolCompletionRequest::new(messages, context.available_tools.clone()) + .with_max_tokens(4096) + .with_temperature(0.7) + .with_tool_choice("auto"); - let response = self.llm.complete(request).await?; + let response = self.llm.complete_with_tools(request).await?; - // Strip any internal thinking tags before returning to user - Ok(strip_thinking_tags(&response.content)) + // If there were tool calls, the content is usually just internal reasoning + // Don't show it - just acknowledge the tool calls (actual execution handled by caller) + if !response.tool_calls.is_empty() { + let tool_info: Vec = response + .tool_calls + .iter() + .map(|tc| format!("`{}({})`", tc.name, tc.arguments)) + .collect(); + return Ok(format!("[Calling tools: {}]", tool_info.join(", "))); + } + + // No tool calls - clean up the response + let content = response + .content + .unwrap_or_else(|| "I'm not sure how to respond to that.".to_string()); + Ok(clean_response(&content)) + } else { + // No tools, use simple completion + let request = CompletionRequest::new(messages) + .with_max_tokens(4096) + .with_temperature(0.7); + + let response = self.llm.complete(request).await?; + Ok(clean_response(&response.content)) + } } fn build_planning_prompt(&self, context: &ReasoningContext) -> String { @@ -292,15 +321,45 @@ Respond with a JSON plan in this format: ) } - fn build_conversation_prompt(&self) -> String { - r#"You are a helpful AI agent assistant. You help users with tasks by: -1. Understanding their requests clearly -2. Asking clarifying questions when needed -3. Providing accurate, helpful responses -4. Being honest about limitations + fn build_conversation_prompt(&self, context: &ReasoningContext) -> String { + let tools_section = if context.available_tools.is_empty() { + String::new() + } else { + let tool_list: Vec = context + .available_tools + .iter() + .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 - don't announce what you're going to do.", + tool_list.join("\n") + ) + }; -Be concise but thorough. If you're unsure, say so."# - .to_string() + format!( + r#"You are NEAR AI Agent, an autonomous assistant. + +CRITICAL: Never output your internal reasoning or thinking process. Your response must contain ONLY the final answer or action. + +FORBIDDEN patterns (never start with these): +- "The user wants..." / "The user is asking..." +- "I need to..." / "I should..." / "I will..." +- "Let me think..." / "Let me first..." +- "This is a request to..." +- Any self-narration about what you're doing + +CORRECT behavior: +- Answer questions directly +- Call tools without announcing it +- Ask clarifying questions if genuinely needed +- Provide code/content without preamble{} + +## Format +- Be concise +- Use markdown where helpful +- Code blocks with language tags"#, + tools_section + ) } fn parse_plan(&self, content: &str) -> Result { @@ -347,6 +406,12 @@ fn extract_json(text: &str) -> Option<&str> { } } +/// Clean up LLM response by stripping thinking tags and reasoning patterns. +fn clean_response(text: &str) -> String { + let text = strip_thinking_tags(text); + strip_reasoning_patterns(&text) +} + /// Strip `...` blocks from LLM output. /// /// Some models (especially Claude with extended thinking) include internal @@ -384,6 +449,71 @@ fn strip_thinking_tags(text: &str) -> String { cleaned } +/// Strip common reasoning/thinking patterns from the start of responses. +/// +/// Models sometimes output their thinking process as plain text despite +/// instructions not to. This strips common patterns like "The user wants...", +/// "Let me think...", "I need to...", etc. +fn strip_reasoning_patterns(text: &str) -> String { + let text = text.trim(); + + // Patterns that indicate internal reasoning (case-insensitive check) + let reasoning_prefixes = [ + "the user wants", + "the user is asking", + "the user would like", + "i need to", + "i should", + "i will", + "i'll", + "let me think", + "let me first", + "let me check", + "let me look", + "let me explore", + "let me search", + "this is a request", + "this request", + "to answer this", + "to help with this", + "first, i", + "okay, so", + "alright, ", + ]; + + // Find where reasoning ends and actual content begins + // Look for paragraph breaks or sentences that start the actual response + let lines: Vec<&str> = text.lines().collect(); + let mut skip_until = 0; + + for (i, line) in lines.iter().enumerate() { + let lower = line.to_lowercase(); + + // Check if this line starts with a reasoning pattern + let is_reasoning = reasoning_prefixes + .iter() + .any(|p| lower.trim_start().starts_with(p)); + + if is_reasoning { + skip_until = i + 1; + } else if !line.trim().is_empty() && skip_until <= i { + // Found non-reasoning content, stop looking + break; + } + } + + if skip_until > 0 && skip_until < lines.len() { + // Skip the reasoning lines and return the rest + let result = lines[skip_until..].join("\n").trim().to_string(); + if !result.is_empty() { + return result; + } + } + + // If we'd strip everything, just return the original (better than empty) + text.to_string() +} + #[cfg(test)] mod tests { use super::*; @@ -450,4 +580,45 @@ Here is my response to your question."#; let output = strip_thinking_tags(input); assert_eq!(output, "Hello"); } + + #[test] + fn test_strip_reasoning_patterns_basic() { + let input = "The user wants me to implement something.\n\nHere's the implementation:"; + let output = strip_reasoning_patterns(input); + assert_eq!(output, "Here's the implementation:"); + } + + #[test] + fn test_strip_reasoning_patterns_multiline() { + let input = r#"The user is asking about Telegram. +I need to think about what this involves. +Let me first check the existing code. + +Here's what I found in the codebase."#; + let output = strip_reasoning_patterns(input); + assert_eq!(output, "Here's what I found in the codebase."); + } + + #[test] + fn test_strip_reasoning_no_patterns() { + let input = "Here's a direct answer to your question."; + let output = strip_reasoning_patterns(input); + assert_eq!(output, "Here's a direct answer to your question."); + } + + #[test] + fn test_strip_reasoning_preserves_all_if_only_reasoning() { + // If stripping would leave nothing, keep the original + let input = "The user wants to know X."; + let output = strip_reasoning_patterns(input); + assert_eq!(output, "The user wants to know X."); + } + + #[test] + fn test_clean_response_combined() { + let input = + "Internal thoughtI need to check this.\n\nActual response here."; + let output = clean_response(input); + assert_eq!(output, "Actual response here."); + } } diff --git a/src/llm/session.rs b/src/llm/session.rs index 054d6c7e..05006e30 100644 --- a/src/llm/session.rs +++ b/src/llm/session.rs @@ -128,15 +128,62 @@ impl SessionManager { /// Ensure we have a valid session, triggering login flow if needed. /// /// If no token exists, triggers the OAuth login flow. If a token exists, - /// it is assumed valid until a 401 response indicates otherwise. + /// validates it by making a test API call. If validation fails, triggers + /// the login flow. pub async fn ensure_authenticated(&self) -> Result<(), LlmError> { - if self.has_token().await { - tracing::debug!("Session token present, assuming valid"); + if !self.has_token().await { + // No token, need to authenticate + return self.initiate_login().await; + } + + // Token exists, validate it by calling /v1/users/me + println!("Validating session..."); + match self.validate_token().await { + Ok(()) => { + println!("Session valid."); + Ok(()) + } + Err(e) => { + println!("Session expired or invalid: {}", e); + self.initiate_login().await + } + } + } + + /// Validate the current token by calling the /v1/users/me endpoint. + async fn validate_token(&self) -> Result<(), LlmError> { + use secrecy::ExposeSecret; + + let token = self.get_token().await?; + let url = format!("{}/v1/users/me", self.config.auth_base_url); + + let response = self + .client + .get(&url) + .header("Authorization", format!("Bearer {}", token.expose_secret())) + .send() + .await + .map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Validation request failed: {}", e), + })?; + + if response.status().is_success() { return Ok(()); } - // No token, need to authenticate - self.initiate_login().await + if response.status().as_u16() == 401 { + return Err(LlmError::SessionExpired { + provider: "nearai".to_string(), + }); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + Err(LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Validation failed: HTTP {}: {}", status, body), + }) } /// Handle an authentication failure (401 response). @@ -184,21 +231,72 @@ impl SessionManager { })?; let callback_url = format!("http://127.0.0.1:{}", port); - // Use GitHub OAuth (Google OAuth may not have the redirect URI configured) - let auth_url = format!( - "{}/v1/auth/github?frontend_callback={}", - self.config.auth_base_url, - urlencoding::encode(&callback_url) - ); - // Print auth URL + // Show auth provider menu println!(); println!("╔════════════════════════════════════════════════════════════════╗"); println!("║ NEAR AI Authentication ║"); println!("╠════════════════════════════════════════════════════════════════╣"); - println!("║ Please open the following URL in your browser to authenticate: ║"); + println!("║ Choose an authentication method: ║"); + println!("║ ║"); + println!("║ [1] GitHub ║"); + println!("║ [2] Google ║"); + println!("║ [3] NEAR Wallet (coming soon) ║"); + println!("║ ║"); println!("╚════════════════════════════════════════════════════════════════╝"); println!(); + print!("Enter choice [1-3]: "); + + // Flush stdout to ensure prompt is displayed + use std::io::Write; + std::io::stdout().flush().ok(); + + // Read user choice + let mut choice = String::new(); + std::io::stdin() + .read_line(&mut choice) + .map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to read input: {}", e), + })?; + + let (auth_provider, auth_url) = match choice.trim() { + "1" | "" => { + let url = format!( + "{}/v1/auth/github?frontend_callback={}", + self.config.auth_base_url, + urlencoding::encode(&callback_url) + ); + ("github", url) + } + "2" => { + let url = format!( + "{}/v1/auth/google?frontend_callback={}", + self.config.auth_base_url, + urlencoding::encode(&callback_url) + ); + ("google", url) + } + "3" => { + println!(); + println!("NEAR Wallet authentication is not yet implemented."); + println!("Please use GitHub or Google for now."); + return Err(LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: "NEAR Wallet auth not yet implemented".to_string(), + }); + } + _ => { + return Err(LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Invalid choice: {}", choice.trim()), + }); + } + }; + + println!(); + println!("Opening {} authentication...", auth_provider); + println!(); println!(" {}", auth_url); println!(); @@ -215,7 +313,8 @@ impl SessionManager { // Wait for callback with timeout // The API redirects to: {frontend_callback}/auth/callback?token=X&session_id=X&expires_at=X&is_new_user=X let timeout = std::time::Duration::from_secs(300); // 5 minutes - let (session_token, auth_provider) = tokio::time::timeout(timeout, async { + let selected_provider = auth_provider.to_string(); + let (session_token, auth_provider) = tokio::time::timeout(timeout, async move { loop { let (mut socket, _) = listener.accept().await.map_err(|e| { LlmError::SessionRenewalFailed { @@ -252,24 +351,82 @@ impl SessionManager { } if let Some(token) = token { - // Send success response + // Send success response with nice styling let response = concat!( "HTTP/1.1 200 OK\r\n", - "Content-Type: text/html\r\n", + "Content-Type: text/html; charset=utf-8\r\n", "Connection: close\r\n", "\r\n", - "NEAR AI Auth", - "", - "

✓ Authentication successful!

", - "

You can close this window and return to the terminal.

", - "" + "\n", + "\n", + "\n", + " \n", + " NEAR AI - Authentication Successful\n", + " \n", + "\n", + "\n", + "
\n", + "
✓
\n", + "

Authentication Successful

\n", + "

You can close this window and return to the terminal.

\n", + "
NEAR AI Agent
\n", + "
\n", + "\n", + "" ); let _ = socket.write_all(response.as_bytes()).await; let _ = socket.shutdown().await; - // Provider is github since we used the github endpoint - return Ok::<_, LlmError>((token, Some("github".to_string()))); + return Ok::<_, LlmError>((token, Some(selected_provider.clone()))); } } } diff --git a/src/main.rs b/src/main.rs index d90ea5ed..fc830e2e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -41,18 +41,34 @@ async fn main() -> anyhow::Result<()> { } } - // Create TUI channel early so we can hook up logging - // (channel is created but not started until agent.run()) + // Load configuration first (before any logging setup) + // so we can do auth before TUI starts + let _ = dotenvy::dotenv(); // Load .env if present + let config = Config::from_env()?; + + // Initialize session manager and authenticate BEFORE TUI setup + // This allows the auth menu to display cleanly without TUI interference + let session_config = SessionConfig { + auth_base_url: config.llm.nearai.auth_base_url.clone(), + session_path: config.llm.nearai.session_path.clone(), + ..Default::default() + }; + let session = create_session_manager(session_config).await; + + // Ensure we're authenticated before proceeding (may trigger login flow) + // This happens before TUI so the menu displays correctly + session.ensure_authenticated().await?; + + // Now create TUI channel and set up logging let tui_channel = TuiChannel::new(); let tui_log_writer = tui_channel.log_writer(); - // Initialize tracing with both stderr (for pre-TUI output) and TUI writer + // Initialize tracing with TUI writer let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("near_agent=info,tower_http=debug")); tracing_subscriber::registry() .with(env_filter) - // TUI layer: sends logs to TUI status line (once TUI is running) .with( tracing_subscriber::fmt::layer() .with_writer(tui_log_writer) @@ -63,10 +79,8 @@ async fn main() -> anyhow::Result<()> { .init(); tracing::info!("Starting NEAR Agent..."); - - // Load configuration - let config = Config::from_env()?; tracing::info!("Loaded configuration for agent: {}", config.agent.name); + tracing::info!("NEAR AI session authenticated"); // Initialize database store (optional for testing) let store = if cli.no_db { @@ -79,18 +93,6 @@ async fn main() -> anyhow::Result<()> { Some(Arc::new(store)) }; - // Initialize session manager for NEAR AI authentication - let session_config = SessionConfig { - auth_base_url: config.llm.nearai.auth_base_url.clone(), - session_path: config.llm.nearai.session_path.clone(), - ..Default::default() - }; - let session = create_session_manager(session_config).await; - - // Ensure we're authenticated before proceeding (may trigger login flow) - session.ensure_authenticated().await?; - tracing::info!("NEAR AI session authenticated"); - // Initialize LLM provider let llm = create_llm_provider(&config.llm, session)?; tracing::info!("LLM provider initialized: {}", llm.model_name()); diff --git a/src/tools/builder/testing.rs b/src/tools/builder/testing.rs index 64b820f9..84fd4f3e 100644 --- a/src/tools/builder/testing.rs +++ b/src/tools/builder/testing.rs @@ -411,6 +411,7 @@ fn get_json_path<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a ser } /// Generate basic test cases for a tool based on its schema. +#[allow(dead_code)] // Public API for auto-generating test cases pub fn generate_basic_tests(name: &str, input_schema: &serde_json::Value) -> TestSuite { let mut suite = TestSuite::new(format!("{}_basic_tests", name)); suite.description = Some("Auto-generated basic tests".to_string()); diff --git a/src/tools/sandbox.rs b/src/tools/sandbox.rs index 891ca92c..ed605429 100644 --- a/src/tools/sandbox.rs +++ b/src/tools/sandbox.rs @@ -51,6 +51,7 @@ pub struct SandboxResult { /// Sandbox for executing untrusted code. pub struct ToolSandbox { + #[allow(dead_code)] // Will be used when sandbox execution is implemented config: SandboxConfig, } diff --git a/src/tools/wasm/limits.rs b/src/tools/wasm/limits.rs index fdd6a847..679ce88d 100644 --- a/src/tools/wasm/limits.rs +++ b/src/tools/wasm/limits.rs @@ -68,10 +68,12 @@ pub struct WasmResourceLimiter { /// Maximum tables allowed. max_tables: u32, /// Current table count. + #[allow(dead_code)] // Reserved for table limit enforcement tables_created: u32, /// Maximum instances allowed. max_instances: u32, /// Current instance count. + #[allow(dead_code)] // Reserved for instance limit enforcement instances_created: u32, } diff --git a/src/workspace/chunker.rs b/src/workspace/chunker.rs index 98f3873d..7ab6f5b5 100644 --- a/src/workspace/chunker.rs +++ b/src/workspace/chunker.rs @@ -116,6 +116,7 @@ pub fn chunk_document(content: &str, config: ChunkConfig) -> Vec { /// Split content by paragraphs first, then chunk. /// /// This is better for preserving semantic boundaries. +#[allow(dead_code)] // Alternative chunking strategy for paragraph-aware indexing pub fn chunk_by_paragraphs(content: &str, config: ChunkConfig) -> Vec { if content.is_empty() { return Vec::new();