simstudioai--sim/apps/sim/lib/copilot/chat/persisted-message.ts

import { generateId } from '@sim/utils/id'
import {
  mergeAndRedactPersistedBlocks,
  redactSensitiveContent,
  redactToolCallResult,
} from '@/lib/copilot/chat/sim-key-redaction'
import {
  MothershipStreamV1CompletionStatus,
  MothershipStreamV1EventType,
  MothershipStreamV1SpanLifecycleEvent,
  MothershipStreamV1SpanPayloadKind,
  type MothershipStreamV1StreamScope,
  MothershipStreamV1TextChannel,
  MothershipStreamV1ToolOutcome,
  MothershipStreamV1ToolPhase,
} from '@/lib/copilot/generated/mothership-stream-v1'
import type {
  ContentBlock,
  LocalToolCallStatus,
  OrchestratorResult,
} from '@/lib/copilot/request/types'

export type PersistedToolState = LocalToolCallStatus | MothershipStreamV1ToolOutcome | 'interrupted'

interface PersistedToolCall {
  id: string
  name: string
  state: PersistedToolState
  params?: Record<string, unknown>
  result?: { success: boolean; output?: unknown; error?: string }
  error?: string
  calledBy?: string
  durationMs?: number
  display?: { title?: string }
}

export interface PersistedContentBlock {
  type: MothershipStreamV1EventType
  lane?: MothershipStreamV1StreamScope['lane']
  channel?: MothershipStreamV1TextChannel
  phase?: MothershipStreamV1ToolPhase
  kind?: MothershipStreamV1SpanPayloadKind
  lifecycle?: MothershipStreamV1SpanLifecycleEvent
  status?: MothershipStreamV1CompletionStatus
  content?: string
  toolCall?: PersistedToolCall
  timestamp?: number
  endedAt?: number
  parentToolCallId?: string
  spanId?: string
  parentSpanId?: string
}

export interface PersistedFileAttachment {
  id: string
  key: string
  filename: string
  media_type: string
  size: number
}

interface PersistedMessageContext {
  kind: string
  label: string
  workflowId?: string
  knowledgeId?: string
  tableId?: string
  fileId?: string
  folderId?: string
  chatId?: string
}

export interface PersistedMessage {
  id: string
  role: 'user' | 'assistant'
  content: string
  timestamp: string
  requestId?: string
  contentBlocks?: PersistedContentBlock[]
  fileAttachments?: PersistedFileAttachment[]
  contexts?: PersistedMessageContext[]
}

/**
 * Drop the `output` of every persisted tool result, keeping `success` and
 * `error`. Tool outputs are never rendered (the chat thread shows only the tool
 * name/title/status) and never replayed to the model (the upstream copilot
 * service owns conversation memory), so storing them only bloats
 * `copilot_messages.content` — a single `get_workflow_logs`/`run_workflow`
 * result can reach hundreds of MB and stall task loads.
 *
 * Applied on both the write path (so new rows never store outputs) and the read
 * path (so already-bloated rows still load fast). Returns the original
 * reference when there is nothing to strip, preserving memoized identity for
 * read-side consumers.
 */
export function stripToolResultOutput(message: PersistedMessage): PersistedMessage {
  if (!message.contentBlocks?.length) return message
  let changed = false
  const contentBlocks = message.contentBlocks.map((block) => {
    const toolCall = block.toolCall
    const result = toolCall?.result
    if (!toolCall || !result || typeof result !== 'object' || !('output' in result)) return block
    changed = true
    const strippedResult: { success: boolean; error?: string } = { success: result.success }
    if (result.error !== undefined) strippedResult.error = result.error
    return { ...block, toolCall: { ...toolCall, result: strippedResult } }
  })
  return changed ? { ...message, contentBlocks } : message
}

// ---------------------------------------------------------------------------
// Write: OrchestratorResult → PersistedMessage
// ---------------------------------------------------------------------------

function resolveToolState(block: ContentBlock): PersistedToolState {
  const tc = block.toolCall
  if (!tc) return 'pending'
  if (tc.result?.success !== undefined) {
    return tc.result.success
      ? MothershipStreamV1ToolOutcome.success
      : MothershipStreamV1ToolOutcome.error
  }
  return tc.status as PersistedToolState
}

/**
 * Copy `timestamp` / `endedAt` from a source object onto a target object.
 * Shared by every block mapper (persist, display, snapshot) so the timing
 * metadata that drives the `Thought for Ns` chip survives the full
 * persist → normalize → display round-trip — and one rule lives in one place.
 */
export function withBlockTiming<T>(target: T, src: { timestamp?: number; endedAt?: number }): T {
  const writable = target as { timestamp?: number; endedAt?: number }
  if (typeof src.timestamp === 'number') writable.timestamp = src.timestamp
  if (typeof src.endedAt === 'number') writable.endedAt = src.endedAt
  return target
}

function withBlockParent<T>(target: T, src: { parentToolCallId?: string }): T {
  if (src.parentToolCallId) {
    ;(target as { parentToolCallId?: string }).parentToolCallId = src.parentToolCallId
  }
  return target
}

/**
 * Carry deterministic span identity (spanId / parentSpanId) across a block
 * mapping so the nesting tree survives the persist → normalize → display round
 * trip. Shared by both the write and read paths.
 */
function withBlockSpan<T>(target: T, src: { spanId?: string; parentSpanId?: string }): T {
  const writable = target as { spanId?: string; parentSpanId?: string }
  if (src.spanId) writable.spanId = src.spanId
  if (src.parentSpanId) writable.parentSpanId = src.parentSpanId
  return target
}

function mapContentBlock(block: ContentBlock): PersistedContentBlock {
  const persisted = mapContentBlockBody(block)
  return withBlockSpan(withBlockParent(withBlockTiming(persisted, block), block), block)
}

function mapContentBlockBody(block: ContentBlock): PersistedContentBlock {
  switch (block.type) {
    case 'text':
      return {
        type: MothershipStreamV1EventType.text,
        channel: MothershipStreamV1TextChannel.assistant,
        content: block.content,
      }
    case 'thinking':
      return {
        type: MothershipStreamV1EventType.text,
        channel: MothershipStreamV1TextChannel.thinking,
        content: block.content,
      }
    case 'subagent':
      return {
        type: MothershipStreamV1EventType.span,
        kind: MothershipStreamV1SpanPayloadKind.subagent,
        lifecycle: MothershipStreamV1SpanLifecycleEvent.start,
        content: block.content,
      }
    case 'subagent_text':
      return {
        type: MothershipStreamV1EventType.text,
        lane: 'subagent',
        channel: MothershipStreamV1TextChannel.assistant,
        content: block.content,
      }
    case 'subagent_thinking':
      return {
        type: MothershipStreamV1EventType.text,
        lane: 'subagent',
        channel: MothershipStreamV1TextChannel.thinking,
        content: block.content,
      }
    case 'tool_call': {
      if (!block.toolCall) {
        return {
          type: MothershipStreamV1EventType.tool,
          phase: MothershipStreamV1ToolPhase.call,
          content: block.content,
        }
      }
      const state = resolveToolState(block)
      const isSubagentTool = !!block.calledBy
      const isNonTerminal =
        state === MothershipStreamV1ToolOutcome.cancelled ||
        state === 'pending' ||
        state === 'executing'

      const redactedResult = redactToolCallResult(block.toolCall.name, block.toolCall.result)

      const toolCall: PersistedToolCall = {
        id: block.toolCall.id,
        name: block.toolCall.name,
        state,
        ...(isSubagentTool && isNonTerminal ? {} : { result: redactedResult }),
        ...(isSubagentTool && isNonTerminal
          ? {}
          : block.toolCall.params
            ? { params: block.toolCall.params }
            : {}),
        ...(block.calledBy ? { calledBy: block.calledBy } : {}),
        ...(block.toolCall.displayTitle
          ? {
              display: {
                title: block.toolCall.displayTitle,
              },
            }
          : {}),
      }

      return {
        type: MothershipStreamV1EventType.tool,
        phase: MothershipStreamV1ToolPhase.call,
        toolCall,
      }
    }
    default:
      return { type: MothershipStreamV1EventType.text, content: block.content }
  }
}

export function buildPersistedAssistantMessage(
  result: OrchestratorResult,
  requestId?: string
): PersistedMessage {
  const message: PersistedMessage = {
    id: generateId(),
    role: 'assistant',
    content: redactSensitiveContent(result.content),
    timestamp: new Date().toISOString(),
  }

  if (requestId || result.requestId) {
    message.requestId = requestId || result.requestId
  }

  if (result.contentBlocks.length > 0) {
    message.contentBlocks = mergeAndRedactPersistedBlocks(result.contentBlocks.map(mapContentBlock))
  }

  return message
}

export function withStoppedContentBlock(message: PersistedMessage): PersistedMessage {
  const contentBlocks = message.contentBlocks ?? []
  const hasAssistantText = contentBlocks.some(
    (block) =>
      block.type === MothershipStreamV1EventType.text &&
      block.channel !== MothershipStreamV1TextChannel.thinking &&
      block.content?.trim()
  )
  if (
    contentBlocks.some(
      (block) =>
        block.type === MothershipStreamV1EventType.complete &&
        block.status === MothershipStreamV1CompletionStatus.cancelled
    )
  ) {
    return message
  }

  return normalizeMessage({
    ...message,
    contentBlocks: [
      ...(hasAssistantText || !message.content.trim()
        ? []
        : [
            {
              type: MothershipStreamV1EventType.text,
              channel: MothershipStreamV1TextChannel.assistant,
              content: message.content,
            },
          ]),
      ...contentBlocks,
      {
        type: MothershipStreamV1EventType.complete,
        status: MothershipStreamV1CompletionStatus.cancelled,
      },
    ],
  })
}

export interface UserMessageParams {
  id: string
  content: string
  fileAttachments?: PersistedFileAttachment[]
  contexts?: PersistedMessageContext[]
}

export function buildPersistedUserMessage(params: UserMessageParams): PersistedMessage {
  const message: PersistedMessage = {
    id: params.id,
    role: 'user',
    content: params.content,
    timestamp: new Date().toISOString(),
  }

  if (params.fileAttachments && params.fileAttachments.length > 0) {
    message.fileAttachments = params.fileAttachments
  }

  if (params.contexts && params.contexts.length > 0) {
    message.contexts = params.contexts.map((c) => ({
      kind: c.kind,
      label: c.label,
      ...(c.workflowId ? { workflowId: c.workflowId } : {}),
      ...(c.knowledgeId ? { knowledgeId: c.knowledgeId } : {}),
      ...(c.tableId ? { tableId: c.tableId } : {}),
      ...(c.fileId ? { fileId: c.fileId } : {}),
      ...(c.folderId ? { folderId: c.folderId } : {}),
      ...(c.chatId ? { chatId: c.chatId } : {}),
    }))
  }

  return message
}

// ---------------------------------------------------------------------------
// Read: raw JSONB → PersistedMessage
// Handles both canonical (type: 'tool', 'text', 'span', 'complete') and
// legacy (type: 'tool_call', 'thinking', 'subagent', 'stopped') blocks.
// ---------------------------------------------------------------------------

const CANONICAL_BLOCK_TYPES: Set<string> = new Set(Object.values(MothershipStreamV1EventType))

interface RawBlock {
  type: string
  lane?: string
  content?: string
  /** Go persists text blocks with key "text" instead of "content" */
  text?: string
  channel?: string
  phase?: string
  kind?: string
  lifecycle?: string
  status?: string
  timestamp?: number
  endedAt?: number
  parentToolCallId?: string
  spanId?: string
  parentSpanId?: string
  toolCall?: {
    id?: string
    name?: string
    state?: string
    params?: Record<string, unknown>
    result?: { success: boolean; output?: unknown; error?: string }
    display?: { text?: string; title?: string; phaseLabel?: string }
    calledBy?: string
    durationMs?: number
    error?: string
  } | null
}

interface LegacyToolCall {
  id: string
  name: string
  status: string
  params?: Record<string, unknown>
  result?: unknown
  error?: string
  durationMs?: number
}

const OUTCOME_NORMALIZATION: Record<string, PersistedToolState> = {
  [MothershipStreamV1ToolOutcome.success]: MothershipStreamV1ToolOutcome.success,
  [MothershipStreamV1ToolOutcome.error]: MothershipStreamV1ToolOutcome.error,
  [MothershipStreamV1ToolOutcome.cancelled]: MothershipStreamV1ToolOutcome.cancelled,
  [MothershipStreamV1ToolOutcome.skipped]: MothershipStreamV1ToolOutcome.skipped,
  [MothershipStreamV1ToolOutcome.rejected]: MothershipStreamV1ToolOutcome.rejected,
  aborted: MothershipStreamV1ToolOutcome.cancelled,
  failed: MothershipStreamV1ToolOutcome.error,
  interrupted: 'interrupted',
  pending: 'pending',
  executing: 'executing',
}

function normalizeToolState(state: string | undefined): PersistedToolState {
  if (!state) return 'pending'
  return OUTCOME_NORMALIZATION[state] ?? MothershipStreamV1ToolOutcome.error
}

function isCanonicalBlock(block: RawBlock): boolean {
  return CANONICAL_BLOCK_TYPES.has(block.type)
}

function normalizeCanonicalBlock(block: RawBlock): PersistedContentBlock {
  const result: PersistedContentBlock = {
    type: block.type as MothershipStreamV1EventType,
  }
  if (block.lane === 'subagent') {
    result.lane = block.lane
  }
  const blockContent = block.content ?? block.text
  if (blockContent !== undefined) result.content = blockContent
  if (block.channel) result.channel = block.channel as MothershipStreamV1TextChannel
  if (block.phase) result.phase = block.phase as MothershipStreamV1ToolPhase
  if (block.kind) result.kind = block.kind as MothershipStreamV1SpanPayloadKind
  if (block.lifecycle) result.lifecycle = block.lifecycle as MothershipStreamV1SpanLifecycleEvent
  if (block.status) result.status = block.status as MothershipStreamV1CompletionStatus
  if (block.parentToolCallId) result.parentToolCallId = block.parentToolCallId
  if (block.toolCall) {
    result.toolCall = {
      id: block.toolCall.id ?? '',
      name: block.toolCall.name ?? '',
      state: normalizeToolState(block.toolCall.state),
      ...(block.toolCall.params ? { params: block.toolCall.params } : {}),
      ...(block.toolCall.result ? { result: block.toolCall.result } : {}),
      ...(block.toolCall.calledBy ? { calledBy: block.toolCall.calledBy } : {}),
      ...(block.toolCall.error ? { error: block.toolCall.error } : {}),
      ...(block.toolCall.durationMs ? { durationMs: block.toolCall.durationMs } : {}),
      ...(block.toolCall.display
        ? {
            display: {
              title:
                block.toolCall.display.title ??
                block.toolCall.display.text ??
                block.toolCall.display.phaseLabel,
            },
          }
        : {}),
    }
  }
  return result
}

function normalizeLegacyBlock(block: RawBlock): PersistedContentBlock {
  if (block.type === 'tool_call' && block.toolCall) {
    return {
      type: MothershipStreamV1EventType.tool,
      phase: MothershipStreamV1ToolPhase.call,
      toolCall: {
        id: block.toolCall.id ?? '',
        name: block.toolCall.name ?? '',
        state: normalizeToolState(block.toolCall.state),
        ...(block.toolCall.params ? { params: block.toolCall.params } : {}),
        ...(block.toolCall.result ? { result: block.toolCall.result } : {}),
        ...(block.toolCall.calledBy ? { calledBy: block.toolCall.calledBy } : {}),
        ...(block.toolCall.display
          ? {
              display: {
                title:
                  block.toolCall.display.title ??
                  block.toolCall.display.text ??
                  block.toolCall.display.phaseLabel,
              },
            }
          : {}),
      },
    }
  }

  if (block.type === 'thinking') {
    return {
      type: MothershipStreamV1EventType.text,
      channel: MothershipStreamV1TextChannel.thinking,
      content: block.content,
    }
  }

  if (block.type === 'subagent' || block.type === 'subagent_text') {
    if (block.type === 'subagent_text') {
      return {
        type: MothershipStreamV1EventType.text,
        lane: 'subagent',
        channel: MothershipStreamV1TextChannel.assistant,
        content: block.content,
      }
    }
    return {
      type: MothershipStreamV1EventType.span,
      kind: MothershipStreamV1SpanPayloadKind.subagent,
      lifecycle: MothershipStreamV1SpanLifecycleEvent.start,
      content: block.content,
    }
  }

  if (block.type === 'subagent_thinking') {
    return {
      type: MothershipStreamV1EventType.text,
      lane: 'subagent',
      channel: MothershipStreamV1TextChannel.thinking,
      content: block.content,
    }
  }

  if (block.type === 'subagent_end') {
    return {
      type: MothershipStreamV1EventType.span,
      kind: MothershipStreamV1SpanPayloadKind.subagent,
      lifecycle: MothershipStreamV1SpanLifecycleEvent.end,
    }
  }

  if (block.type === 'stopped') {
    return {
      type: MothershipStreamV1EventType.complete,
      status: MothershipStreamV1CompletionStatus.cancelled,
    }
  }

  return {
    type: MothershipStreamV1EventType.text,
    channel: MothershipStreamV1TextChannel.assistant,
    content: block.content ?? block.text,
  }
}

function normalizeBlock(block: RawBlock): PersistedContentBlock {
  const result = isCanonicalBlock(block)
    ? normalizeCanonicalBlock(block)
    : normalizeLegacyBlock(block)
  if (typeof block.timestamp === 'number' && result.timestamp === undefined) {
    result.timestamp = block.timestamp
  }
  if (typeof block.endedAt === 'number' && result.endedAt === undefined) {
    result.endedAt = block.endedAt
  }
  if (block.parentToolCallId && result.parentToolCallId === undefined) {
    result.parentToolCallId = block.parentToolCallId
  }
  if (block.spanId && result.spanId === undefined) {
    result.spanId = block.spanId
  }
  if (block.parentSpanId && result.parentSpanId === undefined) {
    result.parentSpanId = block.parentSpanId
  }
  return result
}

function normalizeLegacyToolCall(tc: LegacyToolCall): PersistedContentBlock {
  const state = normalizeToolState(tc.status)
  return {
    type: MothershipStreamV1EventType.tool,
    phase: MothershipStreamV1ToolPhase.call,
    toolCall: {
      id: tc.id,
      name: tc.name,
      state,
      ...(tc.params ? { params: tc.params } : {}),
      ...(tc.result != null
        ? {
            result: {
              success: tc.status === MothershipStreamV1ToolOutcome.success,
              output: tc.result,
              ...(tc.error ? { error: tc.error } : {}),
            },
          }
        : {}),
      ...(tc.durationMs ? { durationMs: tc.durationMs } : {}),
    },
  }
}

function blocksContainTools(blocks: RawBlock[]): boolean {
  return blocks.some((b) => b.type === 'tool_call' || b.type === MothershipStreamV1EventType.tool)
}

function normalizeBlocks(rawBlocks: RawBlock[], messageContent: string): PersistedContentBlock[] {
  const blocks = rawBlocks.map(normalizeBlock)
  const hasAssistantText = blocks.some(
    (b) =>
      b.type === MothershipStreamV1EventType.text &&
      b.channel !== MothershipStreamV1TextChannel.thinking &&
      b.content?.trim()
  )
  if (!hasAssistantText && messageContent.trim()) {
    blocks.push({
      type: MothershipStreamV1EventType.text,
      channel: MothershipStreamV1TextChannel.assistant,
      content: messageContent,
    })
  }
  return blocks
}

export function normalizeMessage(raw: Record<string, unknown>): PersistedMessage {
  const msg: PersistedMessage = {
    id: (raw.id as string) ?? generateId(),
    role: (raw.role as 'user' | 'assistant') ?? 'assistant',
    content: (raw.content as string) ?? '',
    timestamp: (raw.timestamp as string) ?? new Date().toISOString(),
  }

  if (raw.requestId && typeof raw.requestId === 'string') {
    msg.requestId = raw.requestId
  }

  const rawBlocks = raw.contentBlocks as RawBlock[] | undefined
  const rawToolCalls = raw.toolCalls as LegacyToolCall[] | undefined
  const hasBlocks = Array.isArray(rawBlocks) && rawBlocks.length > 0
  const hasToolCalls = Array.isArray(rawToolCalls) && rawToolCalls.length > 0

  if (hasBlocks) {
    msg.contentBlocks = normalizeBlocks(rawBlocks!, msg.content)
    const contentBlocksAlreadyContainTools = blocksContainTools(rawBlocks!)
    if (hasToolCalls && !contentBlocksAlreadyContainTools) {
      msg.contentBlocks.push(...rawToolCalls!.map(normalizeLegacyToolCall))
    }
  } else if (hasToolCalls) {
    msg.contentBlocks = rawToolCalls!.map(normalizeLegacyToolCall)
    if (msg.content.trim()) {
      msg.contentBlocks.push({
        type: MothershipStreamV1EventType.text,
        channel: MothershipStreamV1TextChannel.assistant,
        content: msg.content,
      })
    }
  }

  const rawAttachments = raw.fileAttachments as PersistedFileAttachment[] | undefined
  if (Array.isArray(rawAttachments) && rawAttachments.length > 0) {
    msg.fileAttachments = rawAttachments
  }

  const rawContexts = raw.contexts as PersistedMessageContext[] | undefined
  if (Array.isArray(rawContexts) && rawContexts.length > 0) {
    msg.contexts = rawContexts.map((c) => ({
      kind: c.kind,
      label: c.label,
      ...(c.workflowId ? { workflowId: c.workflowId } : {}),
      ...(c.knowledgeId ? { knowledgeId: c.knowledgeId } : {}),
      ...(c.tableId ? { tableId: c.tableId } : {}),
      ...(c.fileId ? { fileId: c.fileId } : {}),
      ...(c.folderId ? { folderId: c.folderId } : {}),
      ...(c.chatId ? { chatId: c.chatId } : {}),
    }))
  }

  return msg
}