1  import { feature } from 'bun:bundle'
  2  import type { UUID } from 'crypto'
  3  import { dirname } from 'path'
  4  import {
  5    getMainLoopModelOverride,
  6    getSessionId,
  7    setMainLoopModelOverride,
  8    setMainThreadAgentType,
  9    setOriginalCwd,
 10    switchSession,
 11  } from '../bootstrap/state.js'
 12  import { clearSystemPromptSections } from '../constants/systemPromptSections.js'
 13  import { restoreCostStateForSession } from '../cost-tracker.js'
 14  import type { AppState } from '../state/AppState.js'
 15  import type { AgentColorName } from '../tools/AgentTool/agentColorManager.js'
 16  import {
 17    type AgentDefinition,
 18    type AgentDefinitionsResult,
 19    getActiveAgentsFromList,
 20    getAgentDefinitionsWithOverrides,
 21  } from '../tools/AgentTool/loadAgentsDir.js'
 22  import { TODO_WRITE_TOOL_NAME } from '../tools/TodoWriteTool/constants.js'
 23  import { asSessionId } from '../types/ids.js'
 24  import type {
 25    AttributionSnapshotMessage,
 26    ContextCollapseCommitEntry,
 27    ContextCollapseSnapshotEntry,
 28    PersistedWorktreeSession,
 29  } from '../types/logs.js'
 30  import type { Message } from '../types/message.js'
 31  import { renameRecordingForSession } from './asciicast.js'
 32  import { clearMemoryFileCaches } from './claudemd.js'
 33  import {
 34    type AttributionState,
 35    attributionRestoreStateFromLog,
 36    restoreAttributionStateFromSnapshots,
 37  } from './commitAttribution.js'
 38  import { updateSessionName } from './concurrentSessions.js'
 39  import { getCwd } from './cwd.js'
 40  import { logForDebugging } from './debug.js'
 41  import type { FileHistorySnapshot } from './fileHistory.js'
 42  import { fileHistoryRestoreStateFromLog } from './fileHistory.js'
 43  import { createSystemMessage } from './messages.js'
 44  import { parseUserSpecifiedModel } from './model/model.js'
 45  import { getPlansDirectory } from './plans.js'
 46  import { setCwd } from './Shell.js'
 47  import {
 48    adoptResumedSessionFile,
 49    recordContentReplacement,
 50    resetSessionFilePointer,
 51    restoreSessionMetadata,
 52    saveMode,
 53    saveWorktreeState,
 54  } from './sessionStorage.js'
 55  import { isTodoV2Enabled } from './tasks.js'
 56  import type { TodoList } from './todo/types.js'
 57  import { TodoListSchema } from './todo/types.js'
 58  import type { ContentReplacementRecord } from './toolResultStorage.js'
 59  import {
 60    getCurrentWorktreeSession,
 61    restoreWorktreeSession,
 62  } from './worktree.js'
 63  
 64  type ResumeResult = {
 65    messages?: Message[]
 66    fileHistorySnapshots?: FileHistorySnapshot[]
 67    attributionSnapshots?: AttributionSnapshotMessage[]
 68    contextCollapseCommits?: ContextCollapseCommitEntry[]
 69    contextCollapseSnapshot?: ContextCollapseSnapshotEntry
 70  }
 71  
 72  /**
 73   * Scan the transcript for the last TodoWrite tool_use block and return its todos.
 74   * Used to hydrate AppState.todos on SDK --resume so the model's todo list
 75   * survives session restarts without file persistence.
 76   */
 77  function extractTodosFromTranscript(messages: Message[]): TodoList {
 78    for (let i = messages.length - 1; i >= 0; i--) {
 79      const msg = messages[i]
 80      if (msg?.type !== 'assistant') continue
 81      const toolUse = msg.message.content.find(
 82        block => block.type === 'tool_use' && block.name === TODO_WRITE_TOOL_NAME,
 83      )
 84      if (!toolUse || toolUse.type !== 'tool_use') continue
 85      const input = toolUse.input
 86      if (input === null || typeof input !== 'object') return []
 87      const parsed = TodoListSchema().safeParse(
 88        (input as Record<string, unknown>).todos,
 89      )
 90      return parsed.success ? parsed.data : []
 91    }
 92    return []
 93  }
 94  
 95  /**
 96   * Restore session state (file history, attribution, todos) from log on resume.
 97   * Used by both SDK (print.ts) and interactive (REPL.tsx, main.tsx) resume paths.
 98   */
 99  export function restoreSessionStateFromLog(
100    result: ResumeResult,
101    setAppState: (f: (prev: AppState) => AppState) => void,
102  ): void {
103    // Restore file history state
104    if (result.fileHistorySnapshots && result.fileHistorySnapshots.length > 0) {
105      fileHistoryRestoreStateFromLog(result.fileHistorySnapshots, newState => {
106        setAppState(prev => ({ ...prev, fileHistory: newState }))
107      })
108    }
109  
110    // Restore attribution state (ant-only feature)
111    if (
112      feature('COMMIT_ATTRIBUTION') &&
113      result.attributionSnapshots &&
114      result.attributionSnapshots.length > 0
115    ) {
116      attributionRestoreStateFromLog(result.attributionSnapshots, newState => {
117        setAppState(prev => ({ ...prev, attribution: newState }))
118      })
119    }
120  
121    // Restore context-collapse commit log + staged snapshot. Must run before
122    // the first query() so projectView() can rebuild the collapsed view from
123    // the resumed Message[]. Called unconditionally (even with
124    // undefined/empty entries) because restoreFromEntries resets the store
125    // first — without that, an in-session /resume into a session with no
126    // commits would leave the prior session's stale commit log intact.
127    if (feature('CONTEXT_COLLAPSE')) {
128      /* eslint-disable @typescript-eslint/no-require-imports */
129      ;(
130        require('../services/contextCollapse/persist.js') as typeof import('../services/contextCollapse/persist.js')
131      ).restoreFromEntries(
132        result.contextCollapseCommits ?? [],
133        result.contextCollapseSnapshot,
134      )
135      /* eslint-enable @typescript-eslint/no-require-imports */
136    }
137  
138    // Restore TodoWrite state from transcript (SDK/non-interactive only).
139    // Interactive mode uses file-backed v2 tasks, so AppState.todos is unused there.
140    if (!isTodoV2Enabled() && result.messages && result.messages.length > 0) {
141      const todos = extractTodosFromTranscript(result.messages)
142      if (todos.length > 0) {
143        const agentId = getSessionId()
144        setAppState(prev => ({
145          ...prev,
146          todos: { ...prev.todos, [agentId]: todos },
147        }))
148      }
149    }
150  }
151  
152  /**
153   * Compute restored attribution state from log snapshots.
154   * Used for computing initial state before render (e.g., main.tsx --continue).
155   * Returns undefined if attribution feature is disabled or no snapshots exist.
156   */
157  export function computeRestoredAttributionState(
158    result: ResumeResult,
159  ): AttributionState | undefined {
160    if (
161      feature('COMMIT_ATTRIBUTION') &&
162      result.attributionSnapshots &&
163      result.attributionSnapshots.length > 0
164    ) {
165      return restoreAttributionStateFromSnapshots(result.attributionSnapshots)
166    }
167    return undefined
168  }
169  
170  /**
171   * Compute standalone agent context (name/color) for session resume.
172   * Used for computing initial state before render (per CLAUDE.md guidelines).
173   * Returns undefined if no name/color is set on the session.
174   */
175  export function computeStandaloneAgentContext(
176    agentName: string | undefined,
177    agentColor: string | undefined,
178  ): AppState['standaloneAgentContext'] | undefined {
179    if (!agentName && !agentColor) {
180      return undefined
181    }
182    return {
183      name: agentName ?? '',
184      color: (agentColor === 'default' ? undefined : agentColor) as
185        | AgentColorName
186        | undefined,
187    }
188  }
189  
190  /**
191   * Restore agent setting from a resumed session.
192   *
193   * When resuming a conversation that used a custom agent, this re-applies the
194   * agent type and model override (unless the user specified --agent on the CLI).
195   * Mutates bootstrap state via setMainThreadAgentType / setMainLoopModelOverride.
196   *
197   * Returns the restored agent definition and its agentType string, or undefined
198   * if no agent was restored.
199   */
200  export function restoreAgentFromSession(
201    agentSetting: string | undefined,
202    currentAgentDefinition: AgentDefinition | undefined,
203    agentDefinitions: AgentDefinitionsResult,
204  ): {
205    agentDefinition: AgentDefinition | undefined
206    agentType: string | undefined
207  } {
208    // If user already specified --agent on CLI, keep that definition
209    if (currentAgentDefinition) {
210      return { agentDefinition: currentAgentDefinition, agentType: undefined }
211    }
212  
213    // If session had no agent, clear any stale bootstrap state
214    if (!agentSetting) {
215      setMainThreadAgentType(undefined)
216      return { agentDefinition: undefined, agentType: undefined }
217    }
218  
219    const resumedAgent = agentDefinitions.activeAgents.find(
220      agent => agent.agentType === agentSetting,
221    )
222    if (!resumedAgent) {
223      logForDebugging(
224        `Resumed session had agent "${agentSetting}" but it is no longer available. Using default behavior.`,
225      )
226      setMainThreadAgentType(undefined)
227      return { agentDefinition: undefined, agentType: undefined }
228    }
229  
230    setMainThreadAgentType(resumedAgent.agentType)
231  
232    // Apply agent's model if user didn't specify one
233    if (
234      !getMainLoopModelOverride() &&
235      resumedAgent.model &&
236      resumedAgent.model !== 'inherit'
237    ) {
238      setMainLoopModelOverride(parseUserSpecifiedModel(resumedAgent.model))
239    }
240  
241    return { agentDefinition: resumedAgent, agentType: resumedAgent.agentType }
242  }
243  
244  /**
245   * Refresh agent definitions after a coordinator/normal mode switch.
246   *
247   * When resuming a session that was in a different mode (coordinator vs normal),
248   * the built-in agents need to be re-derived to match the new mode. CLI-provided
249   * agents (from --agents flag) are merged back in.
250   */
251  export async function refreshAgentDefinitionsForModeSwitch(
252    modeWasSwitched: boolean,
253    currentCwd: string,
254    cliAgents: AgentDefinition[],
255    currentAgentDefinitions: AgentDefinitionsResult,
256  ): Promise<AgentDefinitionsResult> {
257    if (!feature('COORDINATOR_MODE') || !modeWasSwitched) {
258      return currentAgentDefinitions
259    }
260  
261    // Re-derive agent definitions after mode switch so built-in agents
262    // reflect the new coordinator/normal mode
263    getAgentDefinitionsWithOverrides.cache.clear?.()
264    const freshAgentDefs = await getAgentDefinitionsWithOverrides(currentCwd)
265    const freshAllAgents = [...freshAgentDefs.allAgents, ...cliAgents]
266    return {
267      ...freshAgentDefs,
268      allAgents: freshAllAgents,
269      activeAgents: getActiveAgentsFromList(freshAllAgents),
270    }
271  }
272  
273  /**
274   * Result of processing a resumed/continued conversation for rendering.
275   */
276  export type ProcessedResume = {
277    messages: Message[]
278    fileHistorySnapshots?: FileHistorySnapshot[]
279    contentReplacements?: ContentReplacementRecord[]
280    agentName: string | undefined
281    agentColor: AgentColorName | undefined
282    restoredAgentDef: AgentDefinition | undefined
283    initialState: AppState
284  }
285  
286  /**
287   * Subset of the coordinator mode module API needed for session resume.
288   */
289  type CoordinatorModeApi = {
290    matchSessionMode(mode?: string): string | undefined
291    isCoordinatorMode(): boolean
292  }
293  
294  /**
295   * The loaded conversation data (return type of loadConversationForResume).
296   */
297  type ResumeLoadResult = {
298    messages: Message[]
299    fileHistorySnapshots?: FileHistorySnapshot[]
300    attributionSnapshots?: AttributionSnapshotMessage[]
301    contentReplacements?: ContentReplacementRecord[]
302    contextCollapseCommits?: ContextCollapseCommitEntry[]
303    contextCollapseSnapshot?: ContextCollapseSnapshotEntry
304    sessionId: UUID | undefined
305    agentName?: string
306    agentColor?: string
307    agentSetting?: string
308    customTitle?: string
309    tag?: string
310    mode?: 'coordinator' | 'normal'
311    worktreeSession?: PersistedWorktreeSession | null
312    prNumber?: number
313    prUrl?: string
314    prRepository?: string
315  }
316  
317  /**
318   * Restore the worktree working directory on resume. The transcript records
319   * the last worktree enter/exit; if the session crashed while inside a
320   * worktree (last entry = session object, not null), cd back into it.
321   *
322   * process.chdir is the TOCTOU-safe existence check — it throws ENOENT if
323   * the /exit dialog removed the directory, or if the user deleted it
324   * manually between sessions.
325   *
326   * When --worktree already created a fresh worktree, that takes precedence
327   * over the resumed session's state. restoreSessionMetadata just overwrote
328   * project.currentSessionWorktree with the stale transcript value, so
329   * re-assert the fresh worktree here before adoptResumedSessionFile writes
330   * it back to disk.
331   */
332  export function restoreWorktreeForResume(
333    worktreeSession: PersistedWorktreeSession | null | undefined,
334  ): void {
335    const fresh = getCurrentWorktreeSession()
336    if (fresh) {
337      saveWorktreeState(fresh)
338      return
339    }
340    if (!worktreeSession) return
341  
342    try {
343      process.chdir(worktreeSession.worktreePath)
344    } catch {
345      // Directory is gone. Override the stale cache so the next
346      // reAppendSessionMetadata records "exited" instead of re-persisting
347      // a path that no longer exists.
348      saveWorktreeState(null)
349      return
350    }
351  
352    setCwd(worktreeSession.worktreePath)
353    setOriginalCwd(getCwd())
354    // projectRoot is intentionally NOT set here. The transcript doesn't record
355    // whether the worktree was entered via --worktree (which sets projectRoot)
356    // or EnterWorktreeTool (which doesn't). Leaving projectRoot stable matches
357    // EnterWorktreeTool's behavior — skills/history stay anchored to the
358    // original project.
359    restoreWorktreeSession(worktreeSession)
360    // The /resume slash command calls this mid-session after caches have been
361    // populated against the old cwd. Cheap no-ops for the CLI-flag path
362    // (caches aren't populated yet there).
363    clearMemoryFileCaches()
364    clearSystemPromptSections()
365    getPlansDirectory.cache.clear?.()
366  }
367  
368  /**
369   * Undo restoreWorktreeForResume before a mid-session /resume switches to
370   * another session. Without this, /resume from a worktree session to a
371   * non-worktree session leaves the user in the old worktree directory with
372   * currentWorktreeSession still pointing at the prior session. /resume to a
373   * *different* worktree fails entirely — the getCurrentWorktreeSession()
374   * guard above blocks the switch.
375   *
376   * Not needed by CLI --resume/--continue: those run once at startup where
377   * getCurrentWorktreeSession() is only truthy if --worktree was used (fresh
378   * worktree that should take precedence, handled by the re-assert above).
379   */
380  export function exitRestoredWorktree(): void {
381    const current = getCurrentWorktreeSession()
382    if (!current) return
383  
384    restoreWorktreeSession(null)
385    // Worktree state changed, so cached prompt sections that reference it are
386    // stale whether or not chdir succeeds below.
387    clearMemoryFileCaches()
388    clearSystemPromptSections()
389    getPlansDirectory.cache.clear?.()
390  
391    try {
392      process.chdir(current.originalCwd)
393    } catch {
394      // Original dir is gone (rare). Stay put — restoreWorktreeForResume
395      // will cd into the target worktree next if there is one.
396      return
397    }
398    setCwd(current.originalCwd)
399    setOriginalCwd(getCwd())
400  }
401  
402  /**
403   * Process a loaded conversation for resume/continue.
404   *
405   * Handles coordinator mode matching, session ID setup, agent restoration,
406   * mode persistence, and initial state computation. Called by both --continue
407   * and --resume paths in main.tsx.
408   */
409  export async function processResumedConversation(
410    result: ResumeLoadResult,
411    opts: {
412      forkSession: boolean
413      sessionIdOverride?: string
414      transcriptPath?: string
415      includeAttribution?: boolean
416    },
417    context: {
418      modeApi: CoordinatorModeApi | null
419      mainThreadAgentDefinition: AgentDefinition | undefined
420      agentDefinitions: AgentDefinitionsResult
421      currentCwd: string
422      cliAgents: AgentDefinition[]
423      initialState: AppState
424    },
425  ): Promise<ProcessedResume> {
426    // Match coordinator/normal mode to the resumed session
427    let modeWarning: string | undefined
428    if (feature('COORDINATOR_MODE')) {
429      modeWarning = context.modeApi?.matchSessionMode(result.mode)
430      if (modeWarning) {
431        result.messages.push(createSystemMessage(modeWarning, 'warning'))
432      }
433    }
434  
435    // Reuse the resumed session's ID unless --fork-session is specified
436    if (!opts.forkSession) {
437      const sid = opts.sessionIdOverride ?? result.sessionId
438      if (sid) {
439        // When resuming from a different project directory (git worktrees,
440        // cross-project), transcriptPath points to the actual file; its dirname
441        // is the project dir. Otherwise the session lives in the current project.
442        switchSession(
443          asSessionId(sid),
444          opts.transcriptPath ? dirname(opts.transcriptPath) : null,
445        )
446        // Rename asciicast recording to match the resumed session ID so
447        // getSessionRecordingPaths() can discover it during /share
448        await renameRecordingForSession()
449        await resetSessionFilePointer()
450        restoreCostStateForSession(sid)
451      }
452    } else if (result.contentReplacements?.length) {
453      // --fork-session keeps the fresh startup session ID. useLogMessages will
454      // copy source messages into the new JSONL via recordTranscript, but
455      // content-replacement entries are a separate entry type only written by
456      // recordContentReplacement (which query.ts calls for newlyReplaced, never
457      // the pre-loaded records). Without this seed, `claude -r {newSessionId}`
458      // finds source tool_use_ids in messages but no matching replacement records
459      // → they're classified as FROZEN → full content sent (cache miss, permanent
460      // overage). insertContentReplacement stamps sessionId = getSessionId() =
461      // the fresh ID, so loadTranscriptFile's keyed lookup will match.
462      await recordContentReplacement(result.contentReplacements)
463    }
464  
465    // Restore session metadata so /status shows the saved name and metadata
466    // is re-appended on session exit. Fork doesn't take ownership of the
467    // original session's worktree — a "Remove" on the fork's exit dialog
468    // would delete a worktree the original session still references — so
469    // strip worktreeSession from the fork path so the cache stays unset.
470    restoreSessionMetadata(
471      opts.forkSession ? { ...result, worktreeSession: undefined } : result,
472    )
473  
474    if (!opts.forkSession) {
475      // Cd back into the worktree the session was in when it last exited.
476      // Done after restoreSessionMetadata (which caches the worktree state
477      // from the transcript) so if the directory is gone we can override
478      // the cache before adoptResumedSessionFile writes it.
479      restoreWorktreeForResume(result.worktreeSession)
480  
481      // Point sessionFile at the resumed transcript and re-append metadata
482      // now. resetSessionFilePointer above nulled it (so the old fresh-session
483      // path doesn't leak), but that blocks reAppendSessionMetadata — which
484      // bails on null — from running in the exit cleanup handler. For fork,
485      // useLogMessages populates a *new* file via recordTranscript on REPL
486      // mount; the normal lazy-materialize path is correct there.
487      adoptResumedSessionFile()
488    }
489  
490    // Restore context-collapse commit log + staged snapshot. The interactive
491    // /resume path goes through restoreSessionStateFromLog (REPL.tsx); CLI
492    // --continue/--resume goes through here instead. Called unconditionally
493    // — see the restoreSessionStateFromLog callsite above for why.
494    if (feature('CONTEXT_COLLAPSE')) {
495      /* eslint-disable @typescript-eslint/no-require-imports */
496      ;(
497        require('../services/contextCollapse/persist.js') as typeof import('../services/contextCollapse/persist.js')
498      ).restoreFromEntries(
499        result.contextCollapseCommits ?? [],
500        result.contextCollapseSnapshot,
501      )
502      /* eslint-enable @typescript-eslint/no-require-imports */
503    }
504  
505    // Restore agent setting from resumed session
506    const { agentDefinition: restoredAgent, agentType: resumedAgentType } =
507      restoreAgentFromSession(
508        result.agentSetting,
509        context.mainThreadAgentDefinition,
510        context.agentDefinitions,
511      )
512  
513    // Persist the current mode so future resumes know what mode this session was in
514    if (feature('COORDINATOR_MODE')) {
515      saveMode(context.modeApi?.isCoordinatorMode() ? 'coordinator' : 'normal')
516    }
517  
518    // Compute initial state before render (per CLAUDE.md guidelines)
519    const restoredAttribution = opts.includeAttribution
520      ? computeRestoredAttributionState(result)
521      : undefined
522    const standaloneAgentContext = computeStandaloneAgentContext(
523      result.agentName,
524      result.agentColor,
525    )
526    void updateSessionName(result.agentName)
527    const refreshedAgentDefs = await refreshAgentDefinitionsForModeSwitch(
528      !!modeWarning,
529      context.currentCwd,
530      context.cliAgents,
531      context.agentDefinitions,
532    )
533  
534    return {
535      messages: result.messages,
536      fileHistorySnapshots: result.fileHistorySnapshots,
537      contentReplacements: result.contentReplacements,
538      agentName: result.agentName,
539      agentColor: (result.agentColor === 'default'
540        ? undefined
541        : result.agentColor) as AgentColorName | undefined,
542      restoredAgentDef: restoredAgent,
543      initialState: {
544        ...context.initialState,
545        ...(resumedAgentType && { agent: resumedAgentType }),
546        ...(restoredAttribution && { attribution: restoredAttribution }),
547        ...(standaloneAgentContext && { standaloneAgentContext }),
548        agentDefinitions: refreshedAgentDefs,
549      },
550    }
551  }