feat(agent-relay-claude-code): drain the runner on workspace stop

registry/coder/modules/agent-relay-claude-code/README.md

@@ -16,7 +16,7 @@ on each build and runs the Claude Code self-hosted runner.
 ```tf
 module "claude_code_runner" {
   source   = "registry.coder.com/coder/agent-relay-claude-code/coder"
-  version  = "0.1.1"
+  version  = "0.2.0"
   agent_id = coder_agent.main.id
 
   # Downloads the Claude Code CLI at start when it is not in the image. Bake
@@ -36,6 +36,13 @@ resource "coder_agent" "main" {
 }
 ```
 
+Each runner registers with a label the Anthropic console shows beside it,
+defaulting to `<owner>/<workspace>` so it is identifiable whether or not the
+template sets a hostname. `client_label` overrides it. The label is display
+only: it never affects authorization, and it cannot steer which sessions a
+runner is assigned — routing is per environment, and one pool is one
+environment.
+
 The `agent_relay_status` metadata block is required. It has to live on the
 `coder_agent`, which the module cannot declare; the relay reads it to decide
 when to reap the workspace.
@@ -47,6 +54,9 @@ when to reap the workspace.
   image for the fastest start. `cli_binary` overrides the path.
 - Builds must finish inside Agent Relay's 300s spawn budget: pre-pulled
   images, no persistent volumes.
+- The compute resource must give the workspace time to shut down. Wire
+  `shutdown_grace_seconds` into it; without that the runner is killed
+  mid-session. See [Graceful shutdown](#graceful-shutdown).
 
 ## Parameters
 
@@ -70,9 +80,50 @@ runner. Everything lands under `$HOME/.coder-modules/coder/agent-relay-claude-co
 | `supervise.sh` | the detached supervisor that owns the runner    |
 | `wrapper.sh`   | session wrapper that forces `bypassPermissions` |
 
+The stop step is a plain `coder_script` rather than a coder-utils step, so its
+output goes to the agent's own script log, not to `logs/` above.
+
+## Environment
+
+The module always passes `--capacity`, `--base-dir`, `--exec-path` and
+`--client-label`, passes `--exit-if-unused-min` and
+`--push-outcome-on-release` unless you turn them off, passes
+`--drain-wait-sec` when you set it, and sets
+`SELF_HOSTED_RUNNER_ENVIRONMENT_SECRET` and
+`SELF_HOSTED_RUNNER_LOCK_TO_ACCOUNT` from the parameters the relay stamps.
+Anything else the CLI accepts can be set by the template, because the
+supervisor inherits the agent's environment:
+
+```tf
+resource "coder_env" "hooks_dir" {
+  agent_id = coder_agent.main.id
+  name     = "SELF_HOSTED_RUNNER_HOOKS_DIR"
+  value    = "/etc/claude-hooks"
+}
+```
+
+Three things to know before relying on that:
+
+- **A flag beats its paired environment variable.** Setting
+  `SELF_HOSTED_RUNNER_BASE_DIR` against the module's own `--base-dir` does
+  nothing; use the `base_dir` input instead. The same applies to every flag
+  the module emits, including the optional ones once you enable them — so
+  `SELF_HOSTED_RUNNER_CLIENT_LABEL` has no effect and `client_label` is the
+  only way to change the label.
+- **Duration environment variables are milliseconds**, while the CLI flags
+  they pair with are seconds or minutes, and the names do not always mirror:
+  `--exit-if-unused-min` pairs with `SELF_HOSTED_RUNNER_IDLE_SHUTDOWN_MS`.
+- **Not every flag has an environment variable.** `--capacity` is one, and
+  it is reserved anyway: Agent Relay's one-workspace-per-session model
+  depends on it being 1.
+
+Run `claude self-hosted-runner --help` for what your CLI actually accepts;
+the pairings above are the CLI's contract, not this module's, and move with
+it.
+
 ## Runner lifecycle
 
-The start step launches `claude self-hosted-runner --capacity 1 --exit-if-unused-min 10` detached and exits,
+The start step launches `claude self-hosted-runner` detached and exits,
 so the agent reaches `ready` immediately. The runner is wrapped so every
 session runs with `--permission-mode bypassPermissions`; there is no terminal
 attached, so an approval prompt would hang it. When the runner exits, Agent
@@ -94,3 +145,62 @@ rather than reporting `working` forever.
 | `failed <reason>` | runner could not start, e.g. `runner-agent-missing` when the CLI is absent |
 
 Renaming the `agent_relay_status` key breaks reaping.
+
+## Graceful shutdown
+
+The start step detaches the supervisor with `setsid`, so it lives in its own
+session and never receives the SIGTERM the container's init gets on shutdown.
+The module therefore registers a stop script that relays the signal to the
+runner and waits for the supervisor to record `done <code>`. It sends SIGTERM
+only, never escalates, and never writes the state file.
+
+That half only works if the platform gives the workspace time to use it, which
+the module cannot arrange: it owns no compute resource. Wire the exported
+budget into the one the template owns.
+
+```tf
+resource "docker_container" "workspace" {
+  # ...
+  destroy_grace_seconds = module.claude_code_runner.shutdown_grace_seconds
+}
+```
+
+On Kubernetes the equivalent is the pod spec's
+`termination_grace_period_seconds`.
+
+**Skip it and the drain never happens.** The Docker provider destroys the
+container with a zero stop timeout unless `destroy_grace_seconds` is set, so
+the container is killed before the stop script can finish: the session is
+never released server-side, the post-session hook never runs, and in-flight
+commits are lost. Kubernetes defaults to 30s, which is below the runner's own
+budget. The value is a ceiling rather than a fixed wait, so a workspace whose
+runner has already finished still stops immediately.
+
+The budget is the runner's advertised 80s to stop the Claude process and run
+the post-session hook, plus 20s for a session release already in flight, plus
+the 5s the agent spends shutting down SSH first:
+
+It comes to 105 seconds by default. Turning on the outcome push adds its 30
+seconds, and every second of `drain_wait_sec` is added on top — so both
+together at `drain_wait_sec = 60` make it 195.
+
+`push_outcome_on_release` pushes the session's outcome branch to `origin`
+before the branch is deleted, so commits survive an ephemeral workspace and a
+resumed session continues from them. Without it, work an incomplete session had
+already committed dies with the container.
+
+It is off by default, so upgrading changes nothing until you ask for it. Turn
+it on deliberately: it fires on every runner-initiated incomplete end, not only
+a drain — idle-release and failed sessions push too — so the workspace needs
+git auth and those sessions leave branches behind.
+
+`drain_wait_sec` is off by default. It buys "the current turn may finish", at
+a second of grace period per second of wait, which is the most expensive part
+of the budget.
+
+The budget only counts what the module passes. If you reach for the
+environment escape hatch instead -- `SELF_HOSTED_RUNNER_DRAIN_WAIT_MS`, or
+`SELF_HOSTED_RUNNER_PUSH_OUTCOME_ON_RELEASE` with the input left off -- the
+runner will spend time this number does not know about, and the platform can
+kill it mid-drain. Use the inputs, or add the difference to the grace period
+yourself.

registry/coder/modules/agent-relay-claude-code/main.test.ts

@@ -8,6 +8,7 @@ import {
 } from "bun:test";
 import {
   execContainer,
+  findResourceInstance,
   readFileContainer,
   removeContainer,
   runContainer,
@@ -65,7 +66,7 @@ const setup = async (vars: Record<string, string> = {}) => {
   return { id, scripts, statusScript };
 };
 
-type Scripts = { install: string; start: string };
+type Scripts = { install: string; start: string; stop: string };
 
 // coder-utils owns the coder_script resources; find ours by display name.
 const collectScripts = (state: TerraformState): Scripts => {
@@ -84,7 +85,9 @@ const collectScripts = (state: TerraformState): Scripts => {
       `expected install and start scripts, found ${Object.keys(byDisplayName)}`,
     );
   }
-  return { install, start };
+  // Ours, not coder-utils', so it is addressable by resource name.
+  const stop = findResourceInstance(state, "coder_script", "stop").script;
+  return { install, start, stop };
 };
 
 const stubBinary = async (id: string, path: string, body: string) => {
@@ -123,6 +126,11 @@ const runScripts = async (id: string, scripts: Scripts, env: string[]) => {
 const runDispatched = (id: string, scripts: Scripts) =>
   runScripts(id, scripts, DISPATCH_ENV);
 
+// The agent runs the stop step on workspace shutdown, with no dispatch
+// env in scope: the supervisor already holds the credential.
+const runStop = (id: string, scripts: Scripts) =>
+  execContainer(id, ["bash", "-c", scripts.stop]);
+
 const readState = async (id: string) =>
   (await readFileContainer(id, STATE_FILE)).trim();
 
@@ -258,6 +266,11 @@ describe("agent-relay-claude-code", () => {
     // The CLI's own default is /workspace, which the agent user cannot
     // create; the module points it at a directory it made.
     expect(args[args.indexOf("--base-dir") + 1]).toBe("/root/workspace");
+    // Both optional runner behaviours are off unless asked for.
+    expect(args).not.toContain("--push-outcome-on-release");
+    // A base64 decode that silently produced nothing would otherwise look
+    // like a runner that simply registered with an empty label.
+    expect(args[args.indexOf("--client-label") + 1]).toBe("default/default");
     const baseDir = await execContainer(id, ["test", "-d", "/root/workspace"]);
     expect(baseDir.exitCode).toBe(0);
 
@@ -361,4 +374,82 @@ describe("agent-relay-claude-code", () => {
     const pwned = await execContainer(id, ["test", "-e", "/tmp/PWNED"]);
     expect(pwned.exitCode).not.toBe(0);
   });
+  // The supervisor is setsid'd into its own session, so the SIGTERM the
+  // container's init receives never reaches the runner. Without the stop
+  // step the runner is killed outright and the session is never released.
+  it("relays SIGTERM to the detached runner and lets the supervisor record the exit", async () => {
+    const { id, scripts } = await setup();
+    // `& wait` is load-bearing: bash defers a trap until a foreground
+    // command returns, so a foreground sleep would hang this for 300s.
+    await stubClaude(id, "trap 'sleep 2; exit 7' TERM\nsleep 300 & wait");
+    await runDispatched(id, scripts);
+    await waitForState(id, /^working \d+$/);
+
+    const startedAt = Date.now();
+    const stop = await runStop(id, scripts);
+    const elapsedMs = Date.now() - startedAt;
+
+    expect(stop.exitCode).toBe(0);
+    expect(stop.stdout).toContain("Draining Claude Code runner");
+    // The runner's own exit code, recorded by the supervisor's wait.
+    expect(await waitForState(id, /^done 7$/)).toBe("done 7");
+    // Proves we waited for the drain rather than firing and forgetting.
+    expect(elapsedMs).toBeGreaterThanOrEqual(2000);
+  });
+
+  it("no-ops when the workspace was never dispatched", async () => {
+    const { id, scripts } = await setup();
+    await stubClaude(id, "sleep 300");
+    // No credential: the start step records idle and exits.
+    await runScripts(id, scripts, []);
+    expect(await readState(id)).toBe("idle");
+
+    const stop = await runStop(id, scripts);
+    expect(stop.exitCode).toBe(0);
+    expect(stop.stdout).toContain("Nothing to drain");
+    expect(await readState(id)).toBe("idle");
+  });
+
+  // Three states with no live runner behind them. They share a container:
+  // each only writes the state file and runs the stop step, and a stale
+  // pid is safe to signal only because the container has its own PID
+  // namespace.
+  it("no-ops when there is no runner to drain", async () => {
+    const { id, scripts } = await setup();
+
+    let stop = await runStop(id, scripts);
+    expect(stop.exitCode).toBe(0);
+    expect(stop.stdout).toContain("No runner state");
+
+    await execContainer(id, ["mkdir", "-p", MODULE_DIR]);
+    await writeFileContainer(id, STATE_FILE, "done 3\n", { user: "root" });
+    stop = await runStop(id, scripts);
+    expect(stop.exitCode).toBe(0);
+    expect(stop.stdout).toContain("Nothing to drain");
+    // A terminal state the reaper already grades must survive untouched.
+    expect(await readState(id)).toBe("done 3");
+
+    await writeFileContainer(id, STATE_FILE, "working 999999\n", {
+      user: "root",
+    });
+    stop = await runStop(id, scripts);
+    expect(stop.exitCode).toBe(0);
+    expect(stop.stdout).toContain("already gone");
+  });
+
+  it("passes the optional runner flags when set", async () => {
+    const { id, scripts } = await setup({
+      drain_wait_sec: "30",
+      push_outcome_on_release: "true",
+    });
+    await stubClaude(id, 'printf "%s\\n" "$@" >/tmp/claude-args\nsleep 300');
+    await runDispatched(id, scripts);
+    await waitForState(id, /^working \d+$/);
+
+    const args = (await readFileContainer(id, "/tmp/claude-args"))
+      .trim()
+      .split("\n");
+    expect(args[args.indexOf("--drain-wait-sec") + 1]).toBe("30");
+    expect(args).toContain("--push-outcome-on-release");
+  });
 });