Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
voidborne-d wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
eeb2dea
feat(source): add timeout to add-research imports
voidborne-d Apr 9, 2026
8e8d5b1
docs: fix add-research quick reference row
voidborne-d Apr 9, 2026
8b6428a
docs: simplify timeout option listing
voidborne Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/cli-reference.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`, `--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` |
Expand Down
17 changes: 13 additions & 4 deletions src/notebooklm/cli/source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
Comment on lines +544 to +549
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For consistency with other commands in the source group (such as source wait), consider using type=int for the --timeout option. Additionally, since the command performs both research polling and source importing, the timeout should ideally apply to the entire operation rather than just the import phase.

Suggested change
"--timeout",
type=float,
default=1800.0,
show_default=True,
help="Retry budget in seconds when --import-all is used",
)
"--timeout",
type=int,
default=1800,
show_default=True,
help="Maximum seconds to wait for research and import",

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

Expand Down Expand Up @@ -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:
Expand Down
46 changes: 46 additions & 0 deletions tests/unit/cli/test_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
Loading