1
import { vi } from "vitest";
2
 
3
import {
4
  BranchName,
5
  DEFAULT_DISPATCH_MODE,
6
  DEFAULT_EXECUTION_RUNTIME,
7
  ProjectId,
8
  RunId,
9
  UnixTimestampMs,
10
  type CommitSha as CommitShaType,
11
} from "@/contracts";
12
import { expectTrusted } from "@/worker/contracts";
13
import type { ProjectExecutionMaterial } from "@/worker/durable/project-do/types";
14
import type {
15
  ProjectControl,
16
  RunControl,
17
  RunExecutionContextState,
18
  RunExecutionScope,
19
  RunLogs,
20
  RunRuntime,
21
  RunStore,
22
} from "@/worker/dispatch/shared/run-execution-context/types";
23
import type { RunLeaseControl } from "@/worker/dispatch/shared/run-lease";
24
 
25
const DEFAULT_PROJECT_ID = ProjectId.assertDecode("prj_0000000000000000000000");
26
const DEFAULT_RUN_ID = RunId.assertDecode("run_0000000000000000000000");
27
const DEFAULT_BRANCH = BranchName.assertDecode("main");
28
const DEFAULT_STARTED_AT = expectTrusted(UnixTimestampMs, 1_740_000_000_000, "UnixTimestampMs");
29
const DEFAULT_REPO_ROOT = "/workspace/repo";
30
const DEFAULT_REPO_URL = "https://github.com/example/anvil-spec";
31
const DEFAULT_CONFIG_PATH = ".anvil.yml";
32
 
33
const createUnusedSyncStub = (label = "unused") =>
34
  vi.fn(() => {
35
    throw new Error(label);
36
  });
37
 
38
const createUnusedAsyncStub = (label = "unused") =>
39
  vi.fn(async () => {
40
    throw new Error(label);
41
  });
42
 
43
interface QueueExecutionMaterialOptions extends Partial<ProjectExecutionMaterial> {
44
  projectId?: ProjectExecutionMaterial["projectId"];
45
}
46
 
47
interface QueueScopeOptions {
48
  env?: Env;
49
  projectId?: RunExecutionScope["projectId"];
50
  runId?: RunExecutionScope["runId"];
51
  startedAt?: RunExecutionScope["startedAt"];
52
  snapshotCommitSha?: CommitShaType | null;
53
  executionMaterial?: ProjectExecutionMaterial;
54
  repoRoot?: RunExecutionScope["repoRoot"];
55
  repoUrl?: string;
56
  branch?: string;
57
  configPath?: string;
58
}
59
 
60
interface QueueLeaseStubOptions {
61
  stop?: RunLeaseControl["stop"];
62
  throwIfOwnershipLost?: RunLeaseControl["throwIfOwnershipLost"];
63
  refreshControl?: RunLeaseControl["refreshControl"];
64
  isCancellationRequested?: boolean;
65
  applyCancellationIfNeeded?: RunLeaseControl["applyCancellationIfNeeded"];
66
}
67
 
68
export const createQueueExecutionMaterial = (
69
  options: QueueExecutionMaterialOptions = {},
70
): ProjectExecutionMaterial => ({
71
  projectId: options.projectId ?? DEFAULT_PROJECT_ID,
72
  encryptedRepoToken: options.encryptedRepoToken ?? null,
73
});
74
 
75
export const createQueueScope = (options: QueueScopeOptions = {}): RunExecutionScope => {
76
  const projectId =
77
    options.projectId ??
78
    (options.executionMaterial
79
      ? expectTrusted(ProjectId, options.executionMaterial.projectId, "ProjectId")
80
      : DEFAULT_PROJECT_ID);
81
  const runId = options.runId ?? DEFAULT_RUN_ID;
82
  const startedAt = options.startedAt ?? DEFAULT_STARTED_AT;
83
  const repoUrl = options.repoUrl ?? DEFAULT_REPO_URL;
84
  const branch = options.branch ? expectTrusted(BranchName, options.branch, "BranchName") : DEFAULT_BRANCH;
85
  const configPath = options.configPath ?? DEFAULT_CONFIG_PATH;
86
  const snapshot = {
87
    projectId,
88
    runId,
89
    triggerType: "manual" as const,
90
    triggeredByUserId: null,
91
    repoUrl,
92
    branch,
93
    commitSha: options.snapshotCommitSha ?? null,
94
    configPath,
95
    dispatchMode: DEFAULT_DISPATCH_MODE,
96
    executionRuntime: DEFAULT_EXECUTION_RUNTIME,
97
    queuedAt: startedAt,
98
  };
99
  const executionMaterial =
100
    options.executionMaterial ??
101
    createQueueExecutionMaterial({
102
      projectId,
103
    });
104
 
105
  return {
106
    env: options.env ?? ({} as Env),
107
    executionMaterial,
108
    claim: {
109
      kind: "execute",
110
      snapshot,
111
    },
112
    snapshot,
113
    projectId,
114
    runId,
115
    repoRoot: options.repoRoot ?? DEFAULT_REPO_ROOT,
116
    startedAt,
117
    logContext: {
118
      projectId,
119
      runId,
120
    },
121
  };
122
};
123
 
124
export const createQueueState = (overrides: Partial<RunExecutionContextState> = {}): RunExecutionContextState => ({
125
  phase: "booting",
126
  session: null,
127
  currentProcess: null,
128
  currentStepPosition: null,
129
  cancelRequestedAt: null,
130
  ownershipLost: false,
131
  ownershipLossStatus: null,
132
  softCancelIssued: false,
133
  hardCancelIssued: false,
134
  preservedTerminalStatus: null,
135
  redactionSecrets: [],
136
  ...overrides,
137
});
138
 
139
export const createQueueLeaseStub = (options: QueueLeaseStubOptions = {}): RunLeaseControl => ({
140
  stop: options.stop ?? vi.fn(async () => {}),
141
  throwIfOwnershipLost: options.throwIfOwnershipLost ?? vi.fn(() => {}),
142
  refreshControl: options.refreshControl ?? vi.fn(async () => {}),
143
  isCancellationRequested: () => options.isCancellationRequested ?? false,
144
  applyCancellationIfNeeded: options.applyCancellationIfNeeded ?? vi.fn(async () => {}),
145
});
146
 
147
export const createQueueRunStoreStub = (overrides: Partial<RunStore> = {}): RunStore => ({
148
  getFreshStub: createUnusedSyncStub(),
149
  getMeta: createUnusedAsyncStub(),
150
  updateState: vi.fn(async () => {}),
151
  tryUpdateState: createUnusedAsyncStub(),
152
  repairTerminalState: vi.fn(async () => {}),
153
  replaceSteps: vi.fn(async () => {}),
154
  updateStepState: vi.fn(async () => {}),
155
  appendLogs: vi.fn(async () => {}),
156
  ...overrides,
157
});
158
 
159
export const createQueueProjectControlStub = (overrides: Partial<ProjectControl> = {}): ProjectControl => ({
160
  getFreshStub: createUnusedSyncStub(),
161
  recordHeartbeat: vi.fn(async () => null),
162
  recordResolvedCommit: createUnusedAsyncStub(),
163
  finalizeRunExecution: vi.fn(async (_terminalStatus, _lastError, _sandboxDestroyed) => {
164
    throw new Error("unused");
165
  }),
166
  kickReconciliation: vi.fn(async (_trigger: string) => {}),
167
  ...overrides,
168
});
169
 
170
export const createQueueRunControlStub = (overrides: Partial<RunControl> = {}): RunControl => ({
171
  getRunMeta: createUnusedAsyncStub(),
172
  updateRunFromCurrent: createUnusedAsyncStub(),
173
  preserveTerminalOutcome: vi.fn(),
174
  ensureRunCancelRequested: vi.fn(async () => null),
175
  ensureRunCanceling: vi.fn(async () => null),
176
  markOwnershipLost: vi.fn(),
177
  ...overrides,
178
});
179
 
180
export const createQueueRunLogsStub = (overrides: Partial<RunLogs> = {}): RunLogs => ({
181
  appendSystemLog: vi.fn(async () => {}),
182
  redactMessage: vi.fn((message: string) => message),
183
  ...overrides,
184
});
185
 
186
export const createQueueRunRuntimeStub = (overrides: Partial<RunRuntime> = {}): RunRuntime => ({
187
  sandbox: {} as RunRuntime["sandbox"],
188
  getSession: createUnusedAsyncStub(),
189
  deleteSession: vi.fn(async () => {}),
190
  disposeSession: vi.fn((_session: Parameters<RunRuntime["disposeSession"]>[0]) => {}),
191
  getLiveCurrentProcess: vi.fn(() => null),
192
  isProcessTreeAlive: vi.fn(async () => false),
193
  softCancelProcessTree: vi.fn(async () => {}),
194
  hardCancelProcessTree: vi.fn(async () => {}),
195
  waitForProcessTreeToStop: vi.fn(async () => true),
196
  waitForProcessTreeToStopSafely: vi.fn(async () => true),
197
  destroySandbox: vi.fn(async () => true),
198
  dispose: vi.fn(() => {}),
199
  ...overrides,
200
});