v4 ai improvements

.changeset/honest-files-decide.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Deprecate toolTask and replace with `ai.tool(mySchemaTask)`

.vscode/launch.json

@@ -149,6 +149,14 @@
       "command": "pnpm run test ./src/run-queue/index.test.ts",
       "cwd": "${workspaceFolder}/internal-packages/run-engine",
       "sourceMaps": true
+    },
+    {
+      "type": "node-terminal",
+      "request": "launch",
+      "name": "Debug d3-demo",
+      "command": "pnpm exec trigger dev",
+      "cwd": "${workspaceFolder}/references/d3-demo",
+      "sourceMaps": true
     }
   ]
 }

apps/webapp/app/components/runs/v3/RunIcon.tsx

@@ -73,7 +73,7 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) {
     case "trigger":
       return <TriggerIcon className={cn(className, "text-orange-500")} />;
     case "python":
-      return <PythonLogoIcon />;
+      return <PythonLogoIcon className={className} />;
     //log levels
     case "debug":
     case "log":

apps/webapp/app/v3/otlpExporter.server.ts

@@ -57,7 +57,7 @@ class OTLPExporter {
 
       const enrichedEvents = enrichCreatableEvents(events);
 
-      this.#logEventsVerbose(enrichedEvents);
+      this.#logEventsVerbose(enrichedEvents, "exportTraces");
 
       span.setAttribute("event_count", enrichedEvents.length);
 
@@ -84,7 +84,7 @@ class OTLPExporter {
 
       const enrichedEvents = enrichCreatableEvents(events);
 
-      this.#logEventsVerbose(enrichedEvents);
+      this.#logEventsVerbose(enrichedEvents, "exportLogs");
 
       span.setAttribute("event_count", enrichedEvents.length);
 
@@ -98,11 +98,11 @@ class OTLPExporter {
     });
   }
 
-  #logEventsVerbose(events: CreatableEvent[]) {
+  #logEventsVerbose(events: CreatableEvent[], prefix: string) {
     if (!this._verbose) return;
 
     events.forEach((event) => {
-      logger.debug("Exporting event", { event });
+      logger.debug(`Exporting ${prefix} event`, { event });
     });
   }
 

apps/webapp/app/v3/utils/enrichCreatableEvents.server.ts

@@ -23,7 +23,7 @@ function enrichStyle(event: CreatableEvent) {
   // GenAI System check
   const system = props["gen_ai.system"];
   if (typeof system === "string") {
-    return { ...baseStyle, icon: `tabler-brand-${system}` };
+    return { ...baseStyle, icon: `tabler-brand-${system.split(".")[0]}` };
   }
 
   // Agent workflow check
@@ -32,6 +32,16 @@ function enrichStyle(event: CreatableEvent) {
     return { ...baseStyle, icon: "tabler-brain" };
   }
 
+  const message = event.message;
+
+  if (typeof message === "string" && message === "ai.toolCall") {
+    return { ...baseStyle, icon: "tabler-tool" };
+  }
+
+  if (typeof message === "string" && message.startsWith("ai.")) {
+    return { ...baseStyle, icon: "tabler-sparkles" };
+  }
+
   return baseStyle;
 }
 

apps/webapp/package.json

@@ -258,4 +258,4 @@
   "engines": {
     "node": ">=16.0.0"
   }
-}
+}

packages/cli-v3/src/entryPoints/dev-run-worker.ts

@@ -433,6 +433,8 @@ const zodIpc = new ZodIpcConnection({
             error: {
               type: "INTERNAL_ERROR",
               code: TaskRunErrorCodes.CONFIGURED_INCORRECTLY,
+              message: err instanceof Error ? err.message : String(err),
+              stackTrace: err instanceof Error ? err.stack : undefined,
             },
             usage: {
               durationMs: 0,

packages/core/src/v3/index.ts

@@ -70,7 +70,14 @@ export * from "./utils/imageRef.js";
 export * from "./utils/heartbeat.js";
 
 export * from "./config.js";
-export { getSchemaParseFn, type AnySchemaParseFn, type SchemaParseFn } from "./types/schemas.js";
+export {
+  getSchemaParseFn,
+  type AnySchemaParseFn,
+  type SchemaParseFn,
+  isSchemaZodEsque,
+  isSchemaValibotEsque,
+  isSchemaArkTypeEsque,
+} from "./types/schemas.js";
 
 import { VERSION } from "../version.js";
 

packages/core/src/v3/otel/tracingSDK.ts

@@ -142,6 +142,7 @@ export class TracingSDK {
 
     traceProvider.addSpanProcessor(
       new TaskContextSpanProcessor(
+        traceProvider.getTracer("trigger-dev-worker", VERSION),
         getEnvVar("OTEL_BATCH_PROCESSING_ENABLED") === "1"
           ? new BatchSpanProcessor(spanExporter, {
               maxExportBatchSize: parseInt(getEnvVar("OTEL_SPAN_MAX_EXPORT_BATCH_SIZE") ?? "64"),

packages/core/src/v3/taskContext/otelProcessors.ts

@@ -4,11 +4,14 @@ import { SemanticInternalAttributes } from "../semanticInternalAttributes.js";
 import { Context } from "@opentelemetry/api";
 import { flattenAttributes } from "../utils/flattenAttributes.js";
 import { taskContext } from "../task-context-api.js";
+import { Tracer } from "@opentelemetry/api";
 
 export class TaskContextSpanProcessor implements SpanProcessor {
   private _innerProcessor: SpanProcessor;
+  private _tracer: Tracer;
 
-  constructor(innerProcessor: SpanProcessor) {
+  constructor(tracer: Tracer, innerProcessor: SpanProcessor) {
+    this._tracer = tracer;
     this._innerProcessor = innerProcessor;
   }
 
@@ -26,7 +29,15 @@ export class TaskContextSpanProcessor implements SpanProcessor {
       );
     }
 
-    this._innerProcessor.onStart(span, parentContext);
+    if (!isPartialSpan(span)) {
+      const partialSpan = createPartialSpan(this._tracer, span, parentContext);
+
+      this._innerProcessor.onStart(span, parentContext);
+
+      partialSpan.end();
+    } else {
+      this._innerProcessor.onStart(span, parentContext);
+    }
   }
 
   // Delegate the rest of the methods to the wrapped processor
@@ -44,6 +55,44 @@ export class TaskContextSpanProcessor implements SpanProcessor {
   }
 }
 
+function isPartialSpan(span: Span) {
+  return span.attributes[SemanticInternalAttributes.SPAN_PARTIAL] === true;
+}
+
+function createPartialSpan(tracer: Tracer, span: Span, parentContext: Context) {
+  const partialSpan = tracer.startSpan(
+    span.name,
+    {
+      attributes: {
+        [SemanticInternalAttributes.SPAN_PARTIAL]: true,
+        [SemanticInternalAttributes.SPAN_ID]: span.spanContext().spanId,
+        ...span.attributes,
+      },
+    },
+    parentContext
+  );
+
+  if (taskContext.ctx) {
+    partialSpan.setAttributes(
+      flattenAttributes(
+        {
+          [SemanticInternalAttributes.ATTEMPT_ID]: taskContext.ctx.attempt.id,
+          [SemanticInternalAttributes.ATTEMPT_NUMBER]: taskContext.ctx.attempt.number,
+        },
+        SemanticInternalAttributes.METADATA
+      )
+    );
+  }
+
+  if (span.events) {
+    for (const event of span.events) {
+      partialSpan.addEvent(event.name, event.attributes, event.time);
+    }
+  }
+
+  return partialSpan;
+}
+
 export class TaskContextLogProcessor implements LogRecordProcessor {
   private _innerProcessor: LogRecordProcessor;
 

packages/core/src/v3/tracer.ts

@@ -1,7 +1,6 @@
 import {
   Attributes,
   Context,
-  HrTime,
   SpanOptions,
   SpanStatusCode,
   TimeInput,
@@ -12,12 +11,12 @@ import {
   type Tracer,
 } from "@opentelemetry/api";
 import { Logger, logs } from "@opentelemetry/api-logs";
-import { SemanticInternalAttributes } from "./semanticInternalAttributes.js";
 import { clock } from "./clock-api.js";
-import { usage } from "./usage-api.js";
-import { taskContext } from "./task-context-api.js";
-import { recordSpanException } from "./otel/utils.js";
 import { isCompleteTaskWithOutput } from "./errors.js";
+import { recordSpanException } from "./otel/utils.js";
+import { SemanticInternalAttributes } from "./semanticInternalAttributes.js";
+import { taskContext } from "./task-context-api.js";
+import { usage } from "./usage-api.js";
 
 export type TriggerTracerConfig =
   | {
@@ -98,29 +97,6 @@ export class TriggerTracer {
           }
         });
 
-        if (taskContext.ctx) {
-          const partialSpan = this.tracer.startSpan(
-            name,
-            {
-              ...options,
-              attributes: {
-                ...attributes,
-                [SemanticInternalAttributes.SPAN_PARTIAL]: true,
-                [SemanticInternalAttributes.SPAN_ID]: span.spanContext().spanId,
-              },
-            },
-            parentContext
-          );
-
-          if (options?.events) {
-            for (const event of options.events) {
-              partialSpan.addEvent(event.name, event.attributes, event.startTime);
-            }
-          }
-
-          partialSpan.end();
-        }
-
         if (options?.events) {
           for (const event of options.events) {
             span.addEvent(event.name, event.attributes, event.startTime);
@@ -179,21 +155,6 @@ export class TriggerTracer {
 
     const span = this.tracer.startSpan(name, options, parentContext);
 
-    this.tracer
-      .startSpan(
-        name,
-        {
-          ...options,
-          attributes: {
-            ...attributes,
-            [SemanticInternalAttributes.SPAN_PARTIAL]: true,
-            [SemanticInternalAttributes.SPAN_ID]: span.spanContext().spanId,
-          },
-        },
-        parentContext
-      )
-      .end();
-
     return span;
   }
 }

packages/core/src/v3/types/schemas.ts

@@ -3,6 +3,18 @@ export type SchemaZodEsque<TInput, TParsedInput> = {
   _output: TParsedInput;
 };
 
+export function isSchemaZodEsque<TInput, TParsedInput>(
+  schema: Schema
+): schema is SchemaZodEsque<TInput, TParsedInput> {
+  return (
+    typeof schema === "object" &&
+    "_def" in schema &&
+    "parse" in schema &&
+    "parseAsync" in schema &&
+    "safeParse" in schema
+  );
+}
+
 export type SchemaValibotEsque<TInput, TParsedInput> = {
   schema: {
     _types?: {
@@ -12,11 +24,23 @@ export type SchemaValibotEsque<TInput, TParsedInput> = {
   };
 };
 
+export function isSchemaValibotEsque<TInput, TParsedInput>(
+  schema: Schema
+): schema is SchemaValibotEsque<TInput, TParsedInput> {
+  return typeof schema === "object" && "_types" in schema;
+}
+
 export type SchemaArkTypeEsque<TInput, TParsedInput> = {
   inferIn: TInput;
   infer: TParsedInput;
 };
 
+export function isSchemaArkTypeEsque<TInput, TParsedInput>(
+  schema: Schema
+): schema is SchemaArkTypeEsque<TInput, TParsedInput> {
+  return typeof schema === "object" && "_inferIn" in schema && "_infer" in schema;
+}
+
 export type SchemaMyZodEsque<TInput> = {
   parse: (input: any) => TInput;
 };

packages/core/src/v3/workers/taskExecutor.ts

@@ -972,7 +972,7 @@ export class TaskExecutor {
       },
       {
         attributes: {
-          [SemanticInternalAttributes.STYLE_ICON]: "clock",
+          [SemanticInternalAttributes.STYLE_ICON]: "tabler-clock",
           [SemanticInternalAttributes.COLLAPSED]: true,
         },
       }

packages/trigger-sdk/package.json

@@ -23,7 +23,8 @@
     "exports": {
       "./package.json": "./package.json",
       ".": "./src/v3/index.ts",
-      "./v3": "./src/v3/index.ts"
+      "./v3": "./src/v3/index.ts",
+      "./ai": "./src/v3/ai.ts"
     },
     "sourceDialects": [
       "@triggerdotdev/source"
@@ -33,6 +34,9 @@
     "*": {
       "v3": [
         "dist/commonjs/v3/index.d.ts"
+      ],
+      "ai": [
+        "dist/commonjs/v3/ai.d.ts"
       ]
     }
   },
@@ -66,7 +70,7 @@
     "@types/slug": "^5.0.3",
     "@types/uuid": "^9.0.0",
     "@types/ws": "^8.5.3",
-    "ai": "^4.0.1",
+    "ai": "^4.2.0",
     "encoding": "^0.1.13",
     "rimraf": "^3.0.2",
     "tshy": "^3.0.2",
@@ -75,7 +79,13 @@
     "zod": "3.23.8"
   },
   "peerDependencies": {
-    "zod": "^3.0.0"
+    "zod": "^3.0.0",
+    "ai": "^4.2.0"
+  },
+  "peerDependenciesMeta": {
+    "ai": {
+      "optional": true
+    }
   },
   "engines": {
     "node": ">=18.20.0"
@@ -103,6 +113,17 @@
         "types": "./dist/commonjs/v3/index.d.ts",
         "default": "./dist/commonjs/v3/index.js"
       }
+    },
+    "./ai": {
+      "import": {
+        "@triggerdotdev/source": "./src/v3/ai.ts",
+        "types": "./dist/esm/v3/ai.d.ts",
+        "default": "./dist/esm/v3/ai.js"
+      },
+      "require": {
+        "types": "./dist/commonjs/v3/ai.d.ts",
+        "default": "./dist/commonjs/v3/ai.js"
+      }
     }
   },
   "main": "./dist/commonjs/v3/index.js",