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

features/steps/llm_trace_steps.py

@@ -682,7 +682,11 @@ def step_sqla_repo(context: Context) -> None:
     engine, factory = _create_in_memory_engine_and_session()
     context.sqla_engine = engine
     context.sqla_factory = factory
-    context.sqla_repo = LLMTraceRepository(session_factory=factory)
+    # Use a single session for all repository operations so that
+    # flush() (without commit) keeps data visible across calls.
+    shared_session = factory()
+    context.sqla_shared_session = shared_session
+    context.sqla_repo = LLMTraceRepository(session_factory=lambda: shared_session)
 
 
 @given("a valid LLM trace with tool calls")
@@ -940,9 +944,113 @@ def step_langsmith_enabled_true(context: Context) -> None:
 
 @then("the raw tool_calls_json in the database should be null")
 def step_raw_tool_calls_null(context: Context) -> None:
-    session = context.sqla_factory()
+    session = getattr(context, "sqla_shared_session", None) or context.sqla_factory()
     row = (
         session.query(LLMTraceModel).filter_by(trace_id=context.trace.trace_id).first()
     )
     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)
+
+    # Assign spy functions to session methods
+    # Using object.__setattr__ to bypass type checking for method replacement
+    object.__setattr__(real_session, "flush", spy_flush)
+    object.__setattr__(real_session, "commit", spy_commit)
+
+    repo = LLMTraceRepository(session_factory=lambda: real_session)
+    repo.save(context.trace)
+    # Commit so the data is visible for subsequent queries
+    object.__setattr__(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()

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