Dev/issues 1

[luk384090-cloud]

Description

[Describe what this PR does and why]

Related Issue: Fixes #[issue_number] or Relates to #[issue_number]

Security Considerations: [Check if API keys or sensitive credentials are exposed in code/logs]

Type of Change

• Bug fix
• New feature
• Breaking change
• Documentation
• Refactoring

Component(s) Affected

• Model
• Application
• Common
• Documentation
• Tests
• CI/CD

Checklist

• Pre-commit hooks pass
• Tests pass locally
• Documentation updated (if needed)
• Ready for review

Testing

[How to test these changes]

Additional Notes

[Optional: any other context]

[gemini-code-assist]

Code Review

This pull request introduces a shared aiohttp session pool with a cached SSL context to enable connection reuse across async API calls, refactors request handlers to utilize this pool, adds support for merging incremental application response chunks, and implements chunking in the Qwen tokenizer to prevent stack overflows on long inputs. The review feedback highlights several critical improvement opportunities: addressing thread-safety issues in the shared session pool by introducing a lock, optimizing string concatenation in the tokenizer's chunking logic to avoid O(N^2) complexity, performing case-insensitive lookups for the User-Agent header, and raising an exception instead of silently returning when a server error occurs during end_session.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The global _aio_sessions dictionary is a weakref.WeakKeyDictionary, which is implemented in pure Python and is not thread-safe. In multi-threaded environments (such as web servers like Uvicorn/Gunicorn or multi-threaded workers), concurrent calls to get_shared_aio_session() from different threads will modify _aio_sessions concurrently, leading to race conditions, KeyError, or RuntimeError.

We should introduce a threading.Lock to synchronize access to _aio_sessions and _shared_ssl_context.

import threading

_shared_ssl_context: Optional[ssl.SSLContext] = None
_aio_sessions: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
_lock = threading.Lock()


def get_ssl_context() -> ssl.SSLContext:
    global _shared_ssl_context
    if _shared_ssl_context is None:
        with _lock:
            if _shared_ssl_context is None:
                _shared_ssl_context = ssl.create_default_context(
                    cafile=certifi.where(),
                )
    return _shared_ssl_context


async def get_shared_aio_session() -> aiohttp.ClientSession:
    """Return a shared aiohttp.ClientSession bound to the running event loop.

    The session is lazily created on first use and reused for all
    subsequent calls on the same event loop. Connection pooling (keep-alive)
    is handled by the underlying TCPConnector.
    """
    loop = asyncio.get_running_loop()

    with _lock:
        session = _aio_sessions.get(loop)
        if session is not None and not session.closed:
            return session

        connector = aiohttp.TCPConnector(ssl=get_ssl_context())
        session = aiohttp.ClientSession(connector=connector, trust_env=True)
        _aio_sessions[loop] = session
        return session


async def close_shared_aio_session() -> None:
    """Close the shared session for the current event loop."""
    loop = asyncio.get_running_loop()
    with _lock:
        session = _aio_sessions.pop(loop, None)
    if session is not None and not session.closed:
        await session.close()

[gemini-code-assist]

In _split_text, the chunks are merged using repeated string concatenation (chunks[-1] += part). Because Python strings are immutable, this operation has O(N^2) time complexity relative to the number of parts. For very large inputs (which is exactly when this chunking logic is triggered), this can cause severe performance degradation, high CPU usage, or timeouts.

We should accumulate the parts in a list and join them using "".join() when the chunk size limit is reached. This is a standard Python best practice and runs in O(N) time.

Suggested change

	chunks: List[str] = []
	for part in parts:
	if chunks and len(chunks[-1]) + len(part) <= chunk_size:
	chunks[-1] += part
	else:
	chunks.append(part)
	return chunks
	chunks: List[str] = []
	current_chunk: List[str] = []
	current_len = 0
	for part in parts:
	if current_len + len(part) <= chunk_size:
	current_chunk.append(part)
	current_len += len(part)
	else:
	if current_chunk:
	chunks.append("".join(current_chunk))
	current_chunk = [part]
	current_len = len(part)
	if current_chunk:
	chunks.append("".join(current_chunk))
	return chunks

[gemini-code-assist]

When appending the incremental_to_full flag to the user-agent header, the code assumes the key is exactly lowercase "user-agent". However, users or other parts of the SDK might pass "User-Agent" (capitalized). This case-sensitivity mismatch can result in duplicate User-Agent headers (both "User-Agent" and "user-agent"), which can cause unexpected behavior or failures in HTTP clients.

We should perform a case-insensitive lookup for the User-Agent header key and update it accordingly.

        if "headers" not in kwargs:
            kwargs["headers"] = {}
        flag = "1" if to_merge_incremental_output else "0"
        ua_key = "user-agent"
        existing_ua = ""
        for k, v in list(kwargs["headers"].items()):
            if k.lower() == "user-agent":
                ua_key = k
                existing_ua = v
                break
        new_ua = f"incremental_to_full/{flag}"
        kwargs["headers"][ua_key] = (
            f"{existing_ua} {new_ua}".strip() if existing_ua else new_ua
        )

[gemini-code-assist]

In end_session, if the server returns an error during the disconnect process, the code closes the connection and returns silently. This hides the error from the caller, who will assume the session ended successfully. It is better to raise an exception so the caller is notified of the failure.

Suggested change

	# if the server returned an error or timed out, close the connection
	if error is not None:
	self.close()
	return
	# if the server returned an error or timed out, close the connection
	if error is not None:
	self.close()
	raise RuntimeError(f"Session ended with error: {error}")