diff --git a/.changeset/fix-cursor-sdk-tracing.md b/.changeset/fix-cursor-sdk-tracing.md new file mode 100644 index 000000000..212c77c7c --- /dev/null +++ b/.changeset/fix-cursor-sdk-tracing.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(cursor): improve automatic and manual sdk tracing diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.json b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.json index 05ee68fd7..93d07c65b 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.json +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.json @@ -12,15 +12,8 @@ "type": "task", "children": [], "input": "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files.", - "output": { - "durationMs": 0, - "id": "", - "model": { - "id": "" - }, - "status": "error" - }, "metadata": { + "cursor_sdk.agent_id": "agent-", "cursor_sdk.duration_ms": 0, "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", "cursor_sdk.model": "composer-2", @@ -32,9 +25,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -60,9 +51,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -88,9 +77,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -115,9 +102,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.txt b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.txt index 231f2c85d..6a07930d2 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.txt +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-auto-hook.span-tree.txt @@ -11,15 +11,8 @@ span_tree: │ } │ └── Cursor Agent [task] │ input: "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files." - │ output: { - │ "durationMs": 0, - │ "id": "", - │ "model": { - │ "id": "" - │ }, - │ "status": "error" - │ } │ metadata: { + │ "cursor_sdk.agent_id": "agent-", │ "cursor_sdk.duration_ms": 0, │ "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", │ "cursor_sdk.model": "composer-2", @@ -31,9 +24,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-stream-operation │ metadata: { │ "operation": "stream", @@ -51,9 +42,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-wait-operation │ metadata: { │ "operation": "wait", @@ -71,9 +60,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" └── cursor-sdk-resume-conversation-operation metadata: { "operation": "resume-conversation", @@ -90,6 +77,4 @@ span_tree: "model": "composer-2", "provider": "cursor" } - metrics: { - "duration": 0 - } + error: "Cursor run failed" diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.json b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.json index 05ee68fd7..93d07c65b 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.json +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.json @@ -12,15 +12,8 @@ "type": "task", "children": [], "input": "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files.", - "output": { - "durationMs": 0, - "id": "", - "model": { - "id": "" - }, - "status": "error" - }, "metadata": { + "cursor_sdk.agent_id": "agent-", "cursor_sdk.duration_ms": 0, "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", "cursor_sdk.model": "composer-2", @@ -32,9 +25,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -60,9 +51,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -88,9 +77,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -115,9 +102,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.txt b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.txt index 231f2c85d..6a07930d2 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.txt +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-auto-hook.span-tree.txt @@ -11,15 +11,8 @@ span_tree: │ } │ └── Cursor Agent [task] │ input: "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files." - │ output: { - │ "durationMs": 0, - │ "id": "", - │ "model": { - │ "id": "" - │ }, - │ "status": "error" - │ } │ metadata: { + │ "cursor_sdk.agent_id": "agent-", │ "cursor_sdk.duration_ms": 0, │ "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", │ "cursor_sdk.model": "composer-2", @@ -31,9 +24,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-stream-operation │ metadata: { │ "operation": "stream", @@ -51,9 +42,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-wait-operation │ metadata: { │ "operation": "wait", @@ -71,9 +60,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" └── cursor-sdk-resume-conversation-operation metadata: { "operation": "resume-conversation", @@ -90,6 +77,4 @@ span_tree: "model": "composer-2", "provider": "cursor" } - metrics: { - "duration": 0 - } + error: "Cursor run failed" diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.json b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.json index 20ac33ba5..64e3befc5 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.json +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.json @@ -12,14 +12,6 @@ "type": "task", "children": [], "input": "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files.", - "output": { - "durationMs": 0, - "id": "", - "model": { - "id": "" - }, - "status": "error" - }, "metadata": { "cursor_sdk.duration_ms": 0, "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", @@ -31,9 +23,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -59,9 +49,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -87,9 +75,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -114,9 +100,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.txt b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.txt index 4172c2799..935233ea0 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.txt +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-latest-wrapped.span-tree.txt @@ -11,14 +11,6 @@ span_tree: │ } │ └── Cursor Agent [task] │ input: "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files." - │ output: { - │ "durationMs": 0, - │ "id": "", - │ "model": { - │ "id": "" - │ }, - │ "status": "error" - │ } │ metadata: { │ "cursor_sdk.duration_ms": 0, │ "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", @@ -30,9 +22,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-stream-operation │ metadata: { │ "operation": "stream", @@ -50,9 +40,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-wait-operation │ metadata: { │ "operation": "wait", @@ -70,9 +58,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" └── cursor-sdk-resume-conversation-operation metadata: { "operation": "resume-conversation", @@ -89,6 +75,4 @@ span_tree: "model": "composer-2", "provider": "cursor" } - metrics: { - "duration": 0 - } + error: "Cursor run failed" diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.json b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.json index 20ac33ba5..64e3befc5 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.json +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.json @@ -12,14 +12,6 @@ "type": "task", "children": [], "input": "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files.", - "output": { - "durationMs": 0, - "id": "", - "model": { - "id": "" - }, - "status": "error" - }, "metadata": { "cursor_sdk.duration_ms": 0, "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", @@ -31,9 +23,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -59,9 +49,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -87,9 +75,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { @@ -114,9 +100,7 @@ "model": "composer-2", "provider": "cursor" }, - "metrics": { - "duration": 0 - } + "error": "Cursor run failed" } ], "metadata": { diff --git a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.txt b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.txt index 4172c2799..935233ea0 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.txt +++ b/e2e/scenarios/cursor-sdk-instrumentation/__snapshots__/cursor-sdk-v1-wrapped.span-tree.txt @@ -11,14 +11,6 @@ span_tree: │ } │ └── Cursor Agent [task] │ input: "Reply with exactly: CURSOR_PROMPT_OK. Do not modify files." - │ output: { - │ "durationMs": 0, - │ "id": "", - │ "model": { - │ "id": "" - │ }, - │ "status": "error" - │ } │ metadata: { │ "cursor_sdk.duration_ms": 0, │ "cursor_sdk.local.cwd": "/e2e/scenarios/cursor-sdk-instrumentation/", @@ -30,9 +22,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-stream-operation │ metadata: { │ "operation": "stream", @@ -50,9 +40,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" ├── cursor-sdk-wait-operation │ metadata: { │ "operation": "wait", @@ -70,9 +58,7 @@ span_tree: │ "model": "composer-2", │ "provider": "cursor" │ } - │ metrics: { - │ "duration": 0 - │ } + │ error: "Cursor run failed" └── cursor-sdk-resume-conversation-operation metadata: { "operation": "resume-conversation", @@ -89,6 +75,4 @@ span_tree: "model": "composer-2", "provider": "cursor" } - metrics: { - "duration": 0 - } + error: "Cursor run failed" diff --git a/e2e/scenarios/cursor-sdk-instrumentation/assertions.ts b/e2e/scenarios/cursor-sdk-instrumentation/assertions.ts index 94e5427ce..cf317e9a7 100644 --- a/e2e/scenarios/cursor-sdk-instrumentation/assertions.ts +++ b/e2e/scenarios/cursor-sdk-instrumentation/assertions.ts @@ -6,17 +6,8 @@ import { withScenarioHarness, type ScenarioRunContext, } from "../../helpers/scenario-harness"; -import { - findAllSpans, - findChildSpans, - findLatestSpan, -} from "../../helpers/trace-selectors"; -import { - matchSpanTreeSnapshot, - spanTreeFields, - type SpanTreeEntry, - type SpanTreeFields, -} from "../../helpers/span-tree"; +import { findChildSpans, findLatestSpan } from "../../helpers/trace-selectors"; +import { matchSpanTreeSnapshot } from "../../helpers/span-tree"; import { ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs"; type RunCursorSDKScenario = (harness: { @@ -35,48 +26,6 @@ type RunCursorSDKScenario = (harness: { }) => Promise; }) => Promise; -const METADATA_KEYS = [ - "provider", - "model", - "operation", - "scenario", - "gen_ai.tool.name", - "cursor_sdk.model", - "cursor_sdk.operation", - "cursor_sdk.agent_id", - "cursor_sdk.run_id", - "cursor_sdk.runtime", - "cursor_sdk.status", - "cursor_sdk.step_types", - "cursor_sdk.tool.status", -] as const; - -function snapshotFields(event: CapturedLogEvent): SpanTreeFields { - const fields = spanTreeFields(event); - const metadata = - fields.metadata && - typeof fields.metadata === "object" && - !Array.isArray(fields.metadata) - ? Object.fromEntries( - Object.entries(fields.metadata).filter(([key]) => - METADATA_KEYS.includes(key as (typeof METADATA_KEYS)[number]), - ), - ) - : undefined; - - if (metadata && typeof metadata["cursor_sdk.agent_id"] === "string") { - metadata["cursor_sdk.agent_id"] = ""; - } - if (metadata && typeof metadata["cursor_sdk.run_id"] === "string") { - metadata["cursor_sdk.run_id"] = ""; - } - - return { - ...fields, - metadata, - }; -} - function findOperation(events: CapturedLogEvent[], name: string) { return findLatestSpan(events, name); } @@ -86,86 +35,6 @@ function findCursorTask(events: CapturedLogEvent[], operationName: string) { return findChildSpans(events, "Cursor Agent", operation?.span.id).at(-1); } -function findSubagentTool( - events: CapturedLogEvent[], - parentId: string | undefined, -) { - if (!parentId) { - return undefined; - } - return [...events] - .reverse() - .find( - (event) => - event.span.type === "tool" && - event.span.parentIds.includes(parentId) && - ["tool: Agent", "tool: Task", "tool: task"].includes( - event.span.name ?? "", - ), - ); -} - -function findSubagentTask( - events: CapturedLogEvent[], - parentId: string | undefined, -) { - if (!parentId) { - return undefined; - } - return [...events] - .reverse() - .find( - (event) => - event.span.type === "task" && - event.span.parentIds.includes(parentId) && - event.span.name?.startsWith("Agent:"), - ); -} - -function outputText(event: CapturedLogEvent | undefined): string { - return typeof event?.output === "string" ? event.output : ""; -} - -function summarize(events: CapturedLogEvent[]): SpanTreeEntry[] { - const promptTask = findCursorTask(events, "cursor-sdk-prompt-operation"); - const streamTask = findCursorTask(events, "cursor-sdk-stream-operation"); - const waitTask = findCursorTask(events, "cursor-sdk-wait-operation"); - const conversationTask = findCursorTask( - events, - "cursor-sdk-resume-conversation-operation", - ); - const tool = findAllSpans(events, "tool: shell").at(-1); - const subagentTool = findSubagentTool(events, streamTask?.span.id); - const subagentTask = findSubagentTask(events, subagentTool?.span.id); - - return [ - findLatestSpan(events, ROOT_NAME), - findOperation(events, "cursor-sdk-prompt-operation"), - promptTask, - findOperation(events, "cursor-sdk-stream-operation"), - streamTask, - tool, - subagentTool, - subagentTask, - findOperation(events, "cursor-sdk-wait-operation"), - waitTask, - findOperation(events, "cursor-sdk-resume-conversation-operation"), - conversationTask, - ].flatMap((event) => - event - ? [ - { - event, - fields: snapshotFields(event), - name: event.span.name?.startsWith("Agent:") - ? "Agent: " - : undefined, - }, - ] - : [], - ); -} - export function defineCursorSDKInstrumentationAssertions(options: { name: string; runScenario: RunCursorSDKScenario; @@ -190,13 +59,6 @@ export function defineCursorSDKInstrumentationAssertions(options: { }); }, timeoutMs); - test("captures the root trace", testConfig, () => { - const root = findLatestSpan(events, ROOT_NAME); - - expect(root).toBeDefined(); - expect(root?.row.metadata).toMatchObject({ scenario: SCENARIO_NAME }); - }); - test( "captures Cursor Agent task spans for run-producing APIs", testConfig, @@ -211,6 +73,9 @@ export function defineCursorSDKInstrumentationAssertions(options: { const task = findCursorTask(events, operationName); expect(operation).toBeDefined(); + expect( + findChildSpans(events, "Cursor Agent", operation?.span.id), + ).toHaveLength(1); expect(task).toBeDefined(); expect(task?.span.parentIds).toEqual([operation?.span.id ?? ""]); expect(task?.row.metadata).toMatchObject({ @@ -220,97 +85,13 @@ export function defineCursorSDKInstrumentationAssertions(options: { }, ); - test( - "captures tool spans when Cursor surfaces tool calls", - testConfig, - () => { - const streamTask = findCursorTask( - events, - "cursor-sdk-stream-operation", - ); - expect(streamTask).toBeDefined(); - const toolSpans = events.filter( - (event) => - event.span.type === "tool" && - event.span.parentIds.includes(streamTask?.span.id ?? ""), - ); - - if (toolSpans.length === 0) { - return; - } - expect( - toolSpans.some( - (event) => - event.input !== undefined && - event.output !== undefined && - event.metadata?.["cursor_sdk.tool.status"] === "completed", - ), - ).toBe(true); - expect( - JSON.stringify(toolSpans.map((event) => event.output)), - ).toContain("cursor_tool_ok"); - }, - ); - - test("captures subagent spans when Cursor uses agents", testConfig, () => { - const streamTask = findCursorTask(events, "cursor-sdk-stream-operation"); - const subagentTool = findSubagentTool(events, streamTask?.span.id); - const subagentTask = findSubagentTask(events, subagentTool?.span.id); - - expect(streamTask).toBeDefined(); - if (!subagentTool) { - return; - } - expect(subagentTool).toBeDefined(); - expect(subagentTool?.metadata).toMatchObject({ - "cursor_sdk.tool.status": "completed", - }); - expect(subagentTask).toBeDefined(); - expect(subagentTask?.span.rootId).toBe(streamTask?.span.rootId); - expect(subagentTask?.metadata).toMatchObject({ - "cursor_sdk.tool.status": "completed", - }); - expect(subagentTask?.output).toBeDefined(); - }); - - test( - "preserves user callbacks when Cursor emits updates", - testConfig, - () => { - const waitTask = findCursorTask(events, "cursor-sdk-wait-operation"); - expect(waitTask).toBeDefined(); - - const deltaSpan = findLatestSpan(events, "cursor-sdk-user-on-delta"); - const stepSpan = findLatestSpan(events, "cursor-sdk-user-on-step"); - if (!deltaSpan && !stepSpan) { - return; - } - - expect(deltaSpan).toBeDefined(); - expect(stepSpan).toBeDefined(); - expect(waitTask?.metrics?.["cursor_sdk.step_duration_ms"]).toEqual( - expect.any(Number), - ); - }, - ); - - test("captures conversation task span", testConfig, () => { - const conversationTask = findCursorTask( - events, - "cursor-sdk-resume-conversation-operation", - ); - - expect(conversationTask).toBeDefined(); - const text = outputText(conversationTask); - if (text.length > 0) { - expect(text).toEqual(expect.any(String)); - } - }); - test( "matches the shared span tree snapshot", testConfig, async ({ expect }) => { + const root = findLatestSpan(events, ROOT_NAME); + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ scenario: SCENARIO_NAME }); await matchSpanTreeSnapshot(events, snapshotPath, { snapshotExpect: expect, }); diff --git a/js/src/auto-instrumentations/configs/cursor-sdk.ts b/js/src/auto-instrumentations/configs/cursor-sdk.ts index 65d8d12ed..7973ecd66 100644 --- a/js/src/auto-instrumentations/configs/cursor-sdk.ts +++ b/js/src/auto-instrumentations/configs/cursor-sdk.ts @@ -2,45 +2,21 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { cursorSDKChannels } from "../../instrumentation/plugins/cursor-sdk-interceptors"; const cursorSDKVersionRange = ">=1.0.7 <2.0.0"; - const cursorSDKEntrypoints = ["dist/esm/index.js", "dist/cjs/index.js"]; +const cursorSDKMethods = ["create", "resume", "prompt"] as const; export const cursorSDKConfigs: InstrumentationConfig[] = - cursorSDKEntrypoints.flatMap((filePath) => [ - { - channelName: cursorSDKChannels.create.channelName, - module: { - name: "@cursor/sdk", - versionRange: cursorSDKVersionRange, - filePath, - }, - functionQuery: { - className: "Agent", - methodName: "create", - }, - }, - { - channelName: cursorSDKChannels.resume.channelName, - module: { - name: "@cursor/sdk", - versionRange: cursorSDKVersionRange, - filePath, - }, - functionQuery: { - className: "Agent", - methodName: "resume", - }, - }, - { - channelName: cursorSDKChannels.prompt.channelName, + cursorSDKEntrypoints.flatMap((filePath) => + cursorSDKMethods.map((methodName) => ({ + channelName: cursorSDKChannels[methodName].channelName, module: { name: "@cursor/sdk", versionRange: cursorSDKVersionRange, filePath, }, functionQuery: { - className: "Agent", - methodName: "prompt", + classMethods: [...cursorSDKMethods], + methodName, }, - }, - ]); + })), + ); diff --git a/js/src/auto-instrumentations/orchestrion-js/transformer.ts b/js/src/auto-instrumentations/orchestrion-js/transformer.ts index d05ee0085..b580e704f 100644 --- a/js/src/auto-instrumentations/orchestrion-js/transformer.ts +++ b/js/src/auto-instrumentations/orchestrion-js/transformer.ts @@ -187,10 +187,23 @@ export class Transformer { ); } else if ("methodName" in functionQuery) { const { methodName } = functionQuery; - queries.push( - `ClassBody > [key.name="${methodName}"][key.type=Identifier] > [async]`, - `Property[key.name="${methodName}"][key.type=Identifier] > [async]`, - ); + if (functionQuery.classMethods) { + // Bundlers can rename classes while retaining their public static API. + const shape = functionQuery.classMethods + .map( + (name) => + `:has(> MethodDefinition[static=true][key.name="${name}"])`, + ) + .join(""); + queries.push( + `ClassBody${shape} > MethodDefinition[static=true][key.name="${methodName}"] > FunctionExpression`, + ); + } else { + queries.push( + `ClassBody > [key.name="${methodName}"][key.type=Identifier] > [async]`, + `Property[key.name="${methodName}"][key.type=Identifier] > [async]`, + ); + } } if ("functionName" in functionQuery) { diff --git a/js/src/auto-instrumentations/orchestrion-js/types.ts b/js/src/auto-instrumentations/orchestrion-js/types.ts index 080499137..351c2d139 100644 --- a/js/src/auto-instrumentations/orchestrion-js/types.ts +++ b/js/src/auto-instrumentations/orchestrion-js/types.ts @@ -29,7 +29,7 @@ export type FunctionQuery = { methodName: string; isExportAlias?: boolean; } - | { methodName: string } + | { methodName: string; classMethods?: string[] } | { functionName: string; isExportAlias?: boolean; diff --git a/js/src/instrumentation/plugins/cursor-sdk-interceptors.ts b/js/src/instrumentation/plugins/cursor-sdk-interceptors.ts index 29ab5789c..a07ba2b11 100644 --- a/js/src/instrumentation/plugins/cursor-sdk-interceptors.ts +++ b/js/src/instrumentation/plugins/cursor-sdk-interceptors.ts @@ -40,3 +40,29 @@ export const cursorSDKChannels = defineInterceptor("@cursor/sdk", { channelName: "agent.send", }), }); + +// Shared by manual wrappers and automatic factory hooks. This patch only invokes +// wrapping hooks and remains safe when tracing is disabled. +export function patchCursorAgentInPlace(agent: unknown): void { + if (!agent || typeof agent !== "object") return; + const candidate = agent as CursorSDKAgent; + const marker = Symbol.for("braintrust.cursor-sdk.auto-patched-agent"); + if (candidate[marker] || typeof candidate.send !== "function") return; + const send = candidate.send; + if (!Object.isExtensible(candidate)) return; + try { + Object.defineProperty(candidate, "send", { + configurable: true, + writable: true, + value(...args: Parameters) { + return cursorSDKChannels.send.invoke(send, candidate, args, { + agent: candidate, + operation: "send", + }); + }, + }); + Object.defineProperty(candidate, marker, { value: true }); + } catch { + // A non-writable send method must not prevent use of the agent. + } +} diff --git a/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts b/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts index 7005d5b97..cf0d10948 100644 --- a/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts @@ -1,270 +1,428 @@ +import { AsyncLocalStorage } from "node:async_hooks"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { newGlobalInvocationHook } from "../../global-instrumentation-hooks"; import { invocationController } from "../test-utils/invocation"; + vi.mock("../../global-instrumentation-hooks", async (importOriginal) => ({ ...(await importOriginal< typeof import("../../global-instrumentation-hooks") >()), newGlobalInvocationHook: vi.fn(), })); - -const { mockStartSpan } = vi.hoisted(() => ({ - mockStartSpan: vi.fn(), -})); - +const { mockStartSpan } = vi.hoisted(() => ({ mockStartSpan: vi.fn() })); vi.mock("../../isomorph", () => ({ - default: {}, + default: { newAsyncLocalStorage: () => new AsyncLocalStorage() }, })); - vi.mock("../../logger", () => ({ startSpan: (...args: unknown[]) => mockStartSpan(...args), + withCurrent: (_span: unknown, call: () => unknown) => call(), })); - import { CursorSDKPlugin } from "./cursor-sdk-plugin"; -const mockNewInvocationHook = newGlobalInvocationHook as ReturnType< - typeof vi.fn ->; - describe("CursorSDKPlugin", () => { - let handlersByName: Map; - let spans: Array<{ - end: ReturnType; - export: ReturnType; - log: ReturnType; - name?: string; - }>; - + let handlers: Map>; + let spans: any[]; + let plugin: CursorSDKPlugin; + const settled = () => new Promise((resolve) => setTimeout(resolve, 0)); + const handler = (name: string) => + handlers.get(`orchestrion:@cursor/sdk:${name}`)!; + const root = () => spans.find((s) => s.name === "Cursor Agent"); + const finalLog = () => root().log.mock.calls.at(-1)[0]; + function send(run: any, options?: any) { + let observedOptions: any; + const result = handler("agent.send").call( + { arguments: ["hello", options] }, + (_message?: any, opts?: any) => { + observedOptions = opts; + return Promise.resolve(run); + }, + ); + return { result, options: observedOptions }; + } beforeEach(() => { - handlersByName = new Map(); spans = []; - mockNewInvocationHook.mockImplementation((name: string) => ({ - intercept: vi.fn((interceptor) => { - handlersByName.set(name, invocationController(interceptor)); - return vi.fn(); - }), - invoke: vi.fn((fn, receiver, args) => Reflect.apply(fn, receiver, args)), - unsubscribe: vi.fn(), - })); + handlers = new Map(); + vi.mocked(newGlobalInvocationHook).mockImplementation( + (name: string) => + ({ + intercept: vi.fn((interceptor) => { + handlers.set(name, invocationController(interceptor)); + return vi.fn(); + }), + invoke: vi.fn((fn, receiver, args) => + Reflect.apply(fn, receiver, args), + ), + }) as any, + ); mockStartSpan.mockImplementation((args: any) => { + const id = `span-${spans.length}`; const span = { + ...args, + id, end: vi.fn(), - export: vi.fn(async () => `${args.name}-export-${spans.length}`), + export: vi.fn(async () => id), log: vi.fn(), - name: args.name, }; - if (args.event) { - span.log(args.event); - } + if (args.event) span.log(args.event); spans.push(span); return span; }); + plugin = new CursorSDKPlugin(); + plugin.enable(); }); - afterEach(() => { + plugin.disable(); vi.clearAllMocks(); }); - it("subscribes to Cursor SDK channels", () => { - const plugin = new CursorSDKPlugin(); - - plugin.enable(); - - expect(handlersByName.has("orchestrion:@cursor/sdk:Agent.create")).toBe( - true, - ); - expect(handlersByName.has("orchestrion:@cursor/sdk:Agent.resume")).toBe( - true, - ); - expect(handlersByName.has("orchestrion:@cursor/sdk:Agent.prompt")).toBe( - true, - ); - expect(handlersByName.has("orchestrion:@cursor/sdk:agent.send")).toBe(true); + it("patches factory agents once without replacing their promises", async () => { + const agent = { send: vi.fn(async (_message?: string) => makeRun()) }; + const originalSend = agent.send; + let patchedSend: typeof agent.send | undefined; + for (const [name, args] of [ + ["Agent.create", [{}]], + ["Agent.resume", ["agent-1"]], + ] as const) { + const promise = Promise.resolve(agent); + expect( + handler(name).call( + { arguments: [...args], moduleVersion: "1.0.12" }, + () => promise, + ), + ).toBe(promise); + await promise; + expect(agent.send).not.toBe(originalSend); + if (patchedSend) expect(agent.send).toBe(patchedSend); + patchedSend = agent.send; + } + await agent.send("hi"); + expect(originalSend).toHaveBeenCalledExactlyOnceWith("hi"); + expect(spans).toHaveLength(0); }); - it("patches agents returned by Agent.create and traces send/wait", async () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); - - const createHandlers = handlersByName.get( - "orchestrion:@cursor/sdk:Agent.create", - ); - const sendHandlers = handlersByName.get( - "orchestrion:@cursor/sdk:agent.send", - ); - const run = makeRun(); - const originalSend = vi.fn(async () => run); - const agent: any = { - agentId: "agent-1", - send: originalSend, - }; - - createHandlers.call({ - arguments: [{ local: { cwd: "/tmp/repo" } }], - result: agent, + it("captures cache-inclusive usage without adding cache buckets twice", async () => { + const { result, options } = send(makeRun()); + await result; + options.onDelta({ update: { type: "token-delta", tokens: 500 } }); + options.onDelta({ + update: { + type: "turn-ended", + usage: { + inputTokens: 3, + outputTokens: 4, + cacheReadTokens: 1, + cacheWriteTokens: 2, + }, + }, }); + options.onDelta({ + update: { + type: "turn-ended", + usage: { + inputTokens: 5, + outputTokens: 2, + cacheReadTokens: 0, + cacheWriteTokens: 1, + }, + }, + }); + await ((await result) as any).wait(); + await settled(); + expect(finalLog().metrics).toMatchObject({ + prompt_tokens: 12, + completion_tokens: 6, + tokens: 18, + prompt_cached_tokens: 1, + prompt_cache_creation_tokens: 3, + }); + expect(finalLog().metrics).not.toHaveProperty("cursor_sdk.delta_tokens"); + expect(root().end).toHaveBeenCalledOnce(); + }); - const patchedRun = await agent.send("use a tool", {}); - expect(patchedRun).toBe(run); - expect(originalSend).toHaveBeenCalledTimes(1); - - const sendEvent = { - agent, - arguments: ["use a tool", {}], - result: run, - }; - sendHandlers.begin(sendEvent); - sendHandlers.resolve(sendEvent); - + it("does not fabricate missing totals or accept invalid token counts", async () => { + const { result, options } = send(makeRun()); + const run = (await result) as any; + options.onDelta({ + update: { + type: "turn-ended", + usage: { inputTokens: -1, outputTokens: 2.5, cacheReadTokens: NaN }, + }, + }); + options.onDelta({ + update: { type: "turn-ended", usage: { outputTokens: 3 } }, + }); await run.wait(); - - const rootSpan = spans.find((span) => span.name === "Cursor Agent"); - expect(rootSpan?.log).toHaveBeenCalledWith( - expect.objectContaining({ - input: "use a tool", - metadata: expect.objectContaining({ - "cursor_sdk.agent_id": "agent-1", - provider: "cursor", - }), - }), - ); - expect(rootSpan?.log).toHaveBeenCalledWith( - expect.objectContaining({ - output: "done", - metadata: expect.objectContaining({ - "cursor_sdk.run_id": "run-1", - "cursor_sdk.status": "finished", - }), - }), - ); - expect(rootSpan?.end).toHaveBeenCalledTimes(1); + await settled(); + expect(finalLog().metrics).not.toHaveProperty("tokens"); + expect(finalLog().metrics).not.toHaveProperty("prompt_tokens"); + expect(finalLog().metrics.completion_tokens).toBe(3); }); - it("captures stream tool calls and usage", async () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); - - const sendHandlers = handlersByName.get( - "orchestrion:@cursor/sdk:agent.send", - ); - const run = makeRun([ + it("updates tools, deduplicates event surfaces, and captures errors and truncation", async () => { + const error = new Error("tool failed"); + const messages = [ + { + type: "assistant", + message: { + content: [ + { type: "tool_use", id: "call-1", name: "shell", input: {} }, + ], + }, + }, { type: "tool_call", call_id: "call-1", name: "shell", status: "running", - args: { command: "echo hi" }, + args: { command: "pwd" }, }, { type: "tool_call", call_id: "call-1", name: "shell", - status: "completed", - result: { stdout: "hi\n" }, + status: "error", + result: error, + truncated: { args: false, result: true }, }, - { - type: "assistant", - message: { content: [{ type: "text", text: "done" }] }, + { type: "request", request_id: "req-1" }, + ]; + const run = makeRun(messages); + const { result, options } = send(run); + await result; + options.onDelta({ + update: { + type: "tool-call-started", + callId: "call-1", + toolCall: { type: "shell", args: {} }, }, - ]); - const event = { - agent: { agentId: "agent-1" }, - arguments: [ - "hello", - { - onDelta: vi.fn(), - }, - ], - result: run, - }; - - sendHandlers.begin(event); - await (event.arguments[1] as any).onDelta({ + }); + options.onDelta({ update: { - type: "turn-ended", - usage: { - inputTokens: 3, - outputTokens: 4, - cacheReadTokens: 1, - cacheWriteTokens: 2, - }, + type: "tool-call-completed", + callId: "call-1", + toolCall: { type: "shell", result: error, status: "error" }, }, }); - sendHandlers.resolve(event); - + const stream = run.stream(); + expect(stream).toBe(run.originalStream); const chunks = []; - for await (const chunk of run.stream()) { - chunks.push(chunk); - } - - expect(chunks).toHaveLength(3); - const toolSpan = spans.find((span) => span.name === "tool: shell"); - expect(toolSpan?.log).toHaveBeenCalledWith( - expect.objectContaining({ - input: { command: "echo hi" }, - }), - ); - expect(toolSpan?.log).toHaveBeenCalledWith( - expect.objectContaining({ - output: { stdout: "hi\n" }, - }), + for await (const message of stream) chunks.push(message); + await run.wait(); + await settled(); + const tools = spans.filter((s) => s.name === "shell"); + expect(tools).toHaveLength(1); + expect(tools[0].log).toHaveBeenCalledWith( + expect.objectContaining({ input: { command: "pwd" } }), ); - const rootSpan = spans.find((span) => span.name === "Cursor Agent"); - expect(rootSpan?.log).toHaveBeenCalledWith( + expect(tools[0].log).toHaveBeenCalledWith( expect.objectContaining({ - metrics: expect.objectContaining({ - completion_tokens: 4, - prompt_cache_creation_tokens: 2, - prompt_cached_tokens: 1, - prompt_tokens: 3, - }), - output: "done", + error, + output: error, + metadata: { + "cursor_sdk.tool.status": "error", + "cursor_sdk.tool.args_truncated": false, + "cursor_sdk.tool.result_truncated": true, + }, }), ); + expect(tools[0].end).toHaveBeenCalledOnce(); + expect(finalLog().metadata.request_id).toBe("req-1"); + expect(chunks).toEqual(messages); }); - it("traces Agent.prompt without a nested send span", () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); - - const promptHandlers = handlersByName.get( - "orchestrion:@cursor/sdk:Agent.prompt", - ); - const sendHandlers = handlersByName.get( - "orchestrion:@cursor/sdk:agent.send", + it("records nested normalized and protobuf JSON tools under the subagent", async () => { + const toolCall = { + type: "task", + args: { description: "dynamic description" }, + result: { + status: "success", + value: { + conversationSteps: [ + { + type: "toolCall", + message: { + type: "read", + args: { path: "a" }, + result: { status: "error", error: { message: "missing" } }, + }, + }, + { + toolCall: { + shellToolCall: { + args: { command: "pwd" }, + result: { success: { stdout: "ok" } }, + }, + }, + }, + ], + }, + }, + }; + const { result, options } = send(makeRun()); + const run = (await result) as any; + options.onDelta({ + update: { type: "tool-call-completed", callId: "task-1", toolCall }, + }); + options.onDelta({ + update: { type: "tool-call-completed", callId: "task-1", toolCall }, + }); + await run.wait(); + await settled(); + expect(spans.map((s) => s.name)).toEqual([ + "Cursor Agent", + "task", + "Cursor Subagent", + "read", + "shell", + ]); + const subagent = spans.find((s) => s.name === "Cursor Subagent"); + expect(spans.find((s) => s.name === "read").parent).toBe(subagent.id); + expect(spans.find((s) => s.name === "shell").parent).toBe(subagent.id); + expect(spans.find((s) => s.name === "read").log).toHaveBeenCalledWith( + expect.objectContaining({ error: { message: "missing" } }), ); - const promptEvent = { arguments: ["hello", { local: { cwd: "/tmp" } }] }; + for (const span of spans) expect(span.end).toHaveBeenCalledOnce(); + }); - promptHandlers.begin(promptEvent); - sendHandlers.begin({ arguments: ["nested", {}] }); - promptHandlers.resolve( - Object.assign(promptEvent, { - result: { id: "run-1", result: "done", status: "finished" }, + it("preserves callback receivers, synchronous return values, and thrown errors", async () => { + const error = new Error("callback failure"); + const callback = vi.fn(function (this: unknown) { + expect(this).toBe(receiver); + return 42; + }); + const { result, options } = send(makeRun(), { + onDelta: callback, + onStep() { + throw error; + }, + }); + const receiver = {}; + expect( + options.onDelta.call(receiver, { + update: { type: "text-delta", text: "hi" }, }), - ); + ).toBe(42); + expect(() => options.onStep({ step: {} })).toThrow(error); + const run = (await result) as any; + await run.wait(); + await settled(); + }); + + it.each(["throw", "reject"])( + "preserves promises and error identity when wait fails via %s", + async (failure) => { + const error = new Error("wait failure"); + const waitPromise = + failure === "reject" + ? Object.assign(Promise.reject(error), { helper: () => "ok" }) + : undefined; + const run = { + ...makeRun(), + wait() { + if (waitPromise) return waitPromise; + throw error; + }, + }; + const sendPromise = Object.assign(Promise.resolve(run), { + helper: () => "send", + }); + const returned = handler("agent.send").call( + { arguments: ["hi"] }, + () => sendPromise, + ); + expect(returned).toBe(sendPromise); + await returned; + if (waitPromise) { + expect(run.wait()).toBe(waitPromise); + await expect(waitPromise).rejects.toBe(error); + } else { + expect(() => run.wait()).toThrow(error); + } + await settled(); + expect(finalLog().error).toBe(error); + expect(root().end).toHaveBeenCalledOnce(); + }, + ); + + it.each(["stream", "cancel"])( + "ends open tools on %s cancellation", + async (operation) => { + const run = makeRun([ + { + type: "tool_call", + call_id: "open", + name: "shell", + status: "running", + }, + ]); + const { result } = send(run); + await result; + const stream = run.stream(); + await stream.next(); + if (operation === "stream") await stream.return(); + else await run.cancel(); + await settled(); + expect(finalLog().metadata["cursor_sdk.status"]).toBe("cancelled"); + for (const span of spans) expect(span.end).toHaveBeenCalledOnce(); + }, + ); - expect(spans.filter((span) => span.name === "Cursor Agent")).toHaveLength( - 1, + it("keeps concurrent prompt/send traces independent and captures nested prompt tools", async () => { + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const pending = handler("Agent.prompt").call( + { arguments: ["prompt", {}] }, + async () => { + await gate; + const run = makeRun([ + { + type: "tool_call", + call_id: "nested", + name: "shell", + status: "completed", + result: "ok", + }, + ]); + await handler("agent.send").call({ arguments: ["prompt"] }, () => + Promise.resolve(run), + ); + for await (const _ of run.stream()) { + /* consume */ + } + return run.wait(); + }, ); + const independent = send(makeRun()); + await independent.result; + await ((await independent.result) as any).wait(); + release(); + await pending; + await settled(); + expect(spans.filter((s) => s.name === "Cursor Agent")).toHaveLength(2); + expect(spans.find((s) => s.name === "shell").parent).toBe(spans[0].id); + for (const span of spans) expect(span.end).toHaveBeenCalledOnce(); }); }); function makeRun(messages: unknown[] = []) { + const stream = (async function* () { + for (const message of messages) yield message; + })(); return { agentId: "agent-1", - async conversation() { - return []; - }, id: "run-1", result: "done", status: "finished", - stream: async function* () { - for (const message of messages) { - yield message; - } - }, + originalStream: stream, + stream: () => stream, async wait() { return { id: "run-1", result: "done", status: "finished" }; }, + async conversation() { + return []; + }, + async cancel() {}, }; } diff --git a/js/src/instrumentation/plugins/cursor-sdk-plugin.ts b/js/src/instrumentation/plugins/cursor-sdk-plugin.ts index c4e70c17d..a01cf4632 100644 --- a/js/src/instrumentation/plugins/cursor-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/cursor-sdk-plugin.ts @@ -1,10 +1,12 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { BasePlugin } from "../core"; import { observeResult, runInstrumentation } from "../core/observe-result"; -import { SpanTypeAttribute } from "../../../util/index"; +import { SpanTypeAttribute, isObject } from "../../../util/index"; import { debugLogger } from "../../debug-logger"; import type { Span } from "../../logger"; -import { startSpan as startBaseSpan } from "../../logger"; +import { startSpan as startBaseSpan, withCurrent } from "../../logger"; +import iso from "../../isomorph"; +import { patchStreamIfNeeded } from "../core/stream-patcher"; import { INSTRUMENTATION_NAMES, withSpanInstrumentationName, @@ -27,346 +29,203 @@ import type { CursorSDKUsage, CursorSDKUserMessage, } from "../../vendor-sdk-types/cursor-sdk"; -import type { ChannelMessage } from "../core/tracing-types"; -import { cursorSDKChannels } from "./cursor-sdk-interceptors"; +import { + cursorSDKChannels, + patchCursorAgentInPlace, +} from "./cursor-sdk-interceptors"; -const PATCHED_AGENT = Symbol.for("braintrust.cursor-sdk.auto-patched-agent"); const PATCHED_RUN = Symbol.for("braintrust.cursor-sdk.patched-run"); -const WRAPPED_AGENT = Symbol.for("braintrust.cursor-sdk.wrapped-agent"); type ToolState = { span: Span; subAgentSpan?: Span; + ended?: boolean; }; type CursorRunState = { activeToolSpans: Map; - agent?: CursorSDKAgent; + completedToolSpans: Map; + observation: Promise; + finalizing?: Promise; + error?: unknown; deltaText: string[]; streamText: string[]; stepText: string[]; - conversationOutput?: unknown; conversationText: string[]; finalized: boolean; - input: string | CursorSDKUserMessage; - lastResult?: CursorSDKRunResult; metadata: Record; metrics: Record; + usageComplete: boolean; run?: CursorSDKRun; span: Span; startTime: number; - streamMessages: CursorSDKMessage[]; taskText: string[]; }; -type PromptState = { - metadata: Record; - span: Span; - startTime: number; -}; +// Scope prompt/send coordination to this invocation, including across awaits. +let promptStore: ReturnType>; +function promptContext() { + return (promptStore ??= iso.newAsyncLocalStorage()); +} export class CursorSDKPlugin extends BasePlugin { - private promptDepth = 0; - - protected onEnable(): void { - this.subscribeToAgentFactories(); - this.subscribeToPrompt(); - this.subscribeToSend(); - } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } + for (const unsubscribe of this.unsubscribers) unsubscribe(); this.unsubscribers = []; - this.promptDepth = 0; - } - - private subscribeToAgentFactories(): void { - this.subscribeToAgentFactory(cursorSDKChannels.create); - this.subscribeToAgentFactory(cursorSDKChannels.resume); } - private subscribeToAgentFactory( - channel: typeof cursorSDKChannels.create | typeof cursorSDKChannels.resume, - ): void { + protected onEnable(): void { + for (const channel of [ + cursorSDKChannels.create, + cursorSDKChannels.resume, + ]) { + this.unsubscribers.push( + channel.intercept((target, receiver, args) => + observeResult( + Reflect.apply(target, receiver, args), + patchCursorAgentInPlace, + () => {}, + ), + ), + ); + } + // Preserve the same Cursor run tracing contract for automatic and manual wrapping. this.unsubscribers.push( - channel.intercept((target, receiver, args) => - observeResult( - Reflect.apply(target, receiver, args), - patchCursorAgentInPlace, - () => {}, + cursorSDKChannels.prompt.intercept((target, receiver, args, additional) => + traceCursorPrompt( + () => Reflect.apply(target, receiver, args), + args, + additional.moduleVersion, + ), + ), + cursorSDKChannels.send.intercept((target, receiver, args, additional) => + traceCursorSend( + (options) => Reflect.apply(target, receiver, [args[0], options]), + args, + additional.agent, + additional.moduleVersion, ), ), ); } +} - private subscribeToPrompt(): void { - const channel = cursorSDKChannels.prompt; - const states = new WeakMap(); - - const removeHandlers = channel.intercept( - (target, receiver, args, additional) => { - const event = { - ...additional, - arguments: args, - self: receiver, - } as ChannelMessage; - const prepare = ( - event: ChannelMessage, - ) => { - this.promptDepth += 1; - const message = event.arguments[0]; - const options = event.arguments[1]; - const metadata = { - ...extractAgentOptionsMetadata(options), - "cursor_sdk.operation": "Agent.prompt", - provider: "cursor", - ...(event.moduleVersion - ? { "cursor_sdk.version": event.moduleVersion } - : {}), - }; - const span = startBaseSpan( - withSpanInstrumentationName( - { - name: "Cursor Agent", - spanAttributes: { type: SpanTypeAttribute.TASK }, - }, - INSTRUMENTATION_NAMES.CURSOR_SDK, - ), - ); - const startTime = getCurrentUnixTimestamp(); - safeLog(span, { - input: sanitizeUserMessage(message), - metadata, - }); - states.set(event, { metadata, span, startTime }); - }; - const resolved = ( - event: ChannelMessage, - ) => { - this.promptDepth = Math.max(0, this.promptDepth - 1); - const state = states.get(event); - if (!state) { - return; - } - try { - safeLog(state.span, { - metadata: { - ...state.metadata, - ...extractRunResultMetadata(event.result), - }, - metrics: buildDurationMetrics(state.startTime), - output: event.result?.result ?? event.result, - }); - } finally { - state.span.end(); - states.delete(event); - } - }; - const failed = ( - event: ChannelMessage, - ) => { - this.promptDepth = Math.max(0, this.promptDepth - 1); - const state = states.get(event); - if (!state || !event.error) { - return; - } - safeLog(state.span, { error: event.error.message }); - state.span.end(); - states.delete(event); - }; - runInstrumentation(() => prepare(event)); - let result; - try { - result = Reflect.apply(target, receiver, args); - } catch (error) { - Object.assign(event, { error }); - runInstrumentation(() => failed(event)); - throw error; - } - return observeResult( - result, - (value) => { - Object.assign(event, { result: value }); - resolved(event); - }, - (error) => { - Object.assign(event, { error }); - failed(event); - }, - ); - }, - ); - this.unsubscribers.push(removeHandlers); - } - - private subscribeToSend(): void { - const channel = cursorSDKChannels.send; - const states = new WeakMap(); - - const removeHandlers = channel.intercept( - (target, receiver, args, additional) => { - const event = { - ...additional, - arguments: args, - self: receiver, - } as ChannelMessage; - const prepare = ( - event: ChannelMessage, - ) => { - if (this.promptDepth > 0) { - return; - } - - const message = event.arguments[0]; - const sendOptions = event.arguments[1]; - const agent = event.agent; - const metadata = { - ...extractSendMetadata(sendOptions), - ...(agent ? extractAgentMetadata(agent) : {}), - "cursor_sdk.operation": "agent.send", - provider: "cursor", - ...(event.moduleVersion - ? { "cursor_sdk.version": event.moduleVersion } - : {}), - }; - const span = startBaseSpan( - withSpanInstrumentationName( - { - name: "Cursor Agent", - spanAttributes: { type: SpanTypeAttribute.TASK }, - }, - INSTRUMENTATION_NAMES.CURSOR_SDK, - ), - ); - const startTime = getCurrentUnixTimestamp(); - safeLog(span, { - input: sanitizeUserMessage(message), - metadata, - }); - - const state: CursorRunState = { - activeToolSpans: new Map(), - agent, - conversationText: [], - deltaText: [], - finalized: false, - input: message, - metadata, - metrics: {}, - span, - startTime, - streamMessages: [], - streamText: [], - stepText: [], - taskText: [], - }; - - if (hasCursorCallbacks(sendOptions)) { - event.arguments[1] = wrapSendOptionsCallbacks(sendOptions, state); - } - states.set(event, state); - }; - const resolved = ( - event: ChannelMessage, - ) => { - const state = states.get(event); - if (!state) { - return; - } - - if (!event.result) { - return; - } - state.run = event.result; - state.metadata = { - ...state.metadata, - ...extractRunMetadata(event.result), - }; - patchCursorRun(event.result, state); - }; - const failed = ( - event: ChannelMessage, - ) => { - const state = states.get(event); - if (!state || !event.error) { - return; - } - safeLog(state.span, { error: event.error.message }); - endOpenToolSpans(state, event.error.message); - state.span.end(); - state.finalized = true; - states.delete(event); - }; - runInstrumentation(() => prepare(event)); - let result; - try { - result = Reflect.apply(target, receiver, args); - } catch (error) { - Object.assign(event, { error }); - runInstrumentation(() => failed(event)); - throw error; - } - return observeResult( - result, - (value) => { - Object.assign(event, { result: value }); - resolved(event); - }, - (error) => { - Object.assign(event, { error }); - failed(event); - }, - ); +function createRunState( + message: string | CursorSDKUserMessage, + metadata: Record, +): CursorRunState { + const span = startBaseSpan( + withSpanInstrumentationName( + { + name: "Cursor Agent", + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { input: sanitizeUserMessage(message), metadata }, }, - ); - this.unsubscribers.push(removeHandlers); - } + INSTRUMENTATION_NAMES.CURSOR_SDK, + ), + ); + return { + activeToolSpans: new Map(), + completedToolSpans: new Map(), + observation: Promise.resolve(), + conversationText: [], + deltaText: [], + finalized: false, + metadata, + metrics: {}, + usageComplete: true, + span, + startTime: getCurrentUnixTimestamp(), + streamText: [], + stepText: [], + taskText: [], + }; } -function patchCursorAgentInPlace(agent: unknown): void { - if (!agent || typeof agent !== "object") { - return; - } - const agentRecord = agent as CursorSDKAgent & Record; - if ( - agentRecord[PATCHED_AGENT] || - agentRecord[WRAPPED_AGENT] || - typeof agentRecord.send !== "function" - ) { - return; +function traceCursorPrompt( + call: () => PromiseLike, + args: [string | CursorSDKUserMessage, CursorSDKAgentOptions | undefined], + moduleVersion?: string, +): PromiseLike { + if (promptContext().getStore()) return call(); + const state = runInstrumentation(() => + createRunState(args[0], { + ...extractAgentOptionsMetadata(args[1]), + "cursor_sdk.operation": "Agent.prompt", + provider: "cursor", + ...(moduleVersion ? { "cursor_sdk.version": moduleVersion } : {}), + }), + ); + if (!state) return call(); + let result; + try { + result = promptContext().run(state, () => withCurrent(state.span, call)); + } catch (error) { + void finalizeCursorRun(state, { error }); + throw error; } + return observeResult( + result, + (result) => { + void finalizeCursorRun(state, { result }); + }, + (error) => { + void finalizeCursorRun(state, { error }); + }, + ); +} - const originalSend = agentRecord.send.bind(agentRecord); +function traceCursorSend( + call: ( + options: CursorSDKSendOptions | undefined, + ) => PromiseLike, + args: [string | CursorSDKUserMessage, CursorSDKSendOptions | undefined], + agent?: CursorSDKAgent, + moduleVersion?: string, +): PromiseLike { + const state = + promptContext().getStore() ?? + runInstrumentation(() => + createRunState(args[0], { + ...extractSendMetadata(args[1]), + ...(agent ? extractAgentMetadata(agent) : {}), + "cursor_sdk.operation": "agent.send", + provider: "cursor", + ...(moduleVersion ? { "cursor_sdk.version": moduleVersion } : {}), + }), + ); + if (!state) return call(args[1]); + const options = + runInstrumentation(() => wrapSendOptionsCallbacks(args[1] ?? {}, state)) ?? + args[1]; + let result; try { - Object.defineProperty(agentRecord, PATCHED_AGENT, { - configurable: false, - enumerable: false, - value: true, - }); - Object.defineProperty(agentRecord, "send", { - configurable: true, - value( - message: string | CursorSDKUserMessage, - options?: CursorSDKSendOptions, - ) { - const args = [message, options] as [ - string | CursorSDKUserMessage, - CursorSDKSendOptions | undefined, - ]; - return cursorSDKChannels.send.invoke( - originalSend, - undefined, - [...args], - { agent: agentRecord, operation: "send" }, - ); - }, - writable: true, - }); - } catch { - // Frozen/sealed agents cannot be patched. Leave user behavior untouched. + result = withCurrent(state.span, () => call(options)); + } catch (error) { + void finalizeCursorRun(state, { error }); + throw error; } + return observeResult( + result, + (run) => { + state.run = run; + patchCursorRun(run, state); + }, + (error) => { + void finalizeCursorRun(state, { error }); + }, + ); +} + +function observeUpdate( + state: CursorRunState, + observe: () => void | Promise, +): void { + state.observation = state.observation + .then(async () => { + if (!state.finalized) await observe(); + }) + .catch((error) => logInstrumentationError("Cursor SDK update", error)); } function wrapSendOptionsCallbacks( @@ -378,37 +237,34 @@ function wrapSendOptionsCallbacks( return { ...options, - async onDelta(args) { - try { - await handleInteractionUpdate(state, args.update); - } catch (error) { - logInstrumentationError("Cursor SDK onDelta", error); - } - if (originalOnDelta) { - return originalOnDelta(args); - } + onDelta: function (args) { + observeUpdate(state, () => handleInteractionUpdate(state, args.update)); + return ( + originalOnDelta && + withCurrent(state.span, () => + Reflect.apply(originalOnDelta, this, [args]), + ) + ); }, - async onStep(args) { - try { - handleStepUpdate(state, args.step); - } catch (error) { - logInstrumentationError("Cursor SDK onStep", error); - } - if (originalOnStep) { - return originalOnStep(args); - } + onStep: function (args) { + observeUpdate(state, () => handleStepUpdate(state, args.step)); + return ( + originalOnStep && + withCurrent(state.span, () => + Reflect.apply(originalOnStep, this, [args]), + ) + ); }, }; } -function hasCursorCallbacks( - options: CursorSDKSendOptions | undefined, -): options is CursorSDKSendOptions { - return ( - !!options && - (typeof options.onDelta === "function" || - typeof options.onStep === "function") - ); +function callCursorRun(state: CursorRunState, call: () => T): T { + try { + return withCurrent(state.span, call); + } catch (error) { + void finalizeCursorRun(state, { error }); + throw error; + } } function patchCursorRun(run: CursorSDKRun, state: CursorRunState): void { @@ -428,77 +284,78 @@ function patchCursorRun(run: CursorSDKRun, state: CursorRunState): void { }); if (typeof runRecord.stream === "function") { - const originalStream = runRecord.stream.bind(runRecord); - Object.defineProperty(runRecord, "stream", { - configurable: true, - value() { - const stream = originalStream(); - return patchCursorStream(stream, state); - }, - writable: true, - }); + const originalStream = runRecord.stream; + runRecord.stream = function (...args) { + const stream = callCursorRun(state, () => + Reflect.apply(originalStream, runRecord, args), + ); + patchStreamIfNeeded(stream, { + aroundNext: (next) => withCurrent(state.span, next), + shouldCollect: (message) => { + observeUpdate(state, () => handleStreamMessage(state, message)); + return false; + }, + onComplete: () => finalizeCursorRun(state), + onCancel: () => finalizeCursorRun(state, { cancelled: true }), + onError: (error) => finalizeCursorRun(state, { error }), + }); + return stream; + }; } - if (typeof runRecord.wait === "function") { - const originalWait = runRecord.wait.bind(runRecord); - Object.defineProperty(runRecord, "wait", { - configurable: true, - async value() { - try { - const result = await originalWait(); - state.lastResult = result; - await finalizeCursorRun(state, { result }); - return result; - } catch (error) { - await finalizeCursorRun(state, { error }); - throw error; - } - }, - writable: true, - }); + const originalWait = runRecord.wait; + runRecord.wait = function (...args) { + const result = callCursorRun(state, () => + Reflect.apply(originalWait, runRecord, args), + ); + return observeResult( + result, + (value) => { + void finalizeCursorRun(state, { result: value }); + }, + (error) => { + void finalizeCursorRun(state, { error }); + }, + ); + }; } - if (typeof runRecord.conversation === "function") { - const originalConversation = runRecord.conversation.bind(runRecord); - Object.defineProperty(runRecord, "conversation", { - configurable: true, - async value() { - try { - const conversation = await originalConversation(); - await handleConversation(state, conversation); - await finalizeCursorRun(state); - return conversation; - } catch (error) { - await finalizeCursorRun(state, { error }); - throw error; - } - }, - writable: true, - }); + const originalConversation = runRecord.conversation; + runRecord.conversation = function (...args) { + return observeResult( + callCursorRun(state, () => + Reflect.apply(originalConversation, runRecord, args), + ), + (turns) => { + observeUpdate(state, () => handleConversation(state, turns)); + if (run.status && run.status !== "running") + void finalizeCursorRun(state); + }, + (error) => { + void finalizeCursorRun(state, { error }); + }, + ); + }; } - } catch { - // If the Run object is not patchable, finish the span with available data. - void finalizeCursorRun(state, { output: run }); - } -} - -async function* patchCursorStream( - stream: AsyncGenerator, - state: CursorRunState, -): AsyncGenerator { - try { - for await (const message of stream) { - try { - await handleStreamMessage(state, message); - } catch (error) { - logInstrumentationError("Cursor SDK stream", error); - } - yield message; + if (typeof runRecord.cancel === "function") { + const originalCancel = runRecord.cancel; + runRecord.cancel = function (...args) { + return observeResult( + callCursorRun(state, () => + Reflect.apply(originalCancel, runRecord, args), + ), + () => { + void finalizeCursorRun(state, { cancelled: true }); + }, + (error) => { + void finalizeCursorRun(state, { error }); + }, + ); + }; } - await finalizeCursorRun(state); } catch (error) { - await finalizeCursorRun(state, { error }); - throw error; + logInstrumentationError("Cursor SDK run patch", error); + void finalizeCursorRun(state); } } @@ -508,15 +365,14 @@ async function handleInteractionUpdate( ): Promise { switch (update.type) { case "text-delta": - if (typeof update.text === "string") { + if (typeof update.text === "string" && update.text.length > 0) { + recordFirstToken(state); state.deltaText.push(update.text); } return; case "token-delta": - if (typeof update.tokens === "number") { - state.metrics["cursor_sdk.delta_tokens"] = - (state.metrics["cursor_sdk.delta_tokens"] ?? 0) + update.tokens; - } + // Token deltas do not distinguish input/output or cache categories. + // Only turn-ended usage can establish canonical token totals. return; case "tool-call-started": case "partial-tool-call": @@ -535,25 +391,13 @@ async function handleInteractionUpdate( ); return; case "turn-ended": - addUsageMetrics( - state.metrics, - (update as { usage?: CursorSDKUsage }).usage, - ); + addUsageMetrics(state, (update as { usage?: CursorSDKUsage }).usage); return; case "summary": if (typeof update.summary === "string") { state.taskText.push(update.summary); } return; - case "step-completed": - if (typeof update.stepDurationMs === "number") { - state.metrics["cursor_sdk.step_duration_ms"] = - (state.metrics["cursor_sdk.step_duration_ms"] ?? 0) + - update.stepDurationMs; - } - state.metrics["cursor_sdk.steps"] = - (state.metrics["cursor_sdk.steps"] ?? 0) + 1; - return; default: return; } @@ -566,60 +410,113 @@ async function handleToolUpdate( { type: "tool-call-started" | "partial-tool-call" | "tool-call-completed" } >, ): Promise { - const callId = update.callId; - if (!callId) { + if (!update.callId) return; + await recordTool( + state, + update.callId, + update.toolCall ?? {}, + update.type === "tool-call-completed", + update.status, + ); +} + +function toolMetadata( + toolCall: Pick, + status?: string, +): Record { + return { + ...(status ? { "cursor_sdk.tool.status": status } : {}), + ...(typeof toolCall.truncated?.args === "boolean" + ? { "cursor_sdk.tool.args_truncated": toolCall.truncated.args } + : {}), + ...(typeof toolCall.truncated?.result === "boolean" + ? { "cursor_sdk.tool.result_truncated": toolCall.truncated.result } + : {}), + }; +} + +async function recordTool( + state: CursorRunState, + callId: string, + toolCall: CursorSDKToolCall, + completed: boolean, + reportedStatus?: string, + parent?: Span, + scope = "", + ancestors = new Set(), +): Promise { + const key = `${scope}/${callId}`; + const finished = state.completedToolSpans.get(key); + if (ancestors.has(toolCall)) return; + if (finished && !completed) { + finished.span.log({ + ...(toolCall.args !== undefined ? { input: toolCall.args } : {}), + metadata: toolMetadata(toolCall), + }); return; } - - const toolCall = update.toolCall; - const name = extractToolName(toolCall); - const args = extractToolArgs(toolCall); const result = extractToolResult(toolCall); - - if ( - update.type === "tool-call-started" || - update.type === "partial-tool-call" - ) { - if (!state.activeToolSpans.has(callId)) { - state.activeToolSpans.set( - callId, - await startToolSpan(state, { - args, - callId, - name, - status: "running", - toolCall, - }), + const resultRecord = asRecord(result); + if (resultRecord && ancestors.has(resultRecord)) return; + const failed = + reportedStatus === "error" || + toolCall.status === "error" || + resultRecord?.status === "error"; + const status = failed + ? "error" + : (reportedStatus ?? + toolCall.status ?? + (completed ? "completed" : "running")); + let toolState = finished ?? state.activeToolSpans.get(key); + if (!toolState) { + toolState = await startToolSpan(state, { + args: toolCall.args, + callId, + name: extractToolName(toolCall), + status, + parent, + truncated: toolCall.truncated, + }); + state.activeToolSpans.set(key, toolState); + } else { + toolState.span.log({ + ...(toolCall.args !== undefined ? { input: toolCall.args } : {}), + metadata: toolMetadata(toolCall, status), + }); + } + if (!completed) return; + // Task results contain normalized SDK steps or protobuf JSON steps. + const nested = asRecord(resultRecord?.value)?.conversationSteps; + if (toolState.subAgentSpan && Array.isArray(nested)) { + const seen = new Set(ancestors).add(toolCall); + if (resultRecord) seen.add(resultRecord); + for (const [index, step] of nested.entries()) { + await handleConversationStep( + state, + step, + toolState.subAgentSpan, + key, + index, + seen, ); } - return; } - - const toolState = - state.activeToolSpans.get(callId) ?? - (await startToolSpan(state, { - args, - callId, - name, - status: "completed", - toolCall, - })); - + const error = failed + ? (resultRecord?.error ?? result ?? new Error("Cursor tool failed")) + : undefined; finishToolSpan(toolState, { - error: toolCall?.status === "error" ? toLoggedError(result) : undefined, - metadata: { - "cursor_sdk.tool.status": toolCall?.status ?? "completed", - }, + error, + metadata: toolMetadata(toolCall, status), output: result, }); - state.activeToolSpans.delete(callId); + state.activeToolSpans.delete(key); + state.completedToolSpans.set(key, toolState); } async function handleStreamMessage( state: CursorRunState, message: CursorSDKMessage, ): Promise { - state.streamMessages.push(message); if (message.type === "system") { const systemMessage = message as Extract< CursorSDKMessage, @@ -647,17 +544,19 @@ async function handleStreamMessage( { type: "assistant" } >; for (const block of assistantMessage.message?.content ?? []) { - if (block?.type === "text" && typeof block.text === "string") { + if ( + block?.type === "text" && + typeof block.text === "string" && + block.text.length > 0 + ) { + recordFirstToken(state); state.streamText.push(block.text); } else if (block?.type === "tool_use" && block.id) { - state.activeToolSpans.set( + await recordTool( + state, block.id, - await startToolSpan(state, { - args: block.input, - callId: block.id, - name: block.name, - status: "running", - }), + { name: block.name, args: block.input }, + false, ); } } @@ -677,8 +576,17 @@ async function handleStreamMessage( return; } - if (message.type === "status" && message.status) { - state.metadata["cursor_sdk.status"] = message.status; + if (message.type === "request" && typeof message.request_id === "string") { + state.metadata.request_id = message.request_id; + } + if (message.type === "status" && typeof message.status === "string") { + state.metadata["cursor_sdk.status"] = message.status.toLowerCase(); + if (message.status.toLowerCase() === "error") { + state.error = + typeof message.message === "string" + ? message.message + : new Error("Cursor run failed"); + } } } @@ -686,61 +594,41 @@ async function handleToolMessage( state: CursorRunState, message: CursorSDKToolUseMessage, ): Promise { - const callId = message.call_id; - if (!callId) { - return; - } - - if (message.status === "running") { - if (!state.activeToolSpans.has(callId)) { - state.activeToolSpans.set( - callId, - await startToolSpan(state, { - args: message.args, - callId, - name: message.name, - status: message.status, - truncated: message.truncated, - }), - ); - } - return; - } - - const toolState = - state.activeToolSpans.get(callId) ?? - (await startToolSpan(state, { - args: message.args, - callId, + if (!message.call_id) return; + await recordTool( + state, + message.call_id, + { name: message.name, + args: message.args, + result: message.result, status: message.status, truncated: message.truncated, - })); - finishToolSpan(toolState, { - error: - message.status === "error" ? toLoggedError(message.result) : undefined, - metadata: { - "cursor_sdk.tool.status": message.status, }, - output: message.result, - }); - state.activeToolSpans.delete(callId); + message.status !== "running", + message.status, + ); } async function handleConversation( state: CursorRunState, turns: CursorSDKConversationTurn[], ): Promise { - state.conversationOutput = turns; - for (const turn of turns) { + for (const [turnIndex, turn] of turns.entries()) { if (turn.type === "agentConversationTurn") { - for (const step of turn.turn?.steps ?? []) { - await handleConversationStep(state, step); + for (const [index, step] of (turn.turn?.steps ?? []).entries()) { + await handleConversationStep( + state, + step, + undefined, + `turn:${turnIndex}`, + index, + ); } } else if (turn.type === "shellConversationTurn") { const command = turn.turn?.shellCommand?.command; if (command) { - const callId = `shell:${state.activeToolSpans.size}:${command}`; + const callId = `shell:${turnIndex}`; const toolState = await startToolSpan(state, { args: turn.turn?.shellCommand, callId, @@ -756,50 +644,93 @@ async function handleConversation( } } +// Protobuf JSON preserves oneof field names instead of the SDK discriminant. +// Only the known tool variants are accepted; arbitrary properties are never used. +const protoToolNames = { + shellToolCall: "shell", + readToolCall: "read", + writeToolCall: "write", + editToolCall: "edit", + deleteToolCall: "delete", + globToolCall: "glob", + grepToolCall: "grep", + lsToolCall: "ls", + mcpToolCall: "mcp", + readLintsToolCall: "readLints", + semSearchToolCall: "semSearch", + createPlanToolCall: "createPlan", + updateTodosToolCall: "updateTodos", + taskToolCall: "task", +}; + +function asRecord(value: unknown): Record | undefined { + return isObject(value) ? value : undefined; +} + +function normalizeToolCall(value: unknown): CursorSDKToolCall | undefined { + const record = asRecord(value); + if (!record) return; + if (typeof record.type === "string" || typeof record.name === "string") + return record as CursorSDKToolCall; + for (const [field, name] of Object.entries(protoToolNames)) { + const tool = asRecord(record[field]); + if (!tool) continue; + const result = asRecord(tool.result); + return { + type: name, + args: tool.args, + result: + result && "success" in result + ? { status: "success", value: result.success } + : result && "error" in result + ? { status: "error", error: result.error } + : result, + }; + } +} + async function handleConversationStep( state: CursorRunState, - step: CursorSDKConversationStep, + value: unknown, + parent?: Span, + scope = "", + index = 0, + ancestors = new Set(), ): Promise { - if ( - step.type === "assistantMessage" && - typeof step.message?.text === "string" - ) { - state.conversationText.push(step.message.text); - return; - } - - if (step.type !== "toolCall") { + const step = asRecord(value); + if (!step) return; + const assistant = asRecord( + step.type === "assistantMessage" ? step.message : step.assistantMessage, + ); + if (typeof assistant?.text === "string") { + if (parent) parent.log({ output: assistant.text }); + else state.conversationText.push(assistant.text); return; } - - const toolCall = step.message; + const toolCall = normalizeToolCall( + step.type === "toolCall" ? step.message : step.toolCall, + ); + if (!toolCall) return; const callId = - typeof toolCall?.callId === "string" + typeof toolCall.callId === "string" ? toolCall.callId - : `conversation-tool:${state.activeToolSpans.size}`; - const toolState = await startToolSpan(state, { - args: extractToolArgs(toolCall), + : `conversation-tool:${index}`; + await recordTool( + state, callId, - name: extractToolName(toolCall), - status: toolCall?.status, toolCall, - }); - finishToolSpan(toolState, { - error: - toolCall?.status === "error" ? toLoggedError(toolCall.result) : undefined, - metadata: { - "cursor_sdk.tool.status": toolCall?.status ?? "completed", - }, - output: extractToolResult(toolCall), - }); + true, + undefined, + parent, + !parent && typeof toolCall.callId === "string" ? "" : scope, + ancestors, + ); } function handleStepUpdate( state: CursorRunState, step: CursorSDKConversationStep, ): void { - state.metrics["cursor_sdk.steps"] = - (state.metrics["cursor_sdk.steps"] ?? 0) + 1; if (step.type) { const stepTypes = state.metadata["cursor_sdk.step_types"]; if (Array.isArray(stepTypes)) { @@ -825,22 +756,16 @@ async function startToolSpan( callId: string; name?: string; status?: string; - toolCall?: CursorSDKToolCall; truncated?: { args?: boolean; result?: boolean }; + parent?: Span; }, ): Promise { const name = args.name || "unknown"; const metadata: Record = { - "cursor_sdk.tool.status": args.status, + ...toolMetadata(args, args.status), "gen_ai.tool.call.id": args.callId, "gen_ai.tool.name": name, }; - if (args.truncated?.args !== undefined) { - metadata["cursor_sdk.tool.args_truncated"] = args.truncated.args; - } - if (args.truncated?.result !== undefined) { - metadata["cursor_sdk.tool.result_truncated"] = args.truncated.result; - } const span = startBaseSpan( withSpanInstrumentationName( @@ -849,8 +774,8 @@ async function startToolSpan( input: args.args, metadata, }, - name: `tool: ${name}`, - parent: await state.span.export(), + name, + parent: await (args.parent ?? state.span).export(), spanAttributes: { type: SpanTypeAttribute.TOOL }, }, INSTRUMENTATION_NAMES.CURSOR_SDK, @@ -869,7 +794,7 @@ async function startToolSpan( "gen_ai.tool.name": name, }, }, - name: formatSubAgentSpanName(args.toolCall, args.args), + name: "Cursor Subagent", parent: await span.export(), spanAttributes: { type: SpanTypeAttribute.TASK }, }, @@ -884,108 +809,83 @@ async function startToolSpan( function finishToolSpan( toolState: ToolState, result: { - error?: string; + error?: unknown; metadata?: Record; output?: unknown; }, ): void { - try { - if (result.error) { - safeLog(toolState.span, { - error: result.error, - metadata: result.metadata, - output: result.output, - }); - if (toolState.subAgentSpan) { - safeLog(toolState.subAgentSpan, { - error: result.error, - metadata: result.metadata, - output: result.output, - }); - } - } else { - safeLog(toolState.span, { - metadata: result.metadata, - output: result.output, - }); - if (toolState.subAgentSpan) { - safeLog(toolState.subAgentSpan, { - metadata: result.metadata, - output: result.output, - }); - } - } - } finally { + toolState.span.log(result); + toolState.subAgentSpan?.log(result); + if (!toolState.ended) { + toolState.ended = true; toolState.subAgentSpan?.end(); toolState.span.end(); } } -async function finalizeCursorRun( +function finalizeCursorRun( state: CursorRunState, params: { error?: unknown; - output?: unknown; result?: CursorSDKRunResult; + cancelled?: boolean; } = {}, ): Promise { - if (state.finalized) { - return; - } - state.finalized = true; - - const error = params.error; - const result = params.result ?? state.lastResult; - const output = - params.output ?? - result?.result ?? - state.run?.result ?? - (state.conversationText.length > 0 - ? state.conversationText.join("\n") - : undefined) ?? - state.conversationOutput ?? - (state.streamText.length > 0 ? state.streamText.join("") : undefined) ?? - (state.deltaText.length > 0 ? state.deltaText.join("") : undefined) ?? - (state.stepText.length > 0 ? state.stepText.join("\n") : undefined) ?? - (state.taskText.length > 0 ? state.taskText.join("\n") : undefined); - - try { - const metadata = { - ...state.metadata, - ...(state.run ? extractRunMetadata(state.run) : {}), - ...(result ? extractRunResultMetadata(result) : {}), - }; - if (error) { - safeLog(state.span, { - error: toLoggedError(error), + if (state.finalizing) return state.finalizing; + // Drain queued observations before ending, including un-awaited SDK callbacks. + state.finalizing = state.observation.then(() => { + state.finalized = true; + const event = runInstrumentation(() => { + const result = params.result; + const metadata = { + ...state.metadata, + ...extractRunMetadata(state.run), + ...extractRunResultMetadata(result), + ...(params.cancelled ? { "cursor_sdk.status": "cancelled" } : {}), + }; + const error = + params.error ?? + state.error ?? + (metadata["cursor_sdk.status"] === "error" + ? new Error("Cursor run failed") + : undefined); + const output = + result?.result ?? + state.run?.result ?? + (state.conversationText.length + ? state.conversationText.join("\n") + : undefined) ?? + (state.streamText.length ? state.streamText.join("") : undefined) ?? + (state.deltaText.length ? state.deltaText.join("") : undefined) ?? + (state.stepText.length ? state.stepText.join("\n") : undefined) ?? + (state.taskText.length ? state.taskText.join("\n") : undefined); + return { + error, metadata, - metrics: { - ...cleanMetrics(state.metrics), - ...buildDurationMetrics(state.startTime), - }, output, - }); - } else { - safeLog(state.span, { - metadata, metrics: { - ...cleanMetrics(state.metrics), - ...buildDurationMetrics(state.startTime), + ...state.metrics, + start: state.startTime, + end: getCurrentUnixTimestamp(), }, - output, + }; + }); + if (event) state.span.log(event); + for (const tool of state.activeToolSpans.values()) { + finishToolSpan(tool, { + error: event?.error ?? params.error ?? state.error, + metadata: + params.cancelled || + event?.metadata["cursor_sdk.status"] === "cancelled" + ? { "cursor_sdk.tool.status": "cancelled" } + : undefined, }); } - } finally { - endOpenToolSpans(state); + state.activeToolSpans.clear(); + state.completedToolSpans.clear(); state.span.end(); - } -} - -function endOpenToolSpans(state: CursorRunState, error?: string): void { - for (const [, toolState] of state.activeToolSpans) { - finishToolSpan(toolState, { error }); - } - state.activeToolSpans.clear(); + }); + return state.finalizing; } function sanitizeUserMessage( @@ -995,7 +895,7 @@ function sanitizeUserMessage( return message; } return { - ...message, + text: message.text, images: message.images?.map((image) => { const imageRecord = image as Record; return { @@ -1005,7 +905,14 @@ function sanitizeUserMessage( ...(typeof imageRecord.mimeType === "string" ? { mimeType: imageRecord.mimeType } : {}), - ...(image.dimension ? { dimension: image.dimension } : {}), + ...(image.dimension + ? { + dimension: { + width: image.dimension.width, + height: image.dimension.height, + }, + } + : {}), hasData: typeof imageRecord.data === "string", }; }), @@ -1126,46 +1033,54 @@ function extractModelMetadata( return { model: model.id, "cursor_sdk.model": model.id, - ...(model.params ? { "cursor_sdk.model.params": model.params } : {}), + ...(model.params + ? { + "cursor_sdk.model.params": model.params.map((param) => ({ + id: param.id, + value: param.value, + })), + } + : {}), }; } function addUsageMetrics( - metrics: Record, + state: CursorRunState, usage: CursorSDKUsage | undefined, ): void { - if (!usage) { - return; - } - if (usage.inputTokens !== undefined) { - metrics.prompt_tokens = (metrics.prompt_tokens ?? 0) + usage.inputTokens; - } - if (usage.outputTokens !== undefined) { - metrics.completion_tokens = - (metrics.completion_tokens ?? 0) + usage.outputTokens; - } - if (usage.cacheReadTokens !== undefined) { - metrics.prompt_cached_tokens = - (metrics.prompt_cached_tokens ?? 0) + usage.cacheReadTokens; + if (!usage) return; + const metrics = state.metrics; + const count = (value: unknown): value is number => + typeof value === "number" && Number.isSafeInteger(value) && value >= 0; + const add = (name: string, value: unknown) => { + if (count(value)) metrics[name] = (metrics[name] ?? 0) + value; + }; + // Cursor reports uncached input and cache reads/writes as separate buckets. + // Cache detail metrics are subsets of the inclusive prompt total. + if (count(usage.inputTokens)) { + add( + "prompt_tokens", + usage.inputTokens + + (count(usage.cacheReadTokens) ? usage.cacheReadTokens : 0) + + (count(usage.cacheWriteTokens) ? usage.cacheWriteTokens : 0), + ); } - if (usage.cacheWriteTokens !== undefined) { - metrics.prompt_cache_creation_tokens = - (metrics.prompt_cache_creation_tokens ?? 0) + usage.cacheWriteTokens; + add("completion_tokens", usage.outputTokens); + add("prompt_cached_tokens", usage.cacheReadTokens); + add("prompt_cache_creation_tokens", usage.cacheWriteTokens); + state.usageComplete &&= count(usage.inputTokens) && count(usage.outputTokens); + if (state.usageComplete) { + metrics.tokens = + (metrics.prompt_tokens ?? 0) + (metrics.completion_tokens ?? 0); + } else { + // An incomplete turn makes the run's computed total incomplete too. + delete metrics.tokens; } - metrics.tokens = - (metrics.prompt_tokens ?? 0) + - (metrics.completion_tokens ?? 0) + - (metrics.prompt_cached_tokens ?? 0) + - (metrics.prompt_cache_creation_tokens ?? 0); } -function buildDurationMetrics(startTime: number): Record { - const end = getCurrentUnixTimestamp(); - return { - duration: end - startTime, - end, - start: startTime, - }; +function recordFirstToken(state: CursorRunState): void { + state.metrics.time_to_first_token ??= + getCurrentUnixTimestamp() - state.startTime; } function extractToolName(toolCall: CursorSDKToolCall | undefined): string { @@ -1181,10 +1096,6 @@ function extractToolName(toolCall: CursorSDKToolCall | undefined): string { return "unknown"; } -function extractToolArgs(toolCall: CursorSDKToolCall | undefined): unknown { - return toolCall && "args" in toolCall ? toolCall.args : undefined; -} - function extractToolResult(toolCall: CursorSDKToolCall | undefined): unknown { return toolCall && "result" in toolCall ? toolCall.result : undefined; } @@ -1193,45 +1104,6 @@ function isSubAgentToolName(name: string): boolean { return name === "Agent" || name === "Task" || name === "task"; } -function formatSubAgentSpanName( - toolCall: CursorSDKToolCall | undefined, - args: unknown, -): string { - const details = (toolCall ?? args) as Record | undefined; - const description = - getString(details, "description") ?? - getString(details, "subagent_type") ?? - getString(details, "type") ?? - getString(details, "name"); - return description ? `Agent: ${description}` : "Agent: sub-agent"; -} - -function getString( - obj: Record | undefined, - key: string, -): string | undefined { - const value = obj?.[key]; - return typeof value === "string" ? value : undefined; -} - -function safeLog(span: Span, event: Parameters[0]): void { - try { - span.log(event); - } catch (error) { - logInstrumentationError("Cursor SDK span log", error); - } -} - function logInstrumentationError(context: string, error: unknown): void { debugLogger.error(`Error processing ${context}:`, error); } - -function cleanMetrics(metrics: Record): Record { - const cleaned: Record = {}; - for (const [key, value] of Object.entries(metrics)) { - if (value !== undefined && Number.isFinite(value)) { - cleaned[key] = value; - } - } - return cleaned; -} diff --git a/js/src/instrumentation/plugins/cursor-sdk-wrapping.test.ts b/js/src/instrumentation/plugins/cursor-sdk-wrapping.test.ts new file mode 100644 index 000000000..e760784d6 --- /dev/null +++ b/js/src/instrumentation/plugins/cursor-sdk-wrapping.test.ts @@ -0,0 +1,137 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import { afterEach, describe, expect, it, vi } from "vitest"; +const { spans } = vi.hoisted(() => ({ spans: [] as any[] })); +vi.mock("../../isomorph", () => ({ + default: { newAsyncLocalStorage: () => new AsyncLocalStorage() }, +})); +vi.mock("../../logger", () => ({ + startSpan: (args: any) => { + const span = { + ...args, + end: vi.fn(), + log: vi.fn(), + export: async () => args.name, + }; + spans.push(span); + return span; + }, + withCurrent: (_span: unknown, call: () => unknown) => call(), +})); +import { wrapCursorSDK } from "../../wrappers/cursor-sdk"; +import { CursorSDKPlugin } from "./cursor-sdk-plugin"; +import { newGlobalInvocationHook } from "../../global-instrumentation-hooks"; + +describe("Cursor automatic and manual wrapping", () => { + const plugin = new CursorSDKPlugin(); + afterEach(() => { + plugin.disable(); + spans.length = 0; + }); + + it.each(["manual", "automatic", "combined"] as const)( + "captures one task per prompt/send with %s instrumentation", + async (mode) => { + plugin.enable(); + const events = [ + { + type: "tool_call", + call_id: "call", + name: "shell", + status: "completed", + args: { command: "pwd" }, + result: "ok", + }, + ]; + const agent = { + send: vi.fn(async () => ({ + status: "finished", + stream: async function* () { + for (const event of events) yield event; + }, + async wait() { + return { status: "finished", result: "ok" }; + }, + })), + }; + const createPromise = Promise.resolve(agent); + class Agent { + static #value = "ok"; + static create(_options: unknown) { + return createPromise; + } + static resume() { + return createPromise; + } + static async prompt(_message: string) { + // Cursor's distributed prompt calls its lexical Agent class, not this. + const instance = await Agent.create({}); + const run = await instance.send(); + await run.wait(); + return { result: this.#value }; + } + } + if (mode !== "manual") { + for (const operation of ["create", "resume", "prompt"] as const) { + const original = Agent[operation]; + const hook = newGlobalInvocationHook( + `orchestrion:@cursor/sdk:Agent.${operation}`, + ); + Object.defineProperty(Agent, operation, { + configurable: true, + writable: true, + value(...args: unknown[]) { + return hook.invoke( + original as (...args: unknown[]) => unknown, + Agent, + args, + { moduleVersion: "1.0.12" }, + ); + }, + }); + } + } + const wrapped = + mode === "automatic" ? { Agent } : wrapCursorSDK({ Agent }); + expect(await wrapped.Agent.prompt("hi")).toEqual({ result: "ok" }); + const instance = await wrapped.Agent.create({}); + const run = await instance.send(); + for await (const _ of run.stream()) { + /* consume */ + } + await run.wait(); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(spans.map((span) => span.name)).toEqual([ + "Cursor Agent", + "Cursor Agent", + "shell", + ]); + for (const span of spans) expect(span.end).toHaveBeenCalledOnce(); + }, + ); + + it("keeps installed wrappers usable after disabling instrumentation", async () => { + plugin.enable(); + const error = new Error("provider failed"); + class Agent { + static create() { + return Promise.resolve({ + send() { + throw error; + }, + }); + } + static resume() { + return this.create(); + } + static prompt() { + throw error; + } + } + const wrapped = wrapCursorSDK({ Agent }); + const agent = await wrapped.Agent.create(); + plugin.disable(); + expect(() => agent.send()).toThrow(error); + expect(() => wrapped.Agent.prompt()).toThrow(error); + expect(spans).toHaveLength(0); + }); +}); diff --git a/js/src/vendor-sdk-types/cursor-sdk.ts b/js/src/vendor-sdk-types/cursor-sdk.ts index b11c0f630..b5d43f833 100644 --- a/js/src/vendor-sdk-types/cursor-sdk.ts +++ b/js/src/vendor-sdk-types/cursor-sdk.ts @@ -28,6 +28,8 @@ export interface CursorSDKAgent { close?: () => void; reload?: () => Promise; [Symbol.asyncDispose]?: () => Promise; + listArtifacts?: () => Promise; + downloadArtifact?: (path: string) => Promise; [key: string | symbol]: unknown; } @@ -287,6 +289,7 @@ export interface CursorSDKUsage { } export interface CursorSDKToolCall { + callId?: string; type?: string; name?: string; args?: unknown; diff --git a/js/src/wrappers/cursor-sdk.test.ts b/js/src/wrappers/cursor-sdk.test.ts index b260e4a1d..3486c713e 100644 --- a/js/src/wrappers/cursor-sdk.test.ts +++ b/js/src/wrappers/cursor-sdk.test.ts @@ -37,133 +37,126 @@ describe("wrapCursorSDK", () => { warnSpy.mockRestore(); }); - it("wraps Agent.create and returned agent.send", async () => { - const run = makeRun(); - const agent = { - agentId: "agent-1", - send: vi.fn(async () => run), - }; - const sdk = { - Agent: class { - static async create() { - return agent; - } - }, - }; - - const wrapped = wrapCursorSDK(sdk as any) as any; - const created = await wrapped.Agent.create({ local: { cwd: "/tmp/repo" } }); - const result = await created.send("hello", { - onDelta: vi.fn(), - onStep: vi.fn(), - }); - - expect(result).toBe(run); - expect(agent.send).toHaveBeenCalledWith("hello", expect.any(Object)); - expect(invoke).toHaveBeenCalledTimes(2); - }); - - it("does not wrap send twice when the plugin already patched the returned agent", async () => { - const run = makeRun(); - const agent = { - [Symbol.for("braintrust.cursor-sdk.auto-patched-agent")]: true, - send: vi.fn(async () => run), - }; - const sdk = { - Agent: class { - static async create() { - return agent; - } - }, - }; - const wrapped = wrapCursorSDK(sdk as any) as any; - const created = await wrapped.Agent.create({}); - await expect(created.send("hello")).resolves.toBe(run); - expect(agent.send).toHaveBeenCalledExactlyOnceWith("hello"); - expect(invoke).toHaveBeenCalledOnce(); - }); - - it("wraps Agent.resume and preserves private-field-safe method binding", async () => { + it("wraps create/resume agents once and preserves private-field methods", async () => { + const run = { id: "run-1" }; class PrivateAgent { - #run = makeRun(); - agentId = "agent-2"; - + #run = run; async send() { return this.#run; } - + reload() { + return Promise.resolve(this.#run); + } + close() { + return this.#run.id; + } + listArtifacts() { + return Promise.resolve([this.#run.id]); + } + downloadArtifact(path: string) { + return Promise.resolve(`${this.#run.id}:${path}`); + } async [Symbol.asyncDispose]() { - return undefined; + return this.#run.id; } } - + const agent = new PrivateAgent(); + const originalSend = vi.spyOn(agent, "send"); const sdk = { Agent: class { + static async create() { + return agent; + } static async resume() { - return new PrivateAgent(); + return agent; } }, }; - const wrapped = wrapCursorSDK(sdk as any) as any; - const agent = await wrapped.Agent.resume("agent-2"); - - await expect(agent.send("hello")).resolves.toMatchObject({ id: "run-1" }); - await expect(agent[Symbol.asyncDispose]()).resolves.toBeUndefined(); + const created = await wrapped.Agent.create({}); + const patchedSend = created.send; + const resumed = await wrapped.Agent.resume("agent-1"); + expect(resumed).toBe(created); + expect(resumed.send).toBe(patchedSend); + const { + send, + reload, + close, + listArtifacts, + downloadArtifact, + [Symbol.asyncDispose]: dispose, + } = resumed; + await expect(send("hello")).resolves.toBe(run); + expect(originalSend).toHaveBeenCalledExactlyOnceWith("hello"); + await expect(reload()).resolves.toBe(run); + expect(close()).toBe("run-1"); + await expect(listArtifacts()).resolves.toEqual(["run-1"]); + await expect(downloadArtifact("output.txt")).resolves.toBe( + "run-1:output.txt", + ); + await expect(dispose()).resolves.toBe("run-1"); + expect(invoke).toHaveBeenCalledTimes(3); }); - it("wraps Agent.prompt", async () => { - const sdk = { - Agent: class { - static async prompt(message: string) { - return { id: "run-1", result: message, status: "finished" }; - } - }, - }; - - const wrapped = wrapCursorSDK(sdk as any) as any; - await expect(wrapped.Agent.prompt("hello")).resolves.toMatchObject({ - result: "hello", + it("leaves original factories and unwrapped subclass dispatch unchanged", async () => { + const create = vi.fn(function (this: unknown) { + return Promise.resolve({ owner: this, send: vi.fn() }); }); - expect(invoke).toHaveBeenCalledTimes(1); + class Agent { + static create = create; + static resume = create; + } + const before = Object.getOwnPropertyDescriptors(Agent); + const wrapped = wrapCursorSDK({ Agent }); + const wrappedAgent = wrapped.Agent; + expect(Object.getOwnPropertyDescriptors(Agent)).toEqual(before); + class Derived extends Agent {} + for (const operation of ["create", "resume"] as const) { + const direct = await Agent[operation](); + expect(direct.owner).toBe(Agent); + const derived = await Derived[operation](); + expect(derived.owner).toBe(Derived); + } + expect(invoke).not.toHaveBeenCalled(); + const promise = wrappedAgent.create(); + expect(await promise).toMatchObject({ owner: Agent }); + expect(invoke).toHaveBeenCalledOnce(); + expect(Object.getOwnPropertyDescriptors(Agent)).toEqual(before); }); - it("handles module namespace-like objects", async () => { - const Agent = class { - static async prompt() { - return { status: "finished" }; + it("handles module namespaces and preserves static promises, arguments, and private getters", async () => { + const agent = { send: vi.fn() }; + const promise = Object.assign(Promise.resolve(agent), { + helper: () => "ok", + }); + class Agent { + static #value = "ok"; + static get value() { + return this.#value; } - }; - const sdk = Object.defineProperty({}, "Agent", { + static create(...args: unknown[]) { + expect(args).toEqual([{}]); + return promise; + } + static resume(...args: unknown[]) { + expect(args).toEqual(["agent-1"]); + return promise; + } + static prompt(...args: unknown[]) { + expect(args).toEqual(["hi"]); + return promise; + } + } + const sdk = Object.defineProperty({ Agent }, "Agent", { configurable: false, - enumerable: true, - value: Agent, writable: false, }); - - const wrapped = wrapCursorSDK(sdk as any) as any; - - await expect(wrapped.Agent.prompt("hello")).resolves.toMatchObject({ - status: "finished", - }); + const wrapped = wrapCursorSDK(sdk); + expect(wrapped.Agent.value).toBe("ok"); + expect(wrapped.Agent).toBe(wrapped.Agent); + expect(wrapped.Agent.create({})).toBe(promise); + expect(wrapped.Agent.resume("agent-1")).toBe(promise); + expect(wrapped.Agent.prompt("hi")).toBe(promise); + await promise; }); }); - -function makeRun() { - return { - agentId: "agent-1", - async conversation() { - return []; - }, - id: "run-1", - stream: async function* () { - yield { - type: "assistant", - message: { content: [{ text: "hello", type: "text" }] }, - }; - }, - async wait() { - return { id: "run-1", result: "hello", status: "finished" }; - }, - }; -} diff --git a/js/src/wrappers/cursor-sdk.ts b/js/src/wrappers/cursor-sdk.ts index a2ec37d60..07a1261d1 100644 --- a/js/src/wrappers/cursor-sdk.ts +++ b/js/src/wrappers/cursor-sdk.ts @@ -1,15 +1,17 @@ import { debugLogger } from "../debug-logger"; -import { cursorSDKChannels } from "../instrumentation/plugins/cursor-sdk-interceptors"; +import { observeResult } from "../instrumentation/core/observe-result"; +import { + cursorSDKChannels, + patchCursorAgentInPlace, +} from "../instrumentation/plugins/cursor-sdk-interceptors"; import type { CursorSDKAgent, CursorSDKAgentClass, - CursorSDKAgentOptions, CursorSDKRunResult, - CursorSDKSendOptions, - CursorSDKUserMessage, } from "../vendor-sdk-types/cursor-sdk"; -const WRAPPED_AGENT = Symbol.for("braintrust.cursor-sdk.wrapped-agent"); +const wrappedClasses = new WeakMap(); +const boundAgents = new WeakSet(); /** * Wraps the Cursor TypeScript SDK with Braintrust tracing. The wrapper emits @@ -31,8 +33,8 @@ export function wrapCursorSDK(sdk: T): T { : (sdk as Record); return new Proxy(target, { - get(target, prop, receiver) { - const value = Reflect.get(target, prop, receiver); + get(target, prop) { + const value = Reflect.get(target, prop, target); if (prop === "Agent" && typeof value === "function") { return wrapCursorAgentClass(value as unknown as CursorSDKAgentClass); } @@ -60,55 +62,55 @@ function isModuleNamespace(obj: unknown): boolean { } function wrapCursorAgentClass(Agent: CursorSDKAgentClass): CursorSDKAgentClass { + const existing = wrappedClasses.get(Agent); + if (existing) return existing; const cache = new Map(); - return new Proxy(Agent, { - get(target, prop, receiver) { + const proxy = new Proxy(Agent, { + get(target, prop) { if (cache.has(prop)) { return cache.get(prop); } - const value = Reflect.get(target, prop, receiver); + const value = Reflect.get(target, prop, target); if (prop === "create" && typeof value === "function") { - const wrapped = async function ( - options: CursorSDKAgentOptions, + const wrapped = function ( + ...args: Parameters ): Promise { - const args = [options] as [CursorSDKAgentOptions]; - return wrapCursorAgent( - await cursorSDKChannels.create.invoke(value, target, args, {}), - ); + return observeResult( + cursorSDKChannels.create.invoke(value, target, args, {}), + prepareCursorAgent, + () => {}, + ) as Promise; }; cache.set(prop, wrapped); return wrapped; } if (prop === "resume" && typeof value === "function") { - const wrapped = async function ( - agentId: string, - options?: Partial, + const wrapped = function ( + ...args: Parameters ): Promise { - const args = [agentId, options] as [ - string, - Partial | undefined, - ]; - return wrapCursorAgent( - await cursorSDKChannels.resume.invoke(value, target, args, {}), - ); + return observeResult( + cursorSDKChannels.resume.invoke(value, target, args, {}), + prepareCursorAgent, + () => {}, + ) as Promise; }; cache.set(prop, wrapped); return wrapped; } if (prop === "prompt" && typeof value === "function") { - const wrapped = async function ( - message: string | CursorSDKUserMessage, - options?: CursorSDKAgentOptions, + const wrapped = function ( + ...args: Parameters ): Promise { - const args = [message, options] as [ - string | CursorSDKUserMessage, - CursorSDKAgentOptions | undefined, - ]; - return cursorSDKChannels.prompt.invoke(value, target, args, {}); + return cursorSDKChannels.prompt.invoke( + value, + target, + args, + {}, + ) as Promise; }; cache.set(prop, wrapped); return wrapped; @@ -123,54 +125,41 @@ function wrapCursorAgentClass(Agent: CursorSDKAgentClass): CursorSDKAgentClass { return value; }, }) as CursorSDKAgentClass; + wrappedClasses.set(Agent, proxy); + wrappedClasses.set(proxy, proxy); + return proxy; } -function wrapCursorAgent(agent: CursorSDKAgent): CursorSDKAgent { - if (!agent || typeof agent !== "object") { - return agent; - } - if ((agent as Record)[WRAPPED_AGENT]) { - return agent; +function prepareCursorAgent(agent: unknown): void { + if ( + !agent || + typeof agent !== "object" || + typeof Reflect.get(agent, "send") !== "function" + ) + return; + const candidate = agent as CursorSDKAgent; + patchCursorAgentInPlace(candidate); + if (boundAgents.has(candidate) || !Object.isExtensible(candidate)) return; + + // Keep factory promises and their resolved agent identity intact, while + // retaining the receiver-safe methods previously provided by the agent proxy. + for (const key of [ + "close", + "reload", + Symbol.asyncDispose, + "listArtifacts", + "downloadArtifact", + ] as const) { + const method = candidate[key]; + if (typeof method !== "function") continue; + const own = Object.getOwnPropertyDescriptor(candidate, key); + if (own && !own.configurable && !own.writable) continue; + Object.defineProperty(candidate, key, { + configurable: own?.configurable ?? true, + enumerable: own?.enumerable ?? false, + writable: own?.writable ?? true, + value: method.bind(candidate), + }); } - - const proxy = new Proxy(agent, { - get(target, prop, receiver) { - if (prop === WRAPPED_AGENT) { - return true; - } - - const value = Reflect.get(target, prop, receiver); - if ( - prop === "send" && - typeof value === "function" && - !(target as Record)[ - Symbol.for("braintrust.cursor-sdk.auto-patched-agent") - ] - ) { - return function ( - message: string | CursorSDKUserMessage, - options?: CursorSDKSendOptions, - ) { - const args = [message, options] as [ - string | CursorSDKUserMessage, - CursorSDKSendOptions | undefined, - ]; - return cursorSDKChannels.send.invoke( - value as CursorSDKAgent["send"], - target, - args, - { agent: target, operation: "send" }, - ); - }; - } - - if (typeof value === "function") { - return value.bind(target); - } - - return value; - }, - }); - - return proxy as CursorSDKAgent; + boundAgents.add(candidate); } diff --git a/js/tests/auto-instrumentations/cursor-sdk.test.ts b/js/tests/auto-instrumentations/cursor-sdk.test.ts new file mode 100644 index 000000000..d77400752 --- /dev/null +++ b/js/tests/auto-instrumentations/cursor-sdk.test.ts @@ -0,0 +1,141 @@ +import { transform } from "esbuild"; +import { runInNewContext } from "node:vm"; +import { AsyncLocalStorage } from "node:async_hooks"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { cursorSDKConfigs } from "../../src/auto-instrumentations/configs/cursor-sdk"; +import { InstrumentationMatcher } from "../../src/auto-instrumentations/orchestrion-js"; +import { + GLOBAL_INSTRUMENTATION_HOOKS_KEY, + newGlobalInvocationHook, +} from "../../src/global-instrumentation-hooks"; + +const { spans } = vi.hoisted(() => ({ spans: [] as any[] })); +vi.mock("../../src/isomorph", () => ({ + default: { newAsyncLocalStorage: () => new AsyncLocalStorage() }, +})); +vi.mock("../../src/logger", () => ({ + startSpan: (args: any) => { + const span = { + ...args, + log: vi.fn(), + end: vi.fn(), + export: async () => args.name, + }; + spans.push(span); + return span; + }, + withCurrent: (_span: unknown, call: () => unknown) => call(), +})); +import { CursorSDKPlugin } from "../../src/instrumentation/plugins/cursor-sdk-plugin"; + +// Cursor ships synchronous methods returning promises and minified class names. +const source = ` + class Unrelated { static create() { return "unrelated"; } } + export class Agent { + static #prefix = "agent"; + static create(options) { return Promise.resolve({ id: this.#prefix, options }); } + static resume(id, options) { return Promise.resolve({ id, options }); } + static prompt(message, options) { return Promise.resolve({ result: message, options }); } + static list() { return Promise.resolve(["untouched"]); } + } + export const unrelated = Unrelated.create(); +`; + +describe("Cursor SDK automatic transforms", () => { + const remove: Array<() => void> = []; + const plugin = new CursorSDKPlugin(); + beforeEach(() => { + spans.length = 0; + plugin.enable(); + }); + afterEach(() => { + plugin.disable(); + for (const unsubscribe of remove.splice(0)) unsubscribe(); + }); + it.each(["esm", "cjs"] as const)( + "instruments the minified %s public API and preserves behavior", + async (format) => { + const built = await transform(source, { + format, + minify: true, + target: "es2022", + }); + const matcher = new InstrumentationMatcher(cursorSDKConfigs); + const transformer = matcher.getTransformer( + "@cursor/sdk", + "1.0.12", + `dist/${format}/index.js`, + )!; + const instrumented = transformer.transform(built.code, format); + const calls: Array<{ + operation: string; + receiver: unknown; + args: unknown[]; + }> = []; + for (const operation of ["create", "resume", "prompt"]) { + const hook = newGlobalInvocationHook( + `orchestrion:@cursor/sdk:Agent.${operation}`, + ); + remove.push( + hook.intercept((target, receiver, args) => { + calls.push({ operation, receiver, args }); + return Reflect.apply(target, receiver, args); + }), + ); + } + const executable = + format === "esm" + ? (await transform(instrumented.code, { format: "cjs" })).code + : instrumented.code; + const module = { exports: {} as any }; + runInNewContext(executable, { + module, + exports: module.exports, + [GLOBAL_INSTRUMENTATION_HOOKS_KEY]: (globalThis as any)[ + GLOBAL_INSTRUMENTATION_HOOKS_KEY + ], + Map, + Symbol, + Promise, + }); + const { Agent, unrelated } = module.exports; + const options = { local: { cwd: "/tmp" } }; + expect(await Agent.create(options)).toEqual({ id: "agent", options }); + expect(await Agent.resume("existing", options)).toEqual({ + id: "existing", + options, + }); + expect(await Agent.prompt("hello", options)).toEqual({ + result: "hello", + options, + }); + expect(await Agent.list()).toEqual(["untouched"]); + expect(unrelated).toBe("unrelated"); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + name: "Cursor Agent", + spanAttributes: { type: "task" }, + }); + expect(spans[0].log).toHaveBeenCalledWith( + expect.objectContaining({ output: "hello" }), + ); + expect(spans[0].end).toHaveBeenCalledOnce(); + expect(calls.map((call) => call.operation)).toEqual([ + "create", + "resume", + "prompt", + ]); + expect(calls.every((call) => call.receiver === Agent)).toBe(true); + expect(calls[0].args[0]).toBe(options); + for (const unsubscribe of remove.splice(0)) unsubscribe(); + plugin.disable(); + expect(await Agent.prompt("disabled", options)).toEqual({ + result: "disabled", + options, + }); + expect(calls).toHaveLength(3); + expect(spans).toHaveLength(1); + }, + ); +});