chore(references): demo chat.headStart in ai-chat reference
Adds a /api/chat route handler exporting chat.headStart, splits the tool definitions across two modules so heavy executes never reach the browser bundle, and exposes a sidebar toggle for paired TTFC tests. - src/lib/chat-tools-schemas.ts (new): schema-only tool definitions — imported by both the route handler and the agent task. No `execute`, no heavy deps. Bundle stays small. - src/trigger/chat-tools.ts (renamed): re-exports the schemas with agent-side `execute` fns added (E2B sandbox, turndown, deepResearch subtask, etc.). Only the trigger task imports this. - src/app/api/chat/route.ts (new): exports POST = chat.headStart, runs step 1 streamText with claude-sonnet-4-6 to match the agent's default. - ChatSettingsContext + sidebar gain a "Use handover (1st turn)" toggle; chat-view threads it into the transport's `headStart` URL. - Smoke result: ~53% TTFC reduction on first turn (1561ms vs 3358ms), with persistence + tool execution behaving identically.
This commit is contained in:
Eric Allam
@@ -9,7 +9,7 @@ import type {
|
||||
aiChatSession,
|
||||
upgradeTestAgent,
|
||||
} from "@/trigger/chat";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools-schemas";
|
||||
import { prisma } from "@/lib/prisma";
|
||||
|
||||
/** Short-lived PATs for local testing of expiry + renewal (not for production). */
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* chat.headStart first-turn endpoint.
|
||||
*
|
||||
* The browser transport POSTs first-turn messages here when the
|
||||
* `headStart` option is set on `useTriggerChatTransport`. This
|
||||
* handler:
|
||||
*
|
||||
* 1. Creates the chat.agent session and triggers a `handover-prepare`
|
||||
* run (atomic, one round-trip), so the agent boots in parallel.
|
||||
* 2. Runs `streamText` step 1 right here in the warm Next.js process
|
||||
* and returns the SSE stream directly to the browser — no waiting
|
||||
* on the agent's cold start.
|
||||
* 3. On step 1's tool-call boundary, hands ownership of the durable
|
||||
* session.out stream over to the agent run, which executes tools
|
||||
* and continues from step 2+ (or exits clean for pure-text turns).
|
||||
*
|
||||
* Subsequent turns bypass this endpoint — the transport hydrates the
|
||||
* session PAT from response headers and writes directly to
|
||||
* `session.in` for turn 2 onward.
|
||||
*
|
||||
* The TTFC win: cold-start agent boot (~488ms) + onTurnStart hooks
|
||||
* (~316ms) overlap with the LLM TTFB instead of stacking before it,
|
||||
* so the user-perceived first chunk arrives ~50% sooner. The agent
|
||||
* still owns tool execution and everything after — heavy deps stay
|
||||
* where they belong.
|
||||
*/
|
||||
import { chat } from "@trigger.dev/sdk/chat-server";
|
||||
import { streamText } from "ai";
|
||||
import { anthropic } from "@ai-sdk/anthropic";
|
||||
// ⚠️ Imports MUST come from `chat-tools-schemas` only — see the
|
||||
// header comment in that file for the bundle-isolation rationale.
|
||||
// Importing `src/trigger/chat-tools.ts` here would drag E2B,
|
||||
// turndown, the trigger SDK runtime, etc. into the Next.js bundle
|
||||
// and defeat the whole point of `chat.headStart`.
|
||||
import { headStartTools } from "@/lib/chat-tools-schemas";
|
||||
|
||||
export const POST = chat.headStart({
|
||||
agentId: "ai-chat",
|
||||
run: async ({ chat: chatHelper }) => {
|
||||
return streamText({
|
||||
// `toStreamTextOptions` wires `messages` (converted from
|
||||
// UIMessages), `tools`, `stopWhen: stepCountIs(1)`, and the
|
||||
// combined `abortSignal`. Customer adds model + system prompt on
|
||||
// top — anything else `streamText` accepts is fair game.
|
||||
...chatHelper.toStreamTextOptions({ tools: headStartTools }),
|
||||
// Match the agent's default (`DEFAULT_MODEL` in `lib/models.ts`)
|
||||
// so step 1 and step 2+ run on the same provider — no jarring
|
||||
// tone/style shift mid-turn, and TTFC comparisons stay honest.
|
||||
model: anthropic("claude-sonnet-4-6"),
|
||||
system:
|
||||
"You are a helpful AI assistant. Be concise and friendly. Use the available tools when relevant.",
|
||||
});
|
||||
},
|
||||
});
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
import { generateId } from "ai";
|
||||
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools-schemas";
|
||||
import { useCallback, useEffect, useState } from "react";
|
||||
import { Chat } from "@/components/chat";
|
||||
import { ChatSidebar } from "@/components/chat-sidebar";
|
||||
@@ -169,6 +169,8 @@ export function ChatApp({
|
||||
onIdleTimeoutChange={setIdleTimeoutInSeconds}
|
||||
taskMode={taskMode}
|
||||
onTaskModeChange={onTaskModeChange}
|
||||
useHandover={false}
|
||||
onUseHandoverChange={() => {}}
|
||||
/>
|
||||
<div className="flex-1">
|
||||
{activeChatId ? (
|
||||
|
||||
@@ -7,6 +7,13 @@ type ChatSettings = {
|
||||
setTaskMode: (mode: string) => void;
|
||||
idleTimeoutInSeconds: number;
|
||||
setIdleTimeoutInSeconds: (seconds: number) => void;
|
||||
/**
|
||||
* When true, first-turn messages are POSTed to `/api/chat`
|
||||
* (`chat.handover` route handler) instead of triggering the agent
|
||||
* directly. Subsequent turns bypass the endpoint regardless.
|
||||
*/
|
||||
useHandover: boolean;
|
||||
setUseHandover: (on: boolean) => void;
|
||||
};
|
||||
|
||||
const ChatSettingsContext = createContext<ChatSettings | null>(null);
|
||||
@@ -14,12 +21,15 @@ const ChatSettingsContext = createContext<ChatSettings | null>(null);
|
||||
export function ChatSettingsProvider({ children }: { children: ReactNode }) {
|
||||
const [taskMode, setTaskMode] = useState("ai-chat");
|
||||
const [idleTimeoutInSeconds, setIdleTimeoutInSeconds] = useState(60);
|
||||
const [useHandover, setUseHandover] = useState(false);
|
||||
|
||||
const value: ChatSettings = {
|
||||
taskMode,
|
||||
setTaskMode,
|
||||
idleTimeoutInSeconds,
|
||||
setIdleTimeoutInSeconds,
|
||||
useHandover,
|
||||
setUseHandover,
|
||||
};
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
|
||||
@@ -28,6 +28,8 @@ export function ChatSidebarWrapper({
|
||||
setTaskMode,
|
||||
idleTimeoutInSeconds,
|
||||
setIdleTimeoutInSeconds,
|
||||
useHandover,
|
||||
setUseHandover,
|
||||
} = useChatSettings();
|
||||
|
||||
// Extract active chatId from URL
|
||||
@@ -85,6 +87,8 @@ export function ChatSidebarWrapper({
|
||||
onIdleTimeoutChange={setIdleTimeoutInSeconds}
|
||||
taskMode={taskMode}
|
||||
onTaskModeChange={setTaskMode}
|
||||
useHandover={useHandover}
|
||||
onUseHandoverChange={setUseHandover}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -29,6 +29,8 @@ type ChatSidebarProps = {
|
||||
onIdleTimeoutChange: (seconds: number) => void;
|
||||
taskMode: string;
|
||||
onTaskModeChange: (mode: string) => void;
|
||||
useHandover: boolean;
|
||||
onUseHandoverChange: (on: boolean) => void;
|
||||
};
|
||||
|
||||
export function ChatSidebar({
|
||||
@@ -42,6 +44,8 @@ export function ChatSidebar({
|
||||
onIdleTimeoutChange,
|
||||
taskMode,
|
||||
onTaskModeChange,
|
||||
useHandover,
|
||||
onUseHandoverChange,
|
||||
}: ChatSidebarProps) {
|
||||
const sorted = [...chats].sort((a, b) => b.updatedAt - a.updatedAt);
|
||||
|
||||
@@ -115,6 +119,18 @@ export function ChatSidebar({
|
||||
<option value="upgrade-test">upgrade-test (requestUpgrade after 3 turns)</option>
|
||||
</select>
|
||||
</div>
|
||||
<label
|
||||
className="flex items-center gap-2 text-xs text-gray-500"
|
||||
title="Route first-turn messages through /api/chat (chat.handover) so step 1 streams from the Next.js process while the agent run boots in parallel."
|
||||
>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={useHandover}
|
||||
onChange={(e) => onUseHandoverChange(e.target.checked)}
|
||||
className="h-3 w-3 rounded border-gray-300"
|
||||
/>
|
||||
<span>Use handover (1st turn)</span>
|
||||
</label>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onWipeAll}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"use client";
|
||||
|
||||
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools-schemas";
|
||||
import { Chat } from "@/components/chat";
|
||||
import { useChatSettings } from "@/components/chat-settings-context";
|
||||
import {
|
||||
@@ -34,7 +34,7 @@ export function ChatView({
|
||||
model,
|
||||
}: ChatViewProps) {
|
||||
const router = useRouter();
|
||||
const { taskMode } = useChatSettings();
|
||||
const { taskMode, useHandover } = useChatSettings();
|
||||
|
||||
const [currentSession, setCurrentSession] = useState<SessionInfo | null>(initialSession);
|
||||
|
||||
@@ -67,6 +67,15 @@ export function ChatView({
|
||||
onSessionChange: handleSessionChange,
|
||||
clientData: { userId: "user_123" },
|
||||
multiTab: true,
|
||||
// Head-start URL: opt-in fast-path for the first message of a
|
||||
// brand-new chat. The transport POSTs to `/api/chat` (which
|
||||
// exports `chat.handover({ agentId, run })`) so step 1's LLM
|
||||
// call runs in the warm Next.js process while the trigger agent
|
||||
// run boots in parallel. After turn 1 the transport hydrates
|
||||
// session state from response headers and writes directly to
|
||||
// `session.in` for turn 2 onward — same direct-trigger path as
|
||||
// when `headStart` is unset.
|
||||
headStart: useHandover ? "/api/chat" : undefined,
|
||||
});
|
||||
|
||||
const handleFirstMessage = useCallback(
|
||||
@@ -101,6 +110,7 @@ export function ChatView({
|
||||
projectDashboardPath={process.env.NEXT_PUBLIC_TRIGGER_PROJECT_DASHBOARD_PATH}
|
||||
onFirstMessage={handleFirstMessage}
|
||||
onMessagesChange={handleMessagesChange}
|
||||
handoverEnabled={useHandover}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ import {
|
||||
lastAssistantMessageIsCompleteWithApprovalResponses,
|
||||
lastAssistantMessageIsCompleteWithToolCalls,
|
||||
} from "ai";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools-schemas";
|
||||
import type { TriggerChatTransport } from "@trigger.dev/sdk/chat";
|
||||
|
||||
// Structural type mirroring @trigger.dev/sdk/ai's CompactionChunkData.
|
||||
@@ -323,6 +323,8 @@ type ChatProps = {
|
||||
projectDashboardPath?: string;
|
||||
onFirstMessage?: (chatId: string, text: string) => void;
|
||||
onMessagesChange?: (chatId: string, messages: ChatUiMessage[]) => void;
|
||||
/** Whether the transport is configured to route first-turn through `chat.handover`. */
|
||||
handoverEnabled?: boolean;
|
||||
};
|
||||
|
||||
export function Chat({
|
||||
@@ -338,6 +340,7 @@ export function Chat({
|
||||
projectDashboardPath,
|
||||
onFirstMessage,
|
||||
onMessagesChange,
|
||||
handoverEnabled = false,
|
||||
}: ChatProps) {
|
||||
const [input, setInput] = useState("");
|
||||
const hasCalledFirstMessage = useRef(false);
|
||||
@@ -552,6 +555,8 @@ export function Chat({
|
||||
return transport.getSession(chatId)?.lastEventId ?? null;
|
||||
},
|
||||
chatId,
|
||||
/** True when the transport is configured to route first-turn through `chat.handover`. */
|
||||
handoverEnabled,
|
||||
|
||||
// ── Actions ───────────────────────────────────────────────────
|
||||
steer: (text: string) => actionsRef.current.steer(text),
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
/**
|
||||
* Schema-only tool definitions — shared between the chat.handover
|
||||
* route handler and the trigger.dev agent task.
|
||||
*
|
||||
* ⚠️ HARD CONSTRAINT — bundle isolation
|
||||
*
|
||||
* This file is imported by `app/api/chat/route.ts` (the chat.handover
|
||||
* POST handler) and runs in the Next.js process. Anything imported
|
||||
* here lands in the route-handler bundle.
|
||||
*
|
||||
* Allowed imports: `ai` (for `tool()`), `zod`, type-only AI SDK
|
||||
* imports. Nothing else.
|
||||
*
|
||||
* DO NOT import from this file:
|
||||
* - `@e2b/code-interpreter`, `puppeteer`, `playwright`, native bindings
|
||||
* - `node:child_process`, heavy filesystem ops
|
||||
* - `@trigger.dev/sdk` runtime (`task`, `schemaTask`,
|
||||
* `chat.stream.writer`, etc. — pulls in the whole task runtime)
|
||||
* - `turndown`, image processing libs, anything that pulls weight
|
||||
*
|
||||
* Heavy `execute` fns live in `src/trigger/chat-tools.ts` — that file
|
||||
* imports these schemas and adds executes on top. The agent task
|
||||
* picks up the executes when it runs; the route handler never sees
|
||||
* them and never imports their deps.
|
||||
*
|
||||
* If you need to add a new tool to the chat.agent's schema-only set,
|
||||
* declare its description + inputSchema here, then wire its execute
|
||||
* fn in `src/trigger/chat-tools.ts`.
|
||||
*/
|
||||
import { tool } from "ai";
|
||||
import type { InferUITools, UIDataTypes, UIMessage } from "ai";
|
||||
import { z } from "zod";
|
||||
|
||||
export 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 → src/trigger/chat-tools.ts
|
||||
});
|
||||
|
||||
export const webFetch = tool({
|
||||
description:
|
||||
"Fetch a URL and return the response as text. " +
|
||||
"Use this to retrieve web pages, APIs, or any HTTP resource.",
|
||||
inputSchema: z.object({
|
||||
url: z.string().url().describe("The URL to fetch"),
|
||||
}),
|
||||
// execute → src/trigger/chat-tools.ts (uses turndown)
|
||||
});
|
||||
|
||||
export const deepResearch = tool({
|
||||
description:
|
||||
"Research a topic by fetching multiple URLs and synthesizing the results. " +
|
||||
"Streams progress updates to the chat as it works.",
|
||||
inputSchema: z.object({
|
||||
query: z.string().describe("The research query or topic"),
|
||||
urls: z.array(z.string().url()).describe("URLs to fetch and analyze"),
|
||||
}),
|
||||
// execute → src/trigger/chat-tools.ts (subtask via ai.toolExecute)
|
||||
});
|
||||
|
||||
export const posthogQuery = tool({
|
||||
description:
|
||||
"Query PostHog analytics using HogQL. Use this to answer questions about events, " +
|
||||
"pageviews, user activity, feature flag usage, or any product analytics question. " +
|
||||
"Write a HogQL query (SQL-like syntax over PostHog events).",
|
||||
inputSchema: z.object({
|
||||
query: z
|
||||
.string()
|
||||
.describe(
|
||||
"HogQL query, e.g. SELECT event, count() FROM events WHERE timestamp > now() - interval 1 day GROUP BY event ORDER BY count() DESC LIMIT 10"
|
||||
),
|
||||
}),
|
||||
// execute → src/trigger/chat-tools.ts (HTTP to PostHog)
|
||||
});
|
||||
|
||||
export const executeCode = tool({
|
||||
description:
|
||||
"Run code in an isolated E2B sandbox (Python by default; other languages supported by E2B). " +
|
||||
"Use for calculations, data analysis, or transforming tool outputs (e.g. PostHog query results). " +
|
||||
"The sandbox persists across turns in the same run until the chat idles and suspends.",
|
||||
inputSchema: z.object({
|
||||
code: z.string().describe("Source code to execute in the sandbox"),
|
||||
language: z
|
||||
.string()
|
||||
.optional()
|
||||
.describe("Language id (e.g. python, javascript). Defaults to python."),
|
||||
}),
|
||||
// execute → src/trigger/chat-tools.ts (E2B sandbox — heavy native dep)
|
||||
});
|
||||
|
||||
export const sendEmail = tool({
|
||||
description:
|
||||
"Send an email to a recipient. Requires human approval before sending. " +
|
||||
"Use when the user asks you to send, draft, or compose an email.",
|
||||
inputSchema: z.object({
|
||||
to: z.string().describe("Recipient email address"),
|
||||
subject: z.string().describe("Email subject line"),
|
||||
body: z.string().describe("Email body text"),
|
||||
}),
|
||||
needsApproval: true,
|
||||
// execute → src/trigger/chat-tools.ts
|
||||
});
|
||||
|
||||
export const askUser = tool({
|
||||
description:
|
||||
"Ask the user a question when you need clarification or input before proceeding. " +
|
||||
"Present 2-4 options for the user to choose from. Use when uncertain about the user's intent.",
|
||||
inputSchema: z.object({
|
||||
question: z.string().describe("The question to ask the user"),
|
||||
options: z
|
||||
.array(
|
||||
z.object({
|
||||
id: z.string().describe("Unique option identifier"),
|
||||
label: z.string().describe("Short option title"),
|
||||
description: z.string().optional().describe("Longer explanation"),
|
||||
})
|
||||
)
|
||||
.min(2)
|
||||
.max(4),
|
||||
}),
|
||||
// No execute by design — round-tripped through the frontend's addToolOutput.
|
||||
});
|
||||
|
||||
/**
|
||||
* The schema-only tool set passed to `chat.headStart`'s `streamText`
|
||||
* call. The agent task imports each schema individually and adds the
|
||||
* matching `execute` fn — see `src/trigger/chat-tools.ts`.
|
||||
*/
|
||||
export const headStartTools = {
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
deepResearch,
|
||||
posthogQuery,
|
||||
executeCode,
|
||||
sendEmail,
|
||||
askUser,
|
||||
};
|
||||
|
||||
type ChatToolSet = typeof headStartTools;
|
||||
export type ChatUiTools = InferUITools<ChatToolSet>;
|
||||
export type ChatUiMessage = UIMessage<unknown, UIDataTypes, ChatUiTools>;
|
||||
@@ -0,0 +1,297 @@
|
||||
/**
|
||||
* Tool executes for the trigger.dev agent task.
|
||||
*
|
||||
* These tools wrap the schema-only definitions from
|
||||
* `@/lib/chat-tools-schemas` with their heavy `execute` fns. This
|
||||
* file is ONLY imported from inside the trigger task module
|
||||
* (`src/trigger/chat.ts`); it must NOT be imported from anything that
|
||||
* runs in the Next.js process (route handlers, components, server
|
||||
* actions, etc.).
|
||||
*
|
||||
* See `src/lib/chat-tools-schemas.ts` for why this split matters —
|
||||
* the bundle-isolation constraint is what makes `chat.handover`'s
|
||||
* cold-start win possible.
|
||||
*/
|
||||
import { ai, chat } from "@trigger.dev/sdk/ai";
|
||||
import { schemaTask } from "@trigger.dev/sdk";
|
||||
import { tool, generateId } from "ai";
|
||||
import { z } from "zod";
|
||||
import os from "node:os";
|
||||
import TurndownService from "turndown";
|
||||
import { codeSandboxRun, runWithCodeSandbox } from "@/lib/code-sandbox";
|
||||
import {
|
||||
inspectEnvironment as inspectEnvironmentSchema,
|
||||
webFetch as webFetchSchema,
|
||||
deepResearch as deepResearchSchema,
|
||||
posthogQuery as posthogQuerySchema,
|
||||
executeCode as executeCodeSchema,
|
||||
sendEmail as sendEmailSchema,
|
||||
askUser as askUserSchema,
|
||||
} from "@/lib/chat-tools-schemas";
|
||||
|
||||
const turndown = new TurndownService();
|
||||
|
||||
declare const Bun: unknown;
|
||||
declare const Deno: unknown;
|
||||
|
||||
export const inspectEnvironment = tool({
|
||||
...inspectEnvironmentSchema,
|
||||
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),
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
export const webFetch = tool({
|
||||
...webFetchSchema,
|
||||
execute: async ({ url }) => {
|
||||
const latency = Number(process.env.WEBFETCH_LATENCY_MS);
|
||||
if (latency > 0) {
|
||||
await new Promise((r) => setTimeout(r, latency));
|
||||
}
|
||||
|
||||
const response = await fetch(url);
|
||||
let text = await response.text();
|
||||
const contentType = response.headers.get("content-type") ?? "";
|
||||
|
||||
if (contentType.includes("html")) {
|
||||
text = turndown.turndown(text);
|
||||
}
|
||||
|
||||
return {
|
||||
status: response.status,
|
||||
contentType,
|
||||
body: text.slice(0, 2000),
|
||||
truncated: text.length > 2000,
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
const deepResearchTask = schemaTask({
|
||||
id: "deep-research",
|
||||
description:
|
||||
"Research a topic by fetching multiple URLs and synthesizing the results. " +
|
||||
"Streams progress updates to the chat as it works.",
|
||||
schema: z.object({
|
||||
query: z.string().describe("The research query or topic"),
|
||||
urls: z.array(z.string().url()).describe("URLs to fetch and analyze"),
|
||||
}),
|
||||
run: async ({ query, urls }) => {
|
||||
const partId = generateId();
|
||||
const results: { url: string; status: number; snippet: string }[] = [];
|
||||
|
||||
for (let i = 0; i < urls.length; i++) {
|
||||
const url = urls[i]!;
|
||||
|
||||
const { waitUntilComplete } = chat.stream.writer({
|
||||
target: "root",
|
||||
execute: ({ write }) => {
|
||||
write({
|
||||
type: "data-research-progress",
|
||||
id: partId,
|
||||
data: {
|
||||
status: "fetching" as const,
|
||||
query,
|
||||
current: i + 1,
|
||||
total: urls.length,
|
||||
currentUrl: url,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
await waitUntilComplete();
|
||||
|
||||
try {
|
||||
const response = await fetch(url);
|
||||
let text = await response.text();
|
||||
const contentType = response.headers.get("content-type") ?? "";
|
||||
|
||||
if (contentType.includes("html")) {
|
||||
text = turndown.turndown(text);
|
||||
}
|
||||
|
||||
results.push({
|
||||
url,
|
||||
status: response.status,
|
||||
snippet: text.slice(0, 500),
|
||||
});
|
||||
} catch (err) {
|
||||
results.push({
|
||||
url,
|
||||
status: 0,
|
||||
snippet: `Error: ${err instanceof Error ? err.message : String(err)}`,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const { waitUntilComplete: waitForDone } = chat.stream.writer({
|
||||
target: "root",
|
||||
execute: ({ write }) => {
|
||||
write({
|
||||
type: "data-research-progress",
|
||||
id: partId,
|
||||
data: {
|
||||
status: "done" as const,
|
||||
query,
|
||||
current: urls.length,
|
||||
total: urls.length,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
await waitForDone();
|
||||
|
||||
return { query, results };
|
||||
},
|
||||
});
|
||||
|
||||
/** Task-backed tool: AI SDK `tool()` for shape/types; `ai.toolExecute` for Trigger subtask + metadata. */
|
||||
export const deepResearch = tool({
|
||||
...deepResearchSchema,
|
||||
execute: ai.toolExecute(deepResearchTask),
|
||||
});
|
||||
|
||||
const POSTHOG_API_KEY = process.env.POSTHOG_API_KEY;
|
||||
const POSTHOG_PROJECT_ID = process.env.POSTHOG_PROJECT_ID;
|
||||
const POSTHOG_HOST = process.env.POSTHOG_HOST ?? "https://eu.posthog.com";
|
||||
|
||||
export const posthogQuery = tool({
|
||||
...posthogQuerySchema,
|
||||
execute: async ({ query }) => {
|
||||
if (!POSTHOG_API_KEY || !POSTHOG_PROJECT_ID) {
|
||||
return { error: "PostHog not configured. Set POSTHOG_API_KEY and POSTHOG_PROJECT_ID." };
|
||||
}
|
||||
const response = await fetch(`${POSTHOG_HOST}/api/projects/${POSTHOG_PROJECT_ID}/query/`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${POSTHOG_API_KEY}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({ query: { kind: "HogQLQuery", query } }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
return { error: `PostHog API error ${response.status}: ${text.slice(0, 500)}` };
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
return {
|
||||
columns: data.columns,
|
||||
results: data.results?.slice(0, 50),
|
||||
rowCount: data.results?.length ?? 0,
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
export const executeCode = tool({
|
||||
...executeCodeSchema,
|
||||
execute: async function executeCodeExecute({ code, language }) {
|
||||
const runId = codeSandboxRun.runId;
|
||||
if (!runId?.trim()) {
|
||||
return {
|
||||
error:
|
||||
"Code sandbox run id is not set yet (call from the chat task after onTurnStart), or this tool is not wired to that task.",
|
||||
};
|
||||
}
|
||||
|
||||
const out = await runWithCodeSandbox(runId, async function runInSandbox(sandbox) {
|
||||
const execution = await sandbox.runCode(code, {
|
||||
...(language?.trim() ? { language: language.trim() } : {}),
|
||||
timeoutMs: 60_000,
|
||||
});
|
||||
|
||||
if (execution.error) {
|
||||
return {
|
||||
error: `${execution.error.name}: ${execution.error.value}`,
|
||||
traceback: execution.error.traceback,
|
||||
stdout: execution.logs.stdout.join("\n"),
|
||||
stderr: execution.logs.stderr.join("\n"),
|
||||
};
|
||||
}
|
||||
|
||||
const mainText = execution.text;
|
||||
const resultSnippets = execution.results
|
||||
.map(function mapResult(r) {
|
||||
return r.text ?? r.markdown ?? r.json;
|
||||
})
|
||||
.filter(Boolean)
|
||||
.slice(0, 5);
|
||||
|
||||
return {
|
||||
text: mainText,
|
||||
results: resultSnippets,
|
||||
stdout: execution.logs.stdout.join("\n"),
|
||||
stderr: execution.logs.stderr.join("\n"),
|
||||
};
|
||||
});
|
||||
|
||||
return out;
|
||||
},
|
||||
});
|
||||
|
||||
export const sendEmail = tool({
|
||||
...sendEmailSchema,
|
||||
execute: async ({ to, subject, body }) => {
|
||||
// Simulated — in a real app this would call an email API
|
||||
return { sent: true, to, subject, preview: body.slice(0, 100) };
|
||||
},
|
||||
});
|
||||
|
||||
// askUser has no execute by design — round-tripped via addToolOutput.
|
||||
export const askUser = askUserSchema;
|
||||
|
||||
/** Tool set passed to `streamText` for the main `chat.agent` run. */
|
||||
export const chatTools = {
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
deepResearch,
|
||||
posthogQuery,
|
||||
executeCode,
|
||||
sendEmail,
|
||||
askUser,
|
||||
};
|
||||
@@ -21,8 +21,8 @@ import {
|
||||
deepResearch,
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
type ChatUiMessage,
|
||||
} from "@/lib/chat-tools";
|
||||
} from "./chat-tools";
|
||||
import type { ChatUiMessage } from "@/lib/chat-tools-schemas";
|
||||
import { disposeCodeSandboxForRun, warmCodeSandbox } from "@/lib/code-sandbox";
|
||||
|
||||
const adapter = new PrismaPg({ connectionString: process.env.DATABASE_URL! });
|
||||
@@ -232,7 +232,15 @@ export const aiChat = chat
|
||||
turn,
|
||||
count: messages.length,
|
||||
});
|
||||
return validateUIMessages({ messages, tools: chatTools });
|
||||
// Cast: `chatTools` has executes (output types are real), but
|
||||
// `ChatUiMessage` is derived from the schema-only set in
|
||||
// `chat-tools-schemas.ts` so its tools have `output: never`.
|
||||
// `validateUIMessages` only reads `inputSchema` at runtime, so
|
||||
// the type narrowing is safely sidestepped.
|
||||
return validateUIMessages({
|
||||
messages,
|
||||
tools: chatTools as unknown as Parameters<typeof validateUIMessages>[0]["tools"],
|
||||
});
|
||||
},
|
||||
// #endregion
|
||||
|
||||
|
||||
Reference in New Issue
Block a user