Fix snapshot recovery for component initializer host calls

golem-debugging-service/tests/debug_mode/debug_worker_executor.rs

@@ -92,6 +92,10 @@ impl DebugWorkerExecutorClient {
                             _ => {}
                         }
                     }
+                    Ok(Message::Close(frame)) => {
+                        anyhow::bail!("Debug connection closed: {frame:?}");
+                    }
+                    Err(error) => return Err(error.into()),
                     _ => {
                         if time.elapsed().as_secs() > 10 {
                             break Err(anyhow::anyhow!("Timeout")); // Break with an error

golem-worker-executor/src/worker/mod.rs

@@ -7428,7 +7428,7 @@ impl RunningWorker {
                 );
             }
         };
-        let context = match Ctx::create(
+        let mut context = match Ctx::create(
             worker_metadata.created_by,
             OwnedAgentId::new(worker_metadata.environment_id, &worker_metadata.agent_id),
             parent.parsed_agent_id.clone(),
@@ -7508,6 +7508,12 @@ impl RunningWorker {
                 );
             }
         };
+        if last_snapshot_index.is_some() {
+            // Core initializers run before load-snapshot, but their recorded host calls are
+            // already inside the skipped snapshot history. Recreate that runtime state with
+            // the same durability suppression as snapshot loading, without consuming the tail.
+            context.begin_call_snapshotting_function();
+        }
         let mut hosted = match instance_host.instantiate(context, &component).await {
             Ok(hosted) => hosted,
             Err(error) => {
@@ -7523,6 +7529,9 @@ impl RunningWorker {
             );
         }
         let (instance, mut store) = hosted.into_parts();
+        if last_snapshot_index.is_some() {
+            store.data_mut().end_call_snapshotting_function();
+        }
         if let Some((active_agent, generation)) = entity_generation {
             let interrupt_state = parent.interrupt_signal.lock().await;
             if !interrupt_state.has_interrupt() {

golem-worker-executor/tests/durability.rs

@@ -729,24 +729,55 @@ async fn automatic_snapshot_every_2nd_invocation(
         .start_agent(&component.id, agent_id.clone())
         .await?;
 
+    // Construction counts as an invocation; align subsequent snapshots with even increments.
+    let initial = executor
+        .invoke_and_await_agent(&component, &agent_id, "get", data_value!())
+        .await?;
+    assert_eq!(initial.into_typed::<u32>()?, 0);
+
     for _ in 0..SNAPSHOT_TEST_INVOCATIONS {
         executor
             .invoke_and_await_agent(&component, &agent_id, "increment", data_value!())
             .await?;
     }
 
     let oplog = executor.get_oplog(&worker_id, OplogIndex::INITIAL).await?;
+    assert!(
+        oplog
+            .iter()
+            .take_while(|entry| !matches!(&entry.entry, PublicOplogEntry::Snapshot(_)))
+            .any(|entry| matches!(
+                &entry.entry,
+                PublicOplogEntry::Start(params) if params.function_name == "monotonic_clock::now"
+            )),
+        "core initializer must record a clock call before the snapshot"
+    );
     let snapshot_count = oplog
         .iter()
         .filter(|entry| matches!(&entry.entry, PublicOplogEntry::Snapshot(_)))
         .count();
 
     assert_eq!(
         snapshot_count,
-        SNAPSHOT_TEST_INVOCATIONS / 2,
+        1 + SNAPSHOT_TEST_INVOCATIONS / 2,
         "Expected a snapshot every 2 invocations"
     );
 
+    let tail = executor
+        .invoke_and_await_agent(&component, &agent_id, "increment", data_value!())
+        .await?;
+    assert_eq!(tail.into_typed::<u32>()?, 11);
+    let oplog_with_tail = executor.get_oplog(&worker_id, OplogIndex::INITIAL).await?;
+    assert_eq!(
+        oplog_with_tail
+            .iter()
+            .rposition(|entry| matches!(&entry.entry, PublicOplogEntry::Snapshot(_))),
+        oplog
+            .iter()
+            .rposition(|entry| matches!(&entry.entry, PublicOplogEntry::Snapshot(_))),
+        "The final increment must remain outside the last snapshot"
+    );
+
     drop(executor);
     let executor = start_with_snapshot_policy(
         deps,
@@ -763,8 +794,8 @@ async fn automatic_snapshot_every_2nd_invocation(
 
     assert_eq!(
         result_after_restart.into_typed::<u32>()?,
-        SNAPSHOT_TEST_INVOCATIONS as u32,
-        "Counter should be restored from the automatic snapshot after restart"
+        11,
+        "Counter should include the increment replayed after the automatic snapshot"
     );
 
     drop(executor);

golem-worker-executor/tests/hot_update.rs

@@ -436,12 +436,22 @@ async fn snapshot_after_auto_update_recovers_with_updated_component_context(
         .await?;
     assert_eq!(before_snapshot.into_typed::<u32>()?, 0);
 
-    let snapshot_count = executor
-        .get_oplog(&worker_id, OplogIndex::INITIAL)
-        .await?
-        .iter()
-        .filter(|entry| matches!(&entry.entry, PublicOplogEntry::Snapshot(_)))
-        .count();
+    // Automatic snapshot creation is queued after the invocation result is published.
+    let snapshot_count = tokio::time::timeout(Duration::from_secs(30), async {
+        loop {
+            let count = executor
+                .get_oplog(&worker_id, OplogIndex::INITIAL)
+                .await?
+                .iter()
+                .filter(|entry| matches!(&entry.entry, PublicOplogEntry::Snapshot(_)))
+                .count();
+            if count > snapshots_before_invocation {
+                break Ok::<_, anyhow::Error>(count);
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await??;
     assert_eq!(snapshot_count, snapshots_before_invocation + 1);
 
     drop(executor);

golem-worker-executor/tests/observability.rs

@@ -216,7 +216,8 @@ async fn search_oplog_1(
     }
 
     assert_eq!(result1.len(), 2, "G1002"); // TODO: this is temporarily not working because of using the dynamic invoke API and not having structured information in the oplog
-    assert_eq!(result2.len(), 2, "imported-function");
+    // Includes the core initializer's monotonic clock call.
+    assert_eq!(result2.len(), 3, "imported-function");
     assert_eq!(result3.len(), 0, "id:G1001 OR id:G1000"); // TODO: this is temporarily not working because of using the dynamic invoke API and not having structured information in the oplog
 
     Ok(())

integration-tests/tests/worker.rs

@@ -1020,7 +1020,8 @@ async fn search_oplog_1(deps: &EnvBasedTestDependencies, _tracing: &Tracing) ->
     let result3 = user.search_oplog(&agent_id, "G1001 OR G1000").await?;
 
     assert_eq!(result1.len(), 2, "G1002"); // TODO: this is temporarily not working because of using the dynamic invoke API and not having structured information in the oplog
-    assert_eq!(result2.len(), 2, "imported-function");
+    // Includes the core initializer's monotonic clock call.
+    assert_eq!(result2.len(), 3, "imported-function");
     assert_eq!(result3.len(), 2, "id:G1001 OR id:G1000");
 
     Ok(())

test-components/agent-counters/src/snapshot_test.rs

@@ -1,6 +1,12 @@
 use golem_rust::{agent_definition, agent_implementation};
 use serde::{Deserialize, Serialize};
 
+#[unsafe(export_name = "_initialize")]
+pub extern "C" fn initialize_snapshot_clock() {
+    // The reactor initializer runs during core instantiation, before snapshot loading.
+    std::hint::black_box(std::time::Instant::now());
+}
+
 #[agent_definition(snapshotting = "enabled")]
 trait SnapshotCounter {
     fn new(id: String) -> Self;