sandbox-agent/factory/packages/backend/src/actors/polling.ts

import { Loop } from "rivetkit/workflow";
import { normalizeMessages } from "../services/queue.js";

export interface PollingControlState {
  intervalMs: number;
  running: boolean;
}

export interface PollingControlQueueNames {
  start: string;
  stop: string;
  setInterval: string;
  force: string;
}

export interface PollingQueueMessage {
  name: string;
  body: unknown;
  complete(response: unknown): Promise<void>;
}

interface PollingActorContext<TState extends PollingControlState> {
  state: TState;
  abortSignal: AbortSignal;
  queue: {
    nextBatch(options: {
      names: readonly string[];
      timeout: number;
      count: number;
      completable: true;
    }): Promise<PollingQueueMessage[]>;
  };
}

interface RunPollingOptions<TState extends PollingControlState> {
  control: PollingControlQueueNames;
  onPoll(c: PollingActorContext<TState>): Promise<void>;
}

export async function runPollingControlLoop<TState extends PollingControlState>(
  c: PollingActorContext<TState>,
  options: RunPollingOptions<TState>
): Promise<void> {
  while (!c.abortSignal.aborted) {
    const messages = normalizeMessages(
      await c.queue.nextBatch({
        names: [
          options.control.start,
          options.control.stop,
          options.control.setInterval,
          options.control.force
        ],
        timeout: Math.max(500, c.state.intervalMs),
        count: 16,
        completable: true
      })
    ) as PollingQueueMessage[];

    if (messages.length === 0) {
      if (!c.state.running) {
        continue;
      }
      await options.onPoll(c);
      continue;
    }

    for (const msg of messages) {
      if (msg.name === options.control.start) {
        c.state.running = true;
        await msg.complete({ ok: true });
        continue;
      }

      if (msg.name === options.control.stop) {
        c.state.running = false;
        await msg.complete({ ok: true });
        continue;
      }

      if (msg.name === options.control.setInterval) {
        const intervalMs = Number((msg.body as { intervalMs?: unknown })?.intervalMs);
        c.state.intervalMs = Number.isFinite(intervalMs) ? Math.max(500, intervalMs) : c.state.intervalMs;
        await msg.complete({ ok: true });
        continue;
      }

      if (msg.name === options.control.force) {
        await options.onPoll(c);
        await msg.complete({ ok: true });
      }
    }
  }
}

interface WorkflowPollingActorContext<TState extends PollingControlState> {
  state: TState;
  loop(config: {
    name: string;
    historyEvery: number;
    historyKeep: number;
    run(ctx: WorkflowPollingActorContext<TState>): Promise<unknown>;
  }): Promise<void>;
}

interface WorkflowPollingQueueMessage extends PollingQueueMessage {}

interface WorkflowPollingLoopContext<TState extends PollingControlState> {
  state: TState;
  queue: {
    nextBatch(name: string, options: {
      names: readonly string[];
      timeout: number;
      count: number;
      completable: true;
    }): Promise<WorkflowPollingQueueMessage[]>;
  };
  step<T>(
    nameOrConfig:
      | string
      | {
          name: string;
          timeout?: number;
          run: () => Promise<T>;
        },
    run?: () => Promise<T>,
  ): Promise<T>;
}

export async function runWorkflowPollingLoop<TState extends PollingControlState>(
  ctx: any,
  options: RunPollingOptions<TState> & { loopName: string },
): Promise<void> {
  await ctx.loop(options.loopName, async (loopCtx: WorkflowPollingLoopContext<TState>) => {
    const control = await loopCtx.step("read-control-state", async () => ({
      intervalMs: Math.max(500, Number(loopCtx.state.intervalMs) || 500),
      running: Boolean(loopCtx.state.running),
    }));

    const messages = normalizeMessages(
      await loopCtx.queue.nextBatch("next-polling-control-batch", {
        names: [
          options.control.start,
          options.control.stop,
          options.control.setInterval,
          options.control.force,
        ],
        timeout: control.running ? control.intervalMs : 60_000,
        count: 16,
        completable: true,
      }),
    ) as WorkflowPollingQueueMessage[];

    if (messages.length === 0) {
      if (control.running) {
        await loopCtx.step({
          name: "poll-tick",
          timeout: 5 * 60_000,
          run: async () => {
            await options.onPoll(loopCtx as unknown as PollingActorContext<TState>);
          },
        });
      }
      return Loop.continue(undefined);
    }

    for (const msg of messages) {
      if (msg.name === options.control.start) {
        await loopCtx.step("control-start", async () => {
          loopCtx.state.running = true;
        });
        await msg.complete({ ok: true });
        continue;
      }

        if (msg.name === options.control.stop) {
          await loopCtx.step("control-stop", async () => {
            loopCtx.state.running = false;
          });
          await msg.complete({ ok: true });
          continue;
        }

        if (msg.name === options.control.setInterval) {
          await loopCtx.step("control-set-interval", async () => {
            const intervalMs = Number((msg.body as { intervalMs?: unknown })?.intervalMs);
            loopCtx.state.intervalMs = Number.isFinite(intervalMs)
              ? Math.max(500, intervalMs)
              : loopCtx.state.intervalMs;
          });
          await msg.complete({ ok: true });
          continue;
        }

        if (msg.name === options.control.force) {
          await loopCtx.step({
            name: "control-force",
            timeout: 5 * 60_000,
            run: async () => {
              await options.onPoll(loopCtx as unknown as PollingActorContext<TState>);
            },
          });
          await msg.complete({ ok: true });
        }
      }

    return Loop.continue(undefined);
  });
}