1  import { randomUUID } from 'crypto'
  2  import { rm } from 'fs'
  3  import { appendFile, copyFile, mkdir } from 'fs/promises'
  4  import { dirname, isAbsolute, join, relative } from 'path'
  5  import { getCwdState } from '../../bootstrap/state.js'
  6  import type { CompletionBoundary } from '../../state/AppStateStore.js'
  7  import {
  8    type AppState,
  9    IDLE_SPECULATION_STATE,
 10    type SpeculationResult,
 11    type SpeculationState,
 12  } from '../../state/AppStateStore.js'
 13  import { commandHasAnyCd } from '../../tools/BashTool/bashPermissions.js'
 14  import { checkReadOnlyConstraints } from '../../tools/BashTool/readOnlyValidation.js'
 15  import type { SpeculationAcceptMessage } from '../../types/logs.js'
 16  import type { Message } from '../../types/message.js'
 17  import { createChildAbortController } from '../../utils/abortController.js'
 18  import { count } from '../../utils/array.js'
 19  import { getGlobalConfig } from '../../utils/config.js'
 20  import { logForDebugging } from '../../utils/debug.js'
 21  import { errorMessage } from '../../utils/errors.js'
 22  import {
 23    type FileStateCache,
 24    mergeFileStateCaches,
 25    READ_FILE_STATE_CACHE_SIZE,
 26  } from '../../utils/fileStateCache.js'
 27  import {
 28    type CacheSafeParams,
 29    createCacheSafeParams,
 30    runForkedAgent,
 31  } from '../../utils/forkedAgent.js'
 32  import { formatDuration, formatNumber } from '../../utils/format.js'
 33  import type { REPLHookContext } from '../../utils/hooks/postSamplingHooks.js'
 34  import { logError } from '../../utils/log.js'
 35  import type { SetAppState } from '../../utils/messageQueueManager.js'
 36  import {
 37    createSystemMessage,
 38    createUserMessage,
 39    INTERRUPT_MESSAGE,
 40    INTERRUPT_MESSAGE_FOR_TOOL_USE,
 41  } from '../../utils/messages.js'
 42  import { getClaudeTempDir } from '../../utils/permissions/filesystem.js'
 43  import { extractReadFilesFromMessages } from '../../utils/queryHelpers.js'
 44  import { getTranscriptPath } from '../../utils/sessionStorage.js'
 45  import { jsonStringify } from '../../utils/slowOperations.js'
 46  import {
 47    type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
 48    logEvent,
 49  } from '../analytics/index.js'
 50  import {
 51    generateSuggestion,
 52    getPromptVariant,
 53    getSuggestionSuppressReason,
 54    logSuggestionSuppressed,
 55    shouldFilterSuggestion,
 56  } from './promptSuggestion.js'
 57  
 58  const MAX_SPECULATION_TURNS = 20
 59  const MAX_SPECULATION_MESSAGES = 100
 60  
 61  const WRITE_TOOLS = new Set(['Edit', 'Write', 'NotebookEdit'])
 62  const SAFE_READ_ONLY_TOOLS = new Set([
 63    'Read',
 64    'Glob',
 65    'Grep',
 66    'ToolSearch',
 67    'LSP',
 68    'TaskGet',
 69    'TaskList',
 70  ])
 71  
 72  function safeRemoveOverlay(overlayPath: string): void {
 73    rm(
 74      overlayPath,
 75      { recursive: true, force: true, maxRetries: 3, retryDelay: 100 },
 76      () => {},
 77    )
 78  }
 79  
 80  function getOverlayPath(id: string): string {
 81    return join(getClaudeTempDir(), 'speculation', String(process.pid), id)
 82  }
 83  
 84  function denySpeculation(
 85    message: string,
 86    reason: string,
 87  ): {
 88    behavior: 'deny'
 89    message: string
 90    decisionReason: { type: 'other'; reason: string }
 91  } {
 92    return {
 93      behavior: 'deny',
 94      message,
 95      decisionReason: { type: 'other', reason },
 96    }
 97  }
 98  
 99  async function copyOverlayToMain(
100    overlayPath: string,
101    writtenPaths: Set<string>,
102    cwd: string,
103  ): Promise<boolean> {
104    let allCopied = true
105    for (const rel of writtenPaths) {
106      const src = join(overlayPath, rel)
107      const dest = join(cwd, rel)
108      try {
109        await mkdir(dirname(dest), { recursive: true })
110        await copyFile(src, dest)
111      } catch {
112        allCopied = false
113        logForDebugging(`[Speculation] Failed to copy ${rel} to main`)
114      }
115    }
116    return allCopied
117  }
118  
119  export type ActiveSpeculationState = Extract<
120    SpeculationState,
121    { status: 'active' }
122  >
123  
124  function logSpeculation(
125    id: string,
126    outcome: 'accepted' | 'aborted' | 'error',
127    startTime: number,
128    suggestionLength: number,
129    messages: Message[],
130    boundary: CompletionBoundary | null,
131    extras?: Record<string, string | number | boolean | undefined>,
132  ): void {
133    logEvent('tengu_speculation', {
134      speculation_id:
135        id as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
136      outcome:
137        outcome as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
138      duration_ms: Date.now() - startTime,
139      suggestion_length: suggestionLength,
140      tools_executed: countToolsInMessages(messages),
141      completed: boundary !== null,
142      boundary_type: boundary?.type as
143        | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
144        | undefined,
145      boundary_tool: getBoundaryTool(boundary) as
146        | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
147        | undefined,
148      boundary_detail: getBoundaryDetail(boundary) as
149        | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
150        | undefined,
151      ...extras,
152    })
153  }
154  
155  function countToolsInMessages(messages: Message[]): number {
156    const blocks = messages
157      .filter(isUserMessageWithArrayContent)
158      .flatMap(m => m.message.content)
159      .filter(
160        (b): b is { type: string; is_error?: boolean } =>
161          typeof b === 'object' && b !== null && 'type' in b,
162      )
163    return count(blocks, b => b.type === 'tool_result' && !b.is_error)
164  }
165  
166  function getBoundaryTool(
167    boundary: CompletionBoundary | null,
168  ): string | undefined {
169    if (!boundary) return undefined
170    switch (boundary.type) {
171      case 'bash':
172        return 'Bash'
173      case 'edit':
174      case 'denied_tool':
175        return boundary.toolName
176      case 'complete':
177        return undefined
178    }
179  }
180  
181  function getBoundaryDetail(
182    boundary: CompletionBoundary | null,
183  ): string | undefined {
184    if (!boundary) return undefined
185    switch (boundary.type) {
186      case 'bash':
187        return boundary.command.slice(0, 200)
188      case 'edit':
189        return boundary.filePath
190      case 'denied_tool':
191        return boundary.detail
192      case 'complete':
193        return undefined
194    }
195  }
196  
197  function isUserMessageWithArrayContent(
198    m: Message,
199  ): m is Message & { message: { content: unknown[] } } {
200    return m.type === 'user' && 'message' in m && Array.isArray(m.message.content)
201  }
202  
203  export function prepareMessagesForInjection(messages: Message[]): Message[] {
204    // Find tool_use IDs that have SUCCESSFUL results (not errors/interruptions)
205    // Pending tool_use blocks (no result) and interrupted ones will be stripped
206    type ToolResult = {
207      type: 'tool_result'
208      tool_use_id: string
209      is_error?: boolean
210      content?: unknown
211    }
212    const isToolResult = (b: unknown): b is ToolResult =>
213      typeof b === 'object' &&
214      b !== null &&
215      (b as ToolResult).type === 'tool_result' &&
216      typeof (b as ToolResult).tool_use_id === 'string'
217    const isSuccessful = (b: ToolResult) =>
218      !b.is_error &&
219      !(
220        typeof b.content === 'string' &&
221        b.content.includes(INTERRUPT_MESSAGE_FOR_TOOL_USE)
222      )
223  
224    const toolIdsWithSuccessfulResults = new Set(
225      messages
226        .filter(isUserMessageWithArrayContent)
227        .flatMap(m => m.message.content)
228        .filter(isToolResult)
229        .filter(isSuccessful)
230        .map(b => b.tool_use_id),
231    )
232  
233    const keep = (b: {
234      type: string
235      id?: string
236      tool_use_id?: string
237      text?: string
238    }) =>
239      b.type !== 'thinking' &&
240      b.type !== 'redacted_thinking' &&
241      !(b.type === 'tool_use' && !toolIdsWithSuccessfulResults.has(b.id!)) &&
242      !(
243        b.type === 'tool_result' &&
244        !toolIdsWithSuccessfulResults.has(b.tool_use_id!)
245      ) &&
246      // Abort during speculation yields a standalone interrupt user message
247      // (query.ts createUserInterruptionMessage). Strip it so it isn't surfaced
248      // to the model as real user input.
249      !(
250        b.type === 'text' &&
251        (b.text === INTERRUPT_MESSAGE ||
252          b.text === INTERRUPT_MESSAGE_FOR_TOOL_USE)
253      )
254  
255    return messages
256      .map(msg => {
257        if (!('message' in msg) || !Array.isArray(msg.message.content)) return msg
258        const content = msg.message.content.filter(keep)
259        if (content.length === msg.message.content.length) return msg
260        if (content.length === 0) return null
261        // Drop messages where all remaining blocks are whitespace-only text
262        // (API rejects these with 400: "text content blocks must contain non-whitespace text")
263        const hasNonWhitespaceContent = content.some(
264          (b: { type: string; text?: string }) =>
265            b.type !== 'text' || (b.text !== undefined && b.text.trim() !== ''),
266        )
267        if (!hasNonWhitespaceContent) return null
268        return { ...msg, message: { ...msg.message, content } } as typeof msg
269      })
270      .filter((m): m is Message => m !== null)
271  }
272  
273  function createSpeculationFeedbackMessage(
274    messages: Message[],
275    boundary: CompletionBoundary | null,
276    timeSavedMs: number,
277    sessionTotalMs: number,
278  ): Message | null {
279    if (process.env.USER_TYPE !== 'ant') return null
280  
281    if (messages.length === 0 || timeSavedMs === 0) return null
282  
283    const toolUses = countToolsInMessages(messages)
284    const tokens = boundary?.type === 'complete' ? boundary.outputTokens : null
285  
286    const parts = []
287    if (toolUses > 0) {
288      parts.push(`Speculated ${toolUses} tool ${toolUses === 1 ? 'use' : 'uses'}`)
289    } else {
290      const turns = messages.length
291      parts.push(`Speculated ${turns} ${turns === 1 ? 'turn' : 'turns'}`)
292    }
293  
294    if (tokens !== null) {
295      parts.push(`${formatNumber(tokens)} tokens`)
296    }
297  
298    const savedText = `+${formatDuration(timeSavedMs)} saved`
299    const sessionSuffix =
300      sessionTotalMs !== timeSavedMs
301        ? ` (${formatDuration(sessionTotalMs)} this session)`
302        : ''
303  
304    return createSystemMessage(
305      `[ANT-ONLY] ${parts.join(' · ')} · ${savedText}${sessionSuffix}`,
306      'warning',
307    )
308  }
309  
310  function updateActiveSpeculationState(
311    setAppState: SetAppState,
312    updater: (state: ActiveSpeculationState) => Partial<ActiveSpeculationState>,
313  ): void {
314    setAppState(prev => {
315      if (prev.speculation.status !== 'active') return prev
316      const current = prev.speculation as ActiveSpeculationState
317      const updates = updater(current)
318      // Check if any values actually changed to avoid unnecessary re-renders
319      const hasChanges = Object.entries(updates).some(
320        ([key, value]) => current[key as keyof ActiveSpeculationState] !== value,
321      )
322      if (!hasChanges) return prev
323      return {
324        ...prev,
325        speculation: { ...current, ...updates },
326      }
327    })
328  }
329  
330  function resetSpeculationState(setAppState: SetAppState): void {
331    setAppState(prev => {
332      if (prev.speculation.status === 'idle') return prev
333      return { ...prev, speculation: IDLE_SPECULATION_STATE }
334    })
335  }
336  
337  export function isSpeculationEnabled(): boolean {
338    const enabled =
339      process.env.USER_TYPE === 'ant' &&
340      (getGlobalConfig().speculationEnabled ?? true)
341    logForDebugging(`[Speculation] enabled=${enabled}`)
342    return enabled
343  }
344  
345  async function generatePipelinedSuggestion(
346    context: REPLHookContext,
347    suggestionText: string,
348    speculatedMessages: Message[],
349    setAppState: SetAppState,
350    parentAbortController: AbortController,
351  ): Promise<void> {
352    try {
353      const appState = context.toolUseContext.getAppState()
354      const suppressReason = getSuggestionSuppressReason(appState)
355      if (suppressReason) {
356        logSuggestionSuppressed(`pipeline_${suppressReason}`)
357        return
358      }
359  
360      const augmentedContext: REPLHookContext = {
361        ...context,
362        messages: [
363          ...context.messages,
364          createUserMessage({ content: suggestionText }),
365          ...speculatedMessages,
366        ],
367      }
368  
369      const pipelineAbortController = createChildAbortController(
370        parentAbortController,
371      )
372      if (pipelineAbortController.signal.aborted) return
373  
374      const promptId = getPromptVariant()
375      const { suggestion, generationRequestId } = await generateSuggestion(
376        pipelineAbortController,
377        promptId,
378        createCacheSafeParams(augmentedContext),
379      )
380  
381      if (pipelineAbortController.signal.aborted) return
382      if (shouldFilterSuggestion(suggestion, promptId)) return
383  
384      logForDebugging(
385        `[Speculation] Pipelined suggestion: "${suggestion!.slice(0, 50)}..."`,
386      )
387      updateActiveSpeculationState(setAppState, () => ({
388        pipelinedSuggestion: {
389          text: suggestion!,
390          promptId,
391          generationRequestId,
392        },
393      }))
394    } catch (error) {
395      if (error instanceof Error && error.name === 'AbortError') return
396      logForDebugging(
397        `[Speculation] Pipelined suggestion failed: ${errorMessage(error)}`,
398      )
399    }
400  }
401  
402  export async function startSpeculation(
403    suggestionText: string,
404    context: REPLHookContext,
405    setAppState: (f: (prev: AppState) => AppState) => void,
406    isPipelined = false,
407    cacheSafeParams?: CacheSafeParams,
408  ): Promise<void> {
409    if (!isSpeculationEnabled()) return
410  
411    // Abort any existing speculation before starting a new one
412    abortSpeculation(setAppState)
413  
414    const id = randomUUID().slice(0, 8)
415  
416    const abortController = createChildAbortController(
417      context.toolUseContext.abortController,
418    )
419  
420    if (abortController.signal.aborted) return
421  
422    const startTime = Date.now()
423    const messagesRef = { current: [] as Message[] }
424    const writtenPathsRef = { current: new Set<string>() }
425    const overlayPath = getOverlayPath(id)
426    const cwd = getCwdState()
427  
428    try {
429      await mkdir(overlayPath, { recursive: true })
430    } catch {
431      logForDebugging('[Speculation] Failed to create overlay directory')
432      return
433    }
434  
435    const contextRef = { current: context }
436  
437    setAppState(prev => ({
438      ...prev,
439      speculation: {
440        status: 'active',
441        id,
442        abort: () => abortController.abort(),
443        startTime,
444        messagesRef,
445        writtenPathsRef,
446        boundary: null,
447        suggestionLength: suggestionText.length,
448        toolUseCount: 0,
449        isPipelined,
450        contextRef,
451      },
452    }))
453  
454    logForDebugging(`[Speculation] Starting speculation ${id}`)
455  
456    try {
457      const result = await runForkedAgent({
458        promptMessages: [createUserMessage({ content: suggestionText })],
459        cacheSafeParams: cacheSafeParams ?? createCacheSafeParams(context),
460        skipTranscript: true,
461        canUseTool: async (tool, input) => {
462          const isWriteTool = WRITE_TOOLS.has(tool.name)
463          const isSafeReadOnlyTool = SAFE_READ_ONLY_TOOLS.has(tool.name)
464  
465          // Check permission mode BEFORE allowing file edits
466          if (isWriteTool) {
467            const appState = context.toolUseContext.getAppState()
468            const { mode, isBypassPermissionsModeAvailable } =
469              appState.toolPermissionContext
470  
471            const canAutoAcceptEdits =
472              mode === 'acceptEdits' ||
473              mode === 'bypassPermissions' ||
474              (mode === 'plan' && isBypassPermissionsModeAvailable)
475  
476            if (!canAutoAcceptEdits) {
477              logForDebugging(`[Speculation] Stopping at file edit: ${tool.name}`)
478              const editPath = (
479                'file_path' in input ? input.file_path : undefined
480              ) as string | undefined
481              updateActiveSpeculationState(setAppState, () => ({
482                boundary: {
483                  type: 'edit',
484                  toolName: tool.name,
485                  filePath: editPath ?? '',
486                  completedAt: Date.now(),
487                },
488              }))
489              abortController.abort()
490              return denySpeculation(
491                'Speculation paused: file edit requires permission',
492                'speculation_edit_boundary',
493              )
494            }
495          }
496  
497          // Handle file path rewriting for overlay isolation
498          if (isWriteTool || isSafeReadOnlyTool) {
499            const pathKey =
500              'notebook_path' in input
501                ? 'notebook_path'
502                : 'path' in input
503                  ? 'path'
504                  : 'file_path'
505            const filePath = input[pathKey] as string | undefined
506            if (filePath) {
507              const rel = relative(cwd, filePath)
508              if (isAbsolute(rel) || rel.startsWith('..')) {
509                if (isWriteTool) {
510                  logForDebugging(
511                    `[Speculation] Denied ${tool.name}: path outside cwd: ${filePath}`,
512                  )
513                  return denySpeculation(
514                    'Write outside cwd not allowed during speculation',
515                    'speculation_write_outside_root',
516                  )
517                }
518                return {
519                  behavior: 'allow' as const,
520                  updatedInput: input,
521                  decisionReason: {
522                    type: 'other' as const,
523                    reason: 'speculation_read_outside_root',
524                  },
525                }
526              }
527  
528              if (isWriteTool) {
529                // Copy-on-write: copy original to overlay if not yet there
530                if (!writtenPathsRef.current.has(rel)) {
531                  const overlayFile = join(overlayPath, rel)
532                  await mkdir(dirname(overlayFile), { recursive: true })
533                  try {
534                    await copyFile(join(cwd, rel), overlayFile)
535                  } catch {
536                    // Original may not exist (new file creation) - that's fine
537                  }
538                  writtenPathsRef.current.add(rel)
539                }
540                input = { ...input, [pathKey]: join(overlayPath, rel) }
541              } else {
542                // Read: redirect to overlay if file was previously written
543                if (writtenPathsRef.current.has(rel)) {
544                  input = { ...input, [pathKey]: join(overlayPath, rel) }
545                }
546                // Otherwise read from main (no rewrite)
547              }
548  
549              logForDebugging(
550                `[Speculation] ${isWriteTool ? 'Write' : 'Read'} ${filePath} -> ${input[pathKey]}`,
551              )
552  
553              return {
554                behavior: 'allow' as const,
555                updatedInput: input,
556                decisionReason: {
557                  type: 'other' as const,
558                  reason: 'speculation_file_access',
559                },
560              }
561            }
562            // Read tools without explicit path (e.g. Glob/Grep defaulting to CWD) are safe
563            if (isSafeReadOnlyTool) {
564              return {
565                behavior: 'allow' as const,
566                updatedInput: input,
567                decisionReason: {
568                  type: 'other' as const,
569                  reason: 'speculation_read_default_cwd',
570                },
571              }
572            }
573            // Write tools with undefined path → fall through to default deny
574          }
575  
576          // Stop at non-read-only bash commands
577          if (tool.name === 'Bash') {
578            const command =
579              'command' in input && typeof input.command === 'string'
580                ? input.command
581                : ''
582            if (
583              !command ||
584              checkReadOnlyConstraints({ command }, commandHasAnyCd(command))
585                .behavior !== 'allow'
586            ) {
587              logForDebugging(
588                `[Speculation] Stopping at bash: ${command.slice(0, 50) || 'missing command'}`,
589              )
590              updateActiveSpeculationState(setAppState, () => ({
591                boundary: { type: 'bash', command, completedAt: Date.now() },
592              }))
593              abortController.abort()
594              return denySpeculation(
595                'Speculation paused: bash boundary',
596                'speculation_bash_boundary',
597              )
598            }
599            // Read-only bash command — allow during speculation
600            return {
601              behavior: 'allow' as const,
602              updatedInput: input,
603              decisionReason: {
604                type: 'other' as const,
605                reason: 'speculation_readonly_bash',
606              },
607            }
608          }
609  
610          // Deny all other tools by default
611          logForDebugging(`[Speculation] Stopping at denied tool: ${tool.name}`)
612          const detail = String(
613            ('url' in input && input.url) ||
614              ('file_path' in input && input.file_path) ||
615              ('path' in input && input.path) ||
616              ('command' in input && input.command) ||
617              '',
618          ).slice(0, 200)
619          updateActiveSpeculationState(setAppState, () => ({
620            boundary: {
621              type: 'denied_tool',
622              toolName: tool.name,
623              detail,
624              completedAt: Date.now(),
625            },
626          }))
627          abortController.abort()
628          return denySpeculation(
629            `Tool ${tool.name} not allowed during speculation`,
630            'speculation_unknown_tool',
631          )
632        },
633        querySource: 'speculation',
634        forkLabel: 'speculation',
635        maxTurns: MAX_SPECULATION_TURNS,
636        overrides: { abortController, requireCanUseTool: true },
637        onMessage: msg => {
638          if (msg.type === 'assistant' || msg.type === 'user') {
639            messagesRef.current.push(msg)
640            if (messagesRef.current.length >= MAX_SPECULATION_MESSAGES) {
641              abortController.abort()
642            }
643            if (isUserMessageWithArrayContent(msg)) {
644              const newTools = count(
645                msg.message.content as { type: string; is_error?: boolean }[],
646                b => b.type === 'tool_result' && !b.is_error,
647              )
648              if (newTools > 0) {
649                updateActiveSpeculationState(setAppState, prev => ({
650                  toolUseCount: prev.toolUseCount + newTools,
651                }))
652              }
653            }
654          }
655        },
656      })
657  
658      if (abortController.signal.aborted) return
659  
660      updateActiveSpeculationState(setAppState, () => ({
661        boundary: {
662          type: 'complete' as const,
663          completedAt: Date.now(),
664          outputTokens: result.totalUsage.output_tokens,
665        },
666      }))
667  
668      logForDebugging(
669        `[Speculation] Complete: ${countToolsInMessages(messagesRef.current)} tools`,
670      )
671  
672      // Pipeline: generate the next suggestion while we wait for the user to accept
673      void generatePipelinedSuggestion(
674        contextRef.current,
675        suggestionText,
676        messagesRef.current,
677        setAppState,
678        abortController,
679      )
680    } catch (error) {
681      abortController.abort()
682  
683      if (error instanceof Error && error.name === 'AbortError') {
684        safeRemoveOverlay(overlayPath)
685        resetSpeculationState(setAppState)
686        return
687      }
688  
689      safeRemoveOverlay(overlayPath)
690  
691      // eslint-disable-next-line no-restricted-syntax -- custom fallback message, not toError(e)
692      logError(error instanceof Error ? error : new Error('Speculation failed'))
693  
694      logSpeculation(
695        id,
696        'error',
697        startTime,
698        suggestionText.length,
699        messagesRef.current,
700        null,
701        {
702          error_type: error instanceof Error ? error.name : 'Unknown',
703          error_message: errorMessage(error).slice(
704            0,
705            200,
706          ) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
707          error_phase:
708            'start' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
709          is_pipelined: isPipelined,
710        },
711      )
712  
713      resetSpeculationState(setAppState)
714    }
715  }
716  
717  export async function acceptSpeculation(
718    state: SpeculationState,
719    setAppState: (f: (prev: AppState) => AppState) => void,
720    cleanMessageCount: number,
721  ): Promise<SpeculationResult | null> {
722    if (state.status !== 'active') return null
723  
724    const {
725      id,
726      messagesRef,
727      writtenPathsRef,
728      abort,
729      startTime,
730      suggestionLength,
731      isPipelined,
732    } = state
733    const messages = messagesRef.current
734    const overlayPath = getOverlayPath(id)
735    const acceptedAt = Date.now()
736  
737    abort()
738  
739    if (cleanMessageCount > 0) {
740      await copyOverlayToMain(overlayPath, writtenPathsRef.current, getCwdState())
741    }
742    safeRemoveOverlay(overlayPath)
743  
744    // Use snapshot boundary as default (available since state.status === 'active' was checked above)
745    let boundary: CompletionBoundary | null = state.boundary
746    let timeSavedMs =
747      Math.min(acceptedAt, boundary?.completedAt ?? Infinity) - startTime
748  
749    setAppState(prev => {
750      // Refine with latest React state if speculation is still active
751      if (prev.speculation.status === 'active' && prev.speculation.boundary) {
752        boundary = prev.speculation.boundary
753        const endTime = Math.min(acceptedAt, boundary.completedAt ?? Infinity)
754        timeSavedMs = endTime - startTime
755      }
756      return {
757        ...prev,
758        speculation: IDLE_SPECULATION_STATE,
759        speculationSessionTimeSavedMs:
760          prev.speculationSessionTimeSavedMs + timeSavedMs,
761      }
762    })
763  
764    logForDebugging(
765      boundary === null
766        ? `[Speculation] Accept ${id}: still running, using ${messages.length} messages`
767        : `[Speculation] Accept ${id}: already complete`,
768    )
769  
770    logSpeculation(
771      id,
772      'accepted',
773      startTime,
774      suggestionLength,
775      messages,
776      boundary,
777      {
778        message_count: messages.length,
779        time_saved_ms: timeSavedMs,
780        is_pipelined: isPipelined,
781      },
782    )
783  
784    if (timeSavedMs > 0) {
785      const entry: SpeculationAcceptMessage = {
786        type: 'speculation-accept',
787        timestamp: new Date().toISOString(),
788        timeSavedMs,
789      }
790      void appendFile(getTranscriptPath(), jsonStringify(entry) + '\n', {
791        mode: 0o600,
792      }).catch(() => {
793        logForDebugging(
794          '[Speculation] Failed to write speculation-accept to transcript',
795        )
796      })
797    }
798  
799    return { messages, boundary, timeSavedMs }
800  }
801  
802  export function abortSpeculation(setAppState: SetAppState): void {
803    setAppState(prev => {
804      if (prev.speculation.status !== 'active') return prev
805  
806      const {
807        id,
808        abort,
809        startTime,
810        boundary,
811        suggestionLength,
812        messagesRef,
813        isPipelined,
814      } = prev.speculation
815  
816      logForDebugging(`[Speculation] Aborting ${id}`)
817  
818      logSpeculation(
819        id,
820        'aborted',
821        startTime,
822        suggestionLength,
823        messagesRef.current,
824        boundary,
825        { abort_reason: 'user_typed', is_pipelined: isPipelined },
826      )
827  
828      abort()
829      safeRemoveOverlay(getOverlayPath(id))
830  
831      return { ...prev, speculation: IDLE_SPECULATION_STATE }
832    })
833  }
834  
835  export async function handleSpeculationAccept(
836    speculationState: ActiveSpeculationState,
837    speculationSessionTimeSavedMs: number,
838    setAppState: SetAppState,
839    input: string,
840    deps: {
841      setMessages: (f: (prev: Message[]) => Message[]) => void
842      readFileState: { current: FileStateCache }
843      cwd: string
844    },
845  ): Promise<{ queryRequired: boolean }> {
846    try {
847      const { setMessages, readFileState, cwd } = deps
848  
849      // Clear prompt suggestion state. logOutcomeAtSubmission logged the accept
850      // but was called with skipReset to avoid aborting speculation before we use it.
851      setAppState(prev => {
852        if (
853          prev.promptSuggestion.text === null &&
854          prev.promptSuggestion.promptId === null
855        ) {
856          return prev
857        }
858        return {
859          ...prev,
860          promptSuggestion: {
861            text: null,
862            promptId: null,
863            shownAt: 0,
864            acceptedAt: 0,
865            generationRequestId: null,
866          },
867        }
868      })
869  
870      // Capture speculation messages before any state updates - must be stable reference
871      const speculationMessages = speculationState.messagesRef.current
872      let cleanMessages = prepareMessagesForInjection(speculationMessages)
873  
874      // Inject user message first for instant visual feedback before any async work
875      const userMessage = createUserMessage({ content: input })
876      setMessages(prev => [...prev, userMessage])
877  
878      const result = await acceptSpeculation(
879        speculationState,
880        setAppState,
881        cleanMessages.length,
882      )
883  
884      const isComplete = result?.boundary?.type === 'complete'
885  
886      // When speculation didn't complete, the follow-up query needs the
887      // conversation to end with a user message. Drop trailing assistant
888      // messages — models that don't support prefill
889      // reject conversations ending with an assistant turn. The model will
890      // regenerate this content in the follow-up query.
891      if (!isComplete) {
892        const lastNonAssistant = cleanMessages.findLastIndex(
893          m => m.type !== 'assistant',
894        )
895        cleanMessages = cleanMessages.slice(0, lastNonAssistant + 1)
896      }
897  
898      const timeSavedMs = result?.timeSavedMs ?? 0
899      const newSessionTotal = speculationSessionTimeSavedMs + timeSavedMs
900      const feedbackMessage = createSpeculationFeedbackMessage(
901        cleanMessages,
902        result?.boundary ?? null,
903        timeSavedMs,
904        newSessionTotal,
905      )
906  
907      // Inject speculated messages
908      setMessages(prev => [...prev, ...cleanMessages])
909  
910      const extracted = extractReadFilesFromMessages(
911        cleanMessages,
912        cwd,
913        READ_FILE_STATE_CACHE_SIZE,
914      )
915      readFileState.current = mergeFileStateCaches(
916        readFileState.current,
917        extracted,
918      )
919  
920      if (feedbackMessage) {
921        setMessages(prev => [...prev, feedbackMessage])
922      }
923  
924      logForDebugging(
925        `[Speculation] ${result?.boundary?.type ?? 'incomplete'}, injected ${cleanMessages.length} messages`,
926      )
927  
928      // Promote pipelined suggestion if speculation completed fully
929      if (isComplete && speculationState.pipelinedSuggestion) {
930        const { text, promptId, generationRequestId } =
931          speculationState.pipelinedSuggestion
932        logForDebugging(
933          `[Speculation] Promoting pipelined suggestion: "${text.slice(0, 50)}..."`,
934        )
935        setAppState(prev => ({
936          ...prev,
937          promptSuggestion: {
938            text,
939            promptId,
940            shownAt: Date.now(),
941            acceptedAt: 0,
942            generationRequestId,
943          },
944        }))
945  
946        // Start speculation on the pipelined suggestion
947        const augmentedContext: REPLHookContext = {
948          ...speculationState.contextRef.current,
949          messages: [
950            ...speculationState.contextRef.current.messages,
951            createUserMessage({ content: input }),
952            ...cleanMessages,
953          ],
954        }
955        void startSpeculation(text, augmentedContext, setAppState, true)
956      }
957  
958      return { queryRequired: !isComplete }
959    } catch (error) {
960      // Fail open: log error and fall back to normal query flow
961      /* eslint-disable no-restricted-syntax -- custom fallback message, not toError(e) */
962      logError(
963        error instanceof Error
964          ? error
965          : new Error('handleSpeculationAccept failed'),
966      )
967      /* eslint-enable no-restricted-syntax */
968      logSpeculation(
969        speculationState.id,
970        'error',
971        speculationState.startTime,
972        speculationState.suggestionLength,
973        speculationState.messagesRef.current,
974        speculationState.boundary,
975        {
976          error_type: error instanceof Error ? error.name : 'Unknown',
977          error_message: errorMessage(error).slice(
978            0,
979            200,
980          ) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
981          error_phase:
982            'accept' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
983          is_pipelined: speculationState.isPipelined,
984        },
985      )
986      safeRemoveOverlay(getOverlayPath(speculationState.id))
987      resetSpeculationState(setAppState)
988      // Query required so user's message is processed normally (without speculated work)
989      return { queryRequired: true }
990    }
991  }