Jay Brown
386 lines
14 KiB
Go
386 lines
14 KiB
Go
package bot
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
|
||
"queryorchestration/internal/database/repository"
|
||
|
||
"github.com/google/uuid"
|
||
"github.com/jackc/pgx/v5"
|
||
)
|
||
|
||
// CreateSession inserts a new owner-scoped session and returns the
|
||
// freshly-created Session domain object. Plan §6: title is derived from
|
||
// the first turn's prompt (M4 path), so this method always creates with
|
||
// the default empty title. State defaults to '{}'. Scope is empty.
|
||
func (s *Service) CreateSession(ctx context.Context, clientID, actor string) (*Session, error) {
|
||
row, err := s.dbCfg.GetDBQueries().InsertBotSession(ctx, &repository.InsertBotSessionParams{
|
||
ClientID: clientID,
|
||
CreatedBy: actor,
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("bot: insert session: %w", err)
|
||
}
|
||
return rowToSession(row.ID, row.ClientID, row.CreatedBy, row.Title, row.State,
|
||
row.CreatedAt.Time, row.UpdatedAt.Time, 0, SessionScope{}, nil, false), nil
|
||
}
|
||
|
||
// GetSessionForOwner returns the session if (id, client_id, actor) match
|
||
// and is_deleted=false. Wraps pgx.ErrNoRows in ErrSessionNotFound so the
|
||
// handler can map to 404 without driver knowledge.
|
||
func (s *Service) GetSessionForOwner(ctx context.Context, sessionID uuid.UUID, clientID, actor string) (*Session, error) {
|
||
row, err := s.dbCfg.GetDBQueries().GetBotSessionForOwner(ctx, &repository.GetBotSessionForOwnerParams{
|
||
SessionID: sessionID,
|
||
ClientID: clientID,
|
||
CreatedBy: actor,
|
||
})
|
||
if err != nil {
|
||
if errors.Is(err, pgx.ErrNoRows) {
|
||
return nil, ErrSessionNotFound
|
||
}
|
||
return nil, fmt.Errorf("bot: get session for owner: %w", err)
|
||
}
|
||
scope, err := s.loadScope(ctx, row.ID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return rowToSession(row.ID, row.ClientID, row.CreatedBy, row.Title, row.State,
|
||
row.CreatedAt.Time, row.UpdatedAt.Time, row.LastTerminalOrdinal, scope, nil, row.IsDeleted), nil
|
||
}
|
||
|
||
// GetSessionForSuperAdmin reads a session by (id, client_id) without the
|
||
// created_by predicate.
|
||
func (s *Service) GetSessionForSuperAdmin(ctx context.Context, sessionID uuid.UUID, clientID string) (*Session, error) {
|
||
row, err := s.dbCfg.GetDBQueries().GetBotSessionForSuperAdmin(ctx, &repository.GetBotSessionForSuperAdminParams{
|
||
SessionID: sessionID,
|
||
ClientID: clientID,
|
||
})
|
||
if err != nil {
|
||
if errors.Is(err, pgx.ErrNoRows) {
|
||
return nil, ErrSessionNotFound
|
||
}
|
||
return nil, fmt.Errorf("bot: get session for super-admin: %w", err)
|
||
}
|
||
scope, err := s.loadScope(ctx, row.ID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return rowToSession(row.ID, row.ClientID, row.CreatedBy, row.Title, row.State,
|
||
row.CreatedAt.Time, row.UpdatedAt.Time, row.LastTerminalOrdinal, scope, nil, row.IsDeleted), nil
|
||
}
|
||
|
||
// ListSessionsForOwner returns up to input.Limit sessions owned by actor
|
||
// in client clientID, with each session decorated with up to
|
||
// MinTurnLimit(input.Limit, cfg.RecentTurns) recent turns.
|
||
func (s *Service) ListSessionsForOwner(ctx context.Context, clientID, actor string, input ListSessionsInput) ([]*Session, error) {
|
||
if err := validateListInput(input); err != nil {
|
||
return nil, err
|
||
}
|
||
rows, err := s.dbCfg.GetDBQueries().ListBotSessionsForOwner(ctx, &repository.ListBotSessionsForOwnerParams{
|
||
ClientID: clientID,
|
||
CreatedBy: actor,
|
||
LimitVal: int32(input.Limit), //nolint:gosec // bounded [1, 200] by validateListInput
|
||
OffsetVal: int32(input.Offset), //nolint:gosec // bounded >= 0 by validateListInput
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("bot: list sessions for owner: %w", err)
|
||
}
|
||
return s.decorateSessionList(ctx, rowsOwner(rows), input.Limit)
|
||
}
|
||
|
||
// ListSessionsForSuperAdmin is the cross-user variant of ListSessionsForOwner.
|
||
func (s *Service) ListSessionsForSuperAdmin(ctx context.Context, clientID string, input ListSessionsInput) ([]*Session, error) {
|
||
if err := validateListInput(input); err != nil {
|
||
return nil, err
|
||
}
|
||
rows, err := s.dbCfg.GetDBQueries().ListBotSessionsForSuperAdmin(ctx, &repository.ListBotSessionsForSuperAdminParams{
|
||
ClientID: clientID,
|
||
LimitVal: int32(input.Limit), //nolint:gosec // bounded [1, 200] by validateListInput
|
||
OffsetVal: int32(input.Offset), //nolint:gosec // bounded >= 0 by validateListInput
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("bot: list sessions for super-admin: %w", err)
|
||
}
|
||
return s.decorateSessionList(ctx, rowsSuperAdmin(rows), input.Limit)
|
||
}
|
||
|
||
// DeleteSessionAsSuperAdmin soft-deletes the session and marks any
|
||
// in-flight turns as session_deleted. Returns ErrSessionNotFound when
|
||
// the session does not exist (or is already deleted) under the given
|
||
// clientId.
|
||
func (s *Service) DeleteSessionAsSuperAdmin(ctx context.Context, sessionID uuid.UUID, clientID string) error {
|
||
return s.dbCfg.ExecuteDBTransaction(ctx, func(txCtx context.Context, q *repository.Queries) error {
|
||
affected, err := q.SoftDeleteBotSessionForSuperAdmin(txCtx, &repository.SoftDeleteBotSessionForSuperAdminParams{
|
||
SessionID: sessionID,
|
||
ClientID: clientID,
|
||
})
|
||
if err != nil {
|
||
return fmt.Errorf("bot: soft delete session: %w", err)
|
||
}
|
||
if affected == 0 {
|
||
return ErrSessionNotFound
|
||
}
|
||
if err := q.MarkInflightTurnsSessionDeletedForSession(txCtx, sessionID); err != nil {
|
||
return fmt.Errorf("bot: mark in-flight turns session_deleted: %w", err)
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// PatchScope replaces the persisted scope of an owner-scoped session
|
||
// with (existing - remove) ∪ add for both documents and folders. All
|
||
// validation and the replace semantics live inside one transaction so
|
||
// concurrent PATCH callers either see the union of their writes or
|
||
// none of them.
|
||
func (s *Service) PatchScope(ctx context.Context, sessionID uuid.UUID, clientID, actor string, input PatchScopeInput) (*Session, error) {
|
||
addDocs := dedupUUIDs(input.AddDocuments)
|
||
removeDocs := dedupUUIDs(input.RemoveDocuments)
|
||
addFolders := dedupUUIDs(input.AddFolders)
|
||
removeFolders := dedupUUIDs(input.RemoveFolders)
|
||
|
||
if hasIntersection(addDocs, removeDocs) || hasIntersection(addFolders, removeFolders) {
|
||
return nil, ErrAddRemoveConflict
|
||
}
|
||
|
||
// Verify session ownership BEFORE doing any cross-client checks so
|
||
// other-user / cross-client lookups always present as 404 rather
|
||
// than leaking the existence of the session via 400 messages.
|
||
_, err := s.GetSessionForOwner(ctx, sessionID, clientID, actor)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if err := s.checkDocumentsBelongToClient(ctx, addDocs, clientID); err != nil {
|
||
return nil, err
|
||
}
|
||
if err := s.checkFoldersBelongToClient(ctx, addFolders, clientID); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if err := s.applyScopeReplacement(ctx, sessionID, addDocs, removeDocs, addFolders, removeFolders); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// Re-read the session post-write so the response carries the
|
||
// authoritative persisted scope and a fresh updatedAt.
|
||
return s.GetSessionForOwner(ctx, sessionID, clientID, actor)
|
||
}
|
||
|
||
// applyScopeReplacement runs the replace-set transaction:
|
||
// existing scope is loaded under the session lock, the new union is
|
||
// computed, both join tables are wiped, then the union is inserted.
|
||
// updated_at is bumped at the end so the owner-list ordering reflects
|
||
// the activity.
|
||
func (s *Service) applyScopeReplacement(ctx context.Context, sessionID uuid.UUID,
|
||
addDocs, removeDocs, addFolders, removeFolders []uuid.UUID) error {
|
||
return s.dbCfg.ExecuteDBTransaction(ctx, func(txCtx context.Context, q *repository.Queries) error {
|
||
existingDocs, err := q.ListBotSessionDocumentsForSession(txCtx, sessionID)
|
||
if err != nil {
|
||
return fmt.Errorf("bot: read existing scope documents: %w", err)
|
||
}
|
||
existingFolders, err := q.ListBotSessionFoldersForSession(txCtx, sessionID)
|
||
if err != nil {
|
||
return fmt.Errorf("bot: read existing scope folders: %w", err)
|
||
}
|
||
newDocs := computeUnion(existingDocs, addDocs, removeDocs)
|
||
newFolders := computeUnion(existingFolders, addFolders, removeFolders)
|
||
|
||
if err := q.DeleteBotSessionDocumentsForSession(txCtx, sessionID); err != nil {
|
||
return fmt.Errorf("bot: clear scope documents: %w", err)
|
||
}
|
||
if err := q.DeleteBotSessionFoldersForSession(txCtx, sessionID); err != nil {
|
||
return fmt.Errorf("bot: clear scope folders: %w", err)
|
||
}
|
||
for _, d := range newDocs {
|
||
if err := q.InsertBotSessionDocument(txCtx, &repository.InsertBotSessionDocumentParams{
|
||
SessionID: sessionID,
|
||
DocumentID: d,
|
||
}); err != nil {
|
||
return fmt.Errorf("bot: insert scope document: %w", err)
|
||
}
|
||
}
|
||
for _, f := range newFolders {
|
||
if err := q.InsertBotSessionFolder(txCtx, &repository.InsertBotSessionFolderParams{
|
||
SessionID: sessionID,
|
||
FolderID: f,
|
||
}); err != nil {
|
||
return fmt.Errorf("bot: insert scope folder: %w", err)
|
||
}
|
||
}
|
||
if err := q.TouchBotSessionUpdatedAt(txCtx, sessionID); err != nil {
|
||
return fmt.Errorf("bot: bump updatedAt: %w", err)
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// checkDocumentsBelongToClient validates that every id in candidates is
|
||
// a documents row owned by clientID. Missing rows or rows owned by
|
||
// another client return ErrDocumentCrossClient.
|
||
func (s *Service) checkDocumentsBelongToClient(ctx context.Context, candidates []uuid.UUID, clientID string) error {
|
||
if len(candidates) == 0 {
|
||
return nil
|
||
}
|
||
rows, err := s.dbCfg.GetDBQueries().GetDocumentClientIDsForBotScope(ctx, candidates)
|
||
if err != nil {
|
||
return fmt.Errorf("bot: bulk doc client lookup: %w", err)
|
||
}
|
||
if len(rows) != len(candidates) {
|
||
return ErrDocumentCrossClient
|
||
}
|
||
for _, r := range rows {
|
||
if r.ClientID != clientID {
|
||
return ErrDocumentCrossClient
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// checkFoldersBelongToClient validates that every id in candidates is a
|
||
// folders row owned by clientID. Missing rows or cross-client rows
|
||
// return ErrFolderCrossClient.
|
||
func (s *Service) checkFoldersBelongToClient(ctx context.Context, candidates []uuid.UUID, clientID string) error {
|
||
if len(candidates) == 0 {
|
||
return nil
|
||
}
|
||
rows, err := s.dbCfg.GetDBQueries().GetFolderClientIDsForBotScope(ctx, candidates)
|
||
if err != nil {
|
||
return fmt.Errorf("bot: bulk folder client lookup: %w", err)
|
||
}
|
||
if len(rows) != len(candidates) {
|
||
return ErrFolderCrossClient
|
||
}
|
||
for _, r := range rows {
|
||
if r.ClientID != clientID {
|
||
return ErrFolderCrossClient
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// loadScope fetches the persisted document/folder ids for the session.
|
||
// Returned as plain slices in deterministic order (same as the SQL).
|
||
func (s *Service) loadScope(ctx context.Context, sessionID uuid.UUID) (SessionScope, error) {
|
||
docs, err := s.dbCfg.GetDBQueries().ListBotSessionDocumentsForSession(ctx, sessionID)
|
||
if err != nil {
|
||
return SessionScope{}, fmt.Errorf("bot: load scope documents: %w", err)
|
||
}
|
||
folders, err := s.dbCfg.GetDBQueries().ListBotSessionFoldersForSession(ctx, sessionID)
|
||
if err != nil {
|
||
return SessionScope{}, fmt.Errorf("bot: load scope folders: %w", err)
|
||
}
|
||
if docs == nil {
|
||
docs = []uuid.UUID{}
|
||
}
|
||
if folders == nil {
|
||
folders = []uuid.UUID{}
|
||
}
|
||
return SessionScope{DocumentIDs: docs, FolderIDs: folders}, nil
|
||
}
|
||
|
||
// decorateSessionList enriches a list of session rows with their scope
|
||
// and recent-turn previews. Runs one ListLastTurnsForSessions call for
|
||
// the whole batch so the per-session preview avoids N+1.
|
||
func (s *Service) decorateSessionList(ctx context.Context, rows []sessionListRow, requestLimit int) ([]*Session, error) {
|
||
if len(rows) == 0 {
|
||
return []*Session{}, nil
|
||
}
|
||
|
||
turnLimit := MinTurnLimit(requestLimit, s.cfg.RecentTurns)
|
||
|
||
sessionIDs := make([]uuid.UUID, 0, len(rows))
|
||
for _, r := range rows {
|
||
sessionIDs = append(sessionIDs, r.id)
|
||
}
|
||
|
||
previewsBySession, err := s.loadPreviews(ctx, sessionIDs, turnLimit)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
out := make([]*Session, 0, len(rows))
|
||
for _, r := range rows {
|
||
scope, scopeErr := s.loadScope(ctx, r.id)
|
||
if scopeErr != nil {
|
||
return nil, scopeErr
|
||
}
|
||
previews := previewsBySession[r.id]
|
||
if previews == nil {
|
||
previews = []TurnPreview{}
|
||
}
|
||
out = append(out, rowToSession(r.id, r.clientID, r.createdBy, r.title, r.state,
|
||
r.createdAt, r.updatedAt, r.lastTerminal, scope, previews, r.isDeleted))
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// loadPreviews calls ListLastTurnsForSessions and groups the result by
|
||
// session_id. turnLimit==0 short-circuits to an empty map so no DB call
|
||
// is made.
|
||
func (s *Service) loadPreviews(ctx context.Context, sessionIDs []uuid.UUID, turnLimit int) (map[uuid.UUID][]TurnPreview, error) {
|
||
out := map[uuid.UUID][]TurnPreview{}
|
||
if turnLimit <= 0 || len(sessionIDs) == 0 {
|
||
return out, nil
|
||
}
|
||
rows, err := s.dbCfg.GetDBQueries().ListLastTurnsForSessions(ctx, &repository.ListLastTurnsForSessionsParams{
|
||
SessionIds: sessionIDs,
|
||
TurnLimit: int32(turnLimit), //nolint:gosec // bounded by cfg.RecentTurns in [1, 5]
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("bot: list last turns: %w", err)
|
||
}
|
||
for _, r := range rows {
|
||
out[r.SessionID] = append(out[r.SessionID], TurnPreview{
|
||
Ordinal: r.Ordinal,
|
||
Prompt: r.Prompt,
|
||
Completion: r.Completion,
|
||
Status: r.Status,
|
||
CreatedAt: r.CreatedAt.Time,
|
||
})
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// validateListInput re-checks the [1, 200] / >= 0 bounds plan §3
|
||
// mandates. The handler validates first; this is defense in depth so a
|
||
// non-HTTP caller cannot bypass the bounds.
|
||
func validateListInput(input ListSessionsInput) error {
|
||
if input.Limit < 1 || input.Limit > 200 {
|
||
return ErrInvalidLimit
|
||
}
|
||
if input.Offset < 0 {
|
||
return ErrInvalidOffset
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// rowToSession is the single place service rows are translated into the
|
||
// domain Session struct. State is normalized to a non-nil JSON object
|
||
// so handlers do not have to special-case the empty case when
|
||
// marshaling a map.
|
||
func rowToSession(id uuid.UUID, clientID, createdBy, title string, state []byte,
|
||
createdAt, updatedAt time.Time,
|
||
lastTerminal int32, scope SessionScope, previews []TurnPreview, isDeleted bool) *Session {
|
||
normalizedState := state
|
||
if len(normalizedState) == 0 {
|
||
normalizedState = []byte("{}")
|
||
}
|
||
return &Session{
|
||
ID: id,
|
||
ClientID: clientID,
|
||
CreatedBy: createdBy,
|
||
Title: title,
|
||
State: json.RawMessage(normalizedState),
|
||
LastTurn: lastTerminal,
|
||
CreatedAt: createdAt,
|
||
UpdatedAt: updatedAt,
|
||
Scope: scope,
|
||
RecentTurns: previews,
|
||
IsDeleted: isDeleted,
|
||
}
|
||
}
|