triggerdotdev--trigger.dev/apps/webapp/app/services/dashboardAgent.server.ts

import { signUserActorToken } from "@trigger.dev/rbac";
import { TriggerClient } from "@trigger.dev/sdk";
import { chat } from "@trigger.dev/sdk/ai";
import { Counter } from "prom-client";
import { prisma } from "~/db.server";
import { env } from "~/env.server";
import { metricsRegister } from "~/metrics.server";
import { singleton } from "~/utils/singleton";
import { runStore } from "~/v3/runStore.server";
import { githubApp } from "./gitHub.server";
import { logger } from "./logger.server";

const TASK_ID = "dashboard-agent";

// The wake poll runs once a minute per visible tab. That trade-off is only defensible while it
// stays measurable, so count the requests. singleton: module-scope registration double-registers
// under dev HMR.
export const dashboardAgentWakeFeedCounter = singleton(
  "dashboardAgentWakeFeedCounter",
  () =>
    new Counter({
      name: "dashboard_agent_wake_feed_requests_total",
      help: "Requests to the dashboard agent's wake feed",
      registers: [metricsRegister],
    })
);

// Read-only cap on the agent's delegated user-actor token. `read:apiKeys` is
// what lets it exchange the token for an env JWT (the gate on the exchange
// route); the rest scope the actual reads. No write/admin scopes, so even a
// leaked token can't mutate anything.
export const DASHBOARD_AGENT_UAT_CAP = [
  "read:apiKeys",
  "read:runs",
  "read:deployments",
  "read:environments",
  "read:errors",
  "read:query",
  // Queue metrics ride on `read:query`, but a queue's own row — paused, depth, limit —
  // is a `queues` read, and without it the agent can only see the metrics window.
  "read:queues",
];

// Minted fresh on every turn (the `in` proxy injects it), so the lifetime only
// has to cover a single turn's tool calls. Short by design — a stale token in
// the agent's run payload expires quickly.
const DASHBOARD_AGENT_UAT_TTL_SECONDS = 10 * 60;

// The Trigger instance this webapp runs against — the same origin the agent
// task calls back to (as the user) for its read tools.
export function dashboardAgentApiOrigin(): string {
  return env.API_ORIGIN ?? env.APP_ORIGIN;
}

// Mint a short-lived, read-only delegated token for the signed-in user. Self
// service from the dashboard session (never a PAT), so a user can only ever
// mint a token for themselves. The `in` proxy injects this into the turn's
// metadata so the token reaches the agent without ever touching the browser.
//
// Endpoints that bind something to one environment read `environmentId` off the token,
// so the agent can't name a different one in a request body.
export function mintDashboardAgentUserActorToken(
  userId: string,
  opts: { environmentId: string }
): Promise<string> {
  return signUserActorToken(env.SESSION_SECRET, {
    userId,
    client: "dashboard-agent",
    environmentId: opts.environmentId,
    cap: DASHBOARD_AGENT_UAT_CAP,
    expirationTime: Math.floor(Date.now() / 1000) + DASHBOARD_AGENT_UAT_TTL_SECONDS,
  });
}

// The session is created in whatever env DASHBOARD_AGENT_SECRET_KEY belongs to.
// baseURL is the Trigger instance this webapp runs against (its own API origin).
function dashboardAgentConfig() {
  const accessToken = env.DASHBOARD_AGENT_SECRET_KEY;
  if (!accessToken) return null;
  return { baseURL: dashboardAgentApiOrigin(), accessToken };
}

export function isDashboardAgentConfigured(): boolean {
  return Boolean(env.DASHBOARD_AGENT_SECRET_KEY);
}

// Pins every agent session (and its continuation runs) to a deployed version
// when DASHBOARD_AGENT_VERSION is set; unset runs on the env's current version.
export function dashboardAgentTriggerConfig(): { lockToVersion: string } | undefined {
  return env.DASHBOARD_AGENT_VERSION ? { lockToVersion: env.DASHBOARD_AGENT_VERSION } : undefined;
}

export async function startDashboardAgentSession(params: {
  chatId: string;
  clientData?: Record<string, unknown>;
}): Promise<{ publicAccessToken: string }> {
  const config = dashboardAgentConfig();
  if (!config) throw new Error("DASHBOARD_AGENT_SECRET_KEY is not set");
  const startSession = chat.createStartSessionAction(TASK_ID, {
    apiClient: config,
    triggerConfig: dashboardAgentTriggerConfig(),
  });
  return startSession({ chatId: params.chatId, clientData: params.clientData });
}

export async function mintDashboardAgentToken(chatId: string): Promise<string> {
  const config = dashboardAgentConfig();
  if (!config) throw new Error("DASHBOARD_AGENT_SECRET_KEY is not set");
  const client = new TriggerClient(config);
  return client.auth.createPublicToken({
    scopes: { read: { sessions: chatId }, write: { sessions: chatId } },
    expirationTime: "1h",
  });
}

// A signed, short-lived pointer to the project's connected repo at a commit. Only
// the URL crosses to the agent; the GitHub token stays here. The agent's code
// tools download + extract it on their own filesystem (see @internal/dashboard-agent).
export type DashboardAgentRepoSnapshot = {
  tarballUrl: string;
  owner: string;
  repo: string;
  sha: string;
  defaultBranch?: string;
};

// The GitHub archive redirect URL is valid for a few minutes; cache the resolved
// pointer briefly so multi-turn chats don't re-mint a token + re-resolve on every
// message. Keyed by project + ref.
const repoSnapshotCache = new Map<
  string,
  { snapshot: DashboardAgentRepoSnapshot; expiresAt: number }
>();
const REPO_SNAPSHOT_TTL_MS = 60_000;
const REPO_SNAPSHOT_MAX_ENTRIES = 1_000;

// Drop expired entries (key cardinality grows with each unique project + pinned
// SHA), then evict oldest-first if still over the cap, so the cache can't grow
// unbounded over a process lifetime.
function pruneRepoSnapshotCache(now = Date.now()) {
  for (const [key, value] of repoSnapshotCache) {
    if (value.expiresAt <= now) repoSnapshotCache.delete(key);
  }
  let overflow = repoSnapshotCache.size - REPO_SNAPSHOT_MAX_ENTRIES;
  if (overflow <= 0) return;
  for (const key of repoSnapshotCache.keys()) {
    repoSnapshotCache.delete(key);
    if (--overflow <= 0) break;
  }
}

/**
 * Resolve the code-mode repo snapshot for a project, or null when the GitHub App
 * is disabled / no repo is connected (which keeps the agent in assistant mode).
 *
 * Mints a `contents:read` installation token scoped to the one repo, resolves the
 * signed archive URL, and returns just that URL. The token never leaves the
 * server. `opts.ref` pins a specific commit (run-SHA pinning); without it, the
 * tracked prod branch (or the repo default) head is used.
 */
export async function resolveDashboardAgentRepoSnapshot(
  projectId: string,
  opts: { ref?: string } = {}
): Promise<DashboardAgentRepoSnapshot | null> {
  if (!githubApp) return null;

  // Cache per project + ref so HEAD and each pinned commit are cached separately.
  const cacheKey = `${projectId}:${opts.ref ?? "HEAD"}`;
  const cached = repoSnapshotCache.get(cacheKey);
  if (cached && cached.expiresAt > Date.now()) return cached.snapshot;
  if (cached) repoSnapshotCache.delete(cacheKey);

  const connected = await prisma.connectedGithubRepository.findFirst({
    where: { projectId },
    select: {
      branchTracking: true,
      repository: {
        select: {
          fullName: true,
          defaultBranch: true,
          installation: { select: { appInstallationId: true } },
        },
      },
    },
  });
  if (!connected) return null;

  const [owner, repo] = connected.repository.fullName.split("/");
  if (!owner || !repo) return null;
  const installationId = Number(connected.repository.installation.appInstallationId);
  const defaultBranch = connected.repository.defaultBranch;
  const tracking = connected.branchTracking as { prod?: { branch?: string } } | null;
  // An explicit 40-char commit SHA is used directly (run-SHA pinning); otherwise
  // resolve the requested branch, the tracked prod branch, or the repo default.
  const requested = opts.ref;
  const isSha = !!requested && /^[0-9a-f]{40}$/i.test(requested);
  const branchRef = requested && !isSha ? requested : tracking?.prod?.branch || defaultBranch;

  try {
    const octokit = await githubApp.getInstallationOctokit(installationId);
    const sha = isSha
      ? requested!
      : (await octokit.rest.repos.getBranch({ owner, repo, branch: branchRef })).data.commit.sha;

    const token = await githubApp.octokit.rest.apps.createInstallationAccessToken({
      installation_id: installationId,
      repositories: [repo],
      permissions: { contents: "read" },
    });

    // Resolve the signed archive URL without downloading the bytes server-side.
    const redirect = await fetch(`https://api.github.com/repos/${owner}/${repo}/tarball/${sha}`, {
      headers: {
        Authorization: `Bearer ${token.data.token}`,
        Accept: "application/vnd.github+json",
        "User-Agent": "trigger-dashboard-agent",
      },
      redirect: "manual",
    });
    const tarballUrl = redirect.headers.get("location");
    if (!tarballUrl) return null;

    const snapshot: DashboardAgentRepoSnapshot = { tarballUrl, owner, repo, sha, defaultBranch };
    pruneRepoSnapshotCache();
    repoSnapshotCache.set(cacheKey, { snapshot, expiresAt: Date.now() + REPO_SNAPSHOT_TTL_MS });
    return snapshot;
  } catch (error) {
    logger.error("Failed to resolve dashboard agent repo snapshot", { error, projectId });
    return null;
  }
}

// Map a run (by friendly id) to the commit its deployed version came from, for
// run-SHA pinning. A run locks to a BackgroundWorker (`lockedToVersionId`), whose
// WorkerDeployment carries the commit. Null for runs with no deployed version
// (e.g. dev runs), so the agent falls back to the branch head.
export async function resolveRunCommit(
  environmentId: string,
  runFriendlyId: string
): Promise<{ sha: string; version: string; dirty: boolean } | null> {
  // Read-your-writes: a just-locked run's lockedToVersionId may not have replicated. Read the owning
  // primary so a live, pinned run resolves its commit instead of silently falling back to branch head.
  const run = await runStore.findRunOnPrimary(
    { friendlyId: runFriendlyId, runtimeEnvironmentId: environmentId },
    { select: { lockedToVersionId: true } }
  );
  if (!run?.lockedToVersionId) return null;

  const deployment = await prisma.workerDeployment.findFirst({
    where: { workerId: run.lockedToVersionId },
    select: { commitSHA: true, version: true, git: true },
  });
  if (!deployment?.commitSHA) return null;

  const dirty = (deployment.git as { dirty?: boolean } | null)?.dirty ?? false;
  return { sha: deployment.commitSHA, version: deployment.version, dirty };
}