//! `x.ai/share_session` extension handler. //! //! Loads a local session, exports it, and asks the backend for a public //! share URL. use agent_client_protocol as acp; use super::{ExtResult, parse_params, to_raw_response}; use crate::agent::MvpAgent; use crate::remote::client::BackendClient; use crate::session::export::ExportedSession; use crate::session::info::Info as SessionInfo; use crate::session::persistence::list_summaries; use crate::session::share::{ShareSessionRequest, ShareSessionResponse}; #[tracing::instrument(skip_all, fields(method = %args.method))] pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { match args.method.as_ref() { "x.ai/share_session" => { tracing::info!("handling share session request"); handle_share_session(agent, args).await } _ => Err(acp::Error::method_not_found()), } } async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { let request: ShareSessionRequest = parse_params(args)?; // Get auth - required for sharing. let auth = require_xai_auth_for_share(&agent.auth_manager)?; // Remote settings / feature-flag gate: sharing_enabled defaults to false // and is only enabled for eligible accounts. let sharing_enabled = agent .cfg .borrow() .remote_settings .as_ref() .and_then(|rs| rs.sharing_enabled) .unwrap_or(false); if !sharing_enabled { return Err( acp::Error::invalid_params().data("Session sharing is not available for your account.") ); } // Find session info by searching through summaries let summaries = list_summaries(None).await.map_err(|e| { acp::Error::internal_error().data(format!("Failed to list sessions: {}", e)) })?; let summary = summaries .iter() .find(|s| s.info.id.0.as_ref() == request.session_id.as_str()) .ok_or_else(|| acp::Error::resource_not_found(Some("Session not found".into())))?; let info = SessionInfo { id: acp::SessionId::new(request.session_id.clone()), cwd: summary.info.cwd.clone(), }; // Load and export session let exported = ExportedSession::from_local_session(&info) .await .map_err(|e| acp::Error::internal_error().data(format!("Failed to load session: {}", e)))?; // Check for empty session if exported.messages.is_empty() { return Err(acp::Error::invalid_params().data("No messages to share yet")); } // Upload to backend and get share URL. let client = BackendClient::new().with_auth_manager(agent.auth_manager.clone()); let agent_id = crate::util::agent_id::agent_id(); let share_url = client .share_session(&exported, &agent_id) .await .map_err(|e| { tracing::error!(error = %e, "Failed to share session with backend"); acp::Error::internal_error().data(format!("Failed to share session: {}", e)) })?; let response = ShareSessionResponse { share_url }; to_raw_response(&response) } fn require_xai_auth_for_share( auth_manager: &crate::auth::AuthManager, ) -> Result { super::auth_gate::require_xai_auth( auth_manager, "Authentication required to share session", "Share session is disabled. Run `grok login` to authenticate.", ) } #[cfg(test)] mod tests { use super::*; use crate::auth::KimiCodeConfig; use crate::auth::{AuthMode, KimiAuth}; use chrono::{Duration, Utc}; use std::sync::Arc; use tempfile::tempdir; fn make_auth_manager_with_token_expiring_in( ttl: Duration, ) -> (Arc, tempfile::TempDir) { let dir = tempdir().expect("tempdir for share auth test"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), KimiCodeConfig::default(), )); let expires_at = Utc::now() + ttl; // We must explicitly set oidc_issuer to a first-party xAI issuer. // Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent) // return true from is_xai_auth(). This is required for the share tests to // exercise the happy path through require_xai_auth_for_share. let auth = KimiAuth { auth_mode: AuthMode::OAuth, key: "test-key".into(), expires_at: Some(expires_at), create_time: Utc::now() - Duration::hours(1), ..Default::default() }; mgr.hot_swap(auth); (mgr, dir) } #[test] fn share_works_outside_the_5m_early_invalidation_window() { let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::minutes(10)); assert!(mgr.current().is_some()); assert!(require_xai_auth_for_share(&mgr).is_ok()); } #[test] fn share_succeeds_inside_the_5m_early_invalidation_window() { let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::seconds(1)); // This is exactly the state that triggered the user bug: assert!( mgr.current().is_none(), "current() drops the token inside the buffer" ); assert!(mgr.expired_auth().is_some()); // Now that we use current_or_expired(), this passes. let res = require_xai_auth_for_share(&mgr); assert!( res.is_ok(), "require_xai_auth_for_share must succeed for a still-valid buffered xAI token" ); } #[test] fn share_fails_with_no_auth_at_all() { let dir = tempdir().expect("tempdir"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), KimiCodeConfig::default(), )); assert!(require_xai_auth_for_share(&mgr).is_err()); } #[test] fn share_rejects_non_xai_auth_with_actionable_grok_login_message() { let dir = tempdir().expect("tempdir"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), KimiCodeConfig::default(), )); // API key is the simplest non-xAI credential (External and enterprise OIDC // are also rejected the same way). let non_xai = KimiAuth { auth_mode: AuthMode::ApiKey, key: "xai-test-key".into(), create_time: Utc::now(), ..Default::default() }; mgr.hot_swap(non_xai); let err = require_xai_auth_for_share(&mgr) .expect_err("non-xAI accounts (API key, External, enterprise IdP) must be rejected"); // This is the key assertion the review asked for: we must test the *exact* // actionable data string for the non-xAI path (distinct from the generic // "Authentication required to share session" path). let serialized = serde_json::to_value(&err).expect("acp::Error serializes to JSON-RPC shape"); let data = serialized .get("data") .and_then(|v| v.as_str()) .expect("auth_required error carries a data string"); assert_eq!( data, "Share session is disabled. Run `grok login` to authenticate." ); } }