Skip to content

Commit 6e2b9cd

Browse files
lforstlforst
committed
Merge origin/main into lforst/ai-sdk-embed-spans
2 parents c40d17f + 7338811 commit 6e2b9cd

10 files changed

Lines changed: 368 additions & 22 deletions

File tree

‎e2e/scenarios/ai-sdk-instrumentation/assertions.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -645,6 +645,7 @@ export function defineAISDKInstrumentationAssertions(options: {
645645
agentSpanName?: AgentSpanName;
646646
name: string;
647647
runScenario: RunAISDKScenario;
648+
sdkMajorVersion: number;
648649
snapshotName: string;
649650
supportsAttachmentScenario: boolean;
650651
supportsDenyOutputOverrideScenario: boolean;
@@ -894,6 +895,32 @@ export function defineAISDKInstrumentationAssertions(options: {
894895
});
895896
}
896897

898+
if (options.sdkMajorVersion >= 4) {
899+
test(
900+
"captures sync streamText()/streamObject() paths in v4+",
901+
testConfig,
902+
() => {
903+
const root = findLatestSpan(events, ROOT_NAME);
904+
const streamTrace = findStreamTrace(events);
905+
906+
expectOperationParentedByRoot(streamTrace.operation, root);
907+
expectAISDKParentSpan(streamTrace.parent);
908+
expect(operationName(streamTrace.operation)).toBe("stream");
909+
expect(streamTrace.parent?.span.name).toBe("streamText");
910+
911+
if (options.supportsStreamObject) {
912+
const streamObjectTrace = findStreamObjectTrace(events);
913+
expectOperationParentedByRoot(streamObjectTrace.operation, root);
914+
expectAISDKParentSpan(streamObjectTrace.parent);
915+
expect(operationName(streamObjectTrace.operation)).toBe(
916+
"stream-object",
917+
);
918+
expect(streamObjectTrace.parent?.span.name).toBe("streamObject");
919+
}
920+
},
921+
);
922+
}
923+
897924
if (options.agentSpanName) {
898925
test("captures trace for agent.generate()", testConfig, () => {
899926
const root = findLatestSpan(events, ROOT_NAME);
@@ -922,6 +949,18 @@ export function defineAISDKInstrumentationAssertions(options: {
922949
expect(trace.modelChildren.length).toBeGreaterThanOrEqual(1);
923950
expect(trace.latestChild?.output).toBeDefined();
924951
});
952+
953+
if (options.sdkMajorVersion === 5 && options.agentSpanName === "Agent") {
954+
test("captures Agent.stream() path in v5", testConfig, () => {
955+
const root = findLatestSpan(events, ROOT_NAME);
956+
const trace = findAgentStreamTrace(events, "Agent");
957+
958+
expectOperationParentedByRoot(trace.operation, root);
959+
expectAISDKParentSpan(trace.parent);
960+
expect(operationName(trace.operation)).toBe("agent-stream");
961+
expect(trace.parent?.span.name).toBe("Agent.stream");
962+
});
963+
}
925964
}
926965

927966
if (options.supportsDenyOutputOverrideScenario) {

‎e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ for (const scenario of aiSDKScenarios) {
5353
supportsOutputObjectScenario,
5454
supportsStreamObject: scenario.supportsStreamObject,
5555
supportsToolExecution: scenario.supportsToolExecution,
56+
sdkMajorVersion,
5657
testFileUrl: import.meta.url,
5758
timeoutMs: AI_SDK_SCENARIO_TIMEOUT_MS,
5859
});
@@ -76,6 +77,7 @@ for (const scenario of aiSDKScenarios) {
7677
supportsOutputObjectScenario,
7778
supportsStreamObject: scenario.supportsStreamObject,
7879
supportsToolExecution: scenario.supportsToolExecution,
80+
sdkMajorVersion,
7981
testFileUrl: import.meta.url,
8082
timeoutMs: AI_SDK_SCENARIO_TIMEOUT_MS,
8183
});

‎generated_types.json‎

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1558,6 +1558,13 @@
15581558
"format": "uuid",
15591559
"description": "Identifies the user who created the dataset"
15601560
},
1561+
"tags": {
1562+
"type": ["array", "null"],
1563+
"items": {
1564+
"type": "string"
1565+
},
1566+
"description": "A list of tags for the dataset"
1567+
},
15611568
"metadata": {
15621569
"type": ["object", "null"],
15631570
"additionalProperties": {},
@@ -1745,6 +1752,45 @@
17451752
"root_span_id"
17461753
]
17471754
},
1755+
"DatasetSnapshot": {
1756+
"type": "object",
1757+
"properties": {
1758+
"id": {
1759+
"type": "string",
1760+
"format": "uuid",
1761+
"description": "Unique identifier for the dataset snapshot"
1762+
},
1763+
"dataset_id": {
1764+
"type": "string",
1765+
"format": "uuid",
1766+
"description": "Unique identifier for the dataset that this snapshot belongs to"
1767+
},
1768+
"name": {
1769+
"type": "string",
1770+
"description": "Name of the dataset snapshot"
1771+
},
1772+
"description": {
1773+
"type": ["string", "null"]
1774+
},
1775+
"xact_id": {
1776+
"type": "string",
1777+
"description": "Transaction id of the brainstore version at the time of the snapshot"
1778+
},
1779+
"created": {
1780+
"type": ["string", "null"],
1781+
"format": "date-time",
1782+
"description": "Date of dataset snapshot creation"
1783+
}
1784+
},
1785+
"required": [
1786+
"id",
1787+
"dataset_id",
1788+
"name",
1789+
"description",
1790+
"xact_id",
1791+
"created"
1792+
]
1793+
},
17481794
"EnvVar": {
17491795
"type": "object",
17501796
"properties": {
@@ -3996,7 +4042,7 @@
39964042
"description": "Textual description of the project automation"
39974043
},
39984044
"config": {
3999-
"anyOf": [
4045+
"oneOf": [
40004046
{
40014047
"type": "object",
40024048
"properties": {
@@ -5958,6 +6004,16 @@
59586004
"type": ["string", "null"],
59596005
"description": "Optional BTQL filter applied before topic automation."
59606006
},
6007+
"rerun_seconds": {
6008+
"type": ["number", "null"],
6009+
"minimum": 600,
6010+
"description": "How often to recompute topic maps"
6011+
},
6012+
"relabel_overlap_seconds": {
6013+
"type": ["number", "null"],
6014+
"minimum": 60,
6015+
"description": "How much recent history to relabel after a new topic map version becomes active"
6016+
},
59616017
"backfill_time_range": {
59626018
"anyOf": [
59636019
{
@@ -5979,7 +6035,7 @@
59796035
"type": "null"
59806036
}
59816037
],
5982-
"description": "Optional default time range for backfill operations."
6038+
"description": "Topic window used for classification coverage and initial backfill."
59836039
}
59846040
},
59856041
"required": [
@@ -6060,6 +6116,9 @@
60606116
},
60616117
"description": "Mapping from topic_id to topic name"
60626118
},
6119+
"generation_settings": {
6120+
"$ref": "#/components/schemas/TopicMapGenerationSettings"
6121+
},
60636122
"distance_threshold": {
60646123
"type": "number",
60656124
"description": "Maximum distance to nearest centroid. If exceeded, returns no_match."
@@ -6087,6 +6146,44 @@
60876146
},
60886147
"required": ["function"]
60896148
},
6149+
"TopicMapGenerationSettings": {
6150+
"type": "object",
6151+
"properties": {
6152+
"algorithm": {
6153+
"type": "string",
6154+
"enum": ["hdbscan", "kmeans"]
6155+
},
6156+
"dimension_reduction": {
6157+
"type": "string",
6158+
"enum": ["umap", "pca", "none"]
6159+
},
6160+
"sample_size": {
6161+
"type": "integer",
6162+
"exclusiveMinimum": 0
6163+
},
6164+
"n_clusters": {
6165+
"type": "integer",
6166+
"exclusiveMinimum": 0
6167+
},
6168+
"min_cluster_size": {
6169+
"type": "integer",
6170+
"exclusiveMinimum": 0
6171+
},
6172+
"min_samples": {
6173+
"type": "integer",
6174+
"exclusiveMinimum": 0
6175+
},
6176+
"hierarchy_threshold": {
6177+
"type": "integer",
6178+
"exclusiveMinimum": 0
6179+
},
6180+
"naming_model": {
6181+
"type": "string"
6182+
}
6183+
},
6184+
"required": ["algorithm", "dimension_reduction"],
6185+
"description": "Clustering and naming settings used to generate this topic map"
6186+
},
60906187
"TraceScope": {
60916188
"type": "object",
60926189
"properties": {
@@ -6199,6 +6296,10 @@
61996296
"type": ["string", "null"],
62006297
"format": "date-time",
62016298
"description": "Date of user creation"
6299+
},
6300+
"last_active_at": {
6301+
"type": ["number", "null"],
6302+
"description": "Unix timestamp in milliseconds of the user's last activity, when available"
62026303
}
62036304
},
62046305
"required": ["id"]

‎js/src/auto-instrumentations/configs/ai-sdk.ts‎

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,19 +39,33 @@ export const aiSDKConfigs: InstrumentationConfig[] = [
3939
},
4040
},
4141

42-
// streamText - function returning stream
42+
// streamText - async function (v3 only, before the sync refactor in v4)
4343
{
4444
channelName: aiSDKChannels.streamText.channelName,
4545
module: {
4646
name: "ai",
47-
versionRange: ">=3.0.0",
47+
versionRange: ">=3.0.0 <4.0.0",
4848
filePath: "dist/index.mjs",
4949
},
5050
functionQuery: {
5151
functionName: "streamText",
5252
kind: "Async",
5353
},
5454
},
55+
56+
// streamText - sync function returning stream (v4+)
57+
{
58+
channelName: aiSDKChannels.streamTextSync.channelName,
59+
module: {
60+
name: "ai",
61+
versionRange: ">=4.0.0",
62+
filePath: "dist/index.mjs",
63+
},
64+
functionQuery: {
65+
functionName: "streamText",
66+
kind: "Sync",
67+
},
68+
},
5569
{
5670
channelName: aiSDKChannels.streamText.channelName,
5771
module: {
@@ -143,19 +157,33 @@ export const aiSDKConfigs: InstrumentationConfig[] = [
143157
},
144158
},
145159

146-
// streamObject - function returning stream
160+
// streamObject - async function (v3 only, before the sync refactor in v4)
147161
{
148162
channelName: aiSDKChannels.streamObject.channelName,
149163
module: {
150164
name: "ai",
151-
versionRange: ">=3.0.0",
165+
versionRange: ">=3.0.0 <4.0.0",
152166
filePath: "dist/index.mjs",
153167
},
154168
functionQuery: {
155169
functionName: "streamObject",
156170
kind: "Async",
157171
},
158172
},
173+
174+
// streamObject - sync function returning stream (v4+)
175+
{
176+
channelName: aiSDKChannels.streamObjectSync.channelName,
177+
module: {
178+
name: "ai",
179+
versionRange: ">=4.0.0",
180+
filePath: "dist/index.mjs",
181+
},
182+
functionQuery: {
183+
functionName: "streamObject",
184+
kind: "Sync",
185+
},
186+
},
159187
{
160188
channelName: aiSDKChannels.streamObject.channelName,
161189
module: {
@@ -199,32 +227,32 @@ export const aiSDKConfigs: InstrumentationConfig[] = [
199227
},
200228
},
201229

202-
// Agent.stream - async method (v5 only)
230+
// Agent.stream - sync method (v5 only)
203231
// The compiled AI SDK bundle emits this as an anonymous class method, so we
204-
// target the first async `stream` method in the file instead of a class name.
232+
// target the first sync `stream` method in the file instead of a class name.
205233
{
206-
channelName: aiSDKChannels.agentStream.channelName,
234+
channelName: aiSDKChannels.agentStreamSync.channelName,
207235
module: {
208236
name: "ai",
209237
versionRange: ">=5.0.0 <6.0.0",
210238
filePath: "dist/index.mjs",
211239
},
212240
functionQuery: {
213241
methodName: "stream",
214-
kind: "Async",
242+
kind: "Sync",
215243
index: 0,
216244
},
217245
},
218246
{
219-
channelName: aiSDKChannels.agentStream.channelName,
247+
channelName: aiSDKChannels.agentStreamSync.channelName,
220248
module: {
221249
name: "ai",
222250
versionRange: ">=5.0.0 <6.0.0",
223251
filePath: "dist/index.js",
224252
},
225253
functionQuery: {
226254
methodName: "stream",
227-
kind: "Async",
255+
kind: "Sync",
228256
index: 0,
229257
},
230258
},

‎js/src/edge-light/config.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,13 @@
11
import iso from "../isomorph";
2+
import type {
3+
IsoTracingChannel,
4+
IsoTracingChannelCollection,
5+
} from "../isomorph";
26
import { _internalSetInitialState } from "../logger";
37
import { resolveRuntimeAsyncLocalStorage } from "../runtime-async-local-storage";
8+
import { tracingChannel } from "dc-browser";
9+
import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel";
10+
import { registry } from "../instrumentation/registry";
411

512
let edgeLightConfigured = false;
613

@@ -20,6 +27,13 @@ export function configureEdgeLight(): void {
2027
iso.newAsyncLocalStorage = <T>() => new runtimeAsyncLocalStorage<T>();
2128
}
2229

30+
iso.newTracingChannel = <M = unknown>(
31+
nameOrChannels: string | IsoTracingChannelCollection<M>,
32+
) =>
33+
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
34+
tracingChannel(nameOrChannels as string | object) as IsoTracingChannel<M>;
35+
patchTracingChannel(tracingChannel);
36+
2337
iso.getEnv = (name: string) => {
2438
if (typeof process === "undefined" || typeof process.env === "undefined") {
2539
return undefined;
@@ -39,5 +53,6 @@ export function configureEdgeLight(): void {
3953
};
4054

4155
_internalSetInitialState();
56+
registry.enable();
4257
edgeLightConfigured = true;
4358
}

0 commit comments

Comments
 (0)