Mnovich/temporal foreground tasks (#2895)

Co-authored-by: Carlos M. Lopez <carlopez@squareup.com>

crates/goose/src/session/mod.rs

@@ -3,9 +3,10 @@ pub mod storage;
 
 // Re-export common session types and functions
 pub use storage::{
-    ensure_session_dir, generate_description, generate_session_id, get_most_recent_session,
-    get_path, list_sessions, persist_messages, read_messages, read_metadata, update_metadata,
-    Identifier, SessionMetadata,
+    ensure_session_dir, generate_description, generate_description_with_schedule_id,
+    generate_session_id, get_most_recent_session, get_path, list_sessions, persist_messages,
+    persist_messages_with_schedule_id, read_messages, read_metadata, update_metadata, Identifier,
+    SessionMetadata,
 };
 
 pub use info::{get_session_info, SessionInfo};

crates/goose/src/session/storage.rs

@@ -437,6 +437,19 @@ pub async fn persist_messages(
     session_file: &Path,
     messages: &[Message],
     provider: Option<Arc<dyn Provider>>,
+) -> Result<()> {
+    persist_messages_with_schedule_id(session_file, messages, provider, None).await
+}
+
+/// Write messages to a session file with metadata, including an optional scheduled job ID
+///
+/// Overwrites the file with metadata as the first line, followed by all messages in JSONL format.
+/// If a provider is supplied, it will automatically generate a description when appropriate.
+pub async fn persist_messages_with_schedule_id(
+    session_file: &Path,
+    messages: &[Message],
+    provider: Option<Arc<dyn Provider>>,
+    schedule_id: Option<String>,
 ) -> Result<()> {
     // Count user messages
     let user_message_count = messages
@@ -448,11 +461,16 @@ pub async fn persist_messages(
     match provider {
         Some(provider) if user_message_count < 4 => {
             //generate_description is responsible for writing the messages
-            generate_description(session_file, messages, provider).await
+            generate_description_with_schedule_id(session_file, messages, provider, schedule_id)
+                .await
         }
         _ => {
             // Read existing metadata
-            let metadata = read_metadata(session_file)?;
+            let mut metadata = read_metadata(session_file)?;
+            // Update the schedule_id if provided
+            if schedule_id.is_some() {
+                metadata.schedule_id = schedule_id;
+            }
             // Write the file with metadata and messages
             save_messages_with_metadata(session_file, &metadata, messages)
         }
@@ -492,6 +510,19 @@ pub async fn generate_description(
     session_file: &Path,
     messages: &[Message],
     provider: Arc<dyn Provider>,
+) -> Result<()> {
+    generate_description_with_schedule_id(session_file, messages, provider, None).await
+}
+
+/// Generate a description for the session using the provider, including an optional scheduled job ID
+///
+/// This function is called when appropriate to generate a short description
+/// of the session based on the conversation history.
+pub async fn generate_description_with_schedule_id(
+    session_file: &Path,
+    messages: &[Message],
+    provider: Arc<dyn Provider>,
+    schedule_id: Option<String>,
 ) -> Result<()> {
     // Create a special message asking for a 3-word description
     let mut description_prompt = "Based on the conversation so far, provide a concise description of this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the description".to_string();
@@ -527,8 +558,11 @@ pub async fn generate_description(
     // Read current metadata
     let mut metadata = read_metadata(session_file)?;
 
-    // Update description
+    // Update description and schedule_id
     metadata.description = description;
+    if schedule_id.is_some() {
+        metadata.schedule_id = schedule_id;
+    }
 
     // Update the file with the new metadata and existing messages
     save_messages_with_metadata(session_file, &metadata, messages)