2
0
Эх сурвалжийг харах

refactor(session): remove revert async facade exports (#22339)

Kit Langton 3 өдөр өмнө
parent
commit
67aaecacac

+ 41 - 27
packages/opencode/src/server/instance/session.ts

@@ -551,28 +551,38 @@ export const SessionRoutes = lazy(() =>
       async (c) => {
         const sessionID = c.req.valid("param").sessionID
         const body = c.req.valid("json")
-        const session = await Session.get(sessionID)
-        await SessionRevert.cleanup(session)
-        const msgs = await Session.messages({ sessionID })
-        const defaultAgent = await AppRuntime.runPromise(Agent.Service.use((svc) => svc.defaultAgent()))
-        let currentAgent = defaultAgent
-        for (let i = msgs.length - 1; i >= 0; i--) {
-          const info = msgs[i].info
-          if (info.role === "user") {
-            currentAgent = info.agent || defaultAgent
-            break
-          }
-        }
-        await SessionCompaction.create({
-          sessionID,
-          agent: currentAgent,
-          model: {
-            providerID: body.providerID,
-            modelID: body.modelID,
-          },
-          auto: body.auto,
-        })
-        await SessionPrompt.loop({ sessionID })
+        await AppRuntime.runPromise(
+          Effect.gen(function* () {
+            const session = yield* Session.Service
+            const revert = yield* SessionRevert.Service
+            const compact = yield* SessionCompaction.Service
+            const prompt = yield* SessionPrompt.Service
+            const agent = yield* Agent.Service
+
+            yield* revert.cleanup(yield* session.get(sessionID))
+            const msgs = yield* session.messages({ sessionID })
+            const defaultAgent = yield* agent.defaultAgent()
+            let currentAgent = defaultAgent
+            for (let i = msgs.length - 1; i >= 0; i--) {
+              const info = msgs[i].info
+              if (info.role === "user") {
+                currentAgent = info.agent || defaultAgent
+                break
+              }
+            }
+
+            yield* compact.create({
+              sessionID,
+              agent: currentAgent,
+              model: {
+                providerID: body.providerID,
+                modelID: body.modelID,
+              },
+              auto: body.auto,
+            })
+            yield* prompt.loop({ sessionID })
+          }),
+        )
         return c.json(true)
       },
     )
@@ -990,10 +1000,14 @@ export const SessionRoutes = lazy(() =>
       async (c) => {
         const sessionID = c.req.valid("param").sessionID
         log.info("revert", c.req.valid("json"))
-        const session = await SessionRevert.revert({
-          sessionID,
-          ...c.req.valid("json"),
-        })
+        const session = await AppRuntime.runPromise(
+          SessionRevert.Service.use((svc) =>
+            svc.revert({
+              sessionID,
+              ...c.req.valid("json"),
+            }),
+          ),
+        )
         return c.json(session)
       },
     )
@@ -1023,7 +1037,7 @@ export const SessionRoutes = lazy(() =>
       ),
       async (c) => {
         const sessionID = c.req.valid("param").sessionID
-        const session = await SessionRevert.unrevert({ sessionID })
+        const session = await AppRuntime.runPromise(SessionRevert.Service.use((svc) => svc.unrevert({ sessionID })))
         return c.json(session)
       },
     )

+ 0 - 15
packages/opencode/src/session/revert.ts

@@ -1,6 +1,5 @@
 import z from "zod"
 import { Effect, Layer, Context } from "effect"
-import { makeRuntime } from "@/effect/run-service"
 import { Bus } from "../bus"
 import { Snapshot } from "../snapshot"
 import { Storage } from "@/storage/storage"
@@ -160,18 +159,4 @@ export namespace SessionRevert {
       Layer.provide(SessionSummary.defaultLayer),
     ),
   )
-
-  const { runPromise } = makeRuntime(Service, defaultLayer)
-
-  export async function revert(input: RevertInput) {
-    return runPromise((svc) => svc.revert(input))
-  }
-
-  export async function unrevert(input: { sessionID: SessionID }) {
-    return runPromise((svc) => svc.unrevert(input))
-  }
-
-  export async function cleanup(session: Session.Info) {
-    return runPromise((svc) => svc.cleanup(session))
-  }
 }

+ 558 - 540
packages/opencode/test/session/revert-compact.test.ts

@@ -1,35 +1,47 @@
-import { describe, expect, test, beforeEach, afterEach } from "bun:test"
+import { describe, expect } from "bun:test"
 import fs from "fs/promises"
 import path from "path"
+import { Effect, Layer } from "effect"
 import { Session } from "../../src/session"
 import { ModelID, ProviderID } from "../../src/provider/schema"
 import { SessionRevert } from "../../src/session/revert"
-import { SessionCompaction } from "../../src/session/compaction"
 import { MessageV2 } from "../../src/session/message-v2"
 import { Snapshot } from "../../src/snapshot"
 import { Log } from "../../src/util/log"
-import { Instance } from "../../src/project/instance"
-import { MessageID, PartID } from "../../src/session/schema"
-import { tmpdir } from "../fixture/fixture"
+import { MessageID, PartID, SessionID } from "../../src/session/schema"
+import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
+import { provideTmpdirInstance } from "../fixture/fixture"
+import { testEffect } from "../lib/effect"
 
 Log.init({ print: false })
 
-function user(sessionID: string, agent = "default") {
-  return Session.updateMessage({
+const env = Layer.mergeAll(
+  Session.defaultLayer,
+  SessionRevert.defaultLayer,
+  Snapshot.defaultLayer,
+  CrossSpawnSpawner.defaultLayer,
+)
+
+const it = testEffect(env)
+
+const user = Effect.fn("test.user")(function* (sessionID: SessionID, agent = "default") {
+  const session = yield* Session.Service
+  return yield* session.updateMessage({
     id: MessageID.ascending(),
     role: "user" as const,
-    sessionID: sessionID as any,
+    sessionID,
     agent,
     model: { providerID: ProviderID.make("openai"), modelID: ModelID.make("gpt-4") },
     time: { created: Date.now() },
   })
-}
+})
 
-function assistant(sessionID: string, parentID: string, dir: string) {
-  return Session.updateMessage({
+const assistant = Effect.fn("test.assistant")(function* (sessionID: SessionID, parentID: MessageID, dir: string) {
+  const session = yield* Session.Service
+  return yield* session.updateMessage({
     id: MessageID.ascending(),
     role: "assistant" as const,
-    sessionID: sessionID as any,
+    sessionID,
     mode: "default",
     agent: "default",
     path: { cwd: dir, root: dir },
@@ -37,27 +49,29 @@ function assistant(sessionID: string, parentID: string, dir: string) {
     tokens: { output: 0, input: 0, reasoning: 0, cache: { read: 0, write: 0 } },
     modelID: ModelID.make("gpt-4"),
     providerID: ProviderID.make("openai"),
-    parentID: parentID as any,
+    parentID,
     time: { created: Date.now() },
     finish: "end_turn",
   })
-}
+})
 
-function text(sessionID: string, messageID: string, content: string) {
-  return Session.updatePart({
+const text = Effect.fn("test.text")(function* (sessionID: SessionID, messageID: MessageID, content: string) {
+  const session = yield* Session.Service
+  return yield* session.updatePart({
     id: PartID.ascending(),
-    messageID: messageID as any,
-    sessionID: sessionID as any,
+    messageID,
+    sessionID,
     type: "text" as const,
     text: content,
   })
-}
+})
 
-function tool(sessionID: string, messageID: string) {
-  return Session.updatePart({
+const tool = Effect.fn("test.tool")(function* (sessionID: SessionID, messageID: MessageID) {
+  const session = yield* Session.Service
+  return yield* session.updatePart({
     id: PartID.ascending(),
-    messageID: messageID as any,
-    sessionID: sessionID as any,
+    messageID,
+    sessionID,
     type: "tool" as const,
     tool: "bash",
     callID: "call-1",
@@ -70,7 +84,10 @@ function tool(sessionID: string, messageID: string) {
       time: { start: 0, end: 1 },
     },
   })
-}
+})
+
+const read = (file: string) => Effect.promise(() => fs.readFile(file, "utf-8"))
+const write = (file: string, text: string) => Effect.promise(() => fs.writeFile(file, text))
 
 const tokens = {
   input: 0,
@@ -80,542 +97,543 @@ const tokens = {
 }
 
 describe("revert + compact workflow", () => {
-  test("should properly handle compact command after revert", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        // Create a session
-        const session = await Session.create({})
-        const sessionID = session.id
-
-        // Create a user message
-        const userMsg1 = await Session.updateMessage({
-          id: MessageID.ascending(),
-          role: "user",
-          sessionID,
-          agent: "default",
-          model: {
-            providerID: ProviderID.make("openai"),
+  it.live(
+    "should properly handle compact command after revert",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+
+          const info = yield* session.create({})
+          const sessionID = info.id
+
+          const userMsg1 = yield* session.updateMessage({
+            id: MessageID.ascending(),
+            role: "user",
+            sessionID,
+            agent: "default",
+            model: {
+              providerID: ProviderID.make("openai"),
+              modelID: ModelID.make("gpt-4"),
+            },
+            time: {
+              created: Date.now(),
+            },
+          })
+
+          yield* session.updatePart({
+            id: PartID.ascending(),
+            messageID: userMsg1.id,
+            sessionID,
+            type: "text",
+            text: "Hello, please help me",
+          })
+
+          const assistantMsg1: MessageV2.Assistant = {
+            id: MessageID.ascending(),
+            role: "assistant",
+            sessionID,
+            mode: "default",
+            agent: "default",
+            path: {
+              cwd: dir,
+              root: dir,
+            },
+            cost: 0,
+            tokens: {
+              output: 0,
+              input: 0,
+              reasoning: 0,
+              cache: { read: 0, write: 0 },
+            },
             modelID: ModelID.make("gpt-4"),
-          },
-          time: {
-            created: Date.now(),
-          },
-        })
-
-        // Add a text part to the user message
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: userMsg1.id,
-          sessionID,
-          type: "text",
-          text: "Hello, please help me",
-        })
-
-        // Create an assistant response message
-        const assistantMsg1: MessageV2.Assistant = {
-          id: MessageID.ascending(),
-          role: "assistant",
-          sessionID,
-          mode: "default",
-          agent: "default",
-          path: {
-            cwd: tmp.path,
-            root: tmp.path,
-          },
-          cost: 0,
-          tokens: {
-            output: 0,
-            input: 0,
-            reasoning: 0,
-            cache: { read: 0, write: 0 },
-          },
-          modelID: ModelID.make("gpt-4"),
-          providerID: ProviderID.make("openai"),
-          parentID: userMsg1.id,
-          time: {
-            created: Date.now(),
-          },
-          finish: "end_turn",
-        }
-        await Session.updateMessage(assistantMsg1)
-
-        // Add a text part to the assistant message
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: assistantMsg1.id,
-          sessionID,
-          type: "text",
-          text: "Sure, I'll help you!",
-        })
-
-        // Create another user message
-        const userMsg2 = await Session.updateMessage({
-          id: MessageID.ascending(),
-          role: "user",
-          sessionID,
-          agent: "default",
-          model: {
             providerID: ProviderID.make("openai"),
-            modelID: ModelID.make("gpt-4"),
-          },
-          time: {
-            created: Date.now(),
-          },
-        })
-
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: userMsg2.id,
-          sessionID,
-          type: "text",
-          text: "What's the capital of France?",
-        })
-
-        // Create another assistant response
-        const assistantMsg2: MessageV2.Assistant = {
-          id: MessageID.ascending(),
-          role: "assistant",
-          sessionID,
-          mode: "default",
-          agent: "default",
-          path: {
-            cwd: tmp.path,
-            root: tmp.path,
-          },
-          cost: 0,
-          tokens: {
-            output: 0,
-            input: 0,
-            reasoning: 0,
-            cache: { read: 0, write: 0 },
-          },
-          modelID: ModelID.make("gpt-4"),
-          providerID: ProviderID.make("openai"),
-          parentID: userMsg2.id,
-          time: {
-            created: Date.now(),
-          },
-          finish: "end_turn",
-        }
-        await Session.updateMessage(assistantMsg2)
-
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: assistantMsg2.id,
-          sessionID,
-          type: "text",
-          text: "The capital of France is Paris.",
-        })
-
-        // Verify messages before revert
-        let messages = await Session.messages({ sessionID })
-        expect(messages.length).toBe(4) // 2 user + 2 assistant messages
-        const messageIds = messages.map((m) => m.info.id)
-        expect(messageIds).toContain(userMsg1.id)
-        expect(messageIds).toContain(userMsg2.id)
-        expect(messageIds).toContain(assistantMsg1.id)
-        expect(messageIds).toContain(assistantMsg2.id)
-
-        // Revert the last user message (userMsg2)
-        await SessionRevert.revert({
-          sessionID,
-          messageID: userMsg2.id,
-        })
-
-        // Check that revert state is set
-        let sessionInfo = await Session.get(sessionID)
-        expect(sessionInfo.revert).toBeDefined()
-        const revertMessageID = sessionInfo.revert?.messageID
-        expect(revertMessageID).toBeDefined()
-
-        // Messages should still be in the list (not removed yet, just marked for revert)
-        messages = await Session.messages({ sessionID })
-        expect(messages.length).toBe(4)
-
-        // Now clean up the revert state (this is what the compact endpoint should do)
-        await SessionRevert.cleanup(sessionInfo)
-
-        // After cleanup, the reverted messages (those after the revert point) should be removed
-        messages = await Session.messages({ sessionID })
-        const remainingIds = messages.map((m) => m.info.id)
-        // The revert point is somewhere in the message chain, so we should have fewer messages
-        expect(messages.length).toBeLessThan(4)
-        // userMsg2 and assistantMsg2 should be removed (they come after the revert point)
-        expect(remainingIds).not.toContain(userMsg2.id)
-        expect(remainingIds).not.toContain(assistantMsg2.id)
-
-        // Revert state should be cleared
-        sessionInfo = await Session.get(sessionID)
-        expect(sessionInfo.revert).toBeUndefined()
-
-        // Clean up
-        await Session.remove(sessionID)
-      },
-    })
-  })
+            parentID: userMsg1.id,
+            time: {
+              created: Date.now(),
+            },
+            finish: "end_turn",
+          }
+          yield* session.updateMessage(assistantMsg1)
+
+          yield* session.updatePart({
+            id: PartID.ascending(),
+            messageID: assistantMsg1.id,
+            sessionID,
+            type: "text",
+            text: "Sure, I'll help you!",
+          })
 
-  test("should properly clean up revert state before creating compaction message", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        // Create a session
-        const session = await Session.create({})
-        const sessionID = session.id
-
-        // Create initial messages
-        const userMsg = await Session.updateMessage({
-          id: MessageID.ascending(),
-          role: "user",
-          sessionID,
-          agent: "default",
-          model: {
-            providerID: ProviderID.make("openai"),
+          const userMsg2 = yield* session.updateMessage({
+            id: MessageID.ascending(),
+            role: "user",
+            sessionID,
+            agent: "default",
+            model: {
+              providerID: ProviderID.make("openai"),
+              modelID: ModelID.make("gpt-4"),
+            },
+            time: {
+              created: Date.now(),
+            },
+          })
+
+          yield* session.updatePart({
+            id: PartID.ascending(),
+            messageID: userMsg2.id,
+            sessionID,
+            type: "text",
+            text: "What's the capital of France?",
+          })
+
+          const assistantMsg2: MessageV2.Assistant = {
+            id: MessageID.ascending(),
+            role: "assistant",
+            sessionID,
+            mode: "default",
+            agent: "default",
+            path: {
+              cwd: dir,
+              root: dir,
+            },
+            cost: 0,
+            tokens: {
+              output: 0,
+              input: 0,
+              reasoning: 0,
+              cache: { read: 0, write: 0 },
+            },
             modelID: ModelID.make("gpt-4"),
-          },
-          time: {
-            created: Date.now(),
-          },
-        })
-
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: userMsg.id,
-          sessionID,
-          type: "text",
-          text: "Hello",
-        })
-
-        const assistantMsg: MessageV2.Assistant = {
-          id: MessageID.ascending(),
-          role: "assistant",
-          sessionID,
-          mode: "default",
-          agent: "default",
-          path: {
-            cwd: tmp.path,
-            root: tmp.path,
-          },
-          cost: 0,
-          tokens: {
-            output: 0,
-            input: 0,
-            reasoning: 0,
-            cache: { read: 0, write: 0 },
-          },
-          modelID: ModelID.make("gpt-4"),
-          providerID: ProviderID.make("openai"),
-          parentID: userMsg.id,
-          time: {
-            created: Date.now(),
-          },
-          finish: "end_turn",
-        }
-        await Session.updateMessage(assistantMsg)
-
-        await Session.updatePart({
-          id: PartID.ascending(),
-          messageID: assistantMsg.id,
-          sessionID,
-          type: "text",
-          text: "Hi there!",
-        })
-
-        // Revert the user message
-        await SessionRevert.revert({
-          sessionID,
-          messageID: userMsg.id,
-        })
-
-        // Check that revert state is set
-        let sessionInfo = await Session.get(sessionID)
-        expect(sessionInfo.revert).toBeDefined()
-
-        // Simulate what the compact endpoint does: cleanup revert before creating compaction
-        await SessionRevert.cleanup(sessionInfo)
-
-        // Verify revert state is cleared
-        sessionInfo = await Session.get(sessionID)
-        expect(sessionInfo.revert).toBeUndefined()
-
-        // Verify messages are properly cleaned up
-        const messages = await Session.messages({ sessionID })
-        expect(messages.length).toBe(0) // All messages should be reverted
-
-        // Clean up
-        await Session.remove(sessionID)
-      },
-    })
-  })
+            providerID: ProviderID.make("openai"),
+            parentID: userMsg2.id,
+            time: {
+              created: Date.now(),
+            },
+            finish: "end_turn",
+          }
+          yield* session.updateMessage(assistantMsg2)
+
+          yield* session.updatePart({
+            id: PartID.ascending(),
+            messageID: assistantMsg2.id,
+            sessionID,
+            type: "text",
+            text: "The capital of France is Paris.",
+          })
 
-  test("cleanup with partID removes parts from the revert point onward", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        const session = await Session.create({})
-        const sid = session.id
-
-        const u1 = await user(sid)
-        const p1 = await text(sid, u1.id, "first part")
-        const p2 = await tool(sid, u1.id)
-        const p3 = await text(sid, u1.id, "third part")
-
-        // Set revert state pointing at a specific part
-        await Session.setRevert({
-          sessionID: sid,
-          revert: { messageID: u1.id, partID: p2.id },
-          summary: { additions: 0, deletions: 0, files: 0 },
-        })
-
-        const info = await Session.get(sid)
-        await SessionRevert.cleanup(info)
-
-        const msgs = await Session.messages({ sessionID: sid })
-        expect(msgs.length).toBe(1)
-        // Only the first part should remain (before the revert partID)
-        expect(msgs[0].parts.length).toBe(1)
-        expect(msgs[0].parts[0].id).toBe(p1.id)
-
-        const cleared = await Session.get(sid)
-        expect(cleared.revert).toBeUndefined()
-      },
-    })
-  })
+          let messages = yield* session.messages({ sessionID })
+          expect(messages.length).toBe(4)
+          const messageIds = messages.map((m) => m.info.id)
+          expect(messageIds).toContain(userMsg1.id)
+          expect(messageIds).toContain(userMsg2.id)
+          expect(messageIds).toContain(assistantMsg1.id)
+          expect(messageIds).toContain(assistantMsg2.id)
+
+          yield* revert.revert({
+            sessionID,
+            messageID: userMsg2.id,
+          })
 
-  test("cleanup removes messages after revert point but keeps earlier ones", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        const session = await Session.create({})
-        const sid = session.id
-
-        const u1 = await user(sid)
-        await text(sid, u1.id, "hello")
-        const a1 = await assistant(sid, u1.id, tmp.path)
-        await text(sid, a1.id, "hi back")
-
-        const u2 = await user(sid)
-        await text(sid, u2.id, "second question")
-        const a2 = await assistant(sid, u2.id, tmp.path)
-        await text(sid, a2.id, "second answer")
-
-        // Revert from u2 onward
-        await Session.setRevert({
-          sessionID: sid,
-          revert: { messageID: u2.id },
-          summary: { additions: 0, deletions: 0, files: 0 },
-        })
-
-        const info = await Session.get(sid)
-        await SessionRevert.cleanup(info)
-
-        const msgs = await Session.messages({ sessionID: sid })
-        const ids = msgs.map((m) => m.info.id)
-        expect(ids).toContain(u1.id)
-        expect(ids).toContain(a1.id)
-        expect(ids).not.toContain(u2.id)
-        expect(ids).not.toContain(a2.id)
-      },
-    })
-  })
+          let sessionInfo = yield* session.get(sessionID)
+          expect(sessionInfo.revert).toBeDefined()
+          expect(sessionInfo.revert?.messageID).toBeDefined()
+
+          messages = yield* session.messages({ sessionID })
+          expect(messages.length).toBe(4)
+
+          yield* revert.cleanup(sessionInfo)
+
+          messages = yield* session.messages({ sessionID })
+          const remainingIds = messages.map((m) => m.info.id)
+          expect(messages.length).toBeLessThan(4)
+          expect(remainingIds).not.toContain(userMsg2.id)
+          expect(remainingIds).not.toContain(assistantMsg2.id)
+
+          sessionInfo = yield* session.get(sessionID)
+          expect(sessionInfo.revert).toBeUndefined()
+
+          yield* session.remove(sessionID)
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "should properly clean up revert state before creating compaction message",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+
+          const info = yield* session.create({})
+          const sessionID = info.id
+
+          const userMsg = yield* session.updateMessage({
+            id: MessageID.ascending(),
+            role: "user",
+            sessionID,
+            agent: "default",
+            model: {
+              providerID: ProviderID.make("openai"),
+              modelID: ModelID.make("gpt-4"),
+            },
+            time: {
+              created: Date.now(),
+            },
+          })
 
-  test("cleanup is a no-op when session has no revert state", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        const session = await Session.create({})
-        const sid = session.id
-
-        const u1 = await user(sid)
-        await text(sid, u1.id, "hello")
-
-        const info = await Session.get(sid)
-        expect(info.revert).toBeUndefined()
-        await SessionRevert.cleanup(info)
-
-        const msgs = await Session.messages({ sessionID: sid })
-        expect(msgs.length).toBe(1)
-      },
-    })
-  })
+          yield* session.updatePart({
+            id: PartID.ascending(),
+            messageID: userMsg.id,
+            sessionID,
+            type: "text",
+            text: "Hello",
+          })
 
-  test("restore messages in sequential order", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        await fs.writeFile(path.join(tmp.path, "a.txt"), "a0")
-        await fs.writeFile(path.join(tmp.path, "b.txt"), "b0")
-        await fs.writeFile(path.join(tmp.path, "c.txt"), "c0")
-
-        const session = await Session.create({})
-        const sid = session.id
-
-        const turn = async (file: string, next: string) => {
-          const u = await user(sid)
-          await text(sid, u.id, `${file}:${next}`)
-          const a = await assistant(sid, u.id, tmp.path)
-          const before = await Snapshot.track()
-          if (!before) throw new Error("expected snapshot")
-          await fs.writeFile(path.join(tmp.path, file), next)
-          const after = await Snapshot.track()
-          if (!after) throw new Error("expected snapshot")
-          const patch = await Snapshot.patch(before)
-          await Session.updatePart({
+          const assistantMsg: MessageV2.Assistant = {
+            id: MessageID.ascending(),
+            role: "assistant",
+            sessionID,
+            mode: "default",
+            agent: "default",
+            path: {
+              cwd: dir,
+              root: dir,
+            },
+            cost: 0,
+            tokens: {
+              output: 0,
+              input: 0,
+              reasoning: 0,
+              cache: { read: 0, write: 0 },
+            },
+            modelID: ModelID.make("gpt-4"),
+            providerID: ProviderID.make("openai"),
+            parentID: userMsg.id,
+            time: {
+              created: Date.now(),
+            },
+            finish: "end_turn",
+          }
+          yield* session.updateMessage(assistantMsg)
+
+          yield* session.updatePart({
             id: PartID.ascending(),
-            messageID: a.id,
+            messageID: assistantMsg.id,
+            sessionID,
+            type: "text",
+            text: "Hi there!",
+          })
+
+          yield* revert.revert({
+            sessionID,
+            messageID: userMsg.id,
+          })
+
+          let sessionInfo = yield* session.get(sessionID)
+          expect(sessionInfo.revert).toBeDefined()
+
+          yield* revert.cleanup(sessionInfo)
+
+          sessionInfo = yield* session.get(sessionID)
+          expect(sessionInfo.revert).toBeUndefined()
+
+          const messages = yield* session.messages({ sessionID })
+          expect(messages.length).toBe(0)
+
+          yield* session.remove(sessionID)
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "cleanup with partID removes parts from the revert point onward",
+    provideTmpdirInstance(
+      () =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+
+          const info = yield* session.create({})
+          const sid = info.id
+
+          const u1 = yield* user(sid)
+          const p1 = yield* text(sid, u1.id, "first part")
+          const p2 = yield* tool(sid, u1.id)
+          yield* text(sid, u1.id, "third part")
+
+          yield* session.setRevert({
             sessionID: sid,
-            type: "step-start",
-            snapshot: before,
+            revert: { messageID: u1.id, partID: p2.id },
+            summary: { additions: 0, deletions: 0, files: 0 },
           })
-          await Session.updatePart({
-            id: PartID.ascending(),
-            messageID: a.id,
+
+          const state = yield* session.get(sid)
+          yield* revert.cleanup(state)
+
+          const msgs = yield* session.messages({ sessionID: sid })
+          expect(msgs.length).toBe(1)
+          expect(msgs[0].parts.length).toBe(1)
+          expect(msgs[0].parts[0].id).toBe(p1.id)
+
+          const cleared = yield* session.get(sid)
+          expect(cleared.revert).toBeUndefined()
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "cleanup removes messages after revert point but keeps earlier ones",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+
+          const info = yield* session.create({})
+          const sid = info.id
+
+          const u1 = yield* user(sid)
+          yield* text(sid, u1.id, "hello")
+          const a1 = yield* assistant(sid, u1.id, dir)
+          yield* text(sid, a1.id, "hi back")
+
+          const u2 = yield* user(sid)
+          yield* text(sid, u2.id, "second question")
+          const a2 = yield* assistant(sid, u2.id, dir)
+          yield* text(sid, a2.id, "second answer")
+
+          yield* session.setRevert({
             sessionID: sid,
-            type: "step-finish",
-            reason: "stop",
-            snapshot: after,
-            cost: 0,
-            tokens,
+            revert: { messageID: u2.id },
+            summary: { additions: 0, deletions: 0, files: 0 },
           })
-          await Session.updatePart({
-            id: PartID.ascending(),
-            messageID: a.id,
+
+          const state = yield* session.get(sid)
+          yield* revert.cleanup(state)
+
+          const msgs = yield* session.messages({ sessionID: sid })
+          const ids = msgs.map((m) => m.info.id)
+          expect(ids).toContain(u1.id)
+          expect(ids).toContain(a1.id)
+          expect(ids).not.toContain(u2.id)
+          expect(ids).not.toContain(a2.id)
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "cleanup is a no-op when session has no revert state",
+    provideTmpdirInstance(
+      () =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+
+          const info = yield* session.create({})
+          const sid = info.id
+
+          const u1 = yield* user(sid)
+          yield* text(sid, u1.id, "hello")
+
+          const state = yield* session.get(sid)
+          expect(state.revert).toBeUndefined()
+          yield* revert.cleanup(state)
+
+          const msgs = yield* session.messages({ sessionID: sid })
+          expect(msgs.length).toBe(1)
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "restore messages in sequential order",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+          const snapshot = yield* Snapshot.Service
+
+          yield* write(path.join(dir, "a.txt"), "a0")
+          yield* write(path.join(dir, "b.txt"), "b0")
+          yield* write(path.join(dir, "c.txt"), "c0")
+
+          const info = yield* session.create({})
+          const sid = info.id
+
+          const turn = Effect.fn("test.turn")(function* (file: string, next: string) {
+            const u = yield* user(sid)
+            yield* text(sid, u.id, `${file}:${next}`)
+            const a = yield* assistant(sid, u.id, dir)
+            const before = yield* snapshot.track()
+            if (!before) throw new Error("expected snapshot")
+            yield* write(path.join(dir, file), next)
+            const after = yield* snapshot.track()
+            if (!after) throw new Error("expected snapshot")
+            const patch = yield* snapshot.patch(before)
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "step-start",
+              snapshot: before,
+            })
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "step-finish",
+              reason: "stop",
+              snapshot: after,
+              cost: 0,
+              tokens,
+            })
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "patch",
+              hash: patch.hash,
+              files: patch.files,
+            })
+            return u.id
+          })
+
+          const first = yield* turn("a.txt", "a1")
+          const second = yield* turn("b.txt", "b2")
+          const third = yield* turn("c.txt", "c3")
+
+          yield* revert.revert({
             sessionID: sid,
-            type: "patch",
-            hash: patch.hash,
-            files: patch.files,
+            messageID: first,
           })
-          return u.id
-        }
-
-        const first = await turn("a.txt", "a1")
-        const second = await turn("b.txt", "b2")
-        const third = await turn("c.txt", "c3")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: first,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(first)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a0")
-        expect(await fs.readFile(path.join(tmp.path, "b.txt"), "utf-8")).toBe("b0")
-        expect(await fs.readFile(path.join(tmp.path, "c.txt"), "utf-8")).toBe("c0")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: second,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(second)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a1")
-        expect(await fs.readFile(path.join(tmp.path, "b.txt"), "utf-8")).toBe("b0")
-        expect(await fs.readFile(path.join(tmp.path, "c.txt"), "utf-8")).toBe("c0")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: third,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(third)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a1")
-        expect(await fs.readFile(path.join(tmp.path, "b.txt"), "utf-8")).toBe("b2")
-        expect(await fs.readFile(path.join(tmp.path, "c.txt"), "utf-8")).toBe("c0")
-
-        await SessionRevert.unrevert({
-          sessionID: sid,
-        })
-        expect((await Session.get(sid)).revert).toBeUndefined()
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a1")
-        expect(await fs.readFile(path.join(tmp.path, "b.txt"), "utf-8")).toBe("b2")
-        expect(await fs.readFile(path.join(tmp.path, "c.txt"), "utf-8")).toBe("c3")
-      },
-    })
-  })
+          expect((yield* session.get(sid)).revert?.messageID).toBe(first)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a0")
+          expect(yield* read(path.join(dir, "b.txt"))).toBe("b0")
+          expect(yield* read(path.join(dir, "c.txt"))).toBe("c0")
 
-  test("restore same file in sequential order", async () => {
-    await using tmp = await tmpdir({ git: true })
-    await Instance.provide({
-      directory: tmp.path,
-      fn: async () => {
-        await fs.writeFile(path.join(tmp.path, "a.txt"), "a0")
-
-        const session = await Session.create({})
-        const sid = session.id
-
-        const turn = async (next: string) => {
-          const u = await user(sid)
-          await text(sid, u.id, `a.txt:${next}`)
-          const a = await assistant(sid, u.id, tmp.path)
-          const before = await Snapshot.track()
-          if (!before) throw new Error("expected snapshot")
-          await fs.writeFile(path.join(tmp.path, "a.txt"), next)
-          const after = await Snapshot.track()
-          if (!after) throw new Error("expected snapshot")
-          const patch = await Snapshot.patch(before)
-          await Session.updatePart({
-            id: PartID.ascending(),
-            messageID: a.id,
+          yield* revert.revert({
             sessionID: sid,
-            type: "step-start",
-            snapshot: before,
+            messageID: second,
           })
-          await Session.updatePart({
-            id: PartID.ascending(),
-            messageID: a.id,
+          expect((yield* session.get(sid)).revert?.messageID).toBe(second)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a1")
+          expect(yield* read(path.join(dir, "b.txt"))).toBe("b0")
+          expect(yield* read(path.join(dir, "c.txt"))).toBe("c0")
+
+          yield* revert.revert({
             sessionID: sid,
-            type: "step-finish",
-            reason: "stop",
-            snapshot: after,
-            cost: 0,
-            tokens,
+            messageID: third,
           })
-          await Session.updatePart({
-            id: PartID.ascending(),
-            messageID: a.id,
+          expect((yield* session.get(sid)).revert?.messageID).toBe(third)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a1")
+          expect(yield* read(path.join(dir, "b.txt"))).toBe("b2")
+          expect(yield* read(path.join(dir, "c.txt"))).toBe("c0")
+
+          yield* revert.unrevert({
             sessionID: sid,
-            type: "patch",
-            hash: patch.hash,
-            files: patch.files,
           })
-          return u.id
-        }
-
-        const first = await turn("a1")
-        const second = await turn("a2")
-        const third = await turn("a3")
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a3")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: first,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(first)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a0")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: second,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(second)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a1")
-
-        await SessionRevert.revert({
-          sessionID: sid,
-          messageID: third,
-        })
-        expect((await Session.get(sid)).revert?.messageID).toBe(third)
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a2")
-
-        await SessionRevert.unrevert({
-          sessionID: sid,
-        })
-        expect((await Session.get(sid)).revert).toBeUndefined()
-        expect(await fs.readFile(path.join(tmp.path, "a.txt"), "utf-8")).toBe("a3")
-      },
-    })
-  })
+          expect((yield* session.get(sid)).revert).toBeUndefined()
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a1")
+          expect(yield* read(path.join(dir, "b.txt"))).toBe("b2")
+          expect(yield* read(path.join(dir, "c.txt"))).toBe("c3")
+        }),
+      { git: true },
+    ),
+  )
+
+  it.live(
+    "restore same file in sequential order",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const session = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+          const snapshot = yield* Snapshot.Service
+
+          yield* write(path.join(dir, "a.txt"), "a0")
+
+          const info = yield* session.create({})
+          const sid = info.id
+
+          const turn = Effect.fn("test.turnSame")(function* (next: string) {
+            const u = yield* user(sid)
+            yield* text(sid, u.id, `a.txt:${next}`)
+            const a = yield* assistant(sid, u.id, dir)
+            const before = yield* snapshot.track()
+            if (!before) throw new Error("expected snapshot")
+            yield* write(path.join(dir, "a.txt"), next)
+            const after = yield* snapshot.track()
+            if (!after) throw new Error("expected snapshot")
+            const patch = yield* snapshot.patch(before)
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "step-start",
+              snapshot: before,
+            })
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "step-finish",
+              reason: "stop",
+              snapshot: after,
+              cost: 0,
+              tokens,
+            })
+            yield* session.updatePart({
+              id: PartID.ascending(),
+              messageID: a.id,
+              sessionID: sid,
+              type: "patch",
+              hash: patch.hash,
+              files: patch.files,
+            })
+            return u.id
+          })
+
+          const first = yield* turn("a1")
+          const second = yield* turn("a2")
+          const third = yield* turn("a3")
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a3")
+
+          yield* revert.revert({
+            sessionID: sid,
+            messageID: first,
+          })
+          expect((yield* session.get(sid)).revert?.messageID).toBe(first)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a0")
+
+          yield* revert.revert({
+            sessionID: sid,
+            messageID: second,
+          })
+          expect((yield* session.get(sid)).revert?.messageID).toBe(second)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a1")
+
+          yield* revert.revert({
+            sessionID: sid,
+            messageID: third,
+          })
+          expect((yield* session.get(sid)).revert?.messageID).toBe(third)
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a2")
+
+          yield* revert.unrevert({
+            sessionID: sid,
+          })
+          expect((yield* session.get(sid)).revert).toBeUndefined()
+          expect(yield* read(path.join(dir, "a.txt"))).toBe("a3")
+        }),
+      { git: true },
+    ),
+  )
 })