@raindrop-ai/ai-sdk 0.0.3 → 0.0.5

package/dist/{chunk-3R547TFD.mjs → chunk-5DKYVVJZ.mjs}

@@ -84,7 +84,7 @@ async function postJson(url, body, headers, opts) {
 // package.json
 var package_default = {
   name: "@raindrop-ai/ai-sdk",
-  version: "0.0.3"};
+  version: "0.0.5"};
 
 // src/internal/version.ts
 var libraryName = package_default.name;
@@ -761,9 +761,6 @@ async function* asyncGeneratorWithCurrent(span, gen) {
     nextValue = yield result.value;
   }
 }
-function isAsyncGenerator(value) {
-  return value !== null && typeof value === "object" && Symbol.asyncIterator in value && typeof value.next === "function";
-}
 
 // src/internal/wrap/helpers.ts
 function isRecord(value) {
@@ -848,31 +845,6 @@ function extractObjectOutput(result) {
   if (obj === void 0) return void 0;
   return (_a = safeJson(obj)) != null ? _a : String(obj);
 }
-function extractUsage(result) {
-  if (!isRecord(result)) return void 0;
-  const usage = isRecord(result["totalUsage"]) ? result["totalUsage"] : result["usage"];
-  if (!isRecord(usage)) return void 0;
-  let inputTokens;
-  const inputVal = usage["inputTokens"];
-  if (typeof inputVal === "number") {
-    inputTokens = inputVal;
-  } else if (isRecord(inputVal) && typeof inputVal["total"] === "number") {
-    inputTokens = inputVal["total"];
-  } else if (typeof usage["promptTokens"] === "number") {
-    inputTokens = usage["promptTokens"];
-  }
-  let outputTokens;
-  const outputVal = usage["outputTokens"];
-  if (typeof outputVal === "number") {
-    outputTokens = outputVal;
-  } else if (isRecord(outputVal) && typeof outputVal["total"] === "number") {
-    outputTokens = outputVal["total"];
-  } else if (typeof usage["completionTokens"] === "number") {
-    outputTokens = usage["completionTokens"];
-  }
-  if (inputTokens === void 0 && outputTokens === void 0) return void 0;
-  return { inputTokens, outputTokens };
-}
 function isAgentClass(value) {
   if (typeof value !== "function") return false;
   const proto = value.prototype;
@@ -882,7 +854,13 @@ function isAgentClass(value) {
 function extractModel(result) {
   if (!isRecord(result)) return void 0;
   const model = result["model"];
-  return typeof model === "string" && model.length ? model : void 0;
+  if (typeof model === "string" && model.length) return model;
+  const response = result["response"];
+  if (isRecord(response)) {
+    const modelId = response["modelId"];
+    if (typeof modelId === "string" && modelId.length) return modelId;
+  }
+  return void 0;
 }
 function extractFinishReason(result) {
   if (!isRecord(result)) return void 0;
@@ -1297,10 +1275,391 @@ function mergeContexts(wrapTime, callTime) {
   }
   return result;
 }
+function getCurrentParentSpanContextSync() {
+  return getContextManager().getParentSpanIds();
+}
+function runWithParentSpanContextSync(ctx, fn) {
+  const cm = getContextManager();
+  const span = {
+    traceIdB64: ctx.traceIdB64,
+    spanIdB64: ctx.spanIdB64,
+    eventId: ctx.eventId
+  };
+  return cm.runInContext(span, fn);
+}
+function isAsyncIterable(value) {
+  return value !== null && typeof value === "object" && Symbol.asyncIterator in value;
+}
+function firstFiniteNumber(...values) {
+  for (const value of values) {
+    if (typeof value === "number" && Number.isFinite(value)) {
+      return value;
+    }
+  }
+  return void 0;
+}
+function resolveUsageRecord(result) {
+  if (!isRecord(result)) return void 0;
+  let usage;
+  try {
+    if (isRecord(result["totalUsage"])) {
+      usage = result["totalUsage"];
+    } else if (isRecord(result["usage"])) {
+      usage = result["usage"];
+    }
+  } catch (e) {
+    return void 0;
+  }
+  return isRecord(usage) ? usage : void 0;
+}
+function extractUsageMetrics(result) {
+  const usage = resolveUsageRecord(result);
+  if (!usage) return {};
+  const inputTokenValue = usage["inputTokens"];
+  const outputTokenValue = usage["outputTokens"];
+  const inputTokens = firstFiniteNumber(
+    isRecord(inputTokenValue) ? inputTokenValue["total"] : void 0,
+    inputTokenValue,
+    usage["promptTokens"],
+    usage["prompt_tokens"]
+  );
+  const outputTokens = firstFiniteNumber(
+    isRecord(outputTokenValue) ? outputTokenValue["total"] : void 0,
+    outputTokenValue,
+    usage["completionTokens"],
+    usage["completion_tokens"]
+  );
+  const totalTokens = firstFiniteNumber(
+    usage["totalTokens"],
+    usage["tokens"],
+    usage["total_tokens"],
+    inputTokens !== void 0 && outputTokens !== void 0 ? inputTokens + outputTokens : void 0
+  );
+  const reasoningTokens = firstFiniteNumber(
+    isRecord(outputTokenValue) ? outputTokenValue["reasoning"] : void 0,
+    usage["reasoningTokens"],
+    usage["completionReasoningTokens"],
+    usage["completion_reasoning_tokens"],
+    usage["reasoning_tokens"],
+    usage["thinkingTokens"],
+    usage["thinking_tokens"]
+  );
+  const cachedInputTokens = firstFiniteNumber(
+    isRecord(inputTokenValue) ? inputTokenValue["cacheRead"] : void 0,
+    usage["cachedInputTokens"],
+    usage["promptCachedTokens"],
+    usage["prompt_cached_tokens"]
+  );
+  return {
+    inputTokens,
+    outputTokens,
+    totalTokens,
+    reasoningTokens,
+    cachedInputTokens
+  };
+}
+function isObjectOperation(operation) {
+  return operation === "generateObject" || operation === "streamObject";
+}
+function logFinalizeError(debug, err) {
+  if (debug) {
+    console.warn(
+      `[raindrop-ai/ai-sdk] finalize failed: ${err instanceof Error ? err.message : err}`
+    );
+  }
+}
+async function safeFinalize(finalize, debug, result, error) {
+  try {
+    await finalize(result, error);
+  } catch (finalizeErr) {
+    logFinalizeError(debug, finalizeErr);
+  }
+}
+function runWithRootContextSync(rootSpan, eventId, fn) {
+  if (!rootSpan) return fn();
+  return runWithParentSpanContextSync(
+    { traceIdB64: rootSpan.ids.traceIdB64, spanIdB64: rootSpan.ids.spanIdB64, eventId },
+    fn
+  );
+}
+async function runWithRootContextAsync(rootSpan, eventId, fn) {
+  if (!rootSpan) return await fn();
+  return await runWithParentSpanContext(
+    { traceIdB64: rootSpan.ids.traceIdB64, spanIdB64: rootSpan.ids.spanIdB64, eventId },
+    fn
+  );
+}
+function teeStreamObjectBaseStream(result) {
+  if (!result || !isRecord(result)) return;
+  const baseStream = result["baseStream"];
+  if (!(baseStream && typeof baseStream === "object" && "tee" in baseStream)) return;
+  try {
+    const [consumeStream, userStream] = baseStream.tee();
+    result["baseStream"] = userStream;
+    consumeStream.pipeTo(new WritableStream({ write() {
+    } })).catch(() => {
+    });
+  } catch (e) {
+  }
+}
+function setupOperation(params) {
+  var _a, _b, _c;
+  const { operation, arg, inherited, aiSDK, options, traceShipper, sendTraces } = params;
+  const wrapTimeCtx = resolveContext(options.context, { operation, args: arg });
+  const telemetry = extractExperimentalTelemetry(arg);
+  const callTimeCtx = extractRaindropMetadata(telemetry == null ? void 0 : telemetry.metadata);
+  const mergedCtx = mergeContexts(wrapTimeCtx, callTimeCtx);
+  const eventId = (_c = (_b = (_a = callTimeCtx.eventId) != null ? _a : mergedCtx.eventId) != null ? _b : inherited == null ? void 0 : inherited.eventId) != null ? _c : randomUUID();
+  const ctx = { ...mergedCtx, eventId };
+  const inheritedParent = inherited && inherited.eventId === eventId ? { traceIdB64: inherited.traceIdB64, spanIdB64: inherited.spanIdB64 } : void 0;
+  const outerOperationId = `ai.${operation}`;
+  const { operationName, resourceName } = opName(outerOperationId, telemetry == null ? void 0 : telemetry.functionId);
+  const modelInfoFromArgs = isRecord(arg) ? extractModelInfo(arg["model"]) : {};
+  const rootSpan = sendTraces ? traceShipper.startSpan({
+    name: outerOperationId,
+    parent: inheritedParent,
+    eventId,
+    operationId: outerOperationId,
+    attributes: [
+      attrString("operation.name", operationName),
+      attrString("resource.name", resourceName),
+      attrString("ai.telemetry.functionId", telemetry == null ? void 0 : telemetry.functionId),
+      attrString("ai.model.provider", modelInfoFromArgs.provider),
+      attrString("ai.model.id", modelInfoFromArgs.modelId),
+      ...attrsFromTelemetryMetadata(telemetry == null ? void 0 : telemetry.metadata),
+      ...attrsFromHeaders(isRecord(arg) ? arg["headers"] : void 0),
+      ...attrsFromSettings(arg),
+      ...(telemetry == null ? void 0 : telemetry.recordInputs) === false ? [] : [
+        attrString(
+          "ai.prompt",
+          safeJsonWithUint8({
+            system: isRecord(arg) ? arg["system"] : void 0,
+            prompt: isRecord(arg) ? arg["prompt"] : void 0,
+            messages: isRecord(arg) ? arg["messages"] : void 0
+          })
+        )
+      ]
+    ]
+  }) : void 0;
+  const rootParentForChildren = rootSpan ? { traceIdB64: rootSpan.ids.traceIdB64, spanIdB64: rootSpan.ids.spanIdB64 } : inheritedParent;
+  const wrapCtx = {
+    eventId,
+    telemetry,
+    sendTraces,
+    traceShipper,
+    rootParentForChildren
+  };
+  const toolCalls = [];
+  const argsWithWrappedTools = wrapTools(arg, wrapCtx, toolCalls);
+  const wrappedArgs = wrapModel(argsWithWrappedTools, aiSDK, outerOperationId, wrapCtx);
+  return {
+    eventId,
+    ctx,
+    telemetry,
+    rootSpan,
+    wrappedArgs,
+    toolCalls
+  };
+}
+function createFinalize(params) {
+  const {
+    operation,
+    arg,
+    setup,
+    autoAttachmentEnabled,
+    sendEvents,
+    debug,
+    options,
+    eventShipper,
+    traceShipper
+  } = params;
+  return async (result, error) => {
+    var _a, _b, _c, _d;
+    const usage = extractUsageMetrics(result);
+    const model = extractModel(result);
+    const inputAttachments = autoAttachmentEnabled ? extractInputAttachmentsFromArgs(arg) : void 0;
+    const outputAttachments = autoAttachmentEnabled ? await extractOutputAttachmentsFromResult(result) : void 0;
+    const baseMessages = coerceMessagesFromArgs(arg);
+    const responseMessages = extractResponseMessages(result);
+    const allMessages = [...baseMessages, ...responseMessages];
+    const outputText = extractTextOutput(result);
+    const outputObjectJson = extractObjectOutput(result);
+    const defaultOutput = isObjectOperation(operation) ? outputObjectJson : outputText;
+    const defaultPatch = {
+      eventName: (_a = setup.ctx.eventName) != null ? _a : operation,
+      input: (_b = lastUserMessageTextFromArgs(arg)) != null ? _b : extractInputFromArgs(arg),
+      output: defaultOutput,
+      model,
+      properties: setup.ctx.properties,
+      attachments: mergeAttachments(setup.ctx.attachments, inputAttachments, outputAttachments)
+    };
+    const built = await maybeBuildEvent(options.buildEvent, allMessages);
+    const patch = mergeBuildEventPatch(defaultPatch, built);
+    const output = patch.output;
+    const finalModel = (_c = patch.model) != null ? _c : model;
+    if (setup.rootSpan) {
+      const finishReason = extractFinishReason(result);
+      const providerMetadata = isRecord(result) ? result["providerMetadata"] : void 0;
+      const resultToolCalls = isRecord(result) && Array.isArray(result["toolCalls"]) ? safeJsonWithUint8(result["toolCalls"]) : setup.toolCalls.length ? safeJsonWithUint8(setup.toolCalls) : void 0;
+      traceShipper.endSpan(setup.rootSpan, {
+        attributes: [
+          ...((_d = setup.telemetry) == null ? void 0 : _d.recordOutputs) === false ? [] : [
+            attrString("ai.response.finishReason", finishReason),
+            isObjectOperation(operation) ? attrString("ai.response.object", output) : attrString("ai.response.text", output),
+            attrString("ai.response.toolCalls", resultToolCalls),
+            attrString("ai.response.providerMetadata", safeJsonWithUint8(providerMetadata))
+          ],
+          attrInt("ai.usage.promptTokens", usage == null ? void 0 : usage.inputTokens),
+          attrInt("ai.usage.completionTokens", usage == null ? void 0 : usage.outputTokens),
+          attrInt("ai.usage.inputTokens", usage == null ? void 0 : usage.inputTokens),
+          attrInt("ai.usage.outputTokens", usage == null ? void 0 : usage.outputTokens),
+          attrInt("ai.usage.totalTokens", usage == null ? void 0 : usage.totalTokens),
+          attrInt("ai.usage.reasoningTokens", usage == null ? void 0 : usage.reasoningTokens),
+          attrInt("ai.usage.cachedInputTokens", usage == null ? void 0 : usage.cachedInputTokens),
+          attrInt("ai.toolCall.count", setup.toolCalls.length),
+          ...error ? [
+            attrString(
+              "error.message",
+              error instanceof Error ? error.message : String(error)
+            )
+          ] : []
+        ]
+      });
+    }
+    if (sendEvents) {
+      void eventShipper.patch(setup.eventId, {
+        eventName: patch.eventName,
+        userId: setup.ctx.userId,
+        convoId: setup.ctx.convoId,
+        input: patch.input,
+        output,
+        model: finalModel,
+        properties: patch.properties,
+        attachments: patch.attachments,
+        isPending: false
+      }).catch((err) => {
+        if (debug) {
+          console.warn(
+            `[raindrop-ai/ai-sdk] event patch failed: ${err instanceof Error ? err.message : err}`
+          );
+        }
+      });
+    }
+  };
+}
+function executeStreamingOperation(params) {
+  const {
+    operation,
+    arg,
+    callArgs,
+    aiSDK,
+    original,
+    deps,
+    sendEvents,
+    sendTraces,
+    autoAttachmentEnabled,
+    debug
+  } = params;
+  const setup = setupOperation({
+    operation,
+    arg,
+    inherited: getCurrentParentSpanContextSync(),
+    aiSDK,
+    options: deps.options,
+    traceShipper: deps.traceShipper,
+    sendTraces
+  });
+  const finalize = createFinalize({
+    operation,
+    arg,
+    setup,
+    autoAttachmentEnabled,
+    sendEvents,
+    debug,
+    options: deps.options,
+    eventShipper: deps.eventShipper,
+    traceShipper: deps.traceShipper
+  });
+  const argWithOnFinish = wrapOnFinish(setup.wrappedArgs, async (result) => {
+    await safeFinalize(finalize, debug, result);
+  });
+  const callOriginal = (...args) => {
+    return original.call(aiSDK, ...args);
+  };
+  try {
+    const result = runWithRootContextSync(setup.rootSpan, setup.eventId, () => {
+      const nextArgs = [...callArgs];
+      nextArgs[0] = argWithOnFinish;
+      return callOriginal(...nextArgs);
+    });
+    if (operation === "streamObject") {
+      teeStreamObjectBaseStream(result);
+    }
+    return result;
+  } catch (error) {
+    void safeFinalize(finalize, debug, void 0, error);
+    throw error;
+  }
+}
+async function executeNonStreamingOperation(params) {
+  const {
+    operation,
+    arg,
+    callArgs,
+    aiSDK,
+    original,
+    deps,
+    sendEvents,
+    sendTraces,
+    autoAttachmentEnabled,
+    debug
+  } = params;
+  const inherited = await getCurrentParentSpanContext();
+  const setup = setupOperation({
+    operation,
+    arg,
+    inherited,
+    aiSDK,
+    options: deps.options,
+    traceShipper: deps.traceShipper,
+    sendTraces
+  });
+  const finalize = createFinalize({
+    operation,
+    arg,
+    setup,
+    autoAttachmentEnabled,
+    sendEvents,
+    debug,
+    options: deps.options,
+    eventShipper: deps.eventShipper,
+    traceShipper: deps.traceShipper
+  });
+  const callOriginal = async (...args) => {
+    return await original.call(aiSDK, ...args);
+  };
+  try {
+    const result = await runWithRootContextAsync(setup.rootSpan, setup.eventId, async () => {
+      const nextArgs = [...callArgs];
+      nextArgs[0] = setup.wrappedArgs;
+      return await callOriginal(...nextArgs);
+    });
+    await safeFinalize(finalize, debug, result);
+    return result;
+  } catch (error) {
+    await safeFinalize(finalize, debug, void 0, error);
+    throw error;
+  }
+}
 function wrapAISDK(aiSDK, deps) {
   var _a, _b;
   const debug = deps.eventShipper.isDebugEnabled() || deps.traceShipper.isDebugEnabled();
-  const instrumentedOps = /* @__PURE__ */ new Set(["generateText", "streamText", "generateObject", "streamObject"]);
+  const instrumentedOps = /* @__PURE__ */ new Set([
+    "generateText",
+    "streamText",
+    "generateObject",
+    "streamObject"
+  ]);
   const agentClasses = /* @__PURE__ */ new Set(["ToolLoopAgent"]);
   const sendEvents = ((_a = deps.options.send) == null ? void 0 : _a.events) !== false;
   const sendTraces = ((_b = deps.options.send) == null ? void 0 : _b.traces) !== false;
@@ -1316,224 +1675,35 @@ function wrapAISDK(aiSDK, deps) {
       if (typeof prop !== "string" || !instrumentedOps.has(prop) || !isFunction(original)) {
         return original;
       }
-      return async (...callArgs) => {
-        var _a2, _b2, _c;
-        const arg = callArgs[0];
+      return (...callArgs) => {
         const operation = prop;
-        const wrapTimeCtx = resolveContext(deps.options.context, { operation, args: arg });
-        const telemetry = extractExperimentalTelemetry(arg);
-        const callTimeCtx = extractRaindropMetadata(telemetry == null ? void 0 : telemetry.metadata);
-        const mergedCtx = mergeContexts(wrapTimeCtx, callTimeCtx);
-        const inherited = await getCurrentParentSpanContext();
-        const eventId = (_c = (_b2 = (_a2 = callTimeCtx.eventId) != null ? _a2 : mergedCtx.eventId) != null ? _b2 : inherited == null ? void 0 : inherited.eventId) != null ? _c : randomUUID();
-        const ctx = { ...mergedCtx};
-        const inheritedParent = inherited && inherited.eventId === eventId ? { traceIdB64: inherited.traceIdB64, spanIdB64: inherited.spanIdB64 } : void 0;
-        const outerOperationId = `ai.${operation}`;
-        const { operationName, resourceName } = opName(outerOperationId, telemetry == null ? void 0 : telemetry.functionId);
-        const modelInfoFromArgs = isRecord(arg) ? extractModelInfo(arg["model"]) : {};
-        const rootSpan = sendTraces ? deps.traceShipper.startSpan({
-          name: outerOperationId,
-          parent: inheritedParent,
-          eventId,
-          operationId: outerOperationId,
-          attributes: [
-            attrString("operation.name", operationName),
-            attrString("resource.name", resourceName),
-            attrString("ai.telemetry.functionId", telemetry == null ? void 0 : telemetry.functionId),
-            attrString("ai.model.provider", modelInfoFromArgs.provider),
-            attrString("ai.model.id", modelInfoFromArgs.modelId),
-            ...attrsFromTelemetryMetadata(telemetry == null ? void 0 : telemetry.metadata),
-            ...attrsFromHeaders(isRecord(arg) ? arg["headers"] : void 0),
-            ...attrsFromSettings(arg),
-            ...(telemetry == null ? void 0 : telemetry.recordInputs) === false ? [] : [
-              attrString(
-                "ai.prompt",
-                safeJsonWithUint8({
-                  system: isRecord(arg) ? arg["system"] : void 0,
-                  prompt: isRecord(arg) ? arg["prompt"] : void 0,
-                  messages: isRecord(arg) ? arg["messages"] : void 0
-                })
-              )
-            ]
-          ]
-        }) : void 0;
-        const rootParentForChildren = rootSpan ? { traceIdB64: rootSpan.ids.traceIdB64, spanIdB64: rootSpan.ids.spanIdB64 } : inheritedParent;
-        const wrapCtx = {
-          eventId,
-          telemetry,
-          sendTraces,
-          traceShipper: deps.traceShipper,
-          rootParentForChildren
-        };
-        const toolCalls = [];
-        const argWithWrappedTools = wrapTools(arg, wrapCtx, toolCalls);
-        const argWithWrappedModel = wrapModel(
-          argWithWrappedTools,
-          aiSDK,
-          outerOperationId,
-          wrapCtx
-        );
-        const finalize = async (result, error) => {
-          var _a3, _b3, _c2;
-          const usage = extractUsage(result);
-          const model = extractModel(result);
-          const inputAttachments = autoAttachmentEnabled ? extractInputAttachmentsFromArgs(arg) : void 0;
-          const outputAttachments = autoAttachmentEnabled ? await extractOutputAttachmentsFromResult(result) : void 0;
-          const baseMessages = coerceMessagesFromArgs(arg);
-          const responseMessages = extractResponseMessages(result);
-          const allMessages = [...baseMessages, ...responseMessages];
-          const outputText = extractTextOutput(result);
-          const outputObjectJson = extractObjectOutput(result);
-          const defaultOutput = operation === "generateObject" || operation === "streamObject" ? outputObjectJson : outputText;
-          const defaultPatch = {
-            eventName: (_a3 = ctx.eventName) != null ? _a3 : operation,
-            input: (_b3 = lastUserMessageTextFromArgs(arg)) != null ? _b3 : extractInputFromArgs(arg),
-            output: defaultOutput,
-            model,
-            properties: ctx.properties,
-            attachments: mergeAttachments(ctx.attachments, inputAttachments, outputAttachments)
-          };
-          const built = await maybeBuildEvent(deps.options.buildEvent, allMessages);
-          const patch = mergeBuildEventPatch(defaultPatch, built);
-          const output = patch.output;
-          const finalModel = (_c2 = patch.model) != null ? _c2 : model;
-          if (rootSpan) {
-            const finishReason = extractFinishReason(result);
-            const providerMetadata = isRecord(result) ? result["providerMetadata"] : void 0;
-            const resultToolCalls = isRecord(result) && Array.isArray(result["toolCalls"]) ? safeJsonWithUint8(result["toolCalls"]) : toolCalls.length ? safeJsonWithUint8(toolCalls) : void 0;
-            const usageRec = isRecord(result) && isRecord(result["usage"]) ? result["usage"] : void 0;
-            const totalTokens = typeof (usageRec == null ? void 0 : usageRec["totalTokens"]) === "number" ? usageRec["totalTokens"] : void 0;
-            const reasoningTokens = typeof (usageRec == null ? void 0 : usageRec["reasoningTokens"]) === "number" ? usageRec["reasoningTokens"] : void 0;
-            const cachedInputTokens = typeof (usageRec == null ? void 0 : usageRec["cachedInputTokens"]) === "number" ? usageRec["cachedInputTokens"] : void 0;
-            deps.traceShipper.endSpan(rootSpan, {
-              attributes: [
-                ...(telemetry == null ? void 0 : telemetry.recordOutputs) === false ? [] : [
-                  attrString("ai.response.finishReason", finishReason),
-                  operation === "generateObject" || operation === "streamObject" ? attrString("ai.response.object", output) : attrString("ai.response.text", output),
-                  attrString("ai.response.toolCalls", resultToolCalls),
-                  attrString(
-                    "ai.response.providerMetadata",
-                    safeJsonWithUint8(providerMetadata)
-                  )
-                ],
-                attrInt("ai.usage.promptTokens", usage == null ? void 0 : usage.inputTokens),
-                attrInt("ai.usage.completionTokens", usage == null ? void 0 : usage.outputTokens),
-                attrInt("ai.usage.inputTokens", usage == null ? void 0 : usage.inputTokens),
-                attrInt("ai.usage.outputTokens", usage == null ? void 0 : usage.outputTokens),
-                attrInt("ai.usage.totalTokens", totalTokens),
-                attrInt("ai.usage.reasoningTokens", reasoningTokens),
-                attrInt("ai.usage.cachedInputTokens", cachedInputTokens),
-                attrInt("ai.toolCall.count", toolCalls.length),
-                ...error ? [
-                  attrString(
-                    "error.message",
-                    error instanceof Error ? error.message : String(error)
-                  )
-                ] : []
-              ]
-            });
-          }
-          if (sendEvents) {
-            void deps.eventShipper.patch(eventId, {
-              eventName: patch.eventName,
-              userId: ctx.userId,
-              convoId: ctx.convoId,
-              input: patch.input,
-              output,
-              model: finalModel,
-              properties: patch.properties,
-              attachments: patch.attachments,
-              isPending: false
-            }).catch((err) => {
-              if (debug)
-                console.warn(
-                  `[raindrop-ai/ai-sdk] event patch failed: ${err instanceof Error ? err.message : err}`
-                );
-            });
-          }
-        };
-        const callOriginal = async (...args) => {
-          return await original.call(
-            aiSDK,
-            ...args
-          );
-        };
-        const runWithContext = async (fn) => {
-          if (!rootSpan) return await fn();
-          return await runWithParentSpanContext(
-            { traceIdB64: rootSpan.ids.traceIdB64, spanIdB64: rootSpan.ids.spanIdB64, eventId },
-            fn
-          );
-        };
+        const arg = callArgs[0];
         if (operation === "streamText" || operation === "streamObject") {
-          const argWithOnFinish = wrapOnFinish(argWithWrappedModel, async (result) => {
-            try {
-              await finalize(result);
-            } catch (err) {
-              if (debug)
-                console.warn(
-                  `[raindrop-ai/ai-sdk] finalize failed: ${err instanceof Error ? err.message : err}`
-                );
-            }
-          });
-          try {
-            const result = await runWithContext(async () => {
-              const nextArgs = [...callArgs];
-              nextArgs[0] = argWithOnFinish;
-              return await callOriginal(...nextArgs);
-            });
-            if (operation === "streamObject" && result && isRecord(result)) {
-              const baseStream = result["baseStream"];
-              if (baseStream && typeof baseStream === "object" && "tee" in baseStream) {
-                try {
-                  const [consumeStream, userStream] = baseStream.tee();
-                  result["baseStream"] = userStream;
-                  consumeStream.pipeTo(new WritableStream({ write() {
-                  } })).catch(() => {
-                  });
-                } catch (e) {
-                }
-              }
-            }
-            return result;
-          } catch (error) {
-            try {
-              await finalize(void 0, error);
-            } catch (err) {
-              if (debug)
-                console.warn(
-                  `[raindrop-ai/ai-sdk] finalize failed: ${err instanceof Error ? err.message : err}`
-                );
-            }
-            throw error;
-          }
-        }
-        try {
-          const result = await runWithContext(async () => {
-            const nextArgs = [...callArgs];
-            nextArgs[0] = argWithWrappedModel;
-            return await callOriginal(...nextArgs);
+          return executeStreamingOperation({
+            operation,
+            arg,
+            callArgs,
+            aiSDK,
+            original,
+            deps,
+            sendEvents,
+            sendTraces,
+            autoAttachmentEnabled,
+            debug
           });
-          try {
-            await finalize(result);
-          } catch (err) {
-            if (debug)
-              console.warn(
-                `[raindrop-ai/ai-sdk] finalize failed: ${err instanceof Error ? err.message : err}`
-              );
-          }
-          return result;
-        } catch (error) {
-          try {
-            await finalize(void 0, error);
-          } catch (err) {
-            if (debug)
-              console.warn(
-                `[raindrop-ai/ai-sdk] finalize failed: ${err instanceof Error ? err.message : err}`
-              );
-          }
-          throw error;
         }
+        return executeNonStreamingOperation({
+          operation,
+          arg,
+          callArgs,
+          aiSDK,
+          original,
+          deps,
+          sendEvents,
+          sendTraces,
+          autoAttachmentEnabled,
+          debug
+        });
       };
     }
   });
@@ -1634,10 +1804,17 @@ function wrapAgentGenerate(generate, instance, agentSettings, className, aiSDK,
       rootParentForChildren
     };
     const toolCalls = [];
-    const callParamsWithWrappedTools = callParams ? wrapTools(callParams, wrapCtx, toolCalls) : callParams;
+    const mergedArgsWithWrappedTools = wrapTools(mergedArgs, wrapCtx, toolCalls);
+    const mergedArgsWithWrappedModel = wrapModel(
+      mergedArgsWithWrappedTools,
+      aiSDK,
+      outerOperationId,
+      wrapCtx
+    );
+    const callParamsWithWrappedToolsAndModel = mergedArgsWithWrappedModel != null ? mergedArgsWithWrappedModel : {};
     const finalize = async (result, error) => {
       var _a3, _b3, _c2;
-      const usage = extractUsage(result);
+      const usage = extractUsageMetrics(result);
       const model = extractModel(result);
       const inputAttachments = autoAttachmentEnabled ? extractInputAttachmentsFromArgs(mergedArgs) : void 0;
       const outputAttachments = autoAttachmentEnabled ? await extractOutputAttachmentsFromResult(result) : void 0;
@@ -1661,10 +1838,6 @@ function wrapAgentGenerate(generate, instance, agentSettings, className, aiSDK,
         const finishReason = extractFinishReason(result);
         const providerMetadata = isRecord(result) ? result["providerMetadata"] : void 0;
         const resultToolCalls = isRecord(result) && Array.isArray(result["toolCalls"]) ? safeJsonWithUint8(result["toolCalls"]) : toolCalls.length ? safeJsonWithUint8(toolCalls) : void 0;
-        const usageRec = isRecord(result) && isRecord(result["totalUsage"]) ? result["totalUsage"] : isRecord(result) && isRecord(result["usage"]) ? result["usage"] : void 0;
-        const totalTokens = typeof (usageRec == null ? void 0 : usageRec["totalTokens"]) === "number" ? usageRec["totalTokens"] : void 0;
-        const reasoningTokens = typeof (usageRec == null ? void 0 : usageRec["reasoningTokens"]) === "number" ? usageRec["reasoningTokens"] : void 0;
-        const cachedInputTokens = typeof (usageRec == null ? void 0 : usageRec["cachedInputTokens"]) === "number" ? usageRec["cachedInputTokens"] : void 0;
         deps.traceShipper.endSpan(rootSpan, {
           attributes: [
             ...(telemetry == null ? void 0 : telemetry.recordOutputs) === false ? [] : [
@@ -1677,9 +1850,9 @@ function wrapAgentGenerate(generate, instance, agentSettings, className, aiSDK,
             attrInt("ai.usage.completionTokens", usage == null ? void 0 : usage.outputTokens),
             attrInt("ai.usage.inputTokens", usage == null ? void 0 : usage.inputTokens),
             attrInt("ai.usage.outputTokens", usage == null ? void 0 : usage.outputTokens),
-            attrInt("ai.usage.totalTokens", totalTokens),
-            attrInt("ai.usage.reasoningTokens", reasoningTokens),
-            attrInt("ai.usage.cachedInputTokens", cachedInputTokens),
+            attrInt("ai.usage.totalTokens", usage == null ? void 0 : usage.totalTokens),
+            attrInt("ai.usage.reasoningTokens", usage == null ? void 0 : usage.reasoningTokens),
+            attrInt("ai.usage.cachedInputTokens", usage == null ? void 0 : usage.cachedInputTokens),
             attrInt("ai.toolCall.count", toolCalls.length),
             ...error ? [
               attrString(
@@ -1735,7 +1908,7 @@ function wrapAgentGenerate(generate, instance, agentSettings, className, aiSDK,
     }
     try {
       const result = await runWithContext(async () => {
-        return await generate.call(instance, callParamsWithWrappedTools);
+        return await generate.call(instance, callParamsWithWrappedToolsAndModel);
       });
       if (debug) {
         console.log(`[raindrop-ai/ai-sdk] Agent ${operation} completed, finalizing...`);
@@ -1818,10 +1991,17 @@ function wrapAgentStream(stream, instance, agentSettings, className, aiSDK, deps
       rootParentForChildren
     };
     const toolCalls = [];
-    const callParamsWithWrappedTools = callParams ? wrapTools(callParams, wrapCtx, toolCalls) : callParams;
+    const mergedArgsWithWrappedTools = wrapTools(mergedArgs, wrapCtx, toolCalls);
+    const mergedArgsWithWrappedModel = wrapModel(
+      mergedArgsWithWrappedTools,
+      aiSDK,
+      outerOperationId,
+      wrapCtx
+    );
+    const callParamsWithWrappedToolsAndModel = mergedArgsWithWrappedModel != null ? mergedArgsWithWrappedModel : {};
     const finalize = async (result, error) => {
       var _a3, _b3, _c2;
-      const usage = extractUsage(result);
+      const usage = extractUsageMetrics(result);
       const model = extractModel(result);
       const inputAttachments = autoAttachmentEnabled ? extractInputAttachmentsFromArgs(mergedArgs) : void 0;
       const outputAttachments = autoAttachmentEnabled ? await extractOutputAttachmentsFromResult(result) : void 0;
@@ -1845,10 +2025,6 @@ function wrapAgentStream(stream, instance, agentSettings, className, aiSDK, deps
         const finishReason = extractFinishReason(result);
         const providerMetadata = isRecord(result) ? result["providerMetadata"] : void 0;
         const resultToolCalls = isRecord(result) && Array.isArray(result["toolCalls"]) ? safeJsonWithUint8(result["toolCalls"]) : toolCalls.length ? safeJsonWithUint8(toolCalls) : void 0;
-        const usageRec = isRecord(result) && isRecord(result["totalUsage"]) ? result["totalUsage"] : isRecord(result) && isRecord(result["usage"]) ? result["usage"] : void 0;
-        const totalTokens = typeof (usageRec == null ? void 0 : usageRec["totalTokens"]) === "number" ? usageRec["totalTokens"] : void 0;
-        const reasoningTokens = typeof (usageRec == null ? void 0 : usageRec["reasoningTokens"]) === "number" ? usageRec["reasoningTokens"] : void 0;
-        const cachedInputTokens = typeof (usageRec == null ? void 0 : usageRec["cachedInputTokens"]) === "number" ? usageRec["cachedInputTokens"] : void 0;
         deps.traceShipper.endSpan(rootSpan, {
           attributes: [
             ...(telemetry == null ? void 0 : telemetry.recordOutputs) === false ? [] : [
@@ -1861,9 +2037,9 @@ function wrapAgentStream(stream, instance, agentSettings, className, aiSDK, deps
             attrInt("ai.usage.completionTokens", usage == null ? void 0 : usage.outputTokens),
             attrInt("ai.usage.inputTokens", usage == null ? void 0 : usage.inputTokens),
             attrInt("ai.usage.outputTokens", usage == null ? void 0 : usage.outputTokens),
-            attrInt("ai.usage.totalTokens", totalTokens),
-            attrInt("ai.usage.reasoningTokens", reasoningTokens),
-            attrInt("ai.usage.cachedInputTokens", cachedInputTokens),
+            attrInt("ai.usage.totalTokens", usage == null ? void 0 : usage.totalTokens),
+            attrInt("ai.usage.reasoningTokens", usage == null ? void 0 : usage.reasoningTokens),
+            attrInt("ai.usage.cachedInputTokens", usage == null ? void 0 : usage.cachedInputTokens),
             attrInt("ai.toolCall.count", toolCalls.length),
             ...error ? [
               attrString(
@@ -1918,7 +2094,7 @@ function wrapAgentStream(stream, instance, agentSettings, className, aiSDK, deps
       });
     }
     const callParamsWithOnFinish = wrapOnFinish(
-      callParamsWithWrappedTools != null ? callParamsWithWrappedTools : {},
+      callParamsWithWrappedToolsAndModel != null ? callParamsWithWrappedToolsAndModel : {},
       async (result) => {
         if (debug) {
           console.log(`[raindrop-ai/ai-sdk] Agent ${operation} onFinish callback, finalizing...`);
@@ -2003,20 +2179,26 @@ function wrapToolExecute(name, tool, ctx, toolCalls) {
     spanIdB64: span.ids.spanIdB64,
     eventId: ctx.eventId
   });
-  const wrappedExecute = (...execArgs) => {
+  const wrappedExecute = function(...execArgs) {
     const toolArgs = execArgs[0];
     const execOptions = execArgs.length > 1 ? execArgs[1] : void 0;
     const toolCallId = isRecord(execOptions) && typeof execOptions["toolCallId"] === "string" ? execOptions["toolCallId"] : randomUUID();
-    const result = originalExecute(...execArgs);
-    if (isAsyncGenerator(result)) {
+    const result = originalExecute.apply(this, execArgs);
+    if (isAsyncIterable(result)) {
       return (async function* () {
         const parentCtx = await getCurrentParentSpanContext();
         const parent = parentCtx && parentCtx.eventId === ctx.eventId ? { traceIdB64: parentCtx.traceIdB64, spanIdB64: parentCtx.spanIdB64 } : ctx.rootParentForChildren;
         const toolSpan = createToolSpan(toolCallId, toolArgs, parent);
         try {
           let lastValue;
-          const wrappedGen = toolSpan ? asyncGeneratorWithCurrent(createContextSpan(toolSpan), result) : result;
-          for await (const value of wrappedGen) {
+          const iterator = result[Symbol.asyncIterator]();
+          const wrappedIterable = toolSpan ? asyncGeneratorWithCurrent(
+            createContextSpan(toolSpan),
+            iterator
+          ) : {
+            [Symbol.asyncIterator]: () => iterator
+          };
+          for await (const value of wrappedIterable) {
             lastValue = value;
             yield value;
           }
@@ -2362,14 +2544,19 @@ function wrapOnFinish(args, onFinish) {
   return {
     ...args,
     onFinish: async (result) => {
+      let userError;
       try {
         const maybePromise = existing(result);
         if (maybePromise && typeof maybePromise.then === "function") {
           await maybePromise;
         }
-      } catch (e) {
+      } catch (error) {
+        userError = error;
       }
       await onFinish(result);
+      if (userError !== void 0) {
+        throw userError;
+      }
     }
   };
 }