From 0a6e549e0f9380d5ca7d9b5ba820123adee4b4e0 Mon Sep 17 00:00:00 2001 From: nicktrn <55853254+nicktrn@users.noreply.github.com> Date: Thu, 25 Apr 2024 10:45:09 +0100 Subject: [PATCH] don't mix messages and schemas --- packages/core/src/v3/schemas/api.ts | 2 +- packages/core/src/v3/schemas/fetch.ts | 2 +- packages/core/src/v3/schemas/messages.ts | 647 +++++++++++++++------ packages/core/src/v3/schemas/resources.ts | 2 +- packages/core/src/v3/schemas/schemas.ts | 652 ++++++---------------- 5 files changed, 654 insertions(+), 651 deletions(-) diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 12878790a..78013422c 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -1,6 +1,6 @@ import { z } from "zod"; import { BackgroundWorkerMetadata, ImageDetailsMetadata } from "./resources"; -import { QueueOptions } from "./messages"; +import { QueueOptions } from "./schemas"; export const WhoAmIResponseSchema = z.object({ userId: z.string(), diff --git a/packages/core/src/v3/schemas/fetch.ts b/packages/core/src/v3/schemas/fetch.ts index b901b8d69..ce4fbedfe 100644 --- a/packages/core/src/v3/schemas/fetch.ts +++ b/packages/core/src/v3/schemas/fetch.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { RetryOptions } from "./messages"; +import { RetryOptions } from "./schemas"; import { EventFilter } from "./eventFilter"; import { Prettify } from "../types"; diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index dc3bf6afe..78951be69 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -1,55 +1,15 @@ import { z } from "zod"; import { TaskRunExecution, TaskRunExecutionResult } from "./common"; -import { WaitReason } from "./schemas"; - -export const EnvironmentType = z.enum(["PRODUCTION", "STAGING", "DEVELOPMENT", "PREVIEW"]); -export type EnvironmentType = z.infer; - -export const MachineCpu = z - .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4)]) - .default(0.5); - -export type MachineCpu = z.infer; - -export const MachineMemory = z - .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4), z.literal(8)]) - .default(1); - -export type MachineMemory = z.infer; - -export const Machine = z.object({ - version: z.literal("v1").default("v1"), - cpu: MachineCpu, - memory: MachineMemory, -}); - -export type Machine = z.infer; - -export const TaskRunExecutionPayload = z.object({ - execution: TaskRunExecution, - traceContext: z.record(z.unknown()), - environment: z.record(z.string()).optional(), -}); - -export type TaskRunExecutionPayload = z.infer; - -export const ProdTaskRunExecution = TaskRunExecution.extend({ - worker: z.object({ - id: z.string(), - contentHash: z.string(), - version: z.string(), - }), -}); - -export type ProdTaskRunExecution = z.infer; - -export const ProdTaskRunExecutionPayload = z.object({ - execution: ProdTaskRunExecution, - traceContext: z.record(z.unknown()), - environment: z.record(z.string()).optional(), -}); - -export type ProdTaskRunExecutionPayload = z.infer; +import { + EnvironmentType, + Machine, + ProdTaskRunExecution, + ProdTaskRunExecutionPayload, + TaskMetadataWithFilePath, + TaskRunExecutionPayload, + WaitReason, +} from "./schemas"; +import { TaskResource } from "./resources"; export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [ z.object({ @@ -149,131 +109,6 @@ export const workerToChildMessages = { }), }; -export const FixedWindowRateLimit = z.object({ - type: z.literal("fixed-window"), - limit: z.number(), - window: z.union([ - z.object({ - seconds: z.number(), - }), - z.object({ - minutes: z.number(), - }), - z.object({ - hours: z.number(), - }), - ]), -}); - -export const SlidingWindowRateLimit = z.object({ - type: z.literal("sliding-window"), - limit: z.number(), - window: z.union([ - z.object({ - seconds: z.number(), - }), - z.object({ - minutes: z.number(), - }), - z.object({ - hours: z.number(), - }), - ]), -}); - -export const RateLimitOptions = z.discriminatedUnion("type", [ - FixedWindowRateLimit, - SlidingWindowRateLimit, -]); - -export const RetryOptions = z.object({ - /** The number of attempts before giving up */ - maxAttempts: z.number().int().optional(), - /** The exponential factor to use when calculating the next retry time. - * - * Each subsequent retry will be calculated as `previousTimeout * factor` - */ - factor: z.number().optional(), - /** The minimum time to wait before retrying */ - minTimeoutInMs: z.number().int().optional(), - /** The maximum time to wait before retrying */ - maxTimeoutInMs: z.number().int().optional(), - /** Randomize the timeout between retries. - * - * This can be useful to prevent the thundering herd problem where all retries happen at the same time. - */ - randomize: z.boolean().optional(), -}); - -export type RetryOptions = z.infer; - -export type RateLimitOptions = z.infer; - -export const QueueOptions = z.object({ - /** You can define a shared queue and then pass the name in to your task. - * - * @example - * - * ```ts - * const myQueue = queue({ - name: "my-queue", - concurrencyLimit: 1, - }); - - export const task1 = task({ - id: "task-1", - queue: { - name: "my-queue", - }, - run: async (payload: { message: string }) => { - // ... - }, - }); - - export const task2 = task({ - id: "task-2", - queue: { - name: "my-queue", - }, - run: async (payload: { message: string }) => { - // ... - }, - }); - * ``` - */ - name: z.string().optional(), - /** An optional property that specifies the maximum number of concurrent run executions. - * - * If this property is omitted, the task can potentially use up the full concurrency of an environment. */ - concurrencyLimit: z.number().int().min(0).max(1000).optional(), - /** @deprecated This feature is coming soon */ - rateLimit: RateLimitOptions.optional(), -}); - -export type QueueOptions = z.infer; - -export const TaskMetadata = z.object({ - id: z.string(), - packageVersion: z.string(), - queue: QueueOptions.optional(), - retry: RetryOptions.optional(), - machine: Machine.partial().optional(), - triggerSource: z.string().optional(), -}); - -export type TaskMetadata = z.infer; - -export const TaskFileMetadata = z.object({ - filePath: z.string(), - exportName: z.string(), -}); - -export type TaskFileMetadata = z.infer; - -export const TaskMetadataWithFilePath = TaskMetadata.merge(TaskFileMetadata); - -export type TaskMetadataWithFilePath = z.infer; - export const UncaughtExceptionMessage = z.object({ version: z.literal("v1").default("v1"), error: z.object({ @@ -422,3 +257,465 @@ export const ProdWorkerToChildMessages = { }), }, }; + +export const ProviderToPlatformMessages = { + LOG: { + message: z.object({ + version: z.literal("v1").default("v1"), + data: z.string(), + }), + }, + LOG_WITH_ACK: { + message: z.object({ + version: z.literal("v1").default("v1"), + data: z.string(), + }), + callback: z.object({ + status: z.literal("ok"), + }), + }, + WORKER_CRASHED: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + reason: z.string().optional(), + exitCode: z.number().optional(), + message: z.string().optional(), + logs: z.string().optional(), + }), + }, + INDEXING_FAILED: { + message: z.object({ + version: z.literal("v1").default("v1"), + deploymentId: z.string(), + error: z.object({ + name: z.string(), + message: z.string(), + stack: z.string().optional(), + }), + }), + }, +}; + +export const PlatformToProviderMessages = { + HEALTH: { + message: z.object({ + version: z.literal("v1").default("v1"), + }), + callback: z.object({ + status: z.literal("ok"), + }), + }, + INDEX: { + message: z.object({ + version: z.literal("v1").default("v1"), + imageTag: z.string(), + shortCode: z.string(), + apiKey: z.string(), + apiUrl: z.string(), + // identifiers + envId: z.string(), + envType: EnvironmentType, + orgId: z.string(), + projectId: z.string(), + deploymentId: z.string(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + error: z.object({ + name: z.string(), + message: z.string(), + stack: z.string().optional(), + }), + }), + z.object({ + success: z.literal(true), + }), + ]), + }, + // TODO: this should be a shared queue message instead + RESTORE: { + message: z.object({ + version: z.literal("v1").default("v1"), + type: z.enum(["DOCKER", "KUBERNETES"]), + location: z.string(), + reason: z.string().optional(), + imageRef: z.string(), + machine: Machine, + // identifiers + checkpointId: z.string(), + envId: z.string(), + envType: EnvironmentType, + orgId: z.string(), + projectId: z.string(), + runId: z.string(), + }), + }, + DELETE: { + message: z.object({ + version: z.literal("v1").default("v1"), + name: z.string(), + }), + callback: z.object({ + message: z.string(), + }), + }, + GET: { + message: z.object({ + version: z.literal("v1").default("v1"), + name: z.string(), + }), + }, +}; + +export const CoordinatorToPlatformMessages = { + LOG: { + message: z.object({ + version: z.literal("v1").default("v1"), + metadata: z.any(), + text: z.string(), + }), + }, + CREATE_WORKER: { + message: z.object({ + version: z.literal("v1").default("v1"), + projectRef: z.string(), + envId: z.string(), + deploymentId: z.string(), + metadata: z.object({ + cliPackageVersion: z.string().optional(), + contentHash: z.string(), + packageVersion: z.string(), + tasks: TaskResource.array(), + }), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + }), + z.object({ + success: z.literal(true), + }), + ]), + }, + READY_FOR_EXECUTION: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + totalCompletions: z.number(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + }), + z.object({ + success: z.literal(true), + payload: ProdTaskRunExecutionPayload, + }), + ]), + }, + READY_FOR_RESUME: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string(), + type: WaitReason, + }), + }, + TASK_RUN_COMPLETED: { + message: z.object({ + version: z.literal("v1").default("v1"), + execution: ProdTaskRunExecution, + completion: TaskRunExecutionResult, + checkpoint: z + .object({ + docker: z.boolean(), + location: z.string(), + }) + .optional(), + }), + }, + TASK_HEARTBEAT: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string(), + }), + }, + CHECKPOINT_CREATED: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string(), + docker: z.boolean(), + location: z.string(), + reason: z.discriminatedUnion("type", [ + z.object({ + type: z.literal("WAIT_FOR_DURATION"), + ms: z.number(), + now: z.number(), + }), + z.object({ + type: z.literal("WAIT_FOR_BATCH"), + batchFriendlyId: z.string(), + runFriendlyIds: z.string().array(), + }), + z.object({ + type: z.literal("WAIT_FOR_TASK"), + friendlyId: z.string(), + }), + z.object({ + type: z.literal("RETRYING_AFTER_FAILURE"), + attemptNumber: z.number(), + }), + ]), + }), + }, + INDEXING_FAILED: { + message: z.object({ + version: z.literal("v1").default("v1"), + deploymentId: z.string(), + error: z.object({ + name: z.string(), + message: z.string(), + stack: z.string().optional(), + }), + }), + }, +}; + +export const PlatformToCoordinatorMessages = { + RESUME_AFTER_DEPENDENCY: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + attemptId: z.string(), + attemptFriendlyId: z.string(), + completions: TaskRunExecutionResult.array(), + executions: TaskRunExecution.array(), + }), + }, + RESUME_AFTER_DURATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + attemptFriendlyId: z.string(), + }), + }, + REQUEST_ATTEMPT_CANCELLATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + attemptFriendlyId: z.string(), + }), + }, + READY_FOR_RETRY: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + }), + }, +}; + +export const ClientToSharedQueueMessages = { + READY_FOR_TASKS: { + message: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + }), + }, + BACKGROUND_WORKER_DEPRECATED: { + message: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + }), + }, + BACKGROUND_WORKER_MESSAGE: { + message: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + data: BackgroundWorkerClientMessages, + }), + }, +}; + +export const SharedQueueToClientMessages = { + SERVER_READY: { + message: z.object({ + version: z.literal("v1").default("v1"), + id: z.string(), + }), + }, + BACKGROUND_WORKER_MESSAGE: { + message: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + data: BackgroundWorkerServerMessages, + }), + }, +}; + +export const ProdWorkerToCoordinatorMessages = { + LOG: { + message: z.object({ + version: z.literal("v1").default("v1"), + text: z.string(), + }), + callback: z.void(), + }, + INDEX_TASKS: { + message: z.object({ + version: z.literal("v1").default("v1"), + deploymentId: z.string(), + tasks: TaskResource.array(), + packageVersion: z.string(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + }), + z.object({ + success: z.literal(true), + }), + ]), + }, + READY_FOR_EXECUTION: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + totalCompletions: z.number(), + }), + }, + READY_FOR_RESUME: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string(), + type: WaitReason, + }), + }, + READY_FOR_CHECKPOINT: { + message: z.object({ + version: z.literal("v1").default("v1"), + }), + }, + CANCEL_CHECKPOINT: { + message: z.object({ + version: z.enum(["v1", "v2"]).default("v2"), + }), + callback: z.object({ + checkpointCanceled: z.boolean(), + reason: WaitReason.optional(), + }), + }, + TASK_HEARTBEAT: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string(), + }), + }, + TASK_RUN_COMPLETED: { + message: z.object({ + version: z.literal("v1").default("v1"), + execution: ProdTaskRunExecution, + completion: TaskRunExecutionResult, + }), + callback: z.object({ + willCheckpointAndRestore: z.boolean(), + shouldExit: z.boolean(), + }), + }, + WAIT_FOR_DURATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + ms: z.number(), + now: z.number(), + attemptFriendlyId: z.string(), + }), + callback: z.object({ + willCheckpointAndRestore: z.boolean(), + }), + }, + WAIT_FOR_TASK: { + message: z.object({ + version: z.literal("v1").default("v1"), + friendlyId: z.string(), + // This is the attempt that is waiting + attemptFriendlyId: z.string(), + }), + callback: z.object({ + willCheckpointAndRestore: z.boolean(), + }), + }, + WAIT_FOR_BATCH: { + message: z.object({ + version: z.literal("v1").default("v1"), + batchFriendlyId: z.string(), + runFriendlyIds: z.string().array(), + // This is the attempt that is waiting + attemptFriendlyId: z.string(), + }), + callback: z.object({ + willCheckpointAndRestore: z.boolean(), + }), + }, + INDEXING_FAILED: { + message: z.object({ + version: z.literal("v1").default("v1"), + deploymentId: z.string(), + error: z.object({ + name: z.string(), + message: z.string(), + stack: z.string().optional(), + }), + }), + }, +}; + +export const CoordinatorToProdWorkerMessages = { + RESUME_AFTER_DEPENDENCY: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + completions: TaskRunExecutionResult.array(), + executions: TaskRunExecution.array(), + }), + }, + RESUME_AFTER_DURATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + }), + }, + EXECUTE_TASK_RUN: { + message: z.object({ + version: z.literal("v1").default("v1"), + executionPayload: ProdTaskRunExecutionPayload, + }), + }, + REQUEST_ATTEMPT_CANCELLATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + }), + }, + REQUEST_EXIT: { + message: z.object({ + version: z.literal("v1").default("v1"), + }), + }, + READY_FOR_RETRY: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + }), + }, +}; + +export const ProdWorkerSocketData = z.object({ + contentHash: z.string(), + projectRef: z.string(), + envId: z.string(), + runId: z.string(), + attemptFriendlyId: z.string().optional(), + podName: z.string(), + deploymentId: z.string(), + deploymentVersion: z.string(), +}); diff --git a/packages/core/src/v3/schemas/resources.ts b/packages/core/src/v3/schemas/resources.ts index 55768f2fd..a23d6e6e5 100644 --- a/packages/core/src/v3/schemas/resources.ts +++ b/packages/core/src/v3/schemas/resources.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { Machine, QueueOptions, RetryOptions } from "./messages"; +import { QueueOptions, RetryOptions, Machine } from "./schemas"; export const TaskResource = z.object({ id: z.string(), diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 5c4fd6c12..fe82a1be3 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -1,16 +1,184 @@ import { z } from "zod"; import { RequireKeys } from "../types"; -import { TaskRunExecution, TaskRunExecutionResult } from "./common"; -import { - BackgroundWorkerClientMessages, - BackgroundWorkerServerMessages, - ProdTaskRunExecution, - ProdTaskRunExecutionPayload, - RetryOptions, - Machine, - EnvironmentType, -} from "./messages"; -import { TaskResource } from "./resources"; +import { TaskRunExecution } from "./common"; + +/* + WARNING: Never import anything from ./messages here. If it's needed in both, put it here instead. +*/ + +export const EnvironmentType = z.enum(["PRODUCTION", "STAGING", "DEVELOPMENT", "PREVIEW"]); +export type EnvironmentType = z.infer; + +export const MachineCpu = z + .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4)]) + .default(0.5); + +export type MachineCpu = z.infer; + +export const MachineMemory = z + .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4), z.literal(8)]) + .default(1); + +export type MachineMemory = z.infer; + +export const Machine = z.object({ + version: z.literal("v1").default("v1"), + cpu: MachineCpu, + memory: MachineMemory, +}); + +export type Machine = z.infer; + +export const TaskRunExecutionPayload = z.object({ + execution: TaskRunExecution, + traceContext: z.record(z.unknown()), + environment: z.record(z.string()).optional(), +}); + +export type TaskRunExecutionPayload = z.infer; + +export const ProdTaskRunExecution = TaskRunExecution.extend({ + worker: z.object({ + id: z.string(), + contentHash: z.string(), + version: z.string(), + }), +}); + +export type ProdTaskRunExecution = z.infer; + +export const ProdTaskRunExecutionPayload = z.object({ + execution: ProdTaskRunExecution, + traceContext: z.record(z.unknown()), + environment: z.record(z.string()).optional(), +}); + +export type ProdTaskRunExecutionPayload = z.infer; + +export const FixedWindowRateLimit = z.object({ + type: z.literal("fixed-window"), + limit: z.number(), + window: z.union([ + z.object({ + seconds: z.number(), + }), + z.object({ + minutes: z.number(), + }), + z.object({ + hours: z.number(), + }), + ]), +}); + +export const SlidingWindowRateLimit = z.object({ + type: z.literal("sliding-window"), + limit: z.number(), + window: z.union([ + z.object({ + seconds: z.number(), + }), + z.object({ + minutes: z.number(), + }), + z.object({ + hours: z.number(), + }), + ]), +}); + +export const RateLimitOptions = z.discriminatedUnion("type", [ + FixedWindowRateLimit, + SlidingWindowRateLimit, +]); + +export type RateLimitOptions = z.infer; + +export const RetryOptions = z.object({ + /** The number of attempts before giving up */ + maxAttempts: z.number().int().optional(), + /** The exponential factor to use when calculating the next retry time. + * + * Each subsequent retry will be calculated as `previousTimeout * factor` + */ + factor: z.number().optional(), + /** The minimum time to wait before retrying */ + minTimeoutInMs: z.number().int().optional(), + /** The maximum time to wait before retrying */ + maxTimeoutInMs: z.number().int().optional(), + /** Randomize the timeout between retries. + * + * This can be useful to prevent the thundering herd problem where all retries happen at the same time. + */ + randomize: z.boolean().optional(), +}); + +export type RetryOptions = z.infer; + +export const QueueOptions = z.object({ + /** You can define a shared queue and then pass the name in to your task. + * + * @example + * + * ```ts + * const myQueue = queue({ + name: "my-queue", + concurrencyLimit: 1, + }); + + export const task1 = task({ + id: "task-1", + queue: { + name: "my-queue", + }, + run: async (payload: { message: string }) => { + // ... + }, + }); + + export const task2 = task({ + id: "task-2", + queue: { + name: "my-queue", + }, + run: async (payload: { message: string }) => { + // ... + }, + }); + * ``` + */ + name: z.string().optional(), + /** An optional property that specifies the maximum number of concurrent run executions. + * + * If this property is omitted, the task can potentially use up the full concurrency of an environment. */ + concurrencyLimit: z.number().int().min(0).max(1000).optional(), + /** @deprecated This feature is coming soon */ + rateLimit: RateLimitOptions.optional(), +}); + +export type QueueOptions = z.infer; + +export const TaskMetadata = z.object({ + id: z.string(), + packageVersion: z.string(), + queue: QueueOptions.optional(), + retry: RetryOptions.optional(), + machine: Machine.partial().optional(), + triggerSource: z.string().optional(), +}); + +export type TaskMetadata = z.infer; + +export const TaskFileMetadata = z.object({ + filePath: z.string(), + exportName: z.string(), +}); + +export type TaskFileMetadata = z.infer; + +export const TaskMetadataWithFilePath = TaskMetadata.merge(TaskFileMetadata); + +export type TaskMetadataWithFilePath = z.infer; export const PostStartCauses = z.enum(["index", "create", "restore"]); export type PostStartCauses = z.infer; @@ -55,465 +223,3 @@ export type ResolvedConfig = RequireKeys< export const WaitReason = z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]); export type WaitReason = z.infer; - -export const ProviderToPlatformMessages = { - LOG: { - message: z.object({ - version: z.literal("v1").default("v1"), - data: z.string(), - }), - }, - LOG_WITH_ACK: { - message: z.object({ - version: z.literal("v1").default("v1"), - data: z.string(), - }), - callback: z.object({ - status: z.literal("ok"), - }), - }, - WORKER_CRASHED: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - reason: z.string().optional(), - exitCode: z.number().optional(), - message: z.string().optional(), - logs: z.string().optional(), - }), - }, - INDEXING_FAILED: { - message: z.object({ - version: z.literal("v1").default("v1"), - deploymentId: z.string(), - error: z.object({ - name: z.string(), - message: z.string(), - stack: z.string().optional(), - }), - }), - }, -}; - -export const PlatformToProviderMessages = { - HEALTH: { - message: z.object({ - version: z.literal("v1").default("v1"), - }), - callback: z.object({ - status: z.literal("ok"), - }), - }, - INDEX: { - message: z.object({ - version: z.literal("v1").default("v1"), - imageTag: z.string(), - shortCode: z.string(), - apiKey: z.string(), - apiUrl: z.string(), - // identifiers - envId: z.string(), - envType: EnvironmentType, - orgId: z.string(), - projectId: z.string(), - deploymentId: z.string(), - }), - callback: z.discriminatedUnion("success", [ - z.object({ - success: z.literal(false), - error: z.object({ - name: z.string(), - message: z.string(), - stack: z.string().optional(), - }), - }), - z.object({ - success: z.literal(true), - }), - ]), - }, - // TODO: this should be a shared queue message instead - RESTORE: { - message: z.object({ - version: z.literal("v1").default("v1"), - type: z.enum(["DOCKER", "KUBERNETES"]), - location: z.string(), - reason: z.string().optional(), - imageRef: z.string(), - machine: Machine, - // identifiers - checkpointId: z.string(), - envId: z.string(), - envType: EnvironmentType, - orgId: z.string(), - projectId: z.string(), - runId: z.string(), - }), - }, - DELETE: { - message: z.object({ - version: z.literal("v1").default("v1"), - name: z.string(), - }), - callback: z.object({ - message: z.string(), - }), - }, - GET: { - message: z.object({ - version: z.literal("v1").default("v1"), - name: z.string(), - }), - }, -}; - -export const CoordinatorToPlatformMessages = { - LOG: { - message: z.object({ - version: z.literal("v1").default("v1"), - metadata: z.any(), - text: z.string(), - }), - }, - CREATE_WORKER: { - message: z.object({ - version: z.literal("v1").default("v1"), - projectRef: z.string(), - envId: z.string(), - deploymentId: z.string(), - metadata: z.object({ - cliPackageVersion: z.string().optional(), - contentHash: z.string(), - packageVersion: z.string(), - tasks: TaskResource.array(), - }), - }), - callback: z.discriminatedUnion("success", [ - z.object({ - success: z.literal(false), - }), - z.object({ - success: z.literal(true), - }), - ]), - }, - READY_FOR_EXECUTION: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - totalCompletions: z.number(), - }), - callback: z.discriminatedUnion("success", [ - z.object({ - success: z.literal(false), - }), - z.object({ - success: z.literal(true), - payload: ProdTaskRunExecutionPayload, - }), - ]), - }, - READY_FOR_RESUME: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptFriendlyId: z.string(), - type: WaitReason, - }), - }, - TASK_RUN_COMPLETED: { - message: z.object({ - version: z.literal("v1").default("v1"), - execution: ProdTaskRunExecution, - completion: TaskRunExecutionResult, - checkpoint: z - .object({ - docker: z.boolean(), - location: z.string(), - }) - .optional(), - }), - }, - TASK_HEARTBEAT: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptFriendlyId: z.string(), - }), - }, - CHECKPOINT_CREATED: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptFriendlyId: z.string(), - docker: z.boolean(), - location: z.string(), - reason: z.discriminatedUnion("type", [ - z.object({ - type: z.literal("WAIT_FOR_DURATION"), - ms: z.number(), - now: z.number(), - }), - z.object({ - type: z.literal("WAIT_FOR_BATCH"), - batchFriendlyId: z.string(), - runFriendlyIds: z.string().array(), - }), - z.object({ - type: z.literal("WAIT_FOR_TASK"), - friendlyId: z.string(), - }), - z.object({ - type: z.literal("RETRYING_AFTER_FAILURE"), - attemptNumber: z.number(), - }), - ]), - }), - }, - INDEXING_FAILED: { - message: z.object({ - version: z.literal("v1").default("v1"), - deploymentId: z.string(), - error: z.object({ - name: z.string(), - message: z.string(), - stack: z.string().optional(), - }), - }), - }, -}; - -export const PlatformToCoordinatorMessages = { - RESUME_AFTER_DEPENDENCY: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - attemptId: z.string(), - attemptFriendlyId: z.string(), - completions: TaskRunExecutionResult.array(), - executions: TaskRunExecution.array(), - }), - }, - RESUME_AFTER_DURATION: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptId: z.string(), - attemptFriendlyId: z.string(), - }), - }, - REQUEST_ATTEMPT_CANCELLATION: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptId: z.string(), - attemptFriendlyId: z.string(), - }), - }, - READY_FOR_RETRY: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - }), - }, -}; - -export const ClientToSharedQueueMessages = { - READY_FOR_TASKS: { - message: z.object({ - version: z.literal("v1").default("v1"), - backgroundWorkerId: z.string(), - }), - }, - BACKGROUND_WORKER_DEPRECATED: { - message: z.object({ - version: z.literal("v1").default("v1"), - backgroundWorkerId: z.string(), - }), - }, - BACKGROUND_WORKER_MESSAGE: { - message: z.object({ - version: z.literal("v1").default("v1"), - backgroundWorkerId: z.string(), - data: BackgroundWorkerClientMessages, - }), - }, -}; - -export const SharedQueueToClientMessages = { - SERVER_READY: { - message: z.object({ - version: z.literal("v1").default("v1"), - id: z.string(), - }), - }, - BACKGROUND_WORKER_MESSAGE: { - message: z.object({ - version: z.literal("v1").default("v1"), - backgroundWorkerId: z.string(), - data: BackgroundWorkerServerMessages, - }), - }, -}; - -export const ProdWorkerToCoordinatorMessages = { - LOG: { - message: z.object({ - version: z.literal("v1").default("v1"), - text: z.string(), - }), - callback: z.void(), - }, - INDEX_TASKS: { - message: z.object({ - version: z.literal("v1").default("v1"), - deploymentId: z.string(), - tasks: TaskResource.array(), - packageVersion: z.string(), - }), - callback: z.discriminatedUnion("success", [ - z.object({ - success: z.literal(false), - }), - z.object({ - success: z.literal(true), - }), - ]), - }, - READY_FOR_EXECUTION: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - totalCompletions: z.number(), - }), - }, - READY_FOR_RESUME: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptFriendlyId: z.string(), - type: WaitReason, - }), - }, - READY_FOR_CHECKPOINT: { - message: z.object({ - version: z.literal("v1").default("v1"), - }), - }, - CANCEL_CHECKPOINT: { - message: z.object({ - version: z.enum(["v1", "v2"]).default("v2"), - }), - callback: z.object({ - checkpointCanceled: z.boolean(), - reason: WaitReason.optional(), - }), - }, - TASK_HEARTBEAT: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptFriendlyId: z.string(), - }), - }, - TASK_RUN_COMPLETED: { - message: z.object({ - version: z.literal("v1").default("v1"), - execution: ProdTaskRunExecution, - completion: TaskRunExecutionResult, - }), - callback: z.object({ - willCheckpointAndRestore: z.boolean(), - shouldExit: z.boolean(), - }), - }, - WAIT_FOR_DURATION: { - message: z.object({ - version: z.literal("v1").default("v1"), - ms: z.number(), - now: z.number(), - attemptFriendlyId: z.string(), - }), - callback: z.object({ - willCheckpointAndRestore: z.boolean(), - }), - }, - WAIT_FOR_TASK: { - message: z.object({ - version: z.literal("v1").default("v1"), - friendlyId: z.string(), - // This is the attempt that is waiting - attemptFriendlyId: z.string(), - }), - callback: z.object({ - willCheckpointAndRestore: z.boolean(), - }), - }, - WAIT_FOR_BATCH: { - message: z.object({ - version: z.literal("v1").default("v1"), - batchFriendlyId: z.string(), - runFriendlyIds: z.string().array(), - // This is the attempt that is waiting - attemptFriendlyId: z.string(), - }), - callback: z.object({ - willCheckpointAndRestore: z.boolean(), - }), - }, - INDEXING_FAILED: { - message: z.object({ - version: z.literal("v1").default("v1"), - deploymentId: z.string(), - error: z.object({ - name: z.string(), - message: z.string(), - stack: z.string().optional(), - }), - }), - }, -}; - -export const CoordinatorToProdWorkerMessages = { - RESUME_AFTER_DEPENDENCY: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptId: z.string(), - completions: TaskRunExecutionResult.array(), - executions: TaskRunExecution.array(), - }), - }, - RESUME_AFTER_DURATION: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptId: z.string(), - }), - }, - EXECUTE_TASK_RUN: { - message: z.object({ - version: z.literal("v1").default("v1"), - executionPayload: ProdTaskRunExecutionPayload, - }), - }, - REQUEST_ATTEMPT_CANCELLATION: { - message: z.object({ - version: z.literal("v1").default("v1"), - attemptId: z.string(), - }), - }, - REQUEST_EXIT: { - message: z.object({ - version: z.literal("v1").default("v1"), - }), - }, - READY_FOR_RETRY: { - message: z.object({ - version: z.literal("v1").default("v1"), - runId: z.string(), - }), - }, -}; - -export const ProdWorkerSocketData = z.object({ - contentHash: z.string(), - projectRef: z.string(), - envId: z.string(), - runId: z.string(), - attemptFriendlyId: z.string().optional(), - podName: z.string(), - deploymentId: z.string(), - deploymentVersion: z.string(), -});