server/ApiTaskHandler.ts

import type {Application, Request, Response} from 'express';
import {Router} from 'express';
import {VIDEO_DESCRIPTION_MODEL_ID} from '../shared/clusterModels.js';
import type {CompleteTaskBody} from '../shared/completeTask.js';
import {parseCompleteTaskBody} from './completeTask.js';
import {
  parseHostField,
  validateHostReady,
  type HostValidationFailure,
} from './hostValidation.js';
import {resolveImagePayload} from './imageLoader.js';
import type {BrokerStore} from './store.js';
import type {
  DescribeRequestBody,
  DetectRequestBody,
  DetectionTask,
} from './types.js';

type ApiErrorType =
  | 'invalid_request_error'
  | 'host_unavailable'
  | 'detection_error'
  | 'description_error'
  | 'timeout';

type ResolvedImage = {imageBase64: string; mimeType: string};

type ImageTaskOptions = {
  hostReady?: Parameters<typeof validateHostReady>[2];
  errorType: ApiErrorType;
  fallbackMessage: string;
  createTask: (hostId: string, image: ResolvedImage) => DetectionTask;
  buildSuccess: (completed: DetectionTask) => Record<string, unknown> | null;
};

export class ApiTaskHandler {
  readonly router: Router;

  constructor(
    private readonly store: BrokerStore,
    private readonly taskTimeoutMs: number,
  ) {
    this.router = Router();
    this.router.post('/api/tasks/:taskId/processing', (req, res) =>
      this.markProcessing(req, res),
    );
    this.router.post('/api/tasks/:taskId/complete', (req, res) => this.complete(req, res));
    this.router.get('/v1/tasks/:taskId', (req, res) => this.getTask(req, res));
    this.router.post('/v1/detect', (req, res) => void this.detect(req, res));
    this.router.post('/v1/describe', (req, res) => void this.describe(req, res));
  }

  mount(app: Application): void {
    app.use(this.router);
  }

  private taskId(req: Request): string {
    return String(req.params.taskId);
  }

  private markProcessing(req: Request, res: Response): void {
    this.store.markProcessing(this.taskId(req));
    res.json({status: 'ok'});
  }

  private complete(req: Request, res: Response): void {
    const taskId = this.taskId(req);
    const body = req.body as CompleteTaskBody;

    try {
      const existing = this.store.getTask(taskId);
      if (!existing) {
        res.status(404).json({error: `Task '${taskId}' not found`});
        return;
      }

      const hostModel = this.store.getHost(existing.hostId)?.model;
      const outcome = parseCompleteTaskBody(existing, body, hostModel);
      const task = this.store.completeTask(taskId, outcome);

      res.json({
        status: 'ok',
        task_id: task.id,
        kind: task.kind,
        ...(task.error ? {error: task.error} : {}),
      });
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      const status = message.includes('not found') ? 404 : 400;
      res.status(status).json({error: message});
    }
  }

  private getTask(req: Request, res: Response): void {
    const task = this.store.getTask(this.taskId(req));
    if (!task) {
      res.status(404).json({error: 'Task not found'});
      return;
    }

    res.json(this.serializeTask(task));
  }

  private async detect(req: Request, res: Response): Promise<void> {
    const body = req.body as DetectRequestBody;
    await this.runImageTask(res, body, {
      errorType: 'detection_error',
      fallbackMessage: 'Detection failed',
      createTask: (hostId, image) =>
        this.store.createDetectTask(hostId, {
          ...image,
          threshold: body.threshold ?? 0.5,
        }),
      buildSuccess: (completed) => ({
        task_id: completed.id,
        host: completed.hostId,
        threshold: completed.threshold,
        detections: completed.results ?? [],
      }),
    });
  }

  private async describe(req: Request, res: Response): Promise<void> {
    const body = req.body as DescribeRequestBody;
    await this.runImageTask(res, body, {
      hostReady: {
        requiredModel: VIDEO_DESCRIPTION_MODEL_ID,
        requiredModelLabel: 'Video description',
      },
      errorType: 'description_error',
      fallbackMessage: 'Description failed',
      createTask: (hostId, image) =>
        this.store.createDescribeTask(hostId, {
          ...image,
          instruction: body.instruction?.trim() || 'What do you see?',
          maxNewTokens: body.max_new_tokens ?? 100,
        }),
      buildSuccess: (completed) => {
        const description = completed.description?.trim() ?? '';
        if (!description) {
          this.apiError(
            res,
            500,
            'Host returned an empty description. Re-register the host with "Video description" selected and wait until SmolVLM finishes loading.',
            'description_error',
            {task_id: completed.id},
          );
          return null;
        }

        return {
          task_id: completed.id,
          host: completed.hostId,
          instruction: completed.instruction,
          max_new_tokens: completed.maxNewTokens,
          description,
        };
      },
    });
  }

  private async runImageTask(
    res: Response,
    body: DetectRequestBody | DescribeRequestBody,
    options: ImageTaskOptions,
  ): Promise<void> {
    const hostId = this.resolveReadyHost(res, body, options.hostReady);
    if (!hostId) {
      return;
    }

    try {
      const image = await resolveImagePayload(body);
      const task = options.createTask(hostId, image);
      const completed = await this.store.waitForTask(task.id, this.taskTimeoutMs);
      this.respondWithCompletedTask(res, completed, options);
    } catch (error) {
      this.respondWithTaskFailure(res, error);
    }
  }

  private resolveReadyHost(
    res: Response,
    body: {host?: string},
    hostReady?: Parameters<typeof validateHostReady>[2],
  ): string | undefined {
    const hostField = parseHostField(body);
    if (typeof hostField !== 'string') {
      this.sendHostFailure(res, hostField);
      return undefined;
    }

    const failure = validateHostReady(this.store, hostField, hostReady);
    if (failure) {
      this.sendHostFailure(res, failure);
      return undefined;
    }

    return hostField;
  }

  private sendHostFailure(res: Response, failure: HostValidationFailure): void {
    const type =
      failure.status === 503 ? 'host_unavailable' : 'invalid_request_error';
    this.apiError(res, failure.status, failure.message, type);
  }

  private respondWithCompletedTask(
    res: Response,
    completed: DetectionTask,
    options: ImageTaskOptions,
  ): void {
    if (completed.status === 'error') {
      this.apiError(
        res,
        500,
        completed.error ?? options.fallbackMessage,
        options.errorType,
        {task_id: completed.id},
      );
      return;
    }

    const payload = options.buildSuccess(completed);
    if (payload) {
      res.json(payload);
    }
  }

  private respondWithTaskFailure(res: Response, error: unknown): void {
    const message = error instanceof Error ? error.message : String(error);
    const status = message.includes('timed out') ? 504 : 400;
    this.apiError(res, status, message, status === 504 ? 'timeout' : 'invalid_request_error');
  }

  private apiError(
    res: Response,
    status: number,
    message: string,
    type: ApiErrorType,
    extra?: Record<string, unknown>,
  ): void {
    res.status(status).json({
      error: {message, type},
      ...extra,
    });
  }

  private serializeTask(task: DetectionTask) {
    return {
      id: task.id,
      kind: task.kind,
      host: task.hostId,
      status: task.status,
      results: task.results,
      description: task.description,
      error: task.error,
      created_at: task.createdAt,
      started_at: task.startedAt,
      completed_at: task.completedAt,
    };
  }
}