While loops

import {
  createMapStep,
  createWhileStep,
  createWorkflow,
} from "@ai_kit/core";
import { z } from "zod";

const pollingStateSchema = z.object({
  jobId: z.string().min(1),
  attempt: z.number().int().min(0),
  done: z.boolean().optional(),
});

type PollingState = z.infer<typeof pollingStateSchema>;
type PollingResult = PollingState & { done: boolean };

const checkStatus = createMapStep<PollingState, PollingResult>({
  id: "check-status",
  inputSchema: pollingStateSchema,
  outputSchema: pollingStateSchema.extend({
    done: z.boolean(),
  }),
  handler: async ({ input }) => {
    const attempt = input.attempt;

    const done = await fetch(`/jobs/${input.jobId}/status?attempt=${attempt}`)
      .then(res => res.json())
      .then(body => body.done === true);

    return {
      jobId: input.jobId,
      attempt: attempt + 1,
      done,
    };
  },
});

const pollUntilDone = createWhileStep({
  id: "poll-job",
  description: "Repeat checkStatus until completion or timeout",
  inputSchema: z.object({
    jobId: z.string().min(1),
    attempt: z.number().int().min(0).default(0),
  }),
  loopStep: checkStatus,
  maxIterations: 12,
  condition: ({ lastOutput }) => {
    return lastOutput?.done !== true;
  },
});

export const pollWorkflow = createWorkflow({
  id: "polling",
  inputSchema: z.object({
    jobId: z.string().min(1),
    attempt: z.number().int().min(0).default(0),
  }),
  outputSchema: z.object({
    lastResult: pollingStateSchema.extend({ done: z.boolean() }).optional(),
    allResults: z.array(pollingStateSchema.extend({ done: z.boolean() })),
  }),
})
  .while(pollUntilDone)
  .commit();