@getpochi/cli 0.5.60 → 0.5.61

package/dist/cli.js

@@ -347749,6 +347749,520 @@ var originalGenerateId = createIdGenerator({
   prefix: "aitxt",
   size: 24
 });
+async function generateText({
+  model: modelArg,
+  tools,
+  toolChoice,
+  system,
+  prompt,
+  messages,
+  maxRetries: maxRetriesArg,
+  abortSignal,
+  headers,
+  stopWhen = stepCountIs(1),
+  experimental_output: output,
+  experimental_telemetry: telemetry,
+  providerOptions,
+  experimental_activeTools,
+  activeTools = experimental_activeTools,
+  experimental_prepareStep,
+  prepareStep = experimental_prepareStep,
+  experimental_repairToolCall: repairToolCall,
+  experimental_download: download2,
+  experimental_context,
+  _internal: {
+    generateId: generateId3 = originalGenerateId,
+    currentDate = () => /* @__PURE__ */ new Date
+  } = {},
+  onStepFinish,
+  ...settings
+}) {
+  const model = resolveLanguageModel(modelArg);
+  const stopConditions = asArray(stopWhen);
+  const { maxRetries, retry } = prepareRetries({
+    maxRetries: maxRetriesArg,
+    abortSignal
+  });
+  const callSettings = prepareCallSettings(settings);
+  const headersWithUserAgent = withUserAgentSuffix(headers != null ? headers : {}, `ai/${VERSION3}`);
+  const baseTelemetryAttributes = getBaseTelemetryAttributes({
+    model,
+    telemetry,
+    headers: headersWithUserAgent,
+    settings: { ...callSettings, maxRetries }
+  });
+  const initialPrompt = await standardizePrompt({
+    system,
+    prompt,
+    messages
+  });
+  const tracer = getTracer(telemetry);
+  try {
+    return await recordSpan({
+      name: "ai.generateText",
+      attributes: selectTelemetryAttributes({
+        telemetry,
+        attributes: {
+          ...assembleOperationName({
+            operationId: "ai.generateText",
+            telemetry
+          }),
+          ...baseTelemetryAttributes,
+          "ai.model.provider": model.provider,
+          "ai.model.id": model.modelId,
+          "ai.prompt": {
+            input: () => JSON.stringify({ system, prompt, messages })
+          }
+        }
+      }),
+      tracer,
+      fn: async (span) => {
+        var _a172, _b8, _c, _d, _e, _f, _g;
+        const callSettings2 = prepareCallSettings(settings);
+        let currentModelResponse;
+        let clientToolCalls = [];
+        let clientToolOutputs = [];
+        const responseMessages = [];
+        const steps = [];
+        do {
+          const stepInputMessages = [
+            ...initialPrompt.messages,
+            ...responseMessages
+          ];
+          const prepareStepResult = await (prepareStep == null ? undefined : prepareStep({
+            model,
+            steps,
+            stepNumber: steps.length,
+            messages: stepInputMessages
+          }));
+          const promptMessages = await convertToLanguageModelPrompt({
+            prompt: {
+              system: (_a172 = prepareStepResult == null ? undefined : prepareStepResult.system) != null ? _a172 : initialPrompt.system,
+              messages: (_b8 = prepareStepResult == null ? undefined : prepareStepResult.messages) != null ? _b8 : stepInputMessages
+            },
+            supportedUrls: await model.supportedUrls,
+            download: download2
+          });
+          const stepModel = resolveLanguageModel((_c = prepareStepResult == null ? undefined : prepareStepResult.model) != null ? _c : model);
+          const { toolChoice: stepToolChoice, tools: stepTools } = prepareToolsAndToolChoice({
+            tools,
+            toolChoice: (_d = prepareStepResult == null ? undefined : prepareStepResult.toolChoice) != null ? _d : toolChoice,
+            activeTools: (_e = prepareStepResult == null ? undefined : prepareStepResult.activeTools) != null ? _e : activeTools
+          });
+          currentModelResponse = await retry(() => {
+            var _a18;
+            return recordSpan({
+              name: "ai.generateText.doGenerate",
+              attributes: selectTelemetryAttributes({
+                telemetry,
+                attributes: {
+                  ...assembleOperationName({
+                    operationId: "ai.generateText.doGenerate",
+                    telemetry
+                  }),
+                  ...baseTelemetryAttributes,
+                  "ai.model.provider": stepModel.provider,
+                  "ai.model.id": stepModel.modelId,
+                  "ai.prompt.messages": {
+                    input: () => stringifyForTelemetry(promptMessages)
+                  },
+                  "ai.prompt.tools": {
+                    input: () => stepTools == null ? undefined : stepTools.map((tool3) => JSON.stringify(tool3))
+                  },
+                  "ai.prompt.toolChoice": {
+                    input: () => stepToolChoice != null ? JSON.stringify(stepToolChoice) : undefined
+                  },
+                  "gen_ai.system": stepModel.provider,
+                  "gen_ai.request.model": stepModel.modelId,
+                  "gen_ai.request.frequency_penalty": settings.frequencyPenalty,
+                  "gen_ai.request.max_tokens": settings.maxOutputTokens,
+                  "gen_ai.request.presence_penalty": settings.presencePenalty,
+                  "gen_ai.request.stop_sequences": settings.stopSequences,
+                  "gen_ai.request.temperature": (_a18 = settings.temperature) != null ? _a18 : undefined,
+                  "gen_ai.request.top_k": settings.topK,
+                  "gen_ai.request.top_p": settings.topP
+                }
+              }),
+              tracer,
+              fn: async (span2) => {
+                var _a19, _b22, _c2, _d2, _e2, _f2, _g2, _h;
+                const result2 = await stepModel.doGenerate({
+                  ...callSettings2,
+                  tools: stepTools,
+                  toolChoice: stepToolChoice,
+                  responseFormat: output == null ? undefined : output.responseFormat,
+                  prompt: promptMessages,
+                  providerOptions,
+                  abortSignal,
+                  headers: headersWithUserAgent
+                });
+                const responseData = {
+                  id: (_b22 = (_a19 = result2.response) == null ? undefined : _a19.id) != null ? _b22 : generateId3(),
+                  timestamp: (_d2 = (_c2 = result2.response) == null ? undefined : _c2.timestamp) != null ? _d2 : currentDate(),
+                  modelId: (_f2 = (_e2 = result2.response) == null ? undefined : _e2.modelId) != null ? _f2 : stepModel.modelId,
+                  headers: (_g2 = result2.response) == null ? undefined : _g2.headers,
+                  body: (_h = result2.response) == null ? undefined : _h.body
+                };
+                span2.setAttributes(selectTelemetryAttributes({
+                  telemetry,
+                  attributes: {
+                    "ai.response.finishReason": result2.finishReason,
+                    "ai.response.text": {
+                      output: () => extractTextContent(result2.content)
+                    },
+                    "ai.response.toolCalls": {
+                      output: () => {
+                        const toolCalls = asToolCalls(result2.content);
+                        return toolCalls == null ? undefined : JSON.stringify(toolCalls);
+                      }
+                    },
+                    "ai.response.id": responseData.id,
+                    "ai.response.model": responseData.modelId,
+                    "ai.response.timestamp": responseData.timestamp.toISOString(),
+                    "ai.response.providerMetadata": JSON.stringify(result2.providerMetadata),
+                    "ai.usage.promptTokens": result2.usage.inputTokens,
+                    "ai.usage.completionTokens": result2.usage.outputTokens,
+                    "gen_ai.response.finish_reasons": [result2.finishReason],
+                    "gen_ai.response.id": responseData.id,
+                    "gen_ai.response.model": responseData.modelId,
+                    "gen_ai.usage.input_tokens": result2.usage.inputTokens,
+                    "gen_ai.usage.output_tokens": result2.usage.outputTokens
+                  }
+                }));
+                return { ...result2, response: responseData };
+              }
+            });
+          });
+          const stepToolCalls = await Promise.all(currentModelResponse.content.filter((part) => part.type === "tool-call").map((toolCall) => parseToolCall({
+            toolCall,
+            tools,
+            repairToolCall,
+            system,
+            messages: stepInputMessages
+          })));
+          for (const toolCall of stepToolCalls) {
+            if (toolCall.invalid) {
+              continue;
+            }
+            const tool3 = tools[toolCall.toolName];
+            if ((tool3 == null ? undefined : tool3.onInputAvailable) != null) {
+              await tool3.onInputAvailable({
+                input: toolCall.input,
+                toolCallId: toolCall.toolCallId,
+                messages: stepInputMessages,
+                abortSignal,
+                experimental_context
+              });
+            }
+          }
+          const invalidToolCalls = stepToolCalls.filter((toolCall) => toolCall.invalid && toolCall.dynamic);
+          clientToolOutputs = [];
+          for (const toolCall of invalidToolCalls) {
+            clientToolOutputs.push({
+              type: "tool-error",
+              toolCallId: toolCall.toolCallId,
+              toolName: toolCall.toolName,
+              input: toolCall.input,
+              error: getErrorMessage2(toolCall.error),
+              dynamic: true
+            });
+          }
+          clientToolCalls = stepToolCalls.filter((toolCall) => !toolCall.providerExecuted);
+          if (tools != null) {
+            clientToolOutputs.push(...await executeTools({
+              toolCalls: clientToolCalls.filter((toolCall) => !toolCall.invalid),
+              tools,
+              tracer,
+              telemetry,
+              messages: stepInputMessages,
+              abortSignal,
+              experimental_context
+            }));
+          }
+          const stepContent = asContent({
+            content: currentModelResponse.content,
+            toolCalls: stepToolCalls,
+            toolOutputs: clientToolOutputs
+          });
+          responseMessages.push(...toResponseMessages({
+            content: stepContent,
+            tools
+          }));
+          const currentStepResult = new DefaultStepResult({
+            content: stepContent,
+            finishReason: currentModelResponse.finishReason,
+            usage: currentModelResponse.usage,
+            warnings: currentModelResponse.warnings,
+            providerMetadata: currentModelResponse.providerMetadata,
+            request: (_f = currentModelResponse.request) != null ? _f : {},
+            response: {
+              ...currentModelResponse.response,
+              messages: structuredClone(responseMessages)
+            }
+          });
+          logWarnings((_g = currentModelResponse.warnings) != null ? _g : []);
+          steps.push(currentStepResult);
+          await (onStepFinish == null ? undefined : onStepFinish(currentStepResult));
+        } while (clientToolCalls.length > 0 && clientToolOutputs.length === clientToolCalls.length && !await isStopConditionMet({ stopConditions, steps }));
+        span.setAttributes(selectTelemetryAttributes({
+          telemetry,
+          attributes: {
+            "ai.response.finishReason": currentModelResponse.finishReason,
+            "ai.response.text": {
+              output: () => extractTextContent(currentModelResponse.content)
+            },
+            "ai.response.toolCalls": {
+              output: () => {
+                const toolCalls = asToolCalls(currentModelResponse.content);
+                return toolCalls == null ? undefined : JSON.stringify(toolCalls);
+              }
+            },
+            "ai.response.providerMetadata": JSON.stringify(currentModelResponse.providerMetadata),
+            "ai.usage.promptTokens": currentModelResponse.usage.inputTokens,
+            "ai.usage.completionTokens": currentModelResponse.usage.outputTokens
+          }
+        }));
+        const lastStep = steps[steps.length - 1];
+        return new DefaultGenerateTextResult({
+          steps,
+          resolvedOutput: await (output == null ? undefined : output.parseOutput({ text: lastStep.text }, {
+            response: lastStep.response,
+            usage: lastStep.usage,
+            finishReason: lastStep.finishReason
+          }))
+        });
+      }
+    });
+  } catch (error40) {
+    throw wrapGatewayError(error40);
+  }
+}
+async function executeTools({
+  toolCalls,
+  tools,
+  tracer,
+  telemetry,
+  messages,
+  abortSignal,
+  experimental_context
+}) {
+  const toolOutputs = await Promise.all(toolCalls.map(async ({ toolCallId, toolName, input }) => {
+    const tool3 = tools[toolName];
+    if ((tool3 == null ? undefined : tool3.execute) == null) {
+      return;
+    }
+    return recordSpan({
+      name: "ai.toolCall",
+      attributes: selectTelemetryAttributes({
+        telemetry,
+        attributes: {
+          ...assembleOperationName({
+            operationId: "ai.toolCall",
+            telemetry
+          }),
+          "ai.toolCall.name": toolName,
+          "ai.toolCall.id": toolCallId,
+          "ai.toolCall.args": {
+            output: () => JSON.stringify(input)
+          }
+        }
+      }),
+      tracer,
+      fn: async (span) => {
+        try {
+          const stream = executeTool({
+            execute: tool3.execute.bind(tool3),
+            input,
+            options: {
+              toolCallId,
+              messages,
+              abortSignal,
+              experimental_context
+            }
+          });
+          let output;
+          for await (const part of stream) {
+            if (part.type === "final") {
+              output = part.output;
+            }
+          }
+          try {
+            span.setAttributes(selectTelemetryAttributes({
+              telemetry,
+              attributes: {
+                "ai.toolCall.result": {
+                  output: () => JSON.stringify(output)
+                }
+              }
+            }));
+          } catch (ignored) {}
+          return {
+            type: "tool-result",
+            toolCallId,
+            toolName,
+            input,
+            output,
+            dynamic: tool3.type === "dynamic"
+          };
+        } catch (error40) {
+          recordErrorOnSpan(span, error40);
+          return {
+            type: "tool-error",
+            toolCallId,
+            toolName,
+            input,
+            error: error40,
+            dynamic: tool3.type === "dynamic"
+          };
+        }
+      }
+    });
+  }));
+  return toolOutputs.filter((output) => output != null);
+}
+var DefaultGenerateTextResult = class {
+  constructor(options) {
+    this.steps = options.steps;
+    this.resolvedOutput = options.resolvedOutput;
+  }
+  get finalStep() {
+    return this.steps[this.steps.length - 1];
+  }
+  get content() {
+    return this.finalStep.content;
+  }
+  get text() {
+    return this.finalStep.text;
+  }
+  get files() {
+    return this.finalStep.files;
+  }
+  get reasoningText() {
+    return this.finalStep.reasoningText;
+  }
+  get reasoning() {
+    return this.finalStep.reasoning;
+  }
+  get toolCalls() {
+    return this.finalStep.toolCalls;
+  }
+  get staticToolCalls() {
+    return this.finalStep.staticToolCalls;
+  }
+  get dynamicToolCalls() {
+    return this.finalStep.dynamicToolCalls;
+  }
+  get toolResults() {
+    return this.finalStep.toolResults;
+  }
+  get staticToolResults() {
+    return this.finalStep.staticToolResults;
+  }
+  get dynamicToolResults() {
+    return this.finalStep.dynamicToolResults;
+  }
+  get sources() {
+    return this.finalStep.sources;
+  }
+  get finishReason() {
+    return this.finalStep.finishReason;
+  }
+  get warnings() {
+    return this.finalStep.warnings;
+  }
+  get providerMetadata() {
+    return this.finalStep.providerMetadata;
+  }
+  get response() {
+    return this.finalStep.response;
+  }
+  get request() {
+    return this.finalStep.request;
+  }
+  get usage() {
+    return this.finalStep.usage;
+  }
+  get totalUsage() {
+    return this.steps.reduce((totalUsage, step) => {
+      return addLanguageModelUsage(totalUsage, step.usage);
+    }, {
+      inputTokens: undefined,
+      outputTokens: undefined,
+      totalTokens: undefined,
+      reasoningTokens: undefined,
+      cachedInputTokens: undefined
+    });
+  }
+  get experimental_output() {
+    if (this.resolvedOutput == null) {
+      throw new NoOutputSpecifiedError;
+    }
+    return this.resolvedOutput;
+  }
+};
+function asToolCalls(content) {
+  const parts = content.filter((part) => part.type === "tool-call");
+  if (parts.length === 0) {
+    return;
+  }
+  return parts.map((toolCall) => ({
+    toolCallId: toolCall.toolCallId,
+    toolName: toolCall.toolName,
+    input: toolCall.input
+  }));
+}
+function asContent({
+  content,
+  toolCalls,
+  toolOutputs
+}) {
+  return [
+    ...content.map((part) => {
+      switch (part.type) {
+        case "text":
+        case "reasoning":
+        case "source":
+          return part;
+        case "file": {
+          return {
+            type: "file",
+            file: new DefaultGeneratedFile(part)
+          };
+        }
+        case "tool-call": {
+          return toolCalls.find((toolCall) => toolCall.toolCallId === part.toolCallId);
+        }
+        case "tool-result": {
+          const toolCall = toolCalls.find((toolCall2) => toolCall2.toolCallId === part.toolCallId);
+          if (toolCall == null) {
+            throw new Error(`Tool call ${part.toolCallId} not found.`);
+          }
+          if (part.isError) {
+            return {
+              type: "tool-error",
+              toolCallId: part.toolCallId,
+              toolName: part.toolName,
+              input: toolCall.input,
+              error: part.result,
+              providerExecuted: true,
+              dynamic: toolCall.dynamic
+            };
+          }
+          return {
+            type: "tool-result",
+            toolCallId: part.toolCallId,
+            toolName: part.toolName,
+            input: toolCall.input,
+            output: part.result,
+            providerExecuted: true,
+            dynamic: toolCall.dynamic
+          };
+        }
+      }
+    }),
+    ...toolOutputs
+  ];
+}
 function prepareHeaders(headers, defaultHeaders) {
   const responseHeaders = new Headers(headers != null ? headers : {});
   for (const [key, value2] of Object.entries(defaultHeaders)) {
@@ -359465,7 +359979,6 @@ function createModel(vendorId, opts) {
 
 // ../vendor-pochi/src/model.ts
 function createPochiModel({
-  id,
   modelId,
   getCredentials
 }) {
@@ -359474,11 +359987,28 @@ function createPochiModel({
     provider: "pochi",
     modelId: modelId || "<default>",
     supportedUrls: {},
-    doGenerate: async () => Promise.reject("Not implemented"),
-    doStream: async ({ prompt, abortSignal, stopSequences, tools }) => {
+    doGenerate: async ({ headers, ...options }) => {
+      const apiClient = createApiClient(getCredentials);
+      const resp = await apiClient.api.chat.$post({
+        json: {
+          model: modelId,
+          options
+        }
+      }, {
+        headers: headers ? i5(headers, (x) => x || "") : undefined
+      });
+      const data = await resp.json();
+      return data;
+    },
+    doStream: async ({
+      prompt,
+      abortSignal,
+      stopSequences,
+      tools,
+      headers
+    }) => {
       const apiClient = createApiClient(getCredentials);
       const data = {
-        id,
         model: modelId,
         callOptions: {
           prompt: convertFilePartDataToBase64(prompt),
@@ -359489,6 +360019,7 @@ function createPochiModel({
       const resp = await apiClient.api.chat.stream.$post({
         json: data
       }, {
+        headers: headers ? i5(headers, (x) => x || "") : undefined,
         init: {
           signal: abortSignal
         }
@@ -368440,7 +368971,7 @@ var {
 // package.json
 var package_default = {
   name: "@getpochi/cli",
-  version: "0.5.60",
+  version: "0.5.61",
   type: "module",
   bin: {
     pochi: "src/cli.ts"
@@ -409842,6 +410373,9 @@ function toTaskError(error46) {
   return internalError(error46.message);
 }
 
+// ../common/src/pochi-api.ts
+var PochiTaskIdHeader = "x-pochi-task-id";
+
 // ../livekit/src/chat/llm/generate-task-title.ts
 var logger18 = getLogger("generateTaskTitle");
 async function generateTaskTitle(options6) {
@@ -409853,6 +410387,7 @@ async function generateTaskTitle(options6) {
   return newTitle;
 }
 async function generateTaskTitleImpl({
+  taskId,
   title,
   messages: messages2,
   getModel,
@@ -409870,7 +410405,7 @@ async function generateTaskTitleImpl({
   if (partCount >= 5 && partCount < 20 && !isTitleGeneratedByLlm(title, titleFromMessages)) {
     try {
       const model2 = getModel();
-      const title2 = await generateTitle2(model2, messages2, abortSignal);
+      const title2 = await generateTitle2(taskId, model2, messages2, abortSignal);
       if (title2.length > 0) {
         return title2;
       }
@@ -409896,7 +410431,7 @@ function getTitleFromMessages(messages2) {
 function isTitleGeneratedByLlm(title, titleFromMessages) {
   return title !== titleFromMessages;
 }
-async function generateTitle2(model2, inputMessages, abortSignal) {
+async function generateTitle2(taskId, model2, inputMessages, abortSignal) {
   const messages2 = [
     ...inputMessages,
     {
@@ -409910,14 +410445,17 @@ async function generateTitle2(model2, inputMessages, abortSignal) {
       ]
     }
   ];
-  const stream12 = streamText({
+  const resp = await generateText({
+    headers: {
+      [PochiTaskIdHeader]: taskId
+    },
     model: model2,
     prompt: convertToModelMessages(formatters.llm(messages2, { removeSystemReminder: true })),
     abortSignal,
     maxOutputTokens: 50,
     maxRetries: 0
   });
-  return (await stream12.text).trim();
+  return resp.text;
 }
 
 // ../livekit/src/chat/background-job/manager.ts
@@ -409964,6 +410502,7 @@ async function process16({
     return;
   }
   const newTitle = await generateTaskTitle({
+    taskId,
     title: task.title,
     messages: messages2,
     getModel
@@ -410050,56 +410589,33 @@ async function digest(data) {
 }
 
 // ../livekit/src/chat/llm/repair-tool-call.ts
-var logger21 = getLogger("RepairToolCall");
-var makeRepairToolCall = (model2) => async ({ toolCall, inputSchema: inputSchema2, error: error46 }) => {
+var makeRepairToolCall = (taskId, model2) => async ({ toolCall, inputSchema: inputSchema2, error: error46 }) => {
   if (NoSuchToolError.isInstance(error46)) {
     return null;
   }
-  const result2 = await streamText({
+  const tools = createClientTools();
+  const tool2 = tools[toolCall.toolName];
+  const { object: repairedArgs } = await generateObject({
+    headers: {
+      [PochiTaskIdHeader]: taskId
+    },
     model: model2,
+    schema: tool2.inputSchema,
     prompt: [
-      `The model tried to call the tool "${toolCall.toolName}" with the following arguments:`,
+      `The model tried to call the tool "${toolCall.toolName}" with the following inputs:`,
       JSON.stringify(toolCall.input),
       "The tool accepts the following schema:",
       JSON.stringify(inputSchema2(toolCall)),
-      "Please fix the arguments. Please ONLY output the json string, without any other text (no markdown code block wrapper, either)",
-      "",
-      "<good-example>",
-      '{"path": "./src/file.ts", "content": "console.log("hello");"}',
-      "</good-example>",
-      "",
-      "<bad-example>",
-      "```json",
-      '{"path": "./src/file.ts", "content": "console.log("hello");"}',
-      "```",
-      "</bad-example>",
-      "",
-      "<bad-example>",
-      "Here is the corrected JSON:",
-      '{"path": "./src/file.ts", "content": "console.log("hello");"}',
-      "</bad-example>",
-      "",
-      "<bad-example>",
-      "```",
-      '{"path": "./src/file.ts", "content": "console.log("hello");"}',
-      "```",
-      "</bad-example>",
-      "",
-      "<bad-example>",
-      '{"path": "./src/file.ts", "content": "console.log("hello");"} // Fixed arguments',
-      "</bad-example>"
+      "Please fix the inputs."
     ].join(`
-`),
-    maxOutputTokens: 3000,
-    maxRetries: 0
+`)
   });
-  const input2 = await result2.text;
-  logger21.debug("Repairing tool call:", toolCall.toolName, input2);
-  return { ...toolCall, input: input2 };
+  return { ...toolCall, input: JSON.stringify(repairedArgs) };
 };
 // ../livekit/src/chat/llm/compact-task.ts
-var logger22 = getLogger("compactTask");
+var logger21 = getLogger("compactTask");
 async function compactTask({
+  taskId,
   model: model2,
   messages: messages2,
   abortSignal,
@@ -410110,7 +410626,7 @@ async function compactTask({
     throw new Error("No messages to compact");
   }
   try {
-    const text20 = prompts.inlineCompact(await createSummary(model2, abortSignal, messages2.slice(0, -1)), messages2.length - 1);
+    const text20 = prompts.inlineCompact(await createSummary(taskId, model2, abortSignal, messages2.slice(0, -1)), messages2.length - 1);
     if (inline) {
       lastMessage.parts.unshift({
         type: "text",
@@ -410120,10 +410636,10 @@ async function compactTask({
     }
     return text20;
   } catch (err2) {
-    logger22.warn("Failed to create summary", err2);
+    logger21.warn("Failed to create summary", err2);
   }
 }
-async function createSummary(model2, abortSignal, inputMessages) {
+async function createSummary(taskId, model2, abortSignal, inputMessages) {
   const messages2 = [
     ...inputMessages,
     {
@@ -410137,7 +410653,10 @@ async function createSummary(model2, abortSignal, inputMessages) {
       ]
     }
   ];
-  const stream12 = streamText({
+  const resp = await generateText({
+    headers: {
+      [PochiTaskIdHeader]: taskId
+    },
     model: model2,
     prompt: convertToModelMessages(formatters.llm(messages2, {
       removeSystemReminder: true
@@ -410146,7 +410665,7 @@ async function createSummary(model2, abortSignal, inputMessages) {
     maxOutputTokens: 3000,
     maxRetries: 0
   });
-  return stream12.text;
+  return resp.text;
 }
 // ../livekit/src/chat/mcp-utils.ts
 function parseMcpToolSet(store, mcpToolSet) {
@@ -410871,10 +411390,10 @@ function createStopWordStream(stop3) {
   });
 }
 // ../livekit/src/chat/middlewares/output-schema-middleware.ts
-function createOutputSchemaMiddleware(outputSchema2) {
+function createOutputSchemaMiddleware(taskId, model2, outputSchema2) {
   return {
     middlewareVersion: "v2",
-    wrapStream: async ({ doStream, model: model2 }) => {
+    wrapStream: async ({ doStream }) => {
       const { stream: stream12, ...rest } = await doStream();
       let toolCallId = "";
       const transformedStream = stream12.pipeThrough(new TransformStream({
@@ -410901,7 +411420,7 @@ function createOutputSchemaMiddleware(outputSchema2) {
             const { result: result2 } = parsedResult.value;
             const newInput = {
               ...parsedResult.value,
-              result: await ensureOutputSchema(model2, outputSchema2, result2)
+              result: await ensureOutputSchema(taskId, model2, outputSchema2, result2)
             };
             controller.enqueue({
               ...chunk4,
@@ -410920,9 +411439,12 @@ function createOutputSchemaMiddleware(outputSchema2) {
     }
   };
 }
-async function ensureOutputSchema(model2, schema7, content3) {
+async function ensureOutputSchema(taskId, model2, schema7, content3) {
   try {
     const { object: object4 } = await generateObject({
+      headers: {
+        [PochiTaskIdHeader]: taskId
+      },
       model: model2,
       schema: schema7,
       prompt: [
@@ -411097,7 +411619,7 @@ function createGoogleVertexTuningModel(llm) {
 }
 
 // ../livekit/src/chat/models/openai.ts
-var logger23 = getLogger("openai");
+var logger22 = getLogger("openai");
 var OpenAIRequestParamsSchema = exports_external.object({
   max_tokens: exports_external.number().positive().optional(),
   max_completion_tokens: exports_external.number().positive().optional()
@@ -411147,7 +411669,7 @@ function overrideMaxOutputToken(body3) {
     const json10 = JSON.parse(body3);
     const result2 = OpenAIRequestParamsSchema.safeParse(json10);
     if (!result2.success) {
-      logger23.error("OpenAI request body validation failed:", result2.error);
+      logger22.error("OpenAI request body validation failed:", result2.error);
       return;
     }
     const parsed = result2.data;
@@ -411156,7 +411678,7 @@ function overrideMaxOutputToken(body3) {
     }
     return JSON.stringify(parsed);
   } catch (error46) {
-    logger23.error("Failed to parse OpenAI request body:", error46);
+    logger22.error("Failed to parse OpenAI request body:", error46);
   }
   return;
 }
@@ -411185,12 +411707,9 @@ function createOpenAIResponsesModel(llm) {
 }
 
 // ../livekit/src/chat/models/index.ts
-function createModel2({
-  id: id5,
-  llm
-}) {
+function createModel2({ llm }) {
   if (llm.type === "vendor") {
-    return llm.getModel(id5);
+    return llm.getModel();
   }
   if (llm.type === "openai") {
     return createOpenAIModel(llm);
@@ -411246,6 +411765,7 @@ class FlexibleChatTransport {
       abortSignal,
       getters: this.getters
     });
+    const model2 = createModel2({ llm });
     const middlewares = [];
     if (!this.isSubTask) {
       middlewares.push(createNewTaskMiddleware(this.store, environment2?.info.cwd, chatId, customAgents));
@@ -411254,7 +411774,7 @@ class FlexibleChatTransport {
       middlewares.push(createReasoningMiddleware());
     }
     if (this.outputSchema) {
-      middlewares.push(createOutputSchemaMiddleware(this.outputSchema));
+      middlewares.push(createOutputSchemaMiddleware(chatId, model2, this.outputSchema));
     }
     if (llm.useToolCallMiddleware) {
       middlewares.push(createToolCallMiddleware());
@@ -411274,8 +411794,10 @@ class FlexibleChatTransport {
       return true;
     });
     const preparedMessages = await prepareMessages(messages2, environment2);
-    const model2 = createModel2({ id: chatId, llm });
     const stream12 = streamText({
+      headers: {
+        [PochiTaskIdHeader]: chatId
+      },
       system: prompts.system(environment2?.info?.customRules, this.customAgent, mcpInfo?.instructions),
       messages: convertToModelMessages(formatters.llm(preparedMessages), { tools }),
       model: wrapLanguageModel({
@@ -411286,7 +411808,7 @@ class FlexibleChatTransport {
       tools,
       maxRetries: 0,
       onError: () => {},
-      experimental_repairToolCall: makeRepairToolCall(model2),
+      experimental_repairToolCall: makeRepairToolCall(chatId, model2),
       experimental_download: async (items) => {
         const promises5 = items.map(async ({
           url: url3,
@@ -411366,7 +411888,7 @@ function estimateTotalTokens(messages2) {
 }
 
 // ../livekit/src/chat/live-chat-kit.ts
-var logger24 = getLogger("LiveChatKit");
+var logger23 = getLogger("LiveChatKit");
 
 class LiveChatKit {
   taskId;
@@ -411416,15 +411938,16 @@ class LiveChatKit {
       const lastMessage = messages2.at(-1);
       if (lastMessage?.role === "user" && lastMessage.metadata?.kind === "user" && lastMessage.metadata.compact) {
         try {
-          const model2 = createModel2({ id: taskId, llm: getters.getLLM() });
+          const model2 = createModel2({ llm: getters.getLLM() });
           await compactTask({
+            taskId: this.taskId,
             model: model2,
             messages: messages2,
             abortSignal: abortSignal2,
             inline: true
           });
         } catch (err2) {
-          logger24.error("Failed to compact task", err2);
+          logger23.error("Failed to compact task", err2);
           throw err2;
         }
       }
@@ -411435,8 +411958,9 @@ class LiveChatKit {
     this.spawn = async () => {
       const taskId2 = crypto.randomUUID();
       const { messages: messages2 } = this.chat;
-      const model2 = createModel2({ id: taskId2, llm: getters.getLLM() });
+      const model2 = createModel2({ llm: getters.getLLM() });
       const summary6 = await compactTask({
+        taskId: taskId2,
         model: model2,
         messages: messages2,
         abortSignal
@@ -411518,7 +412042,7 @@ class LiveChatKit {
       if (!task) {
         throw new Error("Task not found");
       }
-      const getModel = () => createModel2({ id: this.taskId, llm: getters.getLLM() });
+      const getModel = () => createModel2({ llm: getters.getLLM() });
       scheduleGenerateTitleJob({
         taskId: this.taskId,
         store,
@@ -411563,7 +412087,7 @@ class LiveChatKit {
     }));
   };
   onError = (error46) => {
-    logger24.error("onError", error46);
+    logger23.error("onError", error46);
     const lastMessage = this.chat.messages.at(-1) || null;
     this.store.commit(events.chatStreamFailed({
       id: this.taskId,
@@ -411716,7 +412240,7 @@ import * as nodePath from "node:path";
 
 // ../common/src/diff-utils.ts
 var import_fast_levenshtein = __toESM(require_levenshtein(), 1);
-var logger25 = getLogger("diffUtils");
+var logger24 = getLogger("diffUtils");
 function normalize6(content3) {
   return content3.replace(/\r\n/g, `
 `).trimEnd();
@@ -411729,7 +412253,7 @@ class DiffError extends Error {
   }
 }
 async function parseDiffAndApply(fileContent3, searchContent, replaceContent, expectedReplacements = 1) {
-  logger25.trace(`Applying diff with expectedReplacements: ${expectedReplacements}`);
+  logger24.trace(`Applying diff with expectedReplacements: ${expectedReplacements}`);
   const isCRLF = fileContent3.includes(`\r
 `);
   const normalizedFileContent = normalize6(fileContent3);
@@ -411746,11 +412270,11 @@ async function parseDiffAndApply(fileContent3, searchContent, replaceContent, ex
   const matches2 = searchContentExact(normalizedFileContent, normalizedSearchContent);
   if (matches2.length < expectedReplacements) {
     matches2.push(...searchContentWithLineTrimmed(normalizedFileContent, normalizedSearchContent));
-    logger25.trace(`Found ${matches2.length} matches after line trimming search strategy`);
+    logger24.trace(`Found ${matches2.length} matches after line trimming search strategy`);
   }
   if (matches2.length < expectedReplacements) {
     matches2.push(...searchContentByBlockAnchor(normalizedFileContent, normalizedSearchContent));
-    logger25.trace(`Found ${matches2.length} matches after block anchor search strategy`);
+    logger24.trace(`Found ${matches2.length} matches after block anchor search strategy`);
   }
   if (matches2.length === 0) {
     throw new DiffError("Search content does not match the file content. Try to reread the file for the latest content.");
@@ -411759,7 +412283,7 @@ async function parseDiffAndApply(fileContent3, searchContent, replaceContent, ex
     throw new DiffError(`Expected ${expectedReplacements} occurrences but found ${matches2.length}. Please verify the search content and expectedReplacements parameter.`);
   }
   const result2 = replaceMatches(normalizedFileContent, matches2, replaceContent);
-  logger25.trace("Successfully applied diff");
+  logger24.trace("Successfully applied diff");
   if (isCRLF) {
     return result2.replace(/\n/g, `\r
 `);
@@ -412113,11 +412637,11 @@ var readFile13 = () => async ({ path: path27, startLine, endLine }, { cwd: cwd2
 
 // src/tools/search-files.ts
 import * as fs14 from "node:fs";
-var logger26 = getLogger("searchFiles");
+var logger25 = getLogger("searchFiles");
 var searchFiles2 = (context15) => async ({ path: path27, regex: regex3, filePattern }, { abortSignal, cwd: cwd2 }) => {
   const rgPath = context15.rg;
   if (!rgPath || !fs14.existsSync(rgPath)) {
-    logger26.error("Ripgrep not found at path", rgPath);
+    logger25.error("Ripgrep not found at path", rgPath);
     throw new Error(`Ripgrep not found at path: ${rgPath}`);
   }
   return await searchFilesWithRipgrep(path27, regex3, rgPath, cwd2, filePattern, abortSignal);
@@ -412195,7 +412719,7 @@ async function executeToolCall(tool2, options6, cwd2, abortSignal) {
 }
 
 // src/task-runner.ts
-var logger27 = getLogger("TaskRunner");
+var logger26 = getLogger("TaskRunner");
 
 class TaskRunner {
   cwd;
@@ -412272,9 +412796,9 @@ class TaskRunner {
     return this.chatKit.task?.shareId;
   }
   async run() {
-    logger27.debug("Starting TaskRunner...");
+    logger26.debug("Starting TaskRunner...");
     try {
-      logger27.trace("Start step loop.");
+      logger26.trace("Start step loop.");
       this.stepCount.reset();
       while (true) {
         const stepResult = await this.step();
@@ -412289,7 +412813,7 @@ class TaskRunner {
       }
     } catch (e11) {
       const error46 = toError2(e11);
-      logger27.trace("Failed:", error46);
+      logger26.trace("Failed:", error46);
     }
   }
   async step() {
@@ -412327,24 +412851,24 @@ class TaskRunner {
       throw new Error("Task is not loaded");
     }
     if ((task.status === "completed" || task.status === "pending-input") && isResultMessage(message)) {
-      logger27.trace("Task is completed or pending input, no more steps to process.");
+      logger26.trace("Task is completed or pending input, no more steps to process.");
       return "finished";
     }
     if (task.status === "failed") {
       if (task.error?.kind === "APICallError" && !task.error.isRetryable) {
         return "finished";
       }
-      logger27.error("Task is failed, trying to resend last message to resume it.", task.error);
+      logger26.error("Task is failed, trying to resend last message to resume it.", task.error);
       return "retry";
     }
     if (message.role !== "assistant") {
-      logger27.trace("Last message is not a assistant message, resending it to resume the task.");
+      logger26.trace("Last message is not a assistant message, resending it to resume the task.");
       return "retry";
     }
     if (isAssistantMessageWithEmptyParts(message) || isAssistantMessageWithPartialToolCalls(message) || isAssistantMessageWithOutputError(message) || lastAssistantMessageIsCompleteWithToolCalls({
       messages: this.chat.messages
     })) {
-      logger27.trace("Last message is assistant with empty parts or partial/completed tool calls, resending it to resume the task.");
+      logger26.trace("Last message is assistant with empty parts or partial/completed tool calls, resending it to resume the task.");
       const processed = prepareLastMessageForRetry(message);
       if (processed) {
         this.chat.appendOrReplaceMessage(processed);
@@ -412352,28 +412876,28 @@ class TaskRunner {
       return "retry";
     }
     if (isAssistantMessageWithNoToolCalls(message)) {
-      logger27.trace("Last message is assistant with no tool calls, sending a new user reminder.");
+      logger26.trace("Last message is assistant with no tool calls, sending a new user reminder.");
       const message2 = createUserMessage(prompts.createSystemReminder("You should use tool calls to answer the question, for example, use attemptCompletion if the job is done, or use askFollowupQuestions to clarify the request."));
       this.chat.appendOrReplaceMessage(message2);
       return "retry";
     }
   }
   async processToolCalls(message) {
-    logger27.trace("Processing tool calls in the last message.");
+    logger26.trace("Processing tool calls in the last message.");
     for (const toolCall of message.parts.filter(isToolUIPart)) {
       if (toolCall.state !== "input-available")
         continue;
       const toolName = getToolName(toolCall);
-      logger27.trace(`Found tool call: ${toolName} with args: ${JSON.stringify(toolCall.input)}`);
+      logger26.trace(`Found tool call: ${toolName} with args: ${JSON.stringify(toolCall.input)}`);
       const toolResult = await executeToolCall(toolCall, this.toolCallOptions, this.cwd);
       await this.chatKit.chat.addToolResult({
         tool: toolName,
         toolCallId: toolCall.toolCallId,
         output: toolResult
       });
-      logger27.trace(`Tool call result: ${JSON.stringify(toolResult)}`);
+      logger26.trace(`Tool call result: ${JSON.stringify(toolResult)}`);
     }
-    logger27.trace("All tool calls processed in the last message.");
+    logger26.trace("All tool calls processed in the last message.");
     return "next";
   }
 }
@@ -412674,8 +413198,8 @@ function registerUpgradeCommand(program5) {
   });
 }
 // src/cli.ts
-var logger28 = getLogger("Pochi");
-logger28.debug(`pochi v${package_default.version}`);
+var logger27 = getLogger("Pochi");
+logger27.debug(`pochi v${package_default.version}`);
 var parsePositiveInt = (input2) => {
   if (!input2) {
     return program5.error("The value for this option must be a positive integer.");
@@ -412798,8 +413322,7 @@ async function createLLMConfigWithVendors(program6, model2) {
     return {
       type: "vendor",
       useToolCallMiddleware: options6.useToolCallMiddleware,
-      getModel: (id5) => createModel(vendorId, {
-        id: id5,
+      getModel: () => createModel(vendorId, {
         modelId,
         getCredentials: vendor2.getCredentials
       })
@@ -412815,8 +413338,7 @@ async function createLLMConfigWithPochi(model2) {
     return {
       type: "vendor",
       useToolCallMiddleware: pochiModelOptions.useToolCallMiddleware,
-      getModel: (id5) => createModel(vendorId, {
-        id: id5,
+      getModel: () => createModel(vendorId, {
         modelId: model2,
         getCredentials: vendor2.getCredentials
       })

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@getpochi/cli",
-  "version": "0.5.60",
+  "version": "0.5.61",
   "type": "module",
   "bin": {
     "pochi": "dist/cli.js"