Submatter/apps/web/server/routers/harnesses.ts

import { and, asc, eq } from "drizzle-orm";
import { TRPCError } from "@trpc/server";
import { z } from "zod";
import { harnessProviders, roleHarnessRoutes } from "@tasks/database/schema";
import { harnessTypeSchema } from "@tasks/shared/worker";
import type { HarnessType } from "@tasks/shared/worker";
import { rateLimit } from "@/server/lib/rate-limit";
import { router, workspaceProcedure } from "@/server/trpc";
import type { WorkspaceProcedureContext } from "@/server/trpc";

export const HARNESS_PRESETS: ReadonlyArray<{
  alias: string;
  displayName: string;
  harnessType: HarnessType;
}> = [
  { alias: "doc-writer", displayName: "Doc writer", harnessType: "minimal" },
  { alias: "minimal", displayName: "Minimal", harnessType: "minimal" },
] as const;

type HarnessRow = typeof harnessProviders.$inferSelect;
type RouteRow = typeof roleHarnessRoutes.$inferSelect;

function assertWorkspaceAdmin(ctx: WorkspaceProcedureContext): void {
  if (ctx.workspace.role !== "owner" && ctx.workspace.role !== "admin") {
    throw new TRPCError({ code: "FORBIDDEN", message: "Only workspace owners and admins can manage harnesses." });
  }
}

function canManageWorkspace(ctx: WorkspaceProcedureContext): boolean {
  return ctx.workspace.role === "owner" || ctx.workspace.role === "admin";
}

function mapHarnessRow(row: HarnessRow) {
  return {
    id: row.id,
    workspaceId: row.workspaceId,
    alias: row.alias,
    displayName: row.displayName,
    harnessType: harnessTypeSchema.parse(row.harnessType),
    endpointUrl: row.endpointUrl,
    authEnvVar: row.authEnvVar,
    config: row.config,
    enabled: row.enabled,
    createdAt: row.createdAt,
    updatedAt: row.updatedAt,
  };
}

function mapRouteRow(row: RouteRow) {
  return {
    id: row.id,
    workspaceId: row.workspaceId,
    roleId: row.roleId,
    harnessProviderId: row.harnessProviderId,
    enabled: row.enabled,
    sortOrder: row.sortOrder,
  };
}

const harnessFieldsSchema = z.object({
  alias: z.string().min(1).max(120),
  displayName: z.string().min(1).max(200),
  harnessType: harnessTypeSchema,
  endpointUrl: z.string().url().nullable().optional(),
  authEnvVar: z.string().min(1).max(120).nullable().optional(),
  config: z.record(z.unknown()).nullable().optional(),
  enabled: z.boolean().optional(),
});

export const harnessesRouter = router({
  list: workspaceProcedure.query(async ({ ctx }) => {
    const rows = await ctx.db
      .select()
      .from(harnessProviders)
      .where(eq(harnessProviders.workspaceId, ctx.workspace.id))
      .orderBy(asc(harnessProviders.alias));
    return { harnesses: rows.map(mapHarnessRow), canManage: canManageWorkspace(ctx) };
  }),

  create: workspaceProcedure.input(harnessFieldsSchema).mutation(async ({ ctx, input }) => {
    assertWorkspaceAdmin(ctx);
    try {
      const [created] = await ctx.db
        .insert(harnessProviders)
        .values({
          workspaceId: ctx.workspace.id,
          alias: input.alias.trim(),
          displayName: input.displayName.trim(),
          harnessType: input.harnessType,
          endpointUrl: input.endpointUrl?.trim() || null,
          authEnvVar: input.authEnvVar?.trim() || null,
          config: input.config ?? null,
          enabled: input.enabled ?? true,
          updatedAt: new Date(),
        })
        .returning();
      if (!created) throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "Failed to create harness provider." });
      return mapHarnessRow(created);
    } catch (e) {
      if (isUniqueViolation(e)) throw new TRPCError({ code: "CONFLICT", message: `Alias "${input.alias}" already exists in this workspace.` });
      throw e;
    }
  }),

  update: workspaceProcedure.input(harnessFieldsSchema.partial().extend({ id: z.string().uuid() })).mutation(async ({ ctx, input }) => {
    assertWorkspaceAdmin(ctx);
    const [existing] = await ctx.db
      .select()
      .from(harnessProviders)
      .where(and(eq(harnessProviders.id, input.id), eq(harnessProviders.workspaceId, ctx.workspace.id)))
      .limit(1);
    if (!existing) throw new TRPCError({ code: "NOT_FOUND", message: "Harness provider not found." });
    try {
      const [updated] = await ctx.db
        .update(harnessProviders)
        .set({
          ...(input.alias !== undefined ? { alias: input.alias.trim() } : {}),
          ...(input.displayName !== undefined ? { displayName: input.displayName.trim() } : {}),
          ...(input.harnessType !== undefined ? { harnessType: input.harnessType } : {}),
          ...(input.endpointUrl !== undefined ? { endpointUrl: input.endpointUrl?.trim() || null } : {}),
          ...(input.authEnvVar !== undefined ? { authEnvVar: input.authEnvVar?.trim() || null } : {}),
          ...(input.config !== undefined ? { config: input.config } : {}),
          ...(input.enabled !== undefined ? { enabled: input.enabled } : {}),
          updatedAt: new Date(),
        })
        .where(and(eq(harnessProviders.id, input.id), eq(harnessProviders.workspaceId, ctx.workspace.id)))
        .returning();
      if (!updated) throw new TRPCError({ code: "NOT_FOUND", message: "Harness provider not found." });
      return mapHarnessRow(updated);
    } catch (e) {
      if (isUniqueViolation(e)) throw new TRPCError({ code: "CONFLICT", message: "Another harness already uses that alias." });
      throw e;
    }
  }),

  delete: workspaceProcedure.input(z.object({ id: z.string().uuid() })).mutation(async ({ ctx, input }) => {
    assertWorkspaceAdmin(ctx);
    const [deleted] = await ctx.db
      .delete(harnessProviders)
      .where(and(eq(harnessProviders.id, input.id), eq(harnessProviders.workspaceId, ctx.workspace.id)))
      .returning({ id: harnessProviders.id });
    if (!deleted) throw new TRPCError({ code: "NOT_FOUND", message: "Harness provider not found." });
    return { id: deleted.id };
  }),

  testHarness: workspaceProcedure.input(z.object({ harnessId: z.string().uuid() })).mutation(async ({ ctx, input }) => {
    assertWorkspaceAdmin(ctx);
    const limit = rateLimit({ key: `harness:test:${ctx.workspace.id}`, limit: 10, windowMs: 60_000 });
    if (!limit.allowed) throw new TRPCError({ code: "TOO_MANY_REQUESTS", message: "Test rate limit exceeded. Try again shortly." });
    const [row] = await ctx.db
      .select()
      .from(harnessProviders)
      .where(and(eq(harnessProviders.id, input.harnessId), eq(harnessProviders.workspaceId, ctx.workspace.id)))
      .limit(1);
    if (!row) throw new TRPCError({ code: "NOT_FOUND", message: "Harness provider not found." });
    if (!row.endpointUrl) {
      return { ok: false as const, latencyMs: 0, error: "No endpoint_url configured for this harness." };
    }
    const started = Date.now();
    try {
      const headers: Record<string, string> = {};
      if (row.authEnvVar) {
        const token = process.env[row.authEnvVar]?.trim();
        if (token) headers["Authorization"] = `Bearer ${token}`;
      }
      const base = row.endpointUrl.replace(/\/+$/, "");
      // Try /health then base itself as fallback
      let res: Response | null = null;
      let lastError = "";
      for (const suffix of ["/health", ""]) {
        try {
          res = await fetch(`${base}${suffix}`, { headers, signal: AbortSignal.timeout(8000) });
          if (res.ok) break;
          lastError = `HTTP ${res.status} ${res.statusText}`;
          // don't break on non-ok if we have another suffix
          if (suffix === "") lastError = `HTTP ${res.status} ${res.statusText}`;
        } catch (e) {
          lastError = e instanceof Error ? e.message : String(e);
        }
      }
      const latencyMs = Date.now() - started;
      if (res && res.ok) return { ok: true as const, latencyMs };
      return { ok: false as const, latencyMs, error: lastError || "Health check failed" };
    } catch (e) {
      return { ok: false as const, latencyMs: Date.now() - started, error: e instanceof Error ? e.message : String(e) };
    }
  }),

  listRoutes: workspaceProcedure.query(async ({ ctx }) => {
    const routes = await ctx.db
      .select({
        route: roleHarnessRoutes,
        harnessAlias: harnessProviders.alias,
        harnessType: harnessProviders.harnessType,
        harnessEnabled: harnessProviders.enabled,
      })
      .from(roleHarnessRoutes)
      .innerJoin(harnessProviders, eq(roleHarnessRoutes.harnessProviderId, harnessProviders.id))
      .where(eq(roleHarnessRoutes.workspaceId, ctx.workspace.id))
      .orderBy(asc(roleHarnessRoutes.sortOrder), asc(roleHarnessRoutes.roleId));
    return {
      routes: routes.map((row) => ({
        ...mapRouteRow(row.route),
        harnessAlias: row.harnessAlias,
        harnessType: row.harnessType,
        harnessEnabled: row.harnessEnabled,
      })),
      canManage: canManageWorkspace(ctx),
    };
  }),

  upsertRoute: workspaceProcedure
    .input(z.object({ roleId: z.string().min(1).max(120), harnessProviderId: z.string().uuid(), enabled: z.boolean().optional(), sortOrder: z.number().int().optional() }))
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);
      const [provider] = await ctx.db
        .select({ id: harnessProviders.id })
        .from(harnessProviders)
        .where(and(eq(harnessProviders.id, input.harnessProviderId), eq(harnessProviders.workspaceId, ctx.workspace.id)))
        .limit(1);
      if (!provider) throw new TRPCError({ code: "NOT_FOUND", message: "Harness provider not found." });
      const [existing] = await ctx.db
        .select({ id: roleHarnessRoutes.id })
        .from(roleHarnessRoutes)
        .where(and(eq(roleHarnessRoutes.workspaceId, ctx.workspace.id), eq(roleHarnessRoutes.roleId, input.roleId)))
        .limit(1);
      if (existing) {
        const [updated] = await ctx.db
          .update(roleHarnessRoutes)
          .set({
            harnessProviderId: input.harnessProviderId,
            enabled: input.enabled ?? true,
            ...(input.sortOrder !== undefined ? { sortOrder: input.sortOrder } : {}),
            updatedAt: new Date(),
          })
          .where(and(eq(roleHarnessRoutes.id, existing.id), eq(roleHarnessRoutes.workspaceId, ctx.workspace.id)))
          .returning();
        if (!updated) throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
        return mapRouteRow(updated);
      }
      const [created] = await ctx.db
        .insert(roleHarnessRoutes)
        .values({
          workspaceId: ctx.workspace.id,
          roleId: input.roleId,
          harnessProviderId: input.harnessProviderId,
          enabled: input.enabled ?? true,
          sortOrder: input.sortOrder ?? 0,
          updatedAt: new Date(),
        })
        .returning();
      if (!created) throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "Failed to create harness route." });
      return mapRouteRow(created);
    }),

  deleteRoute: workspaceProcedure.input(z.object({ id: z.string().uuid() })).mutation(async ({ ctx, input }) => {
    assertWorkspaceAdmin(ctx);
    const [deleted] = await ctx.db.delete(roleHarnessRoutes).where(and(eq(roleHarnessRoutes.id, input.id), eq(roleHarnessRoutes.workspaceId, ctx.workspace.id))).returning({ id: roleHarnessRoutes.id });
    if (!deleted) throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
    return { id: deleted.id };
  }),
});

function isUniqueViolation(error: unknown): boolean {
  return typeof error === "object" && error !== null && "code" in error && (error as { code?: string }).code === "23505";
}