Skip to content

fix(coderd): allow agent auth during workspace shutdown - #21538

Merged
mafredri merged 4 commits into
mainfrom
mafredri/fix-workspace-agent-shutdown-authorizaiton
Jan 21, 2026
Merged

mafredri merged 4 commits into
mainfrom
mafredri/fix-workspace-agent-shutdown-authorizaiton

Conversation

@mafredri

@mafredri mafredri commented Jan 16, 2026 •

Copy link
Copy Markdown
Member

Agents were losing authentication during workspace shutdown, causing
shutdown scripts to fail. The auth query required agents to belong to
the latest build, but during shutdown a stop build becomes latest while
the start build's agents are still running.

Modified the auth query to allow start build agents to authenticate
temporarily during stop execution. The query allows auth when:

  • Agent's start build job succeeded
  • Latest build is stop with pending/running job status
  • Builds are adjacent (stop is build_number + 1)
  • Template versions match

Auth closes once stop completes.

Renamed GetWorkspaceAgentAndLatestBuildByAuthToken to
GetAuthenticatedWorkspaceAgentAndBuildByAuthToken since it returns the
agent's build (not always latest) during shutdown.

Added tests covering happy path, failure cases (completed stops, failed
builds, version mismatches), and edge cases (multiple cycles, wrong
transitions) at both database and middleware levels.

Closes coder/internal#1249
Fixes #19467


Alternative approaches were considered:

  1. Token reuse: Have stop builds reuse agent tokens from start builds.
    However, auth_token lacks unique constraint, so query would return
    multiple agents and break authentication.
  2. Separate validation functions: Split into GetWorkspaceAgentByAuthToken
    (validates token) and GetLatestWorkspaceBuildForWorkspace (gets actual
    latest). (Benefit unclear and requires test refactoring.)

@mafredri
mafredri force-pushed the mafredri/fix-workspace-agent-shutdown-authorizaiton branch from 8e38e41 to f464c6b Compare January 16, 2026 13:06
@mafredri
mafredri marked this pull request as ready for review January 16, 2026 13:19
@mafredri mafredri changed the title feat(coderd): allow agent auth during workspace shutdown fix(coderd): allow agent auth during workspace shutdown Jan 16, 2026
WHERE
workspace_id = workspace_build_with_user.workspace_id
-- Filter out builds that are not the latest, with exception for shutdown case.
AND (

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this benefit from using a CASE statement instead so we can do some form of short-circuiting for the normal case?

ref: https://www.postgresql.org/docs/current/functions-conditional.html#FUNCTIONS-CASE

If no WHEN condition yields true, the value of the CASE expression is the result of the ELSE clause.

@mafredri mafredri Jan 16, 2026 •

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Excellent call-out, I forgot that OR/AND isn't necessarily short-circuited. Definitely doing this 👍🏻

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WHERE l2.workspace_id = latest.workspace_id
)
AND latest.transition = 'stop'
AND latest.template_version_id = workspace_build_with_user.template_version_id

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand the purpose of the version check. Let's say we have a stop build for the workspace at a new version of the template: why does that mean we should refuse to authenticate the pervious agent that is in the process of being shut down?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a defensive check against an unlikely scenario to prevent RBAC mismatches if start and stop use different template versions.

But upon further inspection, agents don't access template resources during shutdown, so the template version in the allowlist is never actually checked.

Removing it.

@spikecurtis

Copy link
Copy Markdown
Contributor

I think it's OK to ship as a stop gap, but this is an incomplete solution.

  1. We rely on some combination of Terraform and the infrastructure provider to send a signal to the Workspace Agent. This works for Kubernetes where Pods are typically signaled and given a grace period before being terminated, but the story for VMs is much dicier. We also support externally provisioned workspaces where we don't touch the underlying compute at all.

  2. The infrastructure might time out the shutdown before the shutdown scripts complete. Sometimes the template author will have access/permission to extend the timeout, but they might not.

  3. Terraform's destroy process for the compute resource might not wait for the workspace agent to finish shut down scripts. If it doesn't, then the STOP build could complete before the script finishes, and then fail to authenticate.

These issues suggest (to me) that we should consider moving to an architecture where Coderd first signals to all (connected) workspace agents that they should start shutting down. This means running scripts, but also terminating SSH and Reconnecting PTY connections, tearing down tailnet tunnels, etc. Then, the agent can signal Coderd it is finished. When all agents are finished, disconnected, or hit some configurable timeout, then we actually kick off the Terraform build.

@mafredri

Copy link
Copy Markdown
Member Author

Great observation @spikecurtis, and I fully agree. We have an issue tracking this, #6175, though we don't have capacity for it right now. I'll update it with the extra scenarios you described, particularly externally provisioned workspaces and permission constraints, as those seem valuable to capture.

This PR just fixes the immediate auth bug. For tasks, we're documenting provider configuration (coder/internal#1271, #6174). It wasn't considered before, but we may also want document, or at least mention, infra-based timeouts there.

Agents were losing authentication during workspace shutdown, causing
shutdown scripts to fail. The auth query required agents to belong to
the latest build, but during shutdown a `stop` build becomes latest while
the `start` build's agents are still running.

Several approaches were considered:

- Reusing agent tokens (no unique constraint, would break auth)
- Separate validation functions (two queries for agent + latest build)
- Grace period with workspace-level tokens (security risk, stale agents)
- Agent lifecycle state tracking (requires migration, complex state management)

Modified the auth query to allow `start` build agents to authenticate
temporarily during `stop` execution. This is simpler than separate functions
and avoids migrations. The query allows auth when:

- Agent's `start` build job succeeded
- Latest build is `stop` with `pending`/`running` job status
- Builds are adjacent (`stop` is `build_number + 1`)
- Template versions match

The version check prevents RBAC issues if someone explicitly changes
template versions between builds. Auth closes once `stop` completes.

Renamed `GetWorkspaceAgentAndLatestBuildByAuthToken` to
`GetAuthenticatedWorkspaceAgentAndBuildByAuthToken` since it returns the
agent's build (not always latest) during shutdown.

Added tests covering happy path, failure cases (completed stops, failed
builds, version mismatches), and edge cases (multiple cycles, wrong
transitions) at both database and middleware levels.

Fixes coder/internal#1249
Fixes #19467
@mafredri
mafredri force-pushed the mafredri/fix-workspace-agent-shutdown-authorizaiton branch from 9b1c5a9 to fb9472b Compare January 21, 2026 12:32
@mafredri
mafredri enabled auto-merge (squash) January 21, 2026 12:36
@mafredri
mafredri merged commit 97e8a5b into main Jan 21, 2026
25 checks passed
@mafredri
mafredri deleted the mafredri/fix-workspace-agent-shutdown-authorizaiton branch January 21, 2026 13:18
@github-actions github-actions Bot locked and limited conversation to collaborators Jan 21, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Fix agent auth during shutdown scripts (#19467) Bug: Agent loses authorization before run_on_stop script completes

3 participants