Codewalkers/apps/server/trpc/routers/chat-session.ts

/**
 * Chat Session Router — send messages, get session, close session
 */

import { TRPCError } from '@trpc/server';
import { z } from 'zod';
import type { ProcedureBuilder } from '../trpc.js';
import {
  requireAgentManager,
  requireInitiativeRepository,
  requirePhaseRepository,
  requirePageRepository,
  requireTaskRepository,
  requireChatSessionRepository,
} from './_helpers.js';
import { gatherInitiativeContext } from './architect.js';
import { buildChatPrompt, type ChatHistoryEntry } from '../../agent/prompts/index.js';

const MAX_HISTORY_MESSAGES = 30;

export function chatSessionProcedures(publicProcedure: ProcedureBuilder) {
  return {
    sendChatMessage: publicProcedure
      .input(z.object({
        targetType: z.enum(['phase', 'task']),
        targetId: z.string().min(1),
        initiativeId: z.string().min(1),
        message: z.string().min(1),
        retry: z.boolean().optional(),
        provider: z.string().optional(),
      }))
      .mutation(async ({ ctx, input }) => {
        const agentManager = requireAgentManager(ctx);
        const chatRepo = requireChatSessionRepository(ctx);
        const initiativeRepo = requireInitiativeRepository(ctx);
        const taskRepo = requireTaskRepository(ctx);

        const initiative = await initiativeRepo.findById(input.initiativeId);
        if (!initiative) {
          throw new TRPCError({ code: 'NOT_FOUND', message: `Initiative '${input.initiativeId}' not found` });
        }

        // Find or create active session
        let session = await chatRepo.findActiveSession(input.targetType, input.targetId);
        if (!session) {
          session = await chatRepo.createSession({
            targetType: input.targetType,
            targetId: input.targetId,
            initiativeId: input.initiativeId,
          });
        }

        // Store user message (skip on retry — message already exists)
        if (!input.retry) {
          await chatRepo.createMessage({
            chatSessionId: session.id,
            role: 'user',
            content: input.message,
          });

          ctx.eventBus.emit({
            type: 'chat:message_created' as const,
            timestamp: new Date(),
            payload: { chatSessionId: session.id, role: 'user' as const },
          });
        }

        // Check if agent exists and is waiting for input
        if (session.agentId) {
          const agent = await agentManager.get(session.agentId);
          if (agent && agent.status === 'waiting_for_input') {
            // Resume the existing agent
            await agentManager.resume(agent.id, { 'chat-response': input.message });
            return { sessionId: session.id, agentId: agent.id, action: 'resumed' as const };
          }

          // Agent exists but not waiting — dismiss it
          if (agent && !agent.userDismissedAt) {
            if (['running'].includes(agent.status)) {
              await agentManager.stop(agent.id);
            }
            await agentManager.dismiss(agent.id);
          }
        }

        // Spawn fresh agent with chat history + context
        const messages = await chatRepo.findMessagesBySessionId(session.id);
        const chatHistory: ChatHistoryEntry[] = messages
          .slice(-MAX_HISTORY_MESSAGES)
          .map(m => ({ role: m.role as 'user' | 'assistant' | 'system', content: m.content }));

        const context = await gatherInitiativeContext(
          ctx.phaseRepository,
          ctx.taskRepository,
          ctx.pageRepository,
          input.initiativeId,
        );

        const prompt = buildChatPrompt(input.targetType, input.targetId, chatHistory, input.message);

        // Create a task for the chat agent
        const targetName = input.targetType === 'phase'
          ? (await ctx.phaseRepository?.findById(input.targetId))?.name ?? input.targetId
          : (await ctx.taskRepository?.findById(input.targetId))?.name ?? input.targetId;

        const task = await taskRepo.create({
          initiativeId: input.initiativeId,
          name: `Chat: ${targetName}`,
          description: `Iterative chat refinement of ${input.targetType}`,
          category: 'discuss',
          status: 'in_progress',
        });

        // Determine target phase/task for input context
        const targetPhase = input.targetType === 'phase'
          ? await ctx.phaseRepository?.findById(input.targetId)
          : undefined;
        const targetTask = input.targetType === 'task'
          ? await ctx.taskRepository?.findById(input.targetId)
          : undefined;

        const agent = await agentManager.spawn({
          taskId: task.id,
          prompt,
          mode: 'chat',
          provider: input.provider,
          initiativeId: input.initiativeId,
          inputContext: {
            initiative,
            phase: targetPhase ?? undefined,
            task: targetTask ?? undefined,
            pages: context.pages.length > 0 ? context.pages : undefined,
            phases: context.phases.length > 0 ? context.phases : undefined,
            tasks: context.tasks.length > 0 ? context.tasks : undefined,
          },
        });

        // Link agent to session
        await chatRepo.updateSession(session.id, { agentId: agent.id });

        return { sessionId: session.id, agentId: agent.id, action: 'spawned' as const };
      }),

    getChatSession: publicProcedure
      .input(z.object({
        targetType: z.enum(['phase', 'task']),
        targetId: z.string().min(1),
      }))
      .query(async ({ ctx, input }) => {
        const chatRepo = requireChatSessionRepository(ctx);
        const session = await chatRepo.findActiveSession(input.targetType, input.targetId);
        if (!session) return null;
        const messages = await chatRepo.findMessagesBySessionId(session.id);
        return { ...session, messages };
      }),

    closeChatSession: publicProcedure
      .input(z.object({ sessionId: z.string().min(1) }))
      .mutation(async ({ ctx, input }) => {
        const chatRepo = requireChatSessionRepository(ctx);
        const agentManager = requireAgentManager(ctx);

        const session = await chatRepo.findSessionById(input.sessionId);
        if (!session) {
          throw new TRPCError({ code: 'NOT_FOUND', message: `Chat session '${input.sessionId}' not found` });
        }

        // Stop and dismiss agent if active
        if (session.agentId) {
          const agent = await agentManager.get(session.agentId);
          if (agent && !agent.userDismissedAt) {
            if (['running', 'waiting_for_input'].includes(agent.status)) {
              await agentManager.stop(agent.id);
            }
            await agentManager.dismiss(agent.id);
          }
        }

        await chatRepo.updateSession(input.sessionId, { status: 'closed' });

        ctx.eventBus.emit({
          type: 'chat:session_closed' as const,
          timestamp: new Date(),
          payload: { chatSessionId: session.id, initiativeId: session.initiativeId },
        });

        return { success: true };
      }),
  };
}