refactor(effect): extract session run state service (#21744)

This commit is contained in:
Kit Langton
2026-04-09 16:03:40 -04:00
committed by GitHub
parent 3199383eef
commit 10441efad1
9 changed files with 184 additions and 111 deletions

View File

@@ -20,7 +20,6 @@ import PROMPT_PLAN from "../session/prompt/plan.txt"
import BUILD_SWITCH from "../session/prompt/build-switch.txt"
import MAX_STEPS from "../session/prompt/max-steps.txt"
import { ToolRegistry } from "../tool/registry"
import { Runner } from "@/effect/runner"
import { MCP } from "../mcp"
import { LSP } from "../lsp"
import { FileTime } from "../file/time"
@@ -48,6 +47,7 @@ import { Cause, Effect, Exit, Layer, Option, Scope, ServiceMap } from "effect"
import { InstanceState } from "@/effect/instance-state"
import { makeRuntime } from "@/effect/run-service"
import { TaskTool } from "@/tool/task"
import { SessionRunState } from "./run-state"
// @ts-ignore
globalThis.AI_SDK_LOG_WARNINGS = false
@@ -66,7 +66,6 @@ export namespace SessionPrompt {
const log = Log.create({ service: "session.prompt" })
export interface Interface {
readonly assertNotBusy: (sessionID: SessionID) => Effect.Effect<void, Session.BusyError>
readonly cancel: (sessionID: SessionID) => Effect.Effect<void>
readonly prompt: (input: PromptInput) => Effect.Effect<MessageV2.WithParts>
readonly loop: (input: z.infer<typeof LoopInput>) => Effect.Effect<MessageV2.WithParts>
@@ -99,55 +98,11 @@ export namespace SessionPrompt {
const spawner = yield* ChildProcessSpawner.ChildProcessSpawner
const scope = yield* Scope.Scope
const instruction = yield* Instruction.Service
const state = yield* InstanceState.make(
Effect.fn("SessionPrompt.state")(function* () {
const runners = new Map<string, Runner<MessageV2.WithParts>>()
yield* Effect.addFinalizer(
Effect.fnUntraced(function* () {
yield* Effect.forEach(runners.values(), (r) => r.cancel, { concurrency: "unbounded", discard: true })
runners.clear()
}),
)
return { runners }
}),
)
const getRunner = (runners: Map<string, Runner<MessageV2.WithParts>>, sessionID: SessionID) => {
const existing = runners.get(sessionID)
if (existing) return existing
const runner = Runner.make<MessageV2.WithParts>(scope, {
onIdle: Effect.gen(function* () {
runners.delete(sessionID)
yield* status.set(sessionID, { type: "idle" })
}),
onBusy: status.set(sessionID, { type: "busy" }),
onInterrupt: lastAssistant(sessionID),
busy: () => {
throw new Session.BusyError(sessionID)
},
})
runners.set(sessionID, runner)
return runner
}
const assertNotBusy: (sessionID: SessionID) => Effect.Effect<void, Session.BusyError> = Effect.fn(
"SessionPrompt.assertNotBusy",
)(function* (sessionID: SessionID) {
const s = yield* InstanceState.get(state)
const runner = s.runners.get(sessionID)
if (runner?.busy) throw new Session.BusyError(sessionID)
})
const state = yield* SessionRunState.Service
const cancel = Effect.fn("SessionPrompt.cancel")(function* (sessionID: SessionID) {
log.info("cancel", { sessionID })
const s = yield* InstanceState.get(state)
const runner = s.runners.get(sessionID)
if (!runner || !runner.busy) {
yield* status.set(sessionID, { type: "idle" })
return
}
yield* runner.cancel
yield* state.cancel(sessionID)
})
const resolvePromptParts = Effect.fn("SessionPrompt.resolvePromptParts")(function* (template: string) {
@@ -1574,16 +1529,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the
const loop: (input: z.infer<typeof LoopInput>) => Effect.Effect<MessageV2.WithParts> = Effect.fn(
"SessionPrompt.loop",
)(function* (input: z.infer<typeof LoopInput>) {
const s = yield* InstanceState.get(state)
const runner = getRunner(s.runners, input.sessionID)
return yield* runner.ensureRunning(runLoop(input.sessionID))
return yield* state.ensureRunning(input.sessionID, lastAssistant(input.sessionID), runLoop(input.sessionID))
})
const shell: (input: ShellInput) => Effect.Effect<MessageV2.WithParts> = Effect.fn("SessionPrompt.shell")(
function* (input: ShellInput) {
const s = yield* InstanceState.get(state)
const runner = getRunner(s.runners, input.sessionID)
return yield* runner.startShell(shellImpl(input))
return yield* state.startShell(input.sessionID, lastAssistant(input.sessionID), shellImpl(input))
},
)
@@ -1704,7 +1655,6 @@ NOTE: At any point in time through this workflow you should feel free to ask the
})
return Service.of({
assertNotBusy,
cancel,
prompt,
loop,
@@ -1718,6 +1668,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
const defaultLayer = Layer.unwrap(
Effect.sync(() =>
layer.pipe(
Layer.provide(SessionRunState.layer),
Layer.provide(SessionStatus.layer),
Layer.provide(SessionCompaction.defaultLayer),
Layer.provide(SessionProcessor.defaultLayer),
@@ -1741,10 +1692,6 @@ NOTE: At any point in time through this workflow you should feel free to ask the
)
const { runPromise } = makeRuntime(Service, defaultLayer)
export async function assertNotBusy(sessionID: SessionID) {
return runPromise((svc) => svc.assertNotBusy(SessionID.zod.parse(sessionID)))
}
export const PromptInput = z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod.optional(),

View File

@@ -9,8 +9,9 @@ import { Log } from "../util/log"
import { Session } from "."
import { MessageV2 } from "./message-v2"
import { SessionID, MessageID, PartID } from "./schema"
import { SessionPrompt } from "./prompt"
import { SessionRunState } from "./run-state"
import { SessionSummary } from "./summary"
import { SessionStatus } from "./status"
export namespace SessionRevert {
const log = Log.create({ service: "session.revert" })
@@ -38,9 +39,10 @@ export namespace SessionRevert {
const storage = yield* Storage.Service
const bus = yield* Bus.Service
const summary = yield* SessionSummary.Service
const state = yield* SessionRunState.Service
const revert = Effect.fn("SessionRevert.revert")(function* (input: RevertInput) {
yield* Effect.promise(() => SessionPrompt.assertNotBusy(input.sessionID))
yield* state.assertNotBusy(input.sessionID)
const all = yield* sessions.messages({ sessionID: input.sessionID })
let lastUser: MessageV2.User | undefined
const session = yield* sessions.get(input.sessionID)
@@ -93,7 +95,7 @@ export namespace SessionRevert {
const unrevert = Effect.fn("SessionRevert.unrevert")(function* (input: { sessionID: SessionID }) {
log.info("unreverting", input)
yield* Effect.promise(() => SessionPrompt.assertNotBusy(input.sessionID))
yield* state.assertNotBusy(input.sessionID)
const session = yield* sessions.get(input.sessionID)
if (!session.revert) return session
if (session.revert.snapshot) yield* snap.restore(session.revert!.snapshot!)
@@ -151,6 +153,8 @@ export namespace SessionRevert {
export const defaultLayer = Layer.unwrap(
Effect.sync(() =>
layer.pipe(
Layer.provide(SessionRunState.layer),
Layer.provide(SessionStatus.layer),
Layer.provide(Session.defaultLayer),
Layer.provide(Snapshot.defaultLayer),
Layer.provide(Storage.defaultLayer),

View File

@@ -0,0 +1,114 @@
import { InstanceState } from "@/effect/instance-state"
import { Runner } from "@/effect/runner"
import { makeRuntime } from "@/effect/run-service"
import { Effect, Layer, Scope, ServiceMap } from "effect"
import { Session } from "."
import { MessageV2 } from "./message-v2"
import { SessionID } from "./schema"
import { SessionStatus } from "./status"
export namespace SessionRunState {
export interface Interface {
readonly assertNotBusy: (sessionID: SessionID) => Effect.Effect<void>
readonly cancel: (sessionID: SessionID) => Effect.Effect<void>
readonly ensureRunning: (
sessionID: SessionID,
onInterrupt: Effect.Effect<MessageV2.WithParts>,
work: Effect.Effect<MessageV2.WithParts>,
) => Effect.Effect<MessageV2.WithParts>
readonly startShell: (
sessionID: SessionID,
onInterrupt: Effect.Effect<MessageV2.WithParts>,
work: Effect.Effect<MessageV2.WithParts>,
) => Effect.Effect<MessageV2.WithParts>
}
export class Service extends ServiceMap.Service<Service, Interface>()("@opencode/SessionRunState") {}
export const layer = Layer.effect(
Service,
Effect.gen(function* () {
const status = yield* SessionStatus.Service
const state = yield* InstanceState.make(
Effect.fn("SessionRunState.state")(function* () {
const scope = yield* Scope.Scope
const runners = new Map<SessionID, Runner<MessageV2.WithParts>>()
yield* Effect.addFinalizer(
Effect.fnUntraced(function* () {
yield* Effect.forEach(runners.values(), (runner) => runner.cancel, {
concurrency: "unbounded",
discard: true,
})
runners.clear()
}),
)
return { runners, scope }
}),
)
const runner = Effect.fn("SessionRunState.runner")(function* (
sessionID: SessionID,
onInterrupt: Effect.Effect<MessageV2.WithParts>,
) {
const data = yield* InstanceState.get(state)
const existing = data.runners.get(sessionID)
if (existing) return existing
const next = Runner.make<MessageV2.WithParts>(data.scope, {
onIdle: Effect.gen(function* () {
data.runners.delete(sessionID)
yield* status.set(sessionID, { type: "idle" })
}),
onBusy: status.set(sessionID, { type: "busy" }),
onInterrupt,
busy: () => {
throw new Session.BusyError(sessionID)
},
})
data.runners.set(sessionID, next)
return next
})
const assertNotBusy = Effect.fn("SessionRunState.assertNotBusy")(function* (sessionID: SessionID) {
const data = yield* InstanceState.get(state)
const existing = data.runners.get(sessionID)
if (existing?.busy) throw new Session.BusyError(sessionID)
})
const cancel = Effect.fn("SessionRunState.cancel")(function* (sessionID: SessionID) {
const data = yield* InstanceState.get(state)
const existing = data.runners.get(sessionID)
if (!existing || !existing.busy) {
yield* status.set(sessionID, { type: "idle" })
return
}
yield* existing.cancel
})
const ensureRunning = Effect.fn("SessionRunState.ensureRunning")(function* (
sessionID: SessionID,
onInterrupt: Effect.Effect<MessageV2.WithParts>,
work: Effect.Effect<MessageV2.WithParts>,
) {
return yield* (yield* runner(sessionID, onInterrupt)).ensureRunning(work)
})
const startShell = Effect.fn("SessionRunState.startShell")(function* (
sessionID: SessionID,
onInterrupt: Effect.Effect<MessageV2.WithParts>,
work: Effect.Effect<MessageV2.WithParts>,
) {
return yield* (yield* runner(sessionID, onInterrupt)).startShell(work)
})
return Service.of({ assertNotBusy, cancel, ensureRunning, startShell })
}),
)
export const defaultLayer = layer.pipe(Layer.provide(SessionStatus.defaultLayer))
const { runPromise } = makeRuntime(Service, defaultLayer)
export async function assertNotBusy(sessionID: SessionID) {
return runPromise((svc) => svc.assertNotBusy(sessionID))
}
}

View File

@@ -85,7 +85,7 @@ export namespace SessionStatus {
}),
)
const defaultLayer = layer.pipe(Layer.provide(Bus.layer))
export const defaultLayer = layer.pipe(Layer.provide(Bus.layer))
const { runPromise } = makeRuntime(Service, defaultLayer)
export async function get(sessionID: SessionID) {