From 19d7579d8284d3f81761fbb3f569183e657b62f3 Mon Sep 17 00:00:00 2001 From: Muhammad Atif Ali Date: Thu, 24 Sep 2026 21:51:34 +0500 Subject: [PATCH] feat(agent-relay-cursor): drain the worker on workspace stop 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 worker was therefore killed outright: it never deregistered from its pool, its sessionEnd hook never ran, and work in flight was lost. This is the same gap the Claude Code module has, in the same place. Add a run_on_stop script that relays the signal and waits for the supervisor to record the exit -- not merely for the pid to disappear, which can happen before the supervisor resumes from `wait` and writes "done ", leaving a state file that reads as orphaned. Tighten the module's worker_alive guard while reusing it. Matching the bare basename accepted any command line containing "agent", including the workspace agent's own, which defeats the pid-reuse protection exactly in the recovery case it exists for; it now matches the argv the supervisor launches. The stop step is the one caller that signals, so the copies in start and status move with it. Export shutdown_grace_seconds for the template to wire into docker_container.destroy_grace_seconds or a pod's termination_grace_period_seconds, since a module owns no compute resource. Unlike the Claude Code runner, this one is an input rather than a computed value. The Cursor CLI advertises no shutdown budget and exposes no drain, retire or push-outcome flag -- verified against `agent worker --help` -- so there is nothing to derive it from. It defaults to 60 seconds and is the operator's to tune. --- .../modules/agent-relay-cursor/README.md | 34 ++++++++ .../coder/modules/agent-relay-cursor/main.tf | 40 ++++++++++ .../agent-relay-cursor/main.tftest.hcl | 79 +++++++++++++++++++ .../modules/agent-relay-cursor/start.sh.tftpl | 9 ++- .../agent-relay-cursor/status.sh.tftpl | 9 ++- .../modules/agent-relay-cursor/stop.sh.tftpl | 79 +++++++++++++++++++ 6 files changed, 244 insertions(+), 6 deletions(-) create mode 100644 registry/coder/modules/agent-relay-cursor/stop.sh.tftpl diff --git a/registry/coder/modules/agent-relay-cursor/README.md b/registry/coder/modules/agent-relay-cursor/README.md index 08f2c9a43..15daef412 100644 --- a/registry/coder/modules/agent-relay-cursor/README.md +++ b/registry/coder/modules/agent-relay-cursor/README.md @@ -140,3 +140,37 @@ disk: a live worker is left alone, a terminal state is left for the relay to act on, and anything else is reset to `pending` before a new worker starts. Liveness is judged by pid and cmdline, so a reused pid reads as `orphaned` rather than `working`. + +## 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 worker and waits for it to exit. It confirms the pid is +still the worker before signalling, sends SIGTERM only, 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.cursor_worker.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 worker +never deregisters from its pool and work in flight is lost. + +Unlike the Claude Code runner, the Cursor CLI advertises no shutdown +budget and exposes no drain flag, so `shutdown_grace_seconds` is not +derived from the CLI — it defaults to 60 and is yours to tune. It is a +ceiling rather than a fixed wait, so a workspace whose worker has already +exited still stops immediately. diff --git a/registry/coder/modules/agent-relay-cursor/main.tf b/registry/coder/modules/agent-relay-cursor/main.tf index 528f8589a..89033d34f 100644 --- a/registry/coder/modules/agent-relay-cursor/main.tf +++ b/registry/coder/modules/agent-relay-cursor/main.tf @@ -81,6 +81,21 @@ variable "log_file" { description = "Path the detached worker's output is written to." } +variable "shutdown_grace_seconds" { + description = <<-EOT + Seconds the stop script waits for the worker to exit after signalling it, and the value a template should grant the workspace so that wait can finish. + + Unlike the Claude Code runner, the Cursor CLI advertises no shutdown budget and exposes no drain flag, so this is not derived from the CLI -- it is a ceiling you tune. A workspace whose worker has already exited stops immediately either way. + EOT + type = number + default = 60 + + validation { + condition = var.shutdown_grace_seconds > 0 && floor(var.shutdown_grace_seconds) == var.shutdown_grace_seconds + error_message = "shutdown_grace_seconds must be a whole number of seconds greater than zero." + } +} + variable "serving_log_pattern" { type = string default = "in use" @@ -268,6 +283,12 @@ locals { install_cli = var.install_cli }) + stop_script = templatefile("${path.module}/stop.sh.tftpl", { + cli_binary = var.cli_binary + state_file = var.state_file + drain_timeout_s = var.shutdown_grace_seconds + }) + start_script = templatefile("${path.module}/start.sh.tftpl", { module_directory = local.module_directory cli_binary = var.cli_binary @@ -278,6 +299,20 @@ locals { }) } +# coder-utils runs install and start steps only, so the stop step is a +# plain coder_script beside it. It is deliberately outside the +# `coder exp sync` ordering the module's other scripts take part in: +# nothing runs after it. +resource "coder_script" "stop" { + agent_id = var.agent_id + display_name = "Cursor worker: Stop Script" + icon = "/icon/cursor.svg" + run_on_start = false + run_on_stop = true + start_blocks_login = false + script = local.stop_script +} + module "coder_utils" { source = "registry.coder.com/coder/coder-utils/coder" version = "0.0.1" @@ -341,3 +376,8 @@ resource "coder_app" "cursor_desktop" { url = "cursor://anysphere.cursor-deeplink/background-agent?bcId=${data.coder_parameter.agent_relay_session_id.value}" external = true } + +output "shutdown_grace_seconds" { + description = "Seconds the platform must give the workspace to shut down for the worker to drain instead of being killed. Wire it into the compute resource the template owns: docker_container.destroy_grace_seconds, or a pod spec's termination_grace_period_seconds. A module owns no compute resource and cannot set this itself. The value is a ceiling, not a fixed wait." + value = var.shutdown_grace_seconds +} diff --git a/registry/coder/modules/agent-relay-cursor/main.tftest.hcl b/registry/coder/modules/agent-relay-cursor/main.tftest.hcl index 354a86165..8e85ab88e 100644 --- a/registry/coder/modules/agent-relay-cursor/main.tftest.hcl +++ b/registry/coder/modules/agent-relay-cursor/main.tftest.hcl @@ -314,3 +314,82 @@ run "serving_log_pattern_is_data" { error_message = "serving_log_pattern must be base64-encoded and matched with grep -F" } } + +run "graceful_shutdown" { + command = plan + + # The supervisor runs under setsid, so the agent's own SIGTERM never + # reaches the worker. The stop script is the only thing that relays it. + assert { + condition = coder_script.stop.run_on_stop == true && coder_script.stop.run_on_start == false + error_message = "the stop script must run on stop and never on start" + } + + # SIGTERM only: the Cursor CLI exposes no drain flag, so escalating + # would just kill work the worker might still be finishing. + assert { + condition = strcontains(local.stop_script, "kill -TERM") && !strcontains(local.stop_script, "kill -9") && !strcontains(local.stop_script, "-KILL") + error_message = "the stop script must send SIGTERM only and never escalate" + } + + # A recycled pid would otherwise be signalled; this is the one caller + # that sends one. + assert { + condition = strcontains(local.stop_script, "worker_alive") && strcontains(local.stop_script, "cmdline") + error_message = "the stop script must confirm the pid is still our worker before signalling it" + } + + # The bare basename "agent" also matches the workspace agent's own + # command line, which is the recycled-pid case the guard exists for. + assert { + condition = strcontains(local.stop_script, ") worker\"") + error_message = "the pid guard must match the worker's argv, not just the binary name" + } + + # supervise.sh is the sole writer of terminal state. + assert { + condition = !strcontains(local.stop_script, "state_file.tmp") + error_message = "the stop script must not write the state file" + } + + # The supervisor is a separate process, so the pid can vanish before + # "done " is written; returning in that gap reads as orphaned. + assert { + condition = strcontains(local.stop_script, "while ! recorded") + error_message = "the stop script must wait for the supervisor to record the exit" + } + + assert { + condition = output.shutdown_grace_seconds == 60 && strcontains(local.stop_script, "budget=60") + error_message = "the advertised budget and the script's own wait must be the same number" + } + + # The stop step is not part of the coder exp sync chain. + assert { + condition = output.scripts == module.coder_utils.scripts + error_message = "the stop script must stay out of the sync ordering" + } +} + +run "shutdown_grace_seconds_overridden" { + command = plan + + variables { + shutdown_grace_seconds = 180 + } + + assert { + condition = output.shutdown_grace_seconds == 180 && strcontains(local.stop_script, "budget=180") + error_message = "an override must reach both the output and the script" + } +} + +run "shutdown_grace_seconds_rejects_zero" { + command = plan + + variables { + shutdown_grace_seconds = 0 + } + + expect_failures = [var.shutdown_grace_seconds] +} diff --git a/registry/coder/modules/agent-relay-cursor/start.sh.tftpl b/registry/coder/modules/agent-relay-cursor/start.sh.tftpl index 5e70a712f..3b8fb1594 100644 --- a/registry/coder/modules/agent-relay-cursor/start.sh.tftpl +++ b/registry/coder/modules/agent-relay-cursor/start.sh.tftpl @@ -26,11 +26,14 @@ write_state() { mv "$state_file.tmp" "$state_file" } -# True when the pid is alive and is our worker. kill -0 alone would -# trust a reused pid; the cmdline check ties it to the CLI we launched. +# True when the pid is alive and is our worker. kill -0 alone would trust +# a reused pid, and the bare basename is not enough either: "agent" also +# appears in the workspace agent's own command line. Match the argv the +# supervisor launches. worker_alive() { [ -n "$${1:-}" ] && kill -0 "$1" 2>/dev/null && - tr '\0' ' ' <"/proc/$1/cmdline" 2>/dev/null | grep -qF -- "$(basename "${cli_binary}")" + tr '\0' ' ' <"/proc/$1/cmdline" 2>/dev/null | + grep -qF -- "$(basename "${cli_binary}") worker" } if [ -z "$${AGENT_RELAY_CURSOR_TOKEN:-}" ]; then diff --git a/registry/coder/modules/agent-relay-cursor/status.sh.tftpl b/registry/coder/modules/agent-relay-cursor/status.sh.tftpl index 276f3ec20..dc4667079 100644 --- a/registry/coder/modules/agent-relay-cursor/status.sh.tftpl +++ b/registry/coder/modules/agent-relay-cursor/status.sh.tftpl @@ -26,11 +26,14 @@ log_file="${log_file}" # string: neither the shell nor grep interprets anything in it. serving_log_pattern="$(echo -n '${serving_log_pattern}' | base64 -d)" -# True when the pid is alive and is our worker. kill -0 alone would -# trust a reused pid; the cmdline check ties it to the CLI we launched. +# True when the pid is alive and is our worker. kill -0 alone would trust +# a reused pid, and the bare basename is not enough either: "agent" also +# appears in the workspace agent's own command line. Match the argv the +# supervisor launches. worker_alive() { [ -n "$${1:-}" ] && kill -0 "$1" 2>/dev/null && - tr '\0' ' ' <"/proc/$1/cmdline" 2>/dev/null | grep -qF -- "$(basename "${cli_binary}")" + tr '\0' ' ' <"/proc/$1/cmdline" 2>/dev/null | + grep -qF -- "$(basename "${cli_binary}") worker" } if [ ! -f "$state_file" ]; then diff --git a/registry/coder/modules/agent-relay-cursor/stop.sh.tftpl b/registry/coder/modules/agent-relay-cursor/stop.sh.tftpl new file mode 100644 index 000000000..627875b56 --- /dev/null +++ b/registry/coder/modules/agent-relay-cursor/stop.sh.tftpl @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Stop step, run by the agent when the workspace shuts down. +# +# The start step launches the supervisor with setsid, so it lives in its +# own session and never receives the SIGTERM the container's init gets. +# Without this script the worker is killed outright: it never deregisters +# from its pool, its sessionEnd hook never runs, and work in flight is +# lost. +# +# This script is only half the fix. The agent runs it with whatever time +# the platform grants, which on Docker is nothing at all unless the +# template sets destroy_grace_seconds. See the module README. +set -euo pipefail + +state_file="${state_file}" +budget=${drain_timeout_s} + +# True when the pid is alive and is our worker. kill -0 alone would trust +# a reused pid, and the bare basename is not enough either: "agent" also +# appears in the workspace agent's own command line. Match the argv the +# supervisor launches. +# It matters most here, because this is the caller that signals. +worker_alive() { + [ -n "$${1:-}" ] && kill -0 "$1" 2>/dev/null && + tr '\0' ' ' <"/proc/$1/cmdline" 2>/dev/null | + grep -qF -- "$(basename "${cli_binary}") worker" +} + +if [ ! -f "$state_file" ]; then + echo "No worker state at $state_file; nothing to drain." + exit 0 +fi + +# "working " is the only state with a live worker behind it. Every +# other value is terminal, or a workspace the relay never dispatched. +read -r state pid <"$state_file" || true +if [ "$${state:-}" != "working" ]; then + echo "Worker is not running (state: $${state:-empty}). Nothing to drain." + exit 0 +fi + +if ! worker_alive "$${pid:-}"; then + echo "Worker $${pid:-unknown} is already gone." + exit 0 +fi + +echo "Draining Cursor worker $pid (up to $${budget}s)..." +# SIGTERM only. Escalating would defeat the point, and the platform +# SIGKILLs us soon enough anyway. +kill -TERM "$pid" 2>/dev/null || true + +# Wait for the supervisor to record the exit, not just for the worker to +# go. The supervisor is a separate process: /proc/ can disappear +# before it resumes from `wait` and writes "done ". Returning in +# that gap leaves the state file saying "working" with a dead pid, which +# the status script reports as orphaned -- a clean drain graded as a +# failure. +recorded() { + read -r s _ <"$state_file" 2>/dev/null || return 1 + [ "$${s:-working}" != working ] +} + +SECONDS=0 +while ! recorded && [ "$SECONDS" -lt "$budget" ]; do + sleep 0.2 +done + +if recorded; then + echo "Worker exited after $${SECONDS}s. State: $(cat "$state_file" 2>/dev/null)" +elif worker_alive "$pid"; then + echo "Worker did not exit within $${budget}s; leaving it to the platform." >&2 +else + echo "Worker is gone but the supervisor never recorded it within $${budget}s." >&2 +fi + +# Always succeed. The supervisor is the only writer of terminal state, and +# a nonzero stop script only buys a shutdown_error lifecycle on a workspace +# that is about to disappear. +exit 0