feat(sessions): start-of-turn ledger rows with completion, and turn-index ordering

Author: mmabrouk

api/oss/src/apis/fastapi/sessions/models.py

@@ -175,6 +175,13 @@ class SessionTurnAppendRequest(BaseModel):
     end_time: Optional[datetime] = None
 
 
+class SessionTurnCompleteRequest(BaseModel):
+    session_id: str
+    turn_index: int
+    agent_session_id: Optional[str] = None
+    end_time: datetime
+
+
 class SessionTurnQueryRequest(BaseModel):
     query: Optional[SessionTurnQuery] = None
     windowing: Optional[Windowing] = None

api/oss/src/apis/fastapi/sessions/router.py

@@ -61,8 +61,9 @@
 from oss.src.core.sessions.interactions.types import InteractionNotFound
 from oss.src.core.sessions.mounts.service import SessionMountsService
 from oss.src.core.sessions.mounts.dtos import SessionMountQuery
-from oss.src.core.sessions.turns.dtos import SessionTurnCreate
+from oss.src.core.sessions.turns.dtos import SessionTurnComplete, SessionTurnCreate
 from oss.src.core.sessions.turns.service import SessionTurnsService
+from oss.src.core.sessions.turns.types import SessionTurnNotFound
 from oss.src.core.sessions.dtos import SessionQuery
 from oss.src.core.sessions.service import SessionsService
 from oss.src.core.mounts.service import MountsService
@@ -107,6 +108,7 @@
     SessionMountsResponse,
     # turns
     SessionTurnAppendRequest,
+    SessionTurnCompleteRequest,
     SessionTurnQueryRequest,
     SessionTurnResponse,
     SessionTurnsResponse,
@@ -1048,6 +1050,15 @@ def __init__(self, *, turns_service: SessionTurnsService) -> None:
             response_model=SessionTurnResponse,
             response_model_exclude_none=True,
         )
+        self.router.add_api_route(
+            "/complete",
+            self.complete_turn,
+            methods=["POST"],
+            operation_id="complete_turn",
+            status_code=status.HTTP_200_OK,
+            response_model=SessionTurnResponse,
+            response_model_exclude_none=True,
+        )
         self.router.add_api_route(
             "/query",
             self.query_turns,
@@ -1102,6 +1113,39 @@ async def append_turn(
         )
         return SessionTurnResponse(count=1, turn=turn)
 
+    @intercept_exceptions()
+    async def complete_turn(
+        self,
+        request: Request,
+        body: SessionTurnCompleteRequest,
+    ) -> SessionTurnResponse:
+        project_id: UUID = request.state.project_id
+        user_id: UUID = request.state.user_id
+
+        if not await check_action_access(
+            user_uid=str(user_id),
+            project_id=str(project_id),
+            permission=Permission.RUN_SESSIONS,
+        ):
+            raise FORBIDDEN_EXCEPTION
+
+        try:
+            turn = await self.turns_service.complete_turn(
+                project_id=project_id,
+                turn=SessionTurnComplete(
+                    session_id=body.session_id,
+                    turn_index=body.turn_index,
+                    agent_session_id=body.agent_session_id,
+                    end_time=body.end_time,
+                ),
+            )
+        except SessionTurnNotFound as e:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail=e.message,
+            ) from e
+        return SessionTurnResponse(count=1, turn=turn)
+
     @intercept_exceptions()
     async def query_turns(
         self,

api/oss/src/core/sessions/turns/dtos.py

@@ -38,6 +38,13 @@ class SessionTurnCreate(BaseModel):
     end_time: Optional[datetime] = None
 
 
+class SessionTurnComplete(BaseModel):
+    session_id: str
+    turn_index: int
+    agent_session_id: Optional[str] = None
+    end_time: datetime
+
+
 class SessionTurnQuery(BaseModel):
     session_id: Optional[str] = None
     stream_id: Optional[UUID] = None

api/oss/src/core/sessions/turns/interfaces.py

@@ -6,6 +6,7 @@
 from oss.src.core.sessions.turns.dtos import (
     HarnessKind,
     SessionTurn,
+    SessionTurnComplete,
     SessionTurnCreate,
     SessionTurnQuery,
 )
@@ -22,6 +23,15 @@ async def append(
         turn: SessionTurnCreate,
     ) -> SessionTurn: ...
 
+    @abstractmethod
+    async def complete(
+        self,
+        *,
+        project_id: UUID,
+        #
+        turn: SessionTurnComplete,
+    ) -> Optional[SessionTurn]: ...
+
     @abstractmethod
     async def fetch_turn(
         self,

api/oss/src/core/sessions/turns/service.py

@@ -12,10 +12,12 @@
 from oss.src.core.sessions.turns.dtos import (
     HarnessKind,
     SessionTurn,
+    SessionTurnComplete,
     SessionTurnCreate,
     SessionTurnQuery,
 )
 from oss.src.core.sessions.turns.interfaces import SessionTurnsDAOInterface
+from oss.src.core.sessions.turns.types import SessionTurnNotFound
 
 
 class SessionTurnsService:
@@ -52,6 +54,21 @@ async def fetch_turn(
             turn_id=turn_id,
         )
 
+    async def complete_turn(
+        self,
+        *,
+        project_id: UUID,
+        #
+        turn: SessionTurnComplete,
+    ) -> SessionTurn:
+        completed = await self._dao.complete(
+            project_id=project_id,
+            turn=turn,
+        )
+        if completed is None:
+            raise SessionTurnNotFound(turn.session_id, turn.turn_index)
+        return completed
+
     async def query_turns(
         self,
         *,

api/oss/src/core/sessions/turns/types.py

@@ -6,7 +6,8 @@ class SessionTurnError(Exception):
 
 
 class SessionTurnNotFound(SessionTurnError):
-    def __init__(self, turn_id: str):
-        self.turn_id = turn_id
-        self.message = f"No turn found with id '{turn_id}'."
+    def __init__(self, session_id: str, turn_index: int):
+        self.session_id = session_id
+        self.turn_index = turn_index
+        self.message = f"No turn {turn_index} found for session '{session_id}'."
         super().__init__(self.message)

api/oss/src/dbs/postgres/sessions/turns/dao.py

@@ -1,12 +1,14 @@
 from typing import List, Optional
 from uuid import UUID
 
-from sqlalchemy import delete as sa_delete, select
+from sqlalchemy import and_, delete as sa_delete, or_, select
+from sqlalchemy import update as sa_update
 from sqlalchemy.exc import IntegrityError
 
 from oss.src.core.sessions.turns.dtos import (
     HarnessKind,
     SessionTurn,
+    SessionTurnComplete,
     SessionTurnCreate,
     SessionTurnQuery,
 )
@@ -23,7 +25,6 @@
     TransactionsEngine,
     get_transactions_engine,
 )
-from oss.src.dbs.postgres.shared.utils import apply_windowing
 
 
 class SessionTurnsDAO(SessionTurnsDAOInterface):
@@ -68,6 +69,48 @@ async def append(
                 raise
         return map_turn_dbe_to_dto(turn_dbe=dbe)
 
+    async def complete(
+        self,
+        *,
+        project_id: UUID,
+        #
+        turn: SessionTurnComplete,
+    ) -> Optional[SessionTurn]:
+        values = {"end_time": turn.end_time}
+        if turn.agent_session_id is not None:
+            values["agent_session_id"] = turn.agent_session_id
+
+        async with self.engine.session() as session:
+            stmt = (
+                sa_update(SessionTurnDBE)
+                .where(
+                    SessionTurnDBE.project_id == project_id,
+                    SessionTurnDBE.session_id == turn.session_id,
+                    SessionTurnDBE.turn_index == turn.turn_index,
+                    SessionTurnDBE.end_time.is_(None),
+                )
+                .values(**values)
+                .returning(SessionTurnDBE)
+            )
+            result = await session.execute(stmt)
+            dbe = result.scalar_one_or_none()
+            if dbe is not None:
+                await session.commit()
+                await session.refresh(dbe)
+                return map_turn_dbe_to_dto(turn_dbe=dbe)
+
+            stmt = select(SessionTurnDBE).where(
+                SessionTurnDBE.project_id == project_id,
+                SessionTurnDBE.session_id == turn.session_id,
+                SessionTurnDBE.turn_index == turn.turn_index,
+            )
+            result = await session.execute(stmt)
+            dbe = result.scalar_one_or_none()
+
+        if dbe is None:
+            return None
+        return map_turn_dbe_to_dto(turn_dbe=dbe)
+
     async def fetch_turn(
         self,
         *,
@@ -119,16 +162,67 @@ async def query_turns(
                             SessionTurnDBE.references.contains(turn_references),
                         )
 
+            descending = windowing is None or windowing.order != "ascending"
+            time_attribute = SessionTurnDBE.start_time
+
             if windowing:
-                stmt = apply_windowing(
-                    stmt=stmt,
-                    DBE=SessionTurnDBE,
-                    attribute="id",
-                    order="descending",
-                    windowing=windowing,
+                if windowing.newest is not None:
+                    if descending and windowing.next is None:
+                        stmt = stmt.where(time_attribute < windowing.newest)
+                    else:
+                        stmt = stmt.where(time_attribute <= windowing.newest)
+                if windowing.oldest is not None:
+                    if not descending and windowing.next is None:
+                        stmt = stmt.where(time_attribute > windowing.oldest)
+                    else:
+                        stmt = stmt.where(time_attribute >= windowing.oldest)
+
+                if windowing.next is not None:
+                    cursor_stmt = select(
+                        SessionTurnDBE.turn_index,
+                        SessionTurnDBE.id,
+                    ).where(
+                        SessionTurnDBE.project_id == project_id,
+                        SessionTurnDBE.id == windowing.next,
+                    )
+                    cursor_result = await session.execute(cursor_stmt)
+                    cursor = cursor_result.one_or_none()
+                    if cursor is not None:
+                        cursor_index, cursor_id = cursor
+                        if descending:
+                            stmt = stmt.where(
+                                or_(
+                                    SessionTurnDBE.turn_index < cursor_index,
+                                    and_(
+                                        SessionTurnDBE.turn_index == cursor_index,
+                                        SessionTurnDBE.id < cursor_id,
+                                    ),
+                                )
+                            )
+                        else:
+                            stmt = stmt.where(
+                                or_(
+                                    SessionTurnDBE.turn_index > cursor_index,
+                                    and_(
+                                        SessionTurnDBE.turn_index == cursor_index,
+                                        SessionTurnDBE.id > cursor_id,
+                                    ),
+                                )
+                            )
+
+            if descending:
+                stmt = stmt.order_by(
+                    SessionTurnDBE.turn_index.desc(),
+                    SessionTurnDBE.id.desc(),
                 )
             else:
-                stmt = stmt.order_by(SessionTurnDBE.created_at.desc())
+                stmt = stmt.order_by(
+                    SessionTurnDBE.turn_index.asc(),
+                    SessionTurnDBE.id.asc(),
+                )
+
+            if windowing and windowing.limit:
+                stmt = stmt.limit(windowing.limit)
 
             result = await session.execute(stmt)
             return [map_turn_dbe_to_dto(turn_dbe=dbe) for dbe in result.scalars().all()]