@sonamu-kit/tasks 0.1.2 → 0.2.0

package/src/database/backend.ts

@@ -20,7 +20,7 @@ import {
   type PaginatedResponse,
   type SleepWorkflowRunParams,
 } from "../backend";
-import { DEFAULT_RETRY_POLICY } from "../core/retry";
+import { mergeRetryPolicy, type SerializableRetryPolicy } from "../core/retry";
 import type { StepAttempt } from "../core/step";
 import type { WorkflowRun } from "../core/workflow";
 import { DEFAULT_SCHEMA, migrate } from "./base";
@@ -38,6 +38,7 @@ interface BackendPostgresOptions {
 }
 
 const logger = getLogger(["sonamu", "internal", "tasks"]);
+const queryLogger = getLogger(["sonamu", "internal", "tasks", "query"]);
 
 /**
  * Manages a connection to a Postgres database for workflow operations.
@@ -54,6 +55,12 @@ export class BackendPostgres implements Backend {
   private get knex(): Knex {
     if (!this._knex) {
       this._knex = knex(this.config);
+      this._knex.on("query", (query) => {
+        queryLogger.debug("SQL: {query}, Values: {bindings}", {
+          query: query.sql,
+          bindings: query.bindings,
+        });
+      });
     }
 
     return this._knex;
@@ -155,6 +162,17 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Creating workflow run: {workflowName}:{version}", {
+      workflowName: params.workflowName,
+      version: params.version,
+    });
+
+    // config에 retryPolicy를 포함시킵니다.
+    const configWithRetryPolicy = {
+      ...(typeof params.config === "object" && params.config !== null ? params.config : {}),
+      retryPolicy: params.retryPolicy ?? undefined,
+    };
+
     const qb = this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -165,7 +183,7 @@ export class BackendPostgres implements Backend {
         version: params.version,
         status: "pending",
         idempotency_key: params.idempotencyKey,
-        config: params.config,
+        config: JSON.stringify(configWithRetryPolicy),
         context: params.context,
         input: params.input,
         attempts: 0,
@@ -190,6 +208,7 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Getting workflow run: {workflowRunId}", { workflowRunId: params.workflowRunId });
     const workflowRun = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -228,6 +247,10 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Listing workflow runs: {after}, {before}", {
+      after: params.after,
+      before: params.before,
+    });
     const limit = params.limit ?? DEFAULT_PAGINATION_PAGE_SIZE;
     const { after, before } = params;
 
@@ -275,6 +298,10 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Claiming workflow run: {workerId}, {leaseDurationMs}", {
+      workerId: params.workerId,
+      leaseDurationMs: params.leaseDurationMs,
+    });
     const claimed = await this.knex
       .with("expired", (qb) =>
         qb
@@ -335,6 +362,11 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Extending workflow run lease: {workflowRunId}, {workerId}, {leaseDurationMs}", {
+      workflowRunId: params.workflowRunId,
+      workerId: params.workerId,
+      leaseDurationMs: params.leaseDurationMs,
+    });
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -361,6 +393,12 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Sleeping workflow run: {workflowRunId}, {workerId}, {availableAt}", {
+      workflowRunId: params.workflowRunId,
+      workerId: params.workerId,
+      availableAt: params.availableAt,
+    });
+
     // 'succeeded' status is deprecated
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
@@ -390,6 +428,12 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Completing workflow run: {workflowRunId}, {workerId}, {output}", {
+      workflowRunId: params.workflowRunId,
+      workerId: params.workerId,
+      output: params.output,
+    });
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -421,8 +465,60 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
-    const { workflowRunId, error } = params;
-    const { initialIntervalMs, backoffCoefficient, maximumIntervalMs } = DEFAULT_RETRY_POLICY;
+    const { workflowRunId, error, forceComplete, customDelayMs } = params;
+
+    logger.info("Failing workflow run: {workflowRunId}, {workerId}, {error}", {
+      workflowRunId: params.workflowRunId,
+      workerId: params.workerId,
+      error: params.error,
+    });
+
+    const workflowRun = await this.knex
+      .withSchema(DEFAULT_SCHEMA)
+      .table("workflow_runs")
+      .where("namespace_id", this.namespaceId)
+      .where("id", workflowRunId)
+      .first();
+
+    if (!workflowRun) {
+      throw new Error("Workflow run not found");
+    }
+
+    const config =
+      typeof workflowRun.config === "string" ? JSON.parse(workflowRun.config) : workflowRun.config;
+    const savedRetryPolicy: SerializableRetryPolicy | undefined = config?.retryPolicy;
+    const retryPolicy = mergeRetryPolicy(savedRetryPolicy);
+
+    const { initialIntervalMs, backoffCoefficient, maximumIntervalMs, maxAttempts } = retryPolicy;
+
+    const currentAttempts = workflowRun.attempts ?? 0;
+    const shouldForceComplete = forceComplete || currentAttempts >= maxAttempts;
+
+    if (shouldForceComplete) {
+      const [updated] = await this.knex
+        .withSchema(DEFAULT_SCHEMA)
+        .table("workflow_runs")
+        .where("namespace_id", this.namespaceId)
+        .where("id", workflowRunId)
+        .where("status", "running")
+        .where("worker_id", params.workerId)
+        .update({
+          status: "failed",
+          available_at: null,
+          finished_at: this.knex.fn.now(),
+          error: JSON.stringify(error),
+          worker_id: null,
+          started_at: null,
+          updated_at: this.knex.fn.now(),
+        })
+        .returning("*");
+
+      if (!updated) {
+        logger.error("Failed to mark workflow run failed: {params}", { params });
+        throw new Error("Failed to mark workflow run failed");
+      }
+      return updated;
+    }
 
     // this beefy query updates a workflow's status, available_at, and
     // finished_at based on the workflow's deadline and retry policy
@@ -430,7 +526,9 @@ export class BackendPostgres implements Backend {
     // if the next retry would exceed the deadline, the run is marked as
     // 'failed' and finalized, otherwise, the run is rescheduled with an updated
     // 'available_at' timestamp for the next retry
-    const retryIntervalExpr = `LEAST(${initialIntervalMs} * POWER(${backoffCoefficient}, "attempts" - 1), ${maximumIntervalMs}) * INTERVAL '1 millisecond'`;
+    const retryIntervalExpr = customDelayMs
+      ? `${customDelayMs} * INTERVAL '1 millisecond'`
+      : `LEAST(${initialIntervalMs} * POWER(${backoffCoefficient}, "attempts" - 1), ${maximumIntervalMs}) * INTERVAL '1 millisecond'`;
     const deadlineExceededCondition = `"deadline_at" IS NOT NULL AND NOW() + (${retryIntervalExpr}) >= "deadline_at"`;
 
     const [updated] = await this.knex
@@ -470,6 +568,8 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Canceling workflow run: {workflowRunId}", { workflowRunId: params.workflowRunId });
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -523,6 +623,12 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Creating step attempt: {workflowRunId}, {stepName}, {kind}", {
+      workflowRunId: params.workflowRunId,
+      stepName: params.stepName,
+      kind: params.kind,
+    });
+
     const [stepAttempt] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts")
@@ -554,6 +660,8 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Getting step attempt: {stepAttemptId}", { stepAttemptId: params.stepAttemptId });
+
     const stepAttempt = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts")
@@ -569,6 +677,12 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Listing step attempts: {workflowRunId}, {after}, {before}", {
+      workflowRunId: params.workflowRunId,
+      after: params.after,
+      before: params.before,
+    });
+
     const limit = params.limit ?? DEFAULT_PAGINATION_PAGE_SIZE;
     const { after, before } = params;
 
@@ -653,11 +767,18 @@ export class BackendPostgres implements Backend {
     };
   }
 
+  // NOTE: 실제 서비스에서 이게 안 되는 것 같은데, 쿼리 등을 체크할 필요가 있음.
   async completeStepAttempt(params: CompleteStepAttemptParams): Promise<StepAttempt> {
     if (!this.initialized) {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Marking step attempt as completed: {workflowRunId}, {stepAttemptId}, {workerId}", {
+      workflowRunId: params.workflowRunId,
+      stepAttemptId: params.stepAttemptId,
+      workerId: params.workerId,
+    });
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts as sa")
@@ -692,6 +813,13 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
+    logger.info("Marking step attempt as failed: {workflowRunId}, {stepAttemptId}, {workerId}", {
+      workflowRunId: params.workflowRunId,
+      stepAttemptId: params.stepAttemptId,
+      workerId: params.workerId,
+    });
+    logger.info("Error: {error.message}", { error: params.error.message });
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts as sa")

package/src/execution.test.ts

@@ -505,6 +505,121 @@ describe("executeWorkflow", () => {
   });
 });
 
+describe("executeWorkflow with dynamic retryPolicy", () => {
+  let backend: BackendPostgres;
+
+  beforeAll(async () => {
+    backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
+      namespaceId: randomUUID(),
+      runMigrations: false,
+    });
+    await backend.initialize();
+  });
+
+  afterAll(async () => {
+    await backend.stop();
+  });
+
+  test("calls failWorkflowRun with forceComplete=true when shouldRetry returns false", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    // shouldRetry가 false를 반환하는 retryPolicy
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-false",
+        retryPolicy: {
+          maxAttempts: 10, // 정적으로는 10번까지 허용하지만
+          shouldRetry: () => ({ shouldRetry: false, delayMs: 0 }), // 동적으로 즉시 거부
+        },
+      },
+      () => {
+        throw new Error("Intentional failure");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry가 false를 반환했으므로 즉시 failed 상태가 되어야 합니다
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("failed");
+    expect(workflowRun?.attempts).toBe(1); // 한 번만 시도하고 종료
+  });
+
+  test("calls failWorkflowRun with forceComplete=false when shouldRetry returns true", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    // shouldRetry가 true를 반환하는 retryPolicy
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-true",
+        retryPolicy: {
+          maxAttempts: 2,
+          shouldRetry: () => ({ shouldRetry: true, delayMs: 1000 }),
+        },
+      },
+      () => {
+        throw new Error("Intentional failure");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry가 true를 반환했으므로 pending 상태로 재시도 대기
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("pending");
+    expect(workflowRun?.attempts).toBe(1);
+  });
+
+  test("receives correct error and attempt number in shouldRetry function", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    let receivedError: unknown = null;
+    let receivedAttempt: number | null = null;
+
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-params",
+        retryPolicy: {
+          maxAttempts: 10,
+          shouldRetry: (error, attempt) => {
+            receivedError = error;
+            receivedAttempt = attempt;
+            return { shouldRetry: false, delayMs: 0 };
+          },
+        },
+      },
+      () => {
+        throw new Error("Test error message");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry 함수가 올바른 파라미터를 받았는지 확인
+    expect(receivedError).not.toBeNull();
+    expect((receivedError as { message?: string }).message).toBe("Test error message");
+    expect(receivedAttempt).toBe(1); // 첫 번째 시도 후이므로 1
+
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("failed");
+  });
+});
+
 function sleep(ms: number): Promise<void> {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }

package/src/execution.ts

@@ -2,6 +2,7 @@ import type { Backend } from "./backend";
 import type { DurationString } from "./core/duration";
 import { serializeError } from "./core/error";
 import type { JsonValue } from "./core/json";
+import { isDynamicRetryPolicy, type RetryPolicy } from "./core/retry";
 import type { StepAttempt, StepAttemptCache } from "./core/step";
 import {
   addToStepAttemptCache,
@@ -186,6 +187,7 @@ export interface ExecuteWorkflowParams {
   workflowFn: WorkflowFunction<unknown, unknown>;
   workflowVersion: string | null;
   workerId: string;
+  retryPolicy?: RetryPolicy;
 }
 
 /**
@@ -198,7 +200,7 @@ export interface ExecuteWorkflowParams {
  * @param params - The execution parameters
  */
 export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>): Promise<void> {
-  const { backend, workflowRun, workflowFn, workflowVersion, workerId } = params;
+  const { backend, workflowRun, workflowFn, workflowVersion, workerId, retryPolicy } = params;
 
   try {
     // load all pages of step history
@@ -281,11 +283,25 @@ export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>):
       return;
     }
 
-    // mark failure
+    // claimWorkflowRun에서 이미 attempts가 증가된 상태입니다.
+    let forceComplete = false;
+    let customDelayMs: number | undefined;
+    if (retryPolicy && isDynamicRetryPolicy(retryPolicy)) {
+      const serializedError = serializeError(error);
+      const decision = retryPolicy.shouldRetry(serializedError, workflowRun.attempts ?? 1);
+      if (!decision.shouldRetry) {
+        forceComplete = true;
+      } else {
+        customDelayMs = decision.delayMs;
+      }
+    }
+
     await backend.failWorkflowRun({
       workflowRunId: workflowRun.id,
       workerId,
       error: serializeError(error),
+      forceComplete,
+      customDelayMs,
     });
   }
 }

package/src/internal.ts

@@ -3,7 +3,27 @@ export type * from "./client";
 export { loadConfig } from "./config";
 export type { DurationString } from "./core/duration";
 export type { JsonValue } from "./core/json";
-export { DEFAULT_RETRY_POLICY } from "./core/retry";
+export {
+  DEFAULT_RETRY_POLICY,
+  mergeRetryPolicy,
+  serializeRetryPolicy,
+  shouldRetryByPolicy,
+  calculateRetryDelayMs,
+  shouldRetry,
+  isDynamicRetryPolicy,
+  isStaticRetryPolicy,
+} from "./core/retry";
+export type {
+  RetryPolicy,
+  StaticRetryPolicy,
+  DynamicRetryPolicy,
+  SerializableRetryPolicy,
+  RetryDecision,
+  RetryDecisionFn,
+  MergedStaticRetryPolicy,
+  MergedDynamicRetryPolicy,
+  MergedRetryPolicy,
+} from "./core/retry";
 export type { StandardSchemaV1 } from "./core/schema";
 export type { StepAttempt } from "./core/step";
 export type { SchemaInput, SchemaOutput, WorkflowRun } from "./core/workflow";

package/src/worker.ts

@@ -202,6 +202,7 @@ export class Worker {
         workflowFn: workflow.fn,
         workflowVersion: execution.workflowRun.version,
         workerId: execution.workerId,
+        retryPolicy: workflow.spec.retryPolicy,
       });
     } catch (error) {
       // specifically for unexpected errors in the execution wrapper itself, not

package/src/workflow.ts

@@ -1,3 +1,4 @@
+import type { RetryPolicy } from "./core/retry";
 import type { StandardSchemaV1 } from "./core/schema";
 import type { WorkflowFunction } from "./execution";
 
@@ -8,6 +9,8 @@ export interface WorkflowSpec<Input, Output, RawInput> {
   readonly version?: string;
   /** The schema used to validate inputs. */
   readonly schema?: StandardSchemaV1<RawInput, Input>;
+  /** The retry policy for the workflow. */
+  readonly retryPolicy?: RetryPolicy;
   /** Phantom type carrier - won't exist at runtime. */
   readonly __types?: {
     output: Output;