Skip to content

fix(oauth): enforce max_agents on auto-provisioned agents - #952

Open
AmirF194 wants to merge 4 commits into
tokencanopy:mainfrom
AmirF194:fix/oauth-autocreate-agent-cap
Open

fix(oauth): enforce max_agents on auto-provisioned agents#952
AmirF194 wants to merge 4 commits into
tokencanopy:mainfrom
AmirF194:fix/oauth-autocreate-agent-cap

Conversation

@AmirF194

Copy link
Copy Markdown
Contributor

Summary

Filed as #951: issueOAuthCodeWithNewAgent (the OAuth auto-provision path) creates an agent via CreateAgentTx with no max_agents check anywhere in the file, unlike POST /v1/agents, which has always called EnforceAgentCreate before inserting (even before #942's atomicity fix). This wires in the same check, right before the transaction that creates the agent.

Client surface checklist

Not applicable: no API or client surface change. The consent flow's response when the cap is already hit changes from a 303 redirect that silently creates the agent anyway to 402 Payment Required, the same status and error shape POST /v1/agents already returns for the identical condition.

Operational risk

No schema change. An OAuth consent request that previously created an agent past the cap now gets 402 limit_exceeded instead, matching the REST endpoint's existing behavior for the same account state.

Test plan

  • New test (internal/agent/oauth_consent_test.go, TestHTTP_Consent_Allow_CreateNew_AtAgentCap) drives the real consent HTTP handler with max_agents: 0. Fails on main (303, agent created), passes on this branch (402, no agent row, no auth code issued), both against a real Postgres 16 container.
  • go test ./internal/agent/... and go test ./... (the Go tests job's own command) both green.
  • gofmt -l . and make fmt-check clean.
  • Not verified: CheckAgentCreate failing for a reason other than the limit itself (a DB error, say) has no dedicated test here, matching the equivalent branch on the REST path, which is also untested today.

Fixes #951

issueOAuthCodeWithNewAgent created agents via CreateAgentTx with no
max_agents check anywhere in the file, unlike POST /v1/agents, which
has always called EnforceAgentCreate first. Wire in the same check
before the transaction that creates the agent.

Fixes tokencanopy#951
@AmirF194
AmirF194 requested a review from jiashuoz as a code owner August 28, 2026 14:24
@jiashuoz

jiashuoz commented Sep 1, 2026

Copy link
Copy Markdown
Member

Blocker: the new cap check is still check-then-act.

issueOAuthCodeWithNewAgent calls CheckAgentCreate before BeginTx/CreateAgentTx. Unlike the REST create path, this sequence has no per-user advisory lock and no max-agent count recheck in the same transaction as the insert. Two concurrent OAuth auto-provision requests—or one OAuth request racing a REST create—can therefore both pass the check and exceed max_agents.

Please move the cap enforcement into the same transaction that inserts the agent, reusing the atomic lock/recheck helper, and add a concurrency regression test (for example, cap 1 with concurrent OAuth auto-creates: at most one agent row and the rest return 402). The current max_agents: 0 test covers only a single request and does not exercise this race.

issueOAuthCodeWithNewAgent's cap check (CheckAgentCreate) ran before
BeginTx/CreateAgentTx as a separate step, so it had no per-user advisory
lock and no recheck inside the same transaction as the insert. Two
concurrent OAuth auto-provision requests, or one racing a REST create,
could both pass the check and exceed max_agents (reviewer-reproduced,
tokencanopy#952).

tokencanopy#942 already closed the equivalent race on the REST create path with
CreateAgentWithLimit (advisory lock + recheck + insert, one transaction).
That helper opens its own tx internally, which does not fit this path:
the OAuth flow needs the cap check and the insert to commit or roll back
together with the authorization-code write on its own already-open tx.
Split CreateAgentWithLimit into a tx-owning wrapper plus
CreateAgentWithLimitTx, which takes a caller-owned tx and does the lock,
count check and insert on it (same relationship CreateAgentTx already has
to CreateAgent). issueOAuthCodeWithNewAgent now calls the Tx variant on
its own transaction instead of the standalone check-then-act helper.

Added a concurrency regression test matching tokencanopy#942's REST-path one: 8
concurrent consent submissions against max_agents=1 must produce exactly
1 created agent and 7 rejections, with the auth-code count matching the
agent count. Confirmed it fails on the pre-fix check-then-act sequence
(non-deterministically over-admits, e.g. 6 created in one run) and passes
with the atomic version. Updated the existing single-request cap test to
seed a real agent under max_agents=1 rather than max_agents=0, since the
atomic path treats max_agents<=0 as unlimited (same convention
CreateAgentWithLimit already documents).

go build and go vet clean on touched files; gofmt clean.

Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>
@AmirF194

AmirF194 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Fixed at bec3759. Split CreateAgentWithLimit into a tx-owning wrapper plus CreateAgentWithLimitTx, which takes a caller-owned transaction and does the advisory lock, count check and insert on it, the same relationship CreateAgentTx already has to CreateAgent. issueOAuthCodeWithNewAgent now calls the Tx variant on its own already-open transaction instead of the old check-then-act CheckAgentCreate.

Added the concurrency regression test you described: 8 concurrent consent submissions against max_agents=1, asserting exactly 1 created agent and 7 rejections, with the auth-code count matching the agent count. Confirmed it fails on the pre-fix sequence (non-deterministically over-admits, one run produced 6 created) and passes with the atomic version. Also updated the existing single-request cap test to seed a real agent under max_agents=1 rather than max_agents=0, since the atomic path treats max_agents<=0 as unlimited.

go build and go vet are clean on the touched files, gofmt is clean. I ran the targeted package tests (both the new race test and the existing cap test) directly against Postgres rather than the full internal/agent suite, which was timing out in this environment independent of this change.

issueOAuthCodeWithNewAgent's maxAgents lookup (a.enforcer.Get) ran after
BeginTx, so each in-flight request held two pool connections at once:
the transaction's own plus a second one for the lookup. Under this
package's own concurrency regression test (8 concurrent requests) that
needs up to 16 connections to make progress against pgxpool's default
MaxConns of 4, and the requests holding a connection while waiting on
a second one that will never free is a real deadlock, not flakiness.

Reproduced: capping pool_max_conns=4 and running
TestHTTP_Consent_ConcurrentCreateNewRespectsMaxAgents (added in this
PR) hangs to the test binary's timeout on every run. Moving the lookup
before BeginTx, the same order the REST create path already uses in
agents_write.go, means only one connection is ever held at a time; the
same test then passes in 3-4s across 20 consecutive runs.

go build, go vet and gofmt clean on the touched file. Full
internal/agent package suite green under -race (41/41).

Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>
@AmirF194

AmirF194 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Update: the new concurrency test I added exposed a real deadlock in CI, not a flake. issueOAuthCodeWithNewAgent looked up the agent limit with a.enforcer.Get after opening the transaction, so each in-flight request held two pool connections at once, the transaction's own plus a second for the lookup. At 8 concurrent requests against pgxpool's default MaxConns of 4, that needs up to 16 connections to make progress, so the job hung to its own timeout instead of finishing.

Reproduced locally by capping pool_max_conns=4: the test hangs on every run.

Fixed at 63f2053: moved the limit lookup before BeginTx, the same order the REST create path already uses in agents_write.go, so only one connection is ever held at a time. The same test now passes in 3-4s across 20 consecutive runs, and the full internal/agent suite is green under -race.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

limits: OAuth auto-provisioned agents bypass max_agents entirely

2 participants