@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.cjs CHANGED
@@ -437,8 +437,18 @@ function wrapOpenAI(client, options = {}) {
437
437
  const completionTokens = response.usage?.completion_tokens ?? 0;
438
438
  const totalTokens = response.usage?.total_tokens ?? 0;
439
439
  let outputContent = "";
440
- if (response.choices?.[0]?.message?.content) {
441
- outputContent = response.choices[0].message.content;
440
+ const toolCalls = [];
441
+ const msg = response.choices?.[0]?.message;
442
+ if (msg?.content) {
443
+ outputContent = msg.content;
444
+ }
445
+ if (msg?.tool_calls) {
446
+ for (const tc of msg.tool_calls) {
447
+ toolCalls.push({
448
+ name: tc.function?.name ?? "unknown",
449
+ arguments: tc.function?.arguments ?? "{}"
450
+ });
451
+ }
442
452
  }
443
453
  const cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
444
454
  trackLLMCall({
@@ -446,6 +456,7 @@ function wrapOpenAI(client, options = {}) {
446
456
  model,
447
457
  messages,
448
458
  output: outputContent,
459
+ toolCalls,
449
460
  promptTokens,
450
461
  completionTokens,
451
462
  totalTokens,
@@ -500,10 +511,16 @@ function wrapAnthropic(client, options = {}) {
500
511
  const completionTokens = response.usage?.output_tokens ?? 0;
501
512
  const totalTokens = promptTokens + completionTokens;
502
513
  let outputContent = "";
514
+ const toolCalls = [];
503
515
  if (response.content) {
504
516
  for (const block of response.content) {
505
517
  if (block.type === "text") {
506
518
  outputContent += block.text;
519
+ } else if (block.type === "tool_use") {
520
+ toolCalls.push({
521
+ name: block.name ?? "unknown",
522
+ arguments: JSON.stringify(block.input ?? {})
523
+ });
507
524
  }
508
525
  }
509
526
  }
@@ -514,6 +531,7 @@ function wrapAnthropic(client, options = {}) {
514
531
  model,
515
532
  messages: fullMessages,
516
533
  output: outputContent,
534
+ toolCalls,
517
535
  promptTokens,
518
536
  completionTokens,
519
537
  totalTokens,
@@ -627,6 +645,7 @@ function wrapLLM(client, provider) {
627
645
  function wrapOpenAIStream(stream, ctx) {
628
646
  let fullContent = "";
629
647
  let usage = null;
648
+ const toolCallMap = /* @__PURE__ */ new Map();
630
649
  let tracked = false;
631
650
  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
632
651
  if (!originalIterator) return stream;
@@ -643,6 +662,7 @@ function wrapOpenAIStream(stream, ctx) {
643
662
  model: ctx.model,
644
663
  messages: ctx.messages,
645
664
  output: fullContent,
665
+ toolCalls: Array.from(toolCallMap.values()),
646
666
  promptTokens,
647
667
  completionTokens,
648
668
  totalTokens,
@@ -661,8 +681,17 @@ function wrapOpenAIStream(stream, ctx) {
661
681
  const result = await iter.next();
662
682
  if (!result.done) {
663
683
  const chunk = result.value;
664
- const delta = chunk.choices?.[0]?.delta?.content;
665
- if (delta) fullContent += delta;
684
+ const delta = chunk.choices?.[0]?.delta;
685
+ if (delta?.content) fullContent += delta.content;
686
+ if (delta?.tool_calls) {
687
+ for (const tc of delta.tool_calls) {
688
+ const idx = tc.index ?? 0;
689
+ const existing = toolCallMap.get(idx) ?? { name: "", arguments: "" };
690
+ if (tc.function?.name) existing.name = tc.function.name;
691
+ if (tc.function?.arguments) existing.arguments += tc.function.arguments;
692
+ toolCallMap.set(idx, existing);
693
+ }
694
+ }
666
695
  if (chunk.usage) usage = chunk.usage;
667
696
  } else {
668
697
  trackResult();
@@ -687,6 +716,8 @@ function wrapAnthropicStream(stream, ctx) {
687
716
  let fullContent = "";
688
717
  let inputTokens = 0;
689
718
  let outputTokens = 0;
719
+ const toolCallsById = /* @__PURE__ */ new Map();
720
+ let currentBlockIdx = -1;
690
721
  let tracked = false;
691
722
  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
692
723
  if (!originalIterator) return stream;
@@ -702,6 +733,7 @@ function wrapAnthropicStream(stream, ctx) {
702
733
  model: ctx.model,
703
734
  messages: fullMessages,
704
735
  output: fullContent,
736
+ toolCalls: Array.from(toolCallsById.values()),
705
737
  promptTokens: inputTokens,
706
738
  completionTokens: outputTokens,
707
739
  totalTokens,
@@ -720,8 +752,26 @@ function wrapAnthropicStream(stream, ctx) {
720
752
  const result = await iter.next();
721
753
  if (!result.done) {
722
754
  const event = result.value;
723
- if (event.type === "content_block_delta" && event.delta?.text) {
724
- fullContent += event.delta.text;
755
+ if (event.type === "content_block_start") {
756
+ currentBlockIdx = event.index ?? currentBlockIdx + 1;
757
+ if (event.content_block?.type === "tool_use") {
758
+ toolCallsById.set(currentBlockIdx, {
759
+ name: event.content_block.name ?? "unknown",
760
+ arguments: ""
761
+ });
762
+ }
763
+ }
764
+ if (event.type === "content_block_delta") {
765
+ if (event.delta?.text) {
766
+ fullContent += event.delta.text;
767
+ }
768
+ if (event.delta?.type === "input_json_delta" && event.delta?.partial_json) {
769
+ const idx = event.index ?? currentBlockIdx;
770
+ const existing = toolCallsById.get(idx);
771
+ if (existing) {
772
+ existing.arguments += event.delta.partial_json;
773
+ }
774
+ }
725
775
  }
726
776
  if (event.type === "message_start" && event.message?.usage) {
727
777
  inputTokens = event.message.usage.input_tokens ?? 0;
@@ -755,6 +805,18 @@ function trackLLMCall(params) {
755
805
  if (!sessionId && !params.trackWithoutSession) {
756
806
  return;
757
807
  }
808
+ const toolOutput = {
809
+ content: params.output,
810
+ tokens: {
811
+ prompt: params.promptTokens,
812
+ completion: params.completionTokens,
813
+ total: params.totalTokens
814
+ },
815
+ cost_usd: params.cost
816
+ };
817
+ if (params.toolCalls && params.toolCalls.length > 0) {
818
+ toolOutput.tool_calls = params.toolCalls;
819
+ }
758
820
  if (sessionId) {
759
821
  client.trackToolCall({
760
822
  sessionId,
@@ -764,15 +826,7 @@ function trackLLMCall(params) {
764
826
  model: params.model,
765
827
  provider: params.provider
766
828
  },
767
- toolOutput: {
768
- content: params.output,
769
- tokens: {
770
- prompt: params.promptTokens,
771
- completion: params.completionTokens,
772
- total: params.totalTokens
773
- },
774
- cost_usd: params.cost
775
- },
829
+ toolOutput,
776
830
  reasoning: `LLM call to ${params.provider} ${params.model}`,
777
831
  estimatedCost: params.cost,
778
832
  tokenCount: params.totalTokens,
@@ -801,15 +855,7 @@ function trackLLMCall(params) {
801
855
  model: params.model,
802
856
  provider: params.provider
803
857
  },
804
- toolOutput: {
805
- content: params.output,
806
- tokens: {
807
- prompt: params.promptTokens,
808
- completion: params.completionTokens,
809
- total: params.totalTokens
810
- },
811
- cost_usd: params.cost
812
- },
858
+ toolOutput,
813
859
  estimatedCost: params.cost,
814
860
  tokenCount: params.totalTokens,
815
861
  metadata: {
@@ -1632,7 +1678,8 @@ function getClient2() {
1632
1678
  }
1633
1679
  function extractModelInfo(model) {
1634
1680
  const modelId = model.modelId || model.id || "unknown";
1635
- const provider = model.provider || guessProvider(modelId);
1681
+ const rawProvider = model.provider || "";
1682
+ const provider = rawProvider.split(".")[0] || guessProvider(modelId);
1636
1683
  return { modelId, provider };
1637
1684
  }
1638
1685
  function guessProvider(modelId) {
@@ -1677,6 +1724,15 @@ function extractInput(params) {
1677
1724
  }
1678
1725
  return "";
1679
1726
  }
1727
+ function normalizeUsage(usage) {
1728
+ if (!usage) return void 0;
1729
+ const u = usage;
1730
+ const promptTokens = (u.inputTokens ?? u.promptTokens ?? 0) || 0;
1731
+ const completionTokens = (u.outputTokens ?? u.completionTokens ?? 0) || 0;
1732
+ const totalTokens = (u.totalTokens ?? promptTokens + completionTokens) || 0;
1733
+ if (promptTokens === 0 && completionTokens === 0 && totalTokens === 0) return void 0;
1734
+ return { promptTokens, completionTokens, totalTokens };
1735
+ }
1680
1736
  function wrapTools(tools, sessionId, client) {
1681
1737
  if (!tools) return void 0;
1682
1738
  const wrappedTools = {};
@@ -1729,34 +1785,37 @@ function wrapToolsAsync(tools, sessionPromise, client) {
1729
1785
  ...tool,
1730
1786
  execute: async (...args) => {
1731
1787
  const startTime = Date.now();
1732
- const sid = await sessionPromise;
1733
1788
  try {
1734
1789
  const result = await originalExecute(...args);
1735
1790
  const durationMs = Date.now() - startTime;
1736
- if (sid) {
1737
- client.trackToolCall({
1738
- sessionId: sid,
1739
- toolName,
1740
- toolInput: args[0],
1741
- toolOutput: result,
1742
- reasoning: `Tool executed in ${durationMs}ms`
1743
- }).catch(() => {
1744
- });
1745
- }
1791
+ sessionPromise.then((sid) => {
1792
+ if (sid) {
1793
+ client.trackToolCall({
1794
+ sessionId: sid,
1795
+ toolName,
1796
+ toolInput: args[0],
1797
+ toolOutput: result,
1798
+ reasoning: `Tool executed in ${durationMs}ms`
1799
+ }).catch(() => {
1800
+ });
1801
+ }
1802
+ });
1746
1803
  return result;
1747
1804
  } catch (error) {
1748
1805
  const durationMs = Date.now() - startTime;
1749
- if (sid) {
1750
- client.trackToolCall({
1751
- sessionId: sid,
1752
- toolName,
1753
- toolInput: args[0],
1754
- toolOutput: {},
1755
- toolError: { message: error instanceof Error ? error.message : "Unknown error" },
1756
- reasoning: `Tool failed after ${durationMs}ms`
1757
- }).catch(() => {
1758
- });
1759
- }
1806
+ sessionPromise.then((sid) => {
1807
+ if (sid) {
1808
+ client.trackToolCall({
1809
+ sessionId: sid,
1810
+ toolName,
1811
+ toolInput: args[0],
1812
+ toolOutput: {},
1813
+ toolError: { message: error instanceof Error ? error.message : "Unknown error" },
1814
+ reasoning: `Tool failed after ${durationMs}ms`
1815
+ }).catch(() => {
1816
+ });
1817
+ }
1818
+ });
1760
1819
  throw error;
1761
1820
  }
1762
1821
  }
@@ -1796,44 +1855,50 @@ function wrapGenerateText(originalFn, client, config) {
1796
1855
  const result = await originalFn(wrappedParams);
1797
1856
  const durationMs = Date.now() - startTime;
1798
1857
  const resolvedModelId = result.response?.modelId || modelId;
1799
- const promptTokens = result.usage?.promptTokens ?? 0;
1800
- const completionTokens = result.usage?.completionTokens ?? 0;
1801
- const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
1858
+ const usage = normalizeUsage(result.usage);
1859
+ const promptTokens = usage?.promptTokens ?? 0;
1860
+ const completionTokens = usage?.completionTokens ?? 0;
1861
+ const totalTokens = usage?.totalTokens ?? 0;
1802
1862
  const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1803
1863
  const steps = result.steps;
1804
1864
  if (steps && steps.length >= 1) {
1805
- const stepPromises = steps.map(
1806
- (step, i) => client.trackEvent({
1865
+ const stepPromises = steps.map((step) => {
1866
+ const su = normalizeUsage(step.usage);
1867
+ return client.trackToolCall({
1807
1868
  sessionId,
1808
- eventType: "llm_call",
1809
- eventData: {
1810
- model: resolvedModelId,
1811
- provider,
1812
- step: i + 1,
1813
- total_steps: steps.length,
1814
- prompt_tokens: step.usage?.promptTokens ?? 0,
1815
- completion_tokens: step.usage?.completionTokens ?? 0,
1816
- total_tokens: step.usage?.totalTokens ?? 0,
1817
- finish_reason: step.finishReason,
1869
+ toolName: `llm:${provider}:${resolvedModelId}`,
1870
+ toolInput: { finishReason: step.finishReason },
1871
+ toolOutput: {
1872
+ text: step.text?.slice(0, 500),
1873
+ tokens: {
1874
+ prompt: su?.promptTokens ?? 0,
1875
+ completion: su?.completionTokens ?? 0
1876
+ }
1877
+ },
1878
+ estimatedCost: calculateCostForCall(provider, resolvedModelId, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
1879
+ tokenCount: su?.totalTokens,
1880
+ metadata: {
1818
1881
  tool_calls: step.toolCalls?.map((tc) => tc.toolName)
1819
1882
  }
1820
1883
  }).catch(() => {
1821
- })
1822
- );
1884
+ });
1885
+ });
1823
1886
  await Promise.all(stepPromises);
1824
1887
  } else {
1825
- await client.trackEvent({
1888
+ await client.trackToolCall({
1826
1889
  sessionId,
1827
- eventType: "llm_call",
1828
- eventData: {
1829
- model: resolvedModelId,
1830
- provider,
1831
- prompt_tokens: promptTokens,
1832
- completion_tokens: completionTokens,
1833
- total_tokens: totalTokens,
1834
- finish_reason: result.finishReason,
1890
+ toolName: `llm:${provider}:${resolvedModelId}`,
1891
+ toolInput: { finishReason: result.finishReason },
1892
+ toolOutput: {
1893
+ text: result.text?.slice(0, 500),
1894
+ tokens: { prompt: promptTokens, completion: completionTokens }
1895
+ },
1896
+ estimatedCost: cost,
1897
+ tokenCount: totalTokens,
1898
+ metadata: {
1835
1899
  tool_calls: result.toolCalls?.map((tc) => tc.toolName)
1836
1900
  }
1901
+ }).catch(() => {
1837
1902
  });
1838
1903
  }
1839
1904
  await client.completeSession({
@@ -1870,7 +1935,11 @@ function wrapStreamText(originalFn, client, config) {
1870
1935
  const { modelId, provider } = extractModelInfo(params.model);
1871
1936
  const input = extractInput(params);
1872
1937
  let sessionId = null;
1873
- const sessionPromise = (async () => {
1938
+ let resolveSessionReady;
1939
+ const sessionReady = new Promise((resolve) => {
1940
+ resolveSessionReady = resolve;
1941
+ });
1942
+ (async () => {
1874
1943
  try {
1875
1944
  const id = await client.createSession({
1876
1945
  name: `streamText: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
@@ -1888,143 +1957,120 @@ function wrapStreamText(originalFn, client, config) {
1888
1957
  client.setInput(id, input).catch(() => {
1889
1958
  });
1890
1959
  }
1891
- return id;
1892
1960
  } catch {
1893
- return null;
1961
+ sessionId = null;
1894
1962
  }
1963
+ resolveSessionReady();
1895
1964
  })();
1965
+ const userOnStepFinish = params.onStepFinish;
1966
+ const userOnFinish = params.onFinish;
1967
+ const userOnError = params.onError;
1896
1968
  const wrappedParams = {
1897
1969
  ...params,
1898
- tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
1899
- };
1900
- const result = originalFn(wrappedParams);
1901
- const originalTextStream = result.textStream;
1902
- let fullText = "";
1903
- let tracked = false;
1904
- async function trackCompletion(text, error) {
1905
- if (tracked) return;
1906
- tracked = true;
1907
- const durationMs = Date.now() - startTime;
1908
- const sid = sessionId || await sessionPromise;
1909
- if (!sid) return;
1910
- if (error) {
1911
- await client.trackError({
1912
- sessionId: sid,
1913
- errorType: error.name || "Error",
1914
- errorMessage: error.message || "Unknown error"
1915
- }).catch(() => {
1916
- });
1917
- await client.completeSession({
1918
- sessionId: sid,
1919
- success: false,
1920
- failureReason: error.message || "Unknown error",
1921
- durationMs
1922
- }).catch(() => {
1923
- });
1924
- return;
1925
- }
1926
- let resolvedModelId = modelId;
1927
- try {
1928
- const resp = result.response ? await result.response : void 0;
1929
- if (resp?.modelId) resolvedModelId = resp.modelId;
1930
- } catch {
1931
- }
1932
- let usage;
1933
- try {
1934
- usage = result.usage ? await result.usage : void 0;
1935
- } catch {
1936
- }
1937
- let steps;
1938
- try {
1939
- steps = result.steps ? await result.steps : void 0;
1940
- } catch {
1941
- }
1942
- if (steps && steps.length >= 1) {
1943
- let totalPrompt = 0, totalCompletion = 0;
1944
- const stepPromises = steps.map((step, i) => {
1945
- const sp = step.usage?.promptTokens ?? 0;
1946
- const sc = step.usage?.completionTokens ?? 0;
1947
- totalPrompt += sp;
1948
- totalCompletion += sc;
1949
- return client.trackEvent({
1950
- sessionId: sid,
1951
- eventType: "llm_call",
1952
- eventData: {
1953
- model: resolvedModelId,
1954
- provider,
1955
- step: i + 1,
1956
- total_steps: steps.length,
1957
- prompt_tokens: sp,
1958
- completion_tokens: sc,
1959
- total_tokens: step.usage?.totalTokens ?? 0,
1960
- finish_reason: step.finishReason,
1970
+ tools: params.tools ? wrapToolsAsync(params.tools, sessionReady.then(() => sessionId), client) : void 0,
1971
+ onStepFinish: async (step) => {
1972
+ await sessionReady;
1973
+ if (sessionId) {
1974
+ const su = normalizeUsage(step.usage);
1975
+ const resolvedStepModel = step.response?.modelId ?? modelId;
1976
+ client.trackToolCall({
1977
+ sessionId,
1978
+ toolName: `llm:${provider}:${resolvedStepModel}`,
1979
+ toolInput: { finishReason: step.finishReason },
1980
+ toolOutput: {
1981
+ text: step.text?.slice(0, 500),
1982
+ tokens: {
1983
+ prompt: su?.promptTokens ?? 0,
1984
+ completion: su?.completionTokens ?? 0
1985
+ }
1986
+ },
1987
+ estimatedCost: calculateCostForCall(provider, resolvedStepModel, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
1988
+ tokenCount: su?.totalTokens,
1989
+ metadata: {
1961
1990
  tool_calls: step.toolCalls?.map((tc) => tc.toolName)
1962
1991
  }
1963
1992
  }).catch(() => {
1964
1993
  });
1965
- });
1966
- await Promise.all(stepPromises);
1967
- const promptTokens = usage?.promptTokens ?? totalPrompt;
1968
- const completionTokens = usage?.completionTokens ?? totalCompletion;
1969
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
1970
- const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1971
- await client.completeSession({
1972
- sessionId: sid,
1973
- success: true,
1974
- output: text,
1975
- durationMs,
1976
- estimatedCost: cost,
1977
- promptTokens,
1978
- completionTokens,
1979
- totalTokens
1980
- }).catch(() => {
1981
- });
1982
- } else {
1983
- const promptTokens = usage?.promptTokens ?? 0;
1984
- const completionTokens = usage?.completionTokens ?? 0;
1985
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
1986
- const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
1987
- await client.trackEvent({
1988
- sessionId: sid,
1989
- eventType: "llm_call",
1990
- eventData: {
1991
- model: resolvedModelId,
1992
- provider,
1993
- prompt_tokens: promptTokens,
1994
- completion_tokens: completionTokens,
1995
- total_tokens: totalTokens
1994
+ }
1995
+ if (userOnStepFinish) {
1996
+ try {
1997
+ userOnStepFinish(step);
1998
+ } catch {
1999
+ }
2000
+ }
2001
+ },
2002
+ onFinish: async (event) => {
2003
+ await sessionReady;
2004
+ if (sessionId) {
2005
+ const durationMs = Date.now() - startTime;
2006
+ const usage = normalizeUsage(event.totalUsage ?? event.usage);
2007
+ const resolvedModelId = event.response?.modelId ?? modelId;
2008
+ const text = event.text ?? "";
2009
+ const promptTokens = usage?.promptTokens ?? 0;
2010
+ const completionTokens = usage?.completionTokens ?? 0;
2011
+ const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
2012
+ const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
2013
+ await client.completeSession({
2014
+ sessionId,
2015
+ success: true,
2016
+ output: text,
2017
+ durationMs,
2018
+ estimatedCost: cost,
2019
+ promptTokens,
2020
+ completionTokens,
2021
+ totalTokens
2022
+ }).catch(() => {
2023
+ });
2024
+ }
2025
+ if (userOnFinish) {
2026
+ try {
2027
+ userOnFinish(event);
2028
+ } catch {
2029
+ }
2030
+ }
2031
+ },
2032
+ onError: async (event) => {
2033
+ await sessionReady;
2034
+ if (sessionId) {
2035
+ const durationMs = Date.now() - startTime;
2036
+ const msg = event.error?.message ?? "Unknown error";
2037
+ await client.trackError({
2038
+ sessionId,
2039
+ errorType: event.error?.name ?? "Error",
2040
+ errorMessage: msg
2041
+ }).catch(() => {
2042
+ });
2043
+ await client.completeSession({
2044
+ sessionId,
2045
+ success: false,
2046
+ failureReason: msg,
2047
+ durationMs
2048
+ }).catch(() => {
2049
+ });
2050
+ }
2051
+ if (userOnError) {
2052
+ try {
2053
+ userOnError(event);
2054
+ } catch {
1996
2055
  }
1997
- }).catch(() => {
1998
- });
1999
- await client.completeSession({
2000
- sessionId: sid,
2001
- success: true,
2002
- output: text,
2003
- durationMs,
2004
- estimatedCost: cost,
2005
- promptTokens,
2006
- completionTokens,
2007
- totalTokens
2008
- }).catch(() => {
2009
- });
2010
- }
2011
- }
2012
- result.textStream = (async function* () {
2013
- try {
2014
- for await (const chunk of originalTextStream) {
2015
- fullText += chunk;
2016
- yield chunk;
2017
2056
  }
2018
- await trackCompletion(fullText);
2019
- } catch (error) {
2020
- await trackCompletion(
2021
- fullText,
2022
- error instanceof Error ? error : new Error(String(error))
2023
- );
2024
- throw error;
2025
2057
  }
2026
- })();
2027
- return result;
2058
+ };
2059
+ try {
2060
+ return originalFn(wrappedParams);
2061
+ } catch (error) {
2062
+ sessionReady.then(() => {
2063
+ if (sessionId) {
2064
+ const durationMs = Date.now() - startTime;
2065
+ const msg = error instanceof Error ? error.message : String(error);
2066
+ client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
2067
+ });
2068
+ client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
2069
+ });
2070
+ }
2071
+ });
2072
+ throw error;
2073
+ }
2028
2074
  };
2029
2075
  }
2030
2076
  function wrapGenerateObject(originalFn, client, config) {
@@ -2051,20 +2097,21 @@ function wrapGenerateObject(originalFn, client, config) {
2051
2097
  const result = await originalFn(params);
2052
2098
  const durationMs = Date.now() - startTime;
2053
2099
  const resolvedModelId = result.response?.modelId || modelId;
2054
- const promptTokens = result.usage?.promptTokens ?? 0;
2055
- const completionTokens = result.usage?.completionTokens ?? 0;
2056
- const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
2100
+ const usage = normalizeUsage(result.usage);
2101
+ const promptTokens = usage?.promptTokens ?? 0;
2102
+ const completionTokens = usage?.completionTokens ?? 0;
2103
+ const totalTokens = usage?.totalTokens ?? 0;
2057
2104
  const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
2058
- await client.trackEvent({
2105
+ await client.trackToolCall({
2059
2106
  sessionId,
2060
- eventType: "llm_call",
2061
- eventData: {
2062
- model: resolvedModelId,
2063
- provider,
2064
- prompt_tokens: promptTokens,
2065
- completion_tokens: completionTokens,
2066
- total_tokens: totalTokens
2067
- }
2107
+ toolName: `llm:${provider}:${resolvedModelId}`,
2108
+ toolInput: { function: "generateObject" },
2109
+ toolOutput: {
2110
+ object: result.object,
2111
+ tokens: { prompt: promptTokens, completion: completionTokens }
2112
+ },
2113
+ estimatedCost: cost,
2114
+ tokenCount: totalTokens
2068
2115
  }).catch(() => {
2069
2116
  });
2070
2117
  await client.completeSession({
@@ -2100,7 +2147,12 @@ function wrapStreamObject(originalFn, client, config) {
2100
2147
  const startTime = Date.now();
2101
2148
  const { modelId, provider } = extractModelInfo(params.model);
2102
2149
  const input = extractInput(params);
2103
- const sessionPromise = (async () => {
2150
+ let sessionId = null;
2151
+ let resolveSessionReady;
2152
+ const sessionReady = new Promise((resolve) => {
2153
+ resolveSessionReady = resolve;
2154
+ });
2155
+ (async () => {
2104
2156
  try {
2105
2157
  const id = await client.createSession({
2106
2158
  name: `streamObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
@@ -2113,86 +2165,113 @@ function wrapStreamObject(originalFn, client, config) {
2113
2165
  function: "streamObject"
2114
2166
  }
2115
2167
  });
2168
+ sessionId = id;
2116
2169
  if (id) {
2117
2170
  client.setInput(id, input).catch(() => {
2118
2171
  });
2119
2172
  }
2120
- return id;
2121
2173
  } catch {
2122
- return null;
2174
+ sessionId = null;
2123
2175
  }
2176
+ resolveSessionReady();
2124
2177
  })();
2125
- const result = originalFn(params);
2126
- async function completeStreamObject(obj, error) {
2127
- const durationMs = Date.now() - startTime;
2128
- const sid = await sessionPromise;
2129
- if (!sid) return;
2130
- if (error) {
2131
- await client.trackError({
2132
- sessionId: sid,
2133
- errorType: error.name || "Error",
2134
- errorMessage: error.message || "Unknown error"
2135
- }).catch(() => {
2136
- });
2137
- await client.completeSession({
2138
- sessionId: sid,
2139
- success: false,
2140
- failureReason: error.message || "Unknown error",
2141
- durationMs
2142
- }).catch(() => {
2143
- });
2144
- return;
2145
- }
2146
- let usage;
2147
- try {
2148
- usage = result.usage ? await result.usage : void 0;
2149
- } catch {
2178
+ const userOnFinish = params.onFinish;
2179
+ const userOnError = params.onError;
2180
+ const wrappedParams = {
2181
+ ...params,
2182
+ onFinish: async (event) => {
2183
+ await sessionReady;
2184
+ if (sessionId) {
2185
+ const durationMs = Date.now() - startTime;
2186
+ const usage = normalizeUsage(event.usage);
2187
+ const resolvedModelId = event.response?.modelId ?? modelId;
2188
+ const promptTokens = usage?.promptTokens ?? 0;
2189
+ const completionTokens = usage?.completionTokens ?? 0;
2190
+ const totalTokens = usage?.totalTokens ?? 0;
2191
+ const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
2192
+ if (event.error) {
2193
+ const errMsg = event.error instanceof Error ? event.error.message : String(event.error);
2194
+ await client.trackError({
2195
+ sessionId,
2196
+ errorType: "SchemaValidationError",
2197
+ errorMessage: errMsg
2198
+ }).catch(() => {
2199
+ });
2200
+ }
2201
+ await client.trackToolCall({
2202
+ sessionId,
2203
+ toolName: `llm:${provider}:${resolvedModelId}`,
2204
+ toolInput: { function: "streamObject" },
2205
+ toolOutput: {
2206
+ object: event.object,
2207
+ tokens: { prompt: promptTokens, completion: completionTokens }
2208
+ },
2209
+ estimatedCost: cost,
2210
+ tokenCount: totalTokens
2211
+ }).catch(() => {
2212
+ });
2213
+ await client.completeSession({
2214
+ sessionId,
2215
+ success: !event.error,
2216
+ output: event.object != null ? JSON.stringify(event.object) : void 0,
2217
+ failureReason: event.error ? String(event.error) : void 0,
2218
+ durationMs,
2219
+ estimatedCost: cost,
2220
+ promptTokens,
2221
+ completionTokens,
2222
+ totalTokens
2223
+ }).catch(() => {
2224
+ });
2225
+ }
2226
+ if (userOnFinish) {
2227
+ try {
2228
+ userOnFinish(event);
2229
+ } catch {
2230
+ }
2231
+ }
2232
+ },
2233
+ onError: async (event) => {
2234
+ await sessionReady;
2235
+ if (sessionId) {
2236
+ const durationMs = Date.now() - startTime;
2237
+ const msg = event.error?.message ?? "Unknown error";
2238
+ await client.trackError({
2239
+ sessionId,
2240
+ errorType: event.error?.name ?? "Error",
2241
+ errorMessage: msg
2242
+ }).catch(() => {
2243
+ });
2244
+ await client.completeSession({
2245
+ sessionId,
2246
+ success: false,
2247
+ failureReason: msg,
2248
+ durationMs
2249
+ }).catch(() => {
2250
+ });
2251
+ }
2252
+ if (userOnError) {
2253
+ try {
2254
+ userOnError(event);
2255
+ } catch {
2256
+ }
2257
+ }
2150
2258
  }
2151
- const promptTokens = usage?.promptTokens ?? 0;
2152
- const completionTokens = usage?.completionTokens ?? 0;
2153
- const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
2154
- const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
2155
- await client.trackEvent({
2156
- sessionId: sid,
2157
- eventType: "llm_call",
2158
- eventData: {
2159
- model: modelId,
2160
- provider,
2161
- prompt_tokens: promptTokens,
2162
- completion_tokens: completionTokens,
2163
- total_tokens: totalTokens
2259
+ };
2260
+ try {
2261
+ return originalFn(wrappedParams);
2262
+ } catch (error) {
2263
+ sessionReady.then(() => {
2264
+ if (sessionId) {
2265
+ const durationMs = Date.now() - startTime;
2266
+ const msg = error instanceof Error ? error.message : String(error);
2267
+ client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
2268
+ });
2269
+ client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
2270
+ });
2164
2271
  }
2165
- }).catch(() => {
2166
- });
2167
- await client.completeSession({
2168
- sessionId: sid,
2169
- success: true,
2170
- output: JSON.stringify(obj),
2171
- durationMs,
2172
- estimatedCost: cost,
2173
- promptTokens,
2174
- completionTokens,
2175
- totalTokens
2176
- }).catch(() => {
2177
- });
2178
- }
2179
- if (result.object) {
2180
- const originalObjectPromise = result.object;
2181
- result.object = originalObjectPromise.then(async (obj) => {
2182
- await completeStreamObject(obj);
2183
- return obj;
2184
- }).catch(async (error) => {
2185
- await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
2186
- throw error;
2187
- });
2188
- } else if (result.usage) {
2189
- result.usage.then(async () => {
2190
- await completeStreamObject(void 0);
2191
- }).catch(async (error) => {
2192
- await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
2193
2272
  });
2273
+ throw error;
2194
2274
  }
2195
- return result;
2196
2275
  };
2197
2276
  }
2198
2277
  function wrapAISDK(ai, options) {