diff --git a/.env.example b/.env.example index de67ea78..43c99d1d 100644 --- a/.env.example +++ b/.env.example @@ -4,9 +4,12 @@ DATABASE_POOL_SIZE=10 # LLM Provider (NEAR AI) # NEAR AI provides a unified interface to all models with user authentication -NEARAI_SESSION_TOKEN=sess_... +# Session token is stored in ~/.near-agent/session.json and managed automatically. +# On first run, the agent will open a browser for OAuth authentication. NEARAI_MODEL=claude-3-5-sonnet-20241022 -NEARAI_BASE_URL=https://private.near.ai +NEARAI_BASE_URL=https://api.near.ai +NEARAI_AUTH_URL=https://private.near.ai +# NEARAI_SESSION_PATH=~/.near-agent/session.json # optional, default shown # Channel Configuration # CLI is always enabled diff --git a/Cargo.lock b/Cargo.lock index 4b6bda9e..8c96241c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1714,6 +1714,25 @@ dependencies = [ "serde", ] +[[package]] +name = "is-docker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928bae27f42bc99b60d9ac7334e3a21d10ad8f1835a4e12ec3ec0464765ed1b3" +dependencies = [ + "once_cell", +] + +[[package]] +name = "is-wsl" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "173609498df190136aa7dea1a91db051746d339e18476eed5ca40521f02d7aa5" +dependencies = [ + "is-docker", + "once_cell", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -1958,6 +1977,7 @@ dependencies = [ "dotenvy", "futures", "hkdf", + "open", "pgvector", "postgres-types", "pretty_assertions", @@ -1983,6 +2003,7 @@ dependencies = [ "tower-http", "tracing", "tracing-subscriber", + "urlencoding", "uuid", "wasmparser 0.220.1", "wasmtime", @@ -2061,6 +2082,17 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" +[[package]] +name = "open" +version = "5.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43bb73a7fa3799b198970490a51174027ba0d4ec504b03cd08caf513d40024bc" +dependencies = [ + "is-wsl", + "libc", + "pathdiff", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -2127,6 +2159,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pathdiff" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -3909,6 +3947,12 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index 196a2023..200bac58 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,12 @@ dirs = "6" # Secrecy for sensitive values secrecy = { version = "0.10", features = ["serde"] } +# URL encoding for OAuth flow +urlencoding = "2" + +# Open URLs in browser +open = "5" + # Vector embeddings for semantic search # The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres) pgvector = { version = "0.4", features = ["postgres"] } diff --git a/src/config.rs b/src/config.rs index 785be960..40dc0cfd 100644 --- a/src/config.rs +++ b/src/config.rs @@ -76,30 +76,42 @@ pub struct LlmConfig { /// NEAR AI chat-api configuration. #[derive(Debug, Clone)] pub struct NearAiConfig { - /// Session token for authentication (format: sess_xxx) - pub session_token: SecretString, /// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o") pub model: String, /// Base URL for the NEAR AI chat-api (default: https://api.near.ai) pub base_url: String, + /// Base URL for auth/refresh endpoints (default: https://private.near.ai) + pub auth_base_url: String, + /// Path to session file (default: ~/.near-agent/session.json) + pub session_path: PathBuf, } impl LlmConfig { fn from_env() -> Result { - let session_token = required_env("NEARAI_SESSION_TOKEN")?; - Ok(Self { nearai: NearAiConfig { - session_token: SecretString::from(session_token), model: optional_env("NEARAI_MODEL")? .unwrap_or_else(|| "claude-3-5-sonnet-20241022".to_string()), base_url: optional_env("NEARAI_BASE_URL")? .unwrap_or_else(|| "https://api.near.ai".to_string()), + auth_base_url: optional_env("NEARAI_AUTH_URL")? + .unwrap_or_else(|| "https://private.near.ai".to_string()), + session_path: optional_env("NEARAI_SESSION_PATH")? + .map(PathBuf::from) + .unwrap_or_else(default_session_path), }, }) } } +/// Get the default session file path (~/.near-agent/session.json). +fn default_session_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".near-agent") + .join("session.json") +} + /// Channel configurations. #[derive(Debug, Clone)] pub struct ChannelsConfig { diff --git a/src/error.rs b/src/error.rs index f41fc82c..0f0d304a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -140,11 +140,20 @@ pub enum LlmError { #[error("Authentication failed for provider {provider}")] AuthFailed { provider: String }, + #[error("Session expired for provider {provider}")] + SessionExpired { provider: String }, + + #[error("Session renewal failed for provider {provider}: {reason}")] + SessionRenewalFailed { provider: String, reason: String }, + #[error("HTTP error: {0}")] Http(#[from] reqwest::Error), #[error("JSON error: {0}")] Json(#[from] serde_json::Error), + + #[error("IO error: {0}")] + Io(#[from] std::io::Error), } /// Tool execution errors. diff --git a/src/llm/mod.rs b/src/llm/mod.rs index bde4a55a..528119c3 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -5,6 +5,7 @@ mod nearai; mod provider; mod reasoning; +pub mod session; pub use nearai::NearAiProvider; pub use provider::{ @@ -12,6 +13,7 @@ pub use provider::{ ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, }; pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, ToolSelection}; +pub use session::{SessionConfig, SessionManager, create_session_manager}; use std::sync::Arc; @@ -19,6 +21,15 @@ use crate::config::LlmConfig; use crate::error::LlmError; /// Create an LLM provider based on configuration. -pub fn create_llm_provider(config: &LlmConfig) -> Result, LlmError> { - Ok(Arc::new(NearAiProvider::new(config.nearai.clone()))) +/// +/// Requires a session manager for authentication. Use `create_session_manager` +/// to create one from the config. +pub fn create_llm_provider( + config: &LlmConfig, + session: Arc, +) -> Result, LlmError> { + Ok(Arc::new(NearAiProvider::new( + config.nearai.clone(), + session, + ))) } diff --git a/src/llm/nearai.rs b/src/llm/nearai.rs index d91a602e..5e34f1f0 100644 --- a/src/llm/nearai.rs +++ b/src/llm/nearai.rs @@ -3,6 +3,8 @@ //! 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::sync::Arc; + use async_trait::async_trait; use reqwest::Client; use rust_decimal::Decimal; @@ -16,23 +18,28 @@ use crate::llm::provider::{ ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, }; +use crate::llm::session::SessionManager; /// NEAR AI Chat API provider. pub struct NearAiProvider { client: Client, config: NearAiConfig, + session: Arc, } impl NearAiProvider { - /// Create a new NEAR AI provider. - pub fn new(config: NearAiConfig) -> Self { - // Create client with reasonable timeout + /// Create a new NEAR AI provider with a session manager. + pub fn new(config: NearAiConfig, session: Arc) -> Self { let client = Client::builder() - .timeout(std::time::Duration::from_secs(120)) // 2 minute timeout for LLM calls + .timeout(std::time::Duration::from_secs(120)) .build() .unwrap_or_else(|_| Client::new()); - Self { client, config } + Self { + client, + config, + session, + } } fn api_url(&self, path: &str) -> String { @@ -43,12 +50,32 @@ impl NearAiProvider { ) } + /// Send a request with automatic session renewal on 401. async fn send_request Deserialize<'de>>( &self, path: &str, body: &T, + ) -> Result { + // Try the request, handling session expiration + match self.send_request_inner(path, body).await { + Ok(result) => Ok(result), + Err(LlmError::SessionExpired { .. }) => { + // Session expired, attempt renewal and retry once + self.session.handle_auth_failure().await?; + self.send_request_inner(path, body).await + } + Err(e) => Err(e), + } + } + + /// Inner request implementation without retry logic. + async fn send_request_inner Deserialize<'de>>( + &self, + path: &str, + body: &T, ) -> Result { let url = self.api_url(path); + let token = self.session.get_token().await?; tracing::debug!("Sending request to NEAR AI: {}", url); tracing::debug!("Request body: {:?}", body); @@ -56,10 +83,7 @@ impl NearAiProvider { let response = self .client .post(&url) - .header( - "Authorization", - format!("Bearer {}", self.config.session_token.expose_secret()), - ) + .header("Authorization", format!("Bearer {}", token.expose_secret())) .header("Content-Type", "application/json") .json(body) .send() @@ -76,6 +100,24 @@ impl NearAiProvider { tracing::debug!("NEAR AI response body: {}", response_text); if !status.is_success() { + // Check for session expiration (401 with specific message patterns) + if status.as_u16() == 401 { + let is_session_expired = response_text.to_lowercase().contains("session") + && (response_text.to_lowercase().contains("expired") + || response_text.to_lowercase().contains("invalid")); + + if is_session_expired { + return Err(LlmError::SessionExpired { + provider: "nearai".to_string(), + }); + } + + // Generic 401 without session expiration indication + return Err(LlmError::AuthFailed { + provider: "nearai".to_string(), + }); + } + // Try to parse as JSON error if let Ok(error) = serde_json::from_str::(&response_text) { if status.as_u16() == 429 { diff --git a/src/llm/session.rs b/src/llm/session.rs new file mode 100644 index 00000000..ce0f498c --- /dev/null +++ b/src/llm/session.rs @@ -0,0 +1,578 @@ +//! Session management for NEAR AI authentication. +//! +//! Handles session token persistence, expiration detection, and renewal via +//! OAuth flow. Tokens are stored in `~/.near-agent/session.json` and refreshed +//! automatically when expired. + +use std::path::PathBuf; +use std::sync::Arc; + +use chrono::{DateTime, Utc}; +use reqwest::Client; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, RwLock}; + +use crate::error::LlmError; + +/// Session data persisted to disk. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionData { + pub session_token: String, + pub created_at: DateTime, + #[serde(default)] + pub auth_provider: Option, +} + +/// Configuration for session management. +#[derive(Debug, Clone)] +pub struct SessionConfig { + /// Base URL for auth endpoints (e.g., https://private.near.ai). + pub auth_base_url: String, + /// Path to session file (e.g., ~/.near-agent/session.json). + pub session_path: PathBuf, + /// Port range for OAuth callback server. + pub callback_port_range: (u16, u16), +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + auth_base_url: "https://private.near.ai".to_string(), + session_path: default_session_path(), + callback_port_range: (9876, 9886), + } + } +} + +/// Get the default session file path (~/.near-agent/session.json). +pub fn default_session_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".near-agent") + .join("session.json") +} + +/// Manages NEAR AI session tokens with persistence and automatic renewal. +pub struct SessionManager { + config: SessionConfig, + client: Client, + /// Current token in memory. + token: RwLock>, + /// Prevents thundering herd during concurrent 401s. + renewal_lock: Mutex<()>, +} + +impl SessionManager { + /// Create a new session manager and load any existing token from disk. + pub fn new(config: SessionConfig) -> Self { + let manager = Self { + config, + client: Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .unwrap_or_else(|_| Client::new()), + token: RwLock::new(None), + renewal_lock: Mutex::new(()), + }; + + // Try to load existing session synchronously during construction + if let Ok(data) = std::fs::read_to_string(&manager.config.session_path) { + if let Ok(session) = serde_json::from_str::(&data) { + // We can't await here, so we use try_write + if let Ok(mut guard) = manager.token.try_write() { + *guard = Some(SecretString::from(session.session_token)); + tracing::info!( + "Loaded session token from {}", + manager.config.session_path.display() + ); + } + } + } + + manager + } + + /// Create a session manager and load token asynchronously. + pub async fn new_async(config: SessionConfig) -> Self { + let manager = Self { + config, + client: Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .unwrap_or_else(|_| Client::new()), + token: RwLock::new(None), + renewal_lock: Mutex::new(()), + }; + + if let Err(e) = manager.load_session().await { + tracing::debug!("No existing session found: {}", e); + } + + manager + } + + /// Get the current session token, returning an error if not authenticated. + pub async fn get_token(&self) -> Result { + let guard = self.token.read().await; + guard.clone().ok_or_else(|| LlmError::AuthFailed { + provider: "nearai".to_string(), + }) + } + + /// Check if we have a valid token (doesn't verify with server). + pub async fn has_token(&self) -> bool { + self.token.read().await.is_some() + } + + /// Ensure we have a valid session, triggering login flow if needed. + /// + /// This proactively validates the token with the server, so we catch + /// expired sessions early rather than failing on the first LLM request. + pub async fn ensure_authenticated(&self) -> Result<(), LlmError> { + if !self.has_token().await { + // No token at all, need to authenticate + return self.initiate_login().await; + } + + // We have a token, but let's validate it's not expired + match self.validate_token().await { + Ok(()) => { + tracing::debug!("Session token validated successfully"); + Ok(()) + } + Err(e) => { + tracing::warn!("Session token validation failed: {}, will re-authenticate", e); + self.initiate_login().await + } + } + } + + /// Validate the current token with the server. + /// + /// Attempts to refresh the token to verify it's still valid. If refresh + /// succeeds, we also get a fresh token as a bonus. + async fn validate_token(&self) -> Result<(), LlmError> { + // Try to refresh - this validates the token and gives us a fresh one + match self.refresh_session().await { + Ok(new_token) => { + let mut guard = self.token.write().await; + *guard = Some(new_token); + Ok(()) + } + Err(e) => Err(e), + } + } + + /// Handle an authentication failure (401 response). + /// + /// First attempts to refresh the session. If refresh fails, initiates + /// a full re-authentication flow. + /// + /// Returns `true` if authentication was recovered, `false` if it failed. + pub async fn handle_auth_failure(&self) -> Result<(), LlmError> { + // Acquire renewal lock to prevent thundering herd + let _guard = self.renewal_lock.lock().await; + + // Double-check: maybe another task already renewed + // (We don't have a way to verify without making a request, + // so we just try to refresh) + + tracing::info!("Session expired, attempting refresh..."); + + // Try refresh first + match self.refresh_session().await { + Ok(new_token) => { + let mut guard = self.token.write().await; + *guard = Some(new_token); + tracing::info!("Session refreshed successfully"); + return Ok(()); + } + Err(e) => { + tracing::warn!( + "Session refresh failed: {}, will need to re-authenticate", + e + ); + } + } + + // Refresh failed, need full re-authentication + self.initiate_login().await + } + + /// Attempt to refresh the session using the current token. + async fn refresh_session(&self) -> Result { + let current_token = self.get_token().await?; + + let url = format!("{}/auth/refresh", self.config.auth_base_url); + tracing::debug!("Attempting session refresh at {}", url); + + let response = self + .client + .post(&url) + .header( + "Authorization", + format!("Bearer {}", current_token.expose_secret()), + ) + .send() + .await + .map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("HTTP request failed: {}", e), + })?; + + if response.status().is_success() { + let body: RefreshResponse = + response + .json() + .await + .map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to parse response: {}", e), + })?; + + let new_token = SecretString::from(body.session_token.clone()); + self.save_session(&body.session_token, None).await?; + return Ok(new_token); + } + + // Refresh endpoint returned non-success + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + Err(LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("HTTP {}: {}", status, body), + }) + } + + /// Start the OAuth login flow. + /// + /// 1. Find an available port for the callback server + /// 2. Print the auth URL and attempt to open browser + /// 3. Wait for OAuth callback with session token + /// 4. Save and return the token + async fn initiate_login(&self) -> Result<(), LlmError> { + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + use tokio::net::TcpListener; + + // Find an available port + let mut listener = None; + let mut port = 0; + + for p in self.config.callback_port_range.0..=self.config.callback_port_range.1 { + match TcpListener::bind(format!("127.0.0.1:{}", p)).await { + Ok(l) => { + listener = Some(l); + port = p; + break; + } + Err(_) => continue, + } + } + + let listener = listener.ok_or_else(|| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!( + "Could not find available port in range {}-{}", + self.config.callback_port_range.0, self.config.callback_port_range.1 + ), + })?; + + let callback_url = format!("http://127.0.0.1:{}", port); + let auth_url = format!( + "{}/v1/auth/google?frontend_callback={}", + self.config.auth_base_url, + urlencoding::encode(&callback_url) + ); + + // Print auth URL + println!(); + println!("╔════════════════════════════════════════════════════════════════╗"); + println!("║ NEAR AI Authentication ║"); + println!("╠════════════════════════════════════════════════════════════════╣"); + println!("║ Please open the following URL in your browser to authenticate: ║"); + println!("╚════════════════════════════════════════════════════════════════╝"); + println!(); + println!(" {}", auth_url); + println!(); + + // Try to open browser automatically + if let Err(e) = open::that(&auth_url) { + tracing::debug!("Could not open browser automatically: {}", e); + println!("(Could not open browser automatically, please copy the URL above)"); + } else { + println!("(Opening browser...)"); + } + println!(); + println!("Waiting for authentication..."); + + // Wait for callback with timeout + // The API redirects to: {frontend_callback}/auth/callback?token=X&session_id=X&expires_at=X&is_new_user=X + let timeout = std::time::Duration::from_secs(300); // 5 minutes + let (session_token, auth_provider) = tokio::time::timeout(timeout, async { + loop { + let (mut socket, _) = listener.accept().await.map_err(|e| { + LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to accept connection: {}", e), + } + })?; + + let mut reader = BufReader::new(&mut socket); + let mut request_line = String::new(); + reader.read_line(&mut request_line).await.map_err(|e| { + LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to read request: {}", e), + } + })?; + + // Parse GET /auth/callback?token=xxx&session_id=xxx&expires_at=xxx&is_new_user=xxx HTTP/1.1 + if let Some(path) = request_line.split_whitespace().nth(1) { + if path.starts_with("/auth/callback") { + // Parse query parameters + if let Some(query) = path.split('?').nth(1) { + let mut token = None; + + for param in query.split('&') { + let parts: Vec<&str> = param.splitn(2, '=').collect(); + if parts.len() == 2 && parts[0] == "token" { + token = Some( + urlencoding::decode(parts[1]) + .unwrap_or_else(|_| parts[1].into()) + .into_owned(), + ); + } + } + + if let Some(token) = token { + // Send success response + let response = concat!( + "HTTP/1.1 200 OK\r\n", + "Content-Type: text/html\r\n", + "Connection: close\r\n", + "\r\n", + "NEAR AI Auth", + "", + "

✓ Authentication successful!

", + "

You can close this window and return to the terminal.

", + "" + ); + + let _ = socket.write_all(response.as_bytes()).await; + let _ = socket.shutdown().await; + + // Provider is google since we used the google endpoint + return Ok::<_, LlmError>((token, Some("google".to_string()))); + } + } + } + } + + // Not the callback we're looking for, send 404 + let response = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n"; + let _ = socket.write_all(response.as_bytes()).await; + } + }) + .await + .map_err(|_| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: "Authentication timed out after 5 minutes".to_string(), + })??; + + // Save the token + self.save_session(&session_token, auth_provider.as_deref()) + .await?; + + // Update in-memory token + { + let mut guard = self.token.write().await; + *guard = Some(SecretString::from(session_token)); + } + + println!(); + println!("✓ Authentication successful!"); + println!(); + + Ok(()) + } + + /// Save session data to disk. + async fn save_session(&self, token: &str, auth_provider: Option<&str>) -> Result<(), LlmError> { + let session = SessionData { + session_token: token.to_string(), + created_at: Utc::now(), + auth_provider: auth_provider.map(String::from), + }; + + // Ensure parent directory exists + 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( + e.kind(), + format!("Failed to create session directory: {}", e), + )) + })?; + } + + let json = + serde_json::to_string_pretty(&session).map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to serialize session: {}", e), + })?; + + tokio::fs::write(&self.config.session_path, json) + .await + .map_err(|e| { + LlmError::Io(std::io::Error::new( + e.kind(), + format!( + "Failed to write session file {}: {}", + self.config.session_path.display(), + e + ), + )) + })?; + + tracing::debug!("Session saved to {}", self.config.session_path.display()); + Ok(()) + } + + /// Load session data from disk. + async fn load_session(&self) -> Result<(), LlmError> { + let data = tokio::fs::read_to_string(&self.config.session_path) + .await + .map_err(|e| { + LlmError::Io(std::io::Error::new( + e.kind(), + format!( + "Failed to read session file {}: {}", + self.config.session_path.display(), + e + ), + )) + })?; + + let session: SessionData = + serde_json::from_str(&data).map_err(|e| LlmError::SessionRenewalFailed { + provider: "nearai".to_string(), + reason: format!("Failed to parse session file: {}", e), + })?; + + { + let mut guard = self.token.write().await; + *guard = Some(SecretString::from(session.session_token)); + } + + tracing::info!( + "Loaded session from {} (created: {})", + self.config.session_path.display(), + session.created_at + ); + + Ok(()) + } + + /// Set token directly (useful for testing or migration from env var). + pub async fn set_token(&self, token: SecretString) { + let mut guard = self.token.write().await; + *guard = Some(token); + } +} + +/// Response from the refresh endpoint. +#[derive(Debug, Deserialize)] +struct RefreshResponse { + session_token: String, +} + +/// Create a session manager from a config, migrating from env var if present. +pub async fn create_session_manager(config: SessionConfig) -> Arc { + let manager = SessionManager::new_async(config).await; + + // Check for legacy env var and migrate if present and no file token + if !manager.has_token().await { + if let Ok(token) = std::env::var("NEARAI_SESSION_TOKEN") { + if !token.is_empty() { + tracing::info!("Migrating session token from NEARAI_SESSION_TOKEN env var to file"); + manager.set_token(SecretString::from(token.clone())).await; + if let Err(e) = manager.save_session(&token, None).await { + tracing::warn!("Failed to save migrated session: {}", e); + } + } + } + } + + Arc::new(manager) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[tokio::test] + async fn test_session_save_load() { + let dir = tempdir().unwrap(); + let session_path = dir.path().join("session.json"); + + let config = SessionConfig { + auth_base_url: "https://example.com".to_string(), + session_path: session_path.clone(), + callback_port_range: (9900, 9910), + }; + + let manager = SessionManager::new_async(config.clone()).await; + + // No token initially + assert!(!manager.has_token().await); + + // Save a token + manager + .save_session("test_token_123", Some("near")) + .await + .unwrap(); + manager + .set_token(SecretString::from("test_token_123")) + .await; + + // Verify it's set + assert!(manager.has_token().await); + let token = manager.get_token().await.unwrap(); + assert_eq!(token.expose_secret(), "test_token_123"); + + // Create new manager and verify it loads the token + let manager2 = SessionManager::new_async(config).await; + assert!(manager2.has_token().await); + let token2 = manager2.get_token().await.unwrap(); + assert_eq!(token2.expose_secret(), "test_token_123"); + + // Verify file contents + let data: SessionData = + serde_json::from_str(&std::fs::read_to_string(&session_path).unwrap()).unwrap(); + assert_eq!(data.session_token, "test_token_123"); + assert_eq!(data.auth_provider, Some("near".to_string())); + } + + #[tokio::test] + async fn test_get_token_without_auth_fails() { + let dir = tempdir().unwrap(); + let config = SessionConfig { + auth_base_url: "https://example.com".to_string(), + session_path: dir.path().join("nonexistent.json"), + callback_port_range: (9900, 9910), + }; + + let manager = SessionManager::new_async(config).await; + let result = manager.get_token().await; + assert!(result.is_err()); + assert!(matches!(result, Err(LlmError::AuthFailed { .. }))); + } + + #[test] + fn test_default_session_path() { + let path = default_session_path(); + assert!(path.ends_with("session.json")); + assert!(path.to_string_lossy().contains(".near-agent")); + } +} diff --git a/src/main.rs b/src/main.rs index baddeaa0..28c71585 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,7 +11,7 @@ use near_agent::{ cli::{Cli, Command, run_tool_command}, config::Config, history::Store, - llm::create_llm_provider, + llm::{SessionConfig, create_llm_provider, create_session_manager}, safety::SafetyLayer, tools::{ ToolRegistry, @@ -79,8 +79,20 @@ 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)?; + let llm = create_llm_provider(&config.llm, session)?; tracing::info!("LLM provider initialized: {}", llm.model_name()); // Initialize safety layer