services/PromptSuggestion/speculation.ts

🖥️ Source preview

import { randomUUID } from 'crypto'
import { rm } from 'fs'
import { appendFile, copyFile, mkdir } from 'fs/promises'
import { dirname, isAbsolute, join, relative } from 'path'
import { getCwdState } from '../../bootstrap/state.js'
import type { CompletionBoundary } from '../../state/AppStateStore.js'
import {
  type AppState,
  IDLE_SPECULATION_STATE,
  type SpeculationResult,
  type SpeculationState,
} from '../../state/AppStateStore.js'
import { commandHasAnyCd } from '../../tools/BashTool/bashPermissions.js'
import { checkReadOnlyConstraints } from '../../tools/BashTool/readOnlyValidation.js'
import type { SpeculationAcceptMessage } from '../../types/logs.js'
import type { Message } from '../../types/message.js'
import { createChildAbortController } from '../../utils/abortController.js'
import { count } from '../../utils/array.js'
import { getGlobalConfig } from '../../utils/config.js'
import { logForDebugging } from '../../utils/debug.js'
import { errorMessage } from '../../utils/errors.js'
import {
  type FileStateCache,
  mergeFileStateCaches,
  READ_FILE_STATE_CACHE_SIZE,
} from '../../utils/fileStateCache.js'
import {
  type CacheSafeParams,
  createCacheSafeParams,
  runForkedAgent,
} from '../../utils/forkedAgent.js'
import { formatDuration, formatNumber } from '../../utils/format.js'
import type { REPLHookContext } from '../../utils/hooks/postSamplingHooks.js'
import { logError } from '../../utils/log.js'
import type { SetAppState } from '../../utils/messageQueueManager.js'
import {
  createSystemMessage,
  createUserMessage,
  INTERRUPT_MESSAGE,
  INTERRUPT_MESSAGE_FOR_TOOL_USE,
} from '../../utils/messages.js'
import { getClaudeTempDir } from '../../utils/permissions/filesystem.js'
import { extractReadFilesFromMessages } from '../../utils/queryHelpers.js'
import { getTranscriptPath } from '../../utils/sessionStorage.js'
import { jsonStringify } from '../../utils/slowOperations.js'
import {
  type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
  logEvent,
} from '../analytics/index.js'
import {
  generateSuggestion,
  getPromptVariant,
  getSuggestionSuppressReason,
  logSuggestionSuppressed,
  shouldFilterSuggestion,
} from './promptSuggestion.js'

const MAX_SPECULATION_TURNS = 20
const MAX_SPECULATION_MESSAGES = 100

const WRITE_TOOLS = new Set(['Edit', 'Write', 'NotebookEdit'])
const SAFE_READ_ONLY_TOOLS = new Set([
  'Read',
  'Glob',
  'Grep',
  'ToolSearch',
  'LSP',
  'TaskGet',
  'TaskList',
])

function safeRemoveOverlay(overlayPath: string): void {
  rm(
    overlayPath,
    { recursive: true, force: true, maxRetries: 3, retryDelay: 100 },
    () => {},
  )
}

function getOverlayPath(id: string): string {
  return join(getClaudeTempDir(), 'speculation', String(process.pid), id)
}

function denySpeculation(
  message: string,
  reason: string,
): {
  behavior: 'deny'
  message: string
  decisionReason: { type: 'other'; reason: string }
} {
  return {
    behavior: 'deny',
    message,
    decisionReason: { type: 'other', reason },
  }
}

async function copyOverlayToMain(
  overlayPath: string,
  writtenPaths: Set<string>,
  cwd: string,
): Promise<boolean> {
  let allCopied = true
  for (const rel of writtenPaths) {
    const src = join(overlayPath, rel)
    const dest = join(cwd, rel)
    try {
      await mkdir(dirname(dest), { recursive: true })
      await copyFile(src, dest)
    } catch {
      allCopied = false
      logForDebugging(`[Speculation] Failed to copy ${rel} to main`)
    }
  }
  return allCopied
}

export type ActiveSpeculationState = Extract<
  SpeculationState,
  { status: 'active' }
>

function logSpeculation(
  id: string,
  outcome: 'accepted' | 'aborted' | 'error',
  startTime: number,
  suggestionLength: number,
  messages: Message[],
  boundary: CompletionBoundary | null,
  extras?: Record<string, string | number | boolean | undefined>,
): void {
  logEvent('tengu_speculation', {
    speculation_id:
      id as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
    outcome:
      outcome as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
    duration_ms: Date.now() - startTime,
    suggestion_length: suggestionLength,
    tools_executed: countToolsInMessages(messages),
    completed: boundary !== null,
    boundary_type: boundary?.type as
      | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
      | undefined,
    boundary_tool: getBoundaryTool(boundary) as
      | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
      | undefined,
    boundary_detail: getBoundaryDetail(boundary) as
      | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
      | undefined,
    ...extras,
  })
}

function countToolsInMessages(messages: Message[]): number {
  const blocks = messages
    .filter(isUserMessageWithArrayContent)
    .flatMap(m => m.message.content)
    .filter(
      (b): b is { type: string; is_error?: boolean } =>
        typeof b === 'object' && b !== null && 'type' in b,
    )
  return count(blocks, b => b.type === 'tool_result' && !b.is_error)
}

function getBoundaryTool(
  boundary: CompletionBoundary | null,
): string | undefined {
  if (!boundary) return undefined
  switch (boundary.type) {
    case 'bash':
      return 'Bash'
    case 'edit':
    case 'denied_tool':
      return boundary.toolName
    case 'complete':
      return undefined
  }
}

function getBoundaryDetail(
  boundary: CompletionBoundary | null,
): string | undefined {
  if (!boundary) return undefined
  switch (boundary.type) {
    case 'bash':
      return boundary.command.slice(0, 200)
    case 'edit':
      return boundary.filePath
    case 'denied_tool':
      return boundary.detail
    case 'complete':
      return undefined
  }
}

function isUserMessageWithArrayContent(
  m: Message,
): m is Message & { message: { content: unknown[] } } {
  return m.type === 'user' && 'message' in m && Array.isArray(m.message.content)
}

export function prepareMessagesForInjection(messages: Message[]): Message[] {
  // Find tool_use IDs that have SUCCESSFUL results (not errors/interruptions)
  // Pending tool_use blocks (no result) and interrupted ones will be stripped
  type ToolResult = {
    type: 'tool_result'
    tool_use_id: string
    is_error?: boolean
    content?: unknown
  }
  const isToolResult = (b: unknown): b is ToolResult =>
    typeof b === 'object' &&
    b !== null &&
    (b as ToolResult).type === 'tool_result' &&
    typeof (b as ToolResult).tool_use_id === 'string'
  const isSuccessful = (b: ToolResult) =>
    !b.is_error &&
    !(
      typeof b.content === 'string' &&
      b.content.includes(INTERRUPT_MESSAGE_FOR_TOOL_USE)
    )

  const toolIdsWithSuccessfulResults = new Set(
    messages
      .filter(isUserMessageWithArrayContent)
      .flatMap(m => m.message.content)
      .filter(isToolResult)
      .filter(isSuccessful)
      .map(b => b.tool_use_id),
  )

  const keep = (b: {
    type: string
    id?: string
    tool_use_id?: string
    text?: string
  }) =>
    b.type !== 'thinking' &&
    b.type !== 'redacted_thinking' &&
    !(b.type === 'tool_use' && !toolIdsWithSuccessfulResults.has(b.id!)) &&
    !(
      b.type === 'tool_result' &&
      !toolIdsWithSuccessfulResults.has(b.tool_use_id!)
    ) &&
    // Abort during speculation yields a standalone interrupt user message
    // (query.ts createUserInterruptionMessage). Strip it so it isn't surfaced
    // to the model as real user input.
    !(
      b.type === 'text' &&
      (b.text === INTERRUPT_MESSAGE ||
        b.text === INTERRUPT_MESSAGE_FOR_TOOL_USE)
    )

  return messages
    .map(msg => {
      if (!('message' in msg) || !Array.isArray(msg.message.content)) return msg
      const content = msg.message.content.filter(keep)
      if (content.length === msg.message.content.length) return msg
      if (content.length === 0) return null
      // Drop messages where all remaining blocks are whitespace-only text
      // (API rejects these with 400: "text content blocks must contain non-whitespace text")
      const hasNonWhitespaceContent = content.some(
        (b: { type: string; text?: string }) =>
          b.type !== 'text' || (b.text !== undefined && b.text.trim() !== ''),
      )
      if (!hasNonWhitespaceContent) return null
      return { ...msg, message: { ...msg.message, content } } as typeof msg
    })
    .filter((m): m is Message => m !== null)
}

function createSpeculationFeedbackMessage(
  messages: Message[],
  boundary: CompletionBoundary | null,
  timeSavedMs: number,
  sessionTotalMs: number,
): Message | null {
  if (process.env.USER_TYPE !== 'ant') return null

  if (messages.length === 0 || timeSavedMs === 0) return null

  const toolUses = countToolsInMessages(messages)
  const tokens = boundary?.type === 'complete' ? boundary.outputTokens : null

  const parts = []
  if (toolUses > 0) {
    parts.push(`Speculated ${toolUses} tool ${toolUses === 1 ? 'use' : 'uses'}`)
  } else {
    const turns = messages.length
    parts.push(`Speculated ${turns} ${turns === 1 ? 'turn' : 'turns'}`)
  }

  if (tokens !== null) {
    parts.push(`${formatNumber(tokens)} tokens`)
  }

  const savedText = `+${formatDuration(timeSavedMs)} saved`
  const sessionSuffix =
    sessionTotalMs !== timeSavedMs
      ? ` (${formatDuration(sessionTotalMs)} this session)`
      : ''

  return createSystemMessage(
    `[ANT-ONLY] ${parts.join(' · ')} · ${savedText}${sessionSuffix}`,
    'warning',
  )
}

function updateActiveSpeculationState(
  setAppState: SetAppState,
  updater: (state: ActiveSpeculationState) => Partial<ActiveSpeculationState>,
): void {
  setAppState(prev => {
    if (prev.speculation.status !== 'active') return prev
    const current = prev.speculation as ActiveSpeculationState
    const updates = updater(current)
    // Check if any values actually changed to avoid unnecessary re-renders
    const hasChanges = Object.entries(updates).some(
      ([key, value]) => current[key as keyof ActiveSpeculationState] !== value,
    )
    if (!hasChanges) return prev
    return {
      ...prev,
      speculation: { ...current, ...updates },
    }
  })
}

function resetSpeculationState(setAppState: SetAppState): void {
  setAppState(prev => {
    if (prev.speculation.status === 'idle') return prev
    return { ...prev, speculation: IDLE_SPECULATION_STATE }
  })
}

export function isSpeculationEnabled(): boolean {
  const enabled =
    process.env.USER_TYPE === 'ant' &&
    (getGlobalConfig().speculationEnabled ?? true)
  logForDebugging(`[Speculation] enabled=${enabled}`)
  return enabled
}

async function generatePipelinedSuggestion(
  context: REPLHookContext,
  suggestionText: string,
  speculatedMessages: Message[],
  setAppState: SetAppState,
  parentAbortController: AbortController,
): Promise<void> {
  try {
    const appState = context.toolUseContext.getAppState()
    const suppressReason = getSuggestionSuppressReason(appState)
    if (suppressReason) {
      logSuggestionSuppressed(`pipeline_${suppressReason}`)
      return
    }

    const augmentedContext: REPLHookContext = {
      ...context,
      messages: [
        ...context.messages,
        createUserMessage({ content: suggestionText }),
        ...speculatedMessages,
      ],
    }

    const pipelineAbortController = createChildAbortController(
      parentAbortController,
    )
    if (pipelineAbortController.signal.aborted) return

    const promptId = getPromptVariant()
    const { suggestion, generationRequestId } = await generateSuggestion(
      pipelineAbortController,
      promptId,
      createCacheSafeParams(augmentedContext),
    )

    if (pipelineAbortController.signal.aborted) return
    if (shouldFilterSuggestion(suggestion, promptId)) return

    logForDebugging(
      `[Speculation] Pipelined suggestion: "${suggestion!.slice(0, 50)}..."`,
    )
    updateActiveSpeculationState(setAppState, () => ({
      pipelinedSuggestion: {
        text: suggestion!,
        promptId,
        generationRequestId,
      },
    }))
  } catch (error) {
    if (error instanceof Error && error.name === 'AbortError') return
    logForDebugging(
      `[Speculation] Pipelined suggestion failed: ${errorMessage(error)}`,
    )
  }
}

export async function startSpeculation(
  suggestionText: string,
  context: REPLHookContext,
  setAppState: (f: (prev: AppState) => AppState) => void,
  isPipelined = false,
  cacheSafeParams?: CacheSafeParams,
): Promise<void> {
  if (!isSpeculationEnabled()) return

  // Abort any existing speculation before starting a new one
  abortSpeculation(setAppState)

  const id = randomUUID().slice(0, 8)

  const abortController = createChildAbortController(
    context.toolUseContext.abortController,
  )

  if (abortController.signal.aborted) return

  const startTime = Date.now()
  const messagesRef = { current: [] as Message[] }
  const writtenPathsRef = { current: new Set<string>() }
  const overlayPath = getOverlayPath(id)
  const cwd = getCwdState()

  try {
    await mkdir(overlayPath, { recursive: true })
  } catch {
    logForDebugging('[Speculation] Failed to create overlay directory')
    return
  }

  const contextRef = { current: context }

  setAppState(prev => ({
    ...prev,
    speculation: {
      status: 'active',
      id,
      abort: () => abortController.abort(),
      startTime,
      messagesRef,
      writtenPathsRef,
      boundary: null,
      suggestionLength: suggestionText.length,
      toolUseCount: 0,
      isPipelined,
      contextRef,
    },
  }))

  logForDebugging(`[Speculation] Starting speculation ${id}`)

  try {
    const result = await runForkedAgent({
      promptMessages: [createUserMessage({ content: suggestionText })],
      cacheSafeParams: cacheSafeParams ?? createCacheSafeParams(context),
      skipTranscript: true,
      canUseTool: async (tool, input) => {
        const isWriteTool = WRITE_TOOLS.has(tool.name)
        const isSafeReadOnlyTool = SAFE_READ_ONLY_TOOLS.has(tool.name)

        // Check permission mode BEFORE allowing file edits
        if (isWriteTool) {
          const appState = context.toolUseContext.getAppState()
          const { mode, isBypassPermissionsModeAvailable } =
            appState.toolPermissionContext

          const canAutoAcceptEdits =
            mode === 'acceptEdits' ||
            mode === 'bypassPermissions' ||
            (mode === 'plan' && isBypassPermissionsModeAvailable)

          if (!canAutoAcceptEdits) {
            logForDebugging(`[Speculation] Stopping at file edit: ${tool.name}`)
            const editPath = (
              'file_path' in input ? input.file_path : undefined
            ) as string | undefined
            updateActiveSpeculationState(setAppState, () => ({
              boundary: {
                type: 'edit',
                toolName: tool.name,
                filePath: editPath ?? '',
                completedAt: Date.now(),
              },
            }))
            abortController.abort()
            return denySpeculation(
              'Speculation paused: file edit requires permission',
              'speculation_edit_boundary',
            )
          }
        }

        // Handle file path rewriting for overlay isolation
        if (isWriteTool || isSafeReadOnlyTool) {
          const pathKey =
            'notebook_path' in input
              ? 'notebook_path'
              : 'path' in input
                ? 'path'
                : 'file_path'
          const filePath = input[pathKey] as string | undefined
          if (filePath) {
            const rel = relative(cwd, filePath)
            if (isAbsolute(rel) || rel.startsWith('..')) {
              if (isWriteTool) {
                logForDebugging(
                  `[Speculation] Denied ${tool.name}: path outside cwd: ${filePath}`,
                )
                return denySpeculation(
                  'Write outside cwd not allowed during speculation',
                  'speculation_write_outside_root',
                )
              }
              return {
                behavior: 'allow' as const,
                updatedInput: input,
                decisionReason: {
                  type: 'other' as const,
                  reason: 'speculation_read_outside_root',
                },
              }
            }

            if (isWriteTool) {
              // Copy-on-write: copy original to overlay if not yet there
              if (!writtenPathsRef.current.has(rel)) {
                const overlayFile = join(overlayPath, rel)
                await mkdir(dirname(overlayFile), { recursive: true })
                try {
                  await copyFile(join(cwd, rel), overlayFile)
                } catch {
                  // Original may not exist (new file creation) - that's fine
                }
                writtenPathsRef.current.add(rel)
              }
              input = { ...input, [pathKey]: join(overlayPath, rel) }
            } else {
              // Read: redirect to overlay if file was previously written
              if (writtenPathsRef.current.has(rel)) {
                input = { ...input, [pathKey]: join(overlayPath, rel) }
              }
              // Otherwise read from main (no rewrite)
            }

            logForDebugging(
              `[Speculation] ${isWriteTool ? 'Write' : 'Read'} ${filePath} -> ${input[pathKey]}`,
            )

            return {
              behavior: 'allow' as const,
              updatedInput: input,
              decisionReason: {
                type: 'other' as const,
                reason: 'speculation_file_access',
              },
            }
          }
          // Read tools without explicit path (e.g. Glob/Grep defaulting to CWD) are safe
          if (isSafeReadOnlyTool) {
            return {
              behavior: 'allow' as const,
              updatedInput: input,
              decisionReason: {
                type: 'other' as const,
                reason: 'speculation_read_default_cwd',
              },
            }
          }
          // Write tools with undefined path → fall through to default deny
        }

        // Stop at non-read-only bash commands
        if (tool.name === 'Bash') {
          const command =
            'command' in input && typeof input.command === 'string'
              ? input.command
              : ''
          if (
            !command ||
            checkReadOnlyConstraints({ command }, commandHasAnyCd(command))
              .behavior !== 'allow'
          ) {
            logForDebugging(
              `[Speculation] Stopping at bash: ${command.slice(0, 50) || 'missing command'}`,
            )
            updateActiveSpeculationState(setAppState, () => ({
              boundary: { type: 'bash', command, completedAt: Date.now() },
            }))
            abortController.abort()
            return denySpeculation(
              'Speculation paused: bash boundary',
              'speculation_bash_boundary',
            )
          }
          // Read-only bash command — allow during speculation
          return {
            behavior: 'allow' as const,
            updatedInput: input,
            decisionReason: {
              type: 'other' as const,
              reason: 'speculation_readonly_bash',
            },
          }
        }

        // Deny all other tools by default
        logForDebugging(`[Speculation] Stopping at denied tool: ${tool.name}`)
        const detail = String(
          ('url' in input && input.url) ||
            ('file_path' in input && input.file_path) ||
            ('path' in input && input.path) ||
            ('command' in input && input.command) ||
            '',
        ).slice(0, 200)
        updateActiveSpeculationState(setAppState, () => ({
          boundary: {
            type: 'denied_tool',
            toolName: tool.name,
            detail,
            completedAt: Date.now(),
          },
        }))
        abortController.abort()
        return denySpeculation(
          `Tool ${tool.name} not allowed during speculation`,
          'speculation_unknown_tool',
        )
      },
      querySource: 'speculation',
      forkLabel: 'speculation',
      maxTurns: MAX_SPECULATION_TURNS,
      overrides: { abortController, requireCanUseTool: true },
      onMessage: msg => {
        if (msg.type === 'assistant' || msg.type === 'user') {
          messagesRef.current.push(msg)
          if (messagesRef.current.length >= MAX_SPECULATION_MESSAGES) {
            abortController.abort()
          }
          if (isUserMessageWithArrayContent(msg)) {
            const newTools = count(
              msg.message.content as { type: string; is_error?: boolean }[],
              b => b.type === 'tool_result' && !b.is_error,
            )
            if (newTools > 0) {
              updateActiveSpeculationState(setAppState, prev => ({
                toolUseCount: prev.toolUseCount + newTools,
              }))
            }
          }
        }
      },
    })

    if (abortController.signal.aborted) return

    updateActiveSpeculationState(setAppState, () => ({
      boundary: {
        type: 'complete' as const,
        completedAt: Date.now(),
        outputTokens: result.totalUsage.output_tokens,
      },
    }))

    logForDebugging(
      `[Speculation] Complete: ${countToolsInMessages(messagesRef.current)} tools`,
    )

    // Pipeline: generate the next suggestion while we wait for the user to accept
    void generatePipelinedSuggestion(
      contextRef.current,
      suggestionText,
      messagesRef.current,
      setAppState,
      abortController,
    )
  } catch (error) {
    abortController.abort()

    if (error instanceof Error && error.name === 'AbortError') {
      safeRemoveOverlay(overlayPath)
      resetSpeculationState(setAppState)
      return
    }

    safeRemoveOverlay(overlayPath)

    // eslint-disable-next-line no-restricted-syntax -- custom fallback message, not toError(e)
    logError(error instanceof Error ? error : new Error('Speculation failed'))

    logSpeculation(
      id,
      'error',
      startTime,
      suggestionText.length,
      messagesRef.current,
      null,
      {
        error_type: error instanceof Error ? error.name : 'Unknown',
        error_message: errorMessage(error).slice(
          0,
          200,
        ) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
        error_phase:
          'start' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
        is_pipelined: isPipelined,
      },
    )

    resetSpeculationState(setAppState)
  }
}

export async function acceptSpeculation(
  state: SpeculationState,
  setAppState: (f: (prev: AppState) => AppState) => void,
  cleanMessageCount: number,
): Promise<SpeculationResult | null> {
  if (state.status !== 'active') return null

  const {
    id,
    messagesRef,
    writtenPathsRef,
    abort,
    startTime,
    suggestionLength,
    isPipelined,
  } = state
  const messages = messagesRef.current
  const overlayPath = getOverlayPath(id)
  const acceptedAt = Date.now()

  abort()

  if (cleanMessageCount > 0) {
    await copyOverlayToMain(overlayPath, writtenPathsRef.current, getCwdState())
  }
  safeRemoveOverlay(overlayPath)

  // Use snapshot boundary as default (available since state.status === 'active' was checked above)
  let boundary: CompletionBoundary | null = state.boundary
  let timeSavedMs =
    Math.min(acceptedAt, boundary?.completedAt ?? Infinity) - startTime

  setAppState(prev => {
    // Refine with latest React state if speculation is still active
    if (prev.speculation.status === 'active' && prev.speculation.boundary) {
      boundary = prev.speculation.boundary
      const endTime = Math.min(acceptedAt, boundary.completedAt ?? Infinity)
      timeSavedMs = endTime - startTime
    }
    return {
      ...prev,
      speculation: IDLE_SPECULATION_STATE,
      speculationSessionTimeSavedMs:
        prev.speculationSessionTimeSavedMs + timeSavedMs,
    }
  })

  logForDebugging(
    boundary === null
      ? `[Speculation] Accept ${id}: still running, using ${messages.length} messages`
      : `[Speculation] Accept ${id}: already complete`,
  )

  logSpeculation(
    id,
    'accepted',
    startTime,
    suggestionLength,
    messages,
    boundary,
    {
      message_count: messages.length,
      time_saved_ms: timeSavedMs,
      is_pipelined: isPipelined,
    },
  )

  if (timeSavedMs > 0) {
    const entry: SpeculationAcceptMessage = {
      type: 'speculation-accept',
      timestamp: new Date().toISOString(),
      timeSavedMs,
    }
    void appendFile(getTranscriptPath(), jsonStringify(entry) + '\n', {
      mode: 0o600,
    }).catch(() => {
      logForDebugging(
        '[Speculation] Failed to write speculation-accept to transcript',
      )
    })
  }

  return { messages, boundary, timeSavedMs }
}

export function abortSpeculation(setAppState: SetAppState): void {
  setAppState(prev => {
    if (prev.speculation.status !== 'active') return prev

    const {
      id,
      abort,
      startTime,
      boundary,
      suggestionLength,
      messagesRef,
      isPipelined,
    } = prev.speculation

    logForDebugging(`[Speculation] Aborting ${id}`)

    logSpeculation(
      id,
      'aborted',
      startTime,
      suggestionLength,
      messagesRef.current,
      boundary,
      { abort_reason: 'user_typed', is_pipelined: isPipelined },
    )

    abort()
    safeRemoveOverlay(getOverlayPath(id))

    return { ...prev, speculation: IDLE_SPECULATION_STATE }
  })
}

export async function handleSpeculationAccept(
  speculationState: ActiveSpeculationState,
  speculationSessionTimeSavedMs: number,
  setAppState: SetAppState,
  input: string,
  deps: {
    setMessages: (f: (prev: Message[]) => Message[]) => void
    readFileState: { current: FileStateCache }
    cwd: string
  },
): Promise<{ queryRequired: boolean }> {
  try {
    const { setMessages, readFileState, cwd } = deps

    // Clear prompt suggestion state. logOutcomeAtSubmission logged the accept
    // but was called with skipReset to avoid aborting speculation before we use it.
    setAppState(prev => {
      if (
        prev.promptSuggestion.text === null &&
        prev.promptSuggestion.promptId === null
      ) {
        return prev
      }
      return {
        ...prev,
        promptSuggestion: {
          text: null,
          promptId: null,
          shownAt: 0,
          acceptedAt: 0,
          generationRequestId: null,
        },
      }
    })

    // Capture speculation messages before any state updates - must be stable reference
    const speculationMessages = speculationState.messagesRef.current
    let cleanMessages = prepareMessagesForInjection(speculationMessages)

    // Inject user message first for instant visual feedback before any async work
    const userMessage = createUserMessage({ content: input })
    setMessages(prev => [...prev, userMessage])

    const result = await acceptSpeculation(
      speculationState,
      setAppState,
      cleanMessages.length,
    )

    const isComplete = result?.boundary?.type === 'complete'

    // When speculation didn't complete, the follow-up query needs the
    // conversation to end with a user message. Drop trailing assistant
    // messages — models that don't support prefill
    // reject conversations ending with an assistant turn. The model will
    // regenerate this content in the follow-up query.
    if (!isComplete) {
      const lastNonAssistant = cleanMessages.findLastIndex(
        m => m.type !== 'assistant',
      )
      cleanMessages = cleanMessages.slice(0, lastNonAssistant + 1)
    }

    const timeSavedMs = result?.timeSavedMs ?? 0
    const newSessionTotal = speculationSessionTimeSavedMs + timeSavedMs
    const feedbackMessage = createSpeculationFeedbackMessage(
      cleanMessages,
      result?.boundary ?? null,
      timeSavedMs,
      newSessionTotal,
    )

    // Inject speculated messages
    setMessages(prev => [...prev, ...cleanMessages])

    const extracted = extractReadFilesFromMessages(
      cleanMessages,
      cwd,
      READ_FILE_STATE_CACHE_SIZE,
    )
    readFileState.current = mergeFileStateCaches(
      readFileState.current,
      extracted,
    )

    if (feedbackMessage) {
      setMessages(prev => [...prev, feedbackMessage])
    }

    logForDebugging(
      `[Speculation] ${result?.boundary?.type ?? 'incomplete'}, injected ${cleanMessages.length} messages`,
    )

    // Promote pipelined suggestion if speculation completed fully
    if (isComplete && speculationState.pipelinedSuggestion) {
      const { text, promptId, generationRequestId } =
        speculationState.pipelinedSuggestion
      logForDebugging(
        `[Speculation] Promoting pipelined suggestion: "${text.slice(0, 50)}..."`,
      )
      setAppState(prev => ({
        ...prev,
        promptSuggestion: {
          text,
          promptId,
          shownAt: Date.now(),
          acceptedAt: 0,
          generationRequestId,
        },
      }))

      // Start speculation on the pipelined suggestion
      const augmentedContext: REPLHookContext = {
        ...speculationState.contextRef.current,
        messages: [
          ...speculationState.contextRef.current.messages,
          createUserMessage({ content: input }),
          ...cleanMessages,
        ],
      }
      void startSpeculation(text, augmentedContext, setAppState, true)
    }

    return { queryRequired: !isComplete }
  } catch (error) {
    // Fail open: log error and fall back to normal query flow
    /* eslint-disable no-restricted-syntax -- custom fallback message, not toError(e) */
    logError(
      error instanceof Error
        ? error
        : new Error('handleSpeculationAccept failed'),
    )
    /* eslint-enable no-restricted-syntax */
    logSpeculation(
      speculationState.id,
      'error',
      speculationState.startTime,
      speculationState.suggestionLength,
      speculationState.messagesRef.current,
      speculationState.boundary,
      {
        error_type: error instanceof Error ? error.name : 'Unknown',
        error_message: errorMessage(error).slice(
          0,
          200,
        ) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
        error_phase:
          'accept' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
        is_pipelined: speculationState.isPipelined,
      },
    )
    safeRemoveOverlay(getOverlayPath(speculationState.id))
    resetSpeculationState(setAppState)
    // Query required so user's message is processed normally (without speculated work)
    return { queryRequired: true }
  }
}