@sentrial/sdk 0.4.2 → 0.4.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -364,8 +364,18 @@ function wrapOpenAI(client, options = {}) {
364
364
  const completionTokens = response.usage?.completion_tokens ?? 0;
365
365
  const totalTokens = response.usage?.total_tokens ?? 0;
366
366
  let outputContent = "";
367
- if (response.choices?.[0]?.message?.content) {
368
- outputContent = response.choices[0].message.content;
367
+ const toolCalls = [];
368
+ const msg = response.choices?.[0]?.message;
369
+ if (msg?.content) {
370
+ outputContent = msg.content;
371
+ }
372
+ if (msg?.tool_calls) {
373
+ for (const tc of msg.tool_calls) {
374
+ toolCalls.push({
375
+ name: tc.function?.name ?? "unknown",
376
+ arguments: tc.function?.arguments ?? "{}"
377
+ });
378
+ }
369
379
  }
370
380
  const cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
371
381
  trackLLMCall({
@@ -373,6 +383,7 @@ function wrapOpenAI(client, options = {}) {
373
383
  model,
374
384
  messages,
375
385
  output: outputContent,
386
+ toolCalls,
376
387
  promptTokens,
377
388
  completionTokens,
378
389
  totalTokens,
@@ -427,10 +438,16 @@ function wrapAnthropic(client, options = {}) {
427
438
  const completionTokens = response.usage?.output_tokens ?? 0;
428
439
  const totalTokens = promptTokens + completionTokens;
429
440
  let outputContent = "";
441
+ const toolCalls = [];
430
442
  if (response.content) {
431
443
  for (const block of response.content) {
432
444
  if (block.type === "text") {
433
445
  outputContent += block.text;
446
+ } else if (block.type === "tool_use") {
447
+ toolCalls.push({
448
+ name: block.name ?? "unknown",
449
+ arguments: JSON.stringify(block.input ?? {})
450
+ });
434
451
  }
435
452
  }
436
453
  }
@@ -441,6 +458,7 @@ function wrapAnthropic(client, options = {}) {
441
458
  model,
442
459
  messages: fullMessages,
443
460
  output: outputContent,
461
+ toolCalls,
444
462
  promptTokens,
445
463
  completionTokens,
446
464
  totalTokens,
@@ -554,6 +572,7 @@ function wrapLLM(client, provider) {
554
572
  function wrapOpenAIStream(stream, ctx) {
555
573
  let fullContent = "";
556
574
  let usage = null;
575
+ const toolCallMap = /* @__PURE__ */ new Map();
557
576
  let tracked = false;
558
577
  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
559
578
  if (!originalIterator) return stream;
@@ -570,6 +589,7 @@ function wrapOpenAIStream(stream, ctx) {
570
589
  model: ctx.model,
571
590
  messages: ctx.messages,
572
591
  output: fullContent,
592
+ toolCalls: Array.from(toolCallMap.values()),
573
593
  promptTokens,
574
594
  completionTokens,
575
595
  totalTokens,
@@ -588,8 +608,17 @@ function wrapOpenAIStream(stream, ctx) {
588
608
  const result = await iter.next();
589
609
  if (!result.done) {
590
610
  const chunk = result.value;
591
- const delta = chunk.choices?.[0]?.delta?.content;
592
- if (delta) fullContent += delta;
611
+ const delta = chunk.choices?.[0]?.delta;
612
+ if (delta?.content) fullContent += delta.content;
613
+ if (delta?.tool_calls) {
614
+ for (const tc of delta.tool_calls) {
615
+ const idx = tc.index ?? 0;
616
+ const existing = toolCallMap.get(idx) ?? { name: "", arguments: "" };
617
+ if (tc.function?.name) existing.name = tc.function.name;
618
+ if (tc.function?.arguments) existing.arguments += tc.function.arguments;
619
+ toolCallMap.set(idx, existing);
620
+ }
621
+ }
593
622
  if (chunk.usage) usage = chunk.usage;
594
623
  } else {
595
624
  trackResult();
@@ -614,6 +643,8 @@ function wrapAnthropicStream(stream, ctx) {
614
643
  let fullContent = "";
615
644
  let inputTokens = 0;
616
645
  let outputTokens = 0;
646
+ const toolCallsById = /* @__PURE__ */ new Map();
647
+ let currentBlockIdx = -1;
617
648
  let tracked = false;
618
649
  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
619
650
  if (!originalIterator) return stream;
@@ -629,6 +660,7 @@ function wrapAnthropicStream(stream, ctx) {
629
660
  model: ctx.model,
630
661
  messages: fullMessages,
631
662
  output: fullContent,
663
+ toolCalls: Array.from(toolCallsById.values()),
632
664
  promptTokens: inputTokens,
633
665
  completionTokens: outputTokens,
634
666
  totalTokens,
@@ -647,8 +679,26 @@ function wrapAnthropicStream(stream, ctx) {
647
679
  const result = await iter.next();
648
680
  if (!result.done) {
649
681
  const event = result.value;
650
- if (event.type === "content_block_delta" && event.delta?.text) {
651
- fullContent += event.delta.text;
682
+ if (event.type === "content_block_start") {
683
+ currentBlockIdx = event.index ?? currentBlockIdx + 1;
684
+ if (event.content_block?.type === "tool_use") {
685
+ toolCallsById.set(currentBlockIdx, {
686
+ name: event.content_block.name ?? "unknown",
687
+ arguments: ""
688
+ });
689
+ }
690
+ }
691
+ if (event.type === "content_block_delta") {
692
+ if (event.delta?.text) {
693
+ fullContent += event.delta.text;
694
+ }
695
+ if (event.delta?.type === "input_json_delta" && event.delta?.partial_json) {
696
+ const idx = event.index ?? currentBlockIdx;
697
+ const existing = toolCallsById.get(idx);
698
+ if (existing) {
699
+ existing.arguments += event.delta.partial_json;
700
+ }
701
+ }
652
702
  }
653
703
  if (event.type === "message_start" && event.message?.usage) {
654
704
  inputTokens = event.message.usage.input_tokens ?? 0;
@@ -682,6 +732,18 @@ function trackLLMCall(params) {
682
732
  if (!sessionId && !params.trackWithoutSession) {
683
733
  return;
684
734
  }
735
+ const toolOutput = {
736
+ content: params.output,
737
+ tokens: {
738
+ prompt: params.promptTokens,
739
+ completion: params.completionTokens,
740
+ total: params.totalTokens
741
+ },
742
+ cost_usd: params.cost
743
+ };
744
+ if (params.toolCalls && params.toolCalls.length > 0) {
745
+ toolOutput.tool_calls = params.toolCalls;
746
+ }
685
747
  if (sessionId) {
686
748
  client.trackToolCall({
687
749
  sessionId,
@@ -691,15 +753,7 @@ function trackLLMCall(params) {
691
753
  model: params.model,
692
754
  provider: params.provider
693
755
  },
694
- toolOutput: {
695
- content: params.output,
696
- tokens: {
697
- prompt: params.promptTokens,
698
- completion: params.completionTokens,
699
- total: params.totalTokens
700
- },
701
- cost_usd: params.cost
702
- },
756
+ toolOutput,
703
757
  reasoning: `LLM call to ${params.provider} ${params.model}`,
704
758
  estimatedCost: params.cost,
705
759
  tokenCount: params.totalTokens,
@@ -728,15 +782,7 @@ function trackLLMCall(params) {
728
782
  model: params.model,
729
783
  provider: params.provider
730
784
  },
731
- toolOutput: {
732
- content: params.output,
733
- tokens: {
734
- prompt: params.promptTokens,
735
- completion: params.completionTokens,
736
- total: params.totalTokens
737
- },
738
- cost_usd: params.cost
739
- },
785
+ toolOutput,
740
786
  estimatedCost: params.cost,
741
787
  tokenCount: params.totalTokens,
742
788
  metadata: {
@@ -1559,7 +1605,8 @@ function getClient2() {
1559
1605
  }
1560
1606
  function extractModelInfo(model) {
1561
1607
  const modelId = model.modelId || model.id || "unknown";
1562
- const provider = model.provider || guessProvider(modelId);
1608
+ const rawProvider = model.provider || "";
1609
+ const provider = rawProvider.split(".")[0] || guessProvider(modelId);
1563
1610
  return { modelId, provider };
1564
1611
  }
1565
1612
  function guessProvider(modelId) {
@@ -1604,6 +1651,15 @@ function extractInput(params) {
1604
1651
  }
1605
1652
  return "";
1606
1653
  }
1654
+ function normalizeUsage(usage) {
1655
+ if (!usage) return void 0;
1656
+ const u = usage;
1657
+ const promptTokens = (u.inputTokens ?? u.promptTokens ?? 0) || 0;
1658
+ const completionTokens = (u.outputTokens ?? u.completionTokens ?? 0) || 0;
1659
+ const totalTokens = (u.totalTokens ?? promptTokens + completionTokens) || 0;
1660
+ if (promptTokens === 0 && completionTokens === 0 && totalTokens === 0) return void 0;
1661
+ return { promptTokens, completionTokens, totalTokens };
1662
+ }
1607
1663
  function wrapTools(tools, sessionId, client) {
1608
1664
  if (!tools) return void 0;
1609
1665
  const wrappedTools = {};
@@ -1656,34 +1712,37 @@ function wrapToolsAsync(tools, sessionPromise, client) {
1656
1712
  ...tool,
1657
1713
  execute: async (...args) => {
1658
1714
  const startTime = Date.now();
1659
- const sid = await sessionPromise;
1660
1715
  try {
1661
1716
  const result = await originalExecute(...args);
1662
1717
  const durationMs = Date.now() - startTime;
1663
- if (sid) {
1664
- client.trackToolCall({
1665
- sessionId: sid,
1666
- toolName,
1667
- toolInput: args[0],
1668
- toolOutput: result,
1669
- reasoning: `Tool executed in ${durationMs}ms`
1670
- }).catch(() => {
1671
- });
1672
- }
1718
+ sessionPromise.then((sid) => {
1719
+ if (sid) {
1720
+ client.trackToolCall({
1721
+ sessionId: sid,
1722
+ toolName,
1723
+ toolInput: args[0],
1724
+ toolOutput: result,
1725
+ reasoning: `Tool executed in ${durationMs}ms`
1726
+ }).catch(() => {
1727
+ });
1728
+ }
1729
+ });
1673
1730
  return result;
1674
1731
  } catch (error) {
1675
1732
  const durationMs = Date.now() - startTime;
1676
- if (sid) {
1677
- client.trackToolCall({
1678
- sessionId: sid,
1679
- toolName,
1680
- toolInput: args[0],
1681
- toolOutput: {},
1682
- toolError: { message: error instanceof Error ? error.message : "Unknown error" },
1683
- reasoning: `Tool failed after ${durationMs}ms`
1684
- }).catch(() => {
1685
- });
1686
- }
1733
+ sessionPromise.then((sid) => {
1734
+ if (sid) {
1735
+ client.trackToolCall({
1736
+ sessionId: sid,
1737
+ toolName,
1738
+ toolInput: args[0],
1739
+ toolOutput: {},
1740
+ toolError: { message: error instanceof Error ? error.message : "Unknown error" },
1741
+ reasoning: `Tool failed after ${durationMs}ms`
1742
+ }).catch(() => {
1743
+ });
1744
+ }
1745
+ });
1687
1746
  throw error;
1688
1747
  }
1689
1748
  }
@@ -1723,44 +1782,50 @@ function wrapGenerateText(originalFn, client, config) {
1723
1782
  const result = await originalFn(wrappedParams);
1724
1783
  const durationMs = Date.now() - startTime;
1725
1784
  const resolvedModelId = result.response?.modelId || modelId;
1726
- const promptTokens = result.usage?.promptTokens ?? 0;
1727
- const completionTokens = result.usage?.completionTokens ?? 0;
1728
- const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
1785
+ const usage = normalizeUsage(result.usage);
1786
+ const promptTokens = usage?.promptTokens ?? 0;
1787
+ const completionTokens = usage?.completionTokens ?? 0;
1788
+ const totalTokens = usage?.totalTokens ?? 0;
1729
1789
  const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1730
1790
  const steps = result.steps;
1731
1791
  if (steps && steps.length >= 1) {
1732
- const stepPromises = steps.map(
1733
- (step, i) => client.trackEvent({
1792
+ const stepPromises = steps.map((step) => {
1793
+ const su = normalizeUsage(step.usage);
1794
+ return client.trackToolCall({
1734
1795
  sessionId,
1735
- eventType: "llm_call",
1736
- eventData: {
1737
- model: resolvedModelId,
1738
- provider,
1739
- step: i + 1,
1740
- total_steps: steps.length,
1741
- prompt_tokens: step.usage?.promptTokens ?? 0,
1742
- completion_tokens: step.usage?.completionTokens ?? 0,
1743
- total_tokens: step.usage?.totalTokens ?? 0,
1744
- finish_reason: step.finishReason,
1796
+ toolName: `llm:${provider}:${resolvedModelId}`,
1797
+ toolInput: { finishReason: step.finishReason },
1798
+ toolOutput: {
1799
+ text: step.text?.slice(0, 500),
1800
+ tokens: {
1801
+ prompt: su?.promptTokens ?? 0,
1802
+ completion: su?.completionTokens ?? 0
1803
+ }
1804
+ },
1805
+ estimatedCost: calculateCostForCall(provider, resolvedModelId, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
1806
+ tokenCount: su?.totalTokens,
1807
+ metadata: {
1745
1808
  tool_calls: step.toolCalls?.map((tc) => tc.toolName)
1746
1809
  }
1747
1810
  }).catch(() => {
1748
- })
1749
- );
1811
+ });
1812
+ });
1750
1813
  await Promise.all(stepPromises);
1751
1814
  } else {
1752
- await client.trackEvent({
1815
+ await client.trackToolCall({
1753
1816
  sessionId,
1754
- eventType: "llm_call",
1755
- eventData: {
1756
- model: resolvedModelId,
1757
- provider,
1758
- prompt_tokens: promptTokens,
1759
- completion_tokens: completionTokens,
1760
- total_tokens: totalTokens,
1761
- finish_reason: result.finishReason,
1817
+ toolName: `llm:${provider}:${resolvedModelId}`,
1818
+ toolInput: { finishReason: result.finishReason },
1819
+ toolOutput: {
1820
+ text: result.text?.slice(0, 500),
1821
+ tokens: { prompt: promptTokens, completion: completionTokens }
1822
+ },
1823
+ estimatedCost: cost,
1824
+ tokenCount: totalTokens,
1825
+ metadata: {
1762
1826
  tool_calls: result.toolCalls?.map((tc) => tc.toolName)
1763
1827
  }
1828
+ }).catch(() => {
1764
1829
  });
1765
1830
  }
1766
1831
  await client.completeSession({
@@ -1797,7 +1862,11 @@ function wrapStreamText(originalFn, client, config) {
1797
1862
  const { modelId, provider } = extractModelInfo(params.model);
1798
1863
  const input = extractInput(params);
1799
1864
  let sessionId = null;
1800
- const sessionPromise = (async () => {
1865
+ let resolveSessionReady;
1866
+ const sessionReady = new Promise((resolve) => {
1867
+ resolveSessionReady = resolve;
1868
+ });
1869
+ (async () => {
1801
1870
  try {
1802
1871
  const id = await client.createSession({
1803
1872
  name: `streamText: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
@@ -1815,143 +1884,120 @@ function wrapStreamText(originalFn, client, config) {
1815
1884
  client.setInput(id, input).catch(() => {
1816
1885
  });
1817
1886
  }
1818
- return id;
1819
1887
  } catch {
1820
- return null;
1888
+ sessionId = null;
1821
1889
  }
1890
+ resolveSessionReady();
1822
1891
  })();
1892
+ const userOnStepFinish = params.onStepFinish;
1893
+ const userOnFinish = params.onFinish;
1894
+ const userOnError = params.onError;
1823
1895
  const wrappedParams = {
1824
1896
  ...params,
1825
- tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
1826
- };
1827
- const result = originalFn(wrappedParams);
1828
- const originalTextStream = result.textStream;
1829
- let fullText = "";
1830
- let tracked = false;
1831
- async function trackCompletion(text, error) {
1832
- if (tracked) return;
1833
- tracked = true;
1834
- const durationMs = Date.now() - startTime;
1835
- const sid = sessionId || await sessionPromise;
1836
- if (!sid) return;
1837
- if (error) {
1838
- await client.trackError({
1839
- sessionId: sid,
1840
- errorType: error.name || "Error",
1841
- errorMessage: error.message || "Unknown error"
1842
- }).catch(() => {
1843
- });
1844
- await client.completeSession({
1845
- sessionId: sid,
1846
- success: false,
1847
- failureReason: error.message || "Unknown error",
1848
- durationMs
1849
- }).catch(() => {
1850
- });
1851
- return;
1852
- }
1853
- let resolvedModelId = modelId;
1854
- try {
1855
- const resp = result.response ? await result.response : void 0;
1856
- if (resp?.modelId) resolvedModelId = resp.modelId;
1857
- } catch {
1858
- }
1859
- let usage;
1860
- try {
1861
- usage = result.usage ? await result.usage : void 0;
1862
- } catch {
1863
- }
1864
- let steps;
1865
- try {
1866
- steps = result.steps ? await result.steps : void 0;
1867
- } catch {
1868
- }
1869
- if (steps && steps.length >= 1) {
1870
- let totalPrompt = 0, totalCompletion = 0;
1871
- const stepPromises = steps.map((step, i) => {
1872
- const sp = step.usage?.promptTokens ?? 0;
1873
- const sc = step.usage?.completionTokens ?? 0;
1874
- totalPrompt += sp;
1875
- totalCompletion += sc;
1876
- return client.trackEvent({
1877
- sessionId: sid,
1878
- eventType: "llm_call",
1879
- eventData: {
1880
- model: resolvedModelId,
1881
- provider,
1882
- step: i + 1,
1883
- total_steps: steps.length,
1884
- prompt_tokens: sp,
1885
- completion_tokens: sc,
1886
- total_tokens: step.usage?.totalTokens ?? 0,
1887
- finish_reason: step.finishReason,
1897
+ tools: params.tools ? wrapToolsAsync(params.tools, sessionReady.then(() => sessionId), client) : void 0,
1898
+ onStepFinish: async (step) => {
1899
+ await sessionReady;
1900
+ if (sessionId) {
1901
+ const su = normalizeUsage(step.usage);
1902
+ const resolvedStepModel = step.response?.modelId ?? modelId;
1903
+ client.trackToolCall({
1904
+ sessionId,
1905
+ toolName: `llm:${provider}:${resolvedStepModel}`,
1906
+ toolInput: { finishReason: step.finishReason },
1907
+ toolOutput: {
1908
+ text: step.text?.slice(0, 500),
1909
+ tokens: {
1910
+ prompt: su?.promptTokens ?? 0,
1911
+ completion: su?.completionTokens ?? 0
1912
+ }
1913
+ },
1914
+ estimatedCost: calculateCostForCall(provider, resolvedStepModel, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
1915
+ tokenCount: su?.totalTokens,
1916
+ metadata: {
1888
1917
  tool_calls: step.toolCalls?.map((tc) => tc.toolName)
1889
1918
  }
1890
1919
  }).catch(() => {
1891
1920
  });
1892
- });
1893
- await Promise.all(stepPromises);
1894
- const promptTokens = usage?.promptTokens ?? totalPrompt;
1895
- const completionTokens = usage?.completionTokens ?? totalCompletion;
1896
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
1897
- const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1898
- await client.completeSession({
1899
- sessionId: sid,
1900
- success: true,
1901
- output: text,
1902
- durationMs,
1903
- estimatedCost: cost,
1904
- promptTokens,
1905
- completionTokens,
1906
- totalTokens
1907
- }).catch(() => {
1908
- });
1909
- } else {
1910
- const promptTokens = usage?.promptTokens ?? 0;
1911
- const completionTokens = usage?.completionTokens ?? 0;
1912
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
1913
- const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1914
- await client.trackEvent({
1915
- sessionId: sid,
1916
- eventType: "llm_call",
1917
- eventData: {
1918
- model: resolvedModelId,
1919
- provider,
1920
- prompt_tokens: promptTokens,
1921
- completion_tokens: completionTokens,
1922
- total_tokens: totalTokens
1921
+ }
1922
+ if (userOnStepFinish) {
1923
+ try {
1924
+ userOnStepFinish(step);
1925
+ } catch {
1926
+ }
1927
+ }
1928
+ },
1929
+ onFinish: async (event) => {
1930
+ await sessionReady;
1931
+ if (sessionId) {
1932
+ const durationMs = Date.now() - startTime;
1933
+ const usage = normalizeUsage(event.totalUsage ?? event.usage);
1934
+ const resolvedModelId = event.response?.modelId ?? modelId;
1935
+ const text = event.text ?? "";
1936
+ const promptTokens = usage?.promptTokens ?? 0;
1937
+ const completionTokens = usage?.completionTokens ?? 0;
1938
+ const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
1939
+ const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1940
+ await client.completeSession({
1941
+ sessionId,
1942
+ success: true,
1943
+ output: text,
1944
+ durationMs,
1945
+ estimatedCost: cost,
1946
+ promptTokens,
1947
+ completionTokens,
1948
+ totalTokens
1949
+ }).catch(() => {
1950
+ });
1951
+ }
1952
+ if (userOnFinish) {
1953
+ try {
1954
+ userOnFinish(event);
1955
+ } catch {
1956
+ }
1957
+ }
1958
+ },
1959
+ onError: async (event) => {
1960
+ await sessionReady;
1961
+ if (sessionId) {
1962
+ const durationMs = Date.now() - startTime;
1963
+ const msg = event.error?.message ?? "Unknown error";
1964
+ await client.trackError({
1965
+ sessionId,
1966
+ errorType: event.error?.name ?? "Error",
1967
+ errorMessage: msg
1968
+ }).catch(() => {
1969
+ });
1970
+ await client.completeSession({
1971
+ sessionId,
1972
+ success: false,
1973
+ failureReason: msg,
1974
+ durationMs
1975
+ }).catch(() => {
1976
+ });
1977
+ }
1978
+ if (userOnError) {
1979
+ try {
1980
+ userOnError(event);
1981
+ } catch {
1923
1982
  }
1924
- }).catch(() => {
1925
- });
1926
- await client.completeSession({
1927
- sessionId: sid,
1928
- success: true,
1929
- output: text,
1930
- durationMs,
1931
- estimatedCost: cost,
1932
- promptTokens,
1933
- completionTokens,
1934
- totalTokens
1935
- }).catch(() => {
1936
- });
1937
- }
1938
- }
1939
- result.textStream = (async function* () {
1940
- try {
1941
- for await (const chunk of originalTextStream) {
1942
- fullText += chunk;
1943
- yield chunk;
1944
1983
  }
1945
- await trackCompletion(fullText);
1946
- } catch (error) {
1947
- await trackCompletion(
1948
- fullText,
1949
- error instanceof Error ? error : new Error(String(error))
1950
- );
1951
- throw error;
1952
1984
  }
1953
- })();
1954
- return result;
1985
+ };
1986
+ try {
1987
+ return originalFn(wrappedParams);
1988
+ } catch (error) {
1989
+ sessionReady.then(() => {
1990
+ if (sessionId) {
1991
+ const durationMs = Date.now() - startTime;
1992
+ const msg = error instanceof Error ? error.message : String(error);
1993
+ client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
1994
+ });
1995
+ client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
1996
+ });
1997
+ }
1998
+ });
1999
+ throw error;
2000
+ }
1955
2001
  };
1956
2002
  }
1957
2003
  function wrapGenerateObject(originalFn, client, config) {
@@ -1978,20 +2024,21 @@ function wrapGenerateObject(originalFn, client, config) {
1978
2024
  const result = await originalFn(params);
1979
2025
  const durationMs = Date.now() - startTime;
1980
2026
  const resolvedModelId = result.response?.modelId || modelId;
1981
- const promptTokens = result.usage?.promptTokens ?? 0;
1982
- const completionTokens = result.usage?.completionTokens ?? 0;
1983
- const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
2027
+ const usage = normalizeUsage(result.usage);
2028
+ const promptTokens = usage?.promptTokens ?? 0;
2029
+ const completionTokens = usage?.completionTokens ?? 0;
2030
+ const totalTokens = usage?.totalTokens ?? 0;
1984
2031
  const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1985
- await client.trackEvent({
2032
+ await client.trackToolCall({
1986
2033
  sessionId,
1987
- eventType: "llm_call",
1988
- eventData: {
1989
- model: resolvedModelId,
1990
- provider,
1991
- prompt_tokens: promptTokens,
1992
- completion_tokens: completionTokens,
1993
- total_tokens: totalTokens
1994
- }
2034
+ toolName: `llm:${provider}:${resolvedModelId}`,
2035
+ toolInput: { function: "generateObject" },
2036
+ toolOutput: {
2037
+ object: result.object,
2038
+ tokens: { prompt: promptTokens, completion: completionTokens }
2039
+ },
2040
+ estimatedCost: cost,
2041
+ tokenCount: totalTokens
1995
2042
  }).catch(() => {
1996
2043
  });
1997
2044
  await client.completeSession({
@@ -2027,7 +2074,12 @@ function wrapStreamObject(originalFn, client, config) {
2027
2074
  const startTime = Date.now();
2028
2075
  const { modelId, provider } = extractModelInfo(params.model);
2029
2076
  const input = extractInput(params);
2030
- const sessionPromise = (async () => {
2077
+ let sessionId = null;
2078
+ let resolveSessionReady;
2079
+ const sessionReady = new Promise((resolve) => {
2080
+ resolveSessionReady = resolve;
2081
+ });
2082
+ (async () => {
2031
2083
  try {
2032
2084
  const id = await client.createSession({
2033
2085
  name: `streamObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
@@ -2040,86 +2092,113 @@ function wrapStreamObject(originalFn, client, config) {
2040
2092
  function: "streamObject"
2041
2093
  }
2042
2094
  });
2095
+ sessionId = id;
2043
2096
  if (id) {
2044
2097
  client.setInput(id, input).catch(() => {
2045
2098
  });
2046
2099
  }
2047
- return id;
2048
2100
  } catch {
2049
- return null;
2101
+ sessionId = null;
2050
2102
  }
2103
+ resolveSessionReady();
2051
2104
  })();
2052
- const result = originalFn(params);
2053
- async function completeStreamObject(obj, error) {
2054
- const durationMs = Date.now() - startTime;
2055
- const sid = await sessionPromise;
2056
- if (!sid) return;
2057
- if (error) {
2058
- await client.trackError({
2059
- sessionId: sid,
2060
- errorType: error.name || "Error",
2061
- errorMessage: error.message || "Unknown error"
2062
- }).catch(() => {
2063
- });
2064
- await client.completeSession({
2065
- sessionId: sid,
2066
- success: false,
2067
- failureReason: error.message || "Unknown error",
2068
- durationMs
2069
- }).catch(() => {
2070
- });
2071
- return;
2072
- }
2073
- let usage;
2074
- try {
2075
- usage = result.usage ? await result.usage : void 0;
2076
- } catch {
2105
+ const userOnFinish = params.onFinish;
2106
+ const userOnError = params.onError;
2107
+ const wrappedParams = {
2108
+ ...params,
2109
+ onFinish: async (event) => {
2110
+ await sessionReady;
2111
+ if (sessionId) {
2112
+ const durationMs = Date.now() - startTime;
2113
+ const usage = normalizeUsage(event.usage);
2114
+ const resolvedModelId = event.response?.modelId ?? modelId;
2115
+ const promptTokens = usage?.promptTokens ?? 0;
2116
+ const completionTokens = usage?.completionTokens ?? 0;
2117
+ const totalTokens = usage?.totalTokens ?? 0;
2118
+ const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
2119
+ if (event.error) {
2120
+ const errMsg = event.error instanceof Error ? event.error.message : String(event.error);
2121
+ await client.trackError({
2122
+ sessionId,
2123
+ errorType: "SchemaValidationError",
2124
+ errorMessage: errMsg
2125
+ }).catch(() => {
2126
+ });
2127
+ }
2128
+ await client.trackToolCall({
2129
+ sessionId,
2130
+ toolName: `llm:${provider}:${resolvedModelId}`,
2131
+ toolInput: { function: "streamObject" },
2132
+ toolOutput: {
2133
+ object: event.object,
2134
+ tokens: { prompt: promptTokens, completion: completionTokens }
2135
+ },
2136
+ estimatedCost: cost,
2137
+ tokenCount: totalTokens
2138
+ }).catch(() => {
2139
+ });
2140
+ await client.completeSession({
2141
+ sessionId,
2142
+ success: !event.error,
2143
+ output: event.object != null ? JSON.stringify(event.object) : void 0,
2144
+ failureReason: event.error ? String(event.error) : void 0,
2145
+ durationMs,
2146
+ estimatedCost: cost,
2147
+ promptTokens,
2148
+ completionTokens,
2149
+ totalTokens
2150
+ }).catch(() => {
2151
+ });
2152
+ }
2153
+ if (userOnFinish) {
2154
+ try {
2155
+ userOnFinish(event);
2156
+ } catch {
2157
+ }
2158
+ }
2159
+ },
2160
+ onError: async (event) => {
2161
+ await sessionReady;
2162
+ if (sessionId) {
2163
+ const durationMs = Date.now() - startTime;
2164
+ const msg = event.error?.message ?? "Unknown error";
2165
+ await client.trackError({
2166
+ sessionId,
2167
+ errorType: event.error?.name ?? "Error",
2168
+ errorMessage: msg
2169
+ }).catch(() => {
2170
+ });
2171
+ await client.completeSession({
2172
+ sessionId,
2173
+ success: false,
2174
+ failureReason: msg,
2175
+ durationMs
2176
+ }).catch(() => {
2177
+ });
2178
+ }
2179
+ if (userOnError) {
2180
+ try {
2181
+ userOnError(event);
2182
+ } catch {
2183
+ }
2184
+ }
2077
2185
  }
2078
- const promptTokens = usage?.promptTokens ?? 0;
2079
- const completionTokens = usage?.completionTokens ?? 0;
2080
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
2081
- const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
2082
- await client.trackEvent({
2083
- sessionId: sid,
2084
- eventType: "llm_call",
2085
- eventData: {
2086
- model: modelId,
2087
- provider,
2088
- prompt_tokens: promptTokens,
2089
- completion_tokens: completionTokens,
2090
- total_tokens: totalTokens
2186
+ };
2187
+ try {
2188
+ return originalFn(wrappedParams);
2189
+ } catch (error) {
2190
+ sessionReady.then(() => {
2191
+ if (sessionId) {
2192
+ const durationMs = Date.now() - startTime;
2193
+ const msg = error instanceof Error ? error.message : String(error);
2194
+ client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
2195
+ });
2196
+ client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
2197
+ });
2091
2198
  }
2092
- }).catch(() => {
2093
- });
2094
- await client.completeSession({
2095
- sessionId: sid,
2096
- success: true,
2097
- output: JSON.stringify(obj),
2098
- durationMs,
2099
- estimatedCost: cost,
2100
- promptTokens,
2101
- completionTokens,
2102
- totalTokens
2103
- }).catch(() => {
2104
- });
2105
- }
2106
- if (result.object) {
2107
- const originalObjectPromise = result.object;
2108
- result.object = originalObjectPromise.then(async (obj) => {
2109
- await completeStreamObject(obj);
2110
- return obj;
2111
- }).catch(async (error) => {
2112
- await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
2113
- throw error;
2114
- });
2115
- } else if (result.usage) {
2116
- result.usage.then(async () => {
2117
- await completeStreamObject(void 0);
2118
- }).catch(async (error) => {
2119
- await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
2120
2199
  });
2200
+ throw error;
2121
2201
  }
2122
- return result;
2123
2202
  };
2124
2203
  }
2125
2204
  function wrapAISDK(ai, options) {