Ensure Gemini Cloud Project sessions are closed #633
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds explicit ADC-authenticated session management and safe cleanup (using try/finally and contextlib.suppress) across project validation, health checks, and both standard and streaming chat completion flows; adds unit tests asserting sessions are closed. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Caller
participant Connector
participant ADC as ADC Auth
participant API as Gemini API
rect rgba(230,240,255,0.5)
note over Connector: _validate_project_access
Caller->>Connector: _validate_project_access()
Connector->>ADC: _get_adc_authorized_session()
ADC-->>Connector: auth_session
Connector->>API: GET project info (auth_session)
API-->>Connector: project details
alt finally
note right of Connector: suppress errors on close
Connector->>ADC: auth_session.close()
end
Connector-->>Caller: result/exception
end
rect rgba(235,255,235,0.5)
note over Connector: _perform_health_check
Caller->>Connector: _perform_health_check()
Connector->>ADC: _get_adc_authorized_session()
ADC-->>Connector: auth_session
Connector->>API: GET /health (auth_session)
API-->>Connector: 200 OK / error
note right of Connector: suppress errors on close
Connector->>ADC: auth_session.close()
Connector-->>Caller: True/False
end
rect rgba(255,245,230,0.5)
note over Connector: _chat_completions_standard
Caller->>Connector: completions(...)
Connector->>ADC: _get_adc_authorized_session()
ADC-->>Connector: auth_session
Connector->>API: POST /chat (auth_session)
API-->>Connector: response
note right of Connector: suppress errors on close
Connector->>ADC: auth_session.close()
Connector-->>Caller: completion
end
rect rgba(255,240,245,0.5)
note over Connector: _chat_completions_streaming
Caller->>Connector: streaming_completions(...)
Connector->>ADC: _get_adc_authorized_session()
ADC-->>Connector: auth_session
Connector->>API: POST /chat:stream (auth_session)
API-->>Connector: stream
loop generator yields
Connector-->>Caller: chunk
end
note right of Connector: on generator end/error, suppress close errors
Connector->>ADC: auth_session.close()
Connector-->>Caller: done
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used📓 Path-based instructions (2)**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
src/**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
🔇 Additional comments (5)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting
| auth_session = self._get_adc_authorized_session() | ||
| try: | ||
| # Use ADC for API calls (matches gemini CLI behavior for project-id auth) | ||
| auth_session = self._get_adc_authorized_session() | ||
|
|
||
| # Ensure project is onboarded for standard-tier | ||
| project_id = await self._ensure_project_onboarded(auth_session) |
There was a problem hiding this comment.
Moving auth_session = self._get_adc_authorized_session() outside the surrounding try means any failure in acquiring Application Default Credentials (for example google.auth.exceptions.DefaultCredentialsError) will now escape without being logged or wrapped in the function’s BackendError fallback. Callers of _chat_completions_standard previously only saw BackendError/AuthenticationError, so this regression will leak raw exceptions and bypass the existing error logging. Consider creating the session inside the try and closing it in finally with a guard so the original error handling remains intact.
Useful? React with 👍 / 👎.
| auth_session = self._get_adc_authorized_session() | ||
| stream_prepared = False | ||
| try: | ||
| # Use ADC for streaming API calls | ||
| auth_session = self._get_adc_authorized_session() | ||
|
|
||
| # Ensure project is onboarded for standard-tier | ||
| project_id = await self._ensure_project_onboarded(auth_session) |
There was a problem hiding this comment.
The streaming variant now obtains the ADC AuthorizedSession before entering the try/except. If _get_adc_authorized_session() raises (e.g. missing credentials or networking errors), the exception bypasses the function’s error logging and conversion to BackendError, and will propagate as a raw Exception. This diverges from prior behaviour and from the non-streaming code paths that callers rely on for consistent error semantics. Acquire the session inside the try and keep the finally guarded so it is still closed when an error occurs.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/connectors/gemini_cloud_project.py(9 hunks)tests/unit/connectors/test_gemini_cloud_project_resource_management.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Avoid using emojis in Python code comments or docstrings
Follow PEP 8 and use type hints for all functions
Import order must be: standard library, third-party, then local imports, separated by blank lines
Naming conventions: snake_case for variables and functions; PascalCase for classes
Error handling: use specific exceptions and include meaningful error messages
Prefer f-strings for string formatting
Files:
tests/unit/connectors/test_gemini_cloud_project_resource_management.pysrc/connectors/gemini_cloud_project.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Catch specific exceptions; avoid broad except Exception blocks
If a broad exception must be caught, log with exc_info=True and re-raise a specific custom exception
Use the most specific exception class from src.core.common.exceptions that accurately describes the error
Provide clear, helpful error messages and include relevant details in the details dictionary of custom exceptions
Files:
src/connectors/gemini_cloud_project.py
🧬 Code graph analysis (2)
tests/unit/connectors/test_gemini_cloud_project_resource_management.py (2)
src/connectors/gemini_cloud_project.py (2)
_validate_project_access(797-842)_perform_health_check(875-933)src/core/services/translation_service.py (1)
TranslationService(20-636)
src/connectors/gemini_cloud_project.py (5)
tests/unit/connectors/test_gemini_cloud_project_resource_management.py (1)
close(26-27)src/connectors/gemini_oauth_personal.py (1)
stream_generator(1620-1782)src/connectors/gemini.py (1)
stream_generator(292-322)src/core/domain/responses.py (1)
StreamingResponseEnvelope(27-53)tests/unit/openai_connector_tests/test_streaming_response.py (1)
content(57-58)
🪛 Ruff (0.13.3)
tests/unit/connectors/test_gemini_cloud_project_resource_management.py
49-49: Unused method argument: args
(ARG002)
49-49: Unused method argument: kwargs
(ARG002)
73-73: Possible hardcoded password assigned to: "token"
(S105)
75-75: Unused method argument: request
(ARG002)
76-76: Possible hardcoded password assigned to: "token"
(S105)
83-83: Unused function argument: url
(ARG001)
83-83: Unused function argument: headers
(ARG001)
83-83: Unused function argument: timeout
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: test (3.10)
| @pytest.fixture() | ||
| def connector() -> GeminiCloudProjectConnector: | ||
| cfg = AppConfig() | ||
| client = httpx.AsyncClient() | ||
| backend = GeminiCloudProjectConnector( | ||
| client, | ||
| cfg, | ||
| translation_service=TranslationService(), | ||
| gcp_project_id="test-project", | ||
| ) | ||
| backend.gemini_api_base_url = "https://example.com" | ||
| return backend |
There was a problem hiding this comment.
Ensure the AsyncClient is closed in the fixture
httpx.AsyncClient keeps background resources alive until aclose() runs; leaving it open in the fixture causes resource leaks and runtime warnings after the tests complete. Convert the fixture to an async fixture that async with-closes the client so each test frees the session.
-@pytest.fixture()
-def connector() -> GeminiCloudProjectConnector:
- cfg = AppConfig()
- client = httpx.AsyncClient()
- backend = GeminiCloudProjectConnector(
- client,
- cfg,
- translation_service=TranslationService(),
- gcp_project_id="test-project",
- )
- backend.gemini_api_base_url = "https://example.com"
- return backend
+@pytest_asyncio.fixture()
+async def connector() -> GeminiCloudProjectConnector:
+ cfg = AppConfig()
+ async with httpx.AsyncClient() as client:
+ backend = GeminiCloudProjectConnector(
+ client,
+ cfg,
+ translation_service=TranslationService(),
+ gcp_project_id="test-project",
+ )
+ backend.gemini_api_base_url = "https://example.com"
+ yield backendRemember to import pytest_asyncio alongside the other pytest imports.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @pytest.fixture() | |
| def connector() -> GeminiCloudProjectConnector: | |
| cfg = AppConfig() | |
| client = httpx.AsyncClient() | |
| backend = GeminiCloudProjectConnector( | |
| client, | |
| cfg, | |
| translation_service=TranslationService(), | |
| gcp_project_id="test-project", | |
| ) | |
| backend.gemini_api_base_url = "https://example.com" | |
| return backend | |
| @pytest_asyncio.fixture() | |
| async def connector() -> GeminiCloudProjectConnector: | |
| cfg = AppConfig() | |
| async with httpx.AsyncClient() as client: | |
| backend = GeminiCloudProjectConnector( | |
| client, | |
| cfg, | |
| translation_service=TranslationService(), | |
| gcp_project_id="test-project", | |
| ) | |
| backend.gemini_api_base_url = "https://example.com" | |
| yield backend |
🤖 Prompt for AI Agents
In tests/unit/connectors/test_gemini_cloud_project_resource_management.py around
lines 30 to 41, the fixture creates an httpx.AsyncClient but never closes it;
convert the fixture to an async fixture (use pytest_asyncio import) and create
the client with an async context manager (async with httpx.AsyncClient() as
client:) or call await client.aclose() in a finally block, then yield the
GeminiCloudProjectConnector instance so tests receive it and the AsyncClient is
properly closed after each test.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
Summary
Testing
https://chatgpt.com/codex/tasks/task_e_68ec2685724083339dafb8594c547032
Summary by CodeRabbit
Bug Fixes
Tests