Skip to content

fix: single-flight Discovery Engine mode detection#6106

Open
he-yufeng wants to merge 1 commit into
google:mainfrom
he-yufeng:fix/discovery-engine-search-mode-singleflight
Open

fix: single-flight Discovery Engine mode detection#6106
he-yufeng wants to merge 1 commit into
google:mainfrom
he-yufeng:fix/discovery-engine-search-mode-singleflight

Conversation

@he-yufeng

Copy link
Copy Markdown
Contributor

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:

DiscoveryEngineSearchTool(search_result_mode=None) auto-detects the result mode by trying CHUNKS first and falling back to DOCUMENTS when Discovery Engine reports a structured datastore. On a fresh shared tool instance, concurrent first calls can all enter that probe path before _search_result_mode is cached, so a structured datastore pays one failing CHUNKS request per concurrent caller.

Solution:

Make the auto-detect probe single-flight per tool instance. The first cold caller detects the datastore mode and caches it; concurrent callers re-check the cached mode after the lock and go straight to the detected mode. The successful CHUNKS case is cached too, since the result mode is a datastore property rather than a query property.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Passed locally:

$env:PYTHONPATH=(Resolve-Path src).Path; python -m pytest tests\unittests\tools\test_discovery_engine_search_tool.py -q
26 passed, 4 warnings in 3.56s

Additional checks:

python -m py_compile src\google\adk\tools\discovery_engine_search_tool.py tests\unittests\tools\test_discovery_engine_search_tool.py
python -m ruff check src\google\adk\tools\discovery_engine_search_tool.py tests\unittests\tools\test_discovery_engine_search_tool.py
python -m pyink --check src\google\adk\tools\discovery_engine_search_tool.py tests\unittests\tools\test_discovery_engine_search_tool.py
git diff --check

Manual End-to-End (E2E) Tests:

Not run. The regression test uses a fake Discovery Engine client to reproduce the concurrent cold-start behavior without live Discovery Engine quota or billing.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

The new test starts multiple threads against one fresh tool instance. The fake client delays and rejects the CHUNKS probe with the structured-datastore error, then succeeds in DOCUMENTS mode. The fixed behavior performs one CHUNKS probe total and one DOCUMENTS request per caller.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DiscoveryEngineSearchTool auto-detect sends redundant CHUNKS searches under concurrent first use

1 participant