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
}

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 AgentTerminalInfo {
  ptySessionId: string | null
  alive: boolean
  bufferSize: number
  connectedClients: number
}

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

export interface AgentSessionState {
  agent: string
  sessionId: string | null
  model: string | null
  cwd: string | null
  transcriptPath: string | null
  permissionMode: string | null
  status: AgentStatus
  currentTool: ActiveTool | null
  lastActivity: number
  lastStopResponse: string | null
  lastError: LastError | null
  pendingApprovals: PendingApproval[]
  terminal: AgentTerminalInfo
  notifications: SessionNotification[]
  hookHistory: HookHistoryEntry[]
  lastHookEvent: string | null
  lastHookDetail: 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'
  agents: Record<string, AgentSessionState>
}

export interface SessionStatePatch {
  type: 'session-state-patch'
  agent: string
  patch: Partial<AgentSessionState>
  event: string
  timestamp: number
}

// ── 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 createDefaultState(agent: string): AgentSessionState {
  return {
    agent,
    sessionId: null,
    model: null,
    cwd: null,
    transcriptPath: null,
    permissionMode: null,
    status: 'idle',
    currentTool: null,
    lastActivity: Date.now(),
    lastStopResponse: null,
    lastError: null,
    pendingApprovals: [],
    terminal: {
      ptySessionId: null,
      alive: false,
      bufferSize: 0,
      connectedClients: 0,
    },
    notifications: [],
    hookHistory: [],
    lastHookEvent: null,
    lastHookDetail: null,
  }
}

class SessionStateManager {
  private agents = new Map<string, AgentSessionState>()

  getOrCreateAgent(agent: string): AgentSessionState {
    let state = this.agents.get(agent)
    if (!state) {
      state = createDefaultState(agent)
      this.agents.set(agent, state)
    }
    return state
  }

  /** Process a raw hook event and return the patch to broadcast */
  processHookEvent(payload: HookPayload): SessionStatePatch {
    const agentName = payload.agent_name || 'claude'
    const state = this.getOrCreateAgent(agentName)
    const now = Date.now()

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

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

    // Only update status/tool if deriveStatus returned a result
    if (derived) {
      const { status, tool } = 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 from PostToolUseFailure
      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 session state
      if (payload.hook_event_name === 'SessionEnd') {
        patch.currentTool = null
        patch.pendingApprovals = []
      }
    }

    // Update session identity fields
    if (payload.session_id) patch.sessionId = 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

    // Last stop response
    if (payload.hook_event_name === 'Stop' && payload.assistant_response) {
      patch.lastStopResponse = payload.assistant_response as string
    }

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

    // Build notification
    const notification = buildNotification(payload)
    if (payload.hook_event_name === 'SessionStart') {
      // Reset notifications on new session, only keep the SessionStart notification
      patch.notifications = [notification]
    } else {
      patch.notifications = [...state.notifications, notification].slice(-MAX_NOTIFICATIONS)
    }

    // Build 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 (payload.hook_event_name === 'SessionStart') {
      patch.hookHistory = [historyEntry]
    } else {
      patch.hookHistory = [...state.hookHistory, historyEntry].slice(-MAX_HOOK_HISTORY)
    }

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

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

  /** Add a pending approval */
  addApproval(agentName: string, approval: PendingApproval): Partial<AgentSessionState> {
    const state = this.getOrCreateAgent(agentName)
    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): { agent: string, patch: Partial<AgentSessionState> } | null {
    for (const [agentName, state] of this.agents) {
      const idx = state.pendingApprovals.findIndex(a => a.requestId === requestId)
      if (idx !== -1) {
        state.pendingApprovals = state.pendingApprovals.filter(a => a.requestId !== requestId)
        // If no more pending approvals, go back to thinking
        if (state.pendingApprovals.length === 0 && state.status === 'permissionRequest') {
          state.status = 'thinking'
        }
        state.lastActivity = Date.now()
        return {
          agent: agentName,
          patch: { pendingApprovals: state.pendingApprovals, status: state.status, lastActivity: state.lastActivity }
        }
      }
    }
    return null
  }

  /** Update terminal info for an agent */
  updateTerminalInfo(agentName: string, info: Partial<AgentTerminalInfo>): Partial<AgentSessionState> {
    const state = this.getOrCreateAgent(agentName)
    state.terminal = { ...state.terminal, ...info }
    return { terminal: state.terminal }
  }

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

  /** Get single agent state */
  getAgentState(agent: string): AgentSessionState | null {
    return this.agents.get(agent) || null
  }

  /** Find agent name by session_id */
  findAgentBySessionId(sessionId: string): string | null {
    for (const [name, state] of this.agents) {
      if (state.sessionId === sessionId) return name
    }
    return null
  }

  /** Find all agents matching a transcript_path */
  findAgentsByTranscript(transcriptPath: string): string[] {
    const matches: string[] = []
    for (const [name, state] of this.agents) {
      if (state.transcriptPath === transcriptPath) matches.push(name)
    }
    return matches
  }

  /** Clean up expired notifications (call periodically) */
  cleanExpiredNotifications(): void {
    const now = Date.now()
    for (const state of this.agents.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)