Skip to content
← Public packages

@jpfi/stories

Create personalized Teddy & Pinguin bedtime stories for your child and get the finished text back.

FollowFork

src/story-run.ts

348 lines · 10.0 KB · TypeScript
Raw
import type { PackageStorage } from "kody:runtime";

import type {
  CheckConnectionInput,
  ConnectionStatus,
  ContentLocale,
  GenerateStoryInput,
  GetStoryResultInput,
  StoryGenerationResult,
  StoryPhase,
} from "./domain";
import {
  DEFAULT_WAIT_SECONDS,
  TOKEN_SETUP_URL,
  deriveIdempotencyKey,
  describeConnection,
  describeFailure,
  describePending,
  normalizeContentLocale,
  normalizeRunId,
  normalizeStoryRequest,
  pollDelayMs,
  requestFingerprint,
  waitBudgetMs,
} from "./domain";
import type { ApiFailure, RunState, TeddyClient } from "./teddy-api";
import { openTeddyClient } from "./teddy-api";

export type StorySession = {
  generate(input: GenerateStoryInput): Promise<StoryGenerationResult>;
  resume(input: GetStoryResultInput): Promise<StoryGenerationResult>;
  checkConnection(input?: CheckConnectionInput): Promise<ConnectionStatus>;
};

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => {
    setTimeout(resolve, ms);
  });
}

function isRetryable(failure: ApiFailure): boolean {
  return failure.kind === "transient" || failure.kind === "rate_limited";
}

function retryDelayMs(failure: ApiFailure, attempt: number): number {
  if (failure.kind === "rate_limited" && failure.retryAfterSeconds !== null) {
    return failure.retryAfterSeconds * 1000;
  }
  if (failure.kind === "transient" && failure.retryAfterSeconds !== null) {
    return failure.retryAfterSeconds * 1000;
  }
  return pollDelayMs(attempt);
}

function capSleep(desiredMs: number, deadlineMs: number): number {
  return Math.max(0, Math.min(desiredMs, deadlineMs - Date.now()));
}

function toResult(
  failure: ApiFailure,
  ctx: { runId: string | null; locale: ContentLocale },
): StoryGenerationResult {
  switch (failure.kind) {
    case "unauthenticated":
      return {
        status: "disconnected",
        message: describeFailure("AUTHENTICATION_REQUIRED", null, ctx.locale),
        tokenSetupUrl: TOKEN_SETUP_URL,
      };
    case "not_found":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: failure.code,
        message: describeFailure(failure.code, null, ctx.locale),
        retryAfterSeconds: null,
      };
    case "quota_exhausted":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: failure.code,
        message: describeFailure(failure.code, null, ctx.locale),
        retryAfterSeconds: null,
      };
    case "rate_limited":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: "RATE_LIMITED",
        message: describeFailure("RATE_LIMITED", null, ctx.locale),
        retryAfterSeconds: failure.retryAfterSeconds,
      };
    case "rejected_request":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: failure.code,
        message: describeFailure(failure.code, null, ctx.locale),
        retryAfterSeconds: null,
      };
    case "transient":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: "TEDDY_UNAVAILABLE",
        message: describeFailure("TEDDY_UNAVAILABLE", null, ctx.locale),
        retryAfterSeconds: failure.retryAfterSeconds,
      };
    case "unexpected":
      return {
        status: "unavailable",
        runId: ctx.runId,
        code: failure.code,
        message: describeFailure(failure.code, null, ctx.locale),
        retryAfterSeconds: null,
      };
    default: {
      const exhaustive: never = failure;
      return exhaustive;
    }
  }
}

function declined(
  runId: string,
  code: string,
  canRetry: boolean,
  locale: ContentLocale,
): StoryGenerationResult {
  return {
    status: "declined",
    runId,
    code,
    message: describeFailure(code, null, locale),
    canRetry,
  };
}

function pendingResult(
  runId: string,
  phase: StoryPhase | null,
  startedAt: number,
  locale: ContentLocale,
): StoryGenerationResult {
  return {
    status: "pending",
    runId,
    phase,
    waitedSeconds: Math.max(0, Math.round((Date.now() - startedAt) / 1000)),
    nextStep: describePending(runId, phase, locale),
  };
}

async function awaitRun(
  client: TeddyClient,
  runId: string,
  locale: ContentLocale,
  deadlineMs: number,
): Promise<StoryGenerationResult> {
  let attempt = 0;
  const startedAt = Date.now();
  let lastPhase: StoryPhase | null = null;

  do {
    const read = await client.readRun(runId);
    if (!read.ok) {
      if (isRetryable(read.failure) && Date.now() < deadlineMs) {
        await sleep(capSleep(retryDelayMs(read.failure, attempt++), deadlineMs));
        continue;
      }
      if (isRetryable(read.failure)) {
        break;
      }
      return toResult(read.failure, { runId, locale });
    }

    lastPhase = read.run.phase;
    const terminal = await settleRun(client, read.run, locale, deadlineMs, attempt);
    if (terminal.kind === "result") {
      return terminal.result;
    }
    if (terminal.kind === "retry") {
      attempt = terminal.attempt;
      await sleep(capSleep(terminal.delayMs, deadlineMs));
      continue;
    }
    if (Date.now() >= deadlineMs) {
      break;
    }
    await sleep(capSleep(pollDelayMs(attempt++), deadlineMs));
  } while (Date.now() < deadlineMs);

  return pendingResult(runId, lastPhase, startedAt, locale);
}

type Settle =
  | { kind: "result"; result: StoryGenerationResult }
  | { kind: "retry"; delayMs: number; attempt: number }
  | { kind: "continue" };

async function settleRun(
  client: TeddyClient,
  run: RunState,
  locale: ContentLocale,
  deadlineMs: number,
  attempt: number,
): Promise<Settle> {
  switch (run.status) {
    case "accepted": {
      if (!run.storyHref) {
        return { kind: "continue" };
      }
      const loaded = await client.loadStory(run.storyHref, locale);
      if (loaded.ok) {
        return {
          kind: "result",
          result: { status: "ready", runId: run.runId, story: loaded.story },
        };
      }
      if (isRetryable(loaded.failure) && Date.now() < deadlineMs) {
        return {
          kind: "retry",
          delayMs: retryDelayMs(loaded.failure, attempt),
          attempt: attempt + 1,
        };
      }
      if (isRetryable(loaded.failure)) {
        return { kind: "continue" };
      }
      return {
        kind: "result",
        result: toResult(loaded.failure, { runId: run.runId, locale }),
      };
    }
    case "rejected":
      return {
        kind: "result",
        result: declined(run.runId, run.error?.code ?? "STORY_REJECTED", true, locale),
      };
    case "cancelled":
      return {
        kind: "result",
        result: declined(run.runId, "STORY_RUN_CANCELLED", false, locale),
      };
    case "failed":
      return {
        kind: "result",
        result: {
          status: "unavailable",
          runId: run.runId,
          code: run.error?.code ?? "STORY_GENERATION_FAILED",
          message: describeFailure(
            run.error?.code ?? "STORY_GENERATION_FAILED",
            run.error?.phase ?? null,
            locale,
          ),
          retryAfterSeconds: null,
        },
      };
    case "queued":
    case "running":
    case "cancelling":
      return { kind: "continue" };
    default: {
      const exhaustive: never = run.status;
      return exhaustive;
    }
  }
}

export async function openStorySession(storage: PackageStorage): Promise<StorySession> {
  const client = await openTeddyClient(storage);
  return {
    async generate(input) {
      const normalized = normalizeStoryRequest(input);
      if (!normalized.ok) {
        return normalized.failure;
      }
      const locale = normalized.request.contentLocale;
      const deadlineMs = Date.now() + waitBudgetMs(input.maxWaitSeconds, DEFAULT_WAIT_SECONDS);
      const fingerprint = await requestFingerprint(client.apiBaseUrl, normalized.request);
      const key = deriveIdempotencyKey(fingerprint, Date.now());
      let seriesId = normalized.request.seriesId;
      if (normalized.request.continueFromStoryId) {
        const created = await client.createSeries({
          contentLocale: locale,
          continueFromStoryId: normalized.request.continueFromStoryId,
        });
        if (!created.ok) {
          return toResult(created.failure, { runId: null, locale });
        }
        seriesId = created.seriesId;
      } else if (normalized.request.seriesStart && !seriesId) {
        const stored = await storage?.get(`series:${fingerprint}`);
        if (typeof stored === "string" && stored) {
          seriesId = stored;
        } else {
          const created = await client.createSeries({ contentLocale: locale });
          if (!created.ok) {
            return toResult(created.failure, { runId: null, locale });
          }
          seriesId = created.seriesId;
          await storage?.set(`series:${fingerprint}`, seriesId);
        }
      }
      const started = await client.startRun(
        { ...normalized.request, seriesId },
        key,
      );
      if (!started.ok) {
        return toResult(started.failure, { runId: null, locale });
      }
      return awaitRun(client, started.run.runId, locale, deadlineMs);
    },
    async resume(input) {
      const parsed = normalizeRunId(input.runId);
      if (!parsed.ok) {
        return parsed.failure;
      }
      const deadlineMs = Date.now() + waitBudgetMs(input.maxWaitSeconds, DEFAULT_WAIT_SECONDS);
      return awaitRun(client, parsed.runId, "de", deadlineMs);
    },
    async checkConnection(input = {}) {
      const locale = normalizeContentLocale(input.contentLocale);
      const probe = await client.probeIdentity();
      if (probe === "authenticated") {
        return {
          status: "connected",
          apiBaseUrl: client.apiBaseUrl,
          message: describeConnection("connected", locale),
        };
      }
      if (probe === "unauthenticated") {
        return {
          status: "disconnected",
          message: describeConnection("disconnected", locale),
          tokenSetupUrl: TOKEN_SETUP_URL,
        };
      }
      return {
        status: "unreachable",
        apiBaseUrl: client.apiBaseUrl,
        message: describeConnection("unreachable", locale),
      };
    },
  };
}