diff --git a/features/llm_trace.feature b/features/llm_trace.feature index feb827178..74ecd1665 100644 --- a/features/llm_trace.feature +++ b/features/llm_trace.feature @@ -274,3 +274,18 @@ Feature: LLM trace observability And a valid LLM trace When I save the trace via the repository Then the raw tool_calls_json in the database should be null + + # --- session contract (flush not commit) -------------------------------- + + Scenario: Repository save() calls flush not commit + Given a SQLAlchemy in-memory repository + And a valid LLM trace + When I save the trace via the repository and spy on the session + Then the session flush should have been called + And the session commit should not have been called + + Scenario: LLM trace rolled back when UnitOfWork transaction rolls back + Given a UnitOfWork with an in-memory database + And a valid LLM trace + When I save the trace inside a UnitOfWork transaction that rolls back + Then the LLM trace should not be persisted in the database diff --git a/features/steps/llm_trace_steps.py b/features/steps/llm_trace_steps.py index dbd635e52..3794ddf2e 100644 --- a/features/steps/llm_trace_steps.py +++ b/features/steps/llm_trace_steps.py @@ -946,3 +946,106 @@ def step_raw_tool_calls_null(context: Context) -> None: ) assert row is not None assert row.tool_calls_json is None + + +# --------------------------------------------------------------------------- +# Session contract: flush not commit +# --------------------------------------------------------------------------- + + +@when("I save the trace via the repository and spy on the session") +def step_save_with_spy(context: Context) -> None: + engine, factory = _create_in_memory_engine_and_session() + context.sqla_engine = engine + context.sqla_factory = factory + + # Create a real session but wrap it to spy on flush/commit calls + real_session = factory() + context.spy_session = real_session + context.flush_called = False + context.commit_called = False + + original_flush = real_session.flush + original_commit = real_session.commit + + def spy_flush(*args: Any, **kwargs: Any) -> None: + context.flush_called = True + return original_flush(*args, **kwargs) + + def spy_commit(*args: Any, **kwargs: Any) -> None: + context.commit_called = True + return original_commit(*args, **kwargs) + + real_session.flush = spy_flush # type: ignore[method-assign] + real_session.commit = spy_commit # type: ignore[method-assign] + + repo = LLMTraceRepository(session_factory=lambda: real_session) + repo.save(context.trace) + # Commit so the data is visible for subsequent queries + real_session.commit = original_commit + real_session.commit() + + +@then("the session flush should have been called") +def step_flush_called(context: Context) -> None: + assert context.flush_called is True, "Expected session.flush() to have been called" + + +@then("the session commit should not have been called") +def step_commit_not_called(context: Context) -> None: + assert context.commit_called is False, ( + "Expected session.commit() NOT to have been called by save(), " + "but it was called" + ) + + +# --------------------------------------------------------------------------- +# UnitOfWork rollback propagation +# --------------------------------------------------------------------------- + + +@given("a UnitOfWork with an in-memory database") +def step_uow_in_memory(context: Context) -> None: + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + ) + Base.metadata.create_all(engine) + context.uow_engine = engine + context.uow_session_factory = sessionmaker( + bind=engine, + expire_on_commit=False, + autoflush=False, + autocommit=False, + ) + + +@when("I save the trace inside a UnitOfWork transaction that rolls back") +def step_save_in_uow_rollback(context: Context) -> None: + session = context.uow_session_factory() + repo = LLMTraceRepository(session_factory=lambda: session) + try: + repo.save(context.trace) + # Simulate a subsequent failure that triggers rollback + raise RuntimeError("Simulated failure after save") + except RuntimeError: + session.rollback() + finally: + session.close() + + +@then("the LLM trace should not be persisted in the database") +def step_trace_not_persisted(context: Context) -> None: + session = context.uow_session_factory() + try: + row = ( + session.query(LLMTraceModel) + .filter_by(trace_id=context.trace.trace_id) + .first() + ) + assert row is None, ( + f"Expected LLM trace {context.trace.trace_id} to be rolled back, " + f"but it was found in the database" + ) + finally: + session.close() diff --git a/src/cleveragents/infrastructure/database/llm_trace_repository.py b/src/cleveragents/infrastructure/database/llm_trace_repository.py index bf6c28e10..70fcfa1ae 100644 --- a/src/cleveragents/infrastructure/database/llm_trace_repository.py +++ b/src/cleveragents/infrastructure/database/llm_trace_repository.py @@ -78,7 +78,7 @@ class LLMTraceRepository: timestamp=trace.timestamp.isoformat(), ) session.add(model) - session.commit() + session.flush() except (SQLAlchemyDatabaseError, OperationalError) as exc: session.rollback() raise DatabaseError(f"Failed to save LLM trace: {exc}") from exc