Support for Replicate task webhooks

[ericallam]

The Replicate API supports creating predictions with a webhook url (see here) and we need to add support for "remote callback" style tasks to achieve the following APIs:

import { Replicate } from "@trigger.dev/replicate";
import { TriggerClient } from "@trigger.dev/sdk";

const replicate = new Replicate({ id: "replicate", apiKey: process.env["REPLICATE_API_KEY"]! });

export const client = new TriggerClient({
  id: "job-catalog",
  apiKey: process.env["TRIGGER_API_KEY"],
  apiUrl: process.env["TRIGGER_API_URL"],
});

client.defineJob({
  id: "replicate-webhook-tasks",
  name: "Replicate Webhook Tasks",
  version: "0.0.1",
  trigger: eventTrigger({
    name: "replicate.tasks",
  }),
  integrations: {
    replicate,
  },
  run: async (payload, io, ctx) => {
    // This will wait for the completed webhook to be called
    const prediction = await io.replicate.predictions.createAndWaitForCompletion(
      "create-prediction",
      {
        version: "d55b9f2d...",
        input: { prompt: "call me later maybe" },
      }
    );

    const training = await replicate.trainings.createAndWaitForCompletion({
      version: "d55b9f2d...",
      destination: "my-username/my-model",
      input: { training_data: "..." },
    });
  },
});

The createAndWaitForCompletion method could be implemented like so:

// This could be implemented like so
class Predictions {
  constructor(private io: Replicate) {}

  createAndWaitForCompletion(key: string, params: { version: string; input: any }) {
    return this.io.runTask(
      key,
      async (replicate, task) => {
        await replicate.predictions.create({
          version,
          input,
          webhook: task.callbackUrl,
          webhook_events_filter: ["completed"],
        });
      },
      {
        callback: { enabled: true, timeoutInSeconds: 3600 }
      }
    );
  }
}

That callback option to runTask will need to be added, and it will have the following effects:

• The return of the task async callback will not be the output of the task. Instead the output will be the JSON body of the first request to the task.callbackUrl
• The callback URL will be in the form of "/api/vi/tasks/Staskid/callback"
• The task.status after calling the this. _apiClient.runTask will be WAITING in

  trigger.dev/packages/trigger-sdk/src/io.ts

  Line 558 in eb4ca0c

  const task = await this._apiClient.runTask(this._id, {

• After the task async callback is executed, the run will exit, to be resumed either after the callbackUrl receives a request or the timeoutInSeconds is reached

To support this the following will need to be changed (possibly not an exhaustive list)

• The RunTaskService (

  trigger.dev/apps/webapp/app/routes/api.v1.runs.$runId.tasks.ts

  Line 84 in eb4ca0c

  export class RunTaskService {

  ) will need to be updated to support the new callback option
  • RunTaskService will set the task status to WAITING when there is a callback option
  • If there is a timeoutInSeconds value, schedule a new graphile worker job with that delay and the task id as the payload. That graphile worker job will look up the task by the id and check it's status. If it's still in WAITING then that will resume the run with a "null" task output and setting the task status to COMPLETED (may change mind on this)
• callbackUrl will need to be added to the Run Task response
• IO.runTask will need to be updated to handle running the task async callback when task.status is WAITING and there is a callbackUrl
• Will need to create the task callback API endpoint and handle setting the task output to the body and updating the task status to COMPLETED and resuming the run

For an example of how to resume the run see how we do it in PerformTaskOperation:

trigger.dev/apps/webapp/app/services/tasks/performTaskOperation.server.ts

Line 223 in eb4ca0c

async #resumeTask(task: NonNullable<FoundTask>, output: any) {

[matt-aitken]

The hardest thing about this change might be the type inference. We want the return type to come from the extra callback, not the normal place. As the regular function is going to throw behind-the-scenes, we could return an empty object cast as the actual return. That would give us the correct type.

[nicktrn]

We want the return type to come from the extra callback, not the normal place.

Do we also want to be able to pass in a handler to verify authenticity, e.g. signatures, and payload validity, e.g. (zod) schema? Think at the very least we should be able to pass a schema. Would shift the return type to the handler or successfully parsed schema instead.

Re auth, we could append a random secret to callbackUrl. No need for client-side handlers. Example: /api/v1/runs/$runId/tasks/$taskId/callback/$secret

If it's still in WAITING then that will resume the run with a "null" task output and setting the task status to COMPLETED (may change mind on this)

Changed your mind on this yet? Can see some potential issues, as the next task will presumably depend on this output for some further action. Anything wrong with just failing the task with a timeout error?

[ericallam]

Yea I think go ahead and fail it with a timeout error 👍

[ericallam]

This has been released now in @trigger.dev/replicate@2.18 and @trigger.dev/sdk@2.18 and server version v2.1.10