Update continue task system and daemon IDs

1 files changed, 28 insertions, 20 deletions

diff --git a/makima/src/server/handlers/mesh_daemon.rs b/makima/src/server/handlers/mesh_daemon.rs
index 9833d51..4bcb5cd 100644
--- a/makima/src/server/handlers/mesh_daemon.rs
+++ b/makima/src/server/handlers/mesh_daemon.rs
@@ -556,9 +556,9 @@ pub async fn daemon_handler(
 async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_result: DaemonAuthResult) {
     let (mut sender, mut receiver) = socket.split();
 
-    // Generate a unique connection ID and daemon ID
+    // Generate a unique connection ID (daemon ID will come from database)
     let connection_id = Uuid::new_v4().to_string();
-    let daemon_id = Uuid::new_v4();
+    let mut daemon_id: Option<Uuid> = None;
     let owner_id = auth_result.owner_id;
 
     // Create command channel for sending commands to this daemon
@@ -580,17 +580,8 @@ async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_re
                                 // API key was already validated via headers, but we use this message
                                 // for backward compatibility to get the machine_id and hostname
 
-                                tracing::info!(
-                                    daemon_id = %daemon_id,
-                                    owner_id = %owner_id,
-                                    hostname = %hostname,
-                                    machine_id = %machine_id,
-                                    max_concurrent_tasks = max_concurrent_tasks,
-                                    "Daemon registered"
-                                );
-
-                                // Register daemon in database
-                                if let Some(ref pool) = state.db_pool {
+                                // Register daemon in database first to get the persistent ID
+                                let db_daemon_id = if let Some(ref pool) = state.db_pool {
                                     match repository::register_daemon(
                                         pool,
                                         owner_id,
@@ -600,24 +591,40 @@ async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_re
                                         max_concurrent_tasks as i32,
                                     ).await {
                                         Ok(db_daemon) => {
-                                            tracing::debug!(
+                                            tracing::info!(
                                                 daemon_id = %db_daemon.id,
+                                                owner_id = %owner_id,
+                                                hostname = %hostname,
+                                                machine_id = %machine_id,
+                                                max_concurrent_tasks = max_concurrent_tasks,
                                                 "Daemon registered in database"
                                             );
+                                            Some(db_daemon.id)
                                         }
                                         Err(e) => {
                                             tracing::error!(
                                                 error = %e,
                                                 "Failed to register daemon in database"
                                             );
+                                            None
                                         }
                                     }
-                                }
+                                } else {
+                                    None
+                                };
+
+                                // If database registration failed, we can't proceed
+                                let Some(actual_daemon_id) = db_daemon_id else {
+                                    tracing::error!("Cannot proceed without database daemon ID");
+                                    break;
+                                };
+
+                                daemon_id = Some(actual_daemon_id);
 
-                                // Register daemon in state with owner_id
+                                // Register daemon in state with owner_id using database ID
                                 state.register_daemon(
                                     connection_id.clone(),
-                                    daemon_id,
+                                    actual_daemon_id,
                                     owner_id,
                                     Some(hostname),
                                     Some(machine_id),
@@ -626,8 +633,8 @@ async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_re
 
                                 registered = true;
 
-                                // Send authentication confirmation
-                                let response = DaemonCommand::Authenticated { daemon_id };
+                                // Send authentication confirmation with database ID
+                                let response = DaemonCommand::Authenticated { daemon_id: actual_daemon_id };
                                 let json = serde_json::to_string(&response).unwrap();
                                 if sender.send(Message::Text(json.into())).await.is_err() {
                                     break;
@@ -672,7 +679,8 @@ async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_re
         return;
     }
 
-    let daemon_uuid = daemon_id;
+    // daemon_id is guaranteed to be Some here since registered is true
+    let daemon_uuid = daemon_id.expect("daemon_id should be set when registered is true");
 
     // Main message loop after authentication
     loop {