platform/src/servers/api/chat/websocket.ts

import type { ServerWebSocket } from 'bun';
import { randomUUID } from 'crypto';
import type { ClientMessage, ServerMessage, Message, TurnMessage, UserSession } from './types';
import { sessionManager } from './session-manager';
import { sendClaudeCodeStreaming } from '@@/channels/send-claude-code';
import { sendOpenCodeStreaming } from '@@/channels/send-opencode';
import { ensureGeneralChatSessionsCwd } from './claude-sessions';
import * as sidecar from '@@/sidecar-registry';
import { join } from 'path';
import { getOwnerHomeDir, getEmailAccountsDir } from '../../../servers/data-path';
import { getUserSettings, getEmailAccounts, getChatEventsSince } from 'officerdb';
import { mkdirSync } from 'node:fs';
import { logger } from './logger';

// Default model when no user preference is set
const DEFAULT_MODEL = 'claude-code';

// Harness selection: the `claude-code` provider runs through the Claude sidecar; every other provider
// (opencode/anthropic/openai/… — all `providerID/modelID` ids) runs through the OpenCode server.
const isClaudeModel = (model: string): boolean => model.startsWith('claude-code');

async function getUserDefaultModel(userId: number): Promise<string | null> {
  try {
    const settings = await getUserSettings(userId);
    const chat = settings?.chat as Record<string, unknown> | undefined;
    return (chat?.defaultModel as string) || null;
  } catch (err) {
    logger.error('Failed to read user settings for default model', { userId, error: String(err) });
  }
  return null;
}

type WSData = {
  userId: number;
  email: string;
  username: string;
  provider: string;
};

const IDLE_TIMEOUT_MS = 60 * 60 * 1000; // 1 hour

const resolveCwd = (email: string, cwd?: string) => {
  const root = getOwnerHomeDir(email);
  if (!cwd || cwd === '~') return root;
  if (cwd.startsWith('~/')) return join(root, cwd.slice(2));
  // The server owner is the only account — absolute paths are theirs to use.
  if (cwd.startsWith('/')) return cwd;
  return join(root, cwd);
};

export const resolveBaseCwd = (email: string, cwd?: string) => resolveCwd(email, cwd);

// The email chat runs from the selected account's storage dir:
//   DATA_PATH/<owner>/email_accounts/<accountEmail>
// `accountEmail` will come from the account selector (msg.contextId) later; for now default to the
// owner's first enabled account. Falls back to the email_accounts root if there are no accounts.
async function resolveEmailCwd(userId: number, ownerEmail: string, accountEmail?: string): Promise<string> {
  let account = accountEmail?.trim();
  if (!account) {
    try {
      const accounts = await getEmailAccounts(userId);
      account = (accounts.find((a) => a.enabled) ?? accounts[0])?.email;
    } catch (err) {
      logger.error('Failed to resolve email account for chat cwd', { userId, error: String(err) });
    }
  }
  const dir = account ? join(getEmailAccountsDir(ownerEmail), account) : getEmailAccountsDir(ownerEmail);
  mkdirSync(dir, { recursive: true });
  return dir;
}

// The working directory a chat turn runs in, by context: email → the account dir; /chat → a chosen
// pwd or the default general_chat_sessions dir; everything else (browser/project/dashboard) → the given cwd.
async function resolveChatCwd(
  msg: { context?: string; contextId?: string; cwd?: string },
  email: string,
  userId: number,
): Promise<string> {
  if (msg.context === 'email') return resolveEmailCwd(userId, email, msg.contextId);
  if (msg.context === 'chat') return msg.cwd?.trim() ? resolveCwd(email, msg.cwd) : ensureGeneralChatSessionsCwd(email);
  return resolveCwd(email, msg.cwd);
}

const wsToSessionMap = new WeakMap<any, string>();

// Per-connection heartbeat. Bun closes a WS idle for `idleTimeout` (60s), and its timer only resets
// on frames *received* from the client — but during a chat turn the client only receives. So we ping
// each connection every 25s; the client auto-pongs at the protocol level, which resets Bun's timer
// (and keeps reverse proxies happy). Genuinely dead sockets still time out (no pong).
const pingTimers = new WeakMap<ServerWebSocket<WSData>, ReturnType<typeof setInterval>>();
const PING_INTERVAL_MS = 25_000;

function sendToClient(ws: ServerWebSocket<WSData> | null, msg: ServerMessage, seq?: number): void {
  if (ws?.readyState === 1) {
    ws.send(JSON.stringify(seq === undefined ? msg : { ...msg, seq }));
  }
}

// Nothing in officer writes to chat_session_events any more. Both harnesses commit their own turn
// output in the sidecar that produced it (`sidecar/claude/session-log.ts`), which is the whole point:
// the durable record does not travel over the socket between the two processes, so officer can restart
// mid-turn without losing it. Officer reads the table on `resume` (getChatEventsSince) and relays what
// the sidecars send. The old `emitToSession` used to live here.

export async function open(ws: ServerWebSocket<WSData>): Promise<void> {
  const timer = setInterval(() => {
    try {
      ws.ping(); // client auto-pongs → resets Bun's idleTimeout
    } catch {
      /* socket already gone */
    }
  }, PING_INTERVAL_MS);
  pingTimers.set(ws, timer);
}

export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void {
  const data = typeof raw === 'string' ? raw : raw.toString();

  (async () => {
    try {
      const clientMsg = JSON.parse(data) as ClientMessage;

      if (clientMsg.type === 'chat') {
        await handleChat(ws, clientMsg);
      } else if (clientMsg.type === 'resume') {
        await handleResume(ws, clientMsg);
      } else if (clientMsg.type === 'stop') {
        await handleStop(ws);
      } else if (clientMsg.type === 'disconnect') {
        await handleDisconnect(ws);
      } else if (clientMsg.type === 'resume-cursor') {
        await handleResumeCursor(ws, clientMsg);
      }
    } catch (err) {
      logger.error('Error handling WebSocket message', { email: ws.data.email, error: String(err) });
      sendToClient(ws, { type: 'error', message: 'Failed to process message' });
    }
  })();
}

export function close(ws: ServerWebSocket<WSData>): void {
  const timer = pingTimers.get(ws);
  if (timer) {
    clearInterval(timer);
    pingTimers.delete(ws);
  }

  const sessionId = wsToSessionMap.get(ws);
  if (sessionId) {
    sessionManager.detachWs(sessionId);
    sessionManager.setIdleTimeout(sessionId, IDLE_TIMEOUT_MS);
  }
}

// ── Claude Code: relay, don't rebuild ──

// The agent sidecar owns Claude's turn output end to end — it translates the parser stream, commits each
// message to chat_session_events and hands us a finished message plus its cursor id. Officer relays it.
// That is what makes a restart survivable: the durable record no longer travels over the socket between
// the two processes, so if this one is down the output is already written and the client replays it.
//
// Officer keeps only the in-memory transcript, which exists to answer a `resume` with sync:messages —
// Claude's own transcript is the real record — so it is folded from the same messages, not rebuilt.
function foldIntoSession(session: UserSession, msg: TurnMessage, model: string): void {
  switch (msg.type) {
    case 'assistant:delta':
      session.streamBuffer += msg.text;
      break;

    case 'assistant:text':
      session.messages.push({ id: randomUUID(), timestamp: Date.now(), role: 'assistant', text: msg.text, model });
      session.meta.messageCount += 1;
      session.streamBuffer = '';
      break;

    case 'tool:start':
      session.messages.push({
        id: randomUUID(),
        timestamp: Date.now(),
        role: 'tool',
        toolCallId: msg.toolCallId,
        toolName: msg.toolName,
        toolInput: msg.toolInput,
      });
      session.meta.messageCount += 1;
      break;

    case 'tool:result':
      for (let i = session.messages.length - 1; i >= 0; i--) {
        const m = session.messages[i]!;
        if (m.role === 'tool' && m.toolCallId === msg.toolCallId) {
          m.output = msg.output;
          m.isError = msg.isError;
          break;
        }
      }
      break;

    case 'result': {
      session.isGenerating = false;
      session.meta.cost.inputTokens += msg.cost.inputTokens;
      session.meta.cost.outputTokens += msg.cost.outputTokens;
      session.meta.cost.totalUSD += msg.cost.totalUSD;
      session.meta.updatedAt = Date.now();
      // The turn's cost belongs to the assistant message it paid for (as it did when officer built these).
      const last = session.messages[session.messages.length - 1];
      if (last?.role === 'assistant' && !last.cost) last.cost = msg.cost;
      break;
    }

    case 'error':
    case 'stopped':
      session.isGenerating = false;
      break;
  }
}

// Shared by both harnesses now: whichever sidecar ran the turn has already translated and committed it.
function createMessageHandler(sessionId: string, model: string) {
  return (msg: TurnMessage, seq?: number): void => {
    const session = sessionManager.getSession(sessionId);
    if (!session) return;
    foldIntoSession(session, msg, model);
    sendToClient(session.ws as ServerWebSocket<WSData> | null, msg, seq);
  };
}

async function handleChat(
  ws: ServerWebSocket<WSData>,
  msg: {
    prompt: string;
    displayText?: string;
    sessionId?: string;
    model?: string;
    cwd?: string;
    cwdRoot?: string;
    groupSlug?: string;
    attachmentIds?: string[];
    thinking?: string;
    context?: string;
    contextId?: string;
    resumeSummary?: string;
    resumeSessionId?: string;
  },
): Promise<void> {
  const { userId } = ws.data;
  const sessionId = msg.sessionId || randomUUID();

  // Prepend resume summary to the prompt if present
  const prompt = msg.resumeSummary
    ? `Here is a summary of a previous conversation to continue from:\n\n${msg.resumeSummary}\n\n---\n\nUser's new message: ${msg.prompt}`
    : msg.prompt;

  // Use provided model, or fall back to user default, or the system default.
  const model = msg.model || (await getUserDefaultModel(userId)) || DEFAULT_MODEL;

  logger.info('Model selected for chat', { sessionId, model, clientModel: msg.model || null });

  // Route by harness: claude-code → Claude sidecar; anything else → OpenCode server.
  return isClaudeModel(model)
    ? handleClaudeCodeChat(ws, sessionId, model, msg, prompt)
    : handleOpenCodeChat(ws, sessionId, model, msg, prompt);
}

async function handleClaudeCodeChat(
  ws: ServerWebSocket<WSData>,
  sessionId: string,
  model: string,
  msg: {
    prompt: string;
    displayText?: string;
    groupSlug?: string;
    context?: string;
    contextId?: string;
    cwd?: string;
    cwdRoot?: string;
    resumeSessionId?: string;
  },
  effectivePrompt: string,
): Promise<void> {
  const { email, username, userId } = ws.data;

  const cwd = await resolveChatCwd(msg, email, userId);

  const groupSlug = msg.groupSlug || null;

  const session = sessionManager.getOrCreate(sessionId, email, cwd, model, groupSlug, msg.context, msg.contextId);
  session.userId = userId;
  sessionManager.attachWs(sessionId, ws);
  wsToSessionMap.set(ws as any, sessionId);

  sendToClient(ws, {
    type: 'session:init',
    sessionId,
    model,
    cwd,
    context: session.meta.context,
    contextId: session.meta.contextId,
  });

  // Add user message to session
  const userMsg: Message = {
    id: randomUUID(),
    timestamp: Date.now(),
    role: 'user',
    text: msg.prompt,
  };
  session.messages.push(userMsg);
  session.meta.messageCount += 1;
  session.meta.updatedAt = Date.now();

  if (!session.meta.title) {
    session.meta.title = (msg.displayText ?? msg.prompt).slice(0, 100);
  }

  session.isGenerating = true;

  const onMessage = createMessageHandler(sessionId, model);

  try {
    if (!session._claudeKill) {
      // First turn of this session: open the persistent session + a SESSION-scoped event subscription
      // (survives turn-end so background task:notifications keep flowing). handle.kill tears both down.
      const handle = await sendClaudeCodeStreaming({
        userId,
        email,
        username,
        prompt: effectivePrompt,
        sessionKey: sessionId,
        cwd,
        model,
        resumeSessionId: msg.resumeSessionId,
        onMessage,
      });
      session.piProcess = sessionId as any;
      session._claudeKill = handle.kill;
    } else {
      // Session already live: push this turn onto the existing persistent session (no new subscription).
      await sidecar.spawnClaudeStreaming({
        userId,
        email,
        username,
        prompt: effectivePrompt,
        sessionKey: sessionId,
        cwd,
        model,
        resumeSessionId: msg.resumeSessionId,
      });
    }
  } catch (err) {
    logger.error('Failed to start Claude Code streaming', { sessionId, error: String(err) });
    sendToClient(ws, { type: 'error', message: 'Failed to start Claude Code' });
    session.isGenerating = false;
  }
}

async function handleOpenCodeChat(
  ws: ServerWebSocket<WSData>,
  sessionId: string,
  model: string,
  msg: {
    prompt: string;
    displayText?: string;
    groupSlug?: string;
    context?: string;
    contextId?: string;
    cwd?: string;
    cwdRoot?: string;
    resumeSessionId?: string;
  },
  effectivePrompt: string,
): Promise<void> {
  const { email, username, userId } = ws.data;

  const cwd = await resolveChatCwd(msg, email, userId);

  const groupSlug = msg.groupSlug || null;

  const session = sessionManager.getOrCreate(sessionId, email, cwd, model, groupSlug, msg.context, msg.contextId);
  session.userId = userId;
  sessionManager.attachWs(sessionId, ws);
  wsToSessionMap.set(ws as any, sessionId);

  sendToClient(ws, {
    type: 'session:init',
    sessionId,
    model,
    cwd,
    context: session.meta.context,
    contextId: session.meta.contextId,
  });

  const userMsg: Message = {
    id: randomUUID(),
    timestamp: Date.now(),
    role: 'user',
    text: msg.prompt,
  };
  session.messages.push(userMsg);
  session.meta.messageCount += 1;
  session.meta.updatedAt = Date.now();

  if (!session.meta.title) {
    session.meta.title = (msg.displayText ?? msg.prompt).slice(0, 100);
  }

  session.isGenerating = true;

  const onMessage = createMessageHandler(sessionId, model);

  try {
    const handle = await sendOpenCodeStreaming({
      userId,
      email,
      username,
      prompt: effectivePrompt,
      sessionKey: sessionId,
      cwd,
      model,
      resumeSessionId: msg.resumeSessionId,
      onMessage,
    });

    // Store the abort handle so handleStop can end the turn (OpenCode is aborted via this handle).
    session.piProcess = sessionId as any;
    session._claudeKill = handle.kill;
  } catch (err) {
    logger.error('Failed to start OpenCode streaming', { sessionId, error: String(err) });
    sendToClient(ws, { type: 'error', message: 'Failed to start OpenCode' });
    session.isGenerating = false;
  }
}

async function handleResume(
  ws: ServerWebSocket<WSData>,
  msg: { sessionId: string; cwd?: string; cwdRoot?: string },
): Promise<void> {
  const { sessionId } = msg;

  try {
    const session = sessionManager.getSession(sessionId);

    if (!session) {
      // Sessions live in memory for the connection's lifetime; there's no disk store to reload from.
      sendToClient(ws, { type: 'error', message: 'Session not found', errorCode: 'SESSION_NOT_FOUND' });
      return;
    }

    sessionManager.attachWs(sessionId, ws);
    wsToSessionMap.set(ws as any, sessionId);

    sendToClient(ws, {
      type: 'session:init',
      sessionId,
      model: session.model,
      cwd: session.cwd,
      context: session.meta.context,
      contextId: session.meta.contextId,
    });

    // Claude resumes lazily: the next chat prompt re-attaches via `--resume <sessionKey>`, so there's
    // no long-lived process to spawn here — just replay the stored transcript to the client.
    sendToClient(ws, {
      type: 'sync:messages',
      sessionId,
      messages: session.messages,
      isGenerating: session.isGenerating,
      streamingText: session.streamBuffer,
    });

    logger.info('Session resumed successfully', { sessionId, messageCount: session.messages.length });
  } catch (err) {
    logger.error('Unexpected error in handleResume', { sessionId, error: String(err) });
    sendToClient(ws, { type: 'error', message: 'Failed to resume session' });
  }
}

async function handleStop(ws: ServerWebSocket<WSData>): Promise<void> {
  const sessionId = wsToSessionMap.get(ws);

  if (sessionId) {
    const session = sessionManager.getSession(sessionId);

    if (session?.piProcess) {
      try {
        if (isClaudeModel(session.model)) {
          // Interrupt the current turn but KEEP the persistent session alive (background tasks + the
          // warm worker survive). Full teardown is 'disconnect' → deleteSession → _claudeKill.
          void sidecar.interruptClaude(sessionId);
          logger.info('Interrupted Claude Code turn via sidecar (session stays warm)', { sessionId });
        } else {
          session._claudeKill?.(); // OpenCode: abort the turn via the stored handle
          logger.info('Aborted OpenCode turn', { sessionId });
        }
        session.isGenerating = false;
      } catch (err) {
        logger.error('Failed to stop process', { sessionId, error: String(err) });
      }
    }
  }

  sendToClient(ws, { type: 'stopped' });
}

// Tear down the whole session (not just the current turn): deleteSession fires _claudeKill (kills any
// in-flight Claude/OpenCode turn) + _sidecarUnsub, clears the idle timer, and drops the session from the
// manager's maps. The WS stays open so the client can immediately start a fresh session.
async function handleDisconnect(ws: ServerWebSocket<WSData>): Promise<void> {
  const sessionId = wsToSessionMap.get(ws);
  if (sessionId) {
    try {
      sessionManager.deleteSession(sessionId);
      wsToSessionMap.delete(ws as any);
      logger.info('Disconnected chat session', { sessionId });
    } catch (err) {
      logger.error('Failed to disconnect session', { sessionId, error: String(err) });
    }
  }
  sendToClient(ws, { type: 'disconnected' });
}

// Reconnect: re-bind this socket to the (possibly still-live) session and replay every durable event
// queued since the client's cursor — so a brief disconnect never loses turn output or a background
// task:notification. attachWs cancels the pending idle-GC. If the in-memory session was already
// GC'd, we still replay history from Postgres (new turns will respawn the session).
async function handleResumeCursor(
  ws: ServerWebSocket<WSData>,
  msg: { sessionId: string; cursor: number },
): Promise<void> {
  const { sessionId, cursor } = msg;
  if (sessionManager.getSession(sessionId)) {
    sessionManager.attachWs(sessionId, ws);
    wsToSessionMap.set(ws as any, sessionId);
  }
  try {
    const events = await getChatEventsSince(sessionId, cursor ?? 0);
    for (const { id, event } of events) {
      sendToClient(ws, event as ServerMessage, id);
    }
  } catch (err) {
    logger.error('Failed to replay chat events on resume-cursor', { sessionId, error: String(err) });
  }
}

export const chatWebsocket = {
  open,
  message,
  close,
  drain() {},
};