diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 2d0f6df0..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 ` | 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`, `--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` | 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, )