aiturk-hermes-ide/apps/desktop/src/app/chat/composer/hooks/use-composer-draft.ts

// Register the built-in draft providers with the suggestion bus (side-effect
// import — the bus itself is provider-agnostic). The repair provider is
// event-driven and registers through the gateway stream instead.
import '@/store/suggestion-providers/cron'
import '@/store/suggestion-providers/github'
import '@/store/suggestion-providers/mcp'
import '@/store/suggestion-providers/skill'

import { useAui, useAuiState, useComposerRuntime } from '@assistant-ui/react'
import { type RefObject, useCallback, useEffect, useLayoutEffect, useRef, useState } from 'react'

import { usePaneVisible } from '@/components/pane-shell/pane-visibility'
import { SLASH_COMMAND_RE } from '@/lib/chat-runtime'
import { sanitizeComposerInput } from '@/lib/composer-input-sanitize'
import {
  type ComposerAttachment,
  type ComposerDraftSyncMode,
  onComposerDraftSyncRequest,
  reloadPersistedDrafts,
  stashSessionDraft,
  takeSessionDraft
} from '@/store/composer'
import { isBrowsingHistory } from '@/store/composer-input-history'
import { clearDraftSuggestions, sampleComposerDraft } from '@/store/composer-suggestions'

import {
  cloneAttachments,
  DRAFT_PERSIST_DEBOUNCE_MS,
  isPendingDraftPersistCurrent,
  type QueueEditState
} from '../composer-utils'
import {
  type ComposerInsertMode,
  focusComposerInput,
  markActiveComposer,
  onComposerFocusRequest,
  onComposerInsertRefsRequest,
  onComposerInsertRequest,
  releaseActiveComposer
} from '../focus'
import { type InlineRefInput, insertInlineRefsIntoEditor } from '../inline-refs'
import {
  composerPlainText,
  normalizeComposerEditorDom,
  placeCaretEnd,
  REF_RE,
  renderComposerContents
} from '../rich-editor'
import { useComposerScope } from '../scope'
import type { ChatBarProps } from '../types'

interface UseComposerDraftArgs {
  activeQueueSessionKey: string | null
  focusKey: ChatBarProps['focusKey']
  inputDisabled: boolean
  queueEditRef: RefObject<QueueEditState | null>
  sessionId: string | null | undefined
}

/**
 * The composer's draft engine — the detached source-of-truth spine. The live
 * text lives in the contentEditable DOM + `draftRef`; React only sees coarse
 * edge selectors, so typing never re-renders the chrome. Owns the imperative
 * composer-runtime subscription (draftRef mirror + external repaint + debounced
 * per-session stash), the edit primitives (append/insert/inline-refs), focus,
 * and per-session load/clear/stash/restore. The contentEditable *event*
 * handlers stay in ChatBar (they bridge into the trigger engine) and drive the
 * primitives exposed here.
 */
export function useComposerDraft({
  activeQueueSessionKey,
  focusKey,
  inputDisabled,
  queueEditRef,
  sessionId
}: UseComposerDraftArgs) {
  const aui = useAui()
  const composerRuntime = useComposerRuntime()
  const paneVisible = usePaneVisible()
  // Which composer this is on the focus bus + which attachment set it owns.
  const { attachments: attachmentScope, target } = useComposerScope()

  // Coarse edges only — these flip rarely (empty↔non-empty, the `?` help sigil,
  // steerable-vs-slash), so typing within a line costs no render.
  const hasText = useAuiState(s => s.composer.text.trim().length > 0)
  const isHelpHint = useAuiState(s => s.composer.text === '?')

  const isSteerableText = useAuiState(s => {
    const trimmed = s.composer.text.trim()

    return trimmed.length > 0 && !SLASH_COMMAND_RE.test(trimmed)
  })

  // assistant-ui's composer mutators throw when the core isn't bound yet (a
  // startup/thread-swap window); the DOM + draftRef hold the text and the
  // subscription reconciles once it binds, so swallow the premature write.
  const setComposerText = useCallback(
    (value: string) => {
      try {
        aui.composer().setText(value)
      } catch {
        // Composer core not bound yet — DOM/draftRef carry the text.
      }
    },
    [aui]
  )

  const editorRef = useRef<HTMLDivElement | null>(null)
  const draftRef = useRef('')
  const pendingDraftPersistRef = useRef<{ scope: string | null; text: string } | null>(null)
  const draftPersistTimerRef = useRef<number | undefined>(undefined)
  const activeQueueSessionKeyRef = useRef(activeQueueSessionKey)
  activeQueueSessionKeyRef.current = activeQueueSessionKey
  // Owned only by the swap effect below — unlike activeQueueSessionKeyRef this
  // does NOT update on every render, so it always reflects the session whose
  // text is actually loaded in the editor. Async work (debounce timers,
  // pagehide flush) must persist against this, not the render-time ref, or a
  // session switch mid-flight files one session's draft under another's key
  // (#54527).
  const draftScopeRef = useRef(activeQueueSessionKey)
  const sessionIdRef = useRef(sessionId)
  sessionIdRef.current = sessionId
  const queueEditStateRef = useRef<QueueEditState | null>(queueEditRef.current)
  queueEditStateRef.current = queueEditRef.current

  const [focusRequestId, setFocusRequestId] = useState(0)

  const focusInput = useCallback(() => {
    focusComposerInput(editorRef.current)
    markActiveComposer(target)
  }, [target])

  const requestMainFocus = useCallback(() => {
    setFocusRequestId(id => id + 1)
  }, [])

  // The single write path for programmatic draft mutations: mirror → AUI state →
  // repaint the editor (caret to end). Repaints even while focused — inserts /
  // restores run mid-focus, and the runtime sync only repaints an unfocused
  // editor — so the visible text never lags the store.
  const paintDraft = useCallback(
    (next: string, focus = true) => {
      draftRef.current = next
      setComposerText(next)

      const editor = editorRef.current

      if (editor) {
        renderComposerContents(editor, next, { trailingCommitted: true })
        placeCaretEnd(editor)
      }

      if (focus) {
        requestMainFocus()
      }
    },
    [requestMainFocus, setComposerText]
  )

  const appendExternalText = useCallback(
    (text: string, mode: ComposerInsertMode) => {
      const value = text.trim()

      if (!value) {
        return
      }

      // 'prefix' puts the value at the START of the draft — slash commands
      // (the skill-suggestion pill) only route when they lead the message.
      if (mode === 'prefix') {
        const rest = draftRef.current.trimStart()

        paintDraft(`${value} ${rest}`.trimEnd())

        return
      }

      const base = mode === 'inline' ? draftRef.current.trimEnd() : draftRef.current
      const sep = mode === 'inline' ? (base ? ' ' : '') : base && !base.endsWith('\n') ? '\n\n' : ''

      paintDraft(`${base}${sep}${value}`)
    },
    [paintDraft]
  )

  // Keep-alive tabs keep this composer mounted. A background session whose
  // turn finished or recovered from a reconnect would otherwise re-run this
  // effect and steal the caret. usePaneVisible defaults true outside a tab
  // stack, so tiles, pop-outs, and secondary windows still auto-focus.
  useEffect(() => {
    if (!inputDisabled && paneVisible) {
      focusInput()
    }
  }, [focusInput, focusKey, focusRequestId, inputDisabled, paneVisible])

  // The mirror of the `markActiveComposer` above: give the key back when this
  // composer goes away (a session tile closing, a pane unmounting). Covers both
  // claim sites for this composer — `focusInput` here and ChatBar's `onFocus` —
  // since they mark the same scope target. Without it `'active'` keeps
  // resolving to a dead tile and every routed focus/insert request is dropped.
  // (Heal-to-visible in focus.ts covers the keep-alive-tab case where the pane
  // stays mounted behind the front tab; this covers true unmounts.)
  useEffect(() => () => releaseActiveComposer(target), [target])

  useEffect(() => {
    if (inputDisabled) {
      return undefined
    }

    const offFocus = onComposerFocusRequest(({ target: requested, typeChar }) => {
      if (requested !== target) {
        return
      }

      // Type-to-focus appends at end; bare Enter just focuses.
      if (typeChar) {
        paintDraft(`${draftRef.current}${typeChar}`, true)

        return
      }

      setFocusRequestId(id => id + 1)
    })

    const offInsert = onComposerInsertRequest(({ mode, target: requested, text }) => {
      if (requested === target) {
        appendExternalText(text, mode)
      }
    })

    return () => {
      offFocus()
      offInsert()
    }
  }, [appendExternalText, inputDisabled, paintDraft, target])

  const stashAt = (scope: string | null, text = draftRef.current, attachments = attachmentScope.$attachments.get()) =>
    stashSessionDraft(scope, text, attachments)

  const loadIntoComposer = (text: string, attachments: ComposerAttachment[]) => {
    // Diagnostic breadcrumb for #59305-class reports: identifies WHAT kind of
    // state got restored into the composer (session switch, queue-edit
    // restore, history browse) without logging any raw content. REF_RE has the
    // global flag — testing against a throwaway clone avoids mutating the
    // shared instance's lastIndex, which would otherwise corrupt this check on
    // the next call.
    if (attachments.length > 0 || new RegExp(REF_RE.source, REF_RE.flags).test(text)) {
      console.debug('[composer-rehydrate]', {
        attachmentCount: attachments.length,
        attachmentKinds: attachments.map(a => a.kind),
        hasTextRefs: new RegExp(REF_RE.source, REF_RE.flags).test(text),
        scope: activeQueueSessionKeyRef.current
      })
    }

    attachmentScope.$attachments.set(cloneAttachments(attachments))
    paintDraft(text, false)
  }

  const clearDraft = useCallback(() => {
    setComposerText('')
    draftRef.current = ''

    if (editorRef.current) {
      renderComposerContents(editorRef.current, '')
      placeCaretEnd(editorRef.current)
    }
  }, [setComposerText])

  // Read the editor's current plain text into draftRef + composer state. This
  // closes the "queued rAF flush hasn't run yet" window so scope-swap/pagehide
  // persistence captures the latest keystrokes.
  const syncDraftFromEditor = useCallback(() => {
    const editor = editorRef.current

    if (!editor) {
      return draftRef.current
    }

    // Same normalize-then-sanitize the rAF flush does. An emptied editor still
    // holds the placeholder <br> that keeps the contenteditable from collapsing
    // to a sliver, and that serializes as "\n" — so an editor the user just
    // cleared would otherwise stash a one-newline draft and come back non-empty.
    normalizeComposerEditorDom(editor)

    const text = sanitizeComposerInput(composerPlainText(editor))

    if (text !== draftRef.current) {
      draftRef.current = text
      setComposerText(text)
    }

    return text
  }, [setComposerText])

  // Imperative draft sync — the spine of the "work only when work is to be
  // performed" model. Subscribing to the composer runtime directly (not
  // `useAuiState(text)` + a `[draft]` effect) keeps per-keystroke text out of
  // React, so typing never re-renders the chrome. On each change we (1) mirror
  // text into draftRef, (2) repaint the editor only when the change came from
  // OUTSIDE it (programmatic clear/restore/insert; the focused editor is the
  // source otherwise), and (3) schedule the debounced per-session stash.
  // Browsing history / editing a queued prompt suppress the stash so recalled
  // text never clobbers the draft.
  // eslint-disable-next-line no-restricted-syntax -- legitimate non-atom ref write (see eslint rule comment)
  useEffect(() => {
    const sync = () => {
      const text = composerRuntime.getState().text
      draftRef.current = text
      // Composer suggestion pills for THIS session's draft (debounced +
      // change-gated in the bus — this is just a timer reset).
      sampleComposerDraft(sessionIdRef.current ?? null, text)

      const editor = editorRef.current

      if (editor && document.activeElement !== editor && composerPlainText(editor) !== text) {
        renderComposerContents(editor, text, { trailingCommitted: true })
      }

      if (isBrowsingHistory(sessionIdRef.current) || queueEditRef.current) {
        return
      }

      const scope = draftScopeRef.current
      const entry = { scope, text }
      pendingDraftPersistRef.current = entry
      window.clearTimeout(draftPersistTimerRef.current)
      draftPersistTimerRef.current = window.setTimeout(() => {
        // Integrity guard (defense-in-depth, #54527): only commit if this is
        // still the pending write on file. A session swap or a newer
        // keystroke clears/replaces it before firing in the normal case; this
        // catches any future call site that skips that bookkeeping instead of
        // silently filing text under the wrong session.
        if (!isPendingDraftPersistCurrent(pendingDraftPersistRef.current, entry)) {
          return
        }

        pendingDraftPersistRef.current = null
        stashAt(scope, text)
      }, DRAFT_PERSIST_DEBOUNCE_MS)
    }

    const unsubscribe = composerRuntime.subscribe(sync)

    return () => {
      unsubscribe()
      window.clearTimeout(draftPersistTimerRef.current)
    }
  }, [composerRuntime, queueEditRef])

  const insertText = (text: string) => {
    const base = draftRef.current
    const sep = base && !base.endsWith('\n') ? '\n' : ''

    paintDraft(`${base}${sep}${text}`)
  }

  // insertInlineRefs mutates the editor in place (chips), so it can't go through
  // paintDraft's re-render — it mirrors the resulting plain text and refocuses.
  const insertInlineRefs = (refs: InlineRefInput[]) => {
    const editor = editorRef.current

    if (!editor) {
      return false
    }

    const nextDraft = insertInlineRefsIntoEditor(editor, refs)

    if (nextDraft === null) {
      return false
    }

    draftRef.current = nextDraft
    setComposerText(nextDraft)
    requestMainFocus()

    return true
  }

  // Latest-closure ref so the once-only subscription always calls the current
  // insertInlineRefs without re-subscribing every render.
  const insertInlineRefsRef = useRef(insertInlineRefs)
  insertInlineRefsRef.current = insertInlineRefs

  useEffect(() => {
    return onComposerInsertRefsRequest(({ refs, target: requested }) => {
      if (requested === target) {
        insertInlineRefsRef.current(refs)
      }
    })
  }, [target])

  // Per-thread draft swap — the composer's only session coupling. Lifecycle
  // never clears composer state; this effect alone stashes on leave, restores
  // on enter. Keyed writes are idempotent, so no skip-sentinel.
  //
  // MUST be a layout effect, not a passive one: it swaps attachmentScope's
  // module-level $attachments atom, and a passive effect fires only after the
  // browser paints the new session's view — leaving a window where the DOM
  // already shows session B while $attachments (and therefore ChatBar's
  // `attachments` prop) still holds session A's chips. A submit fired in that
  // window (e.g. a fast session switch immediately followed by Enter) would
  // ship A's attachments into B's turn (#59305). useLayoutEffect closes the
  // window by running before paint.

  useLayoutEffect(() => {
    // A pending debounce timer from the outgoing session is now stale — its
    // scope was correct when scheduled, but the authoritative stash below
    // (and the cleanup on the way out) already covers that text. Letting it
    // fire later would just clobber with an older snapshot.
    window.clearTimeout(draftPersistTimerRef.current)
    pendingDraftPersistRef.current = null
    draftScopeRef.current = activeQueueSessionKey

    const { attachments, text } = takeSessionDraft(activeQueueSessionKey)
    loadIntoComposer(text, attachments)

    return () => {
      const latestText = syncDraftFromEditor()
      const editing = queueEditStateRef.current

      if (editing?.sessionKey === activeQueueSessionKey) {
        stashAt(activeQueueSessionKey, editing.draft, editing.attachments)
      } else if (!isBrowsingHistory(sessionId)) {
        stashAt(activeQueueSessionKey, latestText)
      }

      // Withdraw the outgoing session's draft suggestions (and any pending
      // sample timer). The incoming session re-earns its own from the draft
      // restore above — without this a leaving session's "Add GitHub" pill
      // lingers in the map and re-appears stale on the way back.
      clearDraftSuggestions(sessionIdRef.current)
    }
  }, [activeQueueSessionKey]) // eslint-disable-line react-hooks/exhaustive-deps

  // The HUD handoff's two verbs. Entering HUD mode flushes this editor's text
  // into the shared stash so the HUD's composer boots with it; leaving repaints
  // from the stash so whatever the HUD typed (or sent, clearing it) is what the
  // app window shows. The per-session swap effect above can't cover either one:
  // the session scope doesn't change, so it never re-consults the stash.
  const syncDraft = (mode: ComposerDraftSyncMode) => {
    if (mode === 'flush') {
      window.clearTimeout(draftPersistTimerRef.current)
      pendingDraftPersistRef.current = null
      stashAt(draftScopeRef.current, syncDraftFromEditor())

      return
    }

    reloadPersistedDrafts()
    const stashed = takeSessionDraft(draftScopeRef.current)
    loadIntoComposer(stashed.text, stashed.attachments)
  }

  const syncDraftRef = useRef(syncDraft)
  syncDraftRef.current = syncDraft

  useEffect(
    () =>
      onComposerDraftSyncRequest(({ mode, target: requested }) => {
        if (requested === target) {
          syncDraftRef.current(mode)
        }
      }),
    [target]
  )

  // pagehide is load-bearing: React skips effect cleanups on reload, so Cmd+R
  // inside the debounce/rAF window would drop trailing keystrokes without this.
  // eslint-disable-next-line no-restricted-syntax -- legitimate non-atom ref write (see eslint rule comment)
  useEffect(() => {
    const flushPendingDraftPersist = () => {
      const scope = draftScopeRef.current
      const editing = queueEditStateRef.current

      if (editing?.sessionKey === scope || isBrowsingHistory(sessionIdRef.current)) {
        return
      }

      const latestText = syncDraftFromEditor()
      pendingDraftPersistRef.current = null
      stashAt(scope, latestText)
    }

    window.addEventListener('pagehide', flushPendingDraftPersist)

    return () => {
      window.removeEventListener('pagehide', flushPendingDraftPersist)
      flushPendingDraftPersist()
    }
  }, [syncDraftFromEditor])

  return {
    activeQueueSessionKeyRef,
    clearDraft,
    draftRef,
    editorRef,
    focusInput,
    hasText,
    insertInlineRefs,
    insertText,
    isHelpHint,
    isSteerableText,
    loadIntoComposer,
    requestMainFocus,
    sessionIdRef,
    setComposerText,
    stashAt,
    syncDraftFromEditor
  }
}