feat(auth): add fresh login and logout flow#270
feat(auth): add fresh login and logout flow#270voidborne-d wants to merge 3 commits intoteng-lin:mainfrom
Conversation
📝 WalkthroughWalkthroughAdds a Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant CLI as Login CLI
participant FSys as File System
participant Browser as Playwright Browser
participant Auth as Auth Handler
rect rgba(100, 150, 200, 0.5)
Note over User,CLI: Fresh login flow
User->>CLI: run `notebooklm login --fresh`
CLI->>FSys: _clear_auth_files(storage_path, browser_profile)
FSys-->>CLI: list of removed paths
CLI->>Browser: launch with `user_data_dir` (profile dir)
Browser->>Auth: perform authentication flow
Auth-->>CLI: "Authentication saved"
end
rect rgba(200, 150, 100, 0.5)
Note over Browser,CLI: Page recovery flow
CLI->>Browser: page.goto(login_url)
Browser-->>CLI: TargetClosedError (page/context/browser closed)
CLI->>Browser: _get_recovered_page(context)
Browser-->>CLI: new live `page` returned
CLI->>CLI: backoff sleep & retry
CLI->>Browser: page.goto(login_url) retry
Browser->>Auth: continue login
Auth-->>CLI: "Authentication saved"
end
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly Related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/notebooklm/cli/session.py (1)
345-369: Consider catching broader OSError in callers.
shutil.rmtreecan raiseOSErrorfor various reasons beyondPermissionError(e.g., read-only files within the directory, filesystem errors). The callers at lines 450 and 1037 only catchPermissionError, which may leave some edge cases unhandled.♻️ Suggested fix: catch OSError in callers
- except PermissionError as e: + except OSError as e: console.print( "[red]Could not clear the saved browser session.[/red]\n" "Close any running NotebookLM/Chromium windows and try again.\n"Apply similar change to the
auth_logoutcommand at line 1037.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/notebooklm/cli/session.py` around lines 345 - 369, The callers of _clear_auth_files currently only catch PermissionError and thus miss other filesystem failures from shutil.rmtree/path.unlink; update the exception handling in the caller around the call to _clear_auth_files (and the auth_logout command) to catch OSError (e.g., except OSError as e) instead of or in addition to PermissionError, and ensure you log or surface the caught exception (include the exception object) so failures like read-only files or other FS errors are handled and diagnosable.tests/unit/cli/test_session.py (1)
690-711: Test methods misplaced insideTestAuthCommandsclass.
test_status_json_output_no_contextandtest_status_handles_corrupted_context_file(lines 690-711) appear to be status command tests that ended up inside theTestAuthCommandsclass. They should be moved toTestStatusCommandfor proper organization.♻️ Suggested fix: move tests to correct class
Move these two test methods from
TestAuthCommandstoTestStatusCommand(after line 648) to maintain logical grouping.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/cli/test_session.py` around lines 690 - 711, Two test methods, test_status_json_output_no_context and test_status_handles_corrupted_context_file, are currently defined inside the TestAuthCommands class but belong to the TestStatusCommand group; move both method definitions out of TestAuthCommands and add them to the TestStatusCommand class so that status-related tests are grouped together (remove the methods from TestAuthCommands and paste them into TestStatusCommand).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/notebooklm/cli/session.py`:
- Around line 345-369: The callers of _clear_auth_files currently only catch
PermissionError and thus miss other filesystem failures from
shutil.rmtree/path.unlink; update the exception handling in the caller around
the call to _clear_auth_files (and the auth_logout command) to catch OSError
(e.g., except OSError as e) instead of or in addition to PermissionError, and
ensure you log or surface the caught exception (include the exception object) so
failures like read-only files or other FS errors are handled and diagnosable.
In `@tests/unit/cli/test_session.py`:
- Around line 690-711: Two test methods, test_status_json_output_no_context and
test_status_handles_corrupted_context_file, are currently defined inside the
TestAuthCommands class but belong to the TestStatusCommand group; move both
method definitions out of TestAuthCommands and add them to the TestStatusCommand
class so that status-related tests are grouped together (remove the methods from
TestAuthCommands and paste them into TestStatusCommand).
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f9b3112e-4764-41a2-b71b-4ad4aec3f69d
📒 Files selected for processing (2)
src/notebooklm/cli/session.pytests/unit/cli/test_session.py
There was a problem hiding this comment.
Code Review
This pull request introduces a --fresh flag to the login command and a new logout command to manage authentication state, alongside a recovery mechanism for Playwright sessions closed during account switching. It also adds corresponding unit tests. Feedback suggests improving file deletion robustness regarding symbolic links, adding return type hints, and simplifying defensive attribute access on Playwright objects.
src/notebooklm/cli/session.py
Outdated
| if path.is_dir(): | ||
| shutil.rmtree(path) | ||
| else: | ||
| path.unlink() |
There was a problem hiding this comment.
Using shutil.rmtree on a path that is a symbolic link to a directory will fail with an OSError in Python 3.x. It is safer to check if the path is a symbolic link first and use unlink() to remove the link itself without affecting the target. Additionally, using missing_ok=True in unlink() (available in Python 3.8+) prevents race conditions if the file is deleted by another process between the existence check and the deletion attempt.
| if path.is_dir(): | |
| shutil.rmtree(path) | |
| else: | |
| path.unlink() | |
| if path.is_symlink(): | |
| path.unlink() | |
| elif path.is_dir(): | |
| shutil.rmtree(path) | |
| else: | |
| path.unlink(missing_ok=True) |
src/notebooklm/cli/session.py
Outdated
| return removed | ||
|
|
||
|
|
||
| def _get_recovered_page(context: Any, current_page: Any): |
There was a problem hiding this comment.
src/notebooklm/cli/session.py
Outdated
|
|
||
| def _get_recovered_page(context: Any, current_page: Any): | ||
| """Recover a live page when the original Playwright page was closed.""" | ||
| pages = [page for page in context.pages if not getattr(page, "is_closed", lambda: False)()] |
There was a problem hiding this comment.
The use of getattr(page, "is_closed", lambda: False)() is unnecessarily defensive since context.pages contains Playwright Page objects which are guaranteed to have an is_closed() method. This can be simplified to a direct method call.
| pages = [page for page in context.pages if not getattr(page, "is_closed", lambda: False)()] | |
| pages = [page for page in context.pages if not page.is_closed()] |
src/notebooklm/cli/session.py
Outdated
| pages = [page for page in context.pages if not getattr(page, "is_closed", lambda: False)()] | ||
| if pages: | ||
| return pages[0] | ||
| if current_page and not getattr(current_page, "is_closed", lambda: False)(): |
There was a problem hiding this comment.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/notebooklm/cli/session.py`:
- Around line 1031-1052: The logout command (auth_logout) only clears file-based
state via _clear_auth_files(get_storage_path(), get_browser_profile_dir()) but
doesn't handle inline env-based auth (NOTEBOOKLM_AUTH_JSON), so update
auth_logout to check os.environ for NOTEBOOKLM_AUTH_JSON after attempting to
remove files and, if present, print a clear warning via console that the env var
takes precedence and the user must unset it (or restart their shell) to fully
log out; keep existing file-removal behavior and messages, but add this
detection and warning when NOTEBOOKLM_AUTH_JSON is set.
- Around line 526-549: The page recovery is only used during initial navigation;
after the user presses ENTER and inside the post-login cookie loop you must
replace the stale page with _get_recovered_page(context, page) before retrying
page.goto or other page ops. Specifically, after the ENTER prompt (where the
user resumes login) call page = _get_recovered_page(context, page) to ensure
page is fresh; and inside the cookie-loop exception handler that catches
TargetClosedError, call page = _get_recovered_page(context, page) (and adjust
retries/backoff as the surrounding login retry logic uses) before attempting
page.goto(url, wait_until="commit") again so the loop doesn't reuse a closed
page. Ensure the updated page variable is used thereafter.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c6db284e-91d9-4069-855d-4872ab1aa7a4
📒 Files selected for processing (2)
src/notebooklm/cli/session.pytests/unit/cli/test_session.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unit/cli/test_session.py
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 `@src/notebooklm/cli/session.py`:
- Around line 345-357: _clear_auth_files currently always appends legacy
home-root auth paths (legacy_storage and legacy_browser) which causes
logout/--fresh on non-default profiles to wipe the default profile; change
_clear_auth_files(signature) to accept a boolean parameter clear_legacy_default
(default False) and only append legacy_storage and legacy_browser when
clear_legacy_default is True, then update the two call sites that invoke
_clear_auth_files to pass clear_legacy_default=True only when the resolved
profile name equals "default" so legacy cleanup is limited to the default
profile.
- Around line 360-368: The loop over paths_to_remove incorrectly skips broken
symlinks because Path.exists() is False for dangling links; change the existence
guard to detect symlinks too (e.g., use path.lexists() or test (path.exists() or
path.is_symlink())) so dangling symlinks are unlinked; update the loop around
paths_to_remove (the block that calls path.unlink(), shutil.rmtree(path), and
path.unlink(missing_ok=True)) to first consider path.is_symlink() or use
lexists() before continuing.
In `@tests/unit/cli/test_session.py`:
- Around line 365-406: The test never exercises the "recover after ENTER" path
because login() reads context.pages[0] before the ENTER prompt but the fixture
already exposes recovered_page there; fix by making mock_context.pages yield
initial_page (closed) on the first access so login() picks the closed page, and
then return recovered_page (open) after the ENTER prompt so the recovery branch
runs; implement this by changing mock_context.pages to a side-effecting
attribute (e.g., a property or a lambda/list factory) that returns
[initial_page] on the first call and [recovered_page] on subsequent calls,
keeping the existing recovered_goto, recovered_page, initial_page, mock_context
and the login() call names intact.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8b6fcf91-881b-41d0-91e9-e46a220c9541
📒 Files selected for processing (2)
src/notebooklm/cli/session.pytests/unit/cli/test_session.py
| def _clear_auth_files(storage_path: Path, browser_profile: Path) -> list[Path]: | ||
| """Delete auth files/directories for the active profile. | ||
|
|
||
| Also checks legacy default-profile paths when the resolved paths are profile-based. | ||
| Returns the paths that were actually removed. | ||
| """ | ||
| paths_to_remove = [storage_path, browser_profile] | ||
|
|
||
| legacy_storage = get_home_dir() / "storage_state.json" | ||
| legacy_browser = get_home_dir() / "browser_profile" | ||
| for legacy_path in (legacy_storage, legacy_browser): | ||
| if legacy_path not in paths_to_remove: | ||
| paths_to_remove.append(legacy_path) |
There was a problem hiding this comment.
Limit legacy cleanup to the default profile.
_clear_auth_files() always appends the legacy home-root auth paths, so login --fresh / auth logout on a non-default profile also wipes the legacy default profile's session. That breaks profile isolation while the command text says it only affects the active profile.
Proposed fix
-def _clear_auth_files(storage_path: Path, browser_profile: Path) -> list[Path]:
+def _clear_auth_files(
+ storage_path: Path, browser_profile: Path, *, clear_legacy_default: bool = False
+) -> list[Path]:
"""Delete auth files/directories for the active profile.
Also checks legacy default-profile paths when the resolved paths are profile-based.
Returns the paths that were actually removed.
"""
paths_to_remove = [storage_path, browser_profile]
- legacy_storage = get_home_dir() / "storage_state.json"
- legacy_browser = get_home_dir() / "browser_profile"
- for legacy_path in (legacy_storage, legacy_browser):
- if legacy_path not in paths_to_remove:
- paths_to_remove.append(legacy_path)
+ if clear_legacy_default:
+ home_dir = get_home_dir()
+ for legacy_path in (home_dir / "storage_state.json", home_dir / "browser_profile"):
+ if legacy_path not in paths_to_remove:
+ paths_to_remove.append(legacy_path)Update the two call sites to pass clear_legacy_default=True only when the resolved profile is "default".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooklm/cli/session.py` around lines 345 - 357, _clear_auth_files
currently always appends legacy home-root auth paths (legacy_storage and
legacy_browser) which causes logout/--fresh on non-default profiles to wipe the
default profile; change _clear_auth_files(signature) to accept a boolean
parameter clear_legacy_default (default False) and only append legacy_storage
and legacy_browser when clear_legacy_default is True, then update the two call
sites that invoke _clear_auth_files to pass clear_legacy_default=True only when
the resolved profile name equals "default" so legacy cleanup is limited to the
default profile.
| for path in paths_to_remove: | ||
| if not path.exists(): | ||
| continue | ||
| if path.is_symlink(): | ||
| path.unlink() | ||
| elif path.is_dir(): | ||
| shutil.rmtree(path) | ||
| else: | ||
| path.unlink(missing_ok=True) |
There was a problem hiding this comment.
Handle broken symlinks before the exists() guard.
A dangling symlink makes Path.exists() return False, so this loop skips it instead of unlinking it. That leaves stale storage_state.json / browser_profile entries behind, and the later login setup can fail with a pre-existing path.
Proposed fix
removed: list[Path] = []
for path in paths_to_remove:
- if not path.exists():
- continue
if path.is_symlink():
path.unlink()
+ removed.append(path)
+ continue
+ if not path.exists():
+ continue
elif path.is_dir():
shutil.rmtree(path)
else:
path.unlink(missing_ok=True)
removed.append(path)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooklm/cli/session.py` around lines 360 - 368, The loop over
paths_to_remove incorrectly skips broken symlinks because Path.exists() is False
for dangling links; change the existence guard to detect symlinks too (e.g., use
path.lexists() or test (path.exists() or path.is_symlink())) so dangling
symlinks are unlinked; update the loop around paths_to_remove (the block that
calls path.unlink(), shutil.rmtree(path), and path.unlink(missing_ok=True)) to
first consider path.is_symlink() or use lexists() before continuing.
| with ( | ||
| patch("notebooklm.cli.session._ensure_chromium_installed"), | ||
| patch("playwright.sync_api.sync_playwright") as mock_pw, | ||
| patch("notebooklm.cli.session.get_storage_path", return_value=storage_file), | ||
| patch("notebooklm.cli.session.get_browser_profile_dir", return_value=browser_profile), | ||
| patch("notebooklm.cli.session._sync_server_language_to_config"), | ||
| patch("builtins.input", return_value=""), | ||
| ): | ||
| mock_context = MagicMock() | ||
| initial_page = MagicMock() | ||
| initial_page.url = "https://notebooklm.google.com/" | ||
| initial_page.is_closed.return_value = True | ||
| recovered_page = MagicMock() | ||
| recovered_page.url = "https://notebooklm.google.com/" | ||
| recovered_page.is_closed.return_value = False | ||
| mock_context.pages = [recovered_page] | ||
| mock_context.storage_state.side_effect = lambda path: Path(path).write_text("{}") | ||
|
|
||
| cookie_retry = {"raised": False} | ||
|
|
||
| def recovered_goto(url, **kwargs): | ||
| if ( | ||
| url == "https://accounts.google.com/" | ||
| and kwargs.get("wait_until") == "commit" | ||
| and not cookie_retry["raised"] | ||
| ): | ||
| cookie_retry["raised"] = True | ||
| raise PlaywrightError( | ||
| "Page.goto: Target page, context or browser has been closed" | ||
| ) | ||
|
|
||
| recovered_page.goto.side_effect = recovered_goto | ||
| mock_launch = ( | ||
| mock_pw.return_value.__enter__.return_value.chromium.launch_persistent_context | ||
| ) | ||
| mock_launch.return_value = mock_context | ||
|
|
||
| result = runner.invoke(cli, ["login"]) | ||
|
|
||
| assert result.exit_code == 0 | ||
| assert recovered_page.goto.call_count >= 3 | ||
| assert "Authentication saved" in result.output |
There was a problem hiding this comment.
This test never reaches the post-ENTER recovery branch.
login() grabs context.pages[0] before the ENTER prompt. Since the fixture already exposes recovered_page there, initial_page is dead setup and Line 581 just returns the same open page. So this test currently verifies the cookie-loop retry only, not the new "recover after ENTER" path.
Proposed fix
with (
patch("notebooklm.cli.session._ensure_chromium_installed"),
patch("playwright.sync_api.sync_playwright") as mock_pw,
patch("notebooklm.cli.session.get_storage_path", return_value=storage_file),
patch("notebooklm.cli.session.get_browser_profile_dir", return_value=browser_profile),
patch("notebooklm.cli.session._sync_server_language_to_config"),
- patch("builtins.input", return_value=""),
+ patch("builtins.input") as mock_input,
):
mock_context = MagicMock()
initial_page = MagicMock()
initial_page.url = "https://notebooklm.google.com/"
initial_page.is_closed.return_value = True
recovered_page = MagicMock()
recovered_page.url = "https://notebooklm.google.com/"
recovered_page.is_closed.return_value = False
- mock_context.pages = [recovered_page]
+ mock_context.pages = [initial_page]
mock_context.storage_state.side_effect = lambda path: Path(path).write_text("{}")
cookie_retry = {"raised": False}
+ def finish_login(_prompt=""):
+ mock_context.pages = [recovered_page]
+ return ""
+
+ mock_input.side_effect = finish_login
+
def recovered_goto(url, **kwargs):
if (
url == "https://accounts.google.com/"
and kwargs.get("wait_until") == "commit"
and not cookie_retry["raised"]📝 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.
| with ( | |
| patch("notebooklm.cli.session._ensure_chromium_installed"), | |
| patch("playwright.sync_api.sync_playwright") as mock_pw, | |
| patch("notebooklm.cli.session.get_storage_path", return_value=storage_file), | |
| patch("notebooklm.cli.session.get_browser_profile_dir", return_value=browser_profile), | |
| patch("notebooklm.cli.session._sync_server_language_to_config"), | |
| patch("builtins.input", return_value=""), | |
| ): | |
| mock_context = MagicMock() | |
| initial_page = MagicMock() | |
| initial_page.url = "https://notebooklm.google.com/" | |
| initial_page.is_closed.return_value = True | |
| recovered_page = MagicMock() | |
| recovered_page.url = "https://notebooklm.google.com/" | |
| recovered_page.is_closed.return_value = False | |
| mock_context.pages = [recovered_page] | |
| mock_context.storage_state.side_effect = lambda path: Path(path).write_text("{}") | |
| cookie_retry = {"raised": False} | |
| def recovered_goto(url, **kwargs): | |
| if ( | |
| url == "https://accounts.google.com/" | |
| and kwargs.get("wait_until") == "commit" | |
| and not cookie_retry["raised"] | |
| ): | |
| cookie_retry["raised"] = True | |
| raise PlaywrightError( | |
| "Page.goto: Target page, context or browser has been closed" | |
| ) | |
| recovered_page.goto.side_effect = recovered_goto | |
| mock_launch = ( | |
| mock_pw.return_value.__enter__.return_value.chromium.launch_persistent_context | |
| ) | |
| mock_launch.return_value = mock_context | |
| result = runner.invoke(cli, ["login"]) | |
| assert result.exit_code == 0 | |
| assert recovered_page.goto.call_count >= 3 | |
| assert "Authentication saved" in result.output | |
| with ( | |
| patch("notebooklm.cli.session._ensure_chromium_installed"), | |
| patch("playwright.sync_api.sync_playwright") as mock_pw, | |
| patch("notebooklm.cli.session.get_storage_path", return_value=storage_file), | |
| patch("notebooklm.cli.session.get_browser_profile_dir", return_value=browser_profile), | |
| patch("notebooklm.cli.session._sync_server_language_to_config"), | |
| patch("builtins.input") as mock_input, | |
| ): | |
| mock_context = MagicMock() | |
| initial_page = MagicMock() | |
| initial_page.url = "https://notebooklm.google.com/" | |
| initial_page.is_closed.return_value = True | |
| recovered_page = MagicMock() | |
| recovered_page.url = "https://notebooklm.google.com/" | |
| recovered_page.is_closed.return_value = False | |
| mock_context.pages = [initial_page] | |
| mock_context.storage_state.side_effect = lambda path: Path(path).write_text("{}") | |
| cookie_retry = {"raised": False} | |
| def finish_login(_prompt=""): | |
| mock_context.pages = [recovered_page] | |
| return "" | |
| mock_input.side_effect = finish_login | |
| def recovered_goto(url, **kwargs): | |
| if ( | |
| url == "https://accounts.google.com/" | |
| and kwargs.get("wait_until") == "commit" | |
| and not cookie_retry["raised"] | |
| ): | |
| cookie_retry["raised"] = True | |
| raise PlaywrightError( | |
| "Page.goto: Target page, context or browser has been closed" | |
| ) | |
| recovered_page.goto.side_effect = recovered_goto | |
| mock_launch = ( | |
| mock_pw.return_value.__enter__.return_value.chromium.launch_persistent_context | |
| ) | |
| mock_launch.return_value = mock_context | |
| result = runner.invoke(cli, ["login"]) | |
| assert result.exit_code == 0 | |
| assert recovered_page.goto.call_count >= 3 | |
| assert "Authentication saved" in result.output |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/cli/test_session.py` around lines 365 - 406, The test never
exercises the "recover after ENTER" path because login() reads context.pages[0]
before the ENTER prompt but the fixture already exposes recovered_page there;
fix by making mock_context.pages yield initial_page (closed) on the first access
so login() picks the closed page, and then return recovered_page (open) after
the ENTER prompt so the recovery branch runs; implement this by changing
mock_context.pages to a side-effecting attribute (e.g., a property or a
lambda/list factory) that returns [initial_page] on the first call and
[recovered_page] on subsequent calls, keeping the existing recovered_goto,
recovered_page, initial_page, mock_context and the login() call names intact.
Summary
notebooklm login --freshto clear saved browser/session state before launching Playwrightnotebooklm auth logoutto remove saved auth files for the active profile, with legacy default-profile cleanup tooTesting
PYTHONPATH=src python3 -m pytest tests/unit/cli/test_session.py -qFixes #246
Summary by CodeRabbit
New Features
--freshflag to login to clear saved authentication state before signing in.auth logoutcommand to remove saved authentication/session data with clear success or "nothing to remove" feedback and a warning if inline auth remains set.Bug Fixes
--freshif account switching repeatedly fails.