AIclinicalresearch/backend/src/modules/ssa/routes/agent-execution.routes.ts

import { FastifyInstance, FastifyRequest } from 'fastify';
import { z } from 'zod';
import { logger } from '../../../common/logging/index.js';
import { prisma } from '../../../config/database.js';
import toolParamConstraints from '../config/tool_param_constraints.json' with { type: 'json' };

type ParamRule = {
  paramType: 'single' | 'multi';
  requiredType: 'any' | 'numeric' | 'categorical';
  hint?: string;
};

type ConstraintsMap = Record<string, Record<string, ParamRule>>;

const Constraints = toolParamConstraints as ConstraintsMap;

function getUserId(request: FastifyRequest): string {
  const userId = (request as any).user?.userId;
  if (!userId) throw new Error('User not authenticated');
  return userId;
}

function normalizeColumnType(raw?: string): 'numeric' | 'categorical' | 'other' {
  const t = (raw || '').toLowerCase();
  if (['numeric', 'number', 'float', 'double', 'integer', 'int'].includes(t)) return 'numeric';
  if (['categorical', 'category', 'factor', 'string', 'text'].includes(t)) return 'categorical';
  return 'other';
}

export default async function agentExecutionRoutes(app: FastifyInstance) {
  app.patch<{
    Params: { executionId: string };
    Body: { steps: Array<{ stepOrder: number; params: Record<string, unknown> }> };
  }>(
    '/:executionId/plan-params',
    async (request, reply) => {
      const userId = getUserId(request);
      const { executionId } = request.params;
      const BodySchema = z.object({
        steps: z.array(
          z.object({
            stepOrder: z.number().int().positive(),
            params: z.record(z.string(), z.unknown()),
          }),
        ).min(1),
      });
      const parsed = BodySchema.safeParse(request.body);
      if (!parsed.success) {
        return reply.status(400).send({
          success: false,
          error: 'Invalid request body',
          details: parsed.error.flatten(),
        });
      }

      const execution = await (prisma as any).ssaAgentExecution.findUnique({
        where: { id: executionId },
        select: {
          id: true,
          status: true,
          reviewResult: true,
          sessionId: true,
          session: { select: { userId: true, dataSchema: true } },
        },
      });

      if (!execution) {
        return reply.status(404).send({ success: false, error: 'Agent execution not found' });
      }
      if (execution.session?.userId !== userId) {
        return reply.status(403).send({ success: false, error: 'No permission to modify this execution' });
      }
      if (execution.status !== 'plan_pending') {
        return reply.status(409).send({
          success: false,
          error: `Cannot modify plan params when status is '${execution.status}'`,
        });
      }

      const review = (execution.reviewResult || {}) as any;
      const steps = Array.isArray(review.steps) ? review.steps : [];
      if (steps.length === 0) {
        return reply.status(400).send({ success: false, error: 'Execution has no structured plan steps' });
      }

      const schemaColumns = ((execution.session?.dataSchema as any)?.columns || []) as Array<any>;
      const columnTypeMap = new Map<string, 'numeric' | 'categorical' | 'other'>(
        schemaColumns.map((c: any) => [c.name, normalizeColumnType(c.type || c.inferred_type)]),
      );

      for (const patch of parsed.data.steps) {
        const target = steps.find((s: any) => Number(s.order) === patch.stepOrder);
        if (!target) {
          return reply.status(400).send({
            success: false,
            error: `Step ${patch.stepOrder} does not exist in reviewResult.steps`,
          });
        }

        const toolCode = String(target.toolCode || target.tool_code || '').trim();
        const ruleSet = toolCode ? Constraints[toolCode] : undefined;

        for (const [paramName, value] of Object.entries(patch.params)) {
          if (typeof value === 'string' && value && columnTypeMap.size > 0 && !columnTypeMap.has(value)) {
            return reply.status(400).send({
              success: false,
              error: `Variable '${value}' in step ${patch.stepOrder}.${paramName} does not exist in dataset`,
            });
          }
          if (Array.isArray(value) && columnTypeMap.size > 0) {
            for (const v of value) {
              if (typeof v === 'string' && !columnTypeMap.has(v)) {
                return reply.status(400).send({
                  success: false,
                  error: `Variable '${v}' in step ${patch.stepOrder}.${paramName} does not exist in dataset`,
                });
              }
            }
          }

          const rule = ruleSet?.[paramName];
          if (!rule || rule.requiredType === 'any') continue;

          if (rule.paramType === 'single' && typeof value === 'string') {
            const t = columnTypeMap.get(value);
            if ((rule.requiredType === 'numeric' && t !== 'numeric')
              || (rule.requiredType === 'categorical' && t !== 'categorical')) {
              return reply.status(400).send({
                success: false,
                error: `Step ${patch.stepOrder}.${paramName} requires ${rule.requiredType} variable`,
                hint: rule.hint,
              });
            }
          }

          if (rule.paramType === 'multi' && Array.isArray(value)) {
            for (const v of value) {
              if (typeof v !== 'string') continue;
              const t = columnTypeMap.get(v);
              if ((rule.requiredType === 'numeric' && t !== 'numeric')
                || (rule.requiredType === 'categorical' && t !== 'categorical')) {
                return reply.status(400).send({
                  success: false,
                  error: `Step ${patch.stepOrder}.${paramName} requires ${rule.requiredType} variables`,
                  hint: rule.hint,
                });
              }
            }
          }
        }
      }

      // 写回 reviewResult.steps[].params
      for (const patch of parsed.data.steps) {
        const idx = steps.findIndex((s: any) => Number(s.order) === patch.stepOrder);
        if (idx >= 0) {
          const currentParams = (steps[idx].params || {}) as Record<string, unknown>;
          steps[idx].params = { ...currentParams, ...patch.params };
        }
      }

      review.steps = steps;
      await (prisma as any).ssaAgentExecution.update({
        where: { id: executionId },
        data: {
          reviewResult: review,
        },
      });

      logger.info('[SSA:AgentExecution] plan params updated', {
        executionId,
        stepsUpdated: parsed.data.steps.length,
      });

      return reply.send({
        success: true,
        stepsUpdated: parsed.data.steps.length,
        reviewResult: review,
      });
    },
  );
}