From eeb2deac39b2b8f4ffe27df8bf5f5dad84373d66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= <258577966+voidborne-d@users.noreply.github.com> Date: Thu, 9 Apr 2026 02:51:42 +0000 Subject: [PATCH 1/3] feat(source): add timeout to add-research imports --- docs/cli-reference.md | 2 +- src/notebooklm/cli/source.py | 17 ++++++++++--- tests/unit/cli/test_source.py | 46 +++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 2d0f6df0..d29e8e0e 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -111,7 +111,7 @@ Supported source types: URLs, YouTube videos, files (PDF, text, Markdown, Word, | `list` | - | - | `source list` | | `add ` | URL/file/text | - | `source add "https://..."` | | `add-drive ` | Drive file ID | - | `source add-drive abc123 "Doc"` | -| `add-research <query>` | Search query | `--mode [fast|deep]`, `--from [web|drive]`, `--import-all`, `--no-wait` | `source add-research "AI" --mode deep --no-wait` | +| `add-research <query>` | Search query | `--mode [fast|deep]`, `--from [web|drive]`, `--import-all`, `--timeout FLOAT`, `--no-wait` | `source add-research "AI" --mode deep --timeout 600` | | `get <id>` | Source ID | - | `source get src123` | | `fulltext <id>` | Source ID | `--json`, `-o FILE` | `source fulltext src123 -o content.txt` | | `guide <id>` | Source ID | `--json` | `source guide src123` | diff --git a/src/notebooklm/cli/source.py b/src/notebooklm/cli/source.py index 1ec9dce3..34a13b55 100644 --- a/src/notebooklm/cli/source.py +++ b/src/notebooklm/cli/source.py @@ -540,6 +540,13 @@ async def _run(): help="Search mode (default: fast)", ) @click.option("--import-all", is_flag=True, help="Import all found sources") +@click.option( + "--timeout", + type=float, + default=1800.0, + show_default=True, + help="Retry budget in seconds when --import-all is used", +) @click.option( "--no-wait", is_flag=True, @@ -547,7 +554,7 @@ async def _run(): ) @with_client def source_add_research( - ctx, query, notebook_id, search_source, mode, import_all, no_wait, client_auth + ctx, query, notebook_id, search_source, mode, import_all, timeout, no_wait, client_auth ): """Search web or drive and add sources from results. @@ -555,9 +562,10 @@ def source_add_research( Examples: source add-research "machine learning" # Search web source add-research "project docs" --from drive # Search Google Drive - source add-research "AI papers" --mode deep # Deep search - source add-research "tutorials" --import-all # Auto-import all results - source add-research "topic" --mode deep --no-wait # Non-blocking deep search + source add-research "AI papers" --mode deep # Deep search + source add-research "tutorials" --import-all # Auto-import all results + source add-research "tutorials" --import-all --timeout 600 # Limit import retry budget + source add-research "topic" --mode deep --no-wait # Non-blocking deep search """ nb_id = require_notebook(notebook_id) @@ -606,6 +614,7 @@ async def _run(): nb_id_resolved, task_id, sources, + max_elapsed=timeout, ) console.print(f"[green]Imported {len(imported)} sources[/green]") else: diff --git a/tests/unit/cli/test_source.py b/tests/unit/cli/test_source.py index e6093f39..a2380084 100644 --- a/tests/unit/cli/test_source.py +++ b/tests/unit/cli/test_source.py @@ -666,6 +666,52 @@ def test_add_research_with_import_all_uses_retry_helper(self, runner, mock_auth) "nb_123", "task_123", [{"title": "Source 1", "url": "http://example.com"}], + max_elapsed=1800.0, + ) + + def test_add_research_with_import_all_passes_custom_timeout(self, runner, mock_auth): + with ( + patch_client_for_module("source") as mock_client_cls, + patch.object(source_module, "import_with_retry", new_callable=AsyncMock) as mock_import, + ): + mock_client = create_mock_client() + mock_client.research.start = AsyncMock(return_value={"task_id": "task_123"}) + mock_client.research.poll = AsyncMock( + return_value={ + "status": "completed", + "task_id": "task_123", + "sources": [{"title": "Source 1", "url": "http://example.com"}], + "report": "# Report", + } + ) + mock_import.return_value = [{"id": "src_1", "title": "Source 1"}] + mock_client_cls.return_value = mock_client + + with patch("notebooklm.cli.helpers.fetch_tokens", new_callable=AsyncMock) as mock_fetch: + mock_fetch.return_value = ("csrf", "session") + result = runner.invoke( + cli, + [ + "source", + "add-research", + "AI papers", + "--mode", + "deep", + "--import-all", + "--timeout", + "600", + "-n", + "nb_123", + ], + ) + + assert result.exit_code == 0 + mock_import.assert_awaited_once_with( + mock_client, + "nb_123", + "task_123", + [{"title": "Source 1", "url": "http://example.com"}], + max_elapsed=600.0, ) From 8e8d5b1038bdb773addfebbcdc8bb94b1de70c66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= <258577966+voidborne-d@users.noreply.github.com> Date: Thu, 9 Apr 2026 06:06:43 +0000 Subject: [PATCH 2/3] docs: fix add-research quick reference row --- docs/cli-reference.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cli-reference.md b/docs/cli-reference.md index d29e8e0e..5e36177a 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -111,7 +111,7 @@ Supported source types: URLs, YouTube videos, files (PDF, text, Markdown, Word, | `list` | - | - | `source list` | | `add <content>` | URL/file/text | - | `source add "https://..."` | | `add-drive <id> <title>` | Drive file ID | - | `source add-drive abc123 "Doc"` | -| `add-research <query>` | Search query | `--mode [fast|deep]`, `--from [web|drive]`, `--import-all`, `--timeout FLOAT`, `--no-wait` | `source add-research "AI" --mode deep --timeout 600` | +| `add-research <query>` | Search query | `--mode [fast\|deep]`, `--from [web\|drive]`, `--import-all`, `--timeout FLOAT`, `--no-wait` | `source add-research "AI" --mode deep --import-all --timeout 600` | | `get <id>` | Source ID | - | `source get src123` | | `fulltext <id>` | Source ID | `--json`, `-o FILE` | `source fulltext src123 -o content.txt` | | `guide <id>` | Source ID | `--json` | `source guide src123` | From 8b6428ac79a52ced9c2df921febb12b58acf1264 Mon Sep 17 00:00:00 2001 From: voidborne-d <voidborne@users.noreply.github.com> Date: Thu, 9 Apr 2026 11:07:00 +0000 Subject: [PATCH 3/3] docs: simplify timeout option listing --- docs/cli-reference.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 5e36177a..eadf9cc3 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -111,7 +111,7 @@ Supported source types: URLs, YouTube videos, files (PDF, text, Markdown, Word, | `list` | - | - | `source list` | | `add <content>` | URL/file/text | - | `source add "https://..."` | | `add-drive <id> <title>` | Drive file ID | - | `source add-drive abc123 "Doc"` | -| `add-research <query>` | Search query | `--mode [fast\|deep]`, `--from [web\|drive]`, `--import-all`, `--timeout FLOAT`, `--no-wait` | `source add-research "AI" --mode deep --import-all --timeout 600` | +| `add-research <query>` | Search query | `--mode [fast\|deep]`, `--from [web\|drive]`, `--import-all`, `--timeout`, `--no-wait` | `source add-research "AI" --mode deep --import-all --timeout 600` | | `get <id>` | Source ID | - | `source get src123` | | `fulltext <id>` | Source ID | `--json`, `-o FILE` | `source fulltext src123 -o content.txt` | | `guide <id>` | Source ID | `--json` | `source guide src123` |