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
This commit is contained in:
Eric Allam
2026-03-08 18:30:45 +00:00
parent acb1f9820e
commit 1cc19f8ca8
2 changed files with 109 additions and 3 deletions
+105 -3
View File
@@ -32,7 +32,24 @@ import {
const METADATA_KEY = "tool.execute.options";
export type ToolCallExecutionOptions = Omit<ToolCallOptions, "abortSignal">;
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<TClientData = unknown> = {
chatId: string;
turn: number;
continuation: boolean;
clientData?: TClientData;
};
const chatTurnContextKey = locals.create<ChatTurnContext>("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<TTaskSchema>, {
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<typeof myChat>();
* // ctx?.clientData is typed based on myChat's clientDataSchema
* ```
*/
function getToolChatContext<TChatTask extends AnyTask = AnyTask>(): ChatTurnContext<InferChatClientData<TChatTask>> | 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<TChatTask>,
};
}
/**
* 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<typeof myChat>();
* // ctx.chatId, ctx.clientData are guaranteed non-null
* ```
*/
function getToolChatContextOrThrow<TChatTask extends AnyTask = AnyTask>(): ChatTurnContext<InferChatClientData<TChatTask>> {
const ctx = getToolChatContext<TChatTask>();
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<any, any, any>
): Schema<unknown> {
@@ -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;
+4
View File
@@ -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<typeof aiChat>();
console.log(`Deep research for chat ${chatId}, user ${clientData?.userId}`);
const partId = generateId();
const results: { url: string; status: number; snippet: string }[] = [];