feat: persist standalone sessions for restore

crates/worker/src/bootstrap.rs

@@ -130,6 +130,29 @@ where
         Ok(PreparedWorker::new(worker, self.layout, self.transport))
     }
 
+    pub async fn prepare_restored(
+        self,
+        worker_name: &str,
+    ) -> Result<PreparedWorker<Box<dyn LlmClient>, St>, WorkerBootstrapError> {
+        let mut worker =
+            Worker::restore_pending_from_worker_metadata_with_context_and_model_client(
+                worker_name,
+                self.manifest,
+                self.store,
+                self.prompt_catalog,
+                self.workspace_context,
+                self.filesystem_authority,
+                self.model_client,
+            )
+            .await
+            .map_err(WorkerBootstrapError::Worker)?;
+
+        if let Some(workdir_session) = self.workdir_session {
+            worker.bind_workdir_session(Some(workdir_session));
+        }
+        Ok(PreparedWorker::new(worker, self.layout, self.transport))
+    }
+
     pub async fn start(self) -> Result<BootstrappedWorker, WorkerBootstrapError> {
         self.prepare().await?.start().await
     }

crates/worker/src/worker.rs

@@ -5515,6 +5515,27 @@ where
         loader: PromptCatalogSource,
         workspace_context: WorkerWorkspaceContext,
         filesystem_authority: WorkerFilesystemAuthority,
+    ) -> Result<Self, WorkerError> {
+        Self::restore_pending_from_worker_metadata_with_context_and_model_client(
+            worker_name,
+            fallback,
+            store,
+            loader,
+            workspace_context,
+            filesystem_authority,
+            None,
+        )
+        .await
+    }
+
+    pub(crate) async fn restore_pending_from_worker_metadata_with_context_and_model_client(
+        worker_name: &str,
+        fallback: WorkerManifest,
+        store: St,
+        loader: PromptCatalogSource,
+        workspace_context: WorkerWorkspaceContext,
+        filesystem_authority: WorkerFilesystemAuthority,
+        model_client: Option<Box<dyn LlmClient>>,
     ) -> Result<Self, WorkerError> {
         let metadata =
             store
@@ -5535,7 +5556,7 @@ where
                 worker_name: worker_name.to_string(),
             })?;
         if let Some(segment_id) = active.segment_id {
-            return Self::restore_from_manifest_with_context(
+            return Self::restore_from_manifest_with_context_and_model_client(
                 active.session_id,
                 segment_id,
                 restore_manifest_from_worker_metadata_snapshot(
@@ -5547,6 +5568,7 @@ where
                 loader,
                 workspace_context,
                 filesystem_authority,
+                model_client,
             )
             .await;
         }
@@ -5557,12 +5579,13 @@ where
         })?;
         let manifest =
             restore_manifest_from_worker_metadata_snapshot(worker_name, Some(snapshot), fallback)?;
-        Self::from_manifest_with_context(
+        Self::from_manifest_with_context_and_model_client(
             manifest,
             store,
             loader,
             workspace_context,
             filesystem_authority,
+            model_client,
         )
         .await
     }
@@ -5614,6 +5637,29 @@ where
         loader: PromptCatalogSource,
         workspace_context: WorkerWorkspaceContext,
         filesystem_authority: WorkerFilesystemAuthority,
+    ) -> Result<Self, WorkerError> {
+        Self::restore_from_manifest_with_context_and_model_client(
+            session_id,
+            segment_id,
+            manifest,
+            store,
+            loader,
+            workspace_context,
+            filesystem_authority,
+            None,
+        )
+        .await
+    }
+
+    pub(crate) async fn restore_from_manifest_with_context_and_model_client(
+        session_id: SessionId,
+        segment_id: SegmentId,
+        manifest: WorkerManifest,
+        store: St,
+        loader: PromptCatalogSource,
+        workspace_context: WorkerWorkspaceContext,
+        filesystem_authority: WorkerFilesystemAuthority,
+        model_client: Option<Box<dyn LlmClient>>,
     ) -> Result<Self, WorkerError> {
         // Read raw entries once so we can both reconstruct state and
         // seed the broadcast sink's mirror with the same prefix that
@@ -5629,13 +5675,14 @@ where
         let mirror_entries: Vec<LogEntry> = raw_entries.clone();
         let scope_config = effective_restore_scope_config(&store, &manifest)?;
 
-        let common = prepare_worker_common_with_context(
+        let common = prepare_worker_common_with_context_and_model_client(
             &manifest,
             &loader,
             /* parse_template */ false,
             workspace_context,
             filesystem_authority,
             scope_config,
+            model_client,
         )?;
 
         // Atomic: register_worker inside install_top_level rejects when