fix(oauth): enforce max_agents on auto-provisioned agents - #952
Conversation
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
|
Blocker: the new cap check is still check-then-act.
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 |
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>
|
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>
|
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. |
Summary
Filed as #951:
issueOAuthCodeWithNewAgent(the OAuth auto-provision path) creates an agent viaCreateAgentTxwith nomax_agentscheck anywhere in the file, unlikePOST /v1/agents, which has always calledEnforceAgentCreatebefore 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 shapePOST /v1/agentsalready 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_exceededinstead, matching the REST endpoint's existing behavior for the same account state.Test plan
internal/agent/oauth_consent_test.go,TestHTTP_Consent_Allow_CreateNew_AtAgentCap) drives the real consent HTTP handler withmax_agents: 0. Fails onmain(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/...andgo test ./...(theGo testsjob's own command) both green.gofmt -l .andmake fmt-checkclean.CheckAgentCreatefailing 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