From a2072532422f399ca0773b9fe90ac431db519a6d Mon Sep 17 00:00:00 2001
From: Shoubhit Dash <shoubhit2005@gmail.com>
Date: Wed, 19 Aug 2026 19:23:10 +0530
Subject: [PATCH] fix(core): await snapshot state readiness (#43435)

---
 packages/core/src/session.ts              | 22 +++++++++------
 packages/core/test/session-prompt.test.ts | 34 ++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts
index a7df3882500..b6cdd8800aa 100644
--- a/packages/core/src/session.ts
+++ b/packages/core/src/session.ts
@@ -905,19 +905,23 @@ const layer = Layer.effect(
           const session = yield* result.get(input.sessionID)
           if ((yield* execution.active).has(input.sessionID))
             return yield* new BusyError({ sessionID: input.sessionID })
-          return yield* SessionRevert.stage({ session, messageID: input.messageID, files: input.files }).pipe(
-            Effect.provideService(Database.Service, database),
-            Effect.provideService(Bus.Service, bus),
-            Effect.provide(locations.get(session.location)),
-          )
+          return yield* Effect.gen(function* () {
+            const plugins = yield* PluginSupervisor.Service
+            yield* plugins.flush
+            return yield* SessionRevert.stage({ session, messageID: input.messageID, files: input.files }).pipe(
+              Effect.provideService(Database.Service, database),
+              Effect.provideService(Bus.Service, bus),
+            )
+          }).pipe(Effect.provide(locations.get(session.location)))
         }),
         clear: Effect.fn("Session.revert.clear")(function* (sessionID) {
           const session = yield* result.get(sessionID)
           if ((yield* execution.active).has(sessionID)) return yield* new BusyError({ sessionID })
-          const revert = yield* SessionRevert.clear(session).pipe(
-            Effect.provideService(Bus.Service, bus),
-            Effect.provide(locations.get(session.location)),
-          )
+          const revert = yield* Effect.gen(function* () {
+            const plugins = yield* PluginSupervisor.Service
+            yield* plugins.flush
+            return yield* SessionRevert.clear(session).pipe(Effect.provideService(Bus.Service, bus))
+          }).pipe(Effect.provide(locations.get(session.location)))
           yield* execution.wake(sessionID)
           return revert
         }),
diff --git a/packages/core/test/session-prompt.test.ts b/packages/core/test/session-prompt.test.ts
index 8f4548eea76..b55b7ee5c51 100644
--- a/packages/core/test/session-prompt.test.ts
+++ b/packages/core/test/session-prompt.test.ts
@@ -28,6 +28,7 @@ import { LocationServiceMap } from "@opencode-ai/core/location-service-map"
 import type { LocationServices } from "@opencode-ai/core/location-services"
 import { Image } from "@opencode-ai/core/image"
 import { PluginSupervisor } from "@opencode-ai/core/plugin/supervisor"
+import { Snapshot } from "@opencode-ai/core/snapshot"
 import { testEffect } from "./lib/effect"
 
 const executionCalls: Session.ID[] = []
@@ -60,7 +61,7 @@ const locations = Layer.effect(
   LocationServiceMap.Service,
   LayerMap.make(
     () =>
-      // Attachment admission only needs image normalization and plugin readiness.
+      // These operations resolve Location services lazily and must wait for plugin-projected state.
       // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion
       Layer.unwrap(
         Effect.sync(() => {
@@ -72,6 +73,12 @@ const locations = Layer.effect(
                   ? Effect.succeed(content.content.length > 5 * 1024 * 1024 ? { ...content, content: "AA==" } : content)
                   : Effect.die(new Error("Image service used before plugins were ready")),
             }),
+            Layer.mock(Snapshot.Service, {
+              capture: () =>
+                ready ? Effect.succeed(undefined) : Effect.die(new Error("Snapshot used before plugins were ready")),
+              restore: () =>
+                ready ? Effect.void : Effect.die(new Error("Snapshot used before plugins were ready")),
+            }),
             Layer.succeed(
               PluginSupervisor.Service,
               PluginSupervisor.Service.of({ flush: Effect.sync(() => (ready = true)) }),
@@ -1051,6 +1058,31 @@ describe("Session.prompt", () => {
   )
 })
 
+describe("Session.revert", () => {
+  it.effect("waits for location plugins before staging", () =>
+    Effect.gen(function* () {
+      yield* setup
+      const { db } = yield* Database.Service
+      const session = yield* Session.Service
+      yield* db.insert(SessionMessageTable).values(assistantRow(messageID, 0)).run().pipe(Effect.orDie)
+      yield* session.revert.stage({ sessionID, messageID })
+    }),
+  )
+
+  it.effect("waits for location plugins before clearing", () =>
+    Effect.gen(function* () {
+      yield* setup
+      const session = yield* Session.Service
+      const bus = yield* Bus.Service
+      yield* bus.publish(SessionEvent.RevertEvent.Staged, {
+        sessionID,
+        revert: { messageID, snapshot: Snapshot.ID.make("tree"), files: [] },
+      })
+      yield* session.revert.clear(sessionID)
+    }),
+  )
+})
+
 describe("Session.inbox", () => {
   it.effect("fails for an unknown session", () =>
     Effect.gen(function* () {