Grok AI Teammates: chat agents with their own cloud computer

import { chat } from "@trigger.dev/sdk/ai";
import { isStepCount } from "ai";
import { z } from "zod";

const clientDataSchema = z.object({ botId: z.string().optional() }).optional();

export const chatAgent = chat
  .withUIMessage<ChatUIMessage>()
  .withClientData({ schema: clientDataSchema })
  .agent({
    id: "chat",
    // Declared here so screenshots stay images when earlier turns are replayed
    tools: async ({ chatId, clientData }) => {
      const { chat, bots, bot } = await getChatBots(chatId, clientData?.botId);

      // Whoever holds the turn; a handoff moves it, and the sandbox tools with it
      let current = bot;
      const sandboxTools = createSandboxTools(async () => current.sandboxId);

      if (chat.kind !== "group") return sandboxTools;

      return {
        ...sandboxTools,
        ...createHandoffTool((botId) => {
          const target = bots.find((bot) => bot.id === botId);
          if (!target) throw new Error("No teammate in this chat has that id.");
          current = target;
          return { botId: target.id, name: target.name };
        }),
      };
    },
    // Fires for stopped and failed turns too, which spent tokens all the same
    onTurnComplete: async ({ chatId, runId, turn, usage }) => {
      const [row] = await db
        .select({ userId: chats.userId })
        .from(chats)
        .where(eq(chats.id, chatId));

      if (usage && row) {
        await recordUsage([
          {
            userId: row.userId,
            kind: "ai",
            amount: chatModelCost(usage),
            idempotencyKey: `ai:${runId}:${turn}`,
          },
        ]);
      }
    },
    run: async ({ chatId, clientData, messages, tools, signal, streamText }) => {
      const { chat: row, bots, bot } = await getChatBots(chatId, clientData?.botId);
      const group = row.kind === "group" ? { chat: row, bots } : undefined;

      const usage = await getUsage(row.userId);
      if (!usage || usage.remaining.ai <= 0) throw new Error(USAGE_LIMIT_ERROR);

      return streamText({
        model: chatModel(),
        system: buildInstructions(bot, group),
        messages,
        tools,
        // After a handoff, the rest of the turn uses the new bot's instructions
        prepareStep: ({ steps }) => {
          const handoff = steps
            .flatMap((step) => step.toolResults)
            .findLast((result) => result.toolName === "handoff");
          const current =
            handoff && bots.find((bot) => bot.id === (handoff.output as Handoff).botId);

          return current ? { instructions: buildInstructions(current, group) } : undefined;
        },
        stopWhen: [
          // Desktop work is one small action per step
          isStepCount(50),
          // A long turn ends once it has spent what was left
          ({ steps }) =>
            steps.reduce((cost, step) => cost + chatModelCost(step.usage), 0) >=
            usage.remaining.ai,
        ],
        abortSignal: signal,
      });
    },
  });

import { schedules } from "@trigger.dev/sdk";
import { inArray } from "drizzle-orm";

const INTERVAL_MINUTES = 5;

export const meterSandboxes = schedules.task({
  id: "meter-sandboxes",
  cron: `*/${INTERVAL_MINUTES} * * * *`,
  maxDuration: 240,
  run: async ({ timestamp }) => {
    const running = await listRunningSandboxIds();
    if (running.length === 0) return;

    const owned = await db
      .select({ sandboxId: bots.sandboxId, userId: bots.userId })
      .from(bots)
      .where(inArray(bots.sandboxId, running));

    await recordUsage(
      owned.map(({ sandboxId, userId }) => ({
        userId,
        kind: "sandbox" as const,
        amount: INTERVAL_MINUTES * 60,
        // A retried run meters the same moment again
        idempotencyKey: `sandbox:${sandboxId}:${timestamp.toISOString()}`,
      }))
    );

    for (const userId of new Set(owned.map((bot) => bot.userId))) {
      const usage = await getUsage(userId);
      if (usage && usage.remaining.sandbox > 0) continue;

      const sandboxIds = owned
        .filter((bot) => bot.userId === userId)
        .map((bot) => bot.sandboxId);
      await Promise.all(sandboxIds.map((id) => stopBotSandbox(id)));
    }
  },
});

import { useTriggerChatTransport, useLoadTranscript } from "@trigger.dev/sdk/chat/react";
import type { chatAgent } from "@/trigger/chat";
import { loadTranscript, mintChatAccessToken, startChatSession } from "@/actions/chat";

function Chat({ chatId, bots }: { chatId: string; bots?: ChatBot[] }) {
  const transport = useTriggerChatTransport<typeof chatAgent>({
    task: "chat",
    accessToken: ({ chatId }) => mintChatAccessToken(chatId),
    startSession: ({ chatId, clientData }) => startChatSession({ chatId, clientData }),
  });
  const { messages, isLoading } = useLoadTranscript(chatId, loadTranscript, { transport });

  // useChat only reads its initial messages on mount, so wait for the transcript
  if (isLoading) return null;

  return (
    <ChatMessages
      key={chatId}
      chatId={chatId}
      bots={bots}
      initialMessages={messages}
      transport={transport}
    />
  );
}

"use server";

import { auth } from "@trigger.dev/sdk";
import { chat, defaultStorage } from "@trigger.dev/sdk/ai";
import type { chatAgent } from "@/trigger/chat";

const startSession = chat.createStartSessionAction<typeof chatAgent>("chat");
const loadStoredTranscript = chat.createLoadTranscriptAction(defaultStorage);

export async function mintChatAccessToken(chatId: string) {
  await assertFeature("bots");
  await assertChatOwner(chatId);

  return auth.createPublicToken({
    scopes: { read: { sessions: chatId }, write: { sessions: chatId } },
    expirationTime: "1h",
  });
}

export async function loadTranscript(params: { chatId: string }) {
  await assertChatOwner(params.chatId);

  return loadStoredTranscript(params);
}