Skip to content

fix(authz): fall back to body principal in _register_in_auth - #373

Open
alvinkam2001 wants to merge 5 commits into
mainfrom
akam/fix-authz-register-in-auth-body-fallback
Open

fix(authz): fall back to body principal in _register_in_auth#373
alvinkam2001 wants to merge 5 commits into
mainfrom
akam/fix-authz-register-in-auth-body-fallback

Conversation

@alvinkam2001

@alvinkam2001 alvinkam2001 commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

The whitelisted /agents/register path clears the middleware principal. The pod SDK ships the manifest-declared identity in the request body, but _register_in_auth was still reading only from the middleware and silently skipping ownership registration. This threads the body-supplied principal into the use case as a fallback, mirroring what #325 (cae5f94) already did for the route's check/grant fallback.

Symptom pre-fix: agent registration succeeds, no error logs, WARN Skipping authorization registration for agent: no creator resolvable fires at agents_use_case.py:83, and the agent has no owner tuple in agentex_permissions — invisible to the SGP UI listing.

Root cause

Bisected across three commits on the same code path:

Commit Date What it did
14796e9 (#270) Jun 3 Moved ownership registration into _register_in_auth, keyed off self.authorization_service.principal_context.
63f89e7 (#292) Jun 9 Route-layer _has_resolvable_creator guard — skips check/grant when middleware principal is missing (i.e. the whitelisted pod path).
cae5f94 (#325) Jun 18 Route-layer fallback to the body's principal_context for check/grant.

cae5f94 recovered ownership grants for authenticated CLI/UI callers via the route layer, but the use case's register_resource call inside _register_in_auth was still reading only from the middleware. Result: every pod-based deploy since Jun 3 that lands on the whitelisted route lost its owner tuple.

What this change does

  1. agentex/src/domain/use_cases/agents_use_case.py
    • _register_in_auth accepts an optional body_principal_context and uses it as fallback when the middleware principal isn't resolvable. Passes the resolved principal explicitly to register_resource(..., principal_context=...) so it doesn't silently default to the (None) middleware principal.
    • Extracts the dict-vs-object identity check into a private _has_resolvable_creator staticmethod (same shape as the one in agents.py:53).
    • register_agent accepts and forwards body_principal_context.
  2. agentex/src/api/routes/agents.py
    • Route passes body_principal_context=request.principal_context when calling the use case, matching the precedence already applied to check/grant.
  3. agentex/tests/integration/use_cases/test_agent_authz_dual_write.py
    • test_create_falls_back_to_body_principal_when_middleware_missing — covers the whitelisted pod path with a body-supplied service account.
    • test_middleware_principal_takes_precedence_over_body — regression guard confirming authenticated CLI/UI callers keep using the middleware identity.

What this change does NOT do

  • register_build is intentionally left alone. That route isn't whitelisted, so the middleware principal is always set for legitimate callers — there's no trigger for the fallback there. Keeping the fix scoped to the actual regression.
  • _has_resolvable_creator is now duplicated between agents.py (route) and agents_use_case.py (use case). De-duping into a shared util would be an unrelated refactor and would need to sit somewhere neutral (not in api/, since the domain shouldn't import from it). Happy to do that in a follow-up if a reviewer prefers.

Symptom (from Rocket beta)

image

The manifest for this agent declared a valid service account:

auth:
  principal:
    service_account_id: {{SERVICE_ACCOUNT_ID}}
    account_id: {{ACCOUNT_ID}}

The SDK correctly base64-encodes this into AUTH_PRINCIPAL_B64 and ships it in the register body's principal_context field (see scale-agentex-python/src/agentex/lib/utils/registration.py). The service was just not reading it in _register_in_auth.

Manual backfill of the agentex_permissions row (via POST /private/v5/agentex/permissions on egp-api-backend) restored visibility, confirming the SGP permissions row is what was missing.

Test plan

  • New: test_create_falls_back_to_body_principal_when_middleware_missing
  • New: test_middleware_principal_takes_precedence_over_body
  • Existing 7 Mock-based tests in TestAgentRegisterOnCreate still pass locally (9/9 passed, 6 Docker-only tests deselected)
  • Verify on beta: deploy a fresh agent from a Rocket manifest, confirm the agentex_permissions row is written automatically (no manual backfill required) and the agent shows up in the SGP UI listing

Greptile Summary

This PR fixes a silent ownership-registration gap on the whitelisted /agents/register path: the middleware principal is cleared for pod self-registration, so _register_in_auth was reading None and skipping the register_resource call, leaving agents with no owner tuple in agentex_permissions and invisible in the SGP UI. The fix threads the request body's principal_context through register_agent_register_in_auth as a fallback, mirroring the route-layer precedence already in place for check/grant.

  • agents_use_case.py: _register_in_auth now falls back to body_principal_context when the middleware principal is unresolvable, returns the resolved principal alongside the registration flag so _safe_deregister can use the same identity for symmetric compensation, and the creator-resolution logic is extracted into a reusable _has_resolvable_creator static method.
  • agents.py: Route passes body_principal_context=request.principal_context to the use case, consistent with the existing check/grant fallback pattern.
  • test_agent_authz_dual_write.py: Three new tests cover the whitelisted-path body-principal fallback, middleware-principal precedence, and deregister symmetry during a duplicate race on the whitelisted path.

Confidence Score: 4/5

The register-path fix is correct and well-tested, but _safe_deregister introduces a behavioral regression on the delete path that would silently pass None to the authorization gateway instead of the authenticated middleware principal.

The ownership-registration fix is well-reasoned and mirrors the existing route-layer pattern exactly. The only defect is in _safe_deregister: its new default is None, and AuthorizationService.deregister_resource uses an Ellipsis sentinel to distinguish 'not passed → use middleware' from 'explicitly None → pass None to gateway.' The delete path calls _safe_deregister(agent.id) with no principal, so None flows through to the gateway on every agent deletion, bypassing the authenticated middleware principal that was correctly used pre-PR. The existing delete tests use AsyncMock and do not assert on the principal_context kwarg, so this does not surface in the test suite.

Files Needing Attention: agentex/src/domain/use_cases/agents_use_case.py — the _safe_deregister default and the delete-path deregister call at line 440.

Important Files Changed

Filename Overview
agentex/src/domain/use_cases/agents_use_case.py Core fix threads body_principal_context through _register_in_auth and _safe_deregister correctly for register paths; however, _safe_deregister's new default of None breaks the delete path by passing None to the gateway instead of the middleware principal (Ellipsis sentinel issue).
agentex/src/api/routes/agents.py Route correctly forwards request.principal_context as body_principal_context to the use case, mirroring the pre-existing check/grant fallback pattern.
agentex/tests/integration/use_cases/test_agent_authz_dual_write.py Three well-constructed new tests cover body-principal fallback, middleware-principal precedence, and duplicate-race compensation on the whitelisted path. The pre-existing _record_existence side-effect function is correctly updated to accept the new principal_context kwarg.

Sequence Diagram

sequenceDiagram
    participant SDK as Pod SDK
    participant Route as /agents/register (route)
    participant UC as AgentsUseCase
    participant AuthSvc as AuthorizationService
    participant Gateway as agentex-auth gateway

    SDK->>Route: "POST /agents/register (body: principal_context=SA)"
    Note over Route: middleware principal = None (whitelisted)
    Route->>Route: _has_resolvable_creator(None) → false
    Route->>Route: "principal_context = request.principal_context (SA)"
    Route->>Route: "enforce_ownership = _has_resolvable_creator(SA) → true"
    Route->>AuthSvc: "check(agent:*, create, principal_context=SA)"
    Route->>UC: "register_agent(..., body_principal_context=SA)"
    UC->>UC: "_register_in_auth(agent.id, body_principal_context=SA)"
    UC->>UC: "_has_resolvable_creator(middleware=None) → false"
    UC->>UC: "principal_context = SA (fallback)"
    UC->>UC: _has_resolvable_creator(SA) → true
    UC->>AuthSvc: "register_resource(agent:id, principal_context=SA)"
    AuthSvc->>Gateway: "register_resource(SA, agent:id, parent=None)"
    UC->>UC: agent_repo.create(agent)
    UC-->>Route: agent_entity
    Route->>AuthSvc: "grant(agent:id, principal_context=SA)"
    Route-->>SDK: RegisterAgentResponse
Loading

Fix All in Cursor Fix All in Claude Code Fix All in Codex

Prompt To Fix All With AI
### Issue 1
agentex/src/domain/use_cases/agents_use_case.py:44-46
`_safe_deregister` now always passes `principal_context=principal_context` to `deregister_resource`. When called from `delete` (line 440) as `_safe_deregister(agent.id)`, the default `None` is forwarded. `AuthorizationService.deregister_resource` uses an Ellipsis sentinel (`principal_context=...`) to mean "fall back to middleware"; receiving `None` instead causes it to pass `None` straight to the gateway, silently bypassing the middleware principal on every delete-path deregistration. Pre-PR, the call carried no `principal_context` kwarg at all, so the middleware principal was used correctly. Propagating the same `...` sentinel as the default restores that behaviour.

```suggestion
    async def _safe_deregister(
        self, agent_id: str, principal_context: Any = ...
    ) -> None:
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Reviews (5): Last reviewed commit: "Merge branch 'main' into akam/fix-authz-..." | Re-trigger Greptile

Greptile also left 1 inline comment on this PR.

@alvinkam2001
alvinkam2001 requested a review from a team as a code owner July 23, 2026 01:23
@alvinkam2001
alvinkam2001 requested a review from asherfink July 23, 2026 01:44

@asherfink asherfink left a comment

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.

verified a few things beyond the happy path:

  • scope is right. _register_in_auth only runs on the genuine first-time create branch, so the gitea/register-build path (agent already exists at deploy time, so register hits an update branch) never hits this. this only bites the older helm-install/self-register model, which is Rocket. so the blast radius matches the description.
  • middleware-takes-precedence is preserved and now test-guarded, so an authenticated caller can't slip a body principal in to reattribute ownership. good, that's the property that matters most.
  • register_resource still can't run unless the route's create check passed (enforce_ownership gates both), so no create-gate bypass.
  • leaving register-build alone is correct (not whitelisted, middleware always set).
  • checked logging since it's forwarding a principal further: the one new log line carries only agent_id, and the authz gateway's error handler logs the truncated response text, not the request payload. principal_context can carry an api_key (the auth cache explicitly refuses to cache principals that do), so keep it out of any debug logging on these paths, but nothing in this PR logs it. clean.

one non-blocking note inline: a register/deregister principal asymmetry in the compensation path. good to merge once you've glanced at it.

await self.authorization_service.register_resource(
AgentexResource.agent(agent_id)
AgentexResource.agent(agent_id),
principal_context=principal_context,

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.

one asymmetry worth closing: register_resource now runs with the resolved (body) principal, but the compensating _safe_deregister still calls deregister_resource(agent(agent_id)) with no principal, so on the whitelisted pod path it deregisters with the None middleware principal. if a multi-pod cold start races and one loses on DuplicateItemError, its _safe_deregister fires with None and (assuming deregister validates the principal the same way register/grant do, per the #292 rationale that a None principal 422s) silently fails and swallows, leaving an orphaned ownership tuple for an agent id that was never persisted. narrow (needs the parallel first-create race on the older helm path) and not a security issue since it's the same account, but worth threading the resolved principal into _safe_deregister so register/deregister stay symmetric. if deregister ignores the principal server-side, disregard.

Comment on lines +44 to +46
async def _safe_deregister(
self, agent_id: str, principal_context: Any = None
) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 _safe_deregister now always passes principal_context=principal_context to deregister_resource. When called from delete (line 440) as _safe_deregister(agent.id), the default None is forwarded. AuthorizationService.deregister_resource uses an Ellipsis sentinel (principal_context=...) to mean "fall back to middleware"; receiving None instead causes it to pass None straight to the gateway, silently bypassing the middleware principal on every delete-path deregistration. Pre-PR, the call carried no principal_context kwarg at all, so the middleware principal was used correctly. Propagating the same ... sentinel as the default restores that behaviour.

Suggested change
async def _safe_deregister(
self, agent_id: str, principal_context: Any = None
) -> None:
async def _safe_deregister(
self, agent_id: str, principal_context: Any = ...
) -> None:
Prompt To Fix With AI
This is a comment left during a code review.
Path: agentex/src/domain/use_cases/agents_use_case.py
Line: 44-46

Comment:
`_safe_deregister` now always passes `principal_context=principal_context` to `deregister_resource`. When called from `delete` (line 440) as `_safe_deregister(agent.id)`, the default `None` is forwarded. `AuthorizationService.deregister_resource` uses an Ellipsis sentinel (`principal_context=...`) to mean "fall back to middleware"; receiving `None` instead causes it to pass `None` straight to the gateway, silently bypassing the middleware principal on every delete-path deregistration. Pre-PR, the call carried no `principal_context` kwarg at all, so the middleware principal was used correctly. Propagating the same `...` sentinel as the default restores that behaviour.

```suggestion
    async def _safe_deregister(
        self, agent_id: str, principal_context: Any = ...
    ) -> None:
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Fix in Cursor Fix in Claude Code Fix in Codex

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.

2 participants