Skip to content

feat: per-run toolkit cloning for multi-user credential isolation#7404

Closed
Mustafa-Esoofally wants to merge 1 commit intofeat/google-auth-db-storagefrom
feat/per-user-creds-v2
Closed

feat: per-run toolkit cloning for multi-user credential isolation#7404
Mustafa-Esoofally wants to merge 1 commit intofeat/google-auth-db-storagefrom
feat/per-user-creds-v2

Conversation

@Mustafa-Esoofally
Copy link
Copy Markdown
Contributor

Summary

Adds _clone_for_run() to the Toolkit base class for per-user credential isolation in multi-tenant deployments (Slack, Telegram, WhatsApp via AgentOS).

Problem: When multiple users send concurrent requests via Slack, all requests share the same toolkit instance. The @google_authenticate decorator stores self.creds and self.service on the shared instance — if User A authenticates first, User B's request reuses A's credentials.

Solution: parse_tools() shallow-copies each Toolkit via _clone_for_run() when user_id is set, so each run gets its own mutable state. The existing auth decorator and all 74 Google tool methods remain unchanged.

Changes

  • toolkit.py_clone_for_run(): copy.copy(self) + reset creds, service, and any *_service attributes to None
  • _tools.py — Clone in parse_tools() when user_id is set + rebind Function entrypoints to the clone
  • auth.pyauthenticate_google defaults to all registered services (one OAuth for everything), get_oauth_router accepts db param
  • Formatting — import sort order fixes in Google toolkit files (ruff)
  • Tests — 14 new tests covering clone isolation, entrypoint rebinding, multi-user scenarios

Design Decisions

  • Clone vs ContextVar: Clone isolates ALL mutable state automatically (creds, service, label_cache, user_email, slides/drive services). ContextVar would require manually wrapping each attribute in properties — miss one and credentials leak silently.
  • Clone vs per-method _get_service(run_context): Clone requires zero changes to 74 tool methods. The alternative touches every method signature.
  • Entrypoint rebinding: Uses registered tool name (not __name__) to handle custom-named and @tool-decorated methods correctly.
  • *_service reset loop: Catches Slides' slides_service/drive_service without hardcoding — any future toolkit with *_service attributes gets reset automatically.

Verification

  • Unit tests: 14 clone tests + 47 existing toolkit tests (61 total, all pass)
  • E2E Slack test: PAL Local bot → OAuth flow → token stored in DB → Gmail API call with per-user creds
  • Concurrent isolation test: 3 simultaneous arun() calls with different user_ids — mustafa gets emails, alice/bob correctly blocked, original toolkit stays clean

Type of change

  • New feature (non-breaking change which adds functionality)

Checklist

  • Format and validate: ./scripts/format.sh && ./scripts/validate.sh
  • Both sync and async paths covered (parse_tools handles both)
  • Backward compatible — no clone when user_id is None (single-user unchanged)
  • No changes to existing tool method signatures

When user_id is set, parse_tools() shallow-copies each Toolkit via
_clone_for_run() so concurrent runs get isolated creds/service state.
Prevents credential leakage between users in multi-tenant deployments
(Slack, Telegram, WhatsApp via AgentOS).

Changes:
- toolkit.py: _clone_for_run() — copy.copy + reset creds/service/*_service
- _tools.py: clone in parse_tools when user_id is set, rebind entrypoints
- auth.py: authenticate_google defaults to all registered services,
  get_oauth_router accepts db param
- Formatting fixes (import sort order) in Google toolkit files
@Mustafa-Esoofally Mustafa-Esoofally requested a review from a team as a code owner April 7, 2026 20:58
Mustafa-Esoofally added a commit to agno-agi/pal that referenced this pull request Apr 7, 2026
- GoogleAuth instance shared across Gmail + Calendar toolkits
- store_token_in_db=True enables DB-backed credential storage
- OAuth callback router mounted at /google/oauth/callback
- Tokens stored per user_id in agno_auth_tokens table

Requires: agno-agi/agno#7404 (per-run toolkit cloning)
@github-actions
Copy link
Copy Markdown
Contributor

This pull request has been automatically marked as stale due to 14 days of inactivity. If this is still relevant, please update it or leave a comment.

@github-actions github-actions Bot added the stale label Apr 22, 2026
Mustafa-Esoofally added a commit that referenced this pull request Apr 22, 2026
Three tests nail down the concurrency guarantee for the Google OAuth port
onto v2.6.0:

1. test_factory_pattern_preserves_isolation — 30 concurrent arun-style calls
   across 3 users, each resolving a fresh toolkit via a callable factory. All
   requests return the correct user's token. This is the intended pattern.

2. test_shared_toolkit_leaks_credentials — documents and enforces the
   failure mode when a single toolkit instance is shared across users. The
   @google_authenticate wrapper (tools/google/auth.py:42) only re-auths when
   self.creds is missing/invalid, so Bob's call silently reuses Alice's
   creds. Deterministic and sequential — no race required to expose the
   leak. Canary test: flips to TOKEN::bob once per-call isolation lands.

3. test_framework_factory_invocation_isolates — exercises the real framework
   path via agno.utils.callables.ainvoke_callable_factory, matching what
   agent.arun does under the hood. Same isolation guarantee.

Together these tests show that on v2.6.0 + PR #7376 (without PR #7404's
_clone_for_run), multi-tenant auth is correct *iff* the user wires toolkits
through a factory callable. They'll also serve as a regression guard if
follow-up work changes the isolation contract.
@Mustafa-Esoofally
Copy link
Copy Markdown
Contributor Author

Consolidated into #7635 together with PR #7376 (DB-backed OAuth). The single unique commit from this branch (f2c1d0ba8Toolkit._clone_for_run() + auto-clone hook in parse_tools + test_clone_for_run.py with 14 tests) cherry-picked onto #7635 as b835f0b69.

#7635 also adds test_google_auth_concurrent.py which exercises the cloning + DB-auth flow end-to-end under real-thread concurrency (16 workers × 48 calls with forced 10ms interleave).

Happy to close this in favor of #7635 once reviewers confirm the consolidation preserves the isolation semantics.

@Mustafa-Esoofally
Copy link
Copy Markdown
Contributor Author

Consolidated into #7635 together with #7376. The single unique commit (f2c1d0b — Toolkit._clone_for_run + auto-clone hook in parse_tools + test_clone_for_run.py with 14 tests) was cherry-picked onto #7635 as 0a7def3.

#7635 also adds test_google_auth_concurrent.py which exercises the cloning + DB-auth flow end-to-end using real google.oauth2.credentials.Credentials + real-thread interleaving (16 workers × 48 calls with forced 10ms yield).

Review thread is continued on #7635.

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

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant