feat: per-run toolkit cloning for multi-user credential isolation#7404
feat: per-run toolkit cloning for multi-user credential isolation#7404Mustafa-Esoofally wants to merge 1 commit intofeat/google-auth-db-storagefrom
Conversation
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
- 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)
|
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. |
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.
|
Consolidated into #7635 together with PR #7376 (DB-backed OAuth). The single unique commit from this branch ( #7635 also adds Happy to close this in favor of #7635 once reviewers confirm the consolidation preserves the isolation semantics. |
|
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. |
Summary
Adds
_clone_for_run()to theToolkitbase 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_authenticatedecorator storesself.credsandself.serviceon 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()whenuser_idis 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)+ resetcreds,service, and any*_serviceattributes to None_tools.py— Clone inparse_tools()whenuser_idis set + rebind Function entrypoints to the cloneauth.py—authenticate_googledefaults to all registered services (one OAuth for everything),get_oauth_routeracceptsdbparamDesign Decisions
_get_service(run_context): Clone requires zero changes to 74 tool methods. The alternative touches every method signature.name(not__name__) to handle custom-named and@tool-decorated methods correctly.*_servicereset loop: Catches Slides'slides_service/drive_servicewithout hardcoding — any future toolkit with*_serviceattributes gets reset automatically.Verification
arun()calls with differentuser_ids— mustafa gets emails, alice/bob correctly blocked, original toolkit stays cleanType of change
Checklist
./scripts/format.sh && ./scripts/validate.sh