Fix OpenAI per-call API key handling #628
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 per-call Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant C as Client
participant OC as OpenAIConnector
participant ID as Identity Resolver
participant HC as Health Check
participant OA as OpenAI API
C->>OC: chat_completions(request_data, identity?, api_key?)
activate OC
note right of OC #E6F4EA: Determine effective_api_key\n(if per-call provided)
OC->>OC: Normalize request_data
OC->>ID: Resolve per-request identity/headers
OC->>OC: Build headers (merge base + per-request + effective_api_key)
OC->>OC: _prepare_payload(...)
OC->>HC: _ensure_healthy(effective_api_key?)
alt Non-streaming
OC->>OA: POST /chat/completions (headers, payload)
OA-->>OC: Response
OC-->>C: ResponseEnvelope
else Streaming
OC->>OA: POST /chat/completions (stream)
OA-->>OC: Event stream
OC-->>C: StreamingResponseEnvelope
end
note left of OC #FFF4E6: finally: restore connector scope (api_key/identity)
deactivate OC
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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
src/connectors/openai.py
Outdated
| processed_messages: list[Any], | ||
| effective_model: str, | ||
| identity: IAppIdentityConfig | None = None, | ||
| api_key: str | None = None, | ||
| **kwargs: Any, | ||
| ) -> ResponseEnvelope | StreamingResponseEnvelope: | ||
| # Perform health check if enabled (for subclasses that support it) | ||
| await self._ensure_healthy() | ||
|
|
||
| # request_data is expected to be a domain ChatRequest (or subclass like CanonicalChatRequest) | ||
| # (the frontend controller converts from frontend-specific format to domain format) | ||
| # Backends should ONLY convert FROM domain TO backend-specific format | ||
| # Type assertion: we know from architectural design that request_data is ChatRequest-like | ||
| from typing import cast | ||
|
|
||
| from src.core.domain.chat import CanonicalChatRequest, ChatRequest | ||
|
|
||
| if not isinstance(request_data, ChatRequest): | ||
| raise TypeError( | ||
| f"Expected ChatRequest or CanonicalChatRequest, got {type(request_data).__name__}. " | ||
| "Backend connectors should only receive domain-format requests." | ||
| ) | ||
| # Cast to CanonicalChatRequest for mypy compatibility with _prepare_payload signature | ||
| domain_request: CanonicalChatRequest = cast(CanonicalChatRequest, request_data) | ||
| # Allow callers to supply a one-off API key (e.g., multi-tenant flows). | ||
| # Temporarily replace the connector-level key for the duration of this | ||
| # call so that header construction and health checks use it. | ||
| original_api_key = self.api_key | ||
| if api_key is not None: | ||
| self.api_key = api_key | ||
|
|
||
| # Ensure identity headers are scoped to the current request only. | ||
| self.identity = identity | ||
|
|
||
| # Prepare the payload using a helper so subclasses and tests can | ||
| # override or patch payload construction logic easily. | ||
| payload = await self._prepare_payload( | ||
| domain_request, processed_messages, effective_model | ||
| ) | ||
| headers_override = kwargs.pop("headers_override", None) | ||
| headers: dict[str, str] | None = None | ||
| # Perform health check if enabled (for subclasses that support it) | ||
| try: | ||
| await self._ensure_healthy() | ||
|
|
||
| if headers_override is not None: | ||
| # Avoid mutating the caller-provided mapping while preserving any | ||
| # Authorization header we compute from the configured API key. | ||
| headers = dict(headers_override) | ||
| # request_data is expected to be a domain ChatRequest (or subclass like CanonicalChatRequest) | ||
| # (the frontend controller converts from frontend-specific format to domain format) | ||
| # Backends should ONLY convert FROM domain TO backend-specific format | ||
| # Type assertion: we know from architectural design that request_data is ChatRequest-like | ||
| from typing import cast | ||
|
|
||
| try: | ||
| base_headers = self.get_headers() | ||
| except Exception: | ||
| base_headers = None | ||
| from src.core.domain.chat import CanonicalChatRequest, ChatRequest | ||
|
|
||
| if base_headers: | ||
| merged_headers = dict(base_headers) | ||
| merged_headers.update(headers) | ||
| headers = merged_headers | ||
| else: | ||
| try: | ||
| # Always update the cached identity so that per-request | ||
| # identity headers do not leak between calls. Downstream | ||
| # callers rely on identity-specific headers being scoped to | ||
| # a single request. | ||
| self.identity = identity | ||
| headers = self.get_headers() | ||
| except Exception: | ||
| headers = None | ||
| if not isinstance(request_data, ChatRequest): | ||
| raise TypeError( | ||
| f"Expected ChatRequest or CanonicalChatRequest, got {type(request_data).__name__}. " | ||
| "Backend connectors should only receive domain-format requests." | ||
| ) | ||
| # Cast to CanonicalChatRequest for mypy compatibility with _prepare_payload signature | ||
| domain_request: CanonicalChatRequest = cast(CanonicalChatRequest, request_data) | ||
|
|
||
| api_base = kwargs.get("openai_url") or self.api_base_url | ||
| url = f"{api_base.rstrip('/')}/chat/completions" | ||
| # Ensure identity headers are scoped to the current request only. | ||
| self.identity = identity | ||
|
|
||
| if domain_request.stream: | ||
| # Return a domain-level streaming envelope (raw bytes iterator) | ||
| try: | ||
| content_iterator = await self._handle_streaming_response( | ||
| url, | ||
| payload, | ||
| headers, | ||
| domain_request.session_id or "", | ||
| "openai", | ||
| ) | ||
| except AuthenticationError as e: | ||
| raise HTTPException(status_code=401, detail=str(e)) | ||
| return StreamingResponseEnvelope( | ||
| content=content_iterator, | ||
| media_type="text/event-stream", | ||
| headers={}, | ||
| ) | ||
| else: | ||
| # Return a domain ResponseEnvelope for non-streaming | ||
| return await self._handle_non_streaming_response( | ||
| url, payload, headers, domain_request.session_id or "" | ||
| # Prepare the payload using a helper so subclasses and tests can | ||
| # override or patch payload construction logic easily. | ||
| payload = await self._prepare_payload( | ||
| domain_request, processed_messages, effective_model | ||
| ) | ||
| headers_override = kwargs.pop("headers_override", None) | ||
| headers: dict[str, str] | None = None | ||
|
|
||
| if headers_override is not None: | ||
| # Avoid mutating the caller-provided mapping while preserving any | ||
| # Authorization header we compute from the configured API key. | ||
| headers = dict(headers_override) | ||
|
|
||
| try: | ||
| base_headers = self.get_headers() | ||
| except Exception: | ||
| base_headers = None | ||
|
|
||
| if base_headers: | ||
| merged_headers = dict(base_headers) | ||
| merged_headers.update(headers) | ||
| headers = merged_headers | ||
| else: | ||
| try: | ||
| # Always update the cached identity so that per-request | ||
| # identity headers do not leak between calls. Downstream | ||
| # callers rely on identity-specific headers being scoped to | ||
| # a single request. | ||
| self.identity = identity | ||
| headers = self.get_headers() | ||
| except Exception: | ||
| headers = None | ||
|
|
||
| api_base = kwargs.get("openai_url") or self.api_base_url | ||
| url = f"{api_base.rstrip('/')}/chat/completions" | ||
|
|
||
| if domain_request.stream: | ||
| # Return a domain-level streaming envelope (raw bytes iterator) | ||
| try: | ||
| content_iterator = await self._handle_streaming_response( | ||
| url, | ||
| payload, | ||
| headers, | ||
| domain_request.session_id or "", | ||
| "openai", | ||
| ) | ||
| except AuthenticationError as e: | ||
| raise HTTPException(status_code=401, detail=str(e)) | ||
| return StreamingResponseEnvelope( | ||
| content=content_iterator, | ||
| media_type="text/event-stream", | ||
| headers={}, | ||
| ) | ||
| else: | ||
| # Return a domain ResponseEnvelope for non-streaming | ||
| return await self._handle_non_streaming_response( | ||
| url, payload, headers, domain_request.session_id or "" | ||
| ) | ||
| finally: | ||
| if api_key is not None: | ||
| self.api_key = original_api_key |
There was a problem hiding this comment.
The new api_key parameter works by temporarily assigning self.api_key and restoring it in a finally block (src/connectors/openai.py lines 222‑321). This introduces a race condition when the same OpenAIConnector instance services concurrent requests with different per-call keys: the attribute is global to the instance, so one request can overwrite another’s key while both are in flight, and whichever finishes first will reset the field to the original value. In a multi-tenant deployment this can produce Authorization headers that use the wrong tenant’s key or fail authentication sporadically. Consider passing the override down to header construction without mutating shared state (e.g., via local variables) or making the key thread-safe.
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/openai.py(1 hunks)tests/unit/connectors/test_precision_payload_mapping.py(2 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:
src/connectors/openai.pytests/unit/connectors/test_precision_payload_mapping.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/openai.py
🧬 Code graph analysis (2)
src/connectors/openai.py (5)
src/core/domain/responses.py (2)
ResponseEnvelope(10-23)StreamingResponseEnvelope(27-53)src/core/domain/chat.py (2)
CanonicalChatRequest(321-324)ChatRequest(141-206)src/connectors/zai.py (1)
get_headers(106-117)src/connectors/openrouter.py (1)
get_headers(93-105)src/core/common/exceptions.py (1)
AuthenticationError(56-65)
tests/unit/connectors/test_precision_payload_mapping.py (2)
src/connectors/openai.py (2)
responses(552-670)chat_completions(222-321)src/core/domain/responses.py (1)
ResponseEnvelope(10-23)
🪛 Ruff (0.13.3)
src/connectors/openai.py
251-254: Avoid specifying long messages outside the exception class
(TRY003)
276-276: Do not catch blind exception: Exception
(BLE001)
291-291: Do not catch blind exception: Exception
(BLE001)
308-308: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
tests/unit/connectors/test_precision_payload_mapping.py
63-63: Unused function argument: self
(ARG001)
64-64: Unused function argument: url
(ARG001)
65-65: Unused function argument: payload
(ARG001)
67-67: Unused function argument: session_id
(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)
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/connectors/openai.py (1)
289-289: Remove redundant identity assignment.Line 289 sets
self.identity = identity, but this was already done on line 259. This assignment is redundant and suggests confused logic flow.Apply this diff:
else: try: - # Always update the cached identity so that per-request - # identity headers do not leak between calls. Downstream - # callers rely on identity-specific headers being scoped to - # a single request. - self.identity = identity headers = self.get_headers(effective_api_key) except Exception: headers = None
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/connectors/openai.py(6 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:
src/connectors/openai.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/openai.py
🔇 Additional comments (1)
src/connectors/openai.py (1)
98-115: Per-call API key handling looks good.The method correctly accepts an optional
api_keyparameter and uses it as an override without mutating instance state. This is the proper pattern and fixes the concurrency issue flagged in previous reviews.
| # Backends should ONLY convert FROM domain TO backend-specific format | ||
| # Type assertion: we know from architectural design that request_data is ChatRequest-like | ||
| from typing import cast | ||
| original_identity = self.identity |
There was a problem hiding this comment.
Critical: Identity mutation creates the same race condition that api_key had.
While the per-call api_key handling correctly avoids mutating instance state, the identity handling still uses the save/mutate/restore pattern that previous reviewers flagged as unsafe:
- Line 233:
original_identity = self.identity - Line 259:
self.identity = identity - Line 289:
self.identity = identity(redundant) - Line 320:
self.identity = original_identity(restore in finally)
If this connector instance services concurrent requests with different identities, the same race condition occurs: one request can overwrite another's identity while both are in flight, and get_headers() (called at lines 275, 290) will read the wrong identity, producing incorrect headers for one or both requests.
Consider passing identity as a parameter to get_headers() and downstream methods rather than mutating self.identity, matching the pattern you correctly used for api_key.
- def get_headers(self, api_key: str | None = None) -> dict[str, str]:
+ def get_headers(self, api_key: str | None = None, identity: IAppIdentityConfig | None = None) -> dict[str, str]:
"""Return request headers including API key and per-request identity."""
headers: dict[str, str] = {}
effective_api_key = api_key if api_key is not None else self.api_key
if effective_api_key:
headers["Authorization"] = f"Bearer {effective_api_key}"
- if self.identity:
+ effective_identity = identity if identity is not None else self.identity
+ if effective_identity:
try:
- identity_headers = self.identity.get_resolved_headers(None)
+ identity_headers = effective_identity.get_resolved_headers(None)
except Exception:
identity_headers = {}
if identity_headers:
headers.update(identity_headers)
return ensure_loop_guard_header(headers)Then update call sites to pass identity explicitly:
- base_headers = self.get_headers(effective_api_key)
+ base_headers = self.get_headers(effective_api_key, identity)Also applies to: 259-259, 289-289, 320-320
🤖 Prompt for AI Agents
In src/connectors/openai.py around lines 233-320, do not mutate self.identity
(originally saved at line 233 and modified at lines 259, 289, restored at 320)
because it creates a race condition for concurrent requests; instead add an
identity parameter to get_headers() and any downstream helper methods, remove
all assignments to self.identity (including the redundant line at 289 and the
restore in finally at 320), and update all call sites within this block to pass
the local identity variable explicitly so each request uses its own identity
without touching instance state.
| try: | ||
| base_headers = self.get_headers(effective_api_key) | ||
| except Exception: | ||
| base_headers = None |
There was a problem hiding this comment.
Catch specific exceptions and log errors.
These blocks catch broad Exception without logging, making debugging difficult when get_headers() fails. Per coding guidelines: "If a broad exception must be caught, log with exc_info=True and re-raise a specific custom exception."
Apply this diff to add logging:
if headers_override is not None:
headers = dict(headers_override)
try:
base_headers = self.get_headers(effective_api_key)
- except Exception:
+ except Exception as e:
+ logger.warning("Failed to get base headers: %s", e, exc_info=True)
base_headers = None
if base_headers:
merged_headers = dict(base_headers)
merged_headers.update(headers)
headers = merged_headers
else:
try:
- # Always update the cached identity so that per-request
- # identity headers do not leak between calls. Downstream
- # callers rely on identity-specific headers being scoped to
- # a single request.
- self.identity = identity
headers = self.get_headers(effective_api_key)
- except Exception:
+ except Exception as e:
+ logger.warning("Failed to get headers: %s", e, exc_info=True)
headers = NoneBased on coding guidelines.
Also applies to: 284-292
Summary
Testing
pytest -o addopts="" tests/unit/connectors/test_precision_payload_mapping.pypytest -o addopts=""(fails: missing pytest_asyncio, pytest_httpx, respx dependencies in the execution environment)https://chatgpt.com/codex/tasks/task_e_68ec25fa5a0c8333923ed0f3169275bf
Summary by CodeRabbit