refactor(session): effectify SessionPrompt service (#19483)

This commit is contained in:
Kit Langton
2026-03-30 16:06:51 -04:00
committed by GitHub
parent fa95a61c4e
commit c5442d418d
20 changed files with 4439 additions and 2142 deletions

View File

@@ -15,7 +15,7 @@ import { Plugin } from "@/plugin"
import { Config } from "@/config/config"
import { NotFoundError } from "@/storage/db"
import { ModelID, ProviderID } from "@/provider/schema"
import { Cause, Effect, Exit, Layer, ServiceMap } from "effect"
import { Effect, Layer, ServiceMap } from "effect"
import { makeRuntime } from "@/effect/run-service"
import { isOverflow as overflow } from "./overflow"
@@ -45,7 +45,6 @@ export namespace SessionCompaction {
parentID: MessageID
messages: MessageV2.WithParts[]
sessionID: SessionID
abort: AbortSignal
auto: boolean
overflow?: boolean
}) => Effect.Effect<"continue" | "stop">
@@ -135,20 +134,28 @@ export namespace SessionCompaction {
parentID: MessageID
messages: MessageV2.WithParts[]
sessionID: SessionID
abort: AbortSignal
auto: boolean
overflow?: boolean
}) {
const userMessage = input.messages.findLast((m) => m.info.id === input.parentID)!.info as MessageV2.User
const parent = input.messages.findLast((m) => m.info.id === input.parentID)
if (!parent || parent.info.role !== "user") {
throw new Error(`Compaction parent must be a user message: ${input.parentID}`)
}
const userMessage = parent.info
let messages = input.messages
let replay: MessageV2.WithParts | undefined
let replay:
| {
info: MessageV2.User
parts: MessageV2.Part[]
}
| undefined
if (input.overflow) {
const idx = input.messages.findIndex((m) => m.info.id === input.parentID)
for (let i = idx - 1; i >= 0; i--) {
const msg = input.messages[i]
if (msg.info.role === "user" && !msg.parts.some((p) => p.type === "compaction")) {
replay = msg
replay = { info: msg.info, parts: msg.parts }
messages = input.messages.slice(0, i)
break
}
@@ -206,7 +213,7 @@ When constructing the summary, try to stick to this template:
const msgs = structuredClone(messages)
yield* plugin.trigger("experimental.chat.messages.transform", {}, { messages: msgs })
const modelMessages = yield* Effect.promise(() => MessageV2.toModelMessages(msgs, model, { stripMedia: true }))
const msg = (yield* session.updateMessage({
const msg: MessageV2.Assistant = {
id: MessageID.ascending(),
role: "assistant",
parentID: input.parentID,
@@ -231,25 +238,17 @@ When constructing the summary, try to stick to this template:
time: {
created: Date.now(),
},
})) as MessageV2.Assistant
}
yield* session.updateMessage(msg)
const processor = yield* processors.create({
assistantMessage: msg,
sessionID: input.sessionID,
model,
abort: input.abort,
})
const cancel = Effect.fn("SessionCompaction.cancel")(function* () {
if (!input.abort.aborted || msg.time.completed) return
msg.error = msg.error ?? new MessageV2.AbortedError({ message: "Aborted" }).toObject()
msg.finish = msg.finish ?? "error"
msg.time.completed = Date.now()
yield* session.updateMessage(msg)
})
const result = yield* processor
.process({
user: userMessage,
agent,
abort: input.abort,
sessionID: input.sessionID,
tools: {},
system: [],
@@ -262,7 +261,7 @@ When constructing the summary, try to stick to this template:
],
model,
})
.pipe(Effect.ensuring(cancel()))
.pipe(Effect.onInterrupt(() => processor.abort()))
if (result === "compact") {
processor.message.error = new MessageV2.ContextOverflowError({
@@ -277,7 +276,7 @@ When constructing the summary, try to stick to this template:
if (result === "continue" && input.auto) {
if (replay) {
const original = replay.info as MessageV2.User
const original = replay.info
const replayMsg = yield* session.updateMessage({
id: MessageID.ascending(),
role: "user",
@@ -386,7 +385,7 @@ When constructing the summary, try to stick to this template:
),
)
const { runPromise, runPromiseExit } = makeRuntime(Service, defaultLayer)
const { runPromise } = makeRuntime(Service, defaultLayer)
export async function isOverflow(input: { tokens: MessageV2.Assistant["tokens"]; model: Provider.Model }) {
return runPromise((svc) => svc.isOverflow(input))
@@ -396,21 +395,16 @@ When constructing the summary, try to stick to this template:
return runPromise((svc) => svc.prune(input))
}
export async function process(input: {
parentID: MessageID
messages: MessageV2.WithParts[]
sessionID: SessionID
abort: AbortSignal
auto: boolean
overflow?: boolean
}) {
const exit = await runPromiseExit((svc) => svc.process(input), { signal: input.abort })
if (Exit.isFailure(exit)) {
if (Cause.hasInterrupts(exit.cause) && input.abort.aborted) return "stop"
throw Cause.squash(exit.cause)
}
return exit.value
}
export const process = fn(
z.object({
parentID: MessageID.zod,
messages: z.custom<MessageV2.WithParts[]>(),
sessionID: SessionID.zod,
auto: z.boolean(),
overflow: z.boolean().optional(),
}),
(input) => runPromise((svc) => svc.process(input)),
)
export const create = fn(
z.object({

View File

@@ -334,14 +334,14 @@ export namespace Session {
readonly messages: (input: { sessionID: SessionID; limit?: number }) => Effect.Effect<MessageV2.WithParts[]>
readonly children: (parentID: SessionID) => Effect.Effect<Info[]>
readonly remove: (sessionID: SessionID) => Effect.Effect<void>
readonly updateMessage: (msg: MessageV2.Info) => Effect.Effect<MessageV2.Info>
readonly updateMessage: <T extends MessageV2.Info>(msg: T) => Effect.Effect<T>
readonly removeMessage: (input: { sessionID: SessionID; messageID: MessageID }) => Effect.Effect<MessageID>
readonly removePart: (input: {
sessionID: SessionID
messageID: MessageID
partID: PartID
}) => Effect.Effect<PartID>
readonly updatePart: (part: MessageV2.Part) => Effect.Effect<MessageV2.Part>
readonly updatePart: <T extends MessageV2.Part>(part: T) => Effect.Effect<T>
readonly updatePartDelta: (input: {
sessionID: SessionID
messageID: MessageID
@@ -469,26 +469,23 @@ export namespace Session {
}
})
const updateMessage = Effect.fn("Session.updateMessage")(function* (msg: MessageV2.Info) {
yield* Effect.sync(() =>
SyncEvent.run(MessageV2.Event.Updated, {
sessionID: msg.sessionID,
info: msg,
}),
)
return msg
})
const updateMessage = <T extends MessageV2.Info>(msg: T): Effect.Effect<T> =>
Effect.gen(function* () {
yield* Effect.sync(() => SyncEvent.run(MessageV2.Event.Updated, { sessionID: msg.sessionID, info: msg }))
return msg
}).pipe(Effect.withSpan("Session.updateMessage"))
const updatePart = Effect.fn("Session.updatePart")(function* (part: MessageV2.Part) {
yield* Effect.sync(() =>
SyncEvent.run(MessageV2.Event.PartUpdated, {
sessionID: part.sessionID,
part: structuredClone(part),
time: Date.now(),
}),
)
return part
})
const updatePart = <T extends MessageV2.Part>(part: T): Effect.Effect<T> =>
Effect.gen(function* () {
yield* Effect.sync(() =>
SyncEvent.run(MessageV2.Event.PartUpdated, {
sessionID: part.sessionID,
part: structuredClone(part),
time: Date.now(),
}),
)
return part
}).pipe(Effect.withSpan("Session.updatePart"))
const create = Effect.fn("Session.create")(function* (input?: {
parentID?: SessionID
@@ -851,7 +848,9 @@ export namespace Session {
export const children = fn(SessionID.zod, (id) => runPromise((svc) => svc.children(id)))
export const remove = fn(SessionID.zod, (id) => runPromise((svc) => svc.remove(id)))
export const updateMessage = fn(MessageV2.Info, (msg) => runPromise((svc) => svc.updateMessage(msg)))
export async function updateMessage<T extends MessageV2.Info>(msg: T): Promise<T> {
return runPromise((svc) => svc.updateMessage(MessageV2.Info.parse(msg) as T))
}
export const removeMessage = fn(z.object({ sessionID: SessionID.zod, messageID: MessageID.zod }), (input) =>
runPromise((svc) => svc.removeMessage(input)),
@@ -862,7 +861,9 @@ export namespace Session {
(input) => runPromise((svc) => svc.removePart(input)),
)
export const updatePart = fn(MessageV2.Part, (part) => runPromise((svc) => svc.updatePart(part)))
export async function updatePart<T extends MessageV2.Part>(part: T): Promise<T> {
return runPromise((svc) => svc.updatePart(MessageV2.Part.parse(part) as T))
}
export const updatePartDelta = fn(
z.object({

View File

@@ -1,6 +1,7 @@
import { Provider } from "@/provider/provider"
import { Log } from "@/util/log"
import { Effect, Layer, ServiceMap } from "effect"
import { Cause, Effect, Layer, Record, ServiceMap } from "effect"
import * as Queue from "effect/Queue"
import * as Stream from "effect/Stream"
import { streamText, wrapLanguageModel, type ModelMessage, type Tool, tool, jsonSchema } from "ai"
import { mergeDeep, pipe } from "remeda"
@@ -28,7 +29,6 @@ export namespace LLM {
agent: Agent.Info
permission?: Permission.Ruleset
system: string[]
abort: AbortSignal
messages: ModelMessage[]
small?: boolean
tools: Record<string, Tool>
@@ -36,6 +36,10 @@ export namespace LLM {
toolChoice?: "auto" | "required" | "none"
}
export type StreamRequest = StreamInput & {
abort: AbortSignal
}
export type Event = Awaited<ReturnType<typeof stream>>["fullStream"] extends AsyncIterable<infer T> ? T : never
export interface Interface {
@@ -49,15 +53,32 @@ export namespace LLM {
Effect.gen(function* () {
return Service.of({
stream(input) {
return Stream.unwrap(
Effect.promise(() => LLM.stream(input)).pipe(
Effect.map((result) =>
Stream.fromAsyncIterable(result.fullStream, (err) => err).pipe(
Stream.mapEffect((event) => Effect.succeed(event)),
),
),
const stream: Stream.Stream<Event, unknown> = Stream.scoped(
Stream.unwrap(
Effect.gen(function* () {
const ctrl = yield* Effect.acquireRelease(
Effect.sync(() => new AbortController()),
(ctrl) => Effect.sync(() => ctrl.abort()),
)
const queue = yield* Queue.unbounded<Event, unknown | Cause.Done>()
yield* Effect.promise(async () => {
const result = await LLM.stream({ ...input, abort: ctrl.signal })
for await (const event of result.fullStream) {
if (!Queue.offerUnsafe(queue, event)) break
}
Queue.endUnsafe(queue)
}).pipe(
Effect.catchCause((cause) => Effect.sync(() => void Queue.failCauseUnsafe(queue, cause))),
Effect.onInterrupt(() => Effect.sync(() => ctrl.abort())),
Effect.forkScoped,
)
return Stream.fromQueue(queue)
}),
),
)
return stream
},
})
}),
@@ -65,7 +86,7 @@ export namespace LLM {
export const defaultLayer = layer
export async function stream(input: StreamInput) {
export async function stream(input: StreamRequest) {
const l = log
.clone()
.tag("providerID", input.model.providerID)
@@ -322,17 +343,12 @@ export namespace LLM {
})
}
async function resolveTools(input: Pick<StreamInput, "tools" | "agent" | "permission" | "user">) {
function resolveTools(input: Pick<StreamInput, "tools" | "agent" | "permission" | "user">) {
const disabled = Permission.disabled(
Object.keys(input.tools),
Permission.merge(input.agent.permission, input.permission ?? []),
)
for (const tool of Object.keys(input.tools)) {
if (input.user.tools?.[tool] === false || disabled.has(tool)) {
delete input.tools[tool]
}
}
return input.tools
return Record.filter(input.tools, (_, k) => input.user.tools?.[k] !== false && !disabled.has(k))
}
// Check if messages contain any tool-call content

View File

@@ -1,8 +1,7 @@
import { Cause, Effect, Exit, Layer, ServiceMap } from "effect"
import { Cause, Effect, Layer, ServiceMap } from "effect"
import * as Stream from "effect/Stream"
import { Agent } from "@/agent/agent"
import { Bus } from "@/bus"
import { makeRuntime } from "@/effect/run-service"
import { Config } from "@/config/config"
import { Permission } from "@/permission"
import { Plugin } from "@/plugin"
@@ -35,17 +34,10 @@ export namespace SessionProcessor {
readonly process: (streamInput: LLM.StreamInput) => Effect.Effect<Result>
}
export interface Info {
readonly message: MessageV2.Assistant
readonly partFromToolCall: (toolCallID: string) => MessageV2.ToolPart | undefined
readonly process: (streamInput: LLM.StreamInput) => Promise<Result>
}
type Input = {
assistantMessage: MessageV2.Assistant
sessionID: SessionID
model: Provider.Model
abort: AbortSignal
}
export interface Interface {
@@ -96,7 +88,6 @@ export namespace SessionProcessor {
assistantMessage: input.assistantMessage,
sessionID: input.sessionID,
model: input.model,
abort: input.abort,
toolcalls: {},
shouldBreak: false,
snapshot: undefined,
@@ -105,11 +96,12 @@ export namespace SessionProcessor {
currentText: undefined,
reasoningMap: {},
}
let aborted = false
const parse = (e: unknown) =>
MessageV2.fromError(e, {
providerID: input.model.providerID,
aborted: input.abort.aborted,
aborted,
})
const handleEvent = Effect.fn("SessionProcessor.handleEvent")(function* (value: StreamEvent) {
@@ -155,7 +147,10 @@ export namespace SessionProcessor {
return
case "tool-input-start":
ctx.toolcalls[value.id] = (yield* session.updatePart({
if (ctx.assistantMessage.summary) {
throw new Error(`Tool call not allowed while generating summary: ${value.toolName}`)
}
ctx.toolcalls[value.id] = yield* session.updatePart({
id: ctx.toolcalls[value.id]?.id ?? PartID.ascending(),
messageID: ctx.assistantMessage.id,
sessionID: ctx.assistantMessage.sessionID,
@@ -163,7 +158,7 @@ export namespace SessionProcessor {
tool: value.toolName,
callID: value.id,
state: { status: "pending", input: {}, raw: "" },
})) as MessageV2.ToolPart
} satisfies MessageV2.ToolPart)
return
case "tool-input-delta":
@@ -173,14 +168,17 @@ export namespace SessionProcessor {
return
case "tool-call": {
if (ctx.assistantMessage.summary) {
throw new Error(`Tool call not allowed while generating summary: ${value.toolName}`)
}
const match = ctx.toolcalls[value.toolCallId]
if (!match) return
ctx.toolcalls[value.toolCallId] = (yield* session.updatePart({
ctx.toolcalls[value.toolCallId] = yield* session.updatePart({
...match,
tool: value.toolName,
state: { status: "running", input: value.input, time: { start: Date.now() } },
metadata: value.providerMetadata,
})) as MessageV2.ToolPart
} satisfies MessageV2.ToolPart)
const parts = yield* Effect.promise(() => MessageV2.parts(ctx.assistantMessage.id))
const recentParts = parts.slice(-DOOM_LOOP_THRESHOLD)
@@ -414,7 +412,7 @@ export namespace SessionProcessor {
})
const halt = Effect.fn("SessionProcessor.halt")(function* (e: unknown) {
log.error("process", { error: e, stack: JSON.stringify((e as any)?.stack) })
log.error("process", { error: e, stack: e instanceof Error ? e.stack : undefined })
const error = parse(e)
if (MessageV2.ContextOverflowError.isInstance(error)) {
ctx.needsCompaction = true
@@ -429,59 +427,6 @@ export namespace SessionProcessor {
yield* status.set(ctx.sessionID, { type: "idle" })
})
const process = Effect.fn("SessionProcessor.process")(function* (streamInput: LLM.StreamInput) {
log.info("process")
ctx.needsCompaction = false
ctx.shouldBreak = (yield* config.get()).experimental?.continue_loop_on_deny !== true
yield* Effect.gen(function* () {
ctx.currentText = undefined
ctx.reasoningMap = {}
const stream = llm.stream(streamInput)
yield* stream.pipe(
Stream.tap((event) =>
Effect.gen(function* () {
input.abort.throwIfAborted()
yield* handleEvent(event)
}),
),
Stream.takeUntil(() => ctx.needsCompaction),
Stream.runDrain,
)
}).pipe(
Effect.catchCauseIf(
(cause) => !Cause.hasInterruptsOnly(cause),
(cause) => Effect.fail(Cause.squash(cause)),
),
Effect.retry(
SessionRetry.policy({
parse,
set: (info) =>
status.set(ctx.sessionID, {
type: "retry",
attempt: info.attempt,
message: info.message,
next: info.next,
}),
}),
),
Effect.catchCause((cause) =>
Cause.hasInterruptsOnly(cause)
? halt(new DOMException("Aborted", "AbortError"))
: halt(Cause.squash(cause)),
),
Effect.ensuring(cleanup()),
)
if (input.abort.aborted && !ctx.assistantMessage.error) {
yield* abort()
}
if (ctx.needsCompaction) return "compact"
if (ctx.blocked || ctx.assistantMessage.error || input.abort.aborted) return "stop"
return "continue"
})
const abort = Effect.fn("SessionProcessor.abort")(() =>
Effect.gen(function* () {
if (!ctx.assistantMessage.error) {
@@ -495,6 +440,53 @@ export namespace SessionProcessor {
}),
)
const process = Effect.fn("SessionProcessor.process")(function* (streamInput: LLM.StreamInput) {
log.info("process")
ctx.needsCompaction = false
ctx.shouldBreak = (yield* config.get()).experimental?.continue_loop_on_deny !== true
return yield* Effect.gen(function* () {
yield* Effect.gen(function* () {
ctx.currentText = undefined
ctx.reasoningMap = {}
const stream = llm.stream(streamInput)
yield* stream.pipe(
Stream.tap((event) => handleEvent(event)),
Stream.takeUntil(() => ctx.needsCompaction),
Stream.runDrain,
)
}).pipe(
Effect.onInterrupt(() => Effect.sync(() => void (aborted = true))),
Effect.catchCauseIf(
(cause) => !Cause.hasInterruptsOnly(cause),
(cause) => Effect.fail(Cause.squash(cause)),
),
Effect.retry(
SessionRetry.policy({
parse,
set: (info) =>
status.set(ctx.sessionID, {
type: "retry",
attempt: info.attempt,
message: info.message,
next: info.next,
}),
}),
),
Effect.catch(halt),
Effect.ensuring(cleanup()),
)
if (aborted && !ctx.assistantMessage.error) {
yield* abort()
}
if (ctx.needsCompaction) return "compact"
if (ctx.blocked || ctx.assistantMessage.error || aborted) return "stop"
return "continue"
}).pipe(Effect.onInterrupt(() => abort().pipe(Effect.asVoid)))
})
return {
get message() {
return ctx.assistantMessage
@@ -526,29 +518,4 @@ export namespace SessionProcessor {
),
),
)
const { runPromise } = makeRuntime(Service, defaultLayer)
export async function create(input: Input): Promise<Info> {
const hit = await runPromise((svc) => svc.create(input))
return {
get message() {
return hit.message
},
partFromToolCall(toolCallID: string) {
return hit.partFromToolCall(toolCallID)
},
async process(streamInput: LLM.StreamInput) {
const exit = await Effect.runPromiseExit(hit.process(streamInput), { signal: input.abort })
if (Exit.isFailure(exit)) {
if (Cause.hasInterrupts(exit.cause) && input.abort.aborted) {
await Effect.runPromise(hit.abort())
return "stop"
}
throw Cause.squash(exit.cause)
}
return exit.value
},
}
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -21,7 +21,7 @@ export namespace SessionRevert {
export type RevertInput = z.infer<typeof RevertInput>
export async function revert(input: RevertInput) {
SessionPrompt.assertNotBusy(input.sessionID)
await SessionPrompt.assertNotBusy(input.sessionID)
const all = await Session.messages({ sessionID: input.sessionID })
let lastUser: MessageV2.User | undefined
const session = await Session.get(input.sessionID)
@@ -80,7 +80,7 @@ export namespace SessionRevert {
export async function unrevert(input: { sessionID: SessionID }) {
log.info("unreverting", input)
SessionPrompt.assertNotBusy(input.sessionID)
await SessionPrompt.assertNotBusy(input.sessionID)
const session = await Session.get(input.sessionID)
if (!session.revert) return session
if (session.revert.snapshot) await Snapshot.restore(session.revert.snapshot)

View File

@@ -3,7 +3,6 @@ import z from "zod"
import { Session } from "."
import { MessageV2 } from "./message-v2"
import { Identifier } from "@/id/id"
import { SessionID, MessageID } from "./schema"
import { Snapshot } from "@/snapshot"
@@ -110,8 +109,8 @@ export namespace SessionSummary {
(m) => m.info.id === input.messageID || (m.info.role === "assistant" && m.info.parentID === input.messageID),
)
const msgWithParts = messages.find((m) => m.info.id === input.messageID)
if (!msgWithParts) return
const userMsg = msgWithParts.info as MessageV2.User
if (!msgWithParts || msgWithParts.info.role !== "user") return
const userMsg = msgWithParts.info
const diffs = await computeDiff({ messages })
userMsg.summary = {
...userMsg.summary,