QueryEngine.ts

/**
 * QueryEngine — session-level orchestrator (docs §03.2).
 *
 * Manages multi-turn conversation history, the API client, tool registry,
 * system prompt, usage tracking, and the file-state cache. Each user message
 * triggers a query() run; the engine yields messages back to the UI/headless
 * layer.
 */
import type { ApiClient } from './services/api/client.js'
import { UsageTracker } from './services/api/usage.js'
import type { Message, MessageCreateParams, Usage } from './services/api/types.js'
import type { BuiltTool, ToolUseContext, FileStateCache } from './Tool.js'
import { createFileStateCache } from './utils/file/readFileState.js'
import { query, type QueryResult, type QueryExitReason } from './query.js'
import { fetchSystemPromptParts } from './context.js'
import type { CanUseTool } from './query/runTools.js'
import { ModelManager } from './cli/modelManager.js'
import type { ModelEntry } from './cli/configFile.js'
import type { HooksRegistry } from './services/hooks/runner.js'
import { runHooks, emptyRegistry } from './services/hooks/runner.js'
import { createSession, appendMessages, loadSession } from './services/session/index.js'
import { setPlanApprovalHandler } from './tools/ExitPlanModeTool/ExitPlanModeTool.js'
import { compactConversation, estimateTokens, DEFAULT_CONTEXT_WINDOW } from './services/compact/compact.js'
import type { PermissionContext, PermissionMode } from './utils/permissions/permissions.js'

export interface QueryEngineOptions {
  client: ApiClient
  tools: BuiltTool[]
  model: string
  maxOutputTokens: number
  maxTurns: number
  cwd: string
  extraDirs?: string[]
  customSystemPrompt?: string
  appendSystemPrompt?: string
  verbose?: boolean
  canUseTool: CanUseTool
  /**
   * Mutable permission context (same ref passed to createCanUseTool). The
   * engine can flip its mode at runtime (e.g. /bypass) and the change takes
   * effect on the next tool check.
   */
  permCtx?: PermissionContext
  autoCompact?: (messages: Message[]) => Promise<Message[] | null>
  /** Small model for summaries (compaction). */
  smallModel?: string
  /** Model catalog (from config file) for /model switching. */
  models?: ModelEntry[]
  /** Hooks registry (PreToolUse/PostToolUse/UserPromptSubmit/SessionStart/SessionEnd/Stop). */
  hooks?: HooksRegistry
  /** Sink for hook warnings. */
  hooksLog?: (msg: string) => void
  /**
   * Session id to resume. When provided, the engine loads its messages at
   * construction and continues appending to it. When omitted, a new session
   * is created (unless `disableSessionPersistence` is set).
   */
  sessionId?: string
  /** Disable session persistence entirely (e.g. for tests). */
  disableSessionPersistence?: boolean
  /** Start in plan mode (read-only tools + ExitPlanMode until approval). */
  startInPlanMode?: boolean
}

export interface SubmitMessageCallbacks {
  onStreamEvent?: (event: unknown) => void
  onTextDelta?: (text: string) => void
  onToolStart?: (toolName: string, input: unknown) => void
  onToolEnd?: (toolName: string, input: unknown, result: unknown, isError: boolean) => void
  onAssistantMessage?: (message: Message) => void
  onUserMessage?: (message: Message) => void
  onExit?: (reason: QueryExitReason, error?: string) => void
  /** Fired after each model call with that round's token usage. */
  onUsage?: (model: string, usage: Usage) => void
  /** Plan-mode: present a plan for user approval. Returns true to proceed. */
  onPlanPresented?: (plan: string) => Promise<boolean>
}

export class QueryEngine {
  private messages: Message[] = []
  private abortController = new AbortController()
  /** Queued user messages typed while the agent was running (non-blocking input). */
  private pendingQueue: Message[] = []
  private readonly readFileState: FileStateCache
  private readonly usageTracker = new UsageTracker()
  private readonly systemPrompt: MessageCreateParams['system']
  private readonly modelManager: ModelManager
  private readonly hooks: HooksRegistry
  private readonly hooksLog: (msg: string) => void
  /** Mutable permission context (shared with canUseTool). */
  private permCtx: PermissionContext | undefined
  /** Active session id (null when persistence is disabled). */
  private sessionId: string | null
  /** Plan-mode flag (read-only tools until ExitPlanMode is approved). */
  private planMode: boolean

  constructor(private opts: QueryEngineOptions) {
    this.readFileState = createFileStateCache()
    this.systemPrompt = fetchSystemPromptParts({
      tools: opts.tools,
      cwd: opts.cwd,
      extraDirs: opts.extraDirs,
      customSystemPrompt: opts.customSystemPrompt,
      appendSystemPrompt: opts.appendSystemPrompt,
      verbose: opts.verbose,
    })
    this.modelManager = new ModelManager({
      model: opts.model,
      smallModel: opts.smallModel ?? opts.model,
      models: opts.models ?? [],
      defaultMaxOutputTokens: opts.maxOutputTokens,
    })
    this.hooks = opts.hooks ?? emptyRegistry()
    this.hooksLog = opts.hooksLog ?? ((m: string) => console.warn(`[hooks] ${m}`))
    this.permCtx = opts.permCtx
    this.planMode = !!opts.startInPlanMode

    // Session persistence: resume an existing session or create a new one.
    if (opts.disableSessionPersistence) {
      this.sessionId = null
    } else if (opts.sessionId) {
      const { messages } = loadSession(opts.cwd, opts.sessionId)
      this.messages = messages
      this.sessionId = opts.sessionId
    } else {
      const meta = createSession(opts.cwd, opts.model)
      this.sessionId = meta.id
    }

    // SessionStart hook (fire-and-forget; observe-only).
    this.fireHooks('SessionStart', { cwd: opts.cwd }).catch(() => { /* observe-only */ })

    // Register the ExitPlanMode approval handler. On approve, exit plan mode so
    // the next turn uses the full tool set.
    this.registerPlanApprovalHandler()
  }

  /** Register the ExitPlanMode approval handler (wraps the UI callback). */
  private registerPlanApprovalHandler(): void {
    setPlanApprovalHandler(async (plan: string) => {
      const cb = (this as unknown as { __planCb?: (plan: string) => Promise<boolean> }).__planCb
      if (!cb) return false
      const approved = await cb(plan)
      if (approved) this.planMode = false
      return approved
    })
  }

  /** Set the plan-approval UI callback (called by the REPL). */
  setPlanApprovalCallback(cb: (plan: string) => Promise<boolean>): void {
    ;(this as unknown as { __planCb?: (plan: string) => Promise<boolean> }).__planCb = cb
  }

  /** Whether the engine is currently in plan mode. */
  isPlanMode(): boolean {
    return this.planMode
  }

  /** Enter plan mode (e.g. via /plan). */
  enterPlanMode(): void {
    this.planMode = true
  }

  /** Exit plan mode (e.g. on rejection after several tries). */
  exitPlanMode(): void {
    this.planMode = false
  }

  /** Current permission mode (null if no permCtx was provided). */
  getPermissionMode(): PermissionMode | null {
    return this.permCtx?.mode ?? null
  }

  /**
   * Switch the permission mode at runtime (e.g. /bypass → 'bypassPermissions').
   * Takes effect on the next tool check. No-op if no permCtx was provided.
   */
  setPermissionMode(mode: PermissionMode): void {
    if (this.permCtx) this.permCtx.mode = mode
  }

  /** Build the tool list for the current mode. */
  private toolsForCurrentMode(): BuiltTool[] {
    if (!this.planMode) {
      // Full mode: exclude ExitPlanMode (it's only for plan mode).
      return this.opts.tools.filter((t) => t.name !== 'ExitPlanMode')
    }
    // Plan mode: read-only tools + ExitPlanMode only. isReadOnly is
    // input-dependent (e.g. BashTool inspects the command); with no input we
    // fail-closed: treat a throwing isReadOnly as non-read-only.
    return this.opts.tools.filter((t) => {
      if (t.name === 'ExitPlanMode') return true
      try {
        return t.isReadOnly?.({}) ?? false
      } catch {
        return false
      }
    })
  }

  /** The active session id (null if persistence disabled). */
  getSessionId(): string | null {
    return this.sessionId
  }

  /**
   * Replace the in-memory message history with a resumed session's messages.
   * Used by /resume. Does not re-persist (subsequent turns append to the
   * resumed session id).
   */
  resumeSession(sessionId: string): { count: number } | null {
    const { messages } = loadSession(this.opts.cwd, sessionId)
    this.messages = messages
    this.sessionId = sessionId
    this.readFileState.clear()
    return { count: messages.length }
  }

  /**
   * Start a fresh conversation: clear in-memory history, file read state, and
   * begin a new persisted session (so the old one remains resumable via
   * /history). Pending queued input is dropped.
   */
  newConversation(): void {
    this.messages = []
    this.pendingQueue = []
    this.readFileState.clear()
    if (!this.opts.disableSessionPersistence) {
      const meta = createSession(this.opts.cwd, this.modelManager.getModel())
      this.sessionId = meta.id
    }
  }

  /** The hooks registry (for /hooks listing). */
  getHooks(): HooksRegistry {
    return this.hooks
  }

  /**
   * Run memory extraction over the current conversation (docs §06.4).
   * Returns a human-readable summary string for the UI.
   */
  async extractMemories(): Promise<string> {
    const { extractMemories } = await import('./services/extractMemories/index.js')
    const result = await extractMemories({
      client: this.opts.client,
      smallModel: this.modelManager.getSmallModel(),
      messages: this.messages,
      cwd: this.opts.cwd,
    })
    if (result.error) return `Memory extraction failed: ${result.error}`
    if (result.written === 0) return 'No new memories to save.'
    return `Wrote ${result.written} memor${result.written === 1 ? 'y' : 'ies'}: ${result.names.join(', ')}${result.skipped ? ` (skipped ${result.skipped} duplicate/invalid)` : ''}`
  }

  /**
   * Manually compact the conversation now (bypassing the auto-compact
   * threshold). Used by /compact. Returns a human-readable summary.
   */
  async compactNow(): Promise<string> {
    if (this.messages.length < 2) return 'Not enough conversation to compact.'
    const compacted = await compactConversation(this.messages, {
      client: this.opts.client,
      model: this.modelManager.getSmallModel(),
      contextWindow: DEFAULT_CONTEXT_WINDOW,
      hooks: this.hooks,
      hooksCwd: this.opts.cwd,
      hooksLog: this.hooksLog,
    }).catch(() => null)
    if (!compacted) return 'Compaction failed.'
    const before = this.messages.length
    this.messages = compacted
    return `Compacted ${before} messages → ${compacted.length}.`
  }

  /** Fire observe-only hooks for an event (no decision honored outside PreToolUse). */
  private async fireHooks(event: 'SessionStart' | 'SessionEnd' | 'Stop' | 'UserPromptSubmit', input: Record<string, unknown>): Promise<void> {
    await runHooks(this.hooks, event, { cwd: this.opts.cwd, ...input } as never, {
      cwd: this.opts.cwd,
      log: this.hooksLog,
    }).catch(() => { /* observe-only */ })
  }

  getUsageTracker(): UsageTracker {
    return this.usageTracker
  }

  /** Rough estimate of the current conversation's token count (chars/4). */
  getContextTokens(): number {
    return estimateTokens(this.messages)
  }

  getMessages(): Message[] {
    return this.messages
  }

  /** Model manager — supports runtime /model switching. */
  getModelManager(): ModelManager {
    return this.modelManager
  }

  /** Interrupt the current query. */
  interrupt(): void {
    this.abortController.abort()
  }

  /**
   * Queue a user message to be injected at the top of the agent's next turn
   * (non-blocking input). Does not interrupt the current turn — the message
   * waits until the in-flight model call / tool execution finishes, then is
   * appended before the next callModel.
   */
  enqueueUserMessage(text: string): void {
    this.pendingQueue.push({ role: 'user', content: text })
  }

  /** Reset the abort controller for a new turn. */
  private resetAbort(): void {
    if (this.abortController.signal.aborted) {
      this.abortController = new AbortController()
    }
  }

  /**
   * Submit a user prompt and run the agent loop to completion.
   */
  async submitMessage(
    prompt: string,
    callbacks: SubmitMessageCallbacks = {},
  ): Promise<QueryResult> {
    this.resetAbort()

    const prevLen = this.messages.length
    const userMessage: Message = { role: 'user', content: prompt }
    this.messages = [...this.messages, userMessage]
    callbacks.onUserMessage?.(userMessage)

    // UserPromptSubmit hook (observe-only).
    await this.fireHooks('UserPromptSubmit', { messages: this.messages })

    // Wire the plan-approval UI callback for this turn (if provided).
    if (callbacks.onPlanPresented) {
      this.setPlanApprovalCallback(callbacks.onPlanPresented)
    }

    const context: ToolUseContext = {
      abortController: this.abortController,
      readFileState: this.readFileState,
      cwd: this.opts.cwd,
      messages: this.messages,
    }

    // Read model + max tokens from the (mutable) model manager so /model
    // switches take effect on the next submitMessage. In plan mode, only
    // read-only tools + ExitPlanMode are exposed.
    const result = await query(this.messages, {
      client: this.opts.client,
      tools: this.toolsForCurrentMode(),
      systemPrompt: this.systemPrompt,
      model: this.modelManager.getModel(),
      maxOutputTokens: this.modelManager.getMaxOutputTokens(),
      maxTurns: this.opts.maxTurns,
      context,
      canUseTool: this.opts.canUseTool,
      autoCompact: this.opts.autoCompact,
      hooks: this.hooks,
      hooksCwd: this.opts.cwd,
      hooksLog: this.hooksLog,
      onStreamEvent: callbacks.onStreamEvent,
      onTextDelta: callbacks.onTextDelta,
      onToolStart: callbacks.onToolStart,
      onToolEnd: callbacks.onToolEnd,
      onUsage: (model, usage) => {
        this.usageTracker.add(model, usage)
        callbacks.onUsage?.(model, usage)
      },
      injectMessages: () => {
        if (this.pendingQueue.length === 0) return []
        const queued = this.pendingQueue.splice(0)
        for (const m of queued) callbacks.onUserMessage?.(m)
        return queued
      },
    })

    // Adopt the updated message history (includes assistant + tool messages).
    this.messages = result.messages

    // Persist the new messages from this turn (best-effort, never block).
    if (this.sessionId) {
      const newMessages = this.messages.slice(prevLen)
      try {
        appendMessages(this.opts.cwd, this.sessionId, newMessages)
      } catch {
        // Persistence failure is non-fatal.
      }
    }

    // Emit the final assistant message if present.
    const lastAssistant = [...this.messages].reverse().find((m) => m.role === 'assistant')
    if (lastAssistant) callbacks.onAssistantMessage?.(lastAssistant)

    // Stop hook (observe-only).
    await this.fireHooks('Stop', { reason: result.reason, messages: this.messages })

    callbacks.onExit?.(result.reason, result.error)
    return result
  }

  /** Fire SessionEnd hooks. Call on REPL exit / process shutdown. */
  async shutdown(): Promise<void> {
    await this.fireHooks('SessionEnd', { messages: this.messages })
  }
}