refactor(BA-5528): align DeploymentChatClient session lifecycle, add Args type

Address review comments on PR #11344:

- Drop _owns_session and the optional session= kwarg on
  DeploymentChatClient. Match BackendAIAuthClient: __init__ takes a
  pre-built session, factory method create() builds one, close() always
  closes. Removes the dual-ownership branch.
- Introduce DeploymentChatClientArgs (frozen dataclass) for connection
  knobs (skip_ssl_verification, connect_timeout, read_timeout).
  Callers use DeploymentChatClient.create(args) instead of passing
  multiple kwargs to the constructor.
- Rename chat_completion's 'request' parameter to 'body'.
- Tests: rename the cache-entry helper to _make_entry, the chat-body
  helper to _make_body. Drop TestExternalSession since the new
  contract is 'whatever you pass to __init__ gets closed'.

Author: jopemachine

src/ai/backend/client/cli/v2/deployment/chat/commands.py

@@ -93,6 +93,7 @@ def chat(
     from ai.backend.client.v2.deployment_chat import (
         DeploymentChatAuthError,
         DeploymentChatClient,
+        DeploymentChatClientArgs,
     )
 
     connection = load_v2_config()
@@ -150,9 +151,10 @@ async def _run() -> None:
             "messages": [{"role": "user", "content": content}],
         }
         api_key = chat_config_store.get_token(deployment_id)
-        async with DeploymentChatClient(
+        client_args = DeploymentChatClientArgs(
             skip_ssl_verification=connection.skip_ssl_verification,
-        ) as client:
+        )
+        async with await DeploymentChatClient.create(client_args) as client:
             try:
                 response = await client.chat_completion(
                     endpoint_entry.endpoint_url,
@@ -281,9 +283,11 @@ async def _discover_model(
     from ai.backend.client.v2.deployment_chat import (
         DeploymentChatAuthError,
         DeploymentChatClient,
+        DeploymentChatClientArgs,
     )
 
-    async with DeploymentChatClient(skip_ssl_verification=skip_ssl_verification) as client:
+    client_args = DeploymentChatClientArgs(skip_ssl_verification=skip_ssl_verification)
+    async with await DeploymentChatClient.create(client_args) as client:
         try:
             payload = await client.list_models(endpoint_url, api_key)
         except (DeploymentChatAuthError, BackendAPIError, BackendClientError):

src/ai/backend/client/v2/deployment_chat.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from types import TracebackType
 from typing import Any, Self
 
@@ -28,31 +29,32 @@ class DeploymentChatAuthError(BackendAPIError):
     """Raised when the inference endpoint rejects the configured API key."""
 
 
+@dataclass(frozen=True)
+class DeploymentChatClientArgs:
+    """Connection knobs for :meth:`DeploymentChatClient.create`."""
+
+    skip_ssl_verification: bool = False
+    connect_timeout: float | None = 10.0
+    read_timeout: float | None = 300.0
+
+
 class DeploymentChatClient:
     """Direct HTTP client for OpenAI-compatible inference endpoints."""
 
     _session: aiohttp.ClientSession
-    _owns_session: bool
 
-    def __init__(
-        self,
-        *,
-        session: aiohttp.ClientSession | None = None,
-        skip_ssl_verification: bool = False,
-        connect_timeout: float | None = 10.0,
-        read_timeout: float | None = 300.0,
-    ) -> None:
-        if session is not None:
-            self._session = session
-            self._owns_session = False
-        else:
-            connector = aiohttp.TCPConnector(ssl=not skip_ssl_verification)
-            timeout = aiohttp.ClientTimeout(
-                sock_connect=connect_timeout,
-                sock_read=read_timeout,
-            )
-            self._session = aiohttp.ClientSession(connector=connector, timeout=timeout)
-            self._owns_session = True
+    def __init__(self, session: aiohttp.ClientSession) -> None:
+        self._session = session
+
+    @classmethod
+    async def create(cls, args: DeploymentChatClientArgs) -> Self:
+        connector = aiohttp.TCPConnector(ssl=not args.skip_ssl_verification)
+        timeout = aiohttp.ClientTimeout(
+            sock_connect=args.connect_timeout,
+            sock_read=args.read_timeout,
+        )
+        session = aiohttp.ClientSession(connector=connector, timeout=timeout)
+        return cls(session)
 
     async def __aenter__(self) -> Self:
         return self
@@ -66,20 +68,20 @@ async def __aexit__(
         await self.close()
 
     async def close(self) -> None:
-        if self._owns_session and not self._session.closed:
+        if not self._session.closed:
             await self._session.close()
 
     async def chat_completion(
         self,
         endpoint_url: str,
         api_key: str | None,
-        request: dict[str, Any],
+        body: dict[str, Any],
         *,
         path: str = DEFAULT_CHAT_PATH,
     ) -> dict[str, Any]:
         """POST a chat completion request to the deployment endpoint."""
         target = self._build_url(endpoint_url, path)
-        return await self._post(target, api_key, request)
+        return await self._post(target, api_key, body)
 
     async def list_models(
         self,

tests/unit/client/cli/test_deployment_chat_types.py

@@ -10,7 +10,7 @@
 )
 
 
-def _entry(
+def _make_entry(
     *,
     endpoint: str = "https://infer.example.test/api",
     default_model: str | None = None,
@@ -24,32 +24,32 @@ def _entry(
 
 class TestEntryFormatSummary:
     def test_format_summary_returns_lines(self) -> None:
-        entry = _entry(default_model="meta/test-model")
+        entry = _make_entry(default_model="meta/test-model")
         lines = entry.format_summary()
         assert any("endpoint_url" in line for line in lines)
         assert any("meta/test-model" in line for line in lines)
         assert any("last_synced_at" in line for line in lines)
 
     def test_format_summary_dash_for_missing_default_model(self) -> None:
-        entry = _entry(default_model=None)
+        entry = _make_entry(default_model=None)
         lines = entry.format_summary()
         assert any("default_model : -" in line for line in lines)
 
 
 class TestCacheMutations:
-    def test_upsert_overwrites_existing_entry(self) -> None:
+    def test_upsert_overwrites_existing_make_entry(self) -> None:
         cache = DeploymentChatCache()
         dep_id = uuid4()
-        cache.upsert(dep_id, _entry(default_model="m1"))
-        cache.upsert(dep_id, _entry(default_model="m2"))
+        cache.upsert(dep_id, _make_entry(default_model="m1"))
+        cache.upsert(dep_id, _make_entry(default_model="m2"))
         stored = cache.get(dep_id)
         assert stored is not None
         assert stored.default_model == "m2"
 
     def test_remove_returns_true_when_present(self) -> None:
         cache = DeploymentChatCache()
         dep_id = uuid4()
-        cache.upsert(dep_id, _entry())
+        cache.upsert(dep_id, _make_entry())
         assert cache.remove(dep_id) is True
         assert cache.get(dep_id) is None
 

tests/unit/client/cli/test_deployment_chat_utils.py

@@ -27,7 +27,7 @@
 )
 
 
-def _entry(
+def _make_entry(
     *,
     endpoint: str = "https://infer.example.test/api",
     default_model: str | None = None,
@@ -44,11 +44,11 @@ def test_load_returns_empty_when_file_missing(self, tmp_path: Path) -> None:
         cache = load_chat_cache(tmp_path / "missing.json")
         assert cache.deployments == {}
 
-    def test_save_then_load_preserves_entry(self, tmp_path: Path) -> None:
+    def test_save_then_load_preserves_make_entry(self, tmp_path: Path) -> None:
         path = tmp_path / "deployment_chat.json"
         cache = DeploymentChatCache()
         dep_id = uuid4()
-        original = _entry(default_model="gpt-test")
+        original = _make_entry(default_model="gpt-test")
         cache.upsert(dep_id, original)
         save_chat_cache(cache, path)
 
@@ -96,7 +96,7 @@ class TestPermissions:
     def test_save_chat_cache_enforces_0600(self, tmp_path: Path) -> None:
         path = tmp_path / "cache.json"
         cache = DeploymentChatCache()
-        cache.upsert(uuid4(), _entry())
+        cache.upsert(uuid4(), _make_entry())
         save_chat_cache(cache, path)
         assert stat.S_IMODE(path.stat().st_mode) == 0o600
 

tests/unit/client/v2/test_deployment_chat_client.py

@@ -3,14 +3,14 @@
 from collections.abc import AsyncIterator, Awaitable, Callable
 from typing import Any
 
-import aiohttp
 import pytest
 from aiohttp import web
 
 from ai.backend.client.exceptions import BackendAPIError, BackendClientError
 from ai.backend.client.v2.deployment_chat import (
     DeploymentChatAuthError,
     DeploymentChatClient,
+    DeploymentChatClientArgs,
 )
 
 HandlerFn = Callable[[web.Request], Awaitable[web.StreamResponse]]
@@ -61,14 +61,14 @@ async def wrapped(request: web.Request) -> web.StreamResponse:
 
 @pytest.fixture
 async def chat_client() -> AsyncIterator[DeploymentChatClient]:
-    client = DeploymentChatClient()
+    client = await DeploymentChatClient.create(DeploymentChatClientArgs())
     try:
         yield client
     finally:
         await client.close()
 
 
-def _request_body() -> dict[str, Any]:
+def _make_body() -> dict[str, Any]:
     return {
         "model": "meta/test-model",
         "messages": [{"role": "user", "content": "hello"}],
@@ -96,7 +96,7 @@ async def handler(_request: web.Request) -> web.Response:
             resp = await chat_client.chat_completion(
                 server.base_url,
                 "sk-test-token",
-                _request_body(),
+                _make_body(),
             )
         finally:
             await server.stop()
@@ -105,7 +105,7 @@ async def handler(_request: web.Request) -> web.Response:
         assert server.recorded["path"] == "/v1/chat/completions"
         assert server.recorded["headers"]["Authorization"] == "Bearer sk-test-token"
         assert server.recorded["headers"]["Content-Type"] == "application/json"
-        assert server.recorded["json"] == _request_body()
+        assert server.recorded["json"] == _make_body()
         assert resp["choices"][0]["message"]["content"] == "hi"
 
     async def test_endpoint_url_already_ending_in_chat_completions(
@@ -117,7 +117,7 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             full_url = f"{server.base_url}/v1/chat/completions"
-            await chat_client.chat_completion(full_url, "sk-x", _request_body())
+            await chat_client.chat_completion(full_url, "sk-x", _make_body())
         finally:
             await server.stop()
         assert server.recorded["path"] == "/v1/chat/completions"
@@ -131,7 +131,7 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             full_url = f"{server.base_url}/v1/chat/completions/"
-            await chat_client.chat_completion(full_url, "sk-x", _request_body())
+            await chat_client.chat_completion(full_url, "sk-x", _make_body())
         finally:
             await server.stop()
         assert server.recorded["path"] == "/v1/chat/completions"
@@ -144,7 +144,7 @@ async def handler(_request: web.Request) -> web.Response:
 
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
-            await chat_client.chat_completion(server.base_url, None, _request_body())
+            await chat_client.chat_completion(server.base_url, None, _make_body())
         finally:
             await server.stop()
         assert "Authorization" not in server.recorded["headers"]
@@ -177,7 +177,7 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             with pytest.raises(DeploymentChatAuthError) as exc_info:
-                await chat_client.chat_completion(server.base_url, "bad", _request_body())
+                await chat_client.chat_completion(server.base_url, "bad", _make_body())
         finally:
             await server.stop()
         assert exc_info.value.status == 401
@@ -191,7 +191,7 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             with pytest.raises(DeploymentChatAuthError):
-                await chat_client.chat_completion(server.base_url, "bad", _request_body())
+                await chat_client.chat_completion(server.base_url, "bad", _make_body())
         finally:
             await server.stop()
 
@@ -206,7 +206,7 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             with pytest.raises(BackendAPIError) as exc_info:
-                await chat_client.chat_completion(server.base_url, "sk", _request_body())
+                await chat_client.chat_completion(server.base_url, "sk", _make_body())
         finally:
             await server.stop()
         assert not isinstance(exc_info.value, DeploymentChatAuthError)
@@ -223,14 +223,6 @@ async def handler(_request: web.Request) -> web.Response:
         server = await _start_server("POST", "/v1/chat/completions", handler)
         try:
             with pytest.raises(BackendClientError):
-                await chat_client.chat_completion(server.base_url, "sk", _request_body())
+                await chat_client.chat_completion(server.base_url, "sk", _make_body())
         finally:
             await server.stop()
-
-
-class TestExternalSession:
-    async def test_does_not_close_externally_owned_session(self) -> None:
-        async with aiohttp.ClientSession() as external:
-            client = DeploymentChatClient(session=external)
-            await client.close()
-            assert external.closed is False