openchamber/packages/ui/src/sync/input-store.ts

/**
 * Input Store — pending input text, synthetic parts, and attached files.
 * Extracted from session-ui-store for subscription isolation.
 */

import { create } from "zustand"
import type { ContextPartMetadata } from '@/lib/messages/contextParts'
import type { AttachedFile } from "@/stores/types/sessionTypes"
import { prepareAttachmentFiles } from "./attachment-files"
import { getChatDraftIdentityKey, type ChatDraftIdentity } from "@/lib/chatDraftPersistence"

const FILE_URI_PREFIX = "file://"
const MAX_ATTACHMENT_PREPARATION_ATTEMPTS = 3
const pendingVSCodeSelectionKeys = new Set<string>()
let attachmentReadGeneration = 0

const encodeFilePath = (filepath: string): string => {
  let normalized = filepath.replace(/\\/g, "/")
  if (/^[A-Za-z]:/.test(normalized)) {
    normalized = `/${normalized}`
  }
  return normalized
    .split("/")
    .map((segment, index) => {
      if (index === 1 && /^[A-Za-z]:$/.test(segment)) return segment
      return encodeURIComponent(segment)
    })
    .join("/")
}

const toFileUrl = (filepath: string): string => {
  const normalized = filepath.replace(/\\/g, "/").trim()
  if (normalized.toLowerCase().startsWith(FILE_URI_PREFIX)) {
    return normalized
  }
  return `${FILE_URI_PREFIX}${encodeFilePath(normalized)}`
}

const getVSCodeSelectionKey = (path: string, filename: string): string => `${path}\u0000${filename}`

const hasGeneratedFilenameCollision = (filenames: string[], attachedFiles: AttachedFile[]): boolean => {
  if (filenames.length === 0) return false
  const attachedFilenames = new Set(attachedFiles.map((attachment) => attachment.filename.toLowerCase()))
  return filenames.some((filename) => attachedFilenames.has(filename.toLowerCase()))
}

const readFileAsDataUrl = (file: File, mime: string): Promise<string> => new Promise((resolve, reject) => {
  const reader = new FileReader()
  reader.onload = () => {
    const value = typeof reader.result === "string" ? reader.result : ""
    const commaIndex = value.indexOf(",")
    resolve(commaIndex === -1 ? value : `data:${mime};base64,${value.slice(commaIndex + 1)}`)
  }
  reader.onerror = () => reject(reader.error ?? new Error("Failed to read file"))
  reader.onabort = () => reject(new Error("File read aborted"))
  reader.readAsDataURL(file)
})

export const prepareLocalAttachments = async (
  file: File,
  reservedFilenames: Iterable<string> = [],
): Promise<AttachedFile[] | undefined> => {
  const preparedOrPending = prepareAttachmentFiles(file, reservedFilenames)
  const preparedFiles = preparedOrPending instanceof Promise ? await preparedOrPending : preparedOrPending
  if (!preparedFiles || preparedFiles.length === 0) return

  const sourceDocumentId = preparedFiles.length > 1
    ? `${Date.now()}-${Math.random().toString(36).slice(2)}`
    : undefined
  const attachedFiles: AttachedFile[] = []
  for (const prepared of preparedFiles) {
    const dataUrl = await readFileAsDataUrl(prepared.file, prepared.mimeType)
    if (!dataUrl) return
    attachedFiles.push({
      id: `${Date.now()}-${Math.random().toString(36).slice(2)}`,
      file: prepared.file,
      dataUrl,
      mimeType: prepared.mimeType,
      filename: prepared.file.name,
      size: prepared.file.size,
      source: "local",
      sourceDocumentId,
    })
  }
  return attachedFiles
}

const getDataUrlByteSize = (url: string): number => {
  if (!url.startsWith("data:")) return 0
  const commaIndex = url.indexOf(",")
  if (commaIndex < 0) return 0
  const metadata = url.slice(0, commaIndex).toLowerCase()
  const payload = url.slice(commaIndex + 1)
  if (!metadata.endsWith(";base64")) return 0
  let padding = 0
  if (payload.endsWith("==")) {
    padding = 2
  } else if (payload.endsWith("=")) {
    padding = 1
  }
  return Math.max(0, Math.floor((payload.length * 3) / 4) - padding)
}

const isSameVSCodeActiveEditorFile = (a: VSCodeActiveEditorFile | null, b: VSCodeActiveEditorFile | null): boolean => {
  if (a === b) return true
  if (!a || !b) return false
  return a.filePath === b.filePath
    && a.fileName === b.fileName
    && a.relativePath === b.relativePath
    && a.fileSize === b.fileSize
    && a.selection?.startLine === b.selection?.startLine
    && a.selection?.endLine === b.selection?.endLine
    && a.selection?.text === b.selection?.text
}

export type SyntheticContextPart = {
  text: string
  attachments?: AttachedFile[]
  synthetic?: boolean
  metadata?: ContextPartMetadata
}

type PendingBtwComposerRequest = {
  parentSessionId: string
  text: string
}

export type VSCodeActiveEditorFile = {
  filePath: string
  fileName: string
  relativePath: string
  fileSize: number | null
  selection: { startLine: number; endLine: number; text: string } | null
}

export type InputState = {
  pendingComposerRestore: {
    target: ChatDraftIdentity
    text: string
    files: Array<{ url: string; mimeType: string; filename: string }>
  } | null
  consumePendingComposerRestore: (target: ChatDraftIdentity | null) => InputState["pendingComposerRestore"]
  pendingInputText: string | null
  pendingInputMode: "replace" | "append" | "append-inline"
  pendingSyntheticParts: SyntheticContextPart[] | null
  /**
   * Text a draft preset chip asked to submit immediately. Set by surfaces that
   * render the chips outside ChatInput (e.g. under the welcome message on
   * narrow layouts); consumed by ChatInput, which owns the command-aware submit.
   */
  pendingPresetSubmit: { text: string; type: "command" | "skill" } | null
  pendingBtwComposerRequest: PendingBtwComposerRequest | null
  attachedFiles: AttachedFile[]
  activeEditorFile: VSCodeActiveEditorFile | null

  setPendingInputText: (text: string | null, mode?: "replace" | "append" | "append-inline") => void
  consumePendingInputText: () => { text: string; mode: "replace" | "append" | "append-inline" } | null
  requestPresetSubmit: (text: string, type: "command" | "skill") => void
  consumePendingPresetSubmit: () => { text: string; type: "command" | "skill" } | null
  requestBtwComposer: (request: PendingBtwComposerRequest) => void
  consumePendingBtwComposerRequest: (parentSessionId: string | null) => PendingBtwComposerRequest | null
  setPendingSyntheticParts: (parts: SyntheticContextPart[] | null) => void
  consumePendingSyntheticParts: () => SyntheticContextPart[] | null
  addAttachedFile: (file: File) => Promise<boolean>
  removeAttachedFile: (id: string) => void
  setAttachedFiles: (files: AttachedFile[]) => void
  clearAttachedFiles: () => void
  addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => void
  addVSCodeSelectionAttachment: (path: string, file: File) => Promise<void>
  setActiveEditorFile: (file: VSCodeActiveEditorFile | null) => void
  /** Add attachments restored from a reverted message (file already on server) */
  addRestoredAttachment: (file: { url: string; mimeType: string; filename: string }) => void
}

export const useInputStore = create<InputState>()((set, get) => ({
  pendingComposerRestore: null,
  consumePendingComposerRestore: (target) => {
    const pending = get().pendingComposerRestore
    if (!pending || !target || getChatDraftIdentityKey(pending.target) !== getChatDraftIdentityKey(target)) return null
    set({ pendingComposerRestore: null })
    return pending
  },
  pendingInputText: null,
  pendingInputMode: "replace",
  pendingSyntheticParts: null,
  pendingPresetSubmit: null,
  pendingBtwComposerRequest: null,
  attachedFiles: [],
  activeEditorFile: null,

  setPendingInputText: (text, mode = "replace") =>
    set({ pendingInputText: text, pendingInputMode: mode }),

  consumePendingInputText: () => {
    const { pendingInputText, pendingInputMode } = get()
    if (pendingInputText === null) return null
    set({ pendingInputText: null, pendingInputMode: "replace" })
    return { text: pendingInputText, mode: pendingInputMode }
  },

  requestPresetSubmit: (text, type) => set({ pendingPresetSubmit: { text, type } }),

  consumePendingPresetSubmit: () => {
    const { pendingPresetSubmit } = get()
    if (pendingPresetSubmit === null) return null
    set({ pendingPresetSubmit: null })
    return pendingPresetSubmit
  },

  requestBtwComposer: (request) => set({ pendingBtwComposerRequest: request }),

  consumePendingBtwComposerRequest: (parentSessionId) => {
    const request = get().pendingBtwComposerRequest
    if (!request || request.parentSessionId !== parentSessionId) return null
    set({ pendingBtwComposerRequest: null })
    return request
  },

  setPendingSyntheticParts: (parts) => set({ pendingSyntheticParts: parts }),

  consumePendingSyntheticParts: () => {
    const { pendingSyntheticParts } = get()
    if (pendingSyntheticParts !== null) {
      set({ pendingSyntheticParts: null })
    }
    return pendingSyntheticParts
  },

  addAttachedFile: async (file: File) => {
    const generation = attachmentReadGeneration
    for (let attempt = 0; attempt < MAX_ATTACHMENT_PREPARATION_ATTEMPTS; attempt += 1) {
      const reservedFilenames = get().attachedFiles.map((attachment) => attachment.filename)
      let attachedFiles: AttachedFile[] | undefined
      try {
        attachedFiles = await prepareLocalAttachments(file, reservedFilenames)
      } catch {
        return false
      }
      if (!attachedFiles || generation !== attachmentReadGeneration) return false

      const generatedFilenames = attachedFiles.slice(1).map((attachment) => attachment.filename)
      if (hasGeneratedFilenameCollision(generatedFilenames, get().attachedFiles)) continue

      set((state) => ({ attachedFiles: [...state.attachedFiles, ...attachedFiles] }))
      return true
    }
    return false
  },

  removeAttachedFile: (id) =>
    set((s) => {
      const target = s.attachedFiles.find((f) => f.id === id)
      if (target?.sourceDocumentId) {
        return { attachedFiles: s.attachedFiles.filter((f) => f.sourceDocumentId !== target.sourceDocumentId) }
      }
      return { attachedFiles: s.attachedFiles.filter((f) => f.id !== id) }
    }),

  setAttachedFiles: (files) => {
    attachmentReadGeneration += 1
    set({ attachedFiles: files })
  },

  clearAttachedFiles: () => {
    attachmentReadGeneration += 1
    set({ attachedFiles: [] })
  },

  addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => {
    const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`
    const isDuplicate = get().attachedFiles.some(
      (f) => f.source === 'vscode' && f.vscodeSource === 'file' && (f.vscodePath || '') === path
    )
    if (isDuplicate) return
    const dataUrl = toFileUrl(path)
    // `file://` URLs are the same contract used by server-source attachments.
    // The submission path passes `dataUrl` as `url` directly to the OpenCode
    // server, which resolves `file://` paths natively. No base64 encoding needed.
    const attached: AttachedFile = {
      id,
      file: new File([], name, { type: 'text/plain' }),
      dataUrl,
      mimeType: 'text/plain',
      filename: name,
      size: fileSize || 0,
      source: 'vscode',
      vscodePath: path,
      vscodeSource: 'file',
    }
    set((s) => ({ attachedFiles: [...s.attachedFiles, attached] }))
  },

  addVSCodeSelectionAttachment: async (path: string, file: File) => {
    const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`
    const generation = attachmentReadGeneration
    const selectionKey = getVSCodeSelectionKey(path, file.name)
    const isDuplicate = get().attachedFiles.some(
      (f) => f.source === 'vscode' && f.vscodeSource === 'selection' && f.filename === file.name && f.vscodePath === path
    )
    if (isDuplicate || pendingVSCodeSelectionKeys.has(selectionKey)) return
    pendingVSCodeSelectionKeys.add(selectionKey)
    let dataUrl: string
    try {
      dataUrl = await readFileAsDataUrl(file, file.type)
    } catch {
      return
    } finally {
      pendingVSCodeSelectionKeys.delete(selectionKey)
    }
    if (generation !== attachmentReadGeneration) return
    const attached: AttachedFile = {
      id,
      file,
      dataUrl,
      mimeType: file.type,
      filename: file.name,
      size: file.size,
      source: 'vscode',
      vscodePath: path,
      vscodeSource: 'selection',
    }
    set((s) => ({ attachedFiles: [...s.attachedFiles, attached] }))
  },

  setActiveEditorFile: (file) => {
    if (isSameVSCodeActiveEditorFile(get().activeEditorFile, file)) return
    set({ activeEditorFile: file })
  },

  addRestoredAttachment: ({ url, mimeType, filename }) => {
    const id = `restored-${Date.now()}-${Math.random().toString(36).slice(2)}`
    // Use "local" source so the file renders in AttachedFilesList.
    // Set serverPath to the URL so ImagePreview can use it as the img src
    // when dataUrl is not a data: URL. sanitizeAttachmentsForSend leaves
    // dataUrl alone for non-server sources, so the URL stays intact on send.
    const attached: AttachedFile = {
      id,
      file: new File([], filename, { type: mimeType }),
      dataUrl: url,
      mimeType,
      filename,
      size: getDataUrlByteSize(url),
      source: "local",
      serverPath: url,
    }
    set((s) => ({ attachedFiles: [...s.attachedFiles, attached] }))
  },
}))