@swarmclawai/swarmclaw 1.2.1 → 1.2.3

package/src/lib/server/chat-execution/chat-turn-partial-persistence.ts

@@ -0,0 +1,221 @@
+import {
+  getSession,
+  saveSession,
+} from '@/lib/server/sessions/session-repository'
+import { notify } from '@/lib/server/ws-hub'
+import type { MessageToolEvent, SSEEvent } from '@/types'
+import { upsertStreamingAssistantArtifact } from '@/lib/chat/chat-streaming-state'
+import { pruneIncompleteToolEvents } from '@/lib/server/chat-execution/chat-streaming-utils'
+import {
+  collectToolEvent,
+  dedupeConsecutiveToolEvents,
+  extractEventJson,
+} from '@/lib/server/chat-execution/chat-execution-tool-events'
+import {
+  getToolEventsSnapshotKey,
+  hasPersistableAssistantPayload,
+} from '@/lib/server/chat-execution/chat-execution-utils'
+import {
+  applyMessageLifecycleHooks,
+  type PreparedExecutableChatTurn,
+} from '@/lib/server/chat-execution/chat-turn-preparation'
+
+export interface PartialAssistantSnapshot {
+  thinkingText: string
+  toolEvents: MessageToolEvent[]
+  streamErrors: string[]
+  accumulatedUsage: {
+    inputTokens: number
+    outputTokens: number
+    estimatedCost: number
+  }
+}
+
+export interface PartialAssistantPersistence {
+  emit: (event: SSEEvent) => void
+  parseAndEmit: (raw: string) => void
+  stop: () => void
+  awaitIdle: () => Promise<void>
+  getToolEvents: () => MessageToolEvent[]
+  getSnapshot: () => PartialAssistantSnapshot
+}
+
+const PARTIAL_SAVE_INTERVAL_MS = 3500
+const PARTIAL_SAVE_MIN_INTERVAL_MS = 400
+
+export function createPartialAssistantPersistence(input: {
+  prepared: PreparedExecutableChatTurn
+  onEvent?: (event: SSEEvent) => void
+}): PartialAssistantPersistence {
+  const { prepared, onEvent } = input
+  const toolEvents: MessageToolEvent[] = []
+  const streamErrors: string[] = []
+  const accumulatedUsage = { inputTokens: 0, outputTokens: 0, estimatedCost: 0 }
+
+  let thinkingText = ''
+  let streamingPartialText = ''
+  let lastPartialSaveAt = 0
+  let lastPartialSnapshotKey = ''
+  let partialSaveTimeout: ReturnType<typeof setTimeout> | null = null
+  let partialPersistenceClosed = false
+  let partialPersistChain: Promise<void> = Promise.resolve()
+
+  const persistStreamingAssistantArtifact = async () => {
+    if (prepared.hideAssistantTranscript) return
+    partialSaveTimeout = null
+    if (partialPersistenceClosed) return
+    const persistedToolEvents = toolEvents.length
+      ? dedupeConsecutiveToolEvents(pruneIncompleteToolEvents([...toolEvents]))
+      : []
+    if (!hasPersistableAssistantPayload(streamingPartialText, thinkingText, persistedToolEvents)) return
+
+    try {
+      const current = getSession(prepared.sessionId)
+      if (!current) return
+      current.messages = Array.isArray(current.messages) ? current.messages : []
+      const partialMsg = await applyMessageLifecycleHooks({
+        session: current,
+        message: {
+          role: 'assistant',
+          text: streamingPartialText,
+          time: Date.now(),
+          streaming: true,
+          runId: prepared.lifecycleRunId,
+          thinking: thinkingText || undefined,
+          toolEvents: persistedToolEvents.length ? persistedToolEvents : undefined,
+        },
+        enabledIds: prepared.extensionsForRun,
+        phase: 'assistant_partial',
+        runId: prepared.lifecycleRunId,
+        isSynthetic: true,
+      })
+      if (!partialMsg) return
+      const snapshotKey = JSON.stringify([
+        partialMsg.text,
+        partialMsg.thinking || '',
+        getToolEventsSnapshotKey(partialMsg.toolEvents || []),
+      ])
+      if (snapshotKey === lastPartialSnapshotKey) return
+      lastPartialSnapshotKey = snapshotKey
+      lastPartialSaveAt = Date.now()
+      upsertStreamingAssistantArtifact(current.messages, partialMsg, {
+        minIndex: prepared.runMessageStartIndex,
+        minTime: prepared.runStartedAt,
+      })
+      saveSession(prepared.sessionId, current)
+      notify(`messages:${prepared.sessionId}`)
+    } catch {
+      // Partial persistence is best-effort.
+    }
+  }
+
+  const triggerPartialAssistantPersist = () => {
+    partialPersistChain = partialPersistChain
+      .catch(() => {})
+      .then(async () => {
+        await persistStreamingAssistantArtifact()
+      })
+  }
+
+  const queuePartialAssistantPersist = (immediate = false) => {
+    if (partialPersistenceClosed) return
+    const now = Date.now()
+    if (immediate || now - lastPartialSaveAt >= PARTIAL_SAVE_MIN_INTERVAL_MS) {
+      if (partialSaveTimeout) {
+        clearTimeout(partialSaveTimeout)
+        partialSaveTimeout = null
+      }
+      triggerPartialAssistantPersist()
+      return
+    }
+    if (partialSaveTimeout) return
+    partialSaveTimeout = setTimeout(() => {
+      triggerPartialAssistantPersist()
+    }, PARTIAL_SAVE_MIN_INTERVAL_MS - (now - lastPartialSaveAt))
+  }
+
+  const emit = (event: SSEEvent) => {
+    let shouldPersistPartial = false
+    let immediatePartialPersist = false
+    if (event.t === 'reset') {
+      streamingPartialText = event.text || ''
+      thinkingText = ''
+      toolEvents.length = 0
+      shouldPersistPartial = true
+      immediatePartialPersist = true
+    }
+    if (event.t === 'd' && typeof event.text === 'string') {
+      streamingPartialText += event.text
+      shouldPersistPartial = true
+      immediatePartialPersist = streamingPartialText.length === event.text.length
+    }
+    if (event.t === 'err' && typeof event.text === 'string') {
+      const trimmed = event.text.trim()
+      if (trimmed) {
+        streamErrors.push(trimmed)
+        if (streamErrors.length > 8) streamErrors.shift()
+      }
+    }
+    if (event.t === 'thinking' && event.text) {
+      thinkingText += event.text
+      shouldPersistPartial = true
+    }
+    if (event.t === 'md' && event.text) {
+      try {
+        const mdPayload = JSON.parse(event.text) as Record<string, unknown>
+        const usage = mdPayload.usage as { inputTokens?: number; outputTokens?: number; estimatedCost?: number } | undefined
+        if (usage) {
+          if (typeof usage.inputTokens === 'number') accumulatedUsage.inputTokens += usage.inputTokens
+          if (typeof usage.outputTokens === 'number') accumulatedUsage.outputTokens += usage.outputTokens
+          if (typeof usage.estimatedCost === 'number') accumulatedUsage.estimatedCost += usage.estimatedCost
+        }
+      } catch {
+        // Ignore non-JSON md events.
+      }
+    }
+    collectToolEvent(event, toolEvents)
+    if (event.t === 'tool_call' || event.t === 'tool_result') {
+      shouldPersistPartial = true
+      immediatePartialPersist = true
+    }
+    if (shouldPersistPartial) queuePartialAssistantPersist(immediatePartialPersist)
+    onEvent?.(event)
+  }
+
+  const periodicTimer = setInterval(() => {
+    persistStreamingAssistantArtifact()
+  }, PARTIAL_SAVE_INTERVAL_MS)
+
+  return {
+    emit,
+    parseAndEmit(raw: string) {
+      const lines = raw.split('\n').filter(Boolean)
+      for (const line of lines) {
+        const event = extractEventJson(line)
+        if (event) emit(event)
+      }
+    },
+    stop() {
+      partialPersistenceClosed = true
+      if (partialSaveTimeout) {
+        clearTimeout(partialSaveTimeout)
+        partialSaveTimeout = null
+      }
+      clearInterval(periodicTimer)
+    },
+    awaitIdle() {
+      return partialPersistChain.catch(() => {})
+    },
+    getToolEvents() {
+      return toolEvents
+    },
+    getSnapshot() {
+      return {
+        thinkingText,
+        toolEvents: [...toolEvents],
+        streamErrors: [...streamErrors],
+        accumulatedUsage: { ...accumulatedUsage },
+      }
+    },
+  }
+}

package/src/lib/server/chat-execution/chat-turn-preflight.ts

@@ -0,0 +1,133 @@
+import { notify } from '@/lib/server/ws-hub'
+import { getSession, saveSession } from '@/lib/server/sessions/session-repository'
+import type { MessageToolEvent, SSEEvent } from '@/types'
+import type { ExecuteChatTurnResult } from './chat-execution-types'
+import {
+  resolveRequestedToolPreflightResponse,
+  runExclusiveDirectMemoryPreflight,
+} from '@/lib/server/chat-execution/chat-turn-tool-routing'
+import {
+  applyMessageLifecycleHooks,
+  type PreparedBlockedChatTurn,
+  type PreparedExecutableChatTurn,
+} from '@/lib/server/chat-execution/chat-turn-preparation'
+
+type DirectMemoryPreflightResult = Awaited<ReturnType<typeof runExclusiveDirectMemoryPreflight>>
+
+export interface ChatTurnPreflightResult {
+  terminalResult?: ExecuteChatTurnResult
+  directMemoryResult?: DirectMemoryPreflightResult
+}
+
+async function completeSyntheticAssistantTurn(params: {
+  runId?: string
+  sessionId: string
+  text: string
+  lifecycleRunId: string
+  enabledIds: string[]
+  shouldPersist: boolean
+  phase: 'assistant_final' | 'heartbeat'
+  error?: string
+  notifyMessages?: boolean
+  notifySessions?: boolean
+}): Promise<ExecuteChatTurnResult> {
+  const session = getSession(params.sessionId)
+  let persisted = false
+  if (session && params.shouldPersist) {
+    const nextAssistantMessage = await applyMessageLifecycleHooks({
+      session,
+      message: {
+        role: 'assistant',
+        text: params.text,
+        time: Date.now(),
+      },
+      enabledIds: params.enabledIds,
+      phase: params.phase,
+      runId: params.lifecycleRunId,
+      isSynthetic: true,
+    })
+    if (nextAssistantMessage) {
+      session.messages = Array.isArray(session.messages) ? session.messages : []
+      session.messages.push(nextAssistantMessage)
+      session.lastActiveAt = Date.now()
+      saveSession(params.sessionId, session)
+      if (params.notifyMessages) notify(`messages:${params.sessionId}`)
+      if (params.notifySessions) notify('sessions')
+      persisted = true
+    }
+  }
+
+  return {
+    runId: params.runId,
+    sessionId: params.sessionId,
+    text: params.text,
+    persisted,
+    toolEvents: [],
+    error: params.error,
+  }
+}
+
+export async function completeBlockedChatTurn(prepared: PreparedBlockedChatTurn): Promise<ExecuteChatTurnResult> {
+  return completeSyntheticAssistantTurn({
+    runId: prepared.runId,
+    sessionId: prepared.sessionId,
+    text: prepared.blockedMessage,
+    lifecycleRunId: prepared.lifecycleRunId,
+    enabledIds: prepared.syntheticEnabledIds,
+    shouldPersist: !prepared.internal,
+    phase: 'assistant_final',
+    error: prepared.blockedMessage,
+  })
+}
+
+export async function runChatTurnPreflight(params: {
+  prepared: PreparedExecutableChatTurn
+  emit: (event: SSEEvent) => void
+  toolEvents: MessageToolEvent[]
+}): Promise<ChatTurnPreflightResult | null> {
+  const { prepared, emit, toolEvents } = params
+
+  const requestedToolPreflightResponse = resolveRequestedToolPreflightResponse({
+    message: prepared.message,
+    enabledExtensions: prepared.extensionsForRun,
+    toolPolicy: prepared.toolPolicy,
+    appSettings: prepared.appSettings,
+    internal: prepared.internal,
+    source: prepared.source,
+    session: prepared.sessionForRun,
+  })
+  if (requestedToolPreflightResponse) {
+    return {
+      terminalResult: await completeSyntheticAssistantTurn({
+        runId: prepared.runId,
+        sessionId: prepared.sessionId,
+        text: requestedToolPreflightResponse,
+        lifecycleRunId: prepared.lifecycleRunId,
+        enabledIds: prepared.extensionsForRun,
+        shouldPersist: !prepared.hideAssistantTranscript,
+        phase: 'assistant_final',
+        notifyMessages: true,
+        notifySessions: true,
+      }),
+    }
+  }
+
+  const directMemoryResult = await runExclusiveDirectMemoryPreflight({
+    session: prepared.sessionForRun,
+    sessionId: prepared.sessionId,
+    message: prepared.message,
+    effectiveMessage: prepared.effectiveMessage,
+    enabledExtensions: prepared.extensionsForRun,
+    toolPolicy: prepared.toolPolicy,
+    appSettings: prepared.appSettings,
+    internal: prepared.internal,
+    source: prepared.source,
+    toolEvents,
+    emit,
+  })
+  if (directMemoryResult) {
+    return { directMemoryResult }
+  }
+
+  return null
+}