subagent-cli 0.3.0 → 0.3.1

package/src/utils/sessionRestore.ts

@@ -1,551 +0,0 @@
-import { feature } from 'src/_stubs/bun-bundle.js'
-import type { UUID } from 'crypto'
-import { dirname } from 'path'
-import {
-  getMainLoopModelOverride,
-  getSessionId,
-  setMainLoopModelOverride,
-  setMainThreadAgentType,
-  setOriginalCwd,
-  switchSession,
-} from '../bootstrap/state.js'
-import { clearSystemPromptSections } from '../constants/systemPromptSections.js'
-import { restoreCostStateForSession } from '../cost-tracker.js'
-import type { AppState } from '../state/AppState.js'
-import type { AgentColorName } from '../tools/AgentTool/agentColorManager.js'
-import {
-  type AgentDefinition,
-  type AgentDefinitionsResult,
-  getActiveAgentsFromList,
-  getAgentDefinitionsWithOverrides,
-} from '../tools/AgentTool/loadAgentsDir.js'
-import { TODO_WRITE_TOOL_NAME } from '../tools/TodoWriteTool/constants.js'
-import { asSessionId } from '../types/ids.js'
-import type {
-  AttributionSnapshotMessage,
-  ContextCollapseCommitEntry,
-  ContextCollapseSnapshotEntry,
-  PersistedWorktreeSession,
-} from '../types/logs.js'
-import type { Message } from '../types/message.js'
-import { renameRecordingForSession } from './asciicast.js'
-import { clearMemoryFileCaches } from './claudemd.js'
-import {
-  type AttributionState,
-  attributionRestoreStateFromLog,
-  restoreAttributionStateFromSnapshots,
-} from './commitAttribution.js'
-import { updateSessionName } from './concurrentSessions.js'
-import { getCwd } from './cwd.js'
-import { logForDebugging } from './debug.js'
-import type { FileHistorySnapshot } from './fileHistory.js'
-import { fileHistoryRestoreStateFromLog } from './fileHistory.js'
-import { createSystemMessage } from './messages.js'
-import { parseUserSpecifiedModel } from './model/model.js'
-import { getPlansDirectory } from './plans.js'
-import { setCwd } from './Shell.js'
-import {
-  adoptResumedSessionFile,
-  recordContentReplacement,
-  resetSessionFilePointer,
-  restoreSessionMetadata,
-  saveMode,
-  saveWorktreeState,
-} from './sessionStorage.js'
-import { isTodoV2Enabled } from './tasks.js'
-import type { TodoList } from './todo/types.js'
-import { TodoListSchema } from './todo/types.js'
-import type { ContentReplacementRecord } from './toolResultStorage.js'
-import {
-  getCurrentWorktreeSession,
-  restoreWorktreeSession,
-} from './worktree.js'
-
-type ResumeResult = {
-  messages?: Message[]
-  fileHistorySnapshots?: FileHistorySnapshot[]
-  attributionSnapshots?: AttributionSnapshotMessage[]
-  contextCollapseCommits?: ContextCollapseCommitEntry[]
-  contextCollapseSnapshot?: ContextCollapseSnapshotEntry
-}
-
-/**
- * Scan the transcript for the last TodoWrite tool_use block and return its todos.
- * Used to hydrate AppState.todos on SDK --resume so the model's todo list
- * survives session restarts without file persistence.
- */
-function extractTodosFromTranscript(messages: Message[]): TodoList {
-  for (let i = messages.length - 1; i >= 0; i--) {
-    const msg = messages[i]
-    if (msg?.type !== 'assistant') continue
-    const toolUse = msg.message.content.find(
-      block => block.type === 'tool_use' && block.name === TODO_WRITE_TOOL_NAME,
-    )
-    if (!toolUse || toolUse.type !== 'tool_use') continue
-    const input = toolUse.input
-    if (input === null || typeof input !== 'object') return []
-    const parsed = TodoListSchema().safeParse(
-      (input as Record<string, unknown>).todos,
-    )
-    return parsed.success ? parsed.data : []
-  }
-  return []
-}
-
-/**
- * Restore session state (file history, attribution, todos) from log on resume.
- * Used by both SDK (print.ts) and interactive (REPL.tsx, main.tsx) resume paths.
- */
-export function restoreSessionStateFromLog(
-  result: ResumeResult,
-  setAppState: (f: (prev: AppState) => AppState) => void,
-): void {
-  // Restore file history state
-  if (result.fileHistorySnapshots && result.fileHistorySnapshots.length > 0) {
-    fileHistoryRestoreStateFromLog(result.fileHistorySnapshots, newState => {
-      setAppState(prev => ({ ...prev, fileHistory: newState }))
-    })
-  }
-
-  // Restore attribution state (ant-only feature)
-  if (
-    feature('COMMIT_ATTRIBUTION') &&
-    result.attributionSnapshots &&
-    result.attributionSnapshots.length > 0
-  ) {
-    attributionRestoreStateFromLog(result.attributionSnapshots, newState => {
-      setAppState(prev => ({ ...prev, attribution: newState }))
-    })
-  }
-
-  // Restore context-collapse commit log + staged snapshot. Must run before
-  // the first query() so projectView() can rebuild the collapsed view from
-  // the resumed Message[]. Called unconditionally (even with
-  // undefined/empty entries) because restoreFromEntries resets the store
-  // first — without that, an in-session /resume into a session with no
-  // commits would leave the prior session's stale commit log intact.
-  if (feature('CONTEXT_COLLAPSE')) {
-    /* eslint-disable @typescript-eslint/no-require-imports */
-    ;(
-      require('../services/contextCollapse/persist.js') as typeof import('../services/contextCollapse/persist.js')
-    ).restoreFromEntries(
-      result.contextCollapseCommits ?? [],
-      result.contextCollapseSnapshot,
-    )
-    /* eslint-enable @typescript-eslint/no-require-imports */
-  }
-
-  // Restore TodoWrite state from transcript (SDK/non-interactive only).
-  // Interactive mode uses file-backed v2 tasks, so AppState.todos is unused there.
-  if (!isTodoV2Enabled() && result.messages && result.messages.length > 0) {
-    const todos = extractTodosFromTranscript(result.messages)
-    if (todos.length > 0) {
-      const agentId = getSessionId()
-      setAppState(prev => ({
-        ...prev,
-        todos: { ...prev.todos, [agentId]: todos },
-      }))
-    }
-  }
-}
-
-/**
- * Compute restored attribution state from log snapshots.
- * Used for computing initial state before render (e.g., main.tsx --continue).
- * Returns undefined if attribution feature is disabled or no snapshots exist.
- */
-export function computeRestoredAttributionState(
-  result: ResumeResult,
-): AttributionState | undefined {
-  if (
-    feature('COMMIT_ATTRIBUTION') &&
-    result.attributionSnapshots &&
-    result.attributionSnapshots.length > 0
-  ) {
-    return restoreAttributionStateFromSnapshots(result.attributionSnapshots)
-  }
-  return undefined
-}
-
-/**
- * Compute standalone agent context (name/color) for session resume.
- * Used for computing initial state before render (per CLAUDE.md guidelines).
- * Returns undefined if no name/color is set on the session.
- */
-export function computeStandaloneAgentContext(
-  agentName: string | undefined,
-  agentColor: string | undefined,
-): AppState['standaloneAgentContext'] | undefined {
-  if (!agentName && !agentColor) {
-    return undefined
-  }
-  return {
-    name: agentName ?? '',
-    color: (agentColor === 'default' ? undefined : agentColor) as
-      | AgentColorName
-      | undefined,
-  }
-}
-
-/**
- * Restore agent setting from a resumed session.
- *
- * When resuming a conversation that used a custom agent, this re-applies the
- * agent type and model override (unless the user specified --agent on the CLI).
- * Mutates bootstrap state via setMainThreadAgentType / setMainLoopModelOverride.
- *
- * Returns the restored agent definition and its agentType string, or undefined
- * if no agent was restored.
- */
-export function restoreAgentFromSession(
-  agentSetting: string | undefined,
-  currentAgentDefinition: AgentDefinition | undefined,
-  agentDefinitions: AgentDefinitionsResult,
-): {
-  agentDefinition: AgentDefinition | undefined
-  agentType: string | undefined
-} {
-  // If user already specified --agent on CLI, keep that definition
-  if (currentAgentDefinition) {
-    return { agentDefinition: currentAgentDefinition, agentType: undefined }
-  }
-
-  // If session had no agent, clear any stale bootstrap state
-  if (!agentSetting) {
-    setMainThreadAgentType(undefined)
-    return { agentDefinition: undefined, agentType: undefined }
-  }
-
-  const resumedAgent = agentDefinitions.activeAgents.find(
-    agent => agent.agentType === agentSetting,
-  )
-  if (!resumedAgent) {
-    logForDebugging(
-      `Resumed session had agent "${agentSetting}" but it is no longer available. Using default behavior.`,
-    )
-    setMainThreadAgentType(undefined)
-    return { agentDefinition: undefined, agentType: undefined }
-  }
-
-  setMainThreadAgentType(resumedAgent.agentType)
-
-  // Apply agent's model if user didn't specify one
-  if (
-    !getMainLoopModelOverride() &&
-    resumedAgent.model &&
-    resumedAgent.model !== 'inherit'
-  ) {
-    setMainLoopModelOverride(parseUserSpecifiedModel(resumedAgent.model))
-  }
-
-  return { agentDefinition: resumedAgent, agentType: resumedAgent.agentType }
-}
-
-/**
- * Refresh agent definitions after a coordinator/normal mode switch.
- *
- * When resuming a session that was in a different mode (coordinator vs normal),
- * the built-in agents need to be re-derived to match the new mode. CLI-provided
- * agents (from --agents flag) are merged back in.
- */
-export async function refreshAgentDefinitionsForModeSwitch(
-  modeWasSwitched: boolean,
-  currentCwd: string,
-  cliAgents: AgentDefinition[],
-  currentAgentDefinitions: AgentDefinitionsResult,
-): Promise<AgentDefinitionsResult> {
-  if (!feature('COORDINATOR_MODE') || !modeWasSwitched) {
-    return currentAgentDefinitions
-  }
-
-  // Re-derive agent definitions after mode switch so built-in agents
-  // reflect the new coordinator/normal mode
-  getAgentDefinitionsWithOverrides.cache.clear?.()
-  const freshAgentDefs = await getAgentDefinitionsWithOverrides(currentCwd)
-  const freshAllAgents = [...freshAgentDefs.allAgents, ...cliAgents]
-  return {
-    ...freshAgentDefs,
-    allAgents: freshAllAgents,
-    activeAgents: getActiveAgentsFromList(freshAllAgents),
-  }
-}
-
-/**
- * Result of processing a resumed/continued conversation for rendering.
- */
-export type ProcessedResume = {
-  messages: Message[]
-  fileHistorySnapshots?: FileHistorySnapshot[]
-  contentReplacements?: ContentReplacementRecord[]
-  agentName: string | undefined
-  agentColor: AgentColorName | undefined
-  restoredAgentDef: AgentDefinition | undefined
-  initialState: AppState
-}
-
-/**
- * Subset of the coordinator mode module API needed for session resume.
- */
-type CoordinatorModeApi = {
-  matchSessionMode(mode?: string): string | undefined
-  isCoordinatorMode(): boolean
-}
-
-/**
- * The loaded conversation data (return type of loadConversationForResume).
- */
-type ResumeLoadResult = {
-  messages: Message[]
-  fileHistorySnapshots?: FileHistorySnapshot[]
-  attributionSnapshots?: AttributionSnapshotMessage[]
-  contentReplacements?: ContentReplacementRecord[]
-  contextCollapseCommits?: ContextCollapseCommitEntry[]
-  contextCollapseSnapshot?: ContextCollapseSnapshotEntry
-  sessionId: UUID | undefined
-  agentName?: string
-  agentColor?: string
-  agentSetting?: string
-  customTitle?: string
-  tag?: string
-  mode?: 'coordinator' | 'normal'
-  worktreeSession?: PersistedWorktreeSession | null
-  prNumber?: number
-  prUrl?: string
-  prRepository?: string
-}
-
-/**
- * Restore the worktree working directory on resume. The transcript records
- * the last worktree enter/exit; if the session crashed while inside a
- * worktree (last entry = session object, not null), cd back into it.
- *
- * process.chdir is the TOCTOU-safe existence check — it throws ENOENT if
- * the /exit dialog removed the directory, or if the user deleted it
- * manually between sessions.
- *
- * When --worktree already created a fresh worktree, that takes precedence
- * over the resumed session's state. restoreSessionMetadata just overwrote
- * project.currentSessionWorktree with the stale transcript value, so
- * re-assert the fresh worktree here before adoptResumedSessionFile writes
- * it back to disk.
- */
-export function restoreWorktreeForResume(
-  worktreeSession: PersistedWorktreeSession | null | undefined,
-): void {
-  const fresh = getCurrentWorktreeSession()
-  if (fresh) {
-    saveWorktreeState(fresh)
-    return
-  }
-  if (!worktreeSession) return
-
-  try {
-    process.chdir(worktreeSession.worktreePath)
-  } catch {
-    // Directory is gone. Override the stale cache so the next
-    // reAppendSessionMetadata records "exited" instead of re-persisting
-    // a path that no longer exists.
-    saveWorktreeState(null)
-    return
-  }
-
-  setCwd(worktreeSession.worktreePath)
-  setOriginalCwd(getCwd())
-  // projectRoot is intentionally NOT set here. The transcript doesn't record
-  // whether the worktree was entered via --worktree (which sets projectRoot)
-  // or EnterWorktreeTool (which doesn't). Leaving projectRoot stable matches
-  // EnterWorktreeTool's behavior — skills/history stay anchored to the
-  // original project.
-  restoreWorktreeSession(worktreeSession)
-  // The /resume slash command calls this mid-session after caches have been
-  // populated against the old cwd. Cheap no-ops for the CLI-flag path
-  // (caches aren't populated yet there).
-  clearMemoryFileCaches()
-  clearSystemPromptSections()
-  getPlansDirectory.cache.clear?.()
-}
-
-/**
- * Undo restoreWorktreeForResume before a mid-session /resume switches to
- * another session. Without this, /resume from a worktree session to a
- * non-worktree session leaves the user in the old worktree directory with
- * currentWorktreeSession still pointing at the prior session. /resume to a
- * *different* worktree fails entirely — the getCurrentWorktreeSession()
- * guard above blocks the switch.
- *
- * Not needed by CLI --resume/--continue: those run once at startup where
- * getCurrentWorktreeSession() is only truthy if --worktree was used (fresh
- * worktree that should take precedence, handled by the re-assert above).
- */
-export function exitRestoredWorktree(): void {
-  const current = getCurrentWorktreeSession()
-  if (!current) return
-
-  restoreWorktreeSession(null)
-  // Worktree state changed, so cached prompt sections that reference it are
-  // stale whether or not chdir succeeds below.
-  clearMemoryFileCaches()
-  clearSystemPromptSections()
-  getPlansDirectory.cache.clear?.()
-
-  try {
-    process.chdir(current.originalCwd)
-  } catch {
-    // Original dir is gone (rare). Stay put — restoreWorktreeForResume
-    // will cd into the target worktree next if there is one.
-    return
-  }
-  setCwd(current.originalCwd)
-  setOriginalCwd(getCwd())
-}
-
-/**
- * Process a loaded conversation for resume/continue.
- *
- * Handles coordinator mode matching, session ID setup, agent restoration,
- * mode persistence, and initial state computation. Called by both --continue
- * and --resume paths in main.tsx.
- */
-export async function processResumedConversation(
-  result: ResumeLoadResult,
-  opts: {
-    forkSession: boolean
-    sessionIdOverride?: string
-    transcriptPath?: string
-    includeAttribution?: boolean
-  },
-  context: {
-    modeApi: CoordinatorModeApi | null
-    mainThreadAgentDefinition: AgentDefinition | undefined
-    agentDefinitions: AgentDefinitionsResult
-    currentCwd: string
-    cliAgents: AgentDefinition[]
-    initialState: AppState
-  },
-): Promise<ProcessedResume> {
-  // Match coordinator/normal mode to the resumed session
-  let modeWarning: string | undefined
-  if (feature('COORDINATOR_MODE')) {
-    modeWarning = context.modeApi?.matchSessionMode(result.mode)
-    if (modeWarning) {
-      result.messages.push(createSystemMessage(modeWarning, 'warning'))
-    }
-  }
-
-  // Reuse the resumed session's ID unless --fork-session is specified
-  if (!opts.forkSession) {
-    const sid = opts.sessionIdOverride ?? result.sessionId
-    if (sid) {
-      // When resuming from a different project directory (git worktrees,
-      // cross-project), transcriptPath points to the actual file; its dirname
-      // is the project dir. Otherwise the session lives in the current project.
-      switchSession(
-        asSessionId(sid),
-        opts.transcriptPath ? dirname(opts.transcriptPath) : null,
-      )
-      // Rename asciicast recording to match the resumed session ID so
-      // getSessionRecordingPaths() can discover it during /share
-      await renameRecordingForSession()
-      await resetSessionFilePointer()
-      restoreCostStateForSession(sid)
-    }
-  } else if (result.contentReplacements?.length) {
-    // --fork-session keeps the fresh startup session ID. useLogMessages will
-    // copy source messages into the new JSONL via recordTranscript, but
-    // content-replacement entries are a separate entry type only written by
-    // recordContentReplacement (which query.ts calls for newlyReplaced, never
-    // the pre-loaded records). Without this seed, `claude -r {newSessionId}`
-    // finds source tool_use_ids in messages but no matching replacement records
-    // → they're classified as FROZEN → full content sent (cache miss, permanent
-    // overage). insertContentReplacement stamps sessionId = getSessionId() =
-    // the fresh ID, so loadTranscriptFile's keyed lookup will match.
-    await recordContentReplacement(result.contentReplacements)
-  }
-
-  // Restore session metadata so /status shows the saved name and metadata
-  // is re-appended on session exit. Fork doesn't take ownership of the
-  // original session's worktree — a "Remove" on the fork's exit dialog
-  // would delete a worktree the original session still references — so
-  // strip worktreeSession from the fork path so the cache stays unset.
-  restoreSessionMetadata(
-    opts.forkSession ? { ...result, worktreeSession: undefined } : result,
-  )
-
-  if (!opts.forkSession) {
-    // Cd back into the worktree the session was in when it last exited.
-    // Done after restoreSessionMetadata (which caches the worktree state
-    // from the transcript) so if the directory is gone we can override
-    // the cache before adoptResumedSessionFile writes it.
-    restoreWorktreeForResume(result.worktreeSession)
-
-    // Point sessionFile at the resumed transcript and re-append metadata
-    // now. resetSessionFilePointer above nulled it (so the old fresh-session
-    // path doesn't leak), but that blocks reAppendSessionMetadata — which
-    // bails on null — from running in the exit cleanup handler. For fork,
-    // useLogMessages populates a *new* file via recordTranscript on REPL
-    // mount; the normal lazy-materialize path is correct there.
-    adoptResumedSessionFile()
-  }
-
-  // Restore context-collapse commit log + staged snapshot. The interactive
-  // /resume path goes through restoreSessionStateFromLog (REPL.tsx); CLI
-  // --continue/--resume goes through here instead. Called unconditionally
-  // — see the restoreSessionStateFromLog callsite above for why.
-  if (feature('CONTEXT_COLLAPSE')) {
-    /* eslint-disable @typescript-eslint/no-require-imports */
-    ;(
-      require('../services/contextCollapse/persist.js') as typeof import('../services/contextCollapse/persist.js')
-    ).restoreFromEntries(
-      result.contextCollapseCommits ?? [],
-      result.contextCollapseSnapshot,
-    )
-    /* eslint-enable @typescript-eslint/no-require-imports */
-  }
-
-  // Restore agent setting from resumed session
-  const { agentDefinition: restoredAgent, agentType: resumedAgentType } =
-    restoreAgentFromSession(
-      result.agentSetting,
-      context.mainThreadAgentDefinition,
-      context.agentDefinitions,
-    )
-
-  // Persist the current mode so future resumes know what mode this session was in
-  if (feature('COORDINATOR_MODE')) {
-    saveMode(context.modeApi?.isCoordinatorMode() ? 'coordinator' : 'normal')
-  }
-
-  // Compute initial state before render (per CLAUDE.md guidelines)
-  const restoredAttribution = opts.includeAttribution
-    ? computeRestoredAttributionState(result)
-    : undefined
-  const standaloneAgentContext = computeStandaloneAgentContext(
-    result.agentName,
-    result.agentColor,
-  )
-  void updateSessionName(result.agentName)
-  const refreshedAgentDefs = await refreshAgentDefinitionsForModeSwitch(
-    !!modeWarning,
-    context.currentCwd,
-    context.cliAgents,
-    context.agentDefinitions,
-  )
-
-  return {
-    messages: result.messages,
-    fileHistorySnapshots: result.fileHistorySnapshots,
-    contentReplacements: result.contentReplacements,
-    agentName: result.agentName,
-    agentColor: (result.agentColor === 'default'
-      ? undefined
-      : result.agentColor) as AgentColorName | undefined,
-    restoredAgentDef: restoredAgent,
-    initialState: {
-      ...context.initialState,
-      ...(resumedAgentType && { agent: resumedAgentType }),
-      ...(restoredAttribution && { attribution: restoredAttribution }),
-      ...(standaloneAgentContext && { standaloneAgentContext }),
-      agentDefinitions: refreshedAgentDefs,
-    },
-  }
-}