refactor(core): capture session inside loop operations

packages/core/agent-loop/src/loop.ts

@@ -95,8 +95,8 @@ export interface LoopHandle {
 /**
  * Drive queued batches as durable turns until disposal. Plugin failures end the
  * current turn without terminating the driver. The caller establishes the
- * `ctx.agents.withInitiator()` boundary before entry; package-private helpers
- * recover that exact Agent from the inherited store.
+ * `ctx.agents.withInitiator()` boundary before entry; package-private
+ * orchestration recovers that exact Agent and captures its Session locally.
  * @param ctx - the plugin context the loop reaches its initiating Agent,
  * events (agent/…, session/flush), and services (systemPrompt, llm, tools)
  * through.
@@ -169,6 +169,13 @@ async function runTurn(
 ): Promise<boolean> {
   const agent = ctx.agents.requireInitiator()
   const { session } = agent
+  const drainSteering = (): boolean => {
+    const messages = handle.inbox.drainSteering()
+    for (const message of messages) {
+      session.append('steering/message', { turn, content: message.content, source: message.source }, { surfaceOp: 'append' })
+    }
+    return messages.length > 0
+  }
 
   // Drain before opening the turn, but append only after `turn/start`.
   const queued = handle.inbox.drainQueued()
@@ -267,7 +274,7 @@ async function runTurn(
 
       // Steering from the previous round's continuation listeners joins before
       // the request.
-      drainSteering(session, handle.inbox, turn)
+      drainSteering()
 
       // The step's AbortController exists BEFORE any async pre-step work so a
       // dispose() or cancel() — in a synchronous turn-start listener or an
@@ -370,7 +377,7 @@ async function runTurn(
       if (stepReason) reason = stepReason
 
       // Steering that arrived during streaming/tool execution.
-      const steered = drainSteering(session, handle.inbox, turn)
+      const steered = drainSteering()
 
       closeStep()
 
@@ -459,15 +466,6 @@ async function runTurn(
   return terminalStopped
 }
 
-/** Drain the steering queue into the session. Returns whether any arrived. */
-function drainSteering(session: Session, inbox: Inbox, turn: number): boolean {
-  const messages = inbox.drainSteering()
-  for (const message of messages) {
-    session.append('steering/message', { turn, content: message.content, source: message.source }, { surfaceOp: 'append' })
-  }
-  return messages.length > 0
-}
-
 /**
  * Run one committed step: transform call config, log the request header, build
  * the request from the cached prefix plus the step-boundary snapshot, stream and
@@ -543,15 +541,47 @@ async function runStep(
   const stepError = finishError(assembler.finish)
   if (stepError) throw stepError
 
+  const recordAssistantMessage = (
+    assembledContent: ContentBlock[],
+    message: Message,
+    preserveReplayState = true,
+  ): void => {
+    session.append(
+      'assistant/message',
+      {
+        turn,
+        step,
+        content: message.content,
+        provenance: assistantProvenance(
+          header.config,
+          assembler.replayState,
+          preserveReplayState && isDeepStrictEqual(message.content, assembledContent),
+        ),
+        ...assembler.usage === undefined ? {} : { usage: assembler.usage },
+      },
+      { surfaceOp: 'append', sourceEventSeqs: chunkSeqs },
+    )
+  }
+
+  // A rejected result still records the successful provider call without retaining rejected output.
+  const processStepResult = async (assembledContent: ContentBlock[], message: Message): Promise<Message> => {
+    try {
+      return await events.waterfall(
+        'agent/step-result', turn, step, message, () => Promise.resolve(message),
+      )
+    } catch (error: unknown) {
+      recordAssistantMessage(assembledContent, { ...message, content: [] }, false)
+      throw error
+    }
+  }
+
   if (assembler.finish.kind === 'max-tokens') {
     const assembled = assembler.message()
     const assembledContent = structuredClone(assembled.content)
     let message: Message = withoutToolCalls(assembled)
-    message = withoutToolCalls(await processStepResult(
-      events, session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs,
-    ))
+    message = withoutToolCalls(await processStepResult(assembledContent, message))
     // Preserve usage even when max-token truncation produced no content.
-    recordAssistantMessage(session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs)
+    recordAssistantMessage(assembledContent, message)
     return { hadToolCalls: false, finish: assembler.finish }
   }
 
@@ -559,13 +589,11 @@ async function runStep(
   const assembled = assembler.message()
   const assembledContent = structuredClone(assembled.content)
   let message: Message = assembled
-  message = await processStepResult(
-    events, session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs,
-  )
+  message = await processStepResult(assembledContent, message)
 
   // Every successful call records its completion anchor, including explicit
   // empty chunk provenance for a contentless, usage-less provider response.
-  recordAssistantMessage(session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs)
+  recordAssistantMessage(assembledContent, message)
 
   // Dispatch may overlap; policy, durable results, and result context stay model-ordered.
   const toolCalls = message.content.filter(block => block.type === 'tool-call')
@@ -578,67 +606,6 @@ async function runStep(
   })
 }
 
-/** Preserve successful-call accounting without retaining output that result processing rejected. */
-async function processStepResult(
-  events: AgentEventDispatch,
-  session: Session,
-  turn: number,
-  step: number,
-  config: LlmCallConfig,
-  assembledContent: ContentBlock[],
-  message: Message,
-  assembler: BlockAssembler,
-  chunkSeqs: number[],
-): Promise<Message> {
-  try {
-    return await events.waterfall(
-      'agent/step-result', turn, step, message, () => Promise.resolve(message),
-    )
-  } catch (error: unknown) {
-    recordAssistantMessage(
-      session,
-      turn,
-      step,
-      config,
-      assembledContent,
-      { ...message, content: [] },
-      assembler,
-      chunkSeqs,
-      false,
-    )
-    throw error
-  }
-}
-
-/** Record one content-or-usage assistant message with replay-safe provenance. */
-function recordAssistantMessage(
-  session: Session,
-  turn: number,
-  step: number,
-  config: LlmCallConfig,
-  assembledContent: ContentBlock[],
-  message: Message,
-  assembler: BlockAssembler,
-  chunkSeqs: number[],
-  preserveReplayState = true,
-): void {
-  session.append(
-    'assistant/message',
-    {
-      turn,
-      step,
-      content: message.content,
-      provenance: assistantProvenance(
-        config,
-        assembler.replayState,
-        preserveReplayState && isDeepStrictEqual(message.content, assembledContent),
-      ),
-      ...assembler.usage === undefined ? {} : { usage: assembler.usage },
-    },
-    { surfaceOp: 'append', sourceEventSeqs: chunkSeqs },
-  )
-}
-
 /** Build durable assistant provenance, dropping replay state after any content rewrite. */
 function assistantProvenance(config: LlmCallConfig, replayState: unknown, contentUnchanged: boolean): NonNullable<Message['provenance']> {
   return {