1
import { afterEach, describe, expect, it, vi } from "vitest";
2
 
3
import { RunId, UnixTimestampMs } from "@/contracts";
4
import { expectTrusted } from "@/worker/contracts";
5
import { CANCEL_GRACE_MS } from "@/worker/dispatch/shared/run-execution-context";
6
import { RunLease } from "@/worker/dispatch/shared/run-lease";
7
 
8
import {
9
  createQueueProjectControlStub,
10
  createQueueRunControlStub,
11
  createQueueRunLogsStub,
12
  createQueueRunRuntimeStub,
13
  createQueueScope,
14
  createQueueState,
15
} from "../../../helpers/dispatch/shared";
16
 
17
type LeaseContext = ConstructorParameters<typeof RunLease>[0];
18
 
19
const makeContext = (): LeaseContext => ({
20
  scope: createQueueScope(),
21
  state: createQueueState(),
22
  projectControl: createQueueProjectControlStub(),
23
  control: createQueueRunControlStub(),
24
  logs: createQueueRunLogsStub(),
25
  runtime: createQueueRunRuntimeStub({
26
    getLiveCurrentProcess: vi.fn(() => ({ id: "proc-1", status: "running" }) as never),
27
    isProcessTreeAlive: vi.fn(async () => true),
28
  }),
29
});
30
 
31
afterEach(() => {
32
  vi.restoreAllMocks();
33
});
34
 
35
describe("run lease", () => {
36
  it("records cancellation requests from heartbeat control", async () => {
37
    const context = makeContext();
38
    context.projectControl.recordHeartbeat = vi.fn(async () => ({
39
      runId: RunId.assertDecode("run_0000000000000000000000"),
40
      status: "cancel_requested" as const,
41
      cancelRequestedAt: expectTrusted(UnixTimestampMs, 1_740_000_000_000, "UnixTimestampMs"),
42
    }));
43
 
44
    const lease = new RunLease(context);
45
    await lease.refreshControl();
46
 
47
    expect(context.state.cancelRequestedAt).toBe(1_740_000_000_000);
48
  });
49
 
50
  it("soft-cancels the active session on the first cancellation pass", async () => {
51
    const context = makeContext();
52
    context.state.cancelRequestedAt = expectTrusted(UnixTimestampMs, 1_740_000_000_000, "UnixTimestampMs");
53
    context.state.session = { id: "session-1" } as never;
54
 
55
    const lease = new RunLease(context);
56
    await lease.applyCancellationIfNeeded();
57
 
58
    expect(context.state.phase).toBe("canceling");
59
    expect(context.state.softCancelIssued).toBe(true);
60
    expect(context.control.ensureRunCancelRequested).toHaveBeenCalledTimes(1);
61
    expect(context.control.ensureRunCanceling).toHaveBeenCalledTimes(1);
62
    expect(context.logs.appendSystemLog).toHaveBeenCalledWith("Cancellation requested. Sending SIGTERM.");
63
    expect(context.runtime.softCancelProcessTree).toHaveBeenCalledWith(
64
      context.state.session,
65
      expect.objectContaining({ id: "proc-1" }),
66
    );
67
  });
68
 
69
  it("escalates to a hard cancel after the grace window expires", async () => {
70
    vi.spyOn(console, "log").mockImplementation(() => {});
71
    const context = makeContext();
72
    context.state.cancelRequestedAt = expectTrusted(
73
      UnixTimestampMs,
74
      Date.now() - CANCEL_GRACE_MS - 1,
75
      "UnixTimestampMs",
76
    );
77
    context.state.session = { id: "session-1" } as never;
78
    context.state.softCancelIssued = true;
79
 
80
    const lease = new RunLease(context);
81
    await lease.applyCancellationIfNeeded();
82
 
83
    expect(context.state.hardCancelIssued).toBe(true);
84
    expect(context.logs.appendSystemLog).toHaveBeenCalledWith("Cancellation grace window expired. Sending SIGKILL.");
85
    expect(context.runtime.hardCancelProcessTree).toHaveBeenCalledWith(
86
      context.state.session,
87
      expect.objectContaining({ id: "proc-1" }),
88
    );
89
  });
90
 
91
  it("soft-cancels the active process tree when canceling without a live process handle", async () => {
92
    const context = makeContext();
93
    context.runtime.getLiveCurrentProcess = vi.fn(() => null);
94
    context.state.cancelRequestedAt = expectTrusted(UnixTimestampMs, 1_740_000_000_000, "UnixTimestampMs");
95
    context.state.session = { id: "session-1" } as never;
96
 
97
    const lease = new RunLease(context);
98
    await lease.applyCancellationIfNeeded();
99
 
100
    expect(context.logs.appendSystemLog).toHaveBeenCalledWith("Cancellation requested. Stopping active process tree.");
101
    expect(context.runtime.softCancelProcessTree).toHaveBeenCalledWith(context.state.session, null);
102
    expect(context.runtime.deleteSession).not.toHaveBeenCalled();
103
  });
104
});