Skip to content

Commit cc21310

Browse files
committed
feat(ai): pass chat context and toolCallId to subtasks, add typed ai.chatContext helpers
- Store chat turn context (chatId, turn, continuation, clientData) in locals for auto-detection - toolFromTask now auto-detects chat context and passes it to subtask metadata - Skip serializing messages array (can be large, rarely needed by subtasks) - Tag subtask runs with toolCallId for dashboard visibility - Add ai.toolCallId() convenience helper - Add ai.chatContext<typeof myChat>() with typed clientData inference - Add ai.chatContextOrThrow<typeof myChat>() that throws if not in a chat context - Update deepResearch example to use ai.chatContextOrThrow - Document all helpers in ai-chat guide
1 parent 2cb1703 commit cc21310

File tree

3 files changed

+153
-3
lines changed

3 files changed

+153
-3
lines changed

docs/guides/ai-chat.mdx

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -884,6 +884,50 @@ The `target` option accepts:
884884
- `"root"` — root task's run (the chat task)
885885
- A specific run ID string
886886

887+
### Accessing tool context in subtasks
888+
889+
When a subtask runs via `ai.tool()`, it can access the tool call context and chat context from the parent:
890+
891+
```ts
892+
import { ai, chat } from "@trigger.dev/sdk/ai";
893+
import type { myChat } from "./chat";
894+
895+
export const mySubtask = schemaTask({
896+
id: "my-subtask",
897+
schema: z.object({ query: z.string() }),
898+
run: async ({ query }) => {
899+
// Get the AI SDK's tool call ID (useful for data-* chunk IDs)
900+
const toolCallId = ai.toolCallId();
901+
902+
// Get typed chat context — pass typeof yourChatTask for typed clientData
903+
const { chatId, clientData } = ai.chatContextOrThrow<typeof myChat>();
904+
// clientData is typed based on myChat's clientDataSchema
905+
906+
// Write a data chunk using the tool call ID
907+
const { waitUntilComplete } = chat.stream.writer({
908+
target: "root",
909+
execute: ({ write }) => {
910+
write({
911+
type: "data-progress",
912+
id: toolCallId,
913+
data: { status: "working", query, userId: clientData?.userId },
914+
});
915+
},
916+
});
917+
await waitUntilComplete();
918+
919+
return { result: "done" };
920+
},
921+
});
922+
```
923+
924+
| Helper | Returns | Description |
925+
|--------|---------|-------------|
926+
| `ai.toolCallId()` | `string \| undefined` | The AI SDK tool call ID |
927+
| `ai.chatContext<typeof myChat>()` | `{ chatId, turn, continuation, clientData } \| undefined` | Chat context with typed `clientData`. Returns `undefined` if not in a chat context. |
928+
| `ai.chatContextOrThrow<typeof myChat>()` | `{ chatId, turn, continuation, clientData }` | Same as above but throws if not in a chat context |
929+
| `ai.currentToolOptions()` | `ToolCallExecutionOptions \| undefined` | Full tool execution options |
930+
887931
## Client data and metadata
888932

889933
### Transport-level client data

packages/trigger-sdk/src/v3/ai.ts

Lines changed: 105 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,24 @@ import {
3232

3333
const METADATA_KEY = "tool.execute.options";
3434

35-
export type ToolCallExecutionOptions = Omit<ToolCallOptions, "abortSignal">;
35+
export type ToolCallExecutionOptions = {
36+
toolCallId: string;
37+
experimental_context?: unknown;
38+
/** Chat context — only present when the tool runs inside a chat.task turn. */
39+
chatId?: string;
40+
turn?: number;
41+
continuation?: boolean;
42+
clientData?: unknown;
43+
};
44+
45+
/** Chat context stored in locals during each chat.task turn for auto-detection. */
46+
type ChatTurnContext<TClientData = unknown> = {
47+
chatId: string;
48+
turn: number;
49+
continuation: boolean;
50+
clientData?: TClientData;
51+
};
52+
const chatTurnContextKey = locals.create<ChatTurnContext>("chat.turnContext");
3653

3754
type ToolResultContent = Array<
3855
| {
@@ -83,13 +100,33 @@ function toolFromTask<
83100
description: task.description,
84101
inputSchema: convertTaskSchemaToToolParameters(task),
85102
execute: async (input, options) => {
86-
const serializedOptions = options ? JSON.parse(JSON.stringify(options)) : undefined;
103+
// Build tool metadata — skip messages (can be large) and abortSignal (non-serializable)
104+
const toolMeta: ToolCallExecutionOptions = {
105+
toolCallId: options?.toolCallId ?? "",
106+
};
107+
if (options?.experimental_context !== undefined) {
108+
try {
109+
toolMeta.experimental_context = JSON.parse(JSON.stringify(options.experimental_context));
110+
} catch {
111+
// Non-serializable context — skip
112+
}
113+
}
114+
115+
// Auto-detect chat context from the parent turn
116+
const chatCtx = locals.get(chatTurnContextKey);
117+
if (chatCtx) {
118+
toolMeta.chatId = chatCtx.chatId;
119+
toolMeta.turn = chatCtx.turn;
120+
toolMeta.continuation = chatCtx.continuation;
121+
toolMeta.clientData = chatCtx.clientData;
122+
}
87123

88124
return await task
89125
.triggerAndWait(input as inferSchemaIn<TTaskSchema>, {
90126
metadata: {
91-
[METADATA_KEY]: serializedOptions,
127+
[METADATA_KEY]: toolMeta as any,
92128
},
129+
tags: options?.toolCallId ? [`toolCallId:${options.toolCallId}`] : undefined,
93130
})
94131
.unwrap();
95132
},
@@ -109,6 +146,57 @@ function getToolOptionsFromMetadata(): ToolCallExecutionOptions | undefined {
109146
return tool as ToolCallExecutionOptions;
110147
}
111148

149+
/**
150+
* Get the current tool call ID from inside a subtask invoked via `ai.tool()`.
151+
* Returns `undefined` if not running as a tool subtask.
152+
*/
153+
function getToolCallId(): string | undefined {
154+
return getToolOptionsFromMetadata()?.toolCallId;
155+
}
156+
157+
/**
158+
* Get the chat context from inside a subtask invoked via `ai.tool()` within a `chat.task`.
159+
* Pass `typeof yourChatTask` as the type parameter to get typed `clientData`.
160+
* Returns `undefined` if the parent is not a chat task.
161+
*
162+
* @example
163+
* ```ts
164+
* const ctx = ai.chatContext<typeof myChat>();
165+
* // ctx?.clientData is typed based on myChat's clientDataSchema
166+
* ```
167+
*/
168+
function getToolChatContext<TChatTask extends AnyTask = AnyTask>(): ChatTurnContext<InferChatClientData<TChatTask>> | undefined {
169+
const opts = getToolOptionsFromMetadata();
170+
if (!opts?.chatId) return undefined;
171+
return {
172+
chatId: opts.chatId,
173+
turn: opts.turn ?? 0,
174+
continuation: opts.continuation ?? false,
175+
clientData: opts.clientData as InferChatClientData<TChatTask>,
176+
};
177+
}
178+
179+
/**
180+
* Get the chat context from inside a subtask, throwing if not in a chat context.
181+
* Pass `typeof yourChatTask` as the type parameter to get typed `clientData`.
182+
*
183+
* @example
184+
* ```ts
185+
* const ctx = ai.chatContextOrThrow<typeof myChat>();
186+
* // ctx.chatId, ctx.clientData are guaranteed non-null
187+
* ```
188+
*/
189+
function getToolChatContextOrThrow<TChatTask extends AnyTask = AnyTask>(): ChatTurnContext<InferChatClientData<TChatTask>> {
190+
const ctx = getToolChatContext<TChatTask>();
191+
if (!ctx) {
192+
throw new Error(
193+
"ai.chatContextOrThrow() called outside of a chat.task context. " +
194+
"This helper can only be used inside a subtask invoked via ai.tool() from a chat.task."
195+
);
196+
}
197+
return ctx;
198+
}
199+
112200
function convertTaskSchemaToToolParameters(
113201
task: AnyTask | TaskWithSchema<any, any, any>
114202
): Schema<unknown> {
@@ -136,6 +224,12 @@ function convertTaskSchemaToToolParameters(
136224
export const ai = {
137225
tool: toolFromTask,
138226
currentToolOptions: getToolOptionsFromMetadata,
227+
/** Get the tool call ID from inside a subtask invoked via `ai.tool()`. */
228+
toolCallId: getToolCallId,
229+
/** Get chat context (chatId, turn, clientData, etc.) from inside a subtask of a `chat.task`. Returns undefined if not in a chat context. */
230+
chatContext: getToolChatContext,
231+
/** Get chat context or throw if not in a chat context. Pass `typeof yourChatTask` for typed clientData. */
232+
chatContextOrThrow: getToolChatContextOrThrow,
139233
};
140234

141235
/**
@@ -756,6 +850,14 @@ function chatTask<
756850
async () => {
757851
locals.set(chatPipeCountKey, 0);
758852

853+
// Store chat context for auto-detection by ai.tool subtasks
854+
locals.set(chatTurnContextKey, {
855+
chatId: currentWirePayload.chatId,
856+
turn,
857+
continuation,
858+
clientData,
859+
});
860+
759861
// Per-turn stop controller (reset each turn)
760862
const stopController = new AbortController();
761863
currentStopController = stopController;

references/ai-chat/src/trigger/chat.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,10 @@ export const deepResearch = schemaTask({
150150
urls: z.array(z.string().url()).describe("URLs to fetch and analyze"),
151151
}),
152152
run: async ({ query, urls }) => {
153+
// Access chat context from the parent chat.task — typed via typeof aiChat
154+
const { chatId, clientData } = ai.chatContextOrThrow<typeof aiChat>();
155+
console.log(`Deep research for chat ${chatId}, user ${clientData?.userId}`);
156+
153157
const partId = generateId();
154158
const results: { url: string; status: number; snippet: string }[] = [];
155159

0 commit comments

Comments
 (0)