feat(core): add session snapshot and revert system (#33226)
This commit is contained in:
parent
25c6abc31a
commit
9bb5370205
42 changed files with 2311 additions and 365 deletions
|
|
@ -8,6 +8,8 @@ import { AbsolutePath, RelativePath } from "../schema"
|
|||
import { WorkspaceV2 } from "../workspace"
|
||||
import { SessionSchema } from "./schema"
|
||||
import { SessionTable } from "./sql"
|
||||
import { SessionMessageID } from "./message-id"
|
||||
import { Snapshot } from "../snapshot"
|
||||
|
||||
export function fromRow(row: typeof SessionTable.$inferSelect): SessionSchema.Info {
|
||||
return SessionSchema.Info.make({
|
||||
|
|
@ -38,6 +40,9 @@ export function fromRow(row: typeof SessionTable.$inferSelect): SessionSchema.In
|
|||
workspaceID: row.workspace_id ? WorkspaceV2.ID.make(row.workspace_id) : undefined,
|
||||
}),
|
||||
subpath: row.path ? RelativePath.make(row.path) : undefined,
|
||||
revert: row.revert
|
||||
? { ...row.revert, messageID: SessionMessageID.ID.make(row.revert.messageID) }
|
||||
: undefined,
|
||||
time: {
|
||||
created: DateTime.makeUnsafe(row.time_created),
|
||||
updated: DateTime.makeUnsafe(row.time_updated),
|
||||
|
|
|
|||
|
|
@ -212,7 +212,12 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
|
|||
draft.finish = event.data.finish
|
||||
draft.cost = event.data.cost
|
||||
draft.tokens = event.data.tokens
|
||||
if (event.data.snapshot) draft.snapshot = { ...draft.snapshot, end: event.data.snapshot }
|
||||
if (event.data.snapshot || event.data.files)
|
||||
draft.snapshot = {
|
||||
...draft.snapshot,
|
||||
end: event.data.snapshot,
|
||||
files: event.data.files ? Array.from(event.data.files) : undefined,
|
||||
}
|
||||
})
|
||||
},
|
||||
"session.next.step.failed": (event) => {
|
||||
|
|
@ -380,6 +385,9 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
|
|||
}),
|
||||
)
|
||||
},
|
||||
"session.next.revert.staged": () => Effect.void,
|
||||
"session.next.revert.cleared": () => Effect.void,
|
||||
"session.next.revert.committed": () => Effect.void,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
export * as SessionProjector from "./projector"
|
||||
|
||||
import { and, desc, eq, sql } from "drizzle-orm"
|
||||
import { and, desc, eq, gt, or, sql } from "drizzle-orm"
|
||||
import { DateTime, Effect, Layer, Schema } from "effect"
|
||||
import { Database } from "../database/database"
|
||||
import { EventV2 } from "../event"
|
||||
|
|
@ -13,8 +13,9 @@ import { SessionMessageUpdater } from "./message-updater"
|
|||
import { SessionInput } from "./input"
|
||||
import { WorkspaceV2 } from "../workspace"
|
||||
import { SessionContextEpoch } from "./context-epoch"
|
||||
import { MessageTable, PartTable, SessionMessageTable, SessionTable } from "./sql"
|
||||
import { MessageTable, PartTable, SessionInputTable, SessionMessageTable, SessionTable } from "./sql"
|
||||
import type { DeepMutable } from "../schema"
|
||||
import { SessionMessageID } from "./message-id"
|
||||
|
||||
type DatabaseService = Database.Interface["db"]
|
||||
|
||||
|
|
@ -66,7 +67,7 @@ function sessionRow(info: SessionV1.SessionInfo): typeof SessionTable.$inferInse
|
|||
tokens_reasoning: (info.tokens ?? { reasoning: 0 }).reasoning,
|
||||
tokens_cache_read: (info.tokens ?? { cache: { read: 0 } }).cache.read,
|
||||
tokens_cache_write: (info.tokens ?? { cache: { write: 0 } }).cache.write,
|
||||
revert: info.revert ?? null,
|
||||
revert: info.revert ? { ...info.revert, messageID: SessionMessageID.ID.make(info.revert.messageID) } : null,
|
||||
permission: info.permission ? [...info.permission] : undefined,
|
||||
time_created: info.time.created,
|
||||
time_updated: info.time.updated,
|
||||
|
|
@ -393,6 +394,40 @@ export const layer = Layer.effectDiscard(
|
|||
yield* events.project(SessionEvent.Reasoning.Ended, (event) => run(db, event))
|
||||
// yield* events.project(SessionEvent.Retried, (event) => run(db, event))
|
||||
yield* events.project(SessionEvent.Compaction.Ended, (event) => run(db, event))
|
||||
yield* events.project(SessionEvent.RevertEvent.Staged, (event) =>
|
||||
db
|
||||
.update(SessionTable)
|
||||
.set({
|
||||
revert: { ...event.data.revert, files: event.data.revert.files ? [...event.data.revert.files] : undefined },
|
||||
time_updated: DateTime.toEpochMillis(event.data.timestamp),
|
||||
})
|
||||
.where(eq(SessionTable.id, event.data.sessionID))
|
||||
.run()
|
||||
.pipe(Effect.orDie, Effect.asVoid),
|
||||
)
|
||||
yield* events.project(SessionEvent.RevertEvent.Cleared, (event) =>
|
||||
db
|
||||
.update(SessionTable)
|
||||
.set({ revert: null, time_updated: DateTime.toEpochMillis(event.data.timestamp) })
|
||||
.where(eq(SessionTable.id, event.data.sessionID))
|
||||
.run()
|
||||
.pipe(Effect.orDie, Effect.asVoid),
|
||||
)
|
||||
yield* events.project(SessionEvent.RevertEvent.Committed, (event) =>
|
||||
Effect.gen(function* () {
|
||||
const boundary = yield* db
|
||||
.select({ seq: SessionMessageTable.seq })
|
||||
.from(SessionMessageTable)
|
||||
.where(and(eq(SessionMessageTable.session_id, event.data.sessionID), eq(SessionMessageTable.id, event.data.messageID)))
|
||||
.get()
|
||||
.pipe(Effect.orDie)
|
||||
if (!boundary) return yield* Effect.die(`Revert boundary message not found: ${event.data.messageID}`)
|
||||
yield* db.delete(SessionMessageTable).where(and(eq(SessionMessageTable.session_id, event.data.sessionID), gt(SessionMessageTable.seq, boundary.seq))).run().pipe(Effect.orDie)
|
||||
yield* db.delete(SessionInputTable).where(and(eq(SessionInputTable.session_id, event.data.sessionID), or(gt(SessionInputTable.admitted_seq, boundary.seq), gt(SessionInputTable.promoted_seq, boundary.seq)))).run().pipe(Effect.orDie)
|
||||
yield* db.update(SessionTable).set({ revert: null, time_updated: DateTime.toEpochMillis(event.data.timestamp) }).where(eq(SessionTable.id, event.data.sessionID)).run().pipe(Effect.orDie)
|
||||
yield* SessionContextEpoch.reset(db, event.data.sessionID)
|
||||
}),
|
||||
)
|
||||
}),
|
||||
)
|
||||
|
||||
|
|
|
|||
118
packages/core/src/session/revert.ts
Normal file
118
packages/core/src/session/revert.ts
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
export * as SessionRevert from "./revert"
|
||||
|
||||
import { and, asc, eq, gt } from "drizzle-orm"
|
||||
import { DateTime, Effect, Schema } from "effect"
|
||||
import { Database } from "../database/database"
|
||||
import { EventV2 } from "../event"
|
||||
import { RelativePath } from "../schema"
|
||||
import { Snapshot } from "../snapshot"
|
||||
import { SessionEvent } from "./event"
|
||||
import { SessionMessage } from "./message"
|
||||
import { SessionSchema } from "./schema"
|
||||
import { SessionMessageTable } from "./sql"
|
||||
|
||||
export class MessageNotFoundError extends Schema.TaggedErrorClass<MessageNotFoundError>()(
|
||||
"Session.MessageNotFoundError",
|
||||
{
|
||||
sessionID: SessionSchema.ID,
|
||||
messageID: SessionMessage.ID,
|
||||
},
|
||||
) {}
|
||||
|
||||
interface BoundaryInput {
|
||||
readonly sessionID: SessionSchema.ID
|
||||
readonly messageID: SessionMessage.ID
|
||||
}
|
||||
|
||||
const plan = Effect.fn("SessionRevert.plan")(function* (input: BoundaryInput) {
|
||||
const db = (yield* Database.Service).db
|
||||
const boundary = yield* db
|
||||
.select({ seq: SessionMessageTable.seq })
|
||||
.from(SessionMessageTable)
|
||||
.where(and(eq(SessionMessageTable.session_id, input.sessionID), eq(SessionMessageTable.id, input.messageID)))
|
||||
.get()
|
||||
.pipe(Effect.orDie)
|
||||
if (!boundary) return yield* new MessageNotFoundError(input)
|
||||
const rows = yield* db
|
||||
.select()
|
||||
.from(SessionMessageTable)
|
||||
.where(
|
||||
and(
|
||||
eq(SessionMessageTable.session_id, input.sessionID),
|
||||
eq(SessionMessageTable.type, "assistant"),
|
||||
gt(SessionMessageTable.seq, boundary.seq),
|
||||
),
|
||||
)
|
||||
.orderBy(asc(SessionMessageTable.seq))
|
||||
.all()
|
||||
.pipe(Effect.orDie)
|
||||
const decode = Schema.decodeUnknownEffect(SessionMessage.Message)
|
||||
const files = new Map<RelativePath, Snapshot.ID>()
|
||||
for (const row of rows) {
|
||||
const message = yield* decode({ ...row.data, id: row.id, type: row.type }).pipe(Effect.orDie)
|
||||
if (message.type !== "assistant" || !message.snapshot?.start) continue
|
||||
for (const file of message.snapshot.files ?? [])
|
||||
if (!files.has(file)) files.set(file, Snapshot.ID.make(message.snapshot.start))
|
||||
}
|
||||
return files
|
||||
})
|
||||
|
||||
export const stage = Effect.fn("SessionRevert.stage")(function* (input: {
|
||||
readonly session: SessionSchema.Info
|
||||
readonly messageID: SessionMessage.ID
|
||||
readonly files?: boolean
|
||||
}) {
|
||||
const snapshot = yield* Snapshot.Service
|
||||
const events = yield* EventV2.Service
|
||||
const original = input.session.revert?.snapshot
|
||||
? Snapshot.ID.make(input.session.revert.snapshot)
|
||||
: (yield* snapshot.capture())
|
||||
const next = yield* plan({ sessionID: input.session.id, messageID: input.messageID })
|
||||
const restore = new Map<RelativePath, Snapshot.ID>()
|
||||
if (original) {
|
||||
for (const file of input.session.revert?.files ?? []) restore.set(file.path, original)
|
||||
}
|
||||
if (input.files !== false) for (const [file, tree] of next) restore.set(file, tree)
|
||||
if (restore.size) yield* snapshot.restore({ files: restore })
|
||||
const paths = input.files === false ? [] : Array.from(next.keys())
|
||||
const files = original
|
||||
? yield* snapshot.diff({ from: original, to: (yield* snapshot.capture()) ?? original, paths })
|
||||
: []
|
||||
const revert = {
|
||||
messageID: input.messageID,
|
||||
snapshot: original,
|
||||
diff: files.map((file) => file.patch).join("").trim(),
|
||||
files,
|
||||
} satisfies SessionSchema.Info["revert"]
|
||||
yield* events.publish(SessionEvent.RevertEvent.Staged, {
|
||||
sessionID: input.session.id,
|
||||
timestamp: yield* DateTime.now,
|
||||
revert,
|
||||
})
|
||||
return revert
|
||||
})
|
||||
|
||||
export const clear = Effect.fn("SessionRevert.clear")(function* (session: SessionSchema.Info) {
|
||||
if (!session.revert) return
|
||||
const snapshot = yield* Snapshot.Service
|
||||
const original = session.revert.snapshot ? Snapshot.ID.make(session.revert.snapshot) : undefined
|
||||
if (original)
|
||||
yield* snapshot.restore({
|
||||
files: new Map((session.revert.files ?? []).map((file) => [file.path, original])),
|
||||
})
|
||||
const events = yield* EventV2.Service
|
||||
yield* events.publish(SessionEvent.RevertEvent.Cleared, {
|
||||
sessionID: session.id,
|
||||
timestamp: yield* DateTime.now,
|
||||
})
|
||||
})
|
||||
|
||||
export const commit = Effect.fn("SessionRevert.commit")(function* (session: SessionSchema.Info) {
|
||||
if (!session.revert) return
|
||||
const events = yield* EventV2.Service
|
||||
yield* events.publish(SessionEvent.RevertEvent.Committed, {
|
||||
sessionID: session.id,
|
||||
messageID: session.revert.messageID,
|
||||
timestamp: yield* DateTime.now,
|
||||
})
|
||||
})
|
||||
|
|
@ -35,6 +35,7 @@ import { SessionRunnerModel } from "./model"
|
|||
import { createLLMEventPublisher } from "./publish-llm-event"
|
||||
import { toLLMMessages } from "./to-llm-message"
|
||||
import { MAX_STEPS_PROMPT } from "./max-steps"
|
||||
import { Snapshot } from "../../snapshot"
|
||||
|
||||
/**
|
||||
* Runs one durable coding-agent Session until it settles.
|
||||
|
|
@ -100,6 +101,7 @@ export const layer = Layer.effect(
|
|||
const skillGuidance = yield* SkillGuidance.Service
|
||||
const referenceGuidance = yield* ReferenceGuidance.Service
|
||||
const config = yield* Config.Service
|
||||
const snapshots = yield* Snapshot.Service
|
||||
const db = (yield* Database.Service).db
|
||||
const compaction = SessionCompaction.make({ events, llm, config: yield* config.entries() })
|
||||
const getSession = Effect.fn("SessionRunner.getSession")(function* (sessionID: SessionSchema.ID) {
|
||||
|
|
@ -205,6 +207,7 @@ export const layer = Layer.effect(
|
|||
})
|
||||
if (yield* compaction.compactIfNeeded({ sessionID: session.id, entries, model, request }))
|
||||
return yield* Effect.die(continueAfterCompaction(currentStep))
|
||||
const startSnapshot = yield* snapshots.capture()
|
||||
const publisher = createLLMEventPublisher(events, {
|
||||
sessionID: session.id,
|
||||
agent: agent.id,
|
||||
|
|
@ -213,6 +216,7 @@ export const layer = Layer.effect(
|
|||
providerID: ProviderV2.ID.make(model.provider),
|
||||
...(session.model?.variant === undefined ? {} : { variant: session.model.variant }),
|
||||
},
|
||||
snapshot: startSnapshot,
|
||||
})
|
||||
const withPublication = Semaphore.makeUnsafe(1).withPermit
|
||||
const publish = (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) =>
|
||||
|
|
@ -302,6 +306,23 @@ export const layer = Layer.effect(
|
|||
const message = failure instanceof Error ? failure.message : String(failure)
|
||||
yield* withPublication(publisher.failUnsettledTools(`Tool execution failed: ${message}`))
|
||||
}
|
||||
const stepSettlement = publisher.stepSettlement()
|
||||
if (stepSettlement && !publisher.hasProviderError()) {
|
||||
const endSnapshot = yield* snapshots.capture()
|
||||
const files = startSnapshot && endSnapshot
|
||||
? yield* snapshots.files({ from: startSnapshot, to: endSnapshot }).pipe(Effect.catch(() => Effect.succeed(undefined)))
|
||||
: undefined
|
||||
yield* withPublication(events.publish(SessionEvent.Step.Ended, {
|
||||
sessionID: session.id,
|
||||
timestamp: yield* DateTime.now,
|
||||
assistantMessageID: yield* publisher.startAssistant(),
|
||||
finish: stepSettlement.finish,
|
||||
cost: 0,
|
||||
tokens: stepSettlement.tokens,
|
||||
snapshot: endSnapshot,
|
||||
files,
|
||||
}))
|
||||
}
|
||||
if (publisher.hasProviderError())
|
||||
yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
|
||||
if (stream._tag === "Success" && !publisher.hasProviderError())
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ type Input = {
|
|||
readonly sessionID: SessionSchema.ID
|
||||
readonly agent: string
|
||||
readonly model: ModelV2.Ref
|
||||
readonly snapshot?: string
|
||||
}
|
||||
|
||||
const safe = (value: number | undefined) => Math.max(0, Number.isFinite(value) ? (value ?? 0) : 0)
|
||||
|
|
@ -68,6 +69,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
let assistantActive = false
|
||||
let assistantFailed = false
|
||||
let providerFailed = false
|
||||
let stepSettlement: { readonly finish: string; readonly tokens: ReturnType<typeof tokens> } | undefined
|
||||
|
||||
const startAssistant = Effect.fnUntraced(function* () {
|
||||
if (assistantMessageID !== undefined) return assistantMessageID
|
||||
|
|
@ -77,6 +79,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
...input,
|
||||
assistantMessageID,
|
||||
timestamp: yield* timestamp,
|
||||
snapshot: input.snapshot,
|
||||
})
|
||||
return assistantMessageID
|
||||
})
|
||||
|
|
@ -393,14 +396,8 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
case "step-finish":
|
||||
yield* flush()
|
||||
assistantActive = false
|
||||
yield* events.publish(SessionEvent.Step.Ended, {
|
||||
sessionID: input.sessionID,
|
||||
timestamp: yield* timestamp,
|
||||
assistantMessageID: yield* startAssistant(),
|
||||
finish: event.reason,
|
||||
cost: 0,
|
||||
tokens: tokens(event.usage),
|
||||
})
|
||||
if (stepSettlement) return yield* Effect.die("Duplicate step finish")
|
||||
stepSettlement = { finish: event.reason, tokens: tokens(event.usage) }
|
||||
return
|
||||
case "finish":
|
||||
return
|
||||
|
|
@ -419,6 +416,8 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
hasActiveAssistant: () => assistantActive,
|
||||
hasAssistantStarted: () => assistantMessageID !== undefined,
|
||||
hasProviderError: () => providerFailed,
|
||||
stepSettlement: () => stepSettlement,
|
||||
startAssistant,
|
||||
assistantMessageID: assistantMessageIDForTool,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import { WorkspaceV2 } from "../workspace"
|
|||
import { Timestamps } from "../database/schema.sql"
|
||||
import type { SystemContext } from "../system-context/index"
|
||||
import { AgentV2 } from "../agent"
|
||||
import type { Revert } from "@opencode-ai/schema/revert"
|
||||
|
||||
type SessionMessageData = Omit<(typeof SessionMessage.Message)["Encoded"], "type" | "id">
|
||||
type V1MessageData = Omit<SessionV1.Info, "id" | "sessionID">
|
||||
|
|
@ -37,7 +38,7 @@ export const SessionTable = sqliteTable(
|
|||
summary_additions: integer(),
|
||||
summary_deletions: integer(),
|
||||
summary_files: integer(),
|
||||
summary_diffs: text({ mode: "json" }).$type<Snapshot.FileDiff[]>(),
|
||||
summary_diffs: text({ mode: "json" }).$type<Snapshot.LegacyFileDiff[]>(),
|
||||
metadata: text({ mode: "json" }).$type<Record<string, unknown>>(),
|
||||
cost: real().notNull().default(0),
|
||||
tokens_input: integer().notNull().default(0),
|
||||
|
|
@ -45,7 +46,7 @@ export const SessionTable = sqliteTable(
|
|||
tokens_reasoning: integer().notNull().default(0),
|
||||
tokens_cache_read: integer().notNull().default(0),
|
||||
tokens_cache_write: integer().notNull().default(0),
|
||||
revert: text({ mode: "json" }).$type<{ messageID: MessageID; partID?: PartID; snapshot?: string; diff?: string }>(),
|
||||
revert: text({ mode: "json" }).$type<Revert.State>(),
|
||||
permission: text({ mode: "json" }).$type<PermissionV1.Ruleset>(),
|
||||
agent: text(),
|
||||
model: text({ mode: "json" }).$type<{
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue