History

execution: add Effect.fn spans for tool-call + code-exec boundaries

Wrap the execution-engine entry points and the tool-invoke dispatch path in Effect.fn spans so every tool call and code execution becomes a tracer span. This lets Axiom dashboards slice generation + execution metrics by tool name, runtime kind, and pause/inline mode without any app-level plumbing. Naming convention: - mcp.execute outermost execute / executeWithPause boundary - mcp.execute.resume resume() boundary (resume a paused execution) - mcp.tool.dispatch per tool-call dispatch - executor.tools.describe describe-tool dispatch (sandbox intrinsic) - executor.tools.search search-tools dispatch (sandbox intrinsic) - executor.sources.list list-sources dispatch (sandbox intrinsic) - executor.code.exec code-executor invocation (engine side) - executor.code.exec.dynamic_worker / .quickjs runtime-side span, annotates executor.runtime Attributes (annotateCurrentSpan): - mcp.tool.name full tool path - mcp.tool.source_id first segment of path (cheap namespace proxy) - mcp.execute.mode "inline" | "pausable" - mcp.execute.code_length source length - mcp.execute.resume.action accept | decline | cancel - executor.runtime "dynamic-worker" | "quickjs" mcp.tool.kind is not annotated at dispatch because resolving it requires a sources.list() lookup per invocation. Callers that already know the source kind (e.g. app-level wrappers) can annotate it on a parent span. org.id is not plumbed into the engine to avoid pulling cloud-only context into shared code — annotate at the caller instead.

Rhys Sullivan committed 7 months ago

Commit 84018a81 parent c64a1594 files · +223 −171Browse files

Unified · Split

4 files
  1. packages/core/execution/src/engine.ts +82 −58
  2. packages/core/execution/src/tool-invoker.ts +131 −111
  3. packages/kernel/runtime-dynamic-worker/src/executor.ts +5 −1
  4. packages/kernel/runtime-quickjs/src/index.ts +5 −1
packages/core/execution/src/engine.ts+82−58
@@ -320,78 +320,102 @@ export const createExecutionEngine = (config: ExecutionEngineConfig): ExecutionE · 319 unmodified lines
320320 * The sandbox is forked as a daemon because paused executions can outlive the
321321 * caller scope that returned the first pause, such as an HTTP request handler.
322322 */
323− const startPausableExecution = (code: string): Effect.Effect<ExecutionResult> =>
324− Effect.gen(function* () {
325− // Ref holds the current pause signal. The elicitation handler reads
326− // it each time it fires, so resume() can swap in a fresh Deferred
327− // before unblocking the fiber.
328− const pauseSignalRef = yield* Ref.make(yield* Deferred.make<InternalPausedExecution>());
329−
330− // Will be set once the fiber is forked.
331− let fiber: Fiber.Fiber<ExecuteResult, unknown>;
332−
333− const elicitationHandler: ElicitationHandler = (ctx) =>
334− Effect.gen(function* () {
335− const responseDeferred = yield* Deferred.make<typeof ElicitationResponse.Type>();
336− const id = `exec_${++nextId}`;
337−
338− const paused: InternalPausedExecution = {
339− id,
340− elicitationContext: ctx,
341− response: responseDeferred,
342− fiber: fiber!,
343− pauseSignalRef,
344− };
345− pausedExecutions.set(id, paused);
346−
347− const currentSignal = yield* Ref.get(pauseSignalRef);
348− yield* Deferred.succeed(currentSignal, paused);
349−
350− // Suspend until resume() completes responseDeferred.
351− return yield* Deferred.await(responseDeferred);
352− });
323+ const startPausableExecution = Effect.fn("mcp.execute")(function* (code: string) {
324+ yield* Effect.annotateCurrentSpan({
325+ "mcp.execute.mode": "pausable",
326+ "mcp.execute.code_length": code.length,
327+ });
353328
354− const invoker = makeFullInvoker(executor, { onElicitation: elicitationHandler });
355− fiber = yield* Effect.forkDaemon(codeExecutor.execute(code, invoker));
329+ // Ref holds the current pause signal. The elicitation handler reads
330+ // it each time it fires, so resume() can swap in a fresh Deferred
331+ // before unblocking the fiber.
332+ const pauseSignalRef = yield* Ref.make(yield* Deferred.make<InternalPausedExecution>());
333+
334+ // Will be set once the fiber is forked.
335+ let fiber: Fiber.Fiber<ExecuteResult, unknown>;
336+
337+ const elicitationHandler: ElicitationHandler = (ctx) =>
338+ Effect.gen(function* () {
339+ const responseDeferred = yield* Deferred.make<typeof ElicitationResponse.Type>();
340+ const id = `exec_${++nextId}`;
341+
342+ const paused: InternalPausedExecution = {
343+ id,
344+ elicitationContext: ctx,
345+ response: responseDeferred,
346+ fiber: fiber!,
347+ pauseSignalRef,
348+ };
349+ pausedExecutions.set(id, paused);
350+
351+ const currentSignal = yield* Ref.get(pauseSignalRef);
352+ yield* Deferred.succeed(currentSignal, paused);
353+
354+ // Suspend until resume() completes responseDeferred.
355+ return yield* Deferred.await(responseDeferred);
356+ });
356357
357− const initialSignal = yield* Ref.get(pauseSignalRef);
358− return yield* awaitCompletionOrPause(fiber, initialSignal);
359− });
358+ const invoker = makeFullInvoker(executor, { onElicitation: elicitationHandler });
359+ fiber = yield* Effect.forkDaemon(
360+ codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec")),
361+ );
362+
363+ const initialSignal = yield* Ref.get(pauseSignalRef);
364+ return (yield* awaitCompletionOrPause(fiber, initialSignal)) as ExecutionResult;
365+ });
360366
361367 /**
362368 * Resume a paused execution. Swaps in a fresh pause signal, completes
363369 * the response Deferred to unblock the fiber, then races completion
364370 * against the next pause.
365371 */
366− const resumeExecution = (
372+ const resumeExecution = Effect.fn("mcp.execute.resume")(function* (
367373 executionId: string,
368374 response: ResumeResponse,
369− ): Effect.Effect<ExecutionResult | null> =>
370− Effect.gen(function* () {
371− const paused = pausedExecutions.get(executionId);
372− if (!paused) return null;
373− pausedExecutions.delete(executionId);
374−
375− // Swap in a fresh pause signal BEFORE unblocking the fiber, so the
376− // next elicitation handler call signals this new Deferred.
377− const nextSignal = yield* Deferred.make<InternalPausedExecution>();
378− yield* Ref.set(paused.pauseSignalRef, nextSignal);
379−
380− yield* Deferred.succeed(paused.response, {
381− action: response.action,
382− content: response.content,
383− });
375+ ) {
376+ yield* Effect.annotateCurrentSpan({
377+ "mcp.execute.resume.action": response.action,
378+ });
379+
380+ const paused = pausedExecutions.get(executionId);
381+ if (!paused) return null;
382+ pausedExecutions.delete(executionId);
383+
384+ // Swap in a fresh pause signal BEFORE unblocking the fiber, so the
385+ // next elicitation handler call signals this new Deferred.
386+ const nextSignal = yield* Deferred.make<InternalPausedExecution>();
387+ yield* Ref.set(paused.pauseSignalRef, nextSignal);
384388
385− return yield* awaitCompletionOrPause(paused.fiber, nextSignal);
389+ yield* Deferred.succeed(paused.response, {
390+ action: response.action,
391+ content: response.content,
386392 });
387393
394+ return (yield* awaitCompletionOrPause(paused.fiber, nextSignal)) as ExecutionResult;
395+ });
396+
397+ /**
398+ * Inline-elicitation execute path. Wrapped so every call produces an
399+ * `mcp.execute` span with the inner `executor.code.exec` as a child.
400+ */
401+ const runInlineExecution = Effect.fn("mcp.execute")(function* (
402+ code: string,
403+ options: { readonly onElicitation: ElicitationHandler },
404+ ) {
405+ yield* Effect.annotateCurrentSpan({
406+ "mcp.execute.mode": "inline",
407+ "mcp.execute.code_length": code.length,
408+ });
409+ const invoker = makeFullInvoker(executor, {
410+ onElicitation: options.onElicitation,
411+ });
412+ return yield* codeExecutor
413+ .execute(code, invoker)
414+ .pipe(Effect.withSpan("executor.code.exec"));
415+ });
416+
388417 return {
389− execute: async (code, options) => {
390− const invoker = makeFullInvoker(executor, {
391− onElicitation: options.onElicitation,
392− });
393− return runEffect(codeExecutor.execute(code, invoker));
394− },
418+ execute: (code, options) => runEffect(runInlineExecution(code, options)),
395419
396420 executeWithPause: (code) => runEffect(startPausableExecution(code)),
397421
packages/core/execution/src/tool-invoker.ts+131−111
@@ -10,41 +10,66 @@ import type { · 9 unmodified lines
1010 import type { SandboxToolInvoker } from "@executor/codemode-core";
1111 import { ExecutionToolError } from "./errors";
1212
13+/**
14+ * Extract the source namespace from a tool path. Tool paths look like
15+ * "<sourceId>.<op>" or "<sourceId>.<group>.<op>" — we take the first
16+ * segment as a cheap, non-lookup stand-in for the source id so the span
17+ * attribute is always populated without hitting `executor.sources.list()`
18+ * per call.
19+ */
20+const extractSourceNamespace = (path: string): string => {
21+ const idx = path.indexOf(".");
22+ return idx === -1 ? path : path.slice(0, idx);
23+};
24+
1325 /**
1426 * Bridges QuickJS `tools.someSource.someOp(args)` calls into
1527 * `executor.tools.invoke(toolId, args)`.
28+ *
29+ * Wrapped in `Effect.fn("mcp.tool.dispatch")` so every tool call becomes a
30+ * span in the Effect tracer. Attributes:
31+ * - `mcp.tool.name` — full tool path (e.g. "github.repos.get")
32+ * - `mcp.tool.source_id` — first segment of the path (namespace)
33+ *
34+ * `mcp.tool.kind` (openapi | mcp | graphql | code) is NOT annotated here
35+ * because it would require a `sources.list()` lookup on every invocation.
36+ * Callers that already know the source kind can annotate at their own span.
1637 */
1738 export const makeExecutorToolInvoker = (
1839 executor: Executor,
1940 options: { readonly invokeOptions: InvokeOptions },
2041 ): SandboxToolInvoker => ({
21− invoke: ({ path, args }) =>
22− Effect.gen(function* () {
23− const result = yield* executor.tools.invoke(path as ToolId, args, options.invokeOptions).pipe(
24− Effect.catchTag("ElicitationDeclinedError", (err) =>
25− Effect.fail(
26− new ExecutionToolError({
27− message: `Tool "${err.toolId}" requires approval but the request was ${err.action === "cancel" ? "cancelled" : "declined"} by the user.`,
28− cause: err,
29− }),
30− ),
42+ invoke: Effect.fn("mcp.tool.dispatch")(function* ({ path, args }) {
43+ yield* Effect.annotateCurrentSpan({
44+ "mcp.tool.name": path,
45+ "mcp.tool.source_id": extractSourceNamespace(path),
46+ });
47+
48+ const result = yield* executor.tools.invoke(path as ToolId, args, options.invokeOptions).pipe(
49+ Effect.catchTag("ElicitationDeclinedError", (err) =>
50+ Effect.fail(
51+ new ExecutionToolError({
52+ message: `Tool "${err.toolId}" requires approval but the request was ${err.action === "cancel" ? "cancelled" : "declined"} by the user.`,
53+ cause: err,
54+ }),
3155 ),
32− );
33− const r = result as { readonly error?: unknown; readonly data?: unknown } | unknown;
34− if (
35− r !== null &&
36− typeof r === "object" &&
37− "error" in r &&
38− (r as { error?: unknown }).error !== null &&
39− (r as { error?: unknown }).error !== undefined
40− ) {
41− return yield* Effect.fail((r as { error: unknown }).error);
42− }
43− if (r !== null && typeof r === "object" && "data" in r) {
44− return (r as { data: unknown }).data;
45− }
46− return r;
47− }),
56+ ),
57+ );
58+ const r = result as { readonly error?: unknown; readonly data?: unknown } | unknown;
59+ if (
60+ r !== null &&
61+ typeof r === "object" &&
62+ "error" in r &&
63+ (r as { error?: unknown }).error !== null &&
64+ (r as { error?: unknown }).error !== undefined
65+ ) {
66+ return yield* Effect.fail((r as { error: unknown }).error);
67+ }
68+ if (r !== null && typeof r === "object" && "data" in r) {
69+ return (r as { data: unknown }).data;
70+ }
71+ return r;
72+ }),
4873 });
4974
5075 export type ToolDiscoveryResult = {
@@ -246,103 +271,98 @@ const scoreToolMatch = (tool: SearchableTool, query: string): ToolDiscoveryResul · 195 unmodified lines
246271 };
247272
248273 /** What `tools.search()` calls inside the sandbox. */
249−export const searchTools = (
274+export const searchTools = Effect.fn("executor.tools.search")(function* (
250275 executor: Executor,
251276 query: string,
252277 limit = 12,
253278 options?: { readonly namespace?: string },
254−): Effect.Effect<ReadonlyArray<ToolDiscoveryResult>> =>
255− Effect.gen(function* () {
256− if (normalizeSearchText(query).length === 0) {
257− return [];
258− }
259−
260− const all = yield* executor.tools.list().pipe(Effect.orDie);
261− return all
262− .filter((tool: Tool) => matchesNamespace(tool, options?.namespace))
263− .map((tool: Tool) => scoreToolMatch(tool, query))
264− .filter((tool): tool is ToolDiscoveryResult => tool !== null)
265− .sort((left, right) => right.score - left.score || left.path.localeCompare(right.path))
266− .slice(0, limit);
279+) {
280+ yield* Effect.annotateCurrentSpan({
281+ "executor.search.query_length": query.length,
282+ "executor.search.limit": limit,
283+ ...(options?.namespace ? { "executor.search.namespace": options.namespace } : {}),
267284 });
268285
286+ if (normalizeSearchText(query).length === 0) {
287+ return [] as ReadonlyArray<ToolDiscoveryResult>;
288+ }
289+
290+ const all = yield* executor.tools.list().pipe(Effect.orDie);
291+ return all
292+ .filter((tool: Tool) => matchesNamespace(tool, options?.namespace))
293+ .map((tool: Tool) => scoreToolMatch(tool, query))
294+ .filter((tool): tool is ToolDiscoveryResult => tool !== null)
295+ .sort((left, right) => right.score - left.score || left.path.localeCompare(right.path))
296+ .slice(0, limit);
297+});
298+
269299 /** What `tools.executor.sources.list()` calls inside the sandbox. */
270−export const listExecutorSources = (
300+export const listExecutorSources = Effect.fn("executor.sources.list")(function* (
271301 executor: Executor,
272302 options?: { readonly query?: string; readonly limit?: number },
273−): Effect.Effect<ReadonlyArray<ExecutorSourceListItem>> =>
274− Effect.gen(function* () {
275− const normalizedQuery = normalizeSearchText(options?.query ?? "");
276− const limit = options?.limit ?? 200;
277− const sources = yield* executor.sources.list().pipe(Effect.orDie);
278−
279− const filtered =
280− normalizedQuery.length === 0
281− ? sources
282− : sources.filter((source: Source) => {
283− const haystack = normalizeSearchText([source.id, source.name, source.kind].join(" "));
284− return tokenizeSearchText(normalizedQuery).every((token) => haystack.includes(token));
285− });
286−
287− // Single query for all tools, then count per source in memory.
288− const allTools = yield* executor.tools.list().pipe(Effect.orDie);
289− const toolCountBySource = new Map<string, number>();
290− for (const tool of allTools) {
291− toolCountBySource.set(tool.sourceId, (toolCountBySource.get(tool.sourceId) ?? 0) + 1);
292− }
293−
294− const withCounts = filtered.map(
295− (source: Source) =>
296− ({
297− id: source.id,
298− name: source.name,
299− kind: source.kind,
300− runtime: source.runtime,
301− canRemove: source.canRemove,
302− canRefresh: source.canRefresh,
303− toolCount: toolCountBySource.get(source.id) ?? 0,
304− }) satisfies ExecutorSourceListItem,
305− );
303+) {
304+ const normalizedQuery = normalizeSearchText(options?.query ?? "");
305+ const limit = options?.limit ?? 200;
306+ const sources = yield* executor.sources.list().pipe(Effect.orDie);
307+
308+ const filtered =
309+ normalizedQuery.length === 0
310+ ? sources
311+ : sources.filter((source: Source) => {
312+ const haystack = normalizeSearchText([source.id, source.name, source.kind].join(" "));
313+ return tokenizeSearchText(normalizedQuery).every((token) => haystack.includes(token));
314+ });
315+
316+ // Single query for all tools, then count per source in memory.
317+ const allTools = yield* executor.tools.list().pipe(Effect.orDie);
318+ const toolCountBySource = new Map<string, number>();
319+ for (const tool of allTools) {
320+ toolCountBySource.set(tool.sourceId, (toolCountBySource.get(tool.sourceId) ?? 0) + 1);
321+ }
306322
307− return withCounts
308− .sort((left, right) => left.name.localeCompare(right.name) || left.id.localeCompare(right.id))
309− .slice(0, limit);
310− });
323+ const withCounts = filtered.map(
324+ (source: Source) =>
325+ ({
326+ id: source.id,
327+ name: source.name,
328+ kind: source.kind,
329+ runtime: source.runtime,
330+ canRemove: source.canRemove,
331+ canRefresh: source.canRefresh,
332+ toolCount: toolCountBySource.get(source.id) ?? 0,
333+ }) satisfies ExecutorSourceListItem,
334+ );
335+
336+ return withCounts
337+ .sort((left, right) => left.name.localeCompare(right.name) || left.id.localeCompare(right.id))
338+ .slice(0, limit);
339+});
311340
312341 /** What `tools.describe.tool()` calls inside the sandbox. */
313−export const describeTool = (
342+export const describeTool = Effect.fn("executor.tools.describe")(function* (
314343 executor: Executor,
315344 path: string,
316−): Effect.Effect<
317− {
318− path: string;
319− name: string;
320− description?: string;
321− inputTypeScript?: string;
322− outputTypeScript?: string;
323− typeScriptDefinitions?: Record<string, string>;
324− },
325− unknown
326−> =>
327− Effect.gen(function* () {
328− // Single tools.schema() call — it already fetches the tool row
329− // internally. No need to also call tools.list() just for name/description.
330− const schema: ToolSchema | null = yield* executor.tools.schema(path);
331−
332− // tools.schema() returns null if the tool doesn't exist. Fall back to
333− // a minimal stub so callers can still render something.
334− if (schema === null) {
335− return { path, name: path };
336− }
345+) {
346+ yield* Effect.annotateCurrentSpan({ "mcp.tool.name": path });
337347
338− // The schema's id is the tool path; name/description come from the
339− // tool row which tools.schema() already loaded.
340− return {
341− path,
342− name: schema.name ?? path,
343− description: schema.description,
344− inputTypeScript: schema.inputTypeScript,
345− outputTypeScript: schema.outputTypeScript,
346− typeScriptDefinitions: schema.typeScriptDefinitions,
347− };
348− });
348+ // Single tools.schema() call — it already fetches the tool row
349+ // internally. No need to also call tools.list() just for name/description.
350+ const schema: ToolSchema | null = yield* executor.tools.schema(path);
351+
352+ // tools.schema() returns null if the tool doesn't exist. Fall back to
353+ // a minimal stub so callers can still render something.
354+ if (schema === null) {
355+ return { path, name: path };
356+ }
357+
358+ // The schema's id is the tool path; name/description come from the
359+ // tool row which tools.schema() already loaded.
360+ return {
361+ path,
362+ name: schema.name ?? path,
363+ description: schema.description,
364+ inputTypeScript: schema.inputTypeScript,
365+ outputTypeScript: schema.outputTypeScript,
366+ typeScriptDefinitions: schema.typeScriptDefinitions,
367+ };
368+});
packages/kernel/runtime-dynamic-worker/src/executor.ts+5−1
@@ -313,7 +313,11 @@ const runInDynamicWorker = ( · 312 unmodified lines
313313 new DynamicWorkerExecutionError({
314314 message: renderTransportMessage(serializeWorkerErrorValue(cause)),
315315 }),
316− });
316+ }).pipe(
317+ Effect.withSpan("executor.code.exec.dynamic_worker", {
318+ attributes: { "executor.runtime": "dynamic-worker" },
319+ }),
320+ );
317321
318322 // ---------------------------------------------------------------------------
319323 // Public API
packages/kernel/runtime-quickjs/src/index.ts+5−1
@@ -390,7 +390,11 @@ const runInQuickJs = ( · 389 unmodified lines
390390 Effect.tryPromise({
391391 try: () => evaluateInQuickJs(options, code, toolInvoker),
392392 catch: (cause) => new QuickJsExecutionError({ message: String(cause) }),
393− });
393+ }).pipe(
394+ Effect.withSpan("executor.code.exec.quickjs", {
395+ attributes: { "executor.runtime": "quickjs" },
396+ }),
397+ );
394398
395399 export const makeQuickJsExecutor = (options: QuickJsExecutorOptions = {}): CodeExecutor => ({
396400 execute: (code: string, toolInvoker: SandboxToolInvoker) =>