agent-ui/server/services/session-state.ts

// ── Centralized Session State Manager ──
// Source of truth for all Claude Code agent session states.
// Lives in Terminal Server (4103) which persists across API server restarts.

// ── Types ──

export type AgentStatus =
  | 'idle'
  | 'thinking'
  | 'reading'
  | 'writing'
  | 'toolUse'
  | 'permissionRequest'
  | 'interrupted'
  | 'error'
  | 'sessionStart'
  | 'sessionEnd'

export interface ActiveTool {
  name: string
  input: unknown
  startedAt: number
}

export interface PendingApproval {
  requestId: string
  type: 'permission' | 'plan'
  toolName?: string
  toolInput?: unknown
  cwd?: string
  lastAssistantText?: string
  timestamp: number
}

export interface HookHistoryEntry {
  event: string
  timestamp: number
  detail?: string
  ptySessionId?: string
}

export interface SessionNotification {
  id: string
  event: string
  title: string
  detail: string
  type: 'info' | 'success' | 'warning' | 'error'
  timestamp: number
  persistent: boolean
  expiresAt: number | null
}

export interface LastError {
  tool: string
  message: string
  interrupted: boolean
  timestamp: number
}

// ── Session state (keyed by ephemeralSessionId / ptySessionId) ──

export interface PtySessionState {
  ptySessionId: string
  agent: string                    // tag/label, NOT primary key
  transcriptSessionId: string | null
  status: AgentStatus
  currentTool: ActiveTool | null
  lastActivity: number
  lastError: LastError | null
  pendingApprovals: PendingApproval[]
  notifications: SessionNotification[]
  hookHistory: HookHistoryEntry[]
  lastHookEvent: string | null
  lastHookDetail: string | null
  sessionActive: boolean
  agentResponding: boolean
  subagentActive: boolean
  compacting: boolean
  // Identity fields from Claude Code
  model: string | null
  cwd: string | null
  transcriptPath: string | null
  permissionMode: string | null
}

export interface HookPayload {
  hook_event_name?: string
  session_id?: string
  tool_name?: string
  tool_input?: unknown
  tool_response?: unknown
  prompt?: string
  cwd?: string
  model?: string
  source?: string
  transcript_path?: string
  message?: string
  notification_type?: string
  stop_hook_active?: boolean
  tool_use_id?: string
  assistant_response?: string
  agent_name?: string
  is_interrupt?: boolean
  error?: string
  reason?: string
  [key: string]: unknown
}

// WS message types
export interface SessionStateSnapshot {
  type: 'session-state-snapshot'
  ptySessions: Record<string, PtySessionState>
}

export interface PtyStatePatch {
  type: 'pty-state-patch'
  ptySessionId: string
  agent: string
  patch: Partial<PtySessionState>
  event: string
  timestamp: number
}

export interface PtyStateSnapshot {
  type: 'pty-state-snapshot'
  ptySessions: Record<string, PtySessionState>
}

// ── Notification builders ──

const MAX_NOTIFICATIONS = 30
const MAX_HOOK_HISTORY = 500
const TTL_INFO = 3500
const TTL_WARNING = 5000

const NOTIFICATION_MAP: Record<string, { title: string, type: SessionNotification['type'], persistent?: boolean }> = {
  SessionStart: { title: 'Session started', type: 'info' },
  SessionEnd: { title: 'Session ended', type: 'info' },
  UserPromptSubmit: { title: 'Processing prompt', type: 'info' },
  PreToolUse: { title: 'Using tool', type: 'info' },
  PostToolUse: { title: 'Tool complete', type: 'success' },
  PostToolUseFailure: { title: 'Tool failed', type: 'error' },
  PermissionRequest: { title: 'Permission required', type: 'warning', persistent: true },
  Notification: { title: 'Notification', type: 'info' },
  Stop: { title: 'Session stopped', type: 'info' },
}

function buildNotification(payload: HookPayload): SessionNotification {
  const event = payload.hook_event_name || 'unknown'
  const config = NOTIFICATION_MAP[event] || { title: event, type: 'info' as const }
  const now = Date.now()

  let detail = ''
  if (payload.tool_name) detail = payload.tool_name
  if (event === 'UserPromptSubmit' && payload.prompt) {
    detail = payload.prompt.length > 80 ? payload.prompt.slice(0, 80) + '...' : payload.prompt
  }
  if (event === 'Notification' && payload.message) {
    detail = payload.message.length > 120 ? payload.message.slice(0, 120) + '...' : payload.message
  }
  if (event === 'Stop' && payload.assistant_response) {
    detail = payload.assistant_response.length > 120 ? payload.assistant_response.slice(0, 120) + '...' : payload.assistant_response
  }

  const persistent = config.persistent ?? false
  const ttl = config.type === 'warning' ? TTL_WARNING : TTL_INFO

  return {
    id: `n_${now}_${Math.random().toString(36).slice(2, 8)}`,
    event,
    title: config.title,
    detail,
    type: config.type,
    timestamp: now,
    persistent,
    expiresAt: persistent ? null : now + ttl,
  }
}

// ── Status derivation (moved from claude-hook.ts) ──

export function deriveStatus(payload: HookPayload): { status: AgentStatus, tool?: string } | null {
  const event = payload.hook_event_name
  const toolName = payload.tool_name

  switch (event) {
    case 'SessionStart':
      return { status: 'sessionStart' }
    case 'SessionEnd':
      return { status: 'sessionEnd' }
    case 'UserPromptSubmit':
      return { status: 'thinking' }
    case 'PreToolUse':
      if (toolName && /^(Read|Glob|Grep)$/.test(toolName))
        return { status: 'reading', tool: toolName }
      if (toolName && /^(Edit|Write)$/.test(toolName))
        return { status: 'writing', tool: toolName }
      return { status: 'toolUse', tool: toolName }
    case 'PostToolUse':
      return { status: 'thinking', tool: toolName }
    case 'PostToolUseFailure':
      if (payload.is_interrupt) return { status: 'interrupted', tool: toolName }
      return { status: 'error', tool: toolName }
    case 'PermissionRequest':
      return { status: 'permissionRequest', tool: toolName }
    case 'Notification':
      return null // Side-effect only, does not change agent status
    case 'Stop':
      return { status: 'idle' }
    default:
      return { status: 'thinking' }
  }
}

// ── SessionStateManager ──

function createDefaultPtyState(ptySessionId: string, agent: string): PtySessionState {
  return {
    ptySessionId,
    agent,
    transcriptSessionId: null,
    status: 'idle',
    currentTool: null,
    lastActivity: Date.now(),
    lastError: null,
    pendingApprovals: [],
    notifications: [],
    hookHistory: [],
    lastHookEvent: null,
    lastHookDetail: null,
    sessionActive: false,
    agentResponding: false,
    subagentActive: false,
    compacting: false,
    model: null,
    cwd: null,
    transcriptPath: null,
    permissionMode: null,
  }
}

class SessionStateManager {
  private ptySessions = new Map<string, PtySessionState>()

  /** Get or create a PTY-scoped state entry */
  getOrCreatePty(ptySessionId: string, agent: string): PtySessionState {
    let state = this.ptySessions.get(ptySessionId)
    if (!state) {
      state = createDefaultPtyState(ptySessionId, agent)
      this.ptySessions.set(ptySessionId, state)
    }
    return state
  }

  /** Get full snapshot for new clients */
  getSnapshot(): Record<string, PtySessionState> {
    const result: Record<string, PtySessionState> = {}
    for (const [id, state] of this.ptySessions) {
      result[id] = { ...state }
    }
    return result
  }

  /** Get single PTY state */
  getPtyState(ptySessionId: string): PtySessionState | null {
    return this.ptySessions.get(ptySessionId) || null
  }

  /** Add a pending approval to a PTY session */
  addApproval(ptySessionId: string, agent: string, approval: PendingApproval): Partial<PtySessionState> {
    const state = this.getOrCreatePty(ptySessionId, agent)
    state.pendingApprovals = [...state.pendingApprovals, approval]
    state.status = 'permissionRequest'
    state.lastActivity = Date.now()
    return { pendingApprovals: state.pendingApprovals, status: state.status, lastActivity: state.lastActivity }
  }

  /** Remove a pending approval by requestId */
  resolveApproval(requestId: string): { ptySessionId: string, agent: string, patch: Partial<PtySessionState> } | null {
    for (const [ptyId, state] of this.ptySessions) {
      const idx = state.pendingApprovals.findIndex(a => a.requestId === requestId)
      if (idx !== -1) {
        state.pendingApprovals = state.pendingApprovals.filter(a => a.requestId !== requestId)
        if (state.pendingApprovals.length === 0 && state.status === 'permissionRequest') {
          state.status = 'thinking'
        }
        state.lastActivity = Date.now()
        return {
          ptySessionId: ptyId,
          agent: state.agent,
          patch: { pendingApprovals: state.pendingApprovals, status: state.status, lastActivity: state.lastActivity }
        }
      }
    }
    return null
  }

  /** Process a raw hook event and return the PTY patch to broadcast. */
  processHookEvent(payload: HookPayload): PtyStatePatch | null {
    const agentName = payload.agent_name || 'claude'
    const ptyId = payload.pty_session_id as string | undefined
    const now = Date.now()

    // Without a PTY ID we can't write state — log and skip
    if (!ptyId) {
      console.warn(`[SessionState] Hook event without ptySessionId, skipping: ${payload.hook_event_name} (agent=${agentName})`)
      return null
    }

    const state = this.getOrCreatePty(ptyId, agentName)

    // Derive status (null means side-effect only, e.g. Notification)
    const derived = deriveStatus(payload)

    // Build patch
    const patch: Partial<PtySessionState> = {
      lastActivity: now,
      lastHookEvent: payload.hook_event_name || null,
      lastHookDetail: payload.tool_name || payload.message || null,
    }

    if (derived) {
      const { status } = derived
      patch.status = status

      // Current tool tracking
      if (status === 'toolUse' || status === 'reading' || status === 'writing' || status === 'permissionRequest') {
        patch.currentTool = {
          name: payload.tool_name || 'unknown',
          input: payload.tool_input,
          startedAt: now,
        }
      } else if (status === 'thinking' || status === 'idle' || status === 'sessionEnd') {
        patch.currentTool = null
      }

      // Track errors
      if (status === 'error' || status === 'interrupted') {
        patch.lastError = {
          tool: payload.tool_name || 'unknown',
          message: (payload.error as string) || '',
          interrupted: status === 'interrupted',
          timestamp: now,
        }
      }

      // Clear lastError on successful flow resumption
      if (status === 'thinking' || status === 'reading' || status === 'writing' || status === 'toolUse') {
        patch.lastError = null
      }

      // SessionEnd: clean up
      if (payload.hook_event_name === 'SessionEnd') {
        patch.currentTool = null
        patch.pendingApprovals = []
      }
    }

    // Session identity fields
    if (payload.session_id) patch.transcriptSessionId = payload.session_id
    if (payload.model) patch.model = payload.model as string
    if (payload.cwd) patch.cwd = payload.cwd as string
    if (payload.transcript_path) patch.transcriptPath = payload.transcript_path as string
    if (payload.permission_mode) patch.permissionMode = payload.permission_mode as string

    // Reset on SessionStart
    if (payload.hook_event_name === 'SessionStart') {
      patch.lastError = null
      patch.pendingApprovals = []
    }

    // Continuous state flags
    const evt = payload.hook_event_name
    switch (evt) {
      case 'SessionStart':
        patch.sessionActive = true
        patch.agentResponding = false
        patch.subagentActive = false
        patch.compacting = false
        break
      case 'SessionEnd':
        patch.sessionActive = false
        patch.agentResponding = false
        patch.subagentActive = false
        patch.compacting = false
        break
      case 'UserPromptSubmit':
        patch.agentResponding = true
        patch.compacting = false
        break
      case 'Stop':
        patch.agentResponding = false
        patch.subagentActive = false
        patch.compacting = false
        break
      case 'SubagentStart':
        patch.subagentActive = true
        break
      case 'SubagentStop':
        patch.subagentActive = false
        break
      case 'PreCompact':
        patch.compacting = true
        break
      case 'PostToolUse':
      case 'PostToolUseFailure':
        patch.compacting = false
        break
    }

    // Notification
    const notification = buildNotification(payload)
    if (payload.hook_event_name === 'SessionStart') {
      patch.notifications = [notification]
    } else {
      patch.notifications = [...state.notifications, notification].slice(-MAX_NOTIFICATIONS)
    }

    // Hook history entry
    const historyEntry: HookHistoryEntry = {
      event: payload.hook_event_name || 'unknown',
      timestamp: now,
    }
    const historyDetail = payload.tool_name || payload.message
    if (historyDetail) historyEntry.detail = historyDetail as string
    if (ptyId) historyEntry.ptySessionId = ptyId

    patch.hookHistory = [...state.hookHistory, historyEntry].slice(-MAX_HOOK_HISTORY)

    // Apply to state
    Object.assign(state, patch)

    return {
      type: 'pty-state-patch',
      ptySessionId: ptyId,
      agent: agentName,
      patch,
      event: payload.hook_event_name || 'unknown',
      timestamp: now,
    }
  }

  /** Find agent name by transcript session_id (searches PTY sessions) */
  findAgentBySessionId(sessionId: string): string | null {
    for (const state of this.ptySessions.values()) {
      if (state.transcriptSessionId === sessionId) return state.agent
    }
    return null
  }

  /** Find all agent names matching a transcript_path (searches PTY sessions) */
  findAgentsByTranscript(transcriptPath: string): string[] {
    const matches = new Set<string>()
    for (const state of this.ptySessions.values()) {
      if (state.transcriptPath === transcriptPath) matches.add(state.agent)
    }
    return Array.from(matches)
  }

  /** Clean up expired notifications */
  cleanExpiredNotifications(): void {
    const now = Date.now()
    for (const state of this.ptySessions.values()) {
      state.notifications = state.notifications.filter(
        n => n.persistent || !n.expiresAt || n.expiresAt > now
      )
    }
  }
}

// Singleton instance
export const sessionState = new SessionStateManager()

// Clean expired notifications every 5s
setInterval(() => {
  sessionState.cleanExpiredNotifications()
}, 5000)