diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts index 2dabfb2d6fba..72a91db0aee8 100644 --- a/packages/core/src/session.ts +++ b/packages/core/src/session.ts @@ -359,29 +359,33 @@ const layer = Layer.effect( }), prompt: Effect.fn("V2Session.prompt")((input) => Effect.uninterruptible( - Effect.gen(function* () { - yield* result.get(input.sessionID) - const prompt = resolvePrompt(input.prompt) - const messageID = input.id ?? SessionMessage.ID.create() - const delivery = input.delivery ?? "steer" - const expected = { sessionID: input.sessionID, messageID, prompt, delivery } - const admitted = yield* SessionInput.admit(db, events, { - id: messageID, - sessionID: input.sessionID, - prompt, - delivery, - }).pipe( - Effect.catchDefect((defect) => - defect instanceof SessionInput.LifecycleConflict - ? new PromptConflictError({ sessionID: input.sessionID, messageID }) - : Effect.die(defect), - ), - ) - if (!SessionInput.equivalent(admitted, expected)) - return yield* new PromptConflictError({ sessionID: input.sessionID, messageID }) - if (input.resume !== false) yield* execution.wake(admitted.sessionID) - return admitted - }), + execution.withLock(input.sessionID)( + Effect.gen(function* () { + const session = yield* result.get(input.sessionID) + if (session.revert) + yield* SessionRevert.commit(session).pipe(Effect.provideService(EventV2.Service, events)) + const prompt = resolvePrompt(input.prompt) + const messageID = input.id ?? SessionMessage.ID.create() + const delivery = input.delivery ?? "steer" + const expected = { sessionID: input.sessionID, messageID, prompt, delivery } + const admitted = yield* SessionInput.admit(db, events, { + id: messageID, + sessionID: input.sessionID, + prompt, + delivery, + }).pipe( + Effect.catchDefect((defect) => + defect instanceof SessionInput.LifecycleConflict + ? new PromptConflictError({ sessionID: input.sessionID, messageID }) + : Effect.die(defect), + ), + ) + if (!SessionInput.equivalent(admitted, expected)) + return yield* new PromptConflictError({ sessionID: input.sessionID, messageID }) + if (input.resume !== false) yield* execution.wake(admitted.sessionID) + return admitted + }), + ), ), ), shell: Effect.fn("V2Session.shell")(function* () { @@ -424,31 +428,51 @@ const layer = Layer.effect( }), active: execution.active, resume: Effect.fn("V2Session.resume")(function* (sessionID) { - yield* result.get(sessionID) - yield* execution.resume(sessionID) + yield* execution.withLock(sessionID)( + Effect.gen(function* () { + const session = yield* result.get(sessionID) + if (session.revert) return + yield* execution.resume(sessionID) + }), + ) }), interrupt: Effect.fn("V2Session.interrupt")((sessionID) => Effect.uninterruptible(execution.interrupt(sessionID)), ), revert: { stage: Effect.fn("V2Session.revert.stage")(function* (input) { - const session = yield* result.get(input.sessionID) - return yield* SessionRevert.stage({ session, messageID: input.messageID, files: input.files }).pipe( - Effect.provideService(Database.Service, database), - Effect.provideService(EventV2.Service, events), - Effect.provide(locations.get(session.location)), + return yield* execution.exclusive( + input.sessionID, + Effect.gen(function* () { + const session = yield* result.get(input.sessionID) + return yield* SessionRevert.stage({ session, messageID: input.messageID, files: input.files }).pipe( + Effect.provideService(Database.Service, database), + Effect.provideService(EventV2.Service, events), + Effect.provide(locations.get(session.location)), + ) + }), ) }), clear: Effect.fn("V2Session.revert.clear")(function* (sessionID) { - const session = yield* result.get(sessionID) - yield* SessionRevert.clear(session).pipe( - Effect.provideService(EventV2.Service, events), - Effect.provide(locations.get(session.location)), + yield* execution.exclusive( + sessionID, + Effect.gen(function* () { + const session = yield* result.get(sessionID) + yield* SessionRevert.clear(session).pipe( + Effect.provideService(EventV2.Service, events), + Effect.provide(locations.get(session.location)), + ) + }), ) }), commit: Effect.fn("V2Session.revert.commit")(function* (sessionID) { - const session = yield* result.get(sessionID) - yield* SessionRevert.commit(session).pipe(Effect.provideService(EventV2.Service, events)) + yield* execution.exclusive( + sessionID, + Effect.gen(function* () { + const session = yield* result.get(sessionID) + yield* SessionRevert.commit(session).pipe(Effect.provideService(EventV2.Service, events)) + }), + ) }), }, }) diff --git a/packages/core/src/session/execution.ts b/packages/core/src/session/execution.ts index 5938c3772639..1d55ea8f6e54 100644 --- a/packages/core/src/session/execution.ts +++ b/packages/core/src/session/execution.ts @@ -15,6 +15,12 @@ export interface Interface { readonly wake: (sessionID: SessionSchema.ID) => Effect.Effect /** Interrupt active work owned by this process. Idle interruption is a no-op. */ readonly interrupt: (sessionID: SessionSchema.ID) => Effect.Effect + /** Serializes short Session state transitions without interrupting active work. */ + readonly withLock: ( + sessionID: SessionSchema.ID, + ) => (effect: Effect.Effect) => Effect.Effect + /** Interrupts active work and prevents new drains while the effect runs. */ + readonly exclusive: (sessionID: SessionSchema.ID, effect: Effect.Effect) => Effect.Effect } /** Routes execution from a Session ID to the runner owned by that Session's Location. */ @@ -30,5 +36,7 @@ export const noopLayer = Layer.succeed( resume: () => Effect.void, wake: () => Effect.void, interrupt: () => Effect.void, + withLock: () => (effect) => effect, + exclusive: (_sessionID, effect) => effect, }), ) diff --git a/packages/core/src/session/execution/local.ts b/packages/core/src/session/execution/local.ts index d874c40832a3..64c10673c5ca 100644 --- a/packages/core/src/session/execution/local.ts +++ b/packages/core/src/session/execution/local.ts @@ -1,4 +1,5 @@ import { Cause, Effect, Layer } from "effect" +import { KeyedMutex } from "../../effect/keyed-mutex" import { LocationServiceMap } from "../../location-service-map" import { makeGlobalNode } from "../../effect/app-node" import { SessionRunCoordinator } from "../run-coordinator" @@ -13,6 +14,7 @@ const layer = Layer.effect( Effect.gen(function* () { const store = yield* SessionStore.Service const locations = yield* LocationServiceMap.Service + const locks = KeyedMutex.makeUnsafe() const coordinator = yield* SessionRunCoordinator.make({ drain: Effect.fnUntraced(function* (sessionID: SessionSchema.ID, force) { const session = yield* store.get(sessionID) @@ -33,6 +35,13 @@ const layer = Layer.effect( interrupt: coordinator.interrupt, resume: coordinator.run, wake: coordinator.wake, + withLock: locks.withLock, + exclusive: (sessionID, effect) => + Effect.acquireUseRelease( + coordinator.pause(sessionID), + () => locks.withLock(sessionID)(effect), + (release) => release, + ), }) }), ) diff --git a/packages/core/src/session/run-coordinator.ts b/packages/core/src/session/run-coordinator.ts index 2f89aff9e3d2..c0d77b9e28a8 100644 --- a/packages/core/src/session/run-coordinator.ts +++ b/packages/core/src/session/run-coordinator.ts @@ -12,6 +12,8 @@ export interface Coordinator { readonly wake: (key: Key) => Effect.Effect /** Stops active execution and waits for its cleanup. */ readonly interrupt: (key: Key) => Effect.Effect + /** Stops active execution and prevents new drains until the returned release effect runs. */ + readonly pause: (key: Key) => Effect.Effect> } type Entry = { @@ -21,11 +23,18 @@ type Entry = { stopping: boolean } +type Pause = { + readonly done: Deferred.Deferred + pendingWake: boolean + users: number +} + export const make = (options: { readonly drain: (key: Key, force: boolean) => Effect.Effect }): Effect.Effect, never, Scope.Scope> => Effect.gen(function* () { const active = new Map>() + const pauses = new Map() const fork = yield* FiberSet.makeRuntime() const makeEntry = (): Entry => ({ @@ -66,6 +75,9 @@ export const make = (options: { const run = (key: Key): Effect.Effect => Effect.uninterruptibleMask((restore) => { + const pause = pauses.get(key) + if (pause !== undefined) return restore(Deferred.await(pause.done).pipe(Effect.andThen(run(key)))) + const entry = active.get(key) if (entry !== undefined) { if (entry.stopping) return restore(Deferred.await(entry.done).pipe(Effect.andThen(run(key)))) @@ -80,6 +92,12 @@ export const make = (options: { const wake = (key: Key) => Effect.sync(() => { + const pause = pauses.get(key) + if (pause !== undefined) { + pause.pendingWake = true + return + } + const entry = active.get(key) if (entry !== undefined) { entry.pendingWake = true @@ -100,5 +118,30 @@ export const make = (options: { return Fiber.interrupt(entry.owner) }) - return { active: Effect.sync(() => new Set(active.keys())), run, wake, interrupt } + const pause = (key: Key): Effect.Effect> => + Effect.uninterruptible( + Effect.gen(function* () { + const current = pauses.get(key) + const entry = + current ?? + ({ + done: Deferred.makeUnsafe(), + pendingWake: false, + users: 0, + } satisfies Pause) + if (!current) pauses.set(key, entry) + entry.users++ + yield* interrupt(key) + return Effect.suspend(() => { + entry.users-- + if (entry.users > 0) return Effect.void + pauses.delete(key) + Deferred.doneUnsafe(entry.done, Effect.void) + if (!entry.pendingWake) return Effect.void + return wake(key) + }) + }), + ) + + return { active: Effect.sync(() => new Set(active.keys())), run, wake, interrupt, pause } }) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 72c761e10d93..7a3d6cdb4fd9 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -384,6 +384,7 @@ const layer = Layer.effect( readonly sessionID: SessionSchema.ID readonly force: boolean }) { + if ((yield* getSession(input.sessionID)).revert) return const hasSteer = yield* SessionInput.hasPending(db, input.sessionID, "steer") const hasQueue = hasSteer ? false : yield* SessionInput.hasPending(db, input.sessionID, "queue") if (!input.force && !hasSteer && !hasQueue) return diff --git a/packages/core/test/session-prompt.test.ts b/packages/core/test/session-prompt.test.ts index c6bc9430b3ac..d6e7e7227c2b 100644 --- a/packages/core/test/session-prompt.test.ts +++ b/packages/core/test/session-prompt.test.ts @@ -40,6 +40,8 @@ const execution = Layer.succeed( Effect.sync(() => { wakeCalls.push(sessionID) }), + withLock: () => (effect) => effect, + exclusive: (_sessionID, effect) => effect, }), ) const it = testEffect( @@ -485,6 +487,52 @@ describe("SessionV2.prompt", () => { }), ) + it.effect("commits a staged revert before admitting a new prompt", () => + Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + const events = yield* EventV2.Service + const boundary = SessionMessage.ID.create() + const reverted = SessionMessage.ID.create() + yield* events.publish(SessionEvent.Prompted, { + sessionID, + messageID: boundary, + timestamp: yield* DateTime.now, + prompt: Prompt.make({ text: "Keep this prompt" }), + delivery: "steer", + }) + yield* events.publish(SessionEvent.Prompted, { + sessionID, + messageID: reverted, + timestamp: yield* DateTime.now, + prompt: Prompt.make({ text: "Discard this prompt" }), + delivery: "steer", + }) + yield* events.publish(SessionEvent.RevertEvent.Staged, { + sessionID, + timestamp: yield* DateTime.now, + revert: { messageID: boundary }, + }) + + const admitted = yield* session.prompt({ + sessionID, + prompt: Prompt.make({ text: "Prompt after revert" }), + resume: false, + }) + const { db } = yield* Database.Service + const stored = yield* db + .select({ revert: SessionTable.revert }) + .from(SessionTable) + .where(eq(SessionTable.id, sessionID)) + .get() + .pipe(Effect.orDie) + + expect(yield* session.messages({ sessionID })).toMatchObject([{ id: boundary, text: "Keep this prompt" }]) + expect(stored?.revert).toBeNull() + expect(yield* SessionInput.find(db, admitted.id)).toMatchObject({ prompt: { text: "Prompt after revert" } }) + }), + ) + it.effect("rejects reuse of one globally unique message ID across sessions", () => Effect.gen(function* () { yield* setup diff --git a/packages/core/test/session-run-coordinator.test.ts b/packages/core/test/session-run-coordinator.test.ts index dfbeda664c54..abc3fcc33170 100644 --- a/packages/core/test/session-run-coordinator.test.ts +++ b/packages/core/test/session-run-coordinator.test.ts @@ -244,6 +244,63 @@ describe("SessionRunCoordinator", () => { ), ) + it.effect("defers wakes received while paused until release", () => + Effect.scoped( + Effect.gen(function* () { + const started = yield* Deferred.make() + const interrupted = yield* Deferred.make() + const resumed = yield* Deferred.make() + let runs = 0 + const coordinator = yield* SessionRunCoordinator.make({ + drain: () => + Effect.sync(() => ++runs).pipe( + Effect.flatMap((run) => + run === 1 + ? Deferred.succeed(started, undefined).pipe( + Effect.andThen(Effect.never), + Effect.onInterrupt(() => Deferred.succeed(interrupted, undefined)), + ) + : Deferred.succeed(resumed, undefined), + ), + ), + }) + + const first = yield* coordinator.run("session").pipe(Effect.forkChild) + yield* Deferred.await(started) + const release = yield* coordinator.pause("session") + yield* Deferred.await(interrupted) + yield* coordinator.wake("session") + yield* Effect.yieldNow + + expect(runs).toBe(1) + yield* release + yield* Deferred.await(resumed) + expect(runs).toBe(2) + expect(yield* Fiber.await(first)).toMatchObject({ _tag: "Failure" }) + }), + ), + ) + + it.effect("defers resumes received while paused until release", () => + Effect.scoped( + Effect.gen(function* () { + const started = yield* Deferred.make() + const coordinator = yield* SessionRunCoordinator.make({ + drain: () => Deferred.succeed(started, undefined), + }) + + const release = yield* coordinator.pause("session") + const resumed = yield* coordinator.run("session").pipe(Effect.forkChild) + yield* Effect.yieldNow + + expect(Array.from(yield* coordinator.active)).toEqual([]) + yield* release + yield* Deferred.await(started) + yield* Fiber.join(resumed) + }), + ), + ) + it.effect("runs a wake registered during interruption cleanup", () => Effect.scoped( Effect.gen(function* () { diff --git a/packages/core/test/session-runner-recorded.test.ts b/packages/core/test/session-runner-recorded.test.ts index d45cc8c73411..f4b54b773e98 100644 --- a/packages/core/test/session-runner-recorded.test.ts +++ b/packages/core/test/session-runner-recorded.test.ts @@ -96,6 +96,8 @@ const execution = Layer.effect( resume: coordinator.run, wake: coordinator.wake, interrupt: coordinator.interrupt, + withLock: () => (effect) => effect, + exclusive: (_sessionID, effect) => effect, }) }), ).pipe(Layer.provide(runnerLayer)) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 0515d55cf5be..a488044bfe84 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -248,6 +248,8 @@ const execution = Layer.effect( resume: coordinator.run, wake: coordinator.wake, interrupt: coordinator.interrupt, + withLock: () => (effect) => effect, + exclusive: (_sessionID, effect) => effect, }) }), ).pipe(Layer.provide(runnerLayer)) @@ -555,6 +557,24 @@ const verifyPartialFlushOnInterruption = (kind: FragmentKind) => }) describe("SessionRunnerLLM", () => { + it.effect("does not request the model while a revert is staged", () => + Effect.gen(function* () { + yield* setup + const { db } = yield* Database.Service + yield* db + .update(SessionTable) + .set({ revert: { messageID: SessionMessage.ID.create() } }) + .where(eq(SessionTable.id, sessionID)) + .run() + .pipe(Effect.orDie) + requests.length = 0 + + yield* (yield* SessionRunner.Service).run({ sessionID, force: true }) + + expect(requests).toHaveLength(0) + }), + ) + it.effect("advertises and executes a globally attached application tool", () => Effect.gen(function* () { yield* setup