@sentrial/sdk 0.3.2 → 0.3.3

package/dist/index.js

@@ -1073,13 +1073,8 @@ function wrapStreamText(originalFn, client, config) {
         return null;
       }
     })();
-    const wrappedParams = {
-      ...params,
-      tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
-    };
-    const result = originalFn(wrappedParams);
     let tracked = false;
-    async function trackCompletion(fullText, error) {
+    async function trackCompletion(fullText, usageInfo, error) {
       if (tracked) return;
       tracked = true;
       const durationMs = Date.now() - startTime;
@@ -1099,14 +1094,9 @@ function wrapStreamText(originalFn, client, config) {
         });
         return;
       }
-      let usage;
-      try {
-        usage = result.usage ? await result.usage : void 0;
-      } catch {
-      }
-      const promptTokens = usage?.promptTokens || 0;
-      const completionTokens = usage?.completionTokens || 0;
-      const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
+      const promptTokens = usageInfo?.promptTokens || 0;
+      const completionTokens = usageInfo?.completionTokens || 0;
+      const totalTokens = usageInfo?.totalTokens || promptTokens + completionTokens;
       const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
       await client.completeSession({
         sessionId: sid,
@@ -1119,48 +1109,30 @@ function wrapStreamText(originalFn, client, config) {
         totalTokens
       });
     }
-    const originalTextStream = result.textStream;
-    let wrappedTextStream = null;
-    const proxied = new Proxy(result, {
-      get(target, prop, receiver) {
-        if (prop === "textStream") {
-          if (!wrappedTextStream) {
-            let fullText = "";
-            wrappedTextStream = (async function* () {
-              try {
-                for await (const chunk of originalTextStream) {
-                  fullText += chunk;
-                  yield chunk;
-                }
-                await trackCompletion(fullText);
-              } catch (error) {
-                await trackCompletion(
-                  fullText,
-                  error instanceof Error ? error : new Error(String(error))
-                );
-                throw error;
-              }
-            })();
-          }
-          return wrappedTextStream;
+    const userOnFinish = params.onFinish;
+    const wrappedParams = {
+      ...params,
+      tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0,
+      onFinish: async (event) => {
+        try {
+          if (userOnFinish) await userOnFinish(event);
+        } catch {
         }
-        return Reflect.get(target, prop, receiver);
+        await trackCompletion(event.text, event.usage);
       }
-    });
+    };
+    const result = originalFn(wrappedParams);
     const textProp = result.text;
-    if (typeof textProp === "string") {
-      trackCompletion(textProp).catch(() => {
-      });
-    } else if (textProp != null && typeof textProp.then === "function") {
+    if (textProp != null && typeof textProp.then === "function") {
       textProp.then((text) => {
         trackCompletion(text).catch(() => {
         });
       }).catch((err) => {
-        trackCompletion("", err instanceof Error ? err : new Error(String(err))).catch(() => {
+        trackCompletion("", void 0, err instanceof Error ? err : new Error(String(err))).catch(() => {
         });
       });
     }
-    return proxied;
+    return result;
   };
 }
 function wrapGenerateObject(originalFn, client, config) {
@@ -1246,68 +1218,52 @@ function wrapStreamObject(originalFn, client, config) {
         return null;
       }
     })();
-    const result = originalFn(params);
-    const originalObjectPromise = result.object;
-    if (originalObjectPromise && typeof originalObjectPromise.then === "function") {
-      let tracked = false;
-      let wrappedObjectPromise = null;
-      return new Proxy(result, {
-        get(target, prop, receiver) {
-          if (prop === "object") {
-            if (!wrappedObjectPromise) {
-              wrappedObjectPromise = originalObjectPromise.then(async (obj) => {
-                if (tracked) return obj;
-                tracked = true;
-                const durationMs = Date.now() - startTime;
-                const sid = await sessionPromise;
-                if (!sid) return obj;
-                let usage;
-                try {
-                  usage = result.usage ? await result.usage : void 0;
-                } catch {
-                }
-                const promptTokens = usage?.promptTokens || 0;
-                const completionTokens = usage?.completionTokens || 0;
-                const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
-                const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
-                await client.completeSession({
-                  sessionId: sid,
-                  success: true,
-                  output: JSON.stringify(obj),
-                  durationMs,
-                  estimatedCost: cost,
-                  promptTokens,
-                  completionTokens,
-                  totalTokens
-                });
-                return obj;
-              }).catch(async (error) => {
-                if (tracked) throw error;
-                tracked = true;
-                const durationMs = Date.now() - startTime;
-                const sid = await sessionPromise;
-                if (!sid) throw error;
-                await client.trackError({
-                  sessionId: sid,
-                  errorType: error instanceof Error ? error.name : "Error",
-                  errorMessage: error instanceof Error ? error.message : "Unknown error"
-                });
-                await client.completeSession({
-                  sessionId: sid,
-                  success: false,
-                  failureReason: error instanceof Error ? error.message : "Unknown error",
-                  durationMs
-                });
-                throw error;
-              });
-            }
-            return wrappedObjectPromise;
-          }
-          return Reflect.get(target, prop, receiver);
+    let tracked = false;
+    const userOnFinish = params.onFinish;
+    const wrappedParams = {
+      ...params,
+      onFinish: async (event) => {
+        try {
+          if (userOnFinish) await userOnFinish(event);
+        } catch {
         }
-      });
-    }
-    return result;
+        if (tracked) return;
+        tracked = true;
+        const durationMs = Date.now() - startTime;
+        const sid = await sessionPromise;
+        if (!sid) return;
+        if (event.error) {
+          const err = event.error instanceof Error ? event.error : new Error(String(event.error));
+          await client.trackError({
+            sessionId: sid,
+            errorType: err.name || "Error",
+            errorMessage: err.message || "Unknown error"
+          });
+          await client.completeSession({
+            sessionId: sid,
+            success: false,
+            failureReason: err.message || "Unknown error",
+            durationMs
+          });
+          return;
+        }
+        const promptTokens = event.usage?.promptTokens || 0;
+        const completionTokens = event.usage?.completionTokens || 0;
+        const totalTokens = event.usage?.totalTokens || promptTokens + completionTokens;
+        const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
+        await client.completeSession({
+          sessionId: sid,
+          success: true,
+          output: event.object != null ? JSON.stringify(event.object) : "",
+          durationMs,
+          estimatedCost: cost,
+          promptTokens,
+          completionTokens,
+          totalTokens
+        });
+      }
+    };
+    return originalFn(wrappedParams);
   };
 }
 function wrapAISDK(ai, options) {