565 lines · 18.4 KB · TypeScript
RawDownload
import {  CheckpointId,  CheckpointRef,  CheckpointScopeId,  NodeId,  OrchestrationV2Checkpoint,  OrchestrationV2CheckpointScope,  ProviderThreadId,  RunId,  ThreadId,} from "@t3tools/contracts";import * as NodeCrypto from "node:crypto";import * as Context from "effect/Context";import * as DateTime from "effect/DateTime";import * as Effect from "effect/Effect";import * as Encoding from "effect/Encoding";import * as Layer from "effect/Layer";import * as Ref from "effect/Ref";import * as Schema from "effect/Schema";import * as Semaphore from "effect/Semaphore"; import { parseTurnDiffFilesFromNumstat } from "../checkpointing/Diffs.ts";import * as CheckpointStore from "../checkpointing/CheckpointStore.ts";import { IdAllocatorV2, type IdAllocatorV2Shape } from "./IdAllocator.ts"; const CHECKPOINT_REFS_PREFIX = "refs/t3/orchestration-v2/checkpoints";const ROOT_CHECKPOINT_SCOPE_NAME = "root"; export class CheckpointRootScopePrepareError extends Schema.TaggedErrorClass<CheckpointRootScopePrepareError>()(  "CheckpointRootScopePrepareError",  {    threadId: ThreadId,    runId: RunId,    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to prepare root checkpoint scope for run ${this.runId}.`;  }} export class CheckpointScopeEnsureError extends Schema.TaggedErrorClass<CheckpointScopeEnsureError>()(  "CheckpointScopeEnsureError",  {    scopeId: CheckpointScopeId,    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to ensure checkpoint scope ${this.scopeId}.`;  }} export class CheckpointBaselineCaptureError extends Schema.TaggedErrorClass<CheckpointBaselineCaptureError>()(  "CheckpointBaselineCaptureError",  {    scopeId: CheckpointScopeId,    ordinalWithinScope: Schema.Number,    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to capture checkpoint baseline ${this.ordinalWithinScope} for scope ${this.scopeId}.`;  }} export class CheckpointCaptureError extends Schema.TaggedErrorClass<CheckpointCaptureError>()(  "CheckpointCaptureError",  {    scopeId: CheckpointScopeId,    parentCheckpointId: Schema.optional(CheckpointId),    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to capture checkpoint for scope ${this.scopeId}.`;  }} export class CheckpointRestoreError extends Schema.TaggedErrorClass<CheckpointRestoreError>()(  "CheckpointRestoreError",  {    scopeId: CheckpointScopeId,    checkpointId: CheckpointId,    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to restore checkpoint ${this.checkpointId} for scope ${this.scopeId}.`;  }} export class CheckpointDeleteStaleRefsError extends Schema.TaggedErrorClass<CheckpointDeleteStaleRefsError>()(  "CheckpointDeleteStaleRefsError",  {    scopeId: CheckpointScopeId,    checkpointIds: Schema.Array(CheckpointId),    cause: Schema.Defect(),  },) {  override get message(): string {    return `Failed to delete stale checkpoint refs for scope ${this.scopeId}.`;  }} export const CheckpointServiceV2Error = Schema.Union([  CheckpointRootScopePrepareError,  CheckpointScopeEnsureError,  CheckpointBaselineCaptureError,  CheckpointCaptureError,  CheckpointRestoreError,  CheckpointDeleteStaleRefsError,]);export type CheckpointServiceV2Error = typeof CheckpointServiceV2Error.Type; const isCheckpointRestoreError = Schema.is(CheckpointRestoreError); export interface CheckpointServiceV2Shape {  readonly prepareRootRunScope: (input: {    readonly threadId: ThreadId;    readonly runId: RunId;    readonly rootNodeId: NodeId;    readonly providerThreadId: ProviderThreadId;    readonly cwd: string;    readonly createdAt: DateTime.Utc;  }) => Effect.Effect<OrchestrationV2CheckpointScope, CheckpointServiceV2Error>;  readonly ensureScope: (    scope: OrchestrationV2CheckpointScope,  ) => Effect.Effect<OrchestrationV2CheckpointScope, CheckpointServiceV2Error>;  readonly captureBaseline: (input: {    readonly scope: OrchestrationV2CheckpointScope;    readonly ordinalWithinScope: number;  }) => Effect.Effect<void, CheckpointServiceV2Error>;  readonly materializeBaselineCheckpoint: (input: {    readonly scope: OrchestrationV2CheckpointScope;    readonly ordinalWithinScope: number;  }) => Effect.Effect<OrchestrationV2Checkpoint, CheckpointServiceV2Error>;  readonly capture: (input: {    readonly scope: OrchestrationV2CheckpointScope;    readonly runId: RunId | null;    readonly nodeId: NodeId;    readonly ordinalWithinScope: number;    readonly appRunOrdinal: number | null;    readonly capturedAt: DateTime.Utc;  }) => Effect.Effect<OrchestrationV2Checkpoint, CheckpointServiceV2Error>;  readonly restore: (input: {    readonly scope: OrchestrationV2CheckpointScope;    readonly checkpoint: OrchestrationV2Checkpoint;  }) => Effect.Effect<void, CheckpointServiceV2Error>;  readonly deleteStaleRefs: (input: {    readonly scope: OrchestrationV2CheckpointScope;    readonly checkpoints: ReadonlyArray<OrchestrationV2Checkpoint>;  }) => Effect.Effect<void, CheckpointServiceV2Error>;} export class CheckpointServiceV2 extends Context.Service<  CheckpointServiceV2,  CheckpointServiceV2Shape>()("t3/orchestration-v2/CheckpointService/CheckpointServiceV2") {} export function checkpointRefForScopeOrdinal(input: {  readonly scopeId: CheckpointScopeId;  readonly ordinalWithinScope: number;}): CheckpointRef {  const scopeKey = NodeCrypto.createHash("sha256").update(input.scopeId).digest("hex").slice(0, 32);  return CheckpointRef.make(    `${CHECKPOINT_REFS_PREFIX}/${Encoding.encodeBase64Url(scopeKey)}/ordinal/${input.ordinalWithinScope}`,  );} function checkpointIdForScopeOrdinal(  idAllocator: IdAllocatorV2Shape,  input: {    readonly scopeId: CheckpointScopeId;    readonly ordinalWithinScope: number;  },) {  return idAllocator.allocate.checkpoint({    checkpointScopeId: input.scopeId,    name: String(input.ordinalWithinScope),  });} function makeRootRunScope(input: {  readonly idAllocator: IdAllocatorV2Shape;  readonly threadId: ThreadId;  readonly runId: RunId;  readonly rootNodeId: NodeId;  readonly providerThreadId: ProviderThreadId;  readonly cwd: string;  readonly createdAt: DateTime.Utc;}) {  return Effect.gen(function* () {    const scopeId = yield* input.idAllocator.allocate.checkpointScope({      threadId: input.threadId,      name: ROOT_CHECKPOINT_SCOPE_NAME,    });    return {      id: scopeId,      threadId: input.threadId,      runId: input.runId,      nodeId: input.rootNodeId,      parentScopeId: null,      providerThreadId: input.providerThreadId,      kind: "root_run",      ordinalWithinParent: 0,      advancesAppRunCount: true,      cwd: input.cwd,      createdAt: input.createdAt,    } satisfies OrchestrationV2CheckpointScope;  });} function makeCheckpoint(input: {  readonly id: CheckpointId;  readonly scope: OrchestrationV2CheckpointScope;  readonly runId: RunId | null;  readonly nodeId: NodeId;  readonly parentCheckpointId: CheckpointId | null;  readonly ordinalWithinScope: number;  readonly appRunOrdinal: number | null;  readonly ref: CheckpointRef;  readonly status: OrchestrationV2Checkpoint["status"];  readonly files: OrchestrationV2Checkpoint["files"];  readonly capturedAt: DateTime.Utc;}): OrchestrationV2Checkpoint {  return {    id: input.id,    threadId: input.scope.threadId,    scopeId: input.scope.id,    runId: input.runId,    nodeId: input.nodeId,    parentCheckpointId: input.parentCheckpointId,    ordinalWithinScope: input.ordinalWithinScope,    appRunOrdinal: input.appRunOrdinal,    ref: input.ref,    status: input.status,    files: input.files,    capturedAt: input.capturedAt,  };} export const layer: Layer.Layer<  CheckpointServiceV2,  never,  CheckpointStore.CheckpointStore | IdAllocatorV2> = Layer.effect(  CheckpointServiceV2,  Effect.gen(function* () {    const checkpointStore = yield* CheckpointStore.CheckpointStore;    const idAllocator = yield* IdAllocatorV2;    const workspaceSemaphores = yield* Ref.make(new Map<string, Semaphore.Semaphore>());     const getWorkspaceSemaphore = (cwd: string) =>      Effect.gen(function* () {        const existing = (yield* Ref.get(workspaceSemaphores)).get(cwd);        if (existing !== undefined) {          return existing;        }         const created = yield* Semaphore.make(1);        return yield* Ref.modify(workspaceSemaphores, (current) => {          const concurrent = current.get(cwd);          if (concurrent !== undefined) {            return [concurrent, current];          }          const updated = new Map(current);          updated.set(cwd, created);          return [created, updated];        });      });     const withWorkspaceLock = <A, E, R>(cwd: string, effect: Effect.Effect<A, E, R>) =>      Effect.flatMap(getWorkspaceSemaphore(cwd), (semaphore) => semaphore.withPermits(1)(effect));     const isGitCheckpointable = (cwd: string) =>      checkpointStore.isGitRepository(cwd).pipe(Effect.orElseSucceed(() => false));     const ensureScope: CheckpointServiceV2Shape["ensureScope"] = (scope) => Effect.succeed(scope);     const captureBaseline: CheckpointServiceV2Shape["captureBaseline"] = (input) =>      withWorkspaceLock(        input.scope.cwd,        Effect.gen(function* () {          if (!(yield* isGitCheckpointable(input.scope.cwd))) {            return;          }           const checkpointRef = checkpointRefForScopeOrdinal({            scopeId: input.scope.id,            ordinalWithinScope: input.ordinalWithinScope,          });          const exists = yield* checkpointStore.hasCheckpointRef({            cwd: input.scope.cwd,            checkpointRef,          });          if (exists) {            return;          }           yield* checkpointStore.captureCheckpoint({            cwd: input.scope.cwd,            checkpointRef,          });        }),      ).pipe(        Effect.mapError(          (cause) =>            new CheckpointBaselineCaptureError({              scopeId: input.scope.id,              ordinalWithinScope: input.ordinalWithinScope,              cause,            }),        ),      );     const materializeBaselineCheckpoint: CheckpointServiceV2Shape["materializeBaselineCheckpoint"] =      (input) =>        withWorkspaceLock(          input.scope.cwd,          Effect.gen(function* () {            const checkpointRef = checkpointRefForScopeOrdinal({              scopeId: input.scope.id,              ordinalWithinScope: input.ordinalWithinScope,            });            const checkpointId = yield* checkpointIdForScopeOrdinal(idAllocator, {              scopeId: input.scope.id,              ordinalWithinScope: input.ordinalWithinScope,            });            const checkpointable = yield* isGitCheckpointable(input.scope.cwd);            const available = checkpointable              ? yield* checkpointStore.hasCheckpointRef({                  cwd: input.scope.cwd,                  checkpointRef,                })              : false;            return makeCheckpoint({              id: checkpointId,              scope: input.scope,              runId: null,              nodeId: input.scope.nodeId,              parentCheckpointId: null,              ordinalWithinScope: input.ordinalWithinScope,              appRunOrdinal: null,              ref: checkpointRef,              status: available ? "ready" : "missing",              files: [],              capturedAt: input.scope.createdAt,            });          }),        ).pipe(          Effect.mapError(            (cause) =>              new CheckpointCaptureError({                scopeId: input.scope.id,                cause,              }),          ),        );     const capture: CheckpointServiceV2Shape["capture"] = (input) =>      withWorkspaceLock(        input.scope.cwd,        Effect.gen(function* () {          const checkpointId = yield* checkpointIdForScopeOrdinal(idAllocator, {            scopeId: input.scope.id,            ordinalWithinScope: input.ordinalWithinScope,          });          const parentCheckpointId =            input.ordinalWithinScope > 0              ? yield* checkpointIdForScopeOrdinal(idAllocator, {                  scopeId: input.scope.id,                  ordinalWithinScope: input.ordinalWithinScope - 1,                })              : null;          const checkpointRef = checkpointRefForScopeOrdinal({            scopeId: input.scope.id,            ordinalWithinScope: input.ordinalWithinScope,          });          const previousCheckpointRef = checkpointRefForScopeOrdinal({            scopeId: input.scope.id,            ordinalWithinScope: Math.max(0, input.ordinalWithinScope - 1),          });           if (!(yield* isGitCheckpointable(input.scope.cwd))) {            return makeCheckpoint({              id: checkpointId,              scope: input.scope,              runId: input.runId,              nodeId: input.nodeId,              parentCheckpointId,              ordinalWithinScope: input.ordinalWithinScope,              appRunOrdinal: input.appRunOrdinal,              ref: checkpointRef,              status: "missing",              files: [],              capturedAt: input.capturedAt,            });          }           const captured = yield* checkpointStore            .captureCheckpoint({              cwd: input.scope.cwd,              checkpointRef,            })            .pipe(              Effect.as(true),              Effect.catch((cause) =>                Effect.logWarning("orchestration V2 checkpoint capture failed", {                  scopeId: input.scope.id,                  checkpointRef,                  cause: String(cause),                }).pipe(Effect.as(false)),              ),            );           if (!captured) {            return makeCheckpoint({              id: checkpointId,              scope: input.scope,              runId: input.runId,              nodeId: input.nodeId,              parentCheckpointId,              ordinalWithinScope: input.ordinalWithinScope,              appRunOrdinal: input.appRunOrdinal,              ref: checkpointRef,              status: "error",              files: [],              capturedAt: input.capturedAt,            });          }           const previousExists = yield* checkpointStore.hasCheckpointRef({            cwd: input.scope.cwd,            checkpointRef: previousCheckpointRef,          });          const files = previousExists            ? yield* checkpointStore                .diffCheckpoints({                  cwd: input.scope.cwd,                  fromCheckpointRef: previousCheckpointRef,                  toCheckpointRef: checkpointRef,                  fallbackFromToHead: false,                  ignoreWhitespace: false,                  format: "numstat",                })                .pipe(                  Effect.map((diff) =>                    parseTurnDiffFilesFromNumstat(diff).map((file) => ({                      path: file.path,                      kind: "modified",                      additions: file.additions,                      deletions: file.deletions,                    })),                  ),                  Effect.catch((cause) =>                    Effect.logWarning("orchestration V2 checkpoint diff summary failed", {                      scopeId: input.scope.id,                      checkpointRef,                      cause: String(cause),                    }).pipe(Effect.as([])),                  ),                )            : [];           return makeCheckpoint({            id: checkpointId,            scope: input.scope,            runId: input.runId,            nodeId: input.nodeId,            parentCheckpointId,            ordinalWithinScope: input.ordinalWithinScope,            appRunOrdinal: input.appRunOrdinal,            ref: checkpointRef,            status: "ready",            files,            capturedAt: input.capturedAt,          });        }),      ).pipe(        Effect.mapError(          (cause) =>            new CheckpointCaptureError({              scopeId: input.scope.id,              cause,            }),        ),      );     const restore: CheckpointServiceV2Shape["restore"] = (input) =>      withWorkspaceLock(        input.scope.cwd,        Effect.gen(function* () {          if (input.checkpoint.status !== "ready") {            return yield* new CheckpointRestoreError({              scopeId: input.scope.id,              checkpointId: input.checkpoint.id,              cause: `Checkpoint status is ${input.checkpoint.status}.`,            });          }

Lines 1–500 of 565.