feat(core): add durable session fork event
This commit is contained in:
parent
360d85a521
commit
b2d46ecd7e
11 changed files with 582 additions and 82 deletions
|
|
@ -3,7 +3,7 @@ export * as EventV2 from "./event"
|
|||
import { Cause, Context, Effect, Layer, Option, PubSub, Queue, Schema, Stream } from "effect"
|
||||
import { Event } from "@opencode-ai/schema/event"
|
||||
import type { Data, Definition, Payload } from "@opencode-ai/schema/event"
|
||||
import { and, asc, eq, gt, inArray } from "drizzle-orm"
|
||||
import { and, asc, eq, gt, inArray, sql } from "drizzle-orm"
|
||||
import { Database } from "./database/database"
|
||||
import { EventSequenceTable, EventTable } from "./event/sql"
|
||||
import { Location } from "./location"
|
||||
|
|
@ -31,6 +31,22 @@ export const latestSequence = Effect.fn("EventV2.latestSequence")(function* (
|
|||
return row?.seq ?? -1
|
||||
})
|
||||
|
||||
export const reserveSequence = Effect.fn("EventV2.reserveSequence")(function* (
|
||||
db: Database.Interface["db"],
|
||||
aggregateID: string,
|
||||
seq: number,
|
||||
) {
|
||||
yield* db
|
||||
.insert(EventSequenceTable)
|
||||
.values([{ aggregate_id: aggregateID, seq }])
|
||||
.onConflictDoUpdate({
|
||||
target: EventSequenceTable.aggregate_id,
|
||||
set: { seq: sql`max(${EventSequenceTable.seq}, ${seq})` },
|
||||
})
|
||||
.run()
|
||||
.pipe(Effect.orDie)
|
||||
})
|
||||
|
||||
export type SerializedEvent = {
|
||||
readonly id: ID
|
||||
readonly type: string
|
||||
|
|
@ -327,7 +343,7 @@ export const layerWith = (options?: LayerOptions) =>
|
|||
.onConflictDoUpdate({
|
||||
target: EventSequenceTable.aggregate_id,
|
||||
set: {
|
||||
seq,
|
||||
seq: sql`max(${EventSequenceTable.seq}, ${seq})`,
|
||||
...(input?.ownerID && row?.ownerID == null ? { owner_id: input.ownerID } : {}),
|
||||
},
|
||||
})
|
||||
|
|
|
|||
|
|
@ -90,6 +90,11 @@ type CompactInput = {
|
|||
sessionID: SessionSchema.ID
|
||||
}
|
||||
|
||||
type ForkInput = {
|
||||
sessionID: SessionSchema.ID
|
||||
messageID?: SessionMessage.ID
|
||||
}
|
||||
|
||||
export class NotFoundError extends Schema.TaggedErrorClass<NotFoundError>()("Session.NotFoundError", {
|
||||
sessionID: SessionSchema.ID,
|
||||
}) {}
|
||||
|
|
@ -113,11 +118,18 @@ export class BusyError extends Schema.TaggedErrorClass<BusyError>()("Session.Bus
|
|||
export const MessageNotFoundError = SessionRevert.MessageNotFoundError
|
||||
export type MessageNotFoundError = SessionRevert.MessageNotFoundError
|
||||
|
||||
export type Error = NotFoundError | MessageDecodeError | OperationUnavailableError | PromptConflictError | BusyError
|
||||
export type Error =
|
||||
| NotFoundError
|
||||
| MessageDecodeError
|
||||
| OperationUnavailableError
|
||||
| PromptConflictError
|
||||
| BusyError
|
||||
| MessageNotFoundError
|
||||
|
||||
export interface Interface {
|
||||
readonly list: (input?: ListInput) => Effect.Effect<SessionSchema.Info[]>
|
||||
readonly create: (input: CreateInput) => Effect.Effect<SessionSchema.Info, NotFoundError>
|
||||
readonly fork: (input: ForkInput) => Effect.Effect<SessionSchema.Info, NotFoundError | MessageNotFoundError>
|
||||
readonly get: (sessionID: SessionSchema.ID) => Effect.Effect<SessionSchema.Info, NotFoundError>
|
||||
readonly messages: (input: {
|
||||
sessionID: SessionSchema.ID
|
||||
|
|
@ -272,6 +284,52 @@ export const layer = Layer.effect(
|
|||
// TODO: Restore recorded sessions onto replacement synchronized workspaces in a future API slice.
|
||||
return yield* result.get(sessionID).pipe(Effect.orDie)
|
||||
}),
|
||||
fork: Effect.fn("V2Session.fork")(function* (input) {
|
||||
const parent = yield* result.get(input.sessionID)
|
||||
const boundary = input.messageID
|
||||
? 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)
|
||||
: undefined
|
||||
if (input.messageID && !boundary)
|
||||
return yield* new MessageNotFoundError({ sessionID: input.sessionID, messageID: input.messageID })
|
||||
const copied = yield* db
|
||||
.select({ seq: SessionMessageTable.seq })
|
||||
.from(SessionMessageTable)
|
||||
.where(
|
||||
and(
|
||||
eq(SessionMessageTable.session_id, input.sessionID),
|
||||
boundary === undefined ? undefined : lt(SessionMessageTable.seq, boundary.seq),
|
||||
),
|
||||
)
|
||||
.orderBy(desc(SessionMessageTable.seq))
|
||||
.limit(1)
|
||||
.get()
|
||||
.pipe(Effect.orDie)
|
||||
const sessionID = SessionSchema.ID.create()
|
||||
yield* events.publish(SessionEvent.Forked, {
|
||||
sessionID,
|
||||
parentID: parent.id,
|
||||
slug: Slug.create(),
|
||||
title: forkTitle(parent.title),
|
||||
agent: parent.agent,
|
||||
model: parent.model,
|
||||
messageID: input.messageID,
|
||||
copiedSeq: copied?.seq ?? 0,
|
||||
timestamp: yield* DateTime.now,
|
||||
}, {
|
||||
commit: (seq) =>
|
||||
copied && copied.seq > seq
|
||||
? EventV2.reserveSequence(db, sessionID, copied.seq)
|
||||
: Effect.void,
|
||||
})
|
||||
return yield* result.get(sessionID).pipe(Effect.orDie)
|
||||
}),
|
||||
get: Effect.fn("V2Session.get")(function* (sessionID) {
|
||||
const session = yield* store.get(sessionID)
|
||||
if (!session) return yield* new NotFoundError({ sessionID })
|
||||
|
|
@ -505,6 +563,12 @@ export const defaultLayer = layer.pipe(
|
|||
Layer.orDie,
|
||||
)
|
||||
|
||||
const forkTitle = (value: string) => {
|
||||
const match = value.match(/^(.+) \(fork #(\d+)\)$/)
|
||||
if (match) return `${match[1]} (fork #${Number.parseInt(match[2], 10) + 1})`
|
||||
return `${value} (fork #1)`
|
||||
}
|
||||
|
||||
const resolvePrompt = (input: PromptInput.Prompt) =>
|
||||
Prompt.make({
|
||||
text: input.text,
|
||||
|
|
|
|||
|
|
@ -124,6 +124,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
|
|||
},
|
||||
"session.next.moved": () => Effect.void,
|
||||
"session.next.renamed": () => Effect.void,
|
||||
"session.next.forked": () => Effect.void,
|
||||
"session.next.prompted": (event) => {
|
||||
return adapter.appendMessage(
|
||||
SessionMessage.User.make({
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
export * as SessionProjector from "./projector"
|
||||
|
||||
import { and, desc, eq, gt, or, sql } from "drizzle-orm"
|
||||
import { and, asc, desc, eq, gt, inArray, lt, or, sql } from "drizzle-orm"
|
||||
import { DateTime, Effect, Layer, Schema } from "effect"
|
||||
import { Database } from "../database/database"
|
||||
import { EventV2 } from "../event"
|
||||
|
|
@ -17,6 +17,7 @@ import { MessageTable, PartTable, SessionInputTable, SessionMessageTable, Sessio
|
|||
import type { DeepMutable } from "../schema"
|
||||
|
||||
type DatabaseService = Database.Interface["db"]
|
||||
type MessageEvent = Exclude<SessionEvent.Event, typeof SessionEvent.Forked.Type>
|
||||
|
||||
const decodeMessage = Schema.decodeUnknownSync(SessionMessage.Message)
|
||||
const encodeMessage = Schema.encodeSync(SessionMessage.Message)
|
||||
|
|
@ -33,6 +34,13 @@ type Usage = {
|
|||
}
|
||||
}
|
||||
|
||||
const ForkBatchSize = 500
|
||||
|
||||
const emptyUsage = (): Usage => ({
|
||||
cost: 0,
|
||||
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
})
|
||||
|
||||
function usage(part: (typeof SessionV1.Event.PartUpdated.Type)["data"]["part"] | unknown): Usage | undefined {
|
||||
if (typeof part !== "object" || part === null) return undefined
|
||||
const value = part as Record<string, unknown>
|
||||
|
|
@ -41,6 +49,22 @@ function usage(part: (typeof SessionV1.Event.PartUpdated.Type)["data"]["part"] |
|
|||
return { cost: value.cost as Usage["cost"], tokens: value.tokens as Usage["tokens"] }
|
||||
}
|
||||
|
||||
function addUsage(target: Usage, value: Usage) {
|
||||
target.cost += value.cost
|
||||
target.tokens.input += value.tokens.input
|
||||
target.tokens.output += value.tokens.output
|
||||
target.tokens.reasoning += value.tokens.reasoning
|
||||
target.tokens.cache.read += value.tokens.cache.read
|
||||
target.tokens.cache.write += value.tokens.cache.write
|
||||
}
|
||||
|
||||
function messageUsage(row: typeof SessionMessageTable.$inferSelect): Usage | undefined {
|
||||
if (row.type !== "assistant") return undefined
|
||||
const message = decodeMessage({ ...row.data, id: row.id, type: row.type })
|
||||
if (message.type !== "assistant" || message.cost === undefined || message.tokens === undefined) return undefined
|
||||
return { cost: message.cost, tokens: message.tokens }
|
||||
}
|
||||
|
||||
function sessionRow(info: SessionV1.SessionInfo): typeof SessionTable.$inferInsert {
|
||||
return {
|
||||
id: info.id,
|
||||
|
|
@ -109,7 +133,150 @@ function applyUsage(
|
|||
.pipe(Effect.orDie)
|
||||
}
|
||||
|
||||
function run(db: DatabaseService, event: SessionEvent.Event) {
|
||||
const projectFork = Effect.fn("SessionProjector.projectFork")(function* (
|
||||
db: DatabaseService,
|
||||
event: typeof SessionEvent.Forked.Type,
|
||||
) {
|
||||
const parent = yield* db
|
||||
.select()
|
||||
.from(SessionTable)
|
||||
.where(eq(SessionTable.id, event.data.parentID))
|
||||
.get()
|
||||
.pipe(Effect.orDie)
|
||||
if (!parent) return yield* Effect.die(`Fork parent session not found: ${event.data.parentID}`)
|
||||
|
||||
const stored = yield* db
|
||||
.insert(SessionTable)
|
||||
.values({
|
||||
id: event.data.sessionID,
|
||||
parent_id: event.data.parentID,
|
||||
project_id: parent.project_id,
|
||||
workspace_id: parent.workspace_id,
|
||||
slug: event.data.slug,
|
||||
directory: parent.directory,
|
||||
path: parent.path,
|
||||
title: event.data.title,
|
||||
agent: event.data.agent,
|
||||
model: event.data.model,
|
||||
version: parent.version,
|
||||
cost: 0,
|
||||
tokens_input: 0,
|
||||
tokens_output: 0,
|
||||
tokens_reasoning: 0,
|
||||
tokens_cache_read: 0,
|
||||
tokens_cache_write: 0,
|
||||
time_created: DateTime.toEpochMillis(event.data.timestamp),
|
||||
time_updated: DateTime.toEpochMillis(event.data.timestamp),
|
||||
})
|
||||
.onConflictDoNothing()
|
||||
.returning({ sessionID: SessionTable.id })
|
||||
.get()
|
||||
.pipe(Effect.orDie)
|
||||
if (!stored) return yield* Effect.die(new SessionAlreadyProjected())
|
||||
|
||||
const usage = emptyUsage()
|
||||
let cursor = -1
|
||||
while (true) {
|
||||
const rows = yield* db
|
||||
.select()
|
||||
.from(SessionMessageTable)
|
||||
.where(
|
||||
and(
|
||||
eq(SessionMessageTable.session_id, event.data.parentID),
|
||||
gt(SessionMessageTable.seq, cursor),
|
||||
event.data.messageID === undefined ? undefined : lt(SessionMessageTable.seq, event.data.copiedSeq + 1),
|
||||
),
|
||||
)
|
||||
.orderBy(asc(SessionMessageTable.seq))
|
||||
.limit(ForkBatchSize)
|
||||
.all()
|
||||
.pipe(Effect.orDie)
|
||||
if (rows.length === 0) break
|
||||
|
||||
const idMap = new Map(rows.map((row) => [row.id, SessionMessage.ID.create()]))
|
||||
yield* db
|
||||
.insert(SessionMessageTable)
|
||||
.values(
|
||||
rows.map((row) => {
|
||||
const id = idMap.get(row.id)
|
||||
if (!id) throw new Error(`Fork message ID mapping missing: ${row.id}`)
|
||||
return {
|
||||
id,
|
||||
session_id: event.data.sessionID,
|
||||
type: row.type,
|
||||
seq: row.seq,
|
||||
time_created: row.time_created,
|
||||
time_updated: row.time_updated,
|
||||
data: row.type === "synthetic" ? { ...row.data, sessionID: event.data.sessionID } : row.data,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.run()
|
||||
.pipe(Effect.orDie)
|
||||
|
||||
const inputRows = yield* db
|
||||
.select()
|
||||
.from(SessionInputTable)
|
||||
.where(
|
||||
and(
|
||||
eq(SessionInputTable.session_id, event.data.parentID),
|
||||
inArray(
|
||||
SessionInputTable.id,
|
||||
rows.map((row) => row.id),
|
||||
),
|
||||
),
|
||||
)
|
||||
.all()
|
||||
.pipe(Effect.orDie)
|
||||
if (inputRows.length > 0) {
|
||||
yield* db
|
||||
.insert(SessionInputTable)
|
||||
.values(
|
||||
inputRows.flatMap((row) => {
|
||||
const id = idMap.get(row.id)
|
||||
return id
|
||||
? [
|
||||
{
|
||||
id,
|
||||
session_id: event.data.sessionID,
|
||||
prompt: row.prompt,
|
||||
delivery: row.delivery,
|
||||
admitted_seq: row.admitted_seq,
|
||||
promoted_seq: row.promoted_seq,
|
||||
time_created: row.time_created,
|
||||
},
|
||||
]
|
||||
: []
|
||||
}),
|
||||
)
|
||||
.run()
|
||||
.pipe(Effect.orDie)
|
||||
}
|
||||
|
||||
for (const row of rows) {
|
||||
const value = messageUsage(row)
|
||||
if (value) addUsage(usage, value)
|
||||
}
|
||||
cursor = rows.at(-1)!.seq
|
||||
}
|
||||
|
||||
yield* db
|
||||
.update(SessionTable)
|
||||
.set({
|
||||
cost: usage.cost,
|
||||
tokens_input: usage.tokens.input,
|
||||
tokens_output: usage.tokens.output,
|
||||
tokens_reasoning: usage.tokens.reasoning,
|
||||
tokens_cache_read: usage.tokens.cache.read,
|
||||
tokens_cache_write: usage.tokens.cache.write,
|
||||
})
|
||||
.where(eq(SessionTable.id, event.data.sessionID))
|
||||
.run()
|
||||
.pipe(Effect.orDie)
|
||||
if (event.data.copiedSeq > 0) yield* EventV2.reserveSequence(db, event.data.sessionID, event.data.copiedSeq)
|
||||
})
|
||||
|
||||
function run(db: DatabaseService, event: MessageEvent) {
|
||||
return Effect.gen(function* () {
|
||||
const decodeRow = (row: typeof SessionMessageTable.$inferSelect) =>
|
||||
decodeMessage({ ...row.data, id: row.id, type: row.type })
|
||||
|
|
@ -355,6 +522,7 @@ export const layer = Layer.effectDiscard(
|
|||
.run()
|
||||
.pipe(Effect.orDie),
|
||||
)
|
||||
yield* events.project(SessionEvent.Forked, (event) => projectFork(db, event))
|
||||
yield* events.project(SessionEvent.Prompted, (event) =>
|
||||
Effect.gen(function* () {
|
||||
if (event.durable === undefined) return yield* Effect.die("Durable Session event is missing aggregate sequence")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue