Login flow

2026-08-25 14:53:34 +00:00 · 2026-02-03 09:20:26 -08:00
parent dedda9c51d
commit 2df4a4f5f0
9 changed files with 737 additions and 20 deletions
@@ -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
@@ -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"
@@ -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"] }
@@ -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<Self, ConfigError> {
-        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 {
@@ -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.
@@ -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<Arc<dyn LlmProvider>, 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<SessionManager>,
+) -> Result<Arc<dyn LlmProvider>, LlmError> {
+    Ok(Arc::new(NearAiProvider::new(
+        config.nearai.clone(),
+        session,
+    )))
 }
@@ -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<SessionManager>,
 }

 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<SessionManager>) -> 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<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
        &self,
        path: &str,
        body: &T,
+    ) -> Result<R, LlmError> {
+        // 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<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
+        &self,
+        path: &str,
+        body: &T,
    ) -> Result<R, LlmError> {
        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::<NearAiErrorResponse>(&response_text) {
                if status.as_u16() == 429 {
@@ -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<Utc>,
+    #[serde(default)]
+    pub auth_provider: Option<String>,
+}
+
+/// 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<Option<SecretString>>,
+    /// 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::<SessionData>(&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<SecretString, LlmError> {
+        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<SecretString, LlmError> {
+        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",
+                                    "<!DOCTYPE html><html><head><title>NEAR AI Auth</title></head>",
+                                    "<body style=\"font-family: sans-serif; text-align: center; padding-top: 50px;\">",
+                                    "<h1>✓ Authentication successful!</h1>",
+                                    "<p>You can close this window and return to the terminal.</p>",
+                                    "</body></html>"
+                                );
+
+                                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<SessionManager> {
+    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"));
+    }
+}
@@ -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