fix(cli): complete OAuth auth flow for CLI mode#505
Conversation
- Remove redundant access_type="offline" from authorization_url call - Broaden import error handler and print to stderr in CLI mode - Pre-start OAuth callback server before tool execution - Add completion signaling (threading.Event) to MinimalOAuthServer - Wait for browser OAuth callback and retry tool on success
start_auth_flow now returns (message, auth_url) tuple so the exception can carry the auth_url attribute needed by the CLI auth flow handler.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughstart_auth_flow now returns (message, auth_url). The OAuth callback server adds CLI signaling (auth_completed, auth_result, wait/reset helpers). CLI handler pre-starts the local server, opens the auth URL, waits for callback, and retries tool execution after authentication. Module import errors now catch all Exceptions. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User (CLI)
participant CLI as CLI Handler
participant OAuthSrv as OAuth Server
participant GoogleAuth as Google Auth
participant Browser as Browser
participant Google as Google OAuth
User->>CLI: Request tool requiring auth
CLI->>OAuthSrv: start local OAuth server
CLI->>GoogleAuth: start_auth_flow(...)
GoogleAuth->>Google: build authorization URL
Google-->>GoogleAuth: authorization URL
GoogleAuth-->>CLI: (message, auth_url)
CLI->>User: show message and open auth_url
User->>Browser: open auth_url
Browser->>Google: user authenticates
Google->>OAuthSrv: callback request (code)
OAuthSrv->>OAuthSrv: process code, set auth_result
OAuthSrv->>OAuthSrv: auth_completed.set()
CLI->>OAuthSrv: wait_for_auth(timeout)
OAuthSrv-->>CLI: return auth_result
CLI->>CLI: retry tool execution
CLI-->>User: tool result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested labels
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
auth/oauth_callback_server.py (1)
82-86:⚠️ Potential issue | 🟠 MajorMissing
auth_completedsignal when client secrets check fails.If
check_client_secrets()returns an error, the handler returns early without settingself.auth_completed. In CLI mode,_handle_cli_auth_flowwill block onwait_for_auth()for the full 5-minute timeout before reporting a generic timeout error, hiding the real cause (missing client secrets).Proposed fix
error_message = check_client_secrets() if error_message: + self.auth_result = {"success": False, "user_id": None, "error": error_message} + self.auth_completed.set() return create_server_error_response(error_message)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@auth/oauth_callback_server.py` around lines 82 - 86, The handler returns early when check_client_secrets() yields an error but doesn't signal self.auth_completed, causing _handle_cli_auth_flow to block on wait_for_auth(); modify the error path in the block around check_client_secrets() so that before returning create_server_error_response(error_message) you set/trigger self.auth_completed (e.g., call self.auth_completed.set() or the appropriate notifier used elsewhere) and store the error (e.g., self.auth_error = error_message) so wait_for_auth()/_handle_cli_auth_flow can immediately wake and report the real missing-client-secrets error.
🧹 Nitpick comments (2)
auth/oauth_callback_server.py (1)
318-337: Minor TOCTOU window in port selection, but acceptable for CLI tooling.Between
get_cli_oauth_port()releasing the probing socket andMinimalOAuthServer.start()binding the same port, another process could claim it. Sincestart()already returns(False, error_msg)on bind failure and the CLI surfaces that error, this is a tolerable race for a local development tool.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@auth/oauth_callback_server.py` around lines 318 - 337, get_cli_oauth_port() introduces a TOCTOU race by probing and releasing a port before MinimalOAuthServer.start() binds it; fix by changing the handshake so the binding is atomic: either (A) remove probing and let MinimalOAuthServer.start() accept port 0 (OS-assigned) and return the actual bound port, or (B) have get_cli_oauth_port() return an already-bound listening socket (or file descriptor) that you pass into MinimalOAuthServer.start() so start uses that socket instead of re-binding; update call sites accordingly (references: get_cli_oauth_port and MinimalOAuthServer.start).core/cli_handler.py (1)
374-398: Environment mutation in_setup_cli_oauth_servercould leak to non-CLI code paths.Setting
os.environ["GOOGLE_OAUTH_REDIRECT_URI"]and callingreload_oauth_config()permanently mutates process-wide state. This is fine for CLI mode (which exits after the tool run), but worth a brief comment to prevent future misuse in long-lived server contexts.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/cli_handler.py` around lines 374 - 398, The function _setup_cli_oauth_server mutates process-wide env via os.environ["GOOGLE_OAUTH_REDIRECT_URI"] and calls reload_oauth_config(), which can leak into non-CLI long-lived processes; fix by saving the original env value (orig = os.environ.get("GOOGLE_OAUTH_REDIRECT_URI")), set the new redirect URI, call reload_oauth_config(), start the MinimalOAuthServer (oauth_server.start()) and register it with set_cli_oauth_server, then restore the original env value and call reload_oauth_config() again so the mutation is transient; also add a brief inline comment near the os.environ mutation explaining this is intentionally scoped for CLI flow and therefore restored afterward.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/cli_handler.py`:
- Around line 401-439: The CLI auth flow can reuse a stale
oauth_server.auth_completed event when retrying, causing
oauth_server.wait_for_auth to return the previous result; before calling await
asyncio.to_thread(oauth_server.wait_for_auth, 300) in _handle_cli_auth_flow,
reset/clear the oauth_server.auth_completed state (or call a provided reset
method on MinimalOAuthServer) so each invocation of wait_for_auth waits for a
fresh browser callback; ensure this reset happens after printing the URL and
before the to_thread call so subsequent run_tool retries trigger a real new auth
flow.
---
Outside diff comments:
In `@auth/oauth_callback_server.py`:
- Around line 82-86: The handler returns early when check_client_secrets()
yields an error but doesn't signal self.auth_completed, causing
_handle_cli_auth_flow to block on wait_for_auth(); modify the error path in the
block around check_client_secrets() so that before returning
create_server_error_response(error_message) you set/trigger self.auth_completed
(e.g., call self.auth_completed.set() or the appropriate notifier used
elsewhere) and store the error (e.g., self.auth_error = error_message) so
wait_for_auth()/_handle_cli_auth_flow can immediately wake and report the real
missing-client-secrets error.
---
Nitpick comments:
In `@auth/oauth_callback_server.py`:
- Around line 318-337: get_cli_oauth_port() introduces a TOCTOU race by probing
and releasing a port before MinimalOAuthServer.start() binds it; fix by changing
the handshake so the binding is atomic: either (A) remove probing and let
MinimalOAuthServer.start() accept port 0 (OS-assigned) and return the actual
bound port, or (B) have get_cli_oauth_port() return an already-bound listening
socket (or file descriptor) that you pass into MinimalOAuthServer.start() so
start uses that socket instead of re-binding; update call sites accordingly
(references: get_cli_oauth_port and MinimalOAuthServer.start).
In `@core/cli_handler.py`:
- Around line 374-398: The function _setup_cli_oauth_server mutates process-wide
env via os.environ["GOOGLE_OAUTH_REDIRECT_URI"] and calls reload_oauth_config(),
which can leak into non-CLI long-lived processes; fix by saving the original env
value (orig = os.environ.get("GOOGLE_OAUTH_REDIRECT_URI")), set the new redirect
URI, call reload_oauth_config(), start the MinimalOAuthServer
(oauth_server.start()) and register it with set_cli_oauth_server, then restore
the original env value and call reload_oauth_config() again so the mutation is
transient; also add a brief inline comment near the os.environ mutation
explaining this is intentionally scoped for CLI flow and therefore restored
afterward.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
auth/google_auth.pyauth/oauth_callback_server.pycore/cli_handler.pycore/server.pymain.py
Without this signal, the CLI blocks for 5 minutes (timeout) instead of immediately reporting the missing credentials error. Closes CodeRabbit review on PR taylorwilsdon#505.
Without this signal, the CLI blocks for 5 minutes (timeout) instead of immediately reporting the missing credentials error. Closes CodeRabbit review on PR taylorwilsdon#505.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@auth/oauth_callback_server.py`:
- Around line 219-234: The auth flow leaves auth_completed set and auth_result
populated so subsequent calls to wait_for_auth return stale success; add a reset
method (e.g., reset_auth_state) that clears self.auth_completed (call clear()),
sets self.auth_result = None (or default failure dict), and ensure callers
invoke this reset before starting a new flow (before opening the auth URL or
calling wait_for_auth); update any code paths that start an auth flow to call
reset_auth_state to guarantee a fresh event/result state.
| def wait_for_auth(self, timeout: float = 300) -> Optional[dict]: | ||
| """ | ||
| Block until OAuth callback is received or timeout. | ||
|
|
||
| Args: | ||
| timeout: Maximum seconds to wait (default 5 minutes) | ||
|
|
||
| Returns: | ||
| Auth result dict {"success": bool, "user_id": str|None, "error": str|None} | ||
| or None if timed out | ||
| """ | ||
| completed = self.auth_completed.wait(timeout=timeout) | ||
| if completed: | ||
| return self.auth_result | ||
| return None | ||
|
|
There was a problem hiding this comment.
Reset auth state between flows to avoid stale success.
auth_completed is never cleared, so a second auth flow in the same process can return an old auth_result immediately. This can cause the CLI to skip a new login or retry prematurely.
Consider adding an explicit reset path (called before opening the auth URL) to clear the event and result.
🔧 Suggested change
- def wait_for_auth(self, timeout: float = 300) -> Optional[dict]:
+ def wait_for_auth(self, timeout: float = 300, reset: bool = False) -> Optional[dict]:
"""
Block until OAuth callback is received or timeout.
@@
- completed = self.auth_completed.wait(timeout=timeout)
+ if reset:
+ self.auth_result = None
+ self.auth_completed.clear()
+ completed = self.auth_completed.wait(timeout=timeout)
if completed:
return self.auth_result
return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@auth/oauth_callback_server.py` around lines 219 - 234, The auth flow leaves
auth_completed set and auth_result populated so subsequent calls to
wait_for_auth return stale success; add a reset method (e.g., reset_auth_state)
that clears self.auth_completed (call clear()), sets self.auth_result = None (or
default failure dict), and ensure callers invoke this reset before starting a
new flow (before opening the auth URL or calling wait_for_auth); update any code
paths that start an auth flow to call reset_auth_state to guarantee a fresh
event/result state.
|
Hm, I've never had the issue you're describing and I know a lot of folks are using cli mode as well. Any idea what about your environment might differ from mine? Will take a look, appreciate the PR! |
|
Thanks for the quick response! The issue shows up specifically when running
If you already have cached credentials (from a previous MCP server session or a prior CLI run), you'd never hit this because Our use case: we're onboarding multiple Google accounts (Gmail) via CLI mode on a headless Docker host, so every first run per account triggers the fresh OAuth flow. That's how we discovered it — the CLI would print "Open this URL..." but immediately exit, never receiving the callback. The fix pre-starts the OAuth server, catches the auth error, waits for the browser callback, then retries the tool. Happy to adjust anything if you'd prefer a different approach! |
|
Appreciate it - if you can sort out the merge conflicts and any outstanding feedback, love to get this in thanks! |
|
Hm, I can reproduce what you're describing now after blowing away my existing cred file but the fix here doesn't seem to work, I get a 404 on the callback. |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
core/cli_handler.py (1)
425-427:⚠️ Potential issue | 🟡 MinorReset OAuth wait state before waiting for a new callback.
If this flow is invoked again in-process,
wait_for_auth()can return a stale prior result because the completion event/result is not reset before waiting.Suggested defensive patch
print("Waiting for authentication (timeout: 5 minutes)...", file=sys.stderr) # Wait for the OAuth callback in a thread (blocking call) + oauth_server.auth_completed.clear() + oauth_server.auth_result = None result = await asyncio.to_thread(oauth_server.wait_for_auth, 300)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/cli_handler.py` around lines 425 - 427, Before calling asyncio.to_thread(oauth_server.wait_for_auth, 300) ensure the OAuth server's previous completion state is cleared so wait_for_auth cannot return a stale result; call an explicit reset/clear method on oauth_server (e.g., oauth_server.reset() or oauth_server.clear_auth_state()) immediately before awaiting, and if such a method does not exist add a small reset method on the OAuth server class that clears its completion event/result so subsequent invocations always wait for a fresh callback.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@core/cli_handler.py`:
- Around line 425-427: Before calling
asyncio.to_thread(oauth_server.wait_for_auth, 300) ensure the OAuth server's
previous completion state is cleared so wait_for_auth cannot return a stale
result; call an explicit reset/clear method on oauth_server (e.g.,
oauth_server.reset() or oauth_server.clear_auth_state()) immediately before
awaiting, and if such a method does not exist add a small reset method on the
OAuth server class that clears its completion event/result so subsequent
invocations always wait for a fresh callback.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
auth/google_auth.pycore/cli_handler.pycore/server.pymain.py
🚧 Files skipped from review as they are similar to previous changes (1)
- main.py
- Add reset_auth_state() to clear previous auth completion before waiting, preventing wait_for_auth from returning stale results immediately - Remove TOCTOU port pre-check in start() — let uvicorn bind directly and detect port conflicts via startup error propagation - Increase server startup timeout from 3s to 5s for reliability - Capture and propagate startup errors from the server thread Addresses CodeRabbit review feedback and 404 callback investigation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Hi @taylorwilsdon, thanks for testing! I've addressed the feedback: Changes in this update:
About the 404: The stale auth state issue (#2) is likely what you hit — if Let me know if you'd like any changes! |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
auth/oauth_callback_server.py (1)
193-210:⚠️ Potential issue | 🔴 CriticalStartup readiness check can falsely succeed on the wrong process.
At Line 203,
connect_exonly proves that something is listening. If another process already owns the port, Line 206 setsis_running=Trueand startup returns success even when this OAuth server failed to bind (leading to callback 404 behavior).🔧 Suggested fix
def run_server(): """Run the server in a separate thread.""" try: @@ except Exception as e: _startup_error[0] = e logger.error(f"Minimal OAuth server error: {e}", exc_info=True) self.is_running = False + # Fail fast if another process already owns this port. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + if probe.connect_ex((hostname, self.port)) == 0: + error_msg = f"OAuth server port already in use: {hostname}:{self.port}" + logger.error(error_msg) + return False, error_msg + # Start server in background thread self.server_thread = threading.Thread(target=run_server, daemon=True) self.server_thread.start() @@ while time.time() - start_time < max_wait: if _startup_error[0]: error_msg = f"OAuth server failed to start: {_startup_error[0]}" logger.error(error_msg) return False, error_msg + if self.server_thread and not self.server_thread.is_alive(): + error_msg = "OAuth server thread exited before startup completed" + logger.error(error_msg) + return False, error_msg try: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: result = s.connect_ex((hostname, self.port)) if result == 0:
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@auth/oauth_callback_server.py`:
- Around line 117-122: Replace leaking raw exception text in the OAuth callback
error path: keep the detailed exception only in the logger (logger.error with
exc_info=True) but change the values placed into self.auth_result (currently
{"success": False, "user_id": None, "error": str(e)}) and the string passed to
create_server_error_response(str(e)) to a generic, non-sensitive message (e.g.,
"An error occurred during authentication" or "Authentication failed"); preserve
setting self.auth_completed and returning the error response, but do not
propagate str(e) to auth_result or client-facing responses.
In `@core/cli_handler.py`:
- Around line 24-28: The CLI path that starts the MinimalOAuthServer (via
get_cli_oauth_port and set_cli_oauth_server) never guarantees teardown on early
returns; modify the run logic so the server is shutdown and deregistered in a
finally/cleanup block: after creating and set_cli_oauth_server(server) ensure
you call the server's stop/shutdown method and then set_cli_oauth_server(None)
on every exit path (including error/early returns) so no threads/sockets or
stale global remain; wrap the run branch that starts the server in try/finally
(or ensure each return first performs the same cleanup) and reference
MinimalOAuthServer, get_cli_oauth_port, and set_cli_oauth_server when making the
changes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3f59ef33-8373-4c04-a6a0-dcbbb6df1540
📒 Files selected for processing (3)
auth/oauth_callback_server.pycore/cli_handler.pycore/server.py
🚧 Files skipped from review as they are similar to previous changes (1)
- core/server.py
| from auth.oauth_callback_server import ( | ||
| MinimalOAuthServer, | ||
| get_cli_oauth_port, | ||
| set_cli_oauth_server, | ||
| ) |
There was a problem hiding this comment.
CLI run path starts OAuth server but never tears it down.
Line 474 starts a callback server, but the run branch has early return paths and no finally cleanup. In long-lived processes, repeated CLI runs can accumulate orphaned server threads/sockets and stale global registration.
🔧 Suggested fix
from auth.oauth_callback_server import (
MinimalOAuthServer,
+ cleanup_oauth_callback_server,
get_cli_oauth_port,
set_cli_oauth_server,
)
@@
if parsed["command"] == "run":
# Pre-start OAuth server before running tools that may need auth
oauth_server = _setup_cli_oauth_server()
# Merge stdin args with inline args (inline takes precedence)
args = read_stdin_args()
args.update(parsed["tool_args"])
- try:
- result = await run_tool(server, parsed["tool_name"], args)
- print(result)
- return 0
- except Exception as e:
- from auth.google_auth import GoogleAuthenticationError
- if isinstance(e, GoogleAuthenticationError) and e.auth_url:
- result = await _handle_cli_auth_flow(
- server, parsed["tool_name"], args, e, oauth_server
- )
- print(result)
- return 0
- raise
+ try:
+ try:
+ result = await run_tool(server, parsed["tool_name"], args)
+ print(result)
+ return 0
+ except Exception as e:
+ from auth.google_auth import GoogleAuthenticationError
+ if isinstance(e, GoogleAuthenticationError) and e.auth_url:
+ result = await _handle_cli_auth_flow(
+ server, parsed["tool_name"], args, e, oauth_server
+ )
+ print(result)
+ return 0
+ raise
+ finally:
+ cleanup_oauth_callback_server()Also applies to: 473-492
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@core/cli_handler.py` around lines 24 - 28, The CLI path that starts the
MinimalOAuthServer (via get_cli_oauth_port and set_cli_oauth_server) never
guarantees teardown on early returns; modify the run logic so the server is
shutdown and deregistered in a finally/cleanup block: after creating and
set_cli_oauth_server(server) ensure you call the server's stop/shutdown method
and then set_cli_oauth_server(None) on every exit path (including error/early
returns) so no threads/sockets or stale global remain; wrap the run branch that
starts the server in try/finally (or ensure each return first performs the same
cleanup) and reference MinimalOAuthServer, get_cli_oauth_port, and
set_cli_oauth_server when making the changes.
- Reset auth state between flows to prevent stale results - Sanitize OAuth error messages exposed to CLI/browser - Add finally block for OAuth server teardown - Fix callback route binding for 404 issue Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/cli_handler.py`:
- Line 437: The current print call prints result['user_id'] to stderr which
leaks PII; update the code that prints authentication feedback (the print(...)
statement referencing result['user_id'] in core/cli_handler.py) to avoid raw
account IDs—either emit a generic success message like "Authenticated
successfully" or print a masked identifier (e.g., replace most characters of the
address with asterisks while keeping only a non-identifying suffix). Ensure the
change touches the print(...) invocation that currently outputs
result['user_id'] and does not expose the full identifier in CLI transcripts or
event streams.
- Around line 472-475: The code currently calls _setup_cli_oauth_server()
unconditionally in the run path which attempts to bind the callback port
immediately; instead, change the flow to only configure/register the CLI OAuth
server object without starting it: have _setup_cli_oauth_server() return an
OAuthServer instance (or factory) but do not call its start()/serve() here, and
remove any immediate bind calls from the run-path; then update the
authentication code paths (the functions that actually perform browser OAuth or
token refresh) to call oauth_server.start() (or instantiate/start the server)
lazily only when a real interactive auth callback is required. Reference:
parsed["command"] == "run", _setup_cli_oauth_server(), and any
oauth_server.start()/start_server calls in auth routines.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a1fdb6fb-d7f1-4374-b591-08474a82b681
📒 Files selected for processing (2)
auth/oauth_callback_server.pycore/cli_handler.py
| def get_cli_oauth_port() -> int: | ||
| """ | ||
| Find an available port for the CLI OAuth callback server. | ||
| Tries ports 8000-8009, falls back to OS-assigned port. | ||
|
|
||
| Returns: | ||
| Available port number | ||
| """ | ||
| for port in range(8000, 8010): | ||
| try: | ||
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: | ||
| s.bind(("localhost", port)) | ||
| return port | ||
| except OSError: | ||
| continue | ||
|
|
||
| # Fallback: let OS assign a port | ||
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: | ||
| s.bind(("localhost", 0)) | ||
| return s.getsockname()[1] |
There was a problem hiding this comment.
Do not pick the callback port with a bind-and-release probe.
get_cli_oauth_port() only proves the port was free momentarily; it does not reserve it. Another local process can claim that port before MinimalOAuthServer.start() runs, and the later readiness probe can even accept that foreign listener as "started" if /oauth2callback returns anything other than 404. Try candidate ports by actually starting the server, then set GOOGLE_OAUTH_REDIRECT_URI from the port that bound successfully.
There was a problem hiding this comment.
Probably more minor than major, the likelihood that a port gets claimed in the moments between check and initialization seems low
There was a problem hiding this comment.
@taylorwilsdon, that's a fair point — the race window is genuinely tiny under normal conditions. On a typical developer machine or a dedicated Docker container, the gap between releasing the socket and uvicorn binding is unlikely to be contested. Downgrading this to minor is reasonable.
The slightly more interesting edge case remains: if some other process already happened to be listening on the chosen port before get_cli_oauth_port() runs, the readiness probe (which accepts any non-404 response) could mistake that foreign listener for the OAuth server. But again, port 8000 is usually free in the described headless-Docker use case, so this is a low-probability path in practice.
Happy to leave this as-is if you'd prefer — it's a pre-existing pattern and the practical risk is low.
🧠 Learnings used
Learnt from: hGriff0n
Repo: taylorwilsdon/google_workspace_mcp PR: 499
File: gkeep/keep_tools.py:417-425
Timestamp: 2026-02-21T18:00:37.834Z
Learning: When using asyncio.to_thread(callable, *args, **kwargs), pass a callable, not the result of a call. For methods on objects, pass the method reference (e.g., await asyncio.to_thread(service.some_method(...).execute)) rather than invoking it upfront (service.some_method(...).execute()). This ensures a callable is provided to to_thread, which will run it in a separate thread. This guideline applies broadly to all Python files in the Google Workspace MCP codebase (and similar Python projects); enforce by scanning for instances where a method is called immediately inside to_thread instead of passing the bound method/reference.
| if not result["success"]: | ||
| raise RuntimeError(f"Authentication failed: {result['error']}") | ||
|
|
||
| print(f"\nAuthenticated as: {result['user_id']}", file=sys.stderr) |
There was a problem hiding this comment.
Don't print the authenticated account ID to stderr.
This sends a raw Google account identifier into CLI transcripts and agent event streams. A generic success message, or at most a masked address, keeps the flow usable without leaking PII.
🔧 Suggested change
- print(f"\nAuthenticated as: {result['user_id']}", file=sys.stderr)
+ print("\nAuthentication completed successfully.", file=sys.stderr)As per coding guidelines "Never log or leak secrets, refresh tokens, or PII in exceptions or event streams."
📝 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.
| print(f"\nAuthenticated as: {result['user_id']}", file=sys.stderr) | |
| print("\nAuthentication completed successfully.", file=sys.stderr) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@core/cli_handler.py` at line 437, The current print call prints
result['user_id'] to stderr which leaks PII; update the code that prints
authentication feedback (the print(...) statement referencing result['user_id']
in core/cli_handler.py) to avoid raw account IDs—either emit a generic success
message like "Authenticated successfully" or print a masked identifier (e.g.,
replace most characters of the address with asterisks while keeping only a
non-identifying suffix). Ensure the change touches the print(...) invocation
that currently outputs result['user_id'] and does not expose the full identifier
in CLI transcripts or event streams.
There was a problem hiding this comment.
CLI mode is explicitly local to a user system so I would not consider this "major" - just a nit
There was a problem hiding this comment.
@taylorwilsdon, that's a fair point — in a strictly local CLI context the user is seeing their own account identifier on their own terminal, which is quite different from leaking PII into a shared server log or agent event stream. Downgrading this to a nit is reasonable. The suggestion to use a generic message still stands as a minor polish item, but it's not a blocker.
🧠 Learnings used
Learnt from: CR
Repo: taylorwilsdon/google_workspace_mcp PR: 0
File: .github/instructions/general.instructions.md:0-0
Timestamp: 2026-03-03T22:54:40.426Z
Learning: Applies to **/*.py : Never log or leak secrets, refresh tokens, or PII in exceptions or event streams.
Learnt from: hGriff0n
Repo: taylorwilsdon/google_workspace_mcp PR: 499
File: gkeep/keep_tools.py:417-425
Timestamp: 2026-02-21T18:00:37.834Z
Learning: When using asyncio.to_thread(callable, *args, **kwargs), pass a callable, not the result of a call. For methods on objects, pass the method reference (e.g., await asyncio.to_thread(service.some_method(...).execute)) rather than invoking it upfront (service.some_method(...).execute()). This ensures a callable is provided to to_thread, which will run it in a separate thread. This guideline applies broadly to all Python files in the Google Workspace MCP codebase (and similar Python projects); enforce by scanning for instances where a method is called immediately inside to_thread instead of passing the bound method/reference.
| if parsed["command"] == "run": | ||
| # Pre-start OAuth server before running tools that may need auth | ||
| oauth_server = _setup_cli_oauth_server() | ||
|
|
There was a problem hiding this comment.
Don't make every CLI tool run depend on starting the callback server.
This now fails the run path before tool execution whenever the local callback server cannot bind, even if cached credentials are already valid and no browser auth would happen. A safer shape is to configure/register the CLI OAuth server here, but defer start() until auth is actually requested.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@core/cli_handler.py` around lines 472 - 475, The code currently calls
_setup_cli_oauth_server() unconditionally in the run path which attempts to bind
the callback port immediately; instead, change the flow to only
configure/register the CLI OAuth server object without starting it: have
_setup_cli_oauth_server() return an OAuthServer instance (or factory) but do not
call its start()/serve() here, and remove any immediate bind calls from the
run-path; then update the authentication code paths (the functions that actually
perform browser OAuth or token refresh) to call oauth_server.start() (or
instantiate/start the server) lazily only when a real interactive auth callback
is required. Reference: parsed["command"] == "run", _setup_cli_oauth_server(),
and any oauth_server.start()/start_server calls in auth routines.
There was a problem hiding this comment.
🧩 Analysis chain
@taylorwilsdon, thanks for confirming. Here's a plan to lazily start the OAuth server only when auth is actually needed:
- Split
_setup_cli_oauth_serverinto a configure phase (port selection, env, reload, server instantiation, registration) and a start phase - Call configure unconditionally in the
runpath so the redirect URI is set before the tool runs - Move
oauth_server.start()into_handle_cli_auth_flow, called only when aGoogleAuthenticationErroris caught - Guard
oauth_server.stop()in thefinallyblock so it's safe to call even if the server was never started
🏁 Script executed:
#!/bin/bash
# Check if MinimalOAuthServer.stop() is already safe to call when not started
ast-grep --pattern 'def stop(self) {
$$$
}'
rg -n "def stop" auth/oauth_callback_server.py -A 10
# Also check if there's an is_running or similar attribute
rg -n "is_running\|_started\|_running\|self\._server" auth/oauth_callback_server.py
Length of output: 378
🏁 Script executed:
#!/bin/bash
rg -n "def stop\|def start\|is_running\|_started\|_running\|self\._server\b" auth/oauth_callback_server.py -A 5
Length of output: 0
🏁 Script executed:
#!/bin/bash
rg -n "is_running|def start|def stop|_started|_running" auth/oauth_callback_server.py -A 3
Length of output: 1495
stop() already guards with if not self.is_running: return, so the finally block is safe as-is. Here's the full lazy-start refactor:
- Rename
_setup_cli_oauth_server→_configure_cli_oauth_server, removing thestart()call - Call the configure function (not start) unconditionally so the redirect URI is registered before tool execution
- Move
oauth_server.start()into_handle_cli_auth_flow, only when auth is actually needed finally: oauth_server.stop()stays unchanged — already a no-op whenis_runningisFalse
🧠 Learnings used
Learnt from: hGriff0n
Repo: taylorwilsdon/google_workspace_mcp PR: 499
File: gkeep/keep_tools.py:417-425
Timestamp: 2026-02-21T18:00:37.834Z
Learning: When using asyncio.to_thread(callable, *args, **kwargs), pass a callable, not the result of a call. For methods on objects, pass the method reference (e.g., await asyncio.to_thread(service.some_method(...).execute)) rather than invoking it upfront (service.some_method(...).execute()). This ensures a callable is provided to to_thread, which will run it in a separate thread. This guideline applies broadly to all Python files in the Google Workspace MCP codebase (and similar Python projects); enforce by scanning for instances where a method is called immediately inside to_thread instead of passing the bound method/reference.
💡 Use @coderabbitai plan to create a plan for code changes
2 participants
Summary
Fixes three bugs that prevent the
--climode from completing OAuth authentication:safe_print()is a no-op in CLI mode, so tool module import failures were silently swallowed. Now prints to stderr in CLI mode and catches allExceptiontypes (not justModuleNotFoundError).sys.exit(0), killing the daemon OAuth callback server before the user could complete the browser flow. Now pre-starts the OAuth server, waits up to 5 minutes for the callback, and retries the tool on success.auth_urlonGoogleAuthenticationError:start_auth_flow()only returned the formatted message string. Theauth_urlattribute on the exception was alwaysNone. Now returns(message, auth_url)tuple so the CLI handler can access the URL directly.Also removes the redundant
access_type="offline"parameter fromflow.authorization_url()(already defaulted by google_auth_oauthlib).Changes
main.pyauth/google_auth.pyaccess_type, return(message, auth_url)tuple fromstart_auth_flow(), passauth_urlto exceptionauth/oauth_callback_server.pythreading.Eventcompletion signaling,wait_for_auth(),get_cli_oauth_port(),set_cli_oauth_server()core/cli_handler.pyGoogleAuthenticationError, wait for auth + retrycore/server.pystart_auth_flow()tuple returnTest plan
workspace-mcp --tools gmail docs sheets --clilists tools (no auth needed)workspace-mcp --tools gmail --cli list_gmail_labels --args '{"user_google_email":"..."}'with no cached credentials shows auth URL, waits for callback, completes auth, retries and prints labelsworkspace-mcp --tools docs sheets) unaffected🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes