refactor(session): simplify execution and state workflows

packages/core/src/form.ts

@@ -109,16 +109,11 @@ export const layer = Layer.effect(
       },
     )
 
-    const find = Effect.fn("Form.find")(function* (id: ID) {
-      return yield* Cache.getSuccess(forms, id).pipe(
-        Effect.flatMap((entry) =>
-          Option.match(entry, {
-            onNone: () => Effect.fail(new NotFoundError({ id })),
-            onSome: Effect.succeed,
-          }),
-        ),
-      )
-    })
+    const requireEntry = Effect.fn("Form.requireEntry")((id: ID) =>
+      Cache.getSuccess(forms, id).pipe(
+        Effect.flatMap((entry) => Effect.fromOption(entry, () => new NotFoundError({ id }))),
+      ),
+    )
 
     const create = Effect.fn("Form.create")((input: CreateInput) =>
       Effect.uninterruptible(
@@ -151,7 +146,7 @@ export const layer = Layer.effect(
       Effect.uninterruptibleMask((restore) =>
         Effect.gen(function* () {
           const form = yield* create(input)
-          const entry = yield* find(form.id).pipe(Effect.orDie)
+          const entry = yield* requireEntry(form.id).pipe(Effect.orDie)
           return yield* restore(Deferred.await(entry.deferred)).pipe(
             Effect.onInterrupt(() => Effect.ignore(cancel(form.id))),
           )
@@ -160,7 +155,7 @@ export const layer = Layer.effect(
     )
 
     const get = Effect.fn("Form.get")(function* (id: ID) {
-      return (yield* find(id)).form
+      return (yield* requireEntry(id)).form
     })
 
     const list = Effect.fn("Form.list")(function* (input?: ListInput) {
@@ -172,13 +167,13 @@ export const layer = Layer.effect(
     })
 
     const state = Effect.fn("Form.state")(function* (id: ID) {
-      return (yield* find(id)).state
+      return (yield* requireEntry(id)).state
     })
 
     const reply = Effect.fn("Form.reply")((input: ReplyInput) =>
       Effect.uninterruptible(
         Effect.gen(function* () {
-          const entry = yield* find(input.id)
+          const entry = yield* requireEntry(input.id)
           if (entry.state.status !== "pending") return yield* new AlreadySettledError({ id: input.id })
           const invalid = validateAnswer(entry.form.fields, input.answer)
           if (invalid) return yield* new InvalidAnswerError({ id: input.id, message: invalid })
@@ -197,7 +192,7 @@ export const layer = Layer.effect(
     const cancel = Effect.fn("Form.cancel")((id: ID) =>
       Effect.uninterruptible(
         Effect.gen(function* () {
-          const entry = yield* find(id)
+          const entry = yield* requireEntry(id)
           if (entry.state.status !== "pending") return yield* new AlreadySettledError({ id })
           const next: TerminalState = { status: "cancelled" }
           yield* bus.publish(Form.Event.Cancelled, { id, sessionID: entry.form.sessionID })

packages/core/src/session/inbox.ts

@@ -157,14 +157,14 @@ export const admit = Effect.fn("SessionInbox.admit")(function* (
       item: request.item,
     })
     .pipe(
-      Effect.flatMap((event) => {
-        const base = {
+      Effect.map((event) =>
+        Info.make({
           id: request.id,
           sessionID: request.sessionID,
           timeCreated: DateTime.makeUnsafe(event.created),
-        }
-        return Effect.succeed(Info.make({ ...base, ...request.item }))
-      }),
+          ...request.item,
+        }),
+      ),
       Effect.catchDefect((defect) =>
         find(db, request.id).pipe(
           Effect.flatMap((stored) =>

packages/core/src/session/message-updater.ts

@@ -42,18 +42,18 @@ export function update(adapter: Adapter, event: SessionEvent.DurableEvent) {
   const updateOwnedAssistant = (messageID: SessionMessage.ID, recipe: (draft: DraftAssistant) => void) =>
     Effect.gen(function* () {
       const assistant = yield* adapter.getAssistant(messageID)
-      if (assistant) yield* adapter.updateAssistant(produce(assistant, recipe))
+      if (!assistant) return
+      yield* adapter.updateAssistant(produce(assistant, recipe))
     })
 
   const clearCurrentRetry = Effect.gen(function* () {
     const assistant = yield* adapter.getCurrentAssistant()
-    if (assistant?.retry) {
-      yield* adapter.updateAssistant(
-        produce(assistant, (draft) => {
-          draft.retry = undefined
-        }),
-      )
-    }
+    if (!assistant?.retry) return
+    yield* adapter.updateAssistant(
+      produce(assistant, (draft) => {
+        draft.retry = undefined
+      }),
+    )
   })
 
   const project = pipe(

packages/core/src/session/model-request.ts

@@ -266,28 +266,28 @@ export const layer = Layer.effect(
           toolChoice: input.toolChoice,
         }),
       )
-      const webSocketEligible =
-        !(yield* hooks.has("session", "http.request", resolved.ref.providerID)) &&
-        !(yield* hooks.has("session", "http.response", resolved.ref.providerID))
+      const hasHttpHooks =
+        (yield* hooks.has("session", "http.request", resolved.ref.providerID)) ||
+        (yield* hooks.has("session", "http.response", resolved.ref.providerID))
       const webSocket =
         resolved.capabilities.responsesWebsockets === true
           ? yield* Config.boolean(responsesWebSocketFlag(resolved.ref.providerID)).pipe(
               Config.withDefault(false),
               Effect.orDie,
             )
           : false
-      const http = webSocketEligible
-        ? undefined
-        : SessionModelHttp.middleware(hooks, {
+      const http = hasHttpHooks
+        ? SessionModelHttp.middleware(hooks, {
             sessionID: session.id,
             agent: input.scope.agentID,
             model: resolved.ref,
           })
+        : undefined
       const options: StreamOptions = {
         ...(http ? { http } : {}),
         ...(input.webSocket === "session" &&
         webSocket &&
-        webSocketEligible &&
+        !hasHttpHooks &&
         resolved.capabilities.responsesWebsockets === true
           ? { webSocket: transport.bind(session.id) }
           : {}),

packages/core/src/session/runner/publish-llm-event.ts

@@ -84,17 +84,15 @@ const hostedContent = (result: ToolResultValue): NonEmptyContent => {
  */
 export const createLLMEventPublisher = (bus: Pick<Bus.Interface, "publish">, input: Input) => {
   const deltaBatchInterval = 100
-  const tools = new Map<
-    string,
-    {
-      readonly assistantMessageID: SessionMessage.ID
-      readonly name: string
-      called: boolean
-      settled: boolean
-      providerExecuted: boolean
-      progress?: Tool.Metadata
-    }
-  >()
+  type ToolState = {
+    readonly assistantMessageID: SessionMessage.ID
+    readonly name: string
+    called: boolean
+    settled: boolean
+    providerExecuted: boolean
+    progress?: Tool.Metadata
+  }
+  const tools = new Map<string, ToolState>()
   const failureSnapshot = (tool: { readonly progress?: Tool.Metadata }, metadata?: Tool.Metadata) => {
     if (tool.progress === undefined) return metadata === undefined ? {} : { metadata }
     if (metadata === undefined) return { metadata: tool.progress }
@@ -263,20 +261,22 @@ export const createLLMEventPublisher = (bus: Pick<Bus.Interface, "publish">, inp
   }) {
     if (tools.has(event.id)) return yield* Effect.die(new Error(`Duplicate tool input start: ${event.id}`))
     const assistantMessageID = yield* startAssistant()
-    tools.set(event.id, {
+    const tool: ToolState = {
       assistantMessageID,
       name: event.name,
       called: false,
       settled: false,
       providerExecuted: event.providerExecuted === true,
-    })
+    }
+    tools.set(event.id, tool)
     yield* toolInput.start(event.id)
     yield* bus.publish(SessionEvent.Tool.Input.Started, {
       sessionID: input.sessionID,
       assistantMessageID,
       id: event.id,
       name: event.name,
     })
+    return tool
   })
 
   const endToolInput = Effect.fnUntraced(function* (
@@ -296,9 +296,8 @@ export const createLLMEventPublisher = (bus: Pick<Bus.Interface, "publish">, inp
     readonly name: string
     readonly raw: string
   }) {
-    if (!tools.has(event.id)) yield* startToolInput(event)
-    const tool = tools.get(event.id)
-    if (!tool || tool.called || tool.settled)
+    const tool = tools.get(event.id) ?? (yield* startToolInput(event))
+    if (tool.called || tool.settled)
       return yield* Effect.die(new Error(`Malformed tool input after call settlement: ${event.id}`))
     if (tool.name !== event.name)
       return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`))
@@ -443,8 +442,7 @@ export const createLLMEventPublisher = (bus: Pick<Bus.Interface, "publish">, inp
         return
       case "tool-call": {
         outputStarted = true
-        if (!tools.has(event.id)) yield* startToolInput(event)
-        const tool = tools.get(event.id)!
+        const tool = tools.get(event.id) ?? (yield* startToolInput(event))
         if (toolInput.has(event.id)) yield* endToolInput(event)
         if (tool.name !== event.name)
           return yield* Effect.die(new Error(`Tool call name changed for ${event.id}: ${tool.name} -> ${event.name}`))

packages/core/src/state.ts

@@ -89,15 +89,14 @@ export function create<State, DraftApi>(options: Options<State, DraftApi>): Inte
     if (options.finalize) yield* options.finalize(options.draft(next))
   })
 
-  const apply = (transform: TransformCallback<DraftApi>, draft: DraftApi) =>
-    Effect.sync(() => {
-      transform(draft)
-    })
-
   const materialize = Effect.fnUntraced(function* () {
     const next = options.initial()
     const api = options.draft(next)
-    for (const transform of transforms) yield* apply(transform.run, api)
+    for (const transform of transforms) {
+      yield* Effect.sync(() => {
+        transform.run(api)
+      })
+    }
     yield* commit(next)
   })
 
@@ -135,7 +134,7 @@ export function create<State, DraftApi>(options: Options<State, DraftApi>): Inte
     return yield* Deferred.await(done)
   })
 
-  const result: Interface<State, DraftApi> = {
+  return {
     get: () => state,
     transform: Effect.fn("State.transform")(function* (update) {
       yield* Effect.annotateCurrentSpan("state", options.name ?? "anonymous")
@@ -176,5 +175,4 @@ export function create<State, DraftApi>(options: Options<State, DraftApi>): Inte
     }),
     reload,
   }
-  return result
 }