import { produce, type WritableDraft } from "immer" import { SessionEvent } from "./session-event" import { SessionMessage } from "./session-message" export type MemoryState = { messages: SessionMessage.Message[] } export interface Adapter { readonly getCurrentAssistant: () => SessionMessage.Assistant | undefined readonly getCurrentCompaction: () => SessionMessage.Compaction | undefined readonly getCurrentShell: (callID: string) => SessionMessage.Shell | undefined readonly updateAssistant: (assistant: SessionMessage.Assistant) => void readonly updateCompaction: (compaction: SessionMessage.Compaction) => void readonly updateShell: (shell: SessionMessage.Shell) => void readonly appendMessage: (message: SessionMessage.Message) => void readonly finish: () => Result } export function memory(state: MemoryState): Adapter { const activeAssistantIndex = () => state.messages.findLastIndex((message) => message.type === "assistant" && !message.time.completed) const activeCompactionIndex = () => state.messages.findLastIndex((message) => message.type === "compaction") const activeShellIndex = (callID: string) => state.messages.findLastIndex((message) => message.type === "shell" && message.callID === callID) return { getCurrentAssistant() { const index = activeAssistantIndex() if (index < 0) return const assistant = state.messages[index] return assistant?.type === "assistant" ? assistant : undefined }, getCurrentCompaction() { const index = activeCompactionIndex() if (index < 0) return const compaction = state.messages[index] return compaction?.type === "compaction" ? compaction : undefined }, getCurrentShell(callID) { const index = activeShellIndex(callID) if (index < 0) return const shell = state.messages[index] return shell?.type === "shell" ? shell : undefined }, updateAssistant(assistant) { const index = activeAssistantIndex() if (index < 0) return const current = state.messages[index] if (current?.type !== "assistant") return state.messages[index] = assistant }, updateCompaction(compaction) { const index = activeCompactionIndex() if (index < 0) return const current = state.messages[index] if (current?.type !== "compaction") return state.messages[index] = compaction }, updateShell(shell) { const index = activeShellIndex(shell.callID) if (index < 0) return const current = state.messages[index] if (current?.type !== "shell") return state.messages[index] = shell }, appendMessage(message) { state.messages.push(message) }, finish() { return state }, } } export function update(adapter: Adapter, event: SessionEvent.Event): Result { const currentAssistant = adapter.getCurrentAssistant() type DraftAssistant = WritableDraft type DraftTool = WritableDraft type DraftText = WritableDraft type DraftReasoning = WritableDraft const latestTool = (assistant: DraftAssistant | undefined, callID?: string) => assistant?.content.findLast( (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 latestReasoning = (assistant: DraftAssistant | undefined, reasoningID: string) => assistant?.content.findLast((item): item is DraftReasoning => item.type === "reasoning" && item.id === reasoningID) SessionEvent.All.match(event, { "session.next.agent.switched": (event) => { adapter.appendMessage( new SessionMessage.AgentSwitched({ id: event.id, type: "agent-switched", metadata: event.metadata, agent: event.data.agent, time: { created: event.data.timestamp }, }), ) }, "session.next.model.switched": (event) => { adapter.appendMessage( new SessionMessage.ModelSwitched({ id: event.id, type: "model-switched", metadata: event.metadata, model: event.data.model, time: { created: event.data.timestamp }, }), ) }, "session.next.prompted": (event) => { adapter.appendMessage( new SessionMessage.User({ id: event.id, type: "user", metadata: event.metadata, text: event.data.prompt.text, files: event.data.prompt.files, agents: event.data.prompt.agents, references: event.data.prompt.references, time: { created: event.data.timestamp }, }), ) }, "session.next.synthetic": (event) => { adapter.appendMessage( new SessionMessage.Synthetic({ sessionID: event.data.sessionID, text: event.data.text, id: event.id, type: "synthetic", time: { created: event.data.timestamp }, }), ) }, "session.next.shell.started": (event) => { adapter.appendMessage( new SessionMessage.Shell({ id: event.id, type: "shell", metadata: event.metadata, callID: event.data.callID, command: event.data.command, output: "", time: { created: event.data.timestamp }, }), ) }, "session.next.shell.ended": (event) => { const currentShell = adapter.getCurrentShell(event.data.callID) if (currentShell) { adapter.updateShell( produce(currentShell, (draft) => { draft.output = event.data.output draft.time.completed = event.data.timestamp }), ) } }, "session.next.step.started": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { draft.time.completed = event.data.timestamp }), ) } adapter.appendMessage( new SessionMessage.Assistant({ id: event.id, type: "assistant", agent: event.data.agent, model: event.data.model, time: { created: event.data.timestamp }, content: [], snapshot: event.data.snapshot ? { start: event.data.snapshot } : undefined, }), ) }, "session.next.step.ended": (event) => { if (currentAssistant) { 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 } }), ) } }, "session.next.step.failed": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { draft.time.completed = event.data.timestamp draft.finish = "error" draft.error = event.data.error }), ) } }, "session.next.text.started": () => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { draft.content.push({ type: "text", text: "", }) }), ) } }, "session.next.text.delta": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { const match = latestText(draft) if (match) match.text += event.data.delta }), ) } }, "session.next.text.ended": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { const match = latestText(draft) if (match) match.text = event.data.text }), ) } }, "session.next.tool.input.started": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { draft.content.push({ type: "tool", id: event.data.callID, name: event.data.name, time: { created: event.data.timestamp, }, state: { status: "pending", input: "", }, }) }), ) } }, "session.next.tool.input.delta": (event) => { if (currentAssistant) { 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.ended": () => {}, "session.next.tool.called": (event) => { if (currentAssistant) { 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: [], } } }), ) } }, "session.next.tool.progress": (event) => { if (currentAssistant) { 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] } }), ) } }, "session.next.tool.success": (event) => { if (currentAssistant) { 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], } } }), ) } }, "session.next.tool.failed": (event) => { if (currentAssistant) { 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, } } }), ) } }, "session.next.reasoning.started": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { draft.content.push({ type: "reasoning", id: event.data.reasoningID, text: "", }) }), ) } }, "session.next.reasoning.delta": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { const match = latestReasoning(draft, event.data.reasoningID) if (match) match.text += event.data.delta }), ) } }, "session.next.reasoning.ended": (event) => { if (currentAssistant) { adapter.updateAssistant( produce(currentAssistant, (draft) => { const match = latestReasoning(draft, event.data.reasoningID) if (match) match.text = event.data.text }), ) } }, "session.next.retried": () => {}, "session.next.compaction.started": (event) => { adapter.appendMessage( new SessionMessage.Compaction({ id: event.id, type: "compaction", metadata: event.metadata, reason: event.data.reason, summary: "", time: { created: event.data.timestamp }, }), ) }, "session.next.compaction.delta": (event) => { const currentCompaction = adapter.getCurrentCompaction() if (currentCompaction) { adapter.updateCompaction( produce(currentCompaction, (draft) => { draft.summary += event.data.text }), ) } }, "session.next.compaction.ended": (event) => { const currentCompaction = adapter.getCurrentCompaction() if (currentCompaction) { adapter.updateCompaction( produce(currentCompaction, (draft) => { draft.summary = event.data.text draft.include = event.data.include }), ) } }, }) return adapter.finish() } export * as SessionMessageUpdater from "./session-message-updater"