diff --git a/src/crud.py b/src/crud.py index bef92afb..93914efc 100644 --- a/src/crud.py +++ b/src/crud.py @@ -1958,10 +1958,10 @@ def _build_status_response( pending_work_units=counts.pending, ) - sessions: dict[str, schemas.DeriverStatus] = {} + sessions: dict[str, schemas.SessionDeriverStatus] = {} for session_id, data in counts.sessions.items(): total = data.completed + data.in_progress + data.pending - sessions[session_id] = schemas.DeriverStatus( + sessions[session_id] = schemas.SessionDeriverStatus( peer_id=peer_name, session_id=session_id, total_work_units=total, diff --git a/src/schemas.py b/src/schemas.py index 3e3837ad..b0ea7f4c 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -329,6 +329,22 @@ class MessageBulkData(BaseModel): workspace_name: str +class SessionDeriverStatus(BaseModel): + peer_id: str | None = Field( + default=None, + description="ID of the peer (optional when filtering by session only)", + ) + session_id: str | None = Field( + default=None, description="Session ID if filtered by session" + ) + total_work_units: int = Field(description="Total work units") + completed_work_units: int = Field(description="Completed work units") + in_progress_work_units: int = Field( + description="Work units currently being processed" + ) + pending_work_units: int = Field(description="Work units waiting to be processed") + + class DeriverStatus(BaseModel): peer_id: str | None = Field( default=None, @@ -343,6 +359,6 @@ class DeriverStatus(BaseModel): description="Work units currently being processed" ) pending_work_units: int = Field(description="Work units waiting to be processed") - sessions: dict[str, "DeriverStatus"] | None = Field( + sessions: dict[str, SessionDeriverStatus] | None = Field( default=None, description="Per-session status when not filtered by session" )