diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c2ffd55be..47cad62b5 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -32,7 +32,24 @@ import { const METADATA_KEY = "tool.execute.options"; -export type ToolCallExecutionOptions = Omit; +export type ToolCallExecutionOptions = { + toolCallId: string; + experimental_context?: unknown; + /** Chat context — only present when the tool runs inside a chat.task turn. */ + chatId?: string; + turn?: number; + continuation?: boolean; + clientData?: unknown; +}; + +/** Chat context stored in locals during each chat.task turn for auto-detection. */ +type ChatTurnContext = { + chatId: string; + turn: number; + continuation: boolean; + clientData?: TClientData; +}; +const chatTurnContextKey = locals.create("chat.turnContext"); type ToolResultContent = Array< | { @@ -83,13 +100,33 @@ function toolFromTask< description: task.description, inputSchema: convertTaskSchemaToToolParameters(task), execute: async (input, options) => { - const serializedOptions = options ? JSON.parse(JSON.stringify(options)) : undefined; + // Build tool metadata — skip messages (can be large) and abortSignal (non-serializable) + const toolMeta: ToolCallExecutionOptions = { + toolCallId: options?.toolCallId ?? "", + }; + if (options?.experimental_context !== undefined) { + try { + toolMeta.experimental_context = JSON.parse(JSON.stringify(options.experimental_context)); + } catch { + // Non-serializable context — skip + } + } + + // Auto-detect chat context from the parent turn + const chatCtx = locals.get(chatTurnContextKey); + if (chatCtx) { + toolMeta.chatId = chatCtx.chatId; + toolMeta.turn = chatCtx.turn; + toolMeta.continuation = chatCtx.continuation; + toolMeta.clientData = chatCtx.clientData; + } return await task .triggerAndWait(input as inferSchemaIn, { metadata: { - [METADATA_KEY]: serializedOptions, + [METADATA_KEY]: toolMeta as any, }, + tags: options?.toolCallId ? [`toolCallId:${options.toolCallId}`] : undefined, }) .unwrap(); }, @@ -109,6 +146,57 @@ function getToolOptionsFromMetadata(): ToolCallExecutionOptions | undefined { return tool as ToolCallExecutionOptions; } +/** + * Get the current tool call ID from inside a subtask invoked via `ai.tool()`. + * Returns `undefined` if not running as a tool subtask. + */ +function getToolCallId(): string | undefined { + return getToolOptionsFromMetadata()?.toolCallId; +} + +/** + * Get the chat context from inside a subtask invoked via `ai.tool()` within a `chat.task`. + * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`. + * Returns `undefined` if the parent is not a chat task. + * + * @example + * ```ts + * const ctx = ai.chatContext(); + * // ctx?.clientData is typed based on myChat's clientDataSchema + * ``` + */ +function getToolChatContext(): ChatTurnContext> | undefined { + const opts = getToolOptionsFromMetadata(); + if (!opts?.chatId) return undefined; + return { + chatId: opts.chatId, + turn: opts.turn ?? 0, + continuation: opts.continuation ?? false, + clientData: opts.clientData as InferChatClientData, + }; +} + +/** + * Get the chat context from inside a subtask, throwing if not in a chat context. + * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`. + * + * @example + * ```ts + * const ctx = ai.chatContextOrThrow(); + * // ctx.chatId, ctx.clientData are guaranteed non-null + * ``` + */ +function getToolChatContextOrThrow(): ChatTurnContext> { + const ctx = getToolChatContext(); + if (!ctx) { + throw new Error( + "ai.chatContextOrThrow() called outside of a chat.task context. " + + "This helper can only be used inside a subtask invoked via ai.tool() from a chat.task." + ); + } + return ctx; +} + function convertTaskSchemaToToolParameters( task: AnyTask | TaskWithSchema ): Schema { @@ -136,6 +224,12 @@ function convertTaskSchemaToToolParameters( export const ai = { tool: toolFromTask, currentToolOptions: getToolOptionsFromMetadata, + /** Get the tool call ID from inside a subtask invoked via `ai.tool()`. */ + toolCallId: getToolCallId, + /** Get chat context (chatId, turn, clientData, etc.) from inside a subtask of a `chat.task`. Returns undefined if not in a chat context. */ + chatContext: getToolChatContext, + /** Get chat context or throw if not in a chat context. Pass `typeof yourChatTask` for typed clientData. */ + chatContextOrThrow: getToolChatContextOrThrow, }; /** @@ -756,6 +850,14 @@ function chatTask< async () => { locals.set(chatPipeCountKey, 0); + // Store chat context for auto-detection by ai.tool subtasks + locals.set(chatTurnContextKey, { + chatId: currentWirePayload.chatId, + turn, + continuation, + clientData, + }); + // Per-turn stop controller (reset each turn) const stopController = new AbortController(); currentStopController = stopController; diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 0eed9cf08..321f279ba 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -150,6 +150,10 @@ export const deepResearch = schemaTask({ urls: z.array(z.string().url()).describe("URLs to fetch and analyze"), }), run: async ({ query, urls }) => { + // Access chat context from the parent chat.task — typed via typeof aiChat + const { chatId, clientData } = ai.chatContextOrThrow(); + console.log(`Deep research for chat ${chatId}, user ${clientData?.userId}`); + const partId = generateId(); const results: { url: string; status: number; snippet: string }[] = [];