spoon/packages/backend/convex/agentJobs.ts

import { ConvexError, v } from 'convex/values';

import type { Doc, Id } from './_generated/dataModel';
import type { MutationCtx } from './_generated/server';
import { internalMutation, mutation, query } from './_generated/server';
import { getOwnedSpoon, getRequiredUserId, optionalText } from './model';

const jobStatus = v.union(
  v.literal('queued'),
  v.literal('claimed'),
  v.literal('preparing'),
  v.literal('running'),
  v.literal('checks_running'),
  v.literal('changes_ready'),
  v.literal('draft_pr_opened'),
  v.literal('failed'),
  v.literal('cancelled'),
  v.literal('timed_out'),
);

const eventLevel = v.union(
  v.literal('debug'),
  v.literal('info'),
  v.literal('warn'),
  v.literal('error'),
);

const eventPhase = v.union(
  v.literal('queued'),
  v.literal('clone'),
  v.literal('plan'),
  v.literal('edit'),
  v.literal('install'),
  v.literal('check'),
  v.literal('test'),
  v.literal('commit'),
  v.literal('push'),
  v.literal('pr'),
  v.literal('cleanup'),
);

const artifactKind = v.union(
  v.literal('plan'),
  v.literal('diff'),
  v.literal('test_output'),
  v.literal('summary'),
  v.literal('error'),
  v.literal('pr_body'),
);

const artifactContentType = v.union(
  v.literal('text/markdown'),
  v.literal('text/plain'),
  v.literal('application/json'),
  v.literal('text/x-diff'),
);

const defaultAgentSettings = {
  enabled: true,
  branchPrefix: 'spoon/agent',
  agentModel: 'gpt-5.1-codex',
  reasoningEffort: 'high' as const,
  maxJobDurationMs: 1_800_000,
  maxOutputBytes: 200_000,
};

const getWorkerToken = () => process.env.SPOON_WORKER_TOKEN?.trim();

const requireWorkerToken = (workerToken: string) => {
  const expected = getWorkerToken();
  if (!expected) throw new ConvexError('SPOON_WORKER_TOKEN is not configured.');
  if (workerToken !== expected) throw new ConvexError('Invalid worker token.');
};

const slugify = (value: string) =>
  value
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, '-')
    .replace(/^-|-$/g, '')
    .slice(0, 36) || 'task';

const shortId = (id: string) => id.replace(/[^a-zA-Z0-9]/g, '').slice(-8);

const buildBranch = (
  requestId: Id<'agentRequests'>,
  prompt: string,
  prefix: string,
  requestedBranchName?: string,
) => {
  const requested = optionalText(requestedBranchName);
  if (requested) return requested.replace(/^\/+|\/+$/g, '');
  return `${prefix.replace(/\/+$/g, '')}/${shortId(requestId)}/${slugify(
    prompt,
  )}`;
};

const getAgentSettings = async (ctx: MutationCtx, spoon: Doc<'spoons'>) => {
  const settings = await ctx.db
    .query('spoonAgentSettings')
    .withIndex('by_spoon', (q) => q.eq('spoonId', spoon._id))
    .first();
  return {
    ...defaultAgentSettings,
    defaultBaseBranch: spoon.forkDefaultBranch ?? spoon.upstreamDefaultBranch,
    ...settings,
  };
};

const assertSecretOwnership = async (
  ctx: MutationCtx,
  spoonId: Id<'spoons'>,
  ownerId: Id<'users'>,
  secretIds: Id<'spoonSecrets'>[],
) => {
  for (const secretId of secretIds) {
    const secret = await ctx.db.get(secretId);
    if (secret?.ownerId !== ownerId || secret.spoonId !== spoonId) {
      throw new ConvexError('Selected secrets must belong to this Spoon.');
    }
  }
};

export const createFromRequest = mutation({
  args: {
    requestId: v.id('agentRequests'),
    selectedSecretIds: v.array(v.id('spoonSecrets')),
    baseBranch: v.optional(v.string()),
    requestedBranchName: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    const ownerId = await getRequiredUserId(ctx);
    const request = await ctx.db.get(args.requestId);
    if (request?.ownerId !== ownerId) {
      throw new ConvexError('Agent request not found.');
    }
    if (request.agentJobId) {
      throw new ConvexError('This request already has an agent job.');
    }
    const spoon = await getOwnedSpoon(ctx, request.spoonId, ownerId);
    if (spoon.provider !== 'github') {
      throw new ConvexError('Agent jobs currently require a GitHub Spoon.');
    }
    if (!spoon.forkOwner || !spoon.forkRepo || !spoon.forkUrl) {
      throw new ConvexError(
        'Add fork repository metadata before queueing a job.',
      );
    }
    const settings = await getAgentSettings(ctx, spoon);
    if (!settings.enabled) {
      throw new ConvexError('Agent jobs are disabled for this Spoon.');
    }
    await assertSecretOwnership(
      ctx,
      spoon._id,
      ownerId,
      args.selectedSecretIds,
    );
    const now = Date.now();
    const baseBranch =
      optionalText(args.baseBranch) ?? settings.defaultBaseBranch;
    const workBranch = buildBranch(
      request._id,
      request.prompt,
      settings.branchPrefix,
      args.requestedBranchName,
    );
    const jobId = await ctx.db.insert('agentJobs', {
      spoonId: spoon._id,
      ownerId,
      agentRequestId: request._id,
      status: 'queued',
      prompt: request.prompt,
      baseBranch,
      workBranch,
      githubInstallationId: spoon.githubInstallationId,
      forkOwner: spoon.forkOwner,
      forkRepo: spoon.forkRepo,
      forkUrl: spoon.forkUrl,
      upstreamOwner: spoon.upstreamOwner,
      upstreamRepo: spoon.upstreamRepo,
      selectedSecretIds: args.selectedSecretIds,
      model: settings.agentModel,
      reasoningEffort: settings.reasoningEffort,
      createdAt: now,
      updatedAt: now,
    });
    await ctx.db.patch(request._id, {
      agentJobId: jobId,
      selectedSecretIds: args.selectedSecretIds,
      baseBranch,
      requestedBranchName: optionalText(args.requestedBranchName),
      status: 'queued',
      updatedAt: now,
    });
    await ctx.db.insert('agentJobEvents', {
      jobId,
      spoonId: spoon._id,
      ownerId,
      level: 'info',
      phase: 'queued',
      message: 'Agent job queued.',
      createdAt: now,
    });
    return jobId;
  },
});

export const listForSpoon = query({
  args: { spoonId: v.id('spoons'), limit: v.optional(v.number()) },
  handler: async (ctx, { spoonId, limit }) => {
    const ownerId = await getRequiredUserId(ctx);
    await getOwnedSpoon(ctx, spoonId, ownerId);
    return await ctx.db
      .query('agentJobs')
      .withIndex('by_spoon', (q) => q.eq('spoonId', spoonId))
      .order('desc')
      .take(limit ?? 25);
  },
});

export const get = query({
  args: { jobId: v.id('agentJobs') },
  handler: async (ctx, { jobId }) => {
    const ownerId = await getRequiredUserId(ctx);
    const job = await ctx.db.get(jobId);
    if (job?.ownerId !== ownerId) throw new ConvexError('Agent job not found.');
    return job;
  },
});

export const listEvents = query({
  args: { jobId: v.id('agentJobs'), limit: v.optional(v.number()) },
  handler: async (ctx, { jobId, limit }) => {
    const ownerId = await getRequiredUserId(ctx);
    const job = await ctx.db.get(jobId);
    if (job?.ownerId !== ownerId) throw new ConvexError('Agent job not found.');
    return await ctx.db
      .query('agentJobEvents')
      .withIndex('by_job', (q) => q.eq('jobId', jobId))
      .order('asc')
      .take(limit ?? 200);
  },
});

export const listArtifacts = query({
  args: { jobId: v.id('agentJobs') },
  handler: async (ctx, { jobId }) => {
    const ownerId = await getRequiredUserId(ctx);
    const job = await ctx.db.get(jobId);
    if (job?.ownerId !== ownerId) throw new ConvexError('Agent job not found.');
    return await ctx.db
      .query('agentJobArtifacts')
      .withIndex('by_job', (q) => q.eq('jobId', jobId))
      .order('asc')
      .collect();
  },
});

export const cancel = mutation({
  args: { jobId: v.id('agentJobs') },
  handler: async (ctx, { jobId }) => {
    const ownerId = await getRequiredUserId(ctx);
    const job = await ctx.db.get(jobId);
    if (job?.ownerId !== ownerId) throw new ConvexError('Agent job not found.');
    if (
      !['queued', 'claimed', 'preparing', 'running', 'checks_running'].includes(
        job.status,
      )
    ) {
      throw new ConvexError('This job cannot be cancelled.');
    }
    const now = Date.now();
    await ctx.db.patch(jobId, {
      status: 'cancelled',
      completedAt: now,
      updatedAt: now,
    });
    await ctx.db.patch(job.agentRequestId, {
      status: 'cancelled',
      updatedAt: now,
    });
    await ctx.db.insert('agentJobEvents', {
      jobId,
      spoonId: job.spoonId,
      ownerId,
      level: 'warn',
      phase: 'cleanup',
      message: 'Agent job cancelled by user.',
      createdAt: now,
    });
    return { success: true };
  },
});

export const claimNextInternal = internalMutation({
  args: { workerId: v.string() },
  handler: async (ctx, { workerId }) => {
    const job = await ctx.db
      .query('agentJobs')
      .withIndex('by_claim', (q) => q.eq('status', 'queued'))
      .order('asc')
      .first();
    if (!job) return null;
    const spoon = await ctx.db.get(job.spoonId);
    const aiSettings = await ctx.db
      .query('userAiSettings')
      .withIndex('by_user_provider', (q) =>
        q.eq('userId', job.ownerId).eq('provider', 'openai'),
      )
      .first();
    const agentSettings = await ctx.db
      .query('spoonAgentSettings')
      .withIndex('by_spoon', (q) => q.eq('spoonId', job.spoonId))
      .first();
    const secrets = [];
    for (const secretId of job.selectedSecretIds) {
      const secret = await ctx.db.get(secretId);
      if (secret?.ownerId === job.ownerId && secret.spoonId === job.spoonId) {
        secrets.push(secret);
      }
    }
    const now = Date.now();
    await ctx.db.patch(job._id, {
      status: 'claimed',
      claimedBy: workerId,
      claimedAt: now,
      updatedAt: now,
    });
    await ctx.db.patch(job.agentRequestId, {
      status: 'running',
      updatedAt: now,
    });
    await ctx.db.insert('agentJobEvents', {
      jobId: job._id,
      spoonId: job.spoonId,
      ownerId: job.ownerId,
      level: 'info',
      phase: 'queued',
      message: `Claimed by ${workerId}.`,
      createdAt: now,
    });
    return {
      job: { ...job, status: 'claimed' as const, claimedBy: workerId },
      spoon,
      aiSettings,
      agentSettings,
      secrets,
    };
  },
});

export const updateStatus = mutation({
  args: {
    workerToken: v.string(),
    workerId: v.string(),
    jobId: v.id('agentJobs'),
    status: jobStatus,
    error: v.optional(v.string()),
    summary: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    requireWorkerToken(args.workerToken);
    const job = await ctx.db.get(args.jobId);
    if (job?.claimedBy !== args.workerId) {
      throw new ConvexError('Agent job not claimed by this worker.');
    }
    const now = Date.now();
    const patch: Partial<Doc<'agentJobs'>> = {
      status: args.status,
      updatedAt: now,
    };
    if (args.status === 'running' && !job.startedAt) patch.startedAt = now;
    if (
      ['failed', 'cancelled', 'timed_out', 'draft_pr_opened'].includes(
        args.status,
      )
    ) {
      patch.completedAt = now;
    }
    if (args.error !== undefined) patch.error = args.error;
    if (args.summary !== undefined) patch.summary = args.summary;
    await ctx.db.patch(args.jobId, patch);
    return { success: true };
  },
});

export const completeWithDraftPr = mutation({
  args: {
    workerToken: v.string(),
    workerId: v.string(),
    jobId: v.id('agentJobs'),
    commitSha: v.string(),
    pullRequestUrl: v.string(),
    pullRequestNumber: v.number(),
    summary: v.string(),
  },
  handler: async (ctx, args) => {
    requireWorkerToken(args.workerToken);
    const job = await ctx.db.get(args.jobId);
    if (job?.claimedBy !== args.workerId) {
      throw new ConvexError('Agent job not claimed by this worker.');
    }
    const now = Date.now();
    await ctx.db.patch(args.jobId, {
      status: 'draft_pr_opened',
      commitSha: args.commitSha,
      pullRequestUrl: args.pullRequestUrl,
      pullRequestNumber: args.pullRequestNumber,
      summary: args.summary,
      completedAt: now,
      updatedAt: now,
    });
    await ctx.db.patch(job.agentRequestId, {
      status: 'merge_request_opened',
      mergeRequestUrl: args.pullRequestUrl,
      summary: args.summary,
      updatedAt: now,
    });
    return { success: true };
  },
});

export const appendEvent = mutation({
  args: {
    workerToken: v.string(),
    workerId: v.string(),
    jobId: v.id('agentJobs'),
    level: eventLevel,
    phase: eventPhase,
    message: v.string(),
    metadata: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    requireWorkerToken(args.workerToken);
    const job = await ctx.db.get(args.jobId);
    if (job?.claimedBy !== args.workerId) {
      throw new ConvexError('Agent job not claimed by this worker.');
    }
    return await ctx.db.insert('agentJobEvents', {
      jobId: args.jobId,
      spoonId: job.spoonId,
      ownerId: job.ownerId,
      level: args.level,
      phase: args.phase,
      message: args.message,
      metadata: args.metadata,
      createdAt: Date.now(),
    });
  },
});

export const addArtifact = mutation({
  args: {
    workerToken: v.string(),
    workerId: v.string(),
    jobId: v.id('agentJobs'),
    kind: artifactKind,
    title: v.string(),
    content: v.string(),
    contentType: artifactContentType,
  },
  handler: async (ctx, args) => {
    requireWorkerToken(args.workerToken);
    const job = await ctx.db.get(args.jobId);
    if (job?.claimedBy !== args.workerId) {
      throw new ConvexError('Agent job not claimed by this worker.');
    }
    return await ctx.db.insert('agentJobArtifacts', {
      jobId: args.jobId,
      spoonId: job.spoonId,
      ownerId: job.ownerId,
      kind: args.kind,
      title: args.title,
      content: args.content,
      contentType: args.contentType,
      createdAt: Date.now(),
    });
  },
});