1
import { asc, eq } from "drizzle-orm";
2
 
3
import { type ProjectId, RunId } from "@/contracts";
4
import { expectTrusted } from "@/worker/contracts";
5
import * as projectSchema from "@/worker/db/durable/schema/project-do";
6
import { isNoContainerInstanceError } from "@/worker/sandbox/container-errors";
7
 
8
import { SANDBOX_CLEANUP_RETRY_DELAYS_MS } from "../constants";
9
import { getSandboxCleanupRetryState, rescheduleAlarmInTransaction, sandboxCleanupRetryKey } from "../sidecar-state";
10
import type { ProjectDoContext, SandboxCleanupRetryState } from "../types";
11
 
12
const buildRetryState = (attempt: number): SandboxCleanupRetryState => ({
13
  attempt,
14
  nextAt:
15
    Date.now() + SANDBOX_CLEANUP_RETRY_DELAYS_MS[Math.min(attempt - 1, SANDBOX_CLEANUP_RETRY_DELAYS_MS.length - 1)],
16
});
17
 
18
export const persistSandboxCleanupRetryState = async (
19
  context: ProjectDoContext,
20
  projectId: ProjectId,
21
  runId: RunId,
22
  value: SandboxCleanupRetryState | null,
23
): Promise<void> => {
24
  await context.ctx.storage.transaction(async (txn) => {
25
    if (value === null) {
26
      await txn.delete(sandboxCleanupRetryKey(runId));
27
    } else {
28
      await txn.put(sandboxCleanupRetryKey(runId), value);
29
    }
30
 
31
    await rescheduleAlarmInTransaction(context, txn, projectId);
32
  });
33
};
34
 
35
export const seedSandboxCleanupRetry = async (
36
  context: ProjectDoContext,
37
  projectId: ProjectId,
38
  runId: RunId,
39
): Promise<void> => {
40
  await persistSandboxCleanupRetryState(context, projectId, runId, buildRetryState(1));
41
};
42
 
43
export const clearSandboxCleanupRetry = async (
44
  context: ProjectDoContext,
45
  projectId: ProjectId,
46
  runId: RunId,
47
): Promise<void> => {
48
  await persistSandboxCleanupRetryState(context, projectId, runId, null);
49
};
50
 
51
export const attemptSandboxCleanup = async (
52
  context: ProjectDoContext,
53
  projectId: ProjectId,
54
  runId: RunId,
55
  trigger: "alarm" | "watchdog",
56
): Promise<boolean> => {
57
  const sandboxStub = context.env.Sandbox.getByName(runId);
58
 
59
  try {
60
    await sandboxStub.setKeepAlive(false);
61
  } catch (error) {
62
    context.logger.warn("sandbox_keep_alive_release_failed", {
63
      projectId,
64
      runId,
65
      trigger,
66
      error: error instanceof Error ? error.message : String(error),
67
    });
68
  }
69
 
70
  try {
71
    await sandboxStub.destroy();
72
    return true;
73
  } catch (error) {
74
    if (isNoContainerInstanceError(error)) {
75
      return true;
76
    }
77
 
78
    context.logger.warn("sandbox_destroy_failed", {
79
      projectId,
80
      runId,
81
      trigger,
82
      error: error instanceof Error ? error.message : String(error),
83
    });
84
    return false;
85
  }
86
};
87
 
88
export const reconcileSandboxCleanup = async (
89
  context: ProjectDoContext,
90
  projectId: ProjectId,
91
): Promise<RunId | null> => {
92
  const rows = await context.db
93
    .select({
94
      runId: projectSchema.projectRuns.runId,
95
    })
96
    .from(projectSchema.projectRuns)
97
    .where(eq(projectSchema.projectRuns.projectId, projectId))
98
    .orderBy(asc(projectSchema.projectRuns.createdAt), asc(projectSchema.projectRuns.runId));
99
 
100
  let candidateRunId: RunId | null = null;
101
  let retryState: SandboxCleanupRetryState | null = null;
102
 
103
  for (const row of rows) {
104
    const runId = expectTrusted(RunId, row.runId, "RunId");
105
    const currentRetryState = await getSandboxCleanupRetryState(context, runId);
106
    if (!currentRetryState || currentRetryState.nextAt > Date.now()) {
107
      continue;
108
    }
109
 
110
    candidateRunId = runId;
111
    retryState = currentRetryState;
112
    break;
113
  }
114
 
115
  if (!candidateRunId || !retryState) {
116
    return null;
117
  }
118
 
119
  if (await attemptSandboxCleanup(context, projectId, candidateRunId, "alarm")) {
120
    await clearSandboxCleanupRetry(context, projectId, candidateRunId);
121
    return candidateRunId;
122
  }
123
 
124
  const nextRetryState = buildRetryState(retryState.attempt + 1);
125
  await persistSandboxCleanupRetryState(context, projectId, candidateRunId, nextRetryState);
126
  context.logger.warn("sandbox_cleanup_retry_scheduled", {
127
    projectId,
128
    runId: candidateRunId,
129
    attempt: nextRetryState.attempt,
130
    nextAt: nextRetryState.nextAt,
131
  });
132
  return null;
133
};