Genarrative/server-node/src/routes/customWorldAgent.ts

import { Router } from 'express';
import { z } from 'zod';

import type {
  CreateCustomWorldAgentSessionRequest,
  CustomWorldAgentActionRequest,
  SendCustomWorldAgentMessageRequest,
} from '../../../packages/shared/src/contracts/customWorldAgent.js';
import type { AppContext } from '../context.js';
import { badRequest, notFound } from '../errors.js';
import { asyncHandler, prepareApiResponse, sendApiResponse } from '../http.js';
import { requireJwtAuth } from '../middleware/auth.js';
import { routeMeta } from '../middleware/routeMeta.js';

const createSessionSchema = z.object({
  seedText: z.string().trim().optional().default(''),
});

const sendMessageSchema = z.object({
  clientMessageId: z.string().trim().min(1),
  text: z.string().trim().min(1),
  quickFillRequested: z.boolean().optional().default(false),
  focusCardId: z.string().trim().nullable().optional().default(null),
  selectedCardIds: z.array(z.string().trim().min(1)).optional().default([]),
});

const actionSchema = z.discriminatedUnion('action', [
  z.object({
    action: z.literal('draft_foundation'),
  }),
  z.object({
    action: z.literal('update_draft_card'),
    cardId: z.string().trim().min(1),
    sections: z
      .array(
        z.object({
          sectionId: z.string().trim().min(1),
          value: z.string(),
        }),
      )
      .min(1),
  }),
  z.object({
    action: z.literal('sync_result_profile'),
    profile: z.record(z.string(), z.unknown()),
  }),
  z.object({
    action: z.literal('generate_characters'),
    count: z.number().int().min(1).max(3),
    promptText: z.string().trim().nullable().optional().default(null),
    anchorCardIds: z.array(z.string().trim().min(1)).optional().default([]),
  }),
  z.object({
    action: z.literal('generate_landmarks'),
    count: z.number().int().min(1).max(3),
    promptText: z.string().trim().nullable().optional().default(null),
    anchorCardIds: z.array(z.string().trim().min(1)).optional().default([]),
  }),
  z.object({
    action: z.literal('generate_role_assets'),
    roleIds: z.array(z.string().trim().min(1)).min(1),
  }),
  z.object({
    action: z.literal('sync_role_assets'),
    roleId: z.string().trim().min(1),
    portraitPath: z.string().trim().min(1),
    generatedVisualAssetId: z.string().trim().min(1),
    generatedAnimationSetId: z.string().trim().nullable().optional(),
    animationMap: z.record(z.string(), z.unknown()).nullable().optional(),
  }),
  z.object({
    action: z.literal('generate_scene_assets'),
    sceneIds: z.array(z.string().trim().min(1)).min(1),
  }),
  z.object({
    action: z.literal('sync_scene_assets'),
    sceneId: z.string().trim().min(1),
    sceneKind: z.enum(['camp', 'landmark']),
    imageSrc: z.string().trim().min(1),
    generatedSceneAssetId: z.string().trim().min(1),
    generatedScenePrompt: z.string().trim().nullable().optional(),
    generatedSceneModel: z.string().trim().nullable().optional(),
  }),
  z.object({
    action: z.literal('expand_long_tail'),
  }),
  z.object({
    action: z.literal('publish_world'),
  }),
  z.object({
    action: z.literal('revert_checkpoint'),
    checkpointId: z.string().trim().min(1),
  }),
]);

function readParam(param: string | string[] | undefined) {
  return Array.isArray(param) ? param[0]?.trim() || '' : param?.trim() || '';
}

export function createCustomWorldAgentRoutes(context: AppContext) {
  const router = Router();
  const requireAuth = requireJwtAuth(context.config, context.userRepository);

  router.use(requireAuth);

  router.post(
    '/sessions',
    routeMeta({ operation: 'runtime.customWorldAgent.createSession' }),
    asyncHandler(async (request, response) => {
      const payload = createSessionSchema.parse(
        request.body,
      ) as CreateCustomWorldAgentSessionRequest;
      sendApiResponse(response, {
        session: await context.customWorldAgentOrchestrator.createSession(
          request.userId!,
          payload,
        ),
      });
    }),
  );

  router.get(
    '/sessions/:sessionId',
    routeMeta({ operation: 'runtime.customWorldAgent.getSession' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }

      const session = await context.customWorldAgentOrchestrator.getSessionSnapshot(
        request.userId!,
        sessionId,
      );
      if (!session) {
        throw notFound('custom world agent session not found');
      }

      sendApiResponse(response, session);
    }),
  );

  router.post(
    '/sessions/:sessionId/messages',
    routeMeta({ operation: 'runtime.customWorldAgent.sendMessage' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }

      const payload = sendMessageSchema.parse(
        request.body,
      ) as SendCustomWorldAgentMessageRequest;
      sendApiResponse(
        response,
        await context.customWorldAgentOrchestrator.submitMessage(
          request.userId!,
          sessionId,
          payload,
        ),
      );
    }),
  );

  router.post(
    '/sessions/:sessionId/messages/stream',
    routeMeta({ operation: 'runtime.customWorldAgent.streamMessage' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }

      const payload = sendMessageSchema.parse(
        request.body,
      ) as SendCustomWorldAgentMessageRequest;
      await context.customWorldAgentOrchestrator.streamMessage({
        request,
        response,
        userId: request.userId!,
        sessionId,
        payload,
      });
    }),
  );

  router.post(
    '/sessions/:sessionId/actions',
    routeMeta({ operation: 'runtime.customWorldAgent.executeAction' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }

      const payload = actionSchema.parse(
        request.body,
      ) as CustomWorldAgentActionRequest;
      sendApiResponse(
        response,
        await context.customWorldAgentOrchestrator.executeAction(
          request.userId!,
          sessionId,
          payload,
        ),
      );
    }),
  );

  router.get(
    '/sessions/:sessionId/operations/:operationId',
    routeMeta({ operation: 'runtime.customWorldAgent.getOperation' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      const operationId = readParam(request.params.operationId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }
      if (!operationId) {
        throw badRequest('operationId is required');
      }

      const operation = await context.customWorldAgentOrchestrator.getOperation(
        request.userId!,
        sessionId,
        operationId,
      );
      if (!operation) {
        throw notFound('custom world agent operation not found');
      }

      prepareApiResponse(request, response, {
        statusCode: 200,
        headers: {
          'Content-Type': 'application/json; charset=utf-8',
        },
      });
      response.end(JSON.stringify({ operation }));
    }),
  );

  router.get(
    '/sessions/:sessionId/cards/:cardId',
    routeMeta({ operation: 'runtime.customWorldAgent.getCardDetail' }),
    asyncHandler(async (request, response) => {
      const sessionId = readParam(request.params.sessionId);
      const cardId = readParam(request.params.cardId);
      if (!sessionId) {
        throw badRequest('sessionId is required');
      }
      if (!cardId) {
        throw badRequest('cardId is required');
      }

      const card = await context.customWorldAgentOrchestrator.getCardDetail(
        request.userId!,
        sessionId,
        cardId,
      );
      if (!card) {
        throw notFound('custom world agent card not found');
      }

      sendApiResponse(response, {
        card,
      });
    }),
  );

  return router;
}