From 3e94a41e8d0f5ff5d7b79c2dec83bedbfc614d20 Mon Sep 17 00:00:00 2001 From: Sourav Jha Date: Fri, 21 Aug 2026 15:47:55 +0530 Subject: [PATCH 1/3] fix: ensure reliable request log persistence by using a dedicated session factory and add integration test for streaming scenarios --- app/routes/chat.py | 23 ++- .../test_stream_log_reliability.py | 134 ++++++++++++++++++ 2 files changed, 144 insertions(+), 13 deletions(-) create mode 100644 tests/integration/test_stream_log_reliability.py diff --git a/app/routes/chat.py b/app/routes/chat.py index 750ea85..ad70a55 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -533,19 +533,16 @@ async def _finalize() -> None: ) from packages.db import session as session_mod - if session_mod._session_factory is not None: - try: - async with session_mod._session_factory() as s: - s.add(log) - await s.commit() - except Exception as commit_err: - logger.warning("request_log_commit_failed", error=str(commit_err)) - else: - db.add(log) - try: - await db.commit() - except Exception as commit_err: - logger.warning("request_log_commit_failed", error=str(commit_err)) + if session_mod._session_factory is None: + from app.config import get_settings + session_mod.init_session_factory(get_settings().database_url) + + try: + async with session_mod._session_factory() as s: + s.add(log) + await s.commit() + except Exception as commit_err: + logger.warning("request_log_commit_failed", error=str(commit_err)) try: async for chunk in _aiter(stream_obj): diff --git a/tests/integration/test_stream_log_reliability.py b/tests/integration/test_stream_log_reliability.py new file mode 100644 index 0000000..9fd6e88 --- /dev/null +++ b/tests/integration/test_stream_log_reliability.py @@ -0,0 +1,134 @@ +"""Streaming chat completions log persistence and teardown resilience tests. + +Validates that streaming completion request logs are reliably persisted to the database +even after the request-scoped dependency session (`get_db()`) has exited and closed. +""" + +from __future__ import annotations + +import json +import time +from unittest.mock import AsyncMock + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app import router_cache +from packages.auth.hashing import hash_api_key +from packages.db import session as session_mod +from packages.db.models.api_key import ApiKey +from packages.db.models.base import Base +from packages.db.models.request_log import RequestLog + + +@pytest.mark.asyncio +async def test_streaming_log_persistence_when_dependency_session_is_closed(monkeypatch): + """Verify that streaming request logs are written to the database after get_db() dependency session closes. + + Root Cause (Issue #2): get_db() yields `db` and closes it when chat_completions returns + StreamingResponse. _finalize() runs post-return during SSE stream output. Using a dedicated + session from session_factory ensures the RequestLog is successfully saved despite `db` being closed. + """ + engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + factory = async_sessionmaker(engine, expire_on_commit=False) + + # Seed API key into isolated test database + raw_key = "sk-orca-test123456789012345678901234" + key_hash = hash_api_key(raw_key) + async with factory() as s: + s.add( + ApiKey( + id="key_streaming_test", + workspace_id="default", + name="test-key", + key_hash=key_hash, + key_prefix="sk-orca-....1234", + is_active=True, + ) + ) + await s.commit() + + # Set _session_factory for authentication middleware and _finalize() + session_mod._session_factory = factory + + chunks = [ + { + "id": "chatcmpl-stream-test", + "object": "chat.completion.chunk", + "model": "gpt-4o-mini", + "created": int(time.time()), + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Test stream chunk"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, + "_orca_meta": {"provider": "openai", "latency_ms": 10}, + } + ] + + fake_client = AsyncMock() + + async def _stream_iter(chunk_list): + for c in chunk_list: + yield c + + fake_client.acompletion = AsyncMock( + return_value=_stream_iter(chunks) + ) + + async def _fake_get_router(_session): + return fake_client + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + from app.deps import get_db + + async def _override_get_db_that_closes_on_return(): + async with factory() as session: + yield session + # Dependency scope ends here when chat_completions returns StreamingResponse, + # explicitly closing session before _finalize() runs! + await session.close() + + from app.main import create_app + + app = create_app() + app.dependency_overrides[get_db] = _override_get_db_that_closes_on_return + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://test", + headers={"Authorization": f"Bearer {raw_key}"}, + ) as c: + response = await c.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + }, + ) + assert response.status_code == 200 + body_text = response.text + assert "data: [DONE]" in body_text + + # Verify that the streaming request log WAS successfully written to DB via session_factory + async with factory() as check_session: + rows = (await check_session.execute(select(RequestLog))).scalars().all() + + assert len(rows) == 1, ( + f"Expected 1 RequestLog row for streaming completion, but found {len(rows)}. " + "Streaming log failed to persist when request-scoped db session was closed." + ) + log = rows[0] + assert log.is_streaming is True + assert log.input_tokens == 5 + assert log.output_tokens == 5 From be61e4c6b6d731dab0ce94c1fbe96078770e39fd Mon Sep 17 00:00:00 2001 From: Sourav Jha Date: Fri, 21 Aug 2026 15:51:55 +0530 Subject: [PATCH 2/3] chore: remove unused json import from stream log reliability integration test --- tests/integration/test_stream_log_reliability.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_stream_log_reliability.py b/tests/integration/test_stream_log_reliability.py index 9fd6e88..383d8a4 100644 --- a/tests/integration/test_stream_log_reliability.py +++ b/tests/integration/test_stream_log_reliability.py @@ -6,7 +6,6 @@ from __future__ import annotations -import json import time from unittest.mock import AsyncMock From 08c72c1bbfcd989e70d74b18f25c1900d2236a41 Mon Sep 17 00:00:00 2001 From: Sourav Jha Date: Fri, 21 Aug 2026 16:08:33 +0530 Subject: [PATCH 3/3] fix: fallback to active db session for request logging if session_factory is unavailable --- app/routes/chat.py | 14 +- .../test_stream_log_reliability.py | 158 +++++++++--------- 2 files changed, 88 insertions(+), 84 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index ad70a55..66f90b8 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -533,14 +533,14 @@ async def _finalize() -> None: ) from packages.db import session as session_mod - if session_mod._session_factory is None: - from app.config import get_settings - session_mod.init_session_factory(get_settings().database_url) - try: - async with session_mod._session_factory() as s: - s.add(log) - await s.commit() + if session_mod._session_factory is not None: + async with session_mod._session_factory() as s: + s.add(log) + await s.commit() + else: + db.add(log) + await db.commit() except Exception as commit_err: logger.warning("request_log_commit_failed", error=str(commit_err)) diff --git a/tests/integration/test_stream_log_reliability.py b/tests/integration/test_stream_log_reliability.py index 383d8a4..2e5914c 100644 --- a/tests/integration/test_stream_log_reliability.py +++ b/tests/integration/test_stream_log_reliability.py @@ -52,82 +52,86 @@ async def test_streaming_log_persistence_when_dependency_session_is_closed(monke ) await s.commit() - # Set _session_factory for authentication middleware and _finalize() - session_mod._session_factory = factory - - chunks = [ - { - "id": "chatcmpl-stream-test", - "object": "chat.completion.chunk", - "model": "gpt-4o-mini", - "created": int(time.time()), - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "Test stream chunk"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, - "_orca_meta": {"provider": "openai", "latency_ms": 10}, - } - ] - - fake_client = AsyncMock() - - async def _stream_iter(chunk_list): - for c in chunk_list: - yield c - - fake_client.acompletion = AsyncMock( - return_value=_stream_iter(chunks) - ) - - async def _fake_get_router(_session): - return fake_client - - monkeypatch.setattr(router_cache, "get_router", _fake_get_router) - - from app.deps import get_db - - async def _override_get_db_that_closes_on_return(): - async with factory() as session: - yield session - # Dependency scope ends here when chat_completions returns StreamingResponse, - # explicitly closing session before _finalize() runs! - await session.close() - - from app.main import create_app - - app = create_app() - app.dependency_overrides[get_db] = _override_get_db_that_closes_on_return - - async with AsyncClient( - transport=ASGITransport(app=app), - base_url="http://test", - headers={"Authorization": f"Bearer {raw_key}"}, - ) as c: - response = await c.post( - "/v1/chat/completions", - json={ + try: + # Set _session_factory for authentication middleware and _finalize() + session_mod._session_factory = factory + + chunks = [ + { + "id": "chatcmpl-stream-test", + "object": "chat.completion.chunk", "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hi"}], - "stream": True, - }, + "created": int(time.time()), + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Test stream chunk"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, + "_orca_meta": {"provider": "openai", "latency_ms": 10}, + } + ] + + fake_client = AsyncMock() + + async def _stream_iter(chunk_list): + for c in chunk_list: + yield c + + fake_client.acompletion = AsyncMock( + return_value=_stream_iter(chunks) ) - assert response.status_code == 200 - body_text = response.text - assert "data: [DONE]" in body_text - - # Verify that the streaming request log WAS successfully written to DB via session_factory - async with factory() as check_session: - rows = (await check_session.execute(select(RequestLog))).scalars().all() - - assert len(rows) == 1, ( - f"Expected 1 RequestLog row for streaming completion, but found {len(rows)}. " - "Streaming log failed to persist when request-scoped db session was closed." - ) - log = rows[0] - assert log.is_streaming is True - assert log.input_tokens == 5 - assert log.output_tokens == 5 + + async def _fake_get_router(_session): + return fake_client + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + from app.deps import get_db + + async def _override_get_db_that_closes_on_return(): + async with factory() as session: + yield session + # Dependency scope ends here when chat_completions returns StreamingResponse, + # explicitly closing session before _finalize() runs! + await session.close() + + from app.main import create_app + + app = create_app() + app.dependency_overrides[get_db] = _override_get_db_that_closes_on_return + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://test", + headers={"Authorization": f"Bearer {raw_key}"}, + ) as c: + response = await c.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + }, + ) + assert response.status_code == 200 + body_text = response.text + assert "data: [DONE]" in body_text + + # Verify that the streaming request log WAS successfully written to DB via session_factory + async with factory() as check_session: + rows = (await check_session.execute(select(RequestLog))).scalars().all() + + assert len(rows) == 1, ( + f"Expected 1 RequestLog row for streaming completion, but found {len(rows)}. " + "Streaming log failed to persist when request-scoped db session was closed." + ) + log = rows[0] + assert log.is_streaming is True + assert log.input_tokens == 5 + assert log.output_tokens == 5 + finally: + session_mod._session_factory = None + await engine.dispose()