Merge pull request #486 from stakpak/perf/async-session-title-generation

perf(session): async session title generation for faster startup

Author: ahmedhesham6

cli/src/commands/agent/run/mode_interactive.rs

@@ -941,6 +941,16 @@ pub async fn run_interactive(
                     Ok(response) => {
                         messages.push(response.choices[0].message.clone());
 
+                        if let Some(session_id) = response
+                            .metadata
+                            .as_ref()
+                            .and_then(|meta| meta.get("session_id"))
+                            .and_then(|value| value.as_str())
+                            .and_then(|value| Uuid::parse_str(value).ok())
+                        {
+                            current_session_id = Some(session_id);
+                        }
+
                         // Accumulate usage from response
                         total_session_usage.prompt_tokens += response.usage.prompt_tokens;
                         total_session_usage.completion_tokens += response.usage.completion_tokens;

cli/src/commands/agent/run/stream.rs

@@ -122,6 +122,7 @@ pub async fn process_responses_stream(
         system_fingerprint: None,
         metadata: None,
     };
+    let mut response_metadata: Option<serde_json::Value> = None;
 
     let mut llm_model: Option<LLMModel> = None;
 
@@ -154,6 +155,9 @@ pub async fn process_responses_stream(
                     // Send usage to TUI for display immediately when we receive it
                     send_input_event(input_tx, InputEvent::StreamUsage(usage.clone())).await?;
                 }
+                if let Some(metadata) = &response.metadata {
+                    response_metadata = Some(metadata.clone());
+                }
 
                 // Skip chunks with no choices (e.g., usage-only events)
                 if response.choices.is_empty() {
@@ -239,6 +243,7 @@ pub async fn process_responses_stream(
         finish_reason: FinishReason::Stop,
         logprobs: None,
     });
+    chat_completion_response.metadata = response_metadata;
 
     // End stream processing loading when stream completes
     send_input_event(

libs/api/src/client/provider.rs

@@ -10,6 +10,7 @@ use crate::models::*;
 use crate::storage::{
     CreateCheckpointRequest as StorageCreateCheckpointRequest,
     CreateSessionRequest as StorageCreateSessionRequest,
+    UpdateSessionRequest as StorageUpdateSessionRequest,
 };
 use async_trait::async_trait;
 use futures_util::Stream;
@@ -664,34 +665,17 @@ impl AgentClient {
             });
         }
 
-        // Create new session
-        // Generate title with fallback - don't fail session creation if title generation fails
-        let title = match self.generate_session_title(messages).await {
-            Ok(title) => title,
-            Err(_) => {
-                // Extract first few words from user message as fallback title
-                messages
-                    .iter()
-                    .find(|m| m.role == Role::User)
-                    .and_then(|m| m.content.as_ref())
-                    .map(|c| {
-                        let text = c.to_string();
-                        text.split_whitespace()
-                            .take(5)
-                            .collect::<Vec<_>>()
-                            .join(" ")
-                    })
-                    .unwrap_or_else(|| "New Session".to_string())
-            }
-        };
+        // Create new session with a fast local title.
+        let fallback_title = Self::fallback_session_title(messages);
 
         // Get current working directory
         let cwd = std::env::current_dir()
             .ok()
             .map(|p| p.to_string_lossy().to_string());
 
         // Create session via storage trait
-        let mut session_request = StorageCreateSessionRequest::new(title, messages.to_vec());
+        let mut session_request =
+            StorageCreateSessionRequest::new(fallback_title.clone(), messages.to_vec());
         if let Some(cwd) = cwd {
             session_request = session_request.with_cwd(cwd);
         }
@@ -702,13 +686,46 @@ impl AgentClient {
             .await
             .map_err(|e| e.to_string())?;
 
+        // Generate a better title asynchronously and update the session when ready.
+        let client = self.clone();
+        let messages_for_title = messages.to_vec();
+        let session_id = result.session_id;
+        tokio::spawn(async move {
+            if let Ok(title) = client.generate_session_title(&messages_for_title).await {
+                let trimmed = title.trim();
+                if !trimmed.is_empty() && trimmed != fallback_title {
+                    let request =
+                        StorageUpdateSessionRequest::new().with_title(trimmed.to_string());
+                    let _ = client
+                        .session_storage
+                        .update_session(session_id, &request)
+                        .await;
+                }
+            }
+        });
+
         Ok(SessionInfo {
             session_id: result.session_id,
             checkpoint_id: result.checkpoint.id,
             checkpoint_created_at: result.checkpoint.created_at,
         })
     }
 
+    fn fallback_session_title(messages: &[ChatMessage]) -> String {
+        messages
+            .iter()
+            .find(|m| m.role == Role::User)
+            .and_then(|m| m.content.as_ref())
+            .map(|c| {
+                let text = c.to_string();
+                text.split_whitespace()
+                    .take(5)
+                    .collect::<Vec<_>>()
+                    .join(" ")
+            })
+            .unwrap_or_else(|| "New Session".to_string())
+    }
+
     /// Save a new checkpoint for the current session
     pub(crate) async fn save_checkpoint(
         &self,