feat(source): add timeout to add-research imports

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`, `--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` |

src/notebooklm/cli/source.py

@@ -540,24 +540,32 @@ 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,
     help="Start research and return immediately (use 'research status/wait' to monitor)",
 )
 @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.
 
     \b
     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:

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,
         )