feat(core): add embedded v2 session runtime and tool foundation (#30632)

2026-06-03 23:02:17 -04:00
parent c35267776a
commit 76ee87ead8
215 changed files with 31344 additions and 3278 deletions
@@ -1,4 +1,4 @@
-import { produce, type WritableDraft } from "immer"
+import { castDraft, produce, type WritableDraft } from "immer"
 import { Effect } from "effect"
 import { SessionEvent } from "./event"
 import { SessionMessage } from "./message"
@@ -9,6 +9,7 @@ export type MemoryState = {

 export interface Adapter {
  readonly getCurrentAssistant: () => Effect.Effect<SessionMessage.Assistant | undefined>
+  readonly getAssistant: (messageID: SessionMessage.ID) => Effect.Effect<SessionMessage.Assistant | undefined>
  readonly getCurrentCompaction: () => Effect.Effect<SessionMessage.Compaction | undefined>
  readonly getCurrentShell: (callID: string) => Effect.Effect<SessionMessage.Shell | undefined>
  readonly updateAssistant: (assistant: SessionMessage.Assistant) => Effect.Effect<void>
@@ -18,8 +19,10 @@ export interface Adapter {
 }

 export function memory(state: MemoryState): Adapter {
-  const activeAssistantIndex = () =>
-    state.messages.findLastIndex((message) => message.type === "assistant" && !message.time.completed)
+  const assistantIndex = (messageID: SessionMessage.ID) =>
+    state.messages.findLastIndex((message) => message.id === messageID)
+  // A newer turn supersedes stale incomplete rows; never resume an older assistant projection.
+  const latestAssistantIndex = () => state.messages.findLastIndex((message) => message.type === "assistant")
  const activeCompactionIndex = () => state.messages.findLastIndex((message) => message.type === "compaction")
  const activeShellIndex = (callID: string) =>
    state.messages.findLastIndex((message) => message.type === "shell" && message.callID === callID)
@@ -27,7 +30,15 @@ export function memory(state: MemoryState): Adapter {
  return {
    getCurrentAssistant() {
      return Effect.sync(() => {
-        const index = activeAssistantIndex()
+        const index = latestAssistantIndex()
+        if (index < 0) return
+        const assistant = state.messages[index]
+        return assistant?.type === "assistant" && !assistant.time.completed ? assistant : undefined
+      })
+    },
+    getAssistant(messageID) {
+      return Effect.sync(() => {
+        const index = assistantIndex(messageID)
        if (index < 0) return
        const assistant = state.messages[index]
        return assistant?.type === "assistant" ? assistant : undefined
@@ -51,7 +62,7 @@ export function memory(state: MemoryState): Adapter {
    },
    updateAssistant(assistant) {
      return Effect.sync(() => {
-        const index = activeAssistantIndex()
+        const index = assistantIndex(assistant.id)
        if (index < 0) return
        const current = state.messages[index]
        if (current?.type !== "assistant") return
@@ -95,12 +106,18 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
      (item): item is DraftTool => item.type === "tool" && (callID === undefined || item.id === callID),
    )

-  const latestText = (assistant: DraftAssistant | undefined) =>
-    assistant?.content.findLast((item): item is DraftText => item.type === "text")
+  const latestText = (assistant: DraftAssistant | undefined, textID: string) =>
+    assistant?.content.findLast((item): item is DraftText => item.type === "text" && item.id === textID)

  const latestReasoning = (assistant: DraftAssistant | undefined, reasoningID: string) =>
    assistant?.content.findLast((item): item is DraftReasoning => item.type === "reasoning" && item.id === reasoningID)

+  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))
+    })
+
  return Effect.gen(function* () {
    yield* SessionEvent.All.match(event, {
      "session.next.agent.switched": (event) => {
@@ -200,42 +217,30 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
        })
      },
      "session.next.step.ended": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                draft.time.completed = event.data.timestamp
-                draft.finish = event.data.finish
-                draft.cost = event.data.cost
-                draft.tokens = event.data.tokens
-                if (event.data.snapshot) draft.snapshot = { ...draft.snapshot, end: event.data.snapshot }
-              }),
-            )
-          }
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          draft.time.completed = event.data.timestamp
+          draft.finish = event.data.finish
+          draft.cost = event.data.cost
+          draft.tokens = event.data.tokens
+          if (event.data.snapshot) draft.snapshot = { ...draft.snapshot, end: event.data.snapshot }
        })
      },
      "session.next.step.failed": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                draft.time.completed = event.data.timestamp
-                draft.finish = "error"
-                draft.error = event.data.error
-              }),
-            )
-          }
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          draft.time.completed = event.data.timestamp
+          draft.finish = "error"
+          draft.error = event.data.error
        })
      },
-      "session.next.text.started": () => {
+      "session.next.text.started": (event) => {
        return Effect.gen(function* () {
          const currentAssistant = yield* adapter.getCurrentAssistant()
          if (currentAssistant) {
            yield* adapter.updateAssistant(
              produce(currentAssistant, (draft) => {
-                draft.content.push(new SessionMessage.AssistantText({ type: "text", text: "" }) as DraftText)
+                draft.content.push(
+                  castDraft(new SessionMessage.AssistantText({ type: "text", id: event.data.textID, text: "" })),
+                )
              }),
            )
          }
@@ -247,7 +252,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
          if (currentAssistant) {
            yield* adapter.updateAssistant(
              produce(currentAssistant, (draft) => {
-                const match = latestText(draft)
+                const match = latestText(draft, event.data.textID)
                if (match) match.text += event.data.delta
              }),
            )
@@ -260,7 +265,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
          if (currentAssistant) {
            yield* adapter.updateAssistant(
              produce(currentAssistant, (draft) => {
-                const match = latestText(draft)
+                const match = latestText(draft, event.data.textID)
                if (match) match.text = event.data.text
              }),
            )
@@ -268,118 +273,93 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
        })
      },
      "session.next.tool.input.started": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                draft.content.push(
-                  new SessionMessage.AssistantTool({
-                    type: "tool",
-                    id: event.data.callID,
-                    name: event.data.name,
-                    time: { created: event.data.timestamp },
-                    state: new SessionMessage.ToolStatePending({ status: "pending", input: "" }),
-                  }) as DraftTool,
-                )
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          draft.content.push(
+            castDraft(
+              new SessionMessage.AssistantTool({
+                type: "tool",
+                id: event.data.callID,
+                name: event.data.name,
+                time: { created: event.data.timestamp },
+                state: new SessionMessage.ToolStatePending({ status: "pending", input: "" }),
              }),
-            )
-          }
+            ),
+          )
        })
      },
-      "session.next.tool.input.delta": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                const match = latestTool(draft, event.data.callID)
-                // oxlint-disable-next-line no-base-to-string -- event.delta is a Schema.String (runtime string)
-                if (match && match.state.status === "pending") match.state.input += event.data.delta
-              }),
-            )
-          }
+      "session.next.tool.input.delta": () => Effect.void,
+      "session.next.tool.input.ended": (event) => {
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          const match = latestTool(draft, event.data.callID)
+          if (match && match.state.status === "pending") match.state.input = event.data.text
        })
      },
-      "session.next.tool.input.ended": () => Effect.void,
      "session.next.tool.called": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                const match = latestTool(draft, event.data.callID)
-                if (match) {
-                  match.provider = event.data.provider
-                  match.time.ran = event.data.timestamp
-                  match.state = {
-                    status: "running",
-                    input: event.data.input,
-                    structured: {},
-                    content: [],
-                  }
-                }
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          const match = latestTool(draft, event.data.callID)
+          if (match) {
+            match.provider = event.data.provider
+            match.time.ran = event.data.timestamp
+            match.state = castDraft(
+              new SessionMessage.ToolStateRunning({
+                status: "running",
+                input: event.data.input,
+                structured: {},
+                content: [],
              }),
            )
          }
        })
      },
      "session.next.tool.progress": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                const match = latestTool(draft, event.data.callID)
-                if (match && match.state.status === "running") {
-                  match.state.structured = event.data.structured
-                  match.state.content = [...event.data.content]
-                }
-              }),
-            )
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          const match = latestTool(draft, event.data.callID)
+          if (match && match.state.status === "running") {
+            match.state.structured = event.data.structured
+            match.state.content = [...event.data.content]
          }
        })
      },
      "session.next.tool.success": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                const match = latestTool(draft, event.data.callID)
-                if (match && match.state.status === "running") {
-                  match.provider = event.data.provider
-                  match.time.completed = event.data.timestamp
-                  match.state = {
-                    status: "completed",
-                    input: match.state.input,
-                    structured: event.data.structured,
-                    content: [...event.data.content],
-                  }
-                }
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          const match = latestTool(draft, event.data.callID)
+          if (match && match.state.status === "running") {
+            match.provider = {
+              executed: event.data.provider.executed || match.provider?.executed === true,
+              metadata: match.provider?.metadata,
+              resultMetadata: event.data.provider.metadata,
+            }
+            match.time.completed = event.data.timestamp
+            match.state = castDraft(
+              new SessionMessage.ToolStateCompleted({
+                status: "completed",
+                input: match.state.input,
+                structured: event.data.structured,
+                content: [...event.data.content],
+                result: event.data.result,
              }),
            )
          }
        })
      },
      "session.next.tool.failed": (event) => {
-        return Effect.gen(function* () {
-          const currentAssistant = yield* adapter.getCurrentAssistant()
-          if (currentAssistant) {
-            yield* adapter.updateAssistant(
-              produce(currentAssistant, (draft) => {
-                const match = latestTool(draft, event.data.callID)
-                if (match && match.state.status === "running") {
-                  match.provider = event.data.provider
-                  match.time.completed = event.data.timestamp
-                  match.state = {
-                    status: "error",
-                    error: event.data.error,
-                    input: match.state.input,
-                    structured: match.state.structured,
-                    content: match.state.content,
-                  }
-                }
+        return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
+          const match = latestTool(draft, event.data.callID)
+          if (match && (match.state.status === "pending" || match.state.status === "running")) {
+            match.provider = {
+              executed: event.data.provider.executed || match.provider?.executed === true,
+              metadata: match.provider?.metadata,
+              resultMetadata: event.data.provider.metadata,
+            }
+            match.time.completed = event.data.timestamp
+            match.state = castDraft(
+              new SessionMessage.ToolStateError({
+                status: "error",
+                error: event.data.error,
+                input: typeof match.state.input === "string" ? {} : match.state.input,
+                structured: match.state.status === "running" ? match.state.structured : {},
+                content: match.state.status === "running" ? match.state.content : [],
+                result: event.data.result,
              }),
            )
          }
@@ -392,11 +372,14 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
            yield* adapter.updateAssistant(
              produce(currentAssistant, (draft) => {
                draft.content.push(
-                  new SessionMessage.AssistantReasoning({
-                    type: "reasoning",
-                    id: event.data.reasoningID,
-                    text: "",
-                  }) as DraftReasoning,
+                  castDraft(
+                    new SessionMessage.AssistantReasoning({
+                      type: "reasoning",
+                      id: event.data.reasoningID,
+                      text: "",
+                      providerMetadata: event.data.providerMetadata,
+                    }),
+                  ),
                )
              }),
            )
@@ -423,7 +406,10 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
            yield* adapter.updateAssistant(
              produce(currentAssistant, (draft) => {
                const match = latestReasoning(draft, event.data.reasoningID)
-                if (match) match.text = event.data.text
+                if (match) {
+                  match.text = event.data.text
+                  if (event.data.providerMetadata !== undefined) match.providerMetadata = event.data.providerMetadata
+                }
              }),
            )
          }