Added tool example

2026-02-21 13:42:09 +00:00
parent 15f48ff05a
commit 19539e7b32
4 changed files with 420 additions and 2 deletions
@@ -916,6 +916,189 @@ describe("TriggerChatTransport", () => {
    });
  });

+  describe("lastEventId tracking", () => {
+    it("should pass lastEventId to SSE subscription on subsequent turns", async () => {
+      const controlChunk = {
+        type: "__trigger_waitpoint_ready",
+        tokenId: "wp_token_eid",
+        publicAccessToken: "wp_access_eid",
+      };
+
+      let triggerCallCount = 0;
+      const streamFetchCalls: { url: string; headers: Record<string, string> }[] = [];
+
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+
+        if (urlStr.includes("/api/v1/tasks/") && urlStr.includes("/trigger")) {
+          triggerCallCount++;
+          return new Response(
+            JSON.stringify({ id: "run_eid" }),
+            {
+              status: 200,
+              headers: {
+                "content-type": "application/json",
+                "x-trigger-jwt": "pub_token_eid",
+              },
+            }
+          );
+        }
+
+        if (urlStr.includes("/api/v1/waitpoints/tokens/") && urlStr.includes("/complete")) {
+          return new Response(
+            JSON.stringify({ success: true }),
+            {
+              status: 200,
+              headers: { "content-type": "application/json" },
+            }
+          );
+        }
+
+        if (urlStr.includes("/realtime/v1/streams/")) {
+          streamFetchCalls.push({
+            url: urlStr,
+            headers: (init?.headers as Record<string, string>) ?? {},
+          });
+
+          const chunks = [
+            ...sampleChunks,
+            { type: "finish" as const, id: "part-1" } as UIMessageChunk,
+            controlChunk,
+          ];
+          return new Response(createSSEStream(sseEncode(chunks)), {
+            status: 200,
+            headers: {
+              "content-type": "text/event-stream",
+              "X-Stream-Version": "v1",
+            },
+          });
+        }
+
+        throw new Error(`Unexpected fetch URL: ${urlStr}`);
+      });
+
+      const transport = new TriggerChatTransport({
+        task: "my-task",
+        accessToken: "token",
+        baseURL: "https://api.test.trigger.dev",
+      });
+
+      // First message — triggers a new run
+      const stream1 = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-eid",
+        messageId: undefined,
+        messages: [createUserMessage("Hello")],
+        abortSignal: undefined,
+      });
+
+      const reader1 = stream1.getReader();
+      while (true) {
+        const { done } = await reader1.read();
+        if (done) break;
+      }
+
+      // Second message — completes the waitpoint
+      const stream2 = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-eid",
+        messageId: undefined,
+        messages: [createUserMessage("Hello"), createAssistantMessage("Hi!"), createUserMessage("What's up?")],
+        abortSignal: undefined,
+      });
+
+      const reader2 = stream2.getReader();
+      while (true) {
+        const { done } = await reader2.read();
+        if (done) break;
+      }
+
+      // The second stream subscription should include a Last-Event-ID header
+      expect(streamFetchCalls.length).toBe(2);
+      const secondStreamHeaders = streamFetchCalls[1]!.headers;
+      // SSEStreamSubscription passes lastEventId as the Last-Event-ID header
+      expect(secondStreamHeaders["Last-Event-ID"]).toBeDefined();
+    });
+  });
+
+  describe("AbortController cleanup", () => {
+    it("should terminate SSE connection after intercepting control chunk", async () => {
+      const controlChunk = {
+        type: "__trigger_waitpoint_ready",
+        tokenId: "wp_token_abort",
+        publicAccessToken: "wp_access_abort",
+      };
+
+      let streamAborted = false;
+
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+
+        if (urlStr.includes("/trigger")) {
+          return new Response(
+            JSON.stringify({ id: "run_abort_cleanup" }),
+            {
+              status: 200,
+              headers: {
+                "content-type": "application/json",
+                "x-trigger-jwt": "pub_token",
+              },
+            }
+          );
+        }
+
+        if (urlStr.includes("/realtime/v1/streams/")) {
+          // Track abort signal
+          const signal = init?.signal;
+          if (signal) {
+            signal.addEventListener("abort", () => {
+              streamAborted = true;
+            });
+          }
+
+          const chunks = [
+            ...sampleChunks,
+            { type: "finish" as const, id: "part-1" } as UIMessageChunk,
+            controlChunk,
+          ];
+          return new Response(createSSEStream(sseEncode(chunks)), {
+            status: 200,
+            headers: {
+              "content-type": "text/event-stream",
+              "X-Stream-Version": "v1",
+            },
+          });
+        }
+
+        throw new Error(`Unexpected fetch URL: ${urlStr}`);
+      });
+
+      const transport = new TriggerChatTransport({
+        task: "my-task",
+        accessToken: "token",
+        baseURL: "https://api.test.trigger.dev",
+      });
+
+      const stream = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-abort-cleanup",
+        messageId: undefined,
+        messages: [createUserMessage("Hello")],
+        abortSignal: undefined,
+      });
+
+      // Consume all chunks
+      const reader = stream.getReader();
+      while (true) {
+        const { done } = await reader.read();
+        if (done) break;
+      }
+
+      // The internal AbortController should have aborted the fetch
+      expect(streamAborted).toBe(true);
+    });
+  });
+
  describe("async accessToken", () => {
    it("should accept an async function for accessToken", async () => {
      let tokenCallCount = 0;
@@ -974,6 +1157,108 @@ describe("TriggerChatTransport", () => {

      expect(tokenCallCount).toBe(1);
    });
+
+    it("should resolve async token for waitpoint completion flow", async () => {
+      const controlChunk = {
+        type: "__trigger_waitpoint_ready",
+        tokenId: "wp_token_async",
+        publicAccessToken: "wp_access_async",
+      };
+
+      let tokenCallCount = 0;
+      let completeWaitpointCalled = false;
+
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+
+        if (urlStr.includes("/api/v1/tasks/") && urlStr.includes("/trigger")) {
+          return new Response(
+            JSON.stringify({ id: "run_async_wp" }),
+            {
+              status: 200,
+              headers: {
+                "content-type": "application/json",
+                "x-trigger-jwt": "stream-token",
+              },
+            }
+          );
+        }
+
+        if (urlStr.includes("/api/v1/waitpoints/tokens/") && urlStr.includes("/complete")) {
+          completeWaitpointCalled = true;
+          return new Response(
+            JSON.stringify({ success: true }),
+            {
+              status: 200,
+              headers: { "content-type": "application/json" },
+            }
+          );
+        }
+
+        if (urlStr.includes("/realtime/v1/streams/")) {
+          const chunks = [
+            ...sampleChunks,
+            { type: "finish" as const, id: "part-1" } as UIMessageChunk,
+            controlChunk,
+          ];
+          return new Response(createSSEStream(sseEncode(chunks)), {
+            status: 200,
+            headers: {
+              "content-type": "text/event-stream",
+              "X-Stream-Version": "v1",
+            },
+          });
+        }
+
+        throw new Error(`Unexpected fetch URL: ${urlStr}`);
+      });
+
+      const transport = new TriggerChatTransport({
+        task: "my-task",
+        accessToken: async () => {
+          tokenCallCount++;
+          await new Promise((r) => setTimeout(r, 1));
+          return `async-wp-token-${tokenCallCount}`;
+        },
+        baseURL: "https://api.test.trigger.dev",
+      });
+
+      // First message — triggers a new run (calls async token)
+      const stream1 = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-async-wp",
+        messageId: undefined,
+        messages: [createUserMessage("Hello")],
+        abortSignal: undefined,
+      });
+
+      const reader1 = stream1.getReader();
+      while (true) {
+        const { done } = await reader1.read();
+        if (done) break;
+      }
+
+      const firstTokenCount = tokenCallCount;
+
+      // Second message — should complete waitpoint (does NOT call async token)
+      const stream2 = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-async-wp",
+        messageId: undefined,
+        messages: [createUserMessage("Hello"), createAssistantMessage("Hi!"), createUserMessage("More")],
+        abortSignal: undefined,
+      });
+
+      const reader2 = stream2.getReader();
+      while (true) {
+        const { done } = await reader2.read();
+        if (done) break;
+      }
+
+      // Token function should NOT have been called again for the waitpoint path
+      expect(tokenCallCount).toBe(firstTokenCount);
+      expect(completeWaitpointCalled).toBe(true);
+    });
  });

  describe("single-run mode (waitpoint loop)", () => {
@@ -15,7 +15,8 @@
    "ai": "^6.0.0",
    "next": "15.3.3",
    "react": "^19.0.0",
-    "react-dom": "^19.0.0"
+    "react-dom": "^19.0.0",
+    "zod": "3.25.76"
  },
  "devDependencies": {
    "@tailwindcss/postcss": "^4",
@@ -5,6 +5,70 @@ import { TriggerChatTransport } from "@trigger.dev/sdk/chat";
 import { useMemo, useState } from "react";
 import { getChatToken } from "@/app/actions";

+function ToolInvocation({ part }: { part: any }) {
+  const [expanded, setExpanded] = useState(false);
+  // Static tools: type is "tool-{name}", dynamic tools have toolName property
+  const toolName =
+    part.type === "dynamic-tool"
+      ? (part.toolName ?? "tool")
+      : part.type.startsWith("tool-")
+        ? part.type.slice(5)
+        : "tool";
+  const state = part.state ?? "input-available";
+  const args = part.input;
+  const result = part.output;
+
+  const isLoading = state === "input-streaming" || state === "input-available";
+  const isError = state === "output-error";
+
+  return (
+    <div className="my-1 rounded border border-gray-200 bg-gray-50 text-xs">
+      <button
+        type="button"
+        onClick={() => setExpanded(!expanded)}
+        className="flex w-full items-center gap-2 px-3 py-1.5 text-left font-medium text-gray-700 hover:bg-gray-100"
+      >
+        {isLoading && (
+          <span className="inline-block h-3 w-3 animate-spin rounded-full border-2 border-gray-300 border-t-gray-600" />
+        )}
+        {!isLoading && !isError && <span className="text-green-600">&#10003;</span>}
+        {isError && <span className="text-red-600">&#10007;</span>}
+        <span>{toolName}</span>
+        <span className="ml-auto text-gray-400">{expanded ? "▲" : "▼"}</span>
+      </button>
+
+      {expanded && (
+        <div className="border-t border-gray-200 px-3 py-2 space-y-2">
+          {args && Object.keys(args).length > 0 && (
+            <div>
+              <div className="mb-1 font-medium text-gray-500">Input</div>
+              <pre className="overflow-x-auto whitespace-pre-wrap rounded bg-white p-2 text-gray-800">
+                {JSON.stringify(args, null, 2)}
+              </pre>
+            </div>
+          )}
+          {state === "output-available" && result !== undefined && (
+            <div>
+              <div className="mb-1 font-medium text-gray-500">Output</div>
+              <pre className="overflow-x-auto whitespace-pre-wrap rounded bg-white p-2 text-gray-800">
+                {JSON.stringify(result, null, 2)}
+              </pre>
+            </div>
+          )}
+          {isError && result !== undefined && (
+            <div>
+              <div className="mb-1 font-medium text-red-500">Error</div>
+              <pre className="overflow-x-auto whitespace-pre-wrap rounded bg-red-50 p-2 text-red-700">
+                {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+              </pre>
+            </div>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
 export function Chat() {
  const [input, setInput] = useState("");

@@ -54,6 +118,12 @@ export function Chat() {
                if (part.type === "text") {
                  return <span key={i}>{part.text}</span>;
                }
+
+                // Static tools: "tool-{toolName}", dynamic tools: "dynamic-tool"
+                if (part.type.startsWith("tool-") || part.type === "dynamic-tool") {
+                  return <ToolInvocation key={i} part={part} />;
+                }
+
                return null;
              })}
            </div>
@@ -1,6 +1,66 @@
 import { chatTask } from "@trigger.dev/sdk/ai";
-import { streamText, convertToModelMessages } from "ai";
+import { streamText, convertToModelMessages, tool } from "ai";
 import { openai } from "@ai-sdk/openai";
+import { z } from "zod";
+import os from "node:os";
+
+const inspectEnvironment = tool({
+  description:
+    "Inspect the current execution environment. Returns runtime info (Node.js/Bun/Deno version), " +
+    "OS details, CPU architecture, memory usage, environment variables, and platform metadata.",
+  inputSchema: z.object({}),
+  execute: async () => {
+    const memUsage = process.memoryUsage();
+
+    return {
+      runtime: {
+        name: typeof Bun !== "undefined" ? "bun" : typeof Deno !== "undefined" ? "deno" : "node",
+        version: process.version,
+        versions: {
+          v8: process.versions.v8,
+          openssl: process.versions.openssl,
+          modules: process.versions.modules,
+        },
+      },
+      os: {
+        platform: process.platform,
+        arch: process.arch,
+        release: os.release(),
+        type: os.type(),
+        hostname: os.hostname(),
+        uptime: `${Math.floor(os.uptime())}s`,
+      },
+      cpus: {
+        count: os.cpus().length,
+        model: os.cpus()[0]?.model,
+      },
+      memory: {
+        total: `${Math.round(os.totalmem() / 1024 / 1024)}MB`,
+        free: `${Math.round(os.freemem() / 1024 / 1024)}MB`,
+        process: {
+          rss: `${Math.round(memUsage.rss / 1024 / 1024)}MB`,
+          heapUsed: `${Math.round(memUsage.heapUsed / 1024 / 1024)}MB`,
+          heapTotal: `${Math.round(memUsage.heapTotal / 1024 / 1024)}MB`,
+        },
+      },
+      env: {
+        NODE_ENV: process.env.NODE_ENV,
+        TZ: process.env.TZ ?? Intl.DateTimeFormat().resolvedOptions().timeZone,
+        LANG: process.env.LANG,
+      },
+      process: {
+        pid: process.pid,
+        cwd: process.cwd(),
+        execPath: process.execPath,
+        argv: process.argv.slice(0, 3),
+      },
+    };
+  },
+});
+
+// Silence TS errors for Bun/Deno global checks
+declare const Bun: unknown;
+declare const Deno: unknown;

 export const chat = chatTask({
  id: "ai-chat",
@@ -9,6 +69,8 @@ export const chat = chatTask({
      model: openai("gpt-4o-mini"),
      system: "You are a helpful assistant. Be concise and friendly.",
      messages: await convertToModelMessages(messages),
+      tools: { inspectEnvironment },
+      maxSteps: 3,
    });
  },
 });