Conversation
WalkthroughThis change refactors both the frontend and backend to remove explicit user ID parameters from database and driver method calls, introducing user-scoped abstractions and agent proxies for server-side operations. On the frontend, query invalidation replaces explicit refetch hooks. The backend now routes driver calls through agent proxies and relies on internal user context. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant NotificationProvider
participant QueryClient
Client->>NotificationProvider: Receive socket message
NotificationProvider->>QueryClient: invalidateQueries (labels/threads) (debounced)
sequenceDiagram
participant API
participant Router
participant getZeroAgent
participant ZeroAgent
participant MailManager
API->>Router: Request (e.g., listThreads)
Router->>getZeroAgent: getZeroAgent(activeConnection.id)
getZeroAgent->>ZeroAgent: Return agent proxy
Router->>ZeroAgent: callDriver('listThreads', args)
ZeroAgent->>MailManager: listThreads(args)
MailManager-->>ZeroAgent: Result
ZeroAgent-->>Router: Result
Router-->>API: Response
Possibly related PRs
Poem
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
This stack of pull requests is managed by Graphite. Learn more about stacking. |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
apps/server/src/routes/chat.ts (1)
25-33: ```shell
#!/bin/bashVerify usage of each imported symbol within the chat route file
file="apps/server/src/routes/chat.ts"
symbols=(getPromptName connection env getPrompt openai and eq McpAgent groq EPrompts)for sym in "${symbols[@]}"; do
echo "=== Searching for '$sym' in $file ==="
rg -n "$sym" "$file" || echo "No occurrences of '$sym' found"
done</blockquote></details> <details> <summary>apps/server/src/lib/server-utils.ts (2)</summary><blockquote> `7-11`: **Ensure proper type safety for the RPC proxy return.** The `getZeroDB` function now returns `rpcTarget` from `stub.setMetaData(userId)`. Consider adding proper TypeScript return type annotation to ensure type safety and better IDE support. ```diff -export const getZeroDB = (userId: string) => { +export const getZeroDB = (userId: string): Promise<RpcDO> => { const stub = env.ZERO_DB.get(env.ZERO_DB.idFromName(userId)); const rpcTarget = stub.setMetaData(userId); return rpcTarget; };
13-16: Add return type annotation for consistency.The new
getZeroAgentfunction should have an explicit return type for better type safety.-export const getZeroAgent = (connectionId: string) => { +export const getZeroAgent = (connectionId: string): DurableObjectStub => { const stub = env.ZERO_AGENT.get(env.ZERO_AGENT.idFromName(connectionId)); return stub; };
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
apps/mail/components/party.tsx(0 hunks)apps/server/src/lib/auth.ts(4 hunks)apps/server/src/lib/driver/utils.ts(1 hunks)apps/server/src/lib/notes-manager.ts(6 hunks)apps/server/src/lib/server-utils.ts(3 hunks)apps/server/src/main.ts(4 hunks)apps/server/src/routes/chat.ts(3 hunks)apps/server/src/trpc/routes/connections.ts(2 hunks)apps/server/src/trpc/routes/drafts.ts(2 hunks)apps/server/src/trpc/routes/label.ts(5 hunks)apps/server/src/trpc/routes/mail.ts(21 hunks)apps/server/src/trpc/routes/settings.ts(2 hunks)apps/server/src/trpc/routes/shortcut.ts(1 hunks)apps/server/src/trpc/trpc.ts(1 hunks)
💤 Files with no reviewable changes (1)
- apps/mail/components/party.tsx
🔇 Additional comments (24)
apps/server/src/trpc/routes/settings.ts (1)
20-20: Here’s a closer look at the trpc route to confirm howdbis wired and whetherfindUserSettings()is scoped correctly:#!/bin/bash # Show the top of settings.ts to see how `db` is obtained and passed sed -n '1,200p' apps/server/src/trpc/routes/settings.tsapps/server/src/lib/driver/utils.ts (1)
19-19: Ensure database method has user scoping to prevent security issues.The call to
db.deleteActiveConnection(activeConnection.id)no longer passes the user ID as a parameter. This could be problematic if the database method doesn't have internal user context, potentially allowing deletion of connections belonging to other users.#!/bin/bash # Verify deleteActiveConnection method implementation echo "Searching for deleteActiveConnection implementation..." ast-grep --pattern 'deleteActiveConnection($_) { $$$ }' echo "Checking if the method validates user ownership of the connection..." rg -A 15 "deleteActiveConnection" --type tsapps/server/src/trpc/routes/shortcut.ts (1)
17-17: Verify user context handling in insertUserHotkeys method.Similar to other database method calls in this refactoring,
db.insertUserHotkeys(shortcuts)no longer receives the user ID parameter. Ensure this method has internal user scoping to prevent shortcuts from being associated with the wrong user.#!/bin/bash # Verify insertUserHotkeys method implementation echo "Searching for insertUserHotkeys implementation..." ast-grep --pattern 'insertUserHotkeys($_) { $$$ }' echo "Checking for user context validation..." rg -A 10 "insertUserHotkeys" --type tsapps/server/src/trpc/trpc.ts (1)
45-45: Verify all usages of ctx.driver have been updated.The removal of
driverfrom the context is a breaking change for procedures usingactiveDriverProcedure. Any code expectingctx.driverwill now fail. Based on the changes indrafts.ts, the new pattern appears to usegetZeroAgent()instead.#!/bin/bash # Find all usages of ctx.driver to ensure they've been updated echo "Searching for ctx.driver usage that may need updating..." rg "ctx\.driver" --type ts echo "Searching for driver property access in context..." rg "\.driver\." --type ts echo "Verifying getZeroAgent usage as replacement..." rg "getZeroAgent" --type tsapps/server/src/trpc/routes/drafts.ts (2)
8-11: LGTM: Consistent refactoring pattern.The refactoring consistently replaces direct driver method calls with the agent proxy pattern while preserving the same input/output behavior. The destructuring of
activeConnectionfrom context and the method call patterns are well-structured.Also applies to: 13-17, 27-31
2-2: ```shell
#!/bin/bashDisplay the Durable Object class and callDriver implementation for ZERO_AGENT in chat.ts
sed -n '1,200p' apps/server/src/routes/chat.ts
</details> <details> <summary>apps/server/src/trpc/routes/connections.ts (3)</summary> `21-21`: **Verify that database method supports the new signature.** The `findManyConnections()` call now omits the user ID parameter. Ensure the database method has been updated to work with implicit user context. ```shell #!/bin/bash # Description: Verify that findManyConnections method signature supports no userId parameter # Search for the method definition to check its current signature ast-grep --pattern 'findManyConnections($_) { $$$ }' # Also search for any calls to this method to see the expected pattern rg -A 3 "findManyConnections"
57-60: The initial search didn’t surface the definitions fordeleteConnectionorupdateUser. Let’s locate their implementations to verify their parameter signatures and any embedded authorization logic:#!/usr/bin/env bash set -e echo "🔍 Searching for deleteConnection usage and definition..." rg "deleteConnection" -n -C 3 echo "" echo "🔍 Searching for updateUser usage and definition..." rg "updateUser" -n -C 3
47-49: ```shell
#!/bin/bash
set -euo pipefailecho "🔍 Searching for all occurrences of findUserConnection..."
rg -n "findUserConnection" --glob '*.ts' --color=never || trueecho "🔍 Searching for all occurrences of updateUser..."
rg -n "updateUser" --glob '*.ts' --color=never || trueecho "🔍 Looking for exported definitions of findUserConnection..."
rg -n "export (async )?(function|const|let|class) findUserConnection" --glob '*.ts' --color=never || trueecho "🔍 Looking for exported definitions of updateUser..."
rg -n "export (async |declare )?(function|const|let|class) updateUser" --glob '*.ts' --color=never || true</details> <details> <summary>apps/server/src/lib/notes-manager.ts (3)</summary> `20-30`: ```shell #!/bin/bash # Search for the definition of findHighestNoteOrder to check for userId parameter rg -n -C5 "findHighestNoteOrder" apps/server # Search for the definition of createNote to check for user context handling rg -n -C5 "createNote" apps/server
9-9: To pinpoint the method’s definition and confirm it properly uses the implicit user context, let’s list every occurrence offindManyNotesByThreadIdand inspect thegetZeroDBimplementation:#!/bin/bash # List all references and definitions of findManyNotesByThreadId echo "=== findManyNotesByThreadId occurrences ===" rg -n "findManyNotesByThreadId" -C3 # Locate getZeroDB implementation to verify how user context is applied echo -e "\n=== getZeroDB implementation ===" rg -n "getZeroDB" -C10
46-52: We still need to confirm that the RPC driver returned bycreateDriver(used under the hood bygetZeroDB) enforcesuserIdscoping on its methods. Please run:#!/bin/bash # 1. Locate createDriver definition echo "Files containing createDriver:" rg -l "createDriver" -g "apps/server/src/lib/driver/**" # 2. Show createDriver implementation echo echo "createDriver implementation:" rg -n "export function createDriver" -C10 $(rg -l "createDriver" -g "apps/server/src/lib/driver/**") # 3. Inspect returned driver methods (e.g., findNoteById, updateNote) echo echo "Driver method definitions in that file:" rg -n -C5 "findNoteById\|updateNote" $(rg -l "createDriver" -g "apps/server/src/lib/driver/**")apps/server/src/lib/auth.ts (2)
119-119: To locate both the declaration and its scoping logic, let’s pull in some context around every occurrence offindManyConnections:#!/bin/bash # Show 10 lines of context around each findManyConnections match rg -C 10 "findManyConnections" -n .
210-221: Let’s inspect howdbis imported and bound inauth.ts, and review the tRPC context to see ifdbcarries the authenticated user ID:#!/bin/bash # Show top of auth.ts (imports, context setup) echo "=== apps/server/src/lib/auth.ts (lines 1–60) ===" sed -n '1,60p' apps/server/src/lib/auth.ts # Show around user-settings usage in auth.ts echo "=== apps/server/src/lib/auth.ts (lines 180–260) ===" sed -n '180,260p' apps/server/src/lib/auth.ts # Show tRPC context definition to see how `db` is constructed echo "=== apps/server/src/trpc/context.ts (lines 1–60) ===" sed -n '1,60p' apps/server/src/trpc/context.tsapps/server/src/trpc/routes/label.ts (4)
72-88: ```shell
#!/bin/bashLocate all occurrences of callDriver and getZeroAgent and display their definitions for error handling inspection
List files containing callDriver and getZeroAgent
echo "Files defining or using callDriver:"
rg -l "callDriver" .echo -e "\nFiles defining or using getZeroAgent:"
rg -l "getZeroAgent" .For each file with callDriver, show the surrounding context
for file in $(rg -l "callDriver"); do
echo -e "\n----- $file -----"
rg -A20 -B5 "callDriver" "$file"
doneFor each file with getZeroAgent, show the surrounding context
for file in $(rg -l "getZeroAgent"); do
echo -e "\n----- $file -----"
rg -A20 -B5 "getZeroAgent" "$file"
done--- `43-49`: ```shell #!/bin/bash # Inspect the create route input schema in labelsRouter rg -n "create: activeDriverProcedure" -A 20 apps/server/src/trpc/routes/label.ts
15-19: I’ll list the driver directory and inspect both the index and utils files to locate anycallDriverimplementation.#!/bin/bash set -e echo "📂 apps/server/src/lib/driver directory structure:" ls -R apps/server/src/lib/driver echo echo "🔍 Searching for 'callDriver' in driver directory:" rg -n "callDriver" apps/server/src/lib/driver || true echo echo "📄 apps/server/src/lib/driver/index.ts (lines 1–200):" sed -n '1,200p' apps/server/src/lib/driver/index.ts echo echo "📄 apps/server/src/lib/driver/utils.ts (lines 1–200):" sed -n '1,200p' apps/server/src/lib/driver/utils.ts
2-2: ```shell
#!/bin/bashExtract the full callDriver method implementation from ZeroAgent class
sed -n '118,200p' apps/server/src/routes/chat.ts
</details> <details> <summary>apps/server/src/routes/chat.ts (1)</summary> `165-168`: ```shell #!/bin/bash # Search for the definition of getPromptName across the codebase rg -n "function getPromptName" -n -C5 . # If not found, look for exported constants or arrow functions rg -n "getPromptName\s*=" -n -C5 .apps/server/src/lib/server-utils.ts (2)
25-32: Clean implementation of user-scoped database operations.The removal of explicit userId parameters from database calls is well-aligned with the new RpcDO proxy pattern. The user context is now properly encapsulated within the database proxy object.
94-94: Good centralization of agent retrieval logic.Using
getZeroAgenthelper function improves code consistency and maintainability.apps/server/src/trpc/routes/mail.ts (1)
23-26: Consistent implementation of the agent proxy pattern.The refactoring to use
agent.callDriveris well-implemented. The pattern of extractingactiveConnection, getting the agent, and calling driver methods is consistent throughout.apps/server/src/main.ts (2)
45-174: Well-designed user-scoped proxy implementation.The
RpcDOclass effectively encapsulates user context and provides a clean API by delegating to the underlyingZeroDBinstance. The pattern ensures that user ID cannot be accidentally passed incorrectly.
179-181: Clean implementation of the metadata setter.The
setMetaDatamethod provides a simple and effective way to create user-scoped RPC proxies.
| } | ||
|
|
||
| await db.deleteUser(user.id); | ||
| await db.deleteUser(); |
There was a problem hiding this comment.
Critical: Verify user deletion is properly scoped.
The deleteUser() call now omits the user ID parameter. This is a critical operation - ensure the method deletes only the authenticated user and not others.
#!/bin/bash
# Description: Verify deleteUser method properly handles user context
# Find the method definition
ast-grep --pattern 'deleteUser() {
$$$
}'
# Look for any user scoping or authorization checks
rg -A 10 "deleteUser.*user\|deleteUser.*authorization"🤖 Prompt for AI Agents
In apps/server/src/lib/auth.ts at line 152, the call to deleteUser() lacks a
user ID parameter, risking deletion of unintended users. Modify the call to pass
the authenticated user's ID explicitly to ensure only that user is deleted.
Also, verify the deleteUser method implementation properly scopes deletion by
user ID and includes authorization checks to prevent unauthorized deletions.
| public async callDriver<TResult = any>( | ||
| method: keyof MailManager, | ||
| ...args: any[] | ||
| ): Promise<TResult> { | ||
| if (!this.driver) { | ||
| throw new Error('Driver not initialized'); | ||
| } | ||
| const driverMethod = this.driver[method] as any; | ||
| if (typeof driverMethod !== 'function') { | ||
| throw new Error(`Method ${String(method)} is not a function on driver`); | ||
| } | ||
| const result = await driverMethod.apply(this.driver, args); | ||
| return result as TResult; | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Improve type safety in callDriver method.
The callDriver method uses any types which reduces type safety. Consider using generic types or method overloads to preserve type information for driver method calls.
- public async callDriver<TResult = any>(
- method: keyof MailManager,
- ...args: any[]
- ): Promise<TResult> {
+ public async callDriver<K extends keyof MailManager>(
+ method: K,
+ ...args: Parameters<MailManager[K]>
+ ): Promise<ReturnType<MailManager[K]>> {
if (!this.driver) {
throw new Error('Driver not initialized');
}
- const driverMethod = this.driver[method] as any;
+ const driverMethod = this.driver[method];
if (typeof driverMethod !== 'function') {
throw new Error(`Method ${String(method)} is not a function on driver`);
}
- const result = await driverMethod.apply(this.driver, args);
- return result as TResult;
+ return await (driverMethod as any).apply(this.driver, args);
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| public async callDriver<TResult = any>( | |
| method: keyof MailManager, | |
| ...args: any[] | |
| ): Promise<TResult> { | |
| if (!this.driver) { | |
| throw new Error('Driver not initialized'); | |
| } | |
| const driverMethod = this.driver[method] as any; | |
| if (typeof driverMethod !== 'function') { | |
| throw new Error(`Method ${String(method)} is not a function on driver`); | |
| } | |
| const result = await driverMethod.apply(this.driver, args); | |
| return result as TResult; | |
| } | |
| public async callDriver<K extends keyof MailManager>( | |
| method: K, | |
| ...args: Parameters<MailManager[K]> | |
| ): Promise<ReturnType<MailManager[K]>> { | |
| if (!this.driver) { | |
| throw new Error('Driver not initialized'); | |
| } | |
| const driverMethod = this.driver[method]; | |
| if (typeof driverMethod !== 'function') { | |
| throw new Error(`Method ${String(method)} is not a function on driver`); | |
| } | |
| return await (driverMethod as any).apply(this.driver, args); | |
| } |
🤖 Prompt for AI Agents
In apps/server/src/routes/chat.ts around lines 118 to 131, the callDriver method
currently uses 'any' types for arguments and result, reducing type safety.
Refactor this method to use generics that infer the argument types and return
type based on the method key provided. This can be done by defining the method
parameter as a key of MailManager and using TypeScript's utility types to
extract the parameter and return types of the corresponding method, ensuring
type-safe calls to driver methods without using 'any'.
| const result = (await agent.callDriver('normalizeIds', threadId)) as { threadIds: string[] }; | ||
| const { threadIds } = result; |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider improving type safety for callDriver results.
Multiple locations cast the result of agent.callDriver calls. This suggests the method returns a generic type that requires casting, which could hide type safety issues.
Consider using generics or method overloads in the callDriver implementation to provide proper type inference:
// In ZeroAgent class
callDriver<T extends keyof DriverMethods>(
method: T,
...args: Parameters<DriverMethods[T]>
): Promise<ReturnType<DriverMethods[T]>>This would eliminate the need for manual type casting and provide compile-time type safety.
Also applies to: 140-142, 185-187
🤖 Prompt for AI Agents
In apps/server/src/trpc/routes/mail.ts around lines 115 to 116, the result of
agent.callDriver is manually cast to a specific type, which reduces type safety.
To fix this, update the callDriver method in the ZeroAgent class to use generics
that infer the return type based on the method name and its parameters. This
involves defining callDriver with a generic type parameter constrained by the
keys of DriverMethods, and returning a Promise of the corresponding method's
return type. This change will remove the need for manual casting and improve
compile-time type safety. Apply the same fix to other callDriver usages around
lines 140-142 and 185-187.
| (await db)[Symbol.dispose]?.(); | ||
| } |
There was a problem hiding this comment.
Verify the await usage in disposal calls.
The disposal pattern (await db)[Symbol.dispose]?.() appears to await a synchronous value. Since getZeroDB returns an RpcDO instance (not a Promise), the await seems unnecessary.
- (await db)[Symbol.dispose]?.();
+ db[Symbol.dispose]?.();Also applies to lines 519-521.
Also applies to: 519-521
🤖 Prompt for AI Agents
In apps/server/src/main.ts around lines 508 to 509 and 519 to 521, the disposal
calls use await on values that are not Promises, specifically `(await
db)[Symbol.dispose]?.()`. Since `getZeroDB` returns an `RpcDO` instances
synchronously, remove the unnecessary await keyword from these disposal calls to
correctly handle synchronous disposal.
# READ CAREFULLY THEN REMOVE Remove bullet points that are not relevant. PLEASE REFRAIN FROM USING AI TO WRITE YOUR CODE AND PR DESCRIPTION. IF YOU DO USE AI TO WRITE YOUR CODE PLEASE PROVIDE A DESCRIPTION AND REVIEW IT CAREFULLY. MAKE SURE YOU UNDERSTAND THE CODE YOU ARE SUBMITTING USING AI. - Pull requests that do not follow these guidelines will be closed without review or comment. - If you use AI to write your PR description your pr will be close without review or comment. - If you are unsure about anything, feel free to ask for clarification. ## Description Please provide a clear description of your changes. --- ## Type of Change Please delete options that are not relevant. - [ ] 🐛 Bug fix (non-breaking change which fixes an issue) - [ ] ✨ New feature (non-breaking change which adds functionality) - [ ] 💥 Breaking change (fix or feature with breaking changes) - [ ] 📝 Documentation update - [ ] 🎨 UI/UX improvement - [ ] 🔒 Security enhancement - [ ] ⚡ Performance improvement ## Areas Affected Please check all that apply: - [ ] Email Integration (Gmail, IMAP, etc.) - [ ] User Interface/Experience - [ ] Authentication/Authorization - [ ] Data Storage/Management - [ ] API Endpoints - [ ] Documentation - [ ] Testing Infrastructure - [ ] Development Workflow - [ ] Deployment/Infrastructure ## Testing Done Describe the tests you've done: - [ ] Unit tests added/updated - [ ] Integration tests added/updated - [ ] Manual testing performed - [ ] Cross-browser testing (if UI changes) - [ ] Mobile responsiveness verified (if UI changes) ## Security Considerations For changes involving data or authentication: - [ ] No sensitive data is exposed - [ ] Authentication checks are in place - [ ] Input validation is implemented - [ ] Rate limiting is considered (if applicable) ## Checklist - [ ] I have read the [CONTRIBUTING](https://github.com/Mail-0/Zero/blob/staging/.github/CONTRIBUTING.md) document - [ ] My code follows the project's style guidelines - [ ] I have performed a self-review of my code - [ ] I have commented my code, particularly in complex areas - [ ] I have updated the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix/feature works - [ ] All tests pass locally - [ ] Any dependent changes are merged and published ## Additional Notes Add any other context about the pull request here. ## Screenshots/Recordings Add screenshots or recordings here if applicable. --- _By submitting this pull request, I confirm that my contribution is made under the terms of the project's license._ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Refactor** - Streamlined database and driver method calls across multiple areas, removing explicit user ID parameters and relying on internal context for user scoping. - Introduced an agent abstraction layer for mail, label, and draft operations, routing all actions through a unified interface. - Enhanced Durable Object interactions with a new user-scoped proxy and improved resource management. - Simplified notification and query invalidation logic in the mail app for better maintainability. - **New Features** - Added a user-scoped proxy class for database operations, enabling more secure and encapsulated data access. - Introduced a public method for dynamic driver calls in chat-related features. <!-- end of auto-generated comment: release notes by coderabbit.ai -->

READ CAREFULLY THEN REMOVE
Remove bullet points that are not relevant.
PLEASE REFRAIN FROM USING AI TO WRITE YOUR CODE AND PR DESCRIPTION. IF YOU DO USE AI TO WRITE YOUR CODE PLEASE PROVIDE A DESCRIPTION AND REVIEW IT CAREFULLY. MAKE SURE YOU UNDERSTAND THE CODE YOU ARE SUBMITTING USING AI.
Description
Please provide a clear description of your changes.
Type of Change
Please delete options that are not relevant.
Areas Affected
Please check all that apply:
Testing Done
Describe the tests you've done:
Security Considerations
For changes involving data or authentication:
Checklist
Additional Notes
Add any other context about the pull request here.
Screenshots/Recordings
Add screenshots or recordings here if applicable.
By submitting this pull request, I confirm that my contribution is made under the terms of the project's license.
Summary by CodeRabbit
Refactor
New Features