mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: chatd refactor (#26270)
Implements the chatd stabilization RFC. Combines: - https://github.com/coder/coder/pull/25908 - https://github.com/coder/coder/pull/25923 - https://github.com/coder/coder/pull/26109 - https://github.com/coder/coder/pull/26110 - https://github.com/coder/coder/pull/26111 - https://github.com/coder/coder/pull/26112
This commit is contained in:
Hugo Dutka
@@ -1692,9 +1692,6 @@ func (q *querier) ArchiveUnusedTemplateVersions(ctx context.Context, arg databas
|
||||
}
|
||||
|
||||
func (q *querier) AutoArchiveInactiveChats(ctx context.Context, arg database.AutoArchiveInactiveChatsParams) ([]database.AutoArchiveInactiveChatsRow, error) {
|
||||
// Background write by dbpurge. The LATERAL read of chat_messages rows
|
||||
// happens below the RBAC boundary; only the chat row itself requires
|
||||
// authorization.
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1718,6 +1715,13 @@ func (q *querier) BackoffChatDiffStatus(ctx context.Context, arg database.Backof
|
||||
return q.db.BackoffChatDiffStatus(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.BatchDeleteChatHeartbeats(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error {
|
||||
// Could be any workspace agent and checking auth to each workspace agent is overkill for
|
||||
// the purpose of this function.
|
||||
@@ -1743,6 +1747,13 @@ func (q *querier) BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg datab
|
||||
return q.db.BatchUpdateWorkspaceNextStartAt(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.BatchUpsertChatHeartbeats(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
|
||||
return err
|
||||
@@ -1857,6 +1868,14 @@ func (q *querier) CountAuditLogs(ctx context.Context, arg database.CountAuditLog
|
||||
return q.db.CountAuthorizedAuditLogs(ctx, arg, prep)
|
||||
}
|
||||
|
||||
func (q *querier) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
_, err := q.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.CountChatQueuedMessages(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) {
|
||||
// Just like the actual query, shortcut if the user is an owner.
|
||||
err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceConnectionLog)
|
||||
@@ -1954,6 +1973,18 @@ func (q *querier) DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) e
|
||||
return q.db.DeleteAPIKeysByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error {
|
||||
chat, err := q.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.DeleteAllChatHeartbeats(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error {
|
||||
chat, err := q.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
@@ -1965,6 +1996,18 @@ func (q *querier) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.U
|
||||
return q.db.DeleteAllChatQueuedMessages(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
|
||||
return nil, err
|
||||
@@ -2043,6 +2086,18 @@ func (q *querier) DeleteChatQueuedMessage(ctx context.Context, arg database.Dele
|
||||
return q.db.DeleteChatQueuedMessage(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.DeleteChatQueuedMessageReturningCount(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return err
|
||||
@@ -2315,6 +2370,13 @@ func (q *querier) DeleteRuntimeConfig(ctx context.Context, key string) error {
|
||||
return q.db.DeleteRuntimeConfig(ctx, key)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.DeleteStaleChatHeartbeats(ctx, staleSeconds)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
|
||||
return database.DeleteTailnetPeerRow{}, err
|
||||
@@ -2826,6 +2888,13 @@ func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUI
|
||||
return q.db.GetAuthorizationUserRoles(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetAutoArchiveInactiveChatCandidates(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceBoundaryLog); err != nil {
|
||||
return database.BoundaryLog{}, err
|
||||
@@ -2877,6 +2946,10 @@ func (q *querier) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat,
|
||||
return fetch(q.log, q.auth, q.db.GetChatByID)(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
return fetch(q.log, q.auth, q.db.GetChatByIDForShare)(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
return fetch(q.log, q.auth, q.db.GetChatByIDForUpdate)(ctx, id)
|
||||
}
|
||||
@@ -3048,6 +3121,18 @@ func (q *querier) GetChatExploreModelOverride(ctx context.Context) (string, erro
|
||||
return q.db.GetChatExploreModelOverride(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) {
|
||||
// This is a read-only query: it returns the chat IDs that belong
|
||||
// to a family. Authorize as Read against the root chat. The
|
||||
// individual SetArchived (or other) transitions that consume
|
||||
// these IDs run their own per-row authorization, so we do not
|
||||
// gate the listing itself on Update permission.
|
||||
if _, err := q.GetChatByID(ctx, id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatFamilyIDsByRootID(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) {
|
||||
file, err := q.db.GetChatFileByID(ctx, id)
|
||||
if err != nil {
|
||||
@@ -3114,6 +3199,14 @@ func (q *querier) GetChatGeneralModelOverride(ctx context.Context) (string, erro
|
||||
return q.db.GetChatGeneralModelOverride(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) {
|
||||
_, err := q.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return database.ChatHeartbeat{}, err
|
||||
}
|
||||
return q.db.GetChatHeartbeat(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) {
|
||||
// The include-default-system-prompt flag is a deployment-wide setting read
|
||||
// during chat creation by every authenticated user, so no RBAC policy
|
||||
@@ -3174,6 +3267,14 @@ func (q *querier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg
|
||||
return q.db.GetChatMessagesByChatIDDescPaginated(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) {
|
||||
_, err := q.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatMessagesByRevisionForStream(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) {
|
||||
// Authorize read on the parent chat.
|
||||
_, err := q.GetChatByID(ctx, chatID)
|
||||
@@ -3222,6 +3323,22 @@ func (q *querier) GetChatPlanModeInstructions(ctx context.Context) (string, erro
|
||||
return q.db.GetChatPlanModeInstructions(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) {
|
||||
_, err := q.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return database.ChatQueuedMessage{}, err
|
||||
}
|
||||
return q.db.GetChatQueuedMessageByID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) {
|
||||
_, err := q.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return database.ChatQueuedMessage{}, err
|
||||
}
|
||||
return q.db.GetChatQueuedMessageHead(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
_, err := q.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
@@ -3230,6 +3347,14 @@ func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (
|
||||
return q.db.GetChatQueuedMessages(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
_, err := q.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatQueuedMessagesByPosition(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatRetentionDays(ctx context.Context) (int32, error) {
|
||||
// Chat retention is a deployment-wide config read by dbpurge.
|
||||
// Only requires a valid actor in context.
|
||||
@@ -3239,6 +3364,13 @@ func (q *querier) GetChatRetentionDays(ctx context.Context) (int32, error) {
|
||||
return q.db.GetChatRetentionDays(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatStreamSyncRows(ctx, ids)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatSystemPrompt(ctx context.Context) (string, error) {
|
||||
// The system prompt is a deployment-wide setting read during chat
|
||||
// creation by every authenticated user, so no RBAC policy check
|
||||
@@ -3311,6 +3443,13 @@ func (q *querier) GetChatUserPromptsByChatID(ctx context.Context, arg database.G
|
||||
return q.db.GetChatUserPromptsByChatID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatWorkerAcquisitionCandidates(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatWorkspaceTTL(ctx context.Context) (string, error) {
|
||||
// The workspace-TTL setting is a deployment-wide value read by any
|
||||
// authenticated chat user. We only require that an explicit actor is
|
||||
@@ -3333,6 +3472,13 @@ func (q *querier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) ([
|
||||
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByChatFileID)(ctx, fileID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatsByIDsForRunnerSync(ctx, ids)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByWorkspaceIDs)(ctx, ids)
|
||||
}
|
||||
@@ -3403,6 +3549,10 @@ func (q *querier) GetDERPMeshKey(ctx context.Context) (string, error) {
|
||||
return q.db.GetDERPMeshKey(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
return q.db.GetDatabaseNow(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) {
|
||||
// Reading the default model config is needed for chat creation.
|
||||
// TODO(CODAGT-161): scope this check when org context is available.
|
||||
@@ -5468,6 +5618,18 @@ func (q *querier) GetWorkspacesForWorkspaceMetrics(ctx context.Context) ([]datab
|
||||
return q.db.GetWorkspacesForWorkspaceMetrics(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, id)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.IncrementChatGenerationAttempt(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) {
|
||||
return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg)
|
||||
}
|
||||
@@ -5638,6 +5800,18 @@ func (q *querier) InsertChatQueuedMessage(ctx context.Context, arg database.Inse
|
||||
return q.db.InsertChatQueuedMessage(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return database.ChatQueuedMessage{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return database.ChatQueuedMessage{}, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.InsertChatQueuedMessageWithCreator(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceCryptoKey); err != nil {
|
||||
return database.CryptoKey{}, err
|
||||
@@ -6210,6 +6384,14 @@ func (q *querier) InsertWorkspaceResourceMetadata(ctx context.Context, arg datab
|
||||
return q.db.InsertWorkspaceResourceMetadata(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) {
|
||||
_, err := q.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return q.db.IsChatHeartbeatStale(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
@@ -6394,6 +6576,18 @@ func (q *querier) ListWorkspaceAgentPortShares(ctx context.Context, workspaceID
|
||||
return q.db.ListWorkspaceAgentPortShares(ctx, workspaceID)
|
||||
}
|
||||
|
||||
func (q *querier) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, id)
|
||||
if err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.LockChatAndBumpSnapshotVersion(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error {
|
||||
resource := rbac.ResourceInboxNotification.WithOwner(arg.UserID.String())
|
||||
|
||||
@@ -6500,6 +6694,18 @@ func (q *querier) ReorderChatQueuedMessageToFront(ctx context.Context, arg datab
|
||||
return q.db.ReorderChatQueuedMessageToFront(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.ReorderChatQueuedMessageToHead(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
|
||||
return database.ResolveUserChatSpendLimitRow{}, err
|
||||
@@ -6742,6 +6948,18 @@ func (q *querier) UpdateChatDebugStep(ctx context.Context, arg database.UpdateCh
|
||||
return q.db.UpdateChatDebugStep(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ID)
|
||||
if err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.UpdateChatExecutionState(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
|
||||
// The batch heartbeat is a system-level operation filtered by
|
||||
// worker_id. Authorization is enforced by the AsChatd context
|
||||
@@ -6864,6 +7082,19 @@ func (q *querier) UpdateChatPlanModeByID(ctx context.Context, arg database.Updat
|
||||
return q.db.UpdateChatPlanModeByID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) {
|
||||
// UpdateChatRetryState is used by the chat processor to publish
|
||||
// transient retry state. It should be called with system context.
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ID)
|
||||
if err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
return q.db.UpdateChatRetryState(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) {
|
||||
// UpdateChatStatus is used by the chat processor to change chat status.
|
||||
// It should be called with system context.
|
||||
@@ -8243,6 +8474,18 @@ func (q *querier) UpsertChatGeneralModelOverride(ctx context.Context, value stri
|
||||
return q.db.UpsertChatGeneralModelOverride(ctx, value)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = chat
|
||||
return q.db.UpsertChatHeartbeat(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return err
|
||||
|
||||
@@ -537,6 +537,21 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().AcquireChats(gomock.Any(), arg).Return([]database.Chat{chat}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat})
|
||||
}))
|
||||
s.Run("GetChatWorkerAcquisitionCandidates", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetChatWorkerAcquisitionCandidatesParams{
|
||||
StaleSeconds: 30,
|
||||
LimitCount: 100,
|
||||
}
|
||||
row := testutil.Fake(s.T(), faker, database.GetChatWorkerAcquisitionCandidatesRow{})
|
||||
dbm.EXPECT().GetChatWorkerAcquisitionCandidates(gomock.Any(), arg).Return([]database.GetChatWorkerAcquisitionCandidatesRow{row}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.GetChatWorkerAcquisitionCandidatesRow{row})
|
||||
}))
|
||||
s.Run("GetChatsByIDsForRunnerSync", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
ids := []uuid.UUID{uuid.New(), uuid.New()}
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{ID: ids[0]})
|
||||
dbm.EXPECT().GetChatsByIDsForRunnerSync(gomock.Any(), ids).Return([]database.Chat{chat}, nil).AnyTimes()
|
||||
check.Args(ids).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat})
|
||||
}))
|
||||
s.Run("DeleteAllChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
@@ -763,6 +778,24 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().GetChatByIDForUpdate(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(chat)
|
||||
}))
|
||||
s.Run("GetChatByIDForShare", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByIDForShare(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(chat)
|
||||
}))
|
||||
s.Run("GetChatStreamSyncRows", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
ids := []uuid.UUID{uuid.New(), uuid.New()}
|
||||
rows := []database.GetChatStreamSyncRowsRow{{ID: ids[0]}}
|
||||
dbm.EXPECT().GetChatStreamSyncRows(gomock.Any(), ids).Return(rows, nil).AnyTimes()
|
||||
check.Args(ids).Asserts(rbac.ResourceChat, policy.ActionRead).Returns(rows)
|
||||
}))
|
||||
s.Run("GetChatFamilyIDsByRootID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
ids := []uuid.UUID{chat.ID}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatFamilyIDsByRootID(gomock.Any(), chat.ID).Return(ids, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(ids)
|
||||
}))
|
||||
s.Run("GetChatsByWorkspaceIDs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chatA := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
chatB := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
@@ -950,9 +983,10 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().UpsertChatAutoArchiveDays(gomock.Any(), int32(90)).Return(nil).AnyTimes()
|
||||
check.Args(int32(90)).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("AutoArchiveInactiveChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().AutoArchiveInactiveChats(gomock.Any(), database.AutoArchiveInactiveChatsParams{}).Return([]database.AutoArchiveInactiveChatsRow{}, nil).AnyTimes()
|
||||
check.Args(database.AutoArchiveInactiveChatsParams{}).Asserts(rbac.ResourceChat, policy.ActionUpdate)
|
||||
s.Run("GetAutoArchiveInactiveChatCandidates", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetAutoArchiveInactiveChatCandidatesParams{LimitCount: 100}
|
||||
dbm.EXPECT().GetAutoArchiveInactiveChatCandidates(gomock.Any(), arg).Return([]database.GetAutoArchiveInactiveChatCandidatesRow{}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.GetAutoArchiveInactiveChatCandidatesRow{})
|
||||
}))
|
||||
s.Run("GetChatMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
@@ -985,6 +1019,14 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().GetChatMessagesByChatIDDescPaginated(gomock.Any(), arg).Return(msgs, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs)
|
||||
}))
|
||||
s.Run("GetChatMessagesByRevisionForStream", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
msgs := []database.ChatMessage{testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID})}
|
||||
arg := database.GetChatMessagesByRevisionForStreamParams{ChatID: chat.ID, AfterRevision: 1}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatMessagesByRevisionForStream(gomock.Any(), arg).Return(msgs, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs)
|
||||
}))
|
||||
s.Run("GetChatUserPromptsByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
rows := []database.GetChatUserPromptsByChatIDRow{{ID: 1, Text: "hello"}}
|
||||
@@ -1214,6 +1256,136 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().UpdateChatACLByID(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionShare).Returns()
|
||||
}))
|
||||
s.Run("LockChatAndBumpSnapshotVersion", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().LockChatAndBumpSnapshotVersion(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(chat)
|
||||
}))
|
||||
s.Run("UpdateChatExecutionState", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.UpdateChatExecutionStateParams{ID: chat.ID, Status: database.ChatStatusRunning}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().UpdateChatExecutionState(gomock.Any(), arg).Return(chat, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat)
|
||||
}))
|
||||
s.Run("IncrementChatGenerationAttempt", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().IncrementChatGenerationAttempt(gomock.Any(), chat.ID).Return(int64(7), nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(int64(7))
|
||||
}))
|
||||
s.Run("UpdateChatRetryState", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.UpdateChatRetryStateParams{ID: chat.ID, RetryState: []byte(`{"attempt":1}`)}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().UpdateChatRetryState(gomock.Any(), arg).Return(chat, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat)
|
||||
}))
|
||||
s.Run("GetDatabaseNow", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
now := time.Now()
|
||||
dbm.EXPECT().GetDatabaseNow(gomock.Any()).Return(now, nil).AnyTimes()
|
||||
check.Args().Asserts().Returns(now)
|
||||
}))
|
||||
s.Run("InsertChatQueuedMessageWithCreator", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageWithCreatorParams{ChatID: chat.ID})
|
||||
qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().InsertChatQueuedMessageWithCreator(gomock.Any(), arg).Return(qm, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(qm)
|
||||
}))
|
||||
s.Run("GetChatQueuedMessagesByPosition", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
qms := []database.ChatQueuedMessage{}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatQueuedMessagesByPosition(gomock.Any(), chat.ID).Return(qms, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qms)
|
||||
}))
|
||||
s.Run("CountChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().CountChatQueuedMessages(gomock.Any(), chat.ID).Return(int64(3), nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(int64(3))
|
||||
}))
|
||||
s.Run("GetChatQueuedMessageHead", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatQueuedMessageHead(gomock.Any(), chat.ID).Return(qm, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qm)
|
||||
}))
|
||||
s.Run("GetChatQueuedMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID})
|
||||
arg := database.GetChatQueuedMessageByIDParams{ID: qm.ID, ChatID: chat.ID}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatQueuedMessageByID(gomock.Any(), arg).Return(qm, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(qm)
|
||||
}))
|
||||
s.Run("DeleteChatQueuedMessageReturningCount", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.DeleteChatQueuedMessageReturningCountParams{ID: 1, ChatID: chat.ID}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteChatQueuedMessageReturningCount(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
s.Run("DeleteAllChatQueuedMessagesReturningCount", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteAllChatQueuedMessagesReturningCount(gomock.Any(), chat.ID).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
s.Run("ReorderChatQueuedMessageToHead", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.ReorderChatQueuedMessageToHeadParams{ChatID: chat.ID, ID: 1}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().ReorderChatQueuedMessageToHead(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
s.Run("UpsertChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.UpsertChatHeartbeatParams{ChatID: chat.ID, RunnerID: uuid.New()}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().UpsertChatHeartbeat(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("BatchUpsertChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.BatchUpsertChatHeartbeatsParams{ChatIds: []uuid.UUID{uuid.New()}, RunnerIds: []uuid.UUID{uuid.New()}}
|
||||
dbm.EXPECT().BatchUpsertChatHeartbeats(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("GetChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.GetChatHeartbeatParams{ChatID: chat.ID, RunnerID: uuid.New()}
|
||||
hb := database.ChatHeartbeat{ChatID: chat.ID, RunnerID: arg.RunnerID}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetChatHeartbeat(gomock.Any(), arg).Return(hb, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(hb)
|
||||
}))
|
||||
s.Run("IsChatHeartbeatStale", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.IsChatHeartbeatStaleParams{ChatID: chat.ID, RunnerID: uuid.New(), StaleSeconds: 30}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().IsChatHeartbeatStale(gomock.Any(), arg).Return(false, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(false)
|
||||
}))
|
||||
s.Run("DeleteAllChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteAllChatHeartbeats(gomock.Any(), chat.ID).Return(nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("BatchDeleteChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.BatchDeleteChatHeartbeatsParams{ChatIds: []uuid.UUID{uuid.New()}, RunnerIds: []uuid.UUID{uuid.New()}}
|
||||
dbm.EXPECT().BatchDeleteChatHeartbeats(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
s.Run("DeleteStaleChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
const staleSeconds int32 = 30
|
||||
dbm.EXPECT().DeleteStaleChatHeartbeats(gomock.Any(), staleSeconds).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(staleSeconds).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
s.Run("UpdateChatByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.UpdateChatByIDParams{
|
||||
@@ -1408,6 +1580,14 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().BackoffChatDiffStatus(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("AutoArchiveInactiveChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.AutoArchiveInactiveChatsParams{
|
||||
ArchiveCutoff: dbtime.Now(),
|
||||
LimitCount: 100,
|
||||
}
|
||||
dbm.EXPECT().AutoArchiveInactiveChats(gomock.Any(), arg).Return([]database.AutoArchiveInactiveChatsRow{}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.AutoArchiveInactiveChatsRow{})
|
||||
}))
|
||||
s.Run("UpsertChatIncludeDefaultSystemPrompt", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().UpsertChatIncludeDefaultSystemPrompt(gomock.Any(), false).Return(nil).AnyTimes()
|
||||
check.Args(false).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
@@ -1697,9 +1877,9 @@ func (s *MethodTestSuite) TestChats() {
|
||||
s.Run("UpdateChatLastTurnSummary", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedUpdatedAt: chat.UpdatedAt,
|
||||
LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true},
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true},
|
||||
}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().UpdateChatLastTurnSummary(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
|
||||
|
||||
@@ -115,13 +115,23 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat
|
||||
if seed.Content.Valid {
|
||||
content = string(seed.Content.RawMessage)
|
||||
}
|
||||
role := takeFirst(seed.Role, database.ChatMessageRoleUser)
|
||||
apiKeyID := seed.APIKeyID.String
|
||||
// Mint a real API key for user turns so the api_key_id foreign key is
|
||||
// satisfied. Without a creator we leave it empty, which the insert query
|
||||
// stores as NULL.
|
||||
if role == database.ChatMessageRoleUser && apiKeyID == "" &&
|
||||
seed.CreatedBy.Valid && seed.CreatedBy.UUID != uuid.Nil {
|
||||
key, _ := APIKey(t, db, database.APIKey{UserID: seed.CreatedBy.UUID})
|
||||
apiKeyID = key.ID
|
||||
}
|
||||
|
||||
msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{
|
||||
ChatID: seed.ChatID,
|
||||
CreatedBy: []uuid.UUID{seed.CreatedBy.UUID},
|
||||
APIKeyID: []string{seed.APIKeyID.String},
|
||||
APIKeyID: []string{apiKeyID},
|
||||
ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID},
|
||||
Role: []database.ChatMessageRole{takeFirst(seed.Role, database.ChatMessageRoleUser)},
|
||||
Role: []database.ChatMessageRole{role},
|
||||
Content: []string{content},
|
||||
ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)},
|
||||
Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)},
|
||||
|
||||
+216
@@ -202,6 +202,14 @@ func (m queryMetricsStore) BackoffChatDiffStatus(ctx context.Context, arg databa
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.BatchDeleteChatHeartbeats(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("BatchDeleteChatHeartbeats").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BatchDeleteChatHeartbeats").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.BatchUpdateWorkspaceAgentMetadata(ctx, arg)
|
||||
@@ -226,6 +234,14 @@ func (m queryMetricsStore) BatchUpdateWorkspaceNextStartAt(ctx context.Context,
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.BatchUpsertChatHeartbeats(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("BatchUpsertChatHeartbeats").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BatchUpsertChatHeartbeats").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.BatchUpsertConnectionLogs(ctx, arg)
|
||||
@@ -322,6 +338,14 @@ func (m queryMetricsStore) CountAuditLogs(ctx context.Context, arg database.Coun
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.CountChatQueuedMessages(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("CountChatQueuedMessages").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "CountChatQueuedMessages").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.CountConnectionLogs(ctx, arg)
|
||||
@@ -418,6 +442,14 @@ func (m queryMetricsStore) DeleteAPIKeysByUserID(ctx context.Context, userID uui
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteAllChatHeartbeats(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("DeleteAllChatHeartbeats").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAllChatHeartbeats").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteAllChatQueuedMessages(ctx, chatID)
|
||||
@@ -426,6 +458,14 @@ func (m queryMetricsStore) DeleteAllChatQueuedMessages(ctx context.Context, chat
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("DeleteAllChatQueuedMessagesReturningCount").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAllChatQueuedMessagesReturningCount").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteAllTailnetTunnels(ctx, arg)
|
||||
@@ -498,6 +538,14 @@ func (m queryMetricsStore) DeleteChatQueuedMessage(ctx context.Context, arg data
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteChatQueuedMessageReturningCount(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("DeleteChatQueuedMessageReturningCount").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatQueuedMessageReturningCount").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteChatUsageLimitGroupOverride(ctx, groupID)
|
||||
@@ -778,6 +826,14 @@ func (m queryMetricsStore) DeleteRuntimeConfig(ctx context.Context, key string)
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteStaleChatHeartbeats(ctx, staleSeconds)
|
||||
m.queryLatencies.WithLabelValues("DeleteStaleChatHeartbeats").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteStaleChatHeartbeats").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteTailnetPeer(ctx, arg)
|
||||
@@ -1274,6 +1330,14 @@ func (m queryMetricsStore) GetAuthorizationUserRoles(ctx context.Context, userID
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAutoArchiveInactiveChatCandidates(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetAutoArchiveInactiveChatCandidates").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAutoArchiveInactiveChatCandidates").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetBoundaryLogByID(ctx, id)
|
||||
@@ -1322,6 +1386,14 @@ func (m queryMetricsStore) GetChatByID(ctx context.Context, id uuid.UUID) (datab
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatByIDForShare(ctx, id)
|
||||
m.queryLatencies.WithLabelValues("GetChatByIDForShare").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatByIDForShare").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatByIDForUpdate(ctx, id)
|
||||
@@ -1450,6 +1522,14 @@ func (m queryMetricsStore) GetChatExploreModelOverride(ctx context.Context) (str
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatFamilyIDsByRootID(ctx, id)
|
||||
m.queryLatencies.WithLabelValues("GetChatFamilyIDsByRootID").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatFamilyIDsByRootID").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatFileByID(ctx, id)
|
||||
@@ -1482,6 +1562,14 @@ func (m queryMetricsStore) GetChatGeneralModelOverride(ctx context.Context) (str
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatHeartbeat(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetChatHeartbeat").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatHeartbeat").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatIncludeDefaultSystemPrompt(ctx)
|
||||
@@ -1530,6 +1618,14 @@ func (m queryMetricsStore) GetChatMessagesByChatIDDescPaginated(ctx context.Cont
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatMessagesByRevisionForStream(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetChatMessagesByRevisionForStream").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByRevisionForStream").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatMessagesForPromptByChatID(ctx, chatID)
|
||||
@@ -1578,6 +1674,22 @@ func (m queryMetricsStore) GetChatPlanModeInstructions(ctx context.Context) (str
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatQueuedMessageByID(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetChatQueuedMessageByID").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessageByID").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatQueuedMessageHead(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("GetChatQueuedMessageHead").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessageHead").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatQueuedMessages(ctx, chatID)
|
||||
@@ -1586,6 +1698,14 @@ func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uui
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatQueuedMessagesByPosition(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("GetChatQueuedMessagesByPosition").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessagesByPosition").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatRetentionDays(ctx context.Context) (int32, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatRetentionDays(ctx)
|
||||
@@ -1594,6 +1714,14 @@ func (m queryMetricsStore) GetChatRetentionDays(ctx context.Context) (int32, err
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatStreamSyncRows(ctx, ids)
|
||||
m.queryLatencies.WithLabelValues("GetChatStreamSyncRows").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatStreamSyncRows").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatSystemPrompt(ctx context.Context) (string, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatSystemPrompt(ctx)
|
||||
@@ -1658,6 +1786,14 @@ func (m queryMetricsStore) GetChatUserPromptsByChatID(ctx context.Context, arg d
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatWorkerAcquisitionCandidates(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetChatWorkerAcquisitionCandidates").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatWorkerAcquisitionCandidates").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatWorkspaceTTL(ctx context.Context) (string, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatWorkspaceTTL(ctx)
|
||||
@@ -1682,6 +1818,14 @@ func (m queryMetricsStore) GetChatsByChatFileID(ctx context.Context, fileID uuid
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatsByIDsForRunnerSync(ctx, ids)
|
||||
m.queryLatencies.WithLabelValues("GetChatsByIDsForRunnerSync").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatsByIDsForRunnerSync").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatsByWorkspaceIDs(ctx, ids)
|
||||
@@ -1754,6 +1898,14 @@ func (m queryMetricsStore) GetDERPMeshKey(ctx context.Context) (string, error) {
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetDatabaseNow(ctx)
|
||||
m.queryLatencies.WithLabelValues("GetDatabaseNow").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetDatabaseNow").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetDefaultChatModelConfig(ctx)
|
||||
@@ -3706,6 +3858,14 @@ func (m queryMetricsStore) GetWorkspacesForWorkspaceMetrics(ctx context.Context)
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.IncrementChatGenerationAttempt(ctx, id)
|
||||
m.queryLatencies.WithLabelValues("IncrementChatGenerationAttempt").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "IncrementChatGenerationAttempt").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.InsertAIBridgeInterception(ctx, arg)
|
||||
@@ -3866,6 +4026,14 @@ func (m queryMetricsStore) InsertChatQueuedMessage(ctx context.Context, arg data
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.InsertChatQueuedMessageWithCreator(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("InsertChatQueuedMessageWithCreator").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatQueuedMessageWithCreator").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.InsertCryptoKey(ctx, arg)
|
||||
@@ -4362,6 +4530,14 @@ func (m queryMetricsStore) InsertWorkspaceResourceMetadata(ctx context.Context,
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.IsChatHeartbeatStale(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("IsChatHeartbeatStale").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "IsChatHeartbeatStale").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.LinkChatFiles(ctx, arg)
|
||||
@@ -4546,6 +4722,14 @@ func (m queryMetricsStore) ListWorkspaceAgentPortShares(ctx context.Context, wor
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.LockChatAndBumpSnapshotVersion(ctx, id)
|
||||
m.queryLatencies.WithLabelValues("LockChatAndBumpSnapshotVersion").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "LockChatAndBumpSnapshotVersion").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.MarkAllInboxNotificationsAsRead(ctx, arg)
|
||||
@@ -4634,6 +4818,14 @@ func (m queryMetricsStore) ReorderChatQueuedMessageToFront(ctx context.Context,
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ReorderChatQueuedMessageToHead(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("ReorderChatQueuedMessageToHead").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ReorderChatQueuedMessageToHead").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID)
|
||||
@@ -4826,6 +5018,14 @@ func (m queryMetricsStore) UpdateChatDebugStep(ctx context.Context, arg database
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpdateChatExecutionState(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpdateChatExecutionState").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatExecutionState").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpdateChatHeartbeats(ctx, arg)
|
||||
@@ -4914,6 +5114,14 @@ func (m queryMetricsStore) UpdateChatPlanModeByID(ctx context.Context, arg datab
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpdateChatRetryState(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpdateChatRetryState").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatRetryState").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpdateChatStatus(ctx, arg)
|
||||
@@ -5866,6 +6074,14 @@ func (m queryMetricsStore) UpsertChatGeneralModelOverride(ctx context.Context, v
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.UpsertChatHeartbeat(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertChatHeartbeat").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatHeartbeat").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.UpsertChatIncludeDefaultSystemPrompt(ctx, includeDefaultSystemPrompt)
|
||||
|
||||
Generated
+402
@@ -223,6 +223,21 @@ func (mr *MockStoreMockRecorder) BackoffChatDiffStatus(ctx, arg any) *gomock.Cal
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BackoffChatDiffStatus", reflect.TypeOf((*MockStore)(nil).BackoffChatDiffStatus), ctx, arg)
|
||||
}
|
||||
|
||||
// BatchDeleteChatHeartbeats mocks base method.
|
||||
func (m *MockStore) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "BatchDeleteChatHeartbeats", ctx, arg)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// BatchDeleteChatHeartbeats indicates an expected call of BatchDeleteChatHeartbeats.
|
||||
func (mr *MockStoreMockRecorder) BatchDeleteChatHeartbeats(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchDeleteChatHeartbeats", reflect.TypeOf((*MockStore)(nil).BatchDeleteChatHeartbeats), ctx, arg)
|
||||
}
|
||||
|
||||
// BatchUpdateWorkspaceAgentMetadata mocks base method.
|
||||
func (m *MockStore) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -265,6 +280,20 @@ func (mr *MockStoreMockRecorder) BatchUpdateWorkspaceNextStartAt(ctx, arg any) *
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpdateWorkspaceNextStartAt", reflect.TypeOf((*MockStore)(nil).BatchUpdateWorkspaceNextStartAt), ctx, arg)
|
||||
}
|
||||
|
||||
// BatchUpsertChatHeartbeats mocks base method.
|
||||
func (m *MockStore) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "BatchUpsertChatHeartbeats", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// BatchUpsertChatHeartbeats indicates an expected call of BatchUpsertChatHeartbeats.
|
||||
func (mr *MockStoreMockRecorder) BatchUpsertChatHeartbeats(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpsertChatHeartbeats", reflect.TypeOf((*MockStore)(nil).BatchUpsertChatHeartbeats), ctx, arg)
|
||||
}
|
||||
|
||||
// BatchUpsertConnectionLogs mocks base method.
|
||||
func (m *MockStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -484,6 +513,21 @@ func (mr *MockStoreMockRecorder) CountAuthorizedConnectionLogs(ctx, arg, prepare
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountAuthorizedConnectionLogs", reflect.TypeOf((*MockStore)(nil).CountAuthorizedConnectionLogs), ctx, arg, prepared)
|
||||
}
|
||||
|
||||
// CountChatQueuedMessages mocks base method.
|
||||
func (m *MockStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "CountChatQueuedMessages", ctx, chatID)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// CountChatQueuedMessages indicates an expected call of CountChatQueuedMessages.
|
||||
func (mr *MockStoreMockRecorder) CountChatQueuedMessages(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).CountChatQueuedMessages), ctx, chatID)
|
||||
}
|
||||
|
||||
// CountConnectionLogs mocks base method.
|
||||
func (m *MockStore) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -660,6 +704,20 @@ func (mr *MockStoreMockRecorder) DeleteAPIKeysByUserID(ctx, userID any) *gomock.
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAPIKeysByUserID", reflect.TypeOf((*MockStore)(nil).DeleteAPIKeysByUserID), ctx, userID)
|
||||
}
|
||||
|
||||
// DeleteAllChatHeartbeats mocks base method.
|
||||
func (m *MockStore) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAllChatHeartbeats", ctx, chatID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteAllChatHeartbeats indicates an expected call of DeleteAllChatHeartbeats.
|
||||
func (mr *MockStoreMockRecorder) DeleteAllChatHeartbeats(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatHeartbeats", reflect.TypeOf((*MockStore)(nil).DeleteAllChatHeartbeats), ctx, chatID)
|
||||
}
|
||||
|
||||
// DeleteAllChatQueuedMessages mocks base method.
|
||||
func (m *MockStore) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -674,6 +732,21 @@ func (mr *MockStoreMockRecorder) DeleteAllChatQueuedMessages(ctx, chatID any) *g
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).DeleteAllChatQueuedMessages), ctx, chatID)
|
||||
}
|
||||
|
||||
// DeleteAllChatQueuedMessagesReturningCount mocks base method.
|
||||
func (m *MockStore) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAllChatQueuedMessagesReturningCount", ctx, chatID)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DeleteAllChatQueuedMessagesReturningCount indicates an expected call of DeleteAllChatQueuedMessagesReturningCount.
|
||||
func (mr *MockStoreMockRecorder) DeleteAllChatQueuedMessagesReturningCount(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatQueuedMessagesReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteAllChatQueuedMessagesReturningCount), ctx, chatID)
|
||||
}
|
||||
|
||||
// DeleteAllTailnetTunnels mocks base method.
|
||||
func (m *MockStore) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -803,6 +876,21 @@ func (mr *MockStoreMockRecorder) DeleteChatQueuedMessage(ctx, arg any) *gomock.C
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessage", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessage), ctx, arg)
|
||||
}
|
||||
|
||||
// DeleteChatQueuedMessageReturningCount mocks base method.
|
||||
func (m *MockStore) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteChatQueuedMessageReturningCount", ctx, arg)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DeleteChatQueuedMessageReturningCount indicates an expected call of DeleteChatQueuedMessageReturningCount.
|
||||
func (mr *MockStoreMockRecorder) DeleteChatQueuedMessageReturningCount(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessageReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessageReturningCount), ctx, arg)
|
||||
}
|
||||
|
||||
// DeleteChatUsageLimitGroupOverride mocks base method.
|
||||
func (m *MockStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -1305,6 +1393,21 @@ func (mr *MockStoreMockRecorder) DeleteRuntimeConfig(ctx, key any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteRuntimeConfig", reflect.TypeOf((*MockStore)(nil).DeleteRuntimeConfig), ctx, key)
|
||||
}
|
||||
|
||||
// DeleteStaleChatHeartbeats mocks base method.
|
||||
func (m *MockStore) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteStaleChatHeartbeats", ctx, staleSeconds)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DeleteStaleChatHeartbeats indicates an expected call of DeleteStaleChatHeartbeats.
|
||||
func (mr *MockStoreMockRecorder) DeleteStaleChatHeartbeats(ctx, staleSeconds any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteStaleChatHeartbeats", reflect.TypeOf((*MockStore)(nil).DeleteStaleChatHeartbeats), ctx, staleSeconds)
|
||||
}
|
||||
|
||||
// DeleteTailnetPeer mocks base method.
|
||||
func (m *MockStore) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2341,6 +2444,21 @@ func (mr *MockStoreMockRecorder) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx,
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthorizedWorkspacesAndAgentsByOwnerID", reflect.TypeOf((*MockStore)(nil).GetAuthorizedWorkspacesAndAgentsByOwnerID), ctx, ownerID, prepared)
|
||||
}
|
||||
|
||||
// GetAutoArchiveInactiveChatCandidates mocks base method.
|
||||
func (m *MockStore) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAutoArchiveInactiveChatCandidates", ctx, arg)
|
||||
ret0, _ := ret[0].([]database.GetAutoArchiveInactiveChatCandidatesRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAutoArchiveInactiveChatCandidates indicates an expected call of GetAutoArchiveInactiveChatCandidates.
|
||||
func (mr *MockStoreMockRecorder) GetAutoArchiveInactiveChatCandidates(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAutoArchiveInactiveChatCandidates", reflect.TypeOf((*MockStore)(nil).GetAutoArchiveInactiveChatCandidates), ctx, arg)
|
||||
}
|
||||
|
||||
// GetBoundaryLogByID mocks base method.
|
||||
func (m *MockStore) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2431,6 +2549,21 @@ func (mr *MockStoreMockRecorder) GetChatByID(ctx, id any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatByID", reflect.TypeOf((*MockStore)(nil).GetChatByID), ctx, id)
|
||||
}
|
||||
|
||||
// GetChatByIDForShare mocks base method.
|
||||
func (m *MockStore) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatByIDForShare", ctx, id)
|
||||
ret0, _ := ret[0].(database.Chat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatByIDForShare indicates an expected call of GetChatByIDForShare.
|
||||
func (mr *MockStoreMockRecorder) GetChatByIDForShare(ctx, id any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatByIDForShare", reflect.TypeOf((*MockStore)(nil).GetChatByIDForShare), ctx, id)
|
||||
}
|
||||
|
||||
// GetChatByIDForUpdate mocks base method.
|
||||
func (m *MockStore) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2671,6 +2804,21 @@ func (mr *MockStoreMockRecorder) GetChatExploreModelOverride(ctx any) *gomock.Ca
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatExploreModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatExploreModelOverride), ctx)
|
||||
}
|
||||
|
||||
// GetChatFamilyIDsByRootID mocks base method.
|
||||
func (m *MockStore) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatFamilyIDsByRootID", ctx, id)
|
||||
ret0, _ := ret[0].([]uuid.UUID)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatFamilyIDsByRootID indicates an expected call of GetChatFamilyIDsByRootID.
|
||||
func (mr *MockStoreMockRecorder) GetChatFamilyIDsByRootID(ctx, id any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatFamilyIDsByRootID", reflect.TypeOf((*MockStore)(nil).GetChatFamilyIDsByRootID), ctx, id)
|
||||
}
|
||||
|
||||
// GetChatFileByID mocks base method.
|
||||
func (m *MockStore) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2731,6 +2879,21 @@ func (mr *MockStoreMockRecorder) GetChatGeneralModelOverride(ctx any) *gomock.Ca
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatGeneralModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatGeneralModelOverride), ctx)
|
||||
}
|
||||
|
||||
// GetChatHeartbeat mocks base method.
|
||||
func (m *MockStore) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatHeartbeat", ctx, arg)
|
||||
ret0, _ := ret[0].(database.ChatHeartbeat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatHeartbeat indicates an expected call of GetChatHeartbeat.
|
||||
func (mr *MockStoreMockRecorder) GetChatHeartbeat(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatHeartbeat", reflect.TypeOf((*MockStore)(nil).GetChatHeartbeat), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatIncludeDefaultSystemPrompt mocks base method.
|
||||
func (m *MockStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2821,6 +2984,21 @@ func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDDescPaginated(ctx, arg a
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDDescPaginated", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDDescPaginated), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatMessagesByRevisionForStream mocks base method.
|
||||
func (m *MockStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatMessagesByRevisionForStream", ctx, arg)
|
||||
ret0, _ := ret[0].([]database.ChatMessage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatMessagesByRevisionForStream indicates an expected call of GetChatMessagesByRevisionForStream.
|
||||
func (mr *MockStoreMockRecorder) GetChatMessagesByRevisionForStream(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByRevisionForStream", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByRevisionForStream), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatMessagesForPromptByChatID mocks base method.
|
||||
func (m *MockStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2911,6 +3089,36 @@ func (mr *MockStoreMockRecorder) GetChatPlanModeInstructions(ctx any) *gomock.Ca
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatPlanModeInstructions", reflect.TypeOf((*MockStore)(nil).GetChatPlanModeInstructions), ctx)
|
||||
}
|
||||
|
||||
// GetChatQueuedMessageByID mocks base method.
|
||||
func (m *MockStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatQueuedMessageByID", ctx, arg)
|
||||
ret0, _ := ret[0].(database.ChatQueuedMessage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatQueuedMessageByID indicates an expected call of GetChatQueuedMessageByID.
|
||||
func (mr *MockStoreMockRecorder) GetChatQueuedMessageByID(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessageByID", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessageByID), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatQueuedMessageHead mocks base method.
|
||||
func (m *MockStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatQueuedMessageHead", ctx, chatID)
|
||||
ret0, _ := ret[0].(database.ChatQueuedMessage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatQueuedMessageHead indicates an expected call of GetChatQueuedMessageHead.
|
||||
func (mr *MockStoreMockRecorder) GetChatQueuedMessageHead(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessageHead", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessageHead), ctx, chatID)
|
||||
}
|
||||
|
||||
// GetChatQueuedMessages mocks base method.
|
||||
func (m *MockStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2926,6 +3134,21 @@ func (mr *MockStoreMockRecorder) GetChatQueuedMessages(ctx, chatID any) *gomock.
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessages), ctx, chatID)
|
||||
}
|
||||
|
||||
// GetChatQueuedMessagesByPosition mocks base method.
|
||||
func (m *MockStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatQueuedMessagesByPosition", ctx, chatID)
|
||||
ret0, _ := ret[0].([]database.ChatQueuedMessage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatQueuedMessagesByPosition indicates an expected call of GetChatQueuedMessagesByPosition.
|
||||
func (mr *MockStoreMockRecorder) GetChatQueuedMessagesByPosition(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessagesByPosition", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessagesByPosition), ctx, chatID)
|
||||
}
|
||||
|
||||
// GetChatRetentionDays mocks base method.
|
||||
func (m *MockStore) GetChatRetentionDays(ctx context.Context) (int32, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -2941,6 +3164,21 @@ func (mr *MockStoreMockRecorder) GetChatRetentionDays(ctx any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatRetentionDays", reflect.TypeOf((*MockStore)(nil).GetChatRetentionDays), ctx)
|
||||
}
|
||||
|
||||
// GetChatStreamSyncRows mocks base method.
|
||||
func (m *MockStore) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatStreamSyncRows", ctx, ids)
|
||||
ret0, _ := ret[0].([]database.GetChatStreamSyncRowsRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatStreamSyncRows indicates an expected call of GetChatStreamSyncRows.
|
||||
func (mr *MockStoreMockRecorder) GetChatStreamSyncRows(ctx, ids any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatStreamSyncRows", reflect.TypeOf((*MockStore)(nil).GetChatStreamSyncRows), ctx, ids)
|
||||
}
|
||||
|
||||
// GetChatSystemPrompt mocks base method.
|
||||
func (m *MockStore) GetChatSystemPrompt(ctx context.Context) (string, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -3061,6 +3299,21 @@ func (mr *MockStoreMockRecorder) GetChatUserPromptsByChatID(ctx, arg any) *gomoc
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUserPromptsByChatID", reflect.TypeOf((*MockStore)(nil).GetChatUserPromptsByChatID), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatWorkerAcquisitionCandidates mocks base method.
|
||||
func (m *MockStore) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatWorkerAcquisitionCandidates", ctx, arg)
|
||||
ret0, _ := ret[0].([]database.GetChatWorkerAcquisitionCandidatesRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatWorkerAcquisitionCandidates indicates an expected call of GetChatWorkerAcquisitionCandidates.
|
||||
func (mr *MockStoreMockRecorder) GetChatWorkerAcquisitionCandidates(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatWorkerAcquisitionCandidates", reflect.TypeOf((*MockStore)(nil).GetChatWorkerAcquisitionCandidates), ctx, arg)
|
||||
}
|
||||
|
||||
// GetChatWorkspaceTTL mocks base method.
|
||||
func (m *MockStore) GetChatWorkspaceTTL(ctx context.Context) (string, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -3106,6 +3359,21 @@ func (mr *MockStoreMockRecorder) GetChatsByChatFileID(ctx, fileID any) *gomock.C
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatsByChatFileID", reflect.TypeOf((*MockStore)(nil).GetChatsByChatFileID), ctx, fileID)
|
||||
}
|
||||
|
||||
// GetChatsByIDsForRunnerSync mocks base method.
|
||||
func (m *MockStore) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatsByIDsForRunnerSync", ctx, ids)
|
||||
ret0, _ := ret[0].([]database.Chat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatsByIDsForRunnerSync indicates an expected call of GetChatsByIDsForRunnerSync.
|
||||
func (mr *MockStoreMockRecorder) GetChatsByIDsForRunnerSync(ctx, ids any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatsByIDsForRunnerSync", reflect.TypeOf((*MockStore)(nil).GetChatsByIDsForRunnerSync), ctx, ids)
|
||||
}
|
||||
|
||||
// GetChatsByWorkspaceIDs mocks base method.
|
||||
func (m *MockStore) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -3241,6 +3509,21 @@ func (mr *MockStoreMockRecorder) GetDERPMeshKey(ctx any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDERPMeshKey", reflect.TypeOf((*MockStore)(nil).GetDERPMeshKey), ctx)
|
||||
}
|
||||
|
||||
// GetDatabaseNow mocks base method.
|
||||
func (m *MockStore) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetDatabaseNow", ctx)
|
||||
ret0, _ := ret[0].(time.Time)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetDatabaseNow indicates an expected call of GetDatabaseNow.
|
||||
func (mr *MockStoreMockRecorder) GetDatabaseNow(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDatabaseNow", reflect.TypeOf((*MockStore)(nil).GetDatabaseNow), ctx)
|
||||
}
|
||||
|
||||
// GetDefaultChatModelConfig mocks base method.
|
||||
func (m *MockStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -6945,6 +7228,21 @@ func (mr *MockStoreMockRecorder) InTx(arg0, arg1 any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InTx", reflect.TypeOf((*MockStore)(nil).InTx), arg0, arg1)
|
||||
}
|
||||
|
||||
// IncrementChatGenerationAttempt mocks base method.
|
||||
func (m *MockStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "IncrementChatGenerationAttempt", ctx, id)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// IncrementChatGenerationAttempt indicates an expected call of IncrementChatGenerationAttempt.
|
||||
func (mr *MockStoreMockRecorder) IncrementChatGenerationAttempt(ctx, id any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementChatGenerationAttempt", reflect.TypeOf((*MockStore)(nil).IncrementChatGenerationAttempt), ctx, id)
|
||||
}
|
||||
|
||||
// InsertAIBridgeInterception mocks base method.
|
||||
func (m *MockStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -7245,6 +7543,21 @@ func (mr *MockStoreMockRecorder) InsertChatQueuedMessage(ctx, arg any) *gomock.C
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatQueuedMessage", reflect.TypeOf((*MockStore)(nil).InsertChatQueuedMessage), ctx, arg)
|
||||
}
|
||||
|
||||
// InsertChatQueuedMessageWithCreator mocks base method.
|
||||
func (m *MockStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "InsertChatQueuedMessageWithCreator", ctx, arg)
|
||||
ret0, _ := ret[0].(database.ChatQueuedMessage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// InsertChatQueuedMessageWithCreator indicates an expected call of InsertChatQueuedMessageWithCreator.
|
||||
func (mr *MockStoreMockRecorder) InsertChatQueuedMessageWithCreator(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatQueuedMessageWithCreator", reflect.TypeOf((*MockStore)(nil).InsertChatQueuedMessageWithCreator), ctx, arg)
|
||||
}
|
||||
|
||||
// InsertCryptoKey mocks base method.
|
||||
func (m *MockStore) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -8160,6 +8473,21 @@ func (mr *MockStoreMockRecorder) InsertWorkspaceResourceMetadata(ctx, arg any) *
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceResourceMetadata", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceResourceMetadata), ctx, arg)
|
||||
}
|
||||
|
||||
// IsChatHeartbeatStale mocks base method.
|
||||
func (m *MockStore) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "IsChatHeartbeatStale", ctx, arg)
|
||||
ret0, _ := ret[0].(bool)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// IsChatHeartbeatStale indicates an expected call of IsChatHeartbeatStale.
|
||||
func (mr *MockStoreMockRecorder) IsChatHeartbeatStale(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsChatHeartbeatStale", reflect.TypeOf((*MockStore)(nil).IsChatHeartbeatStale), ctx, arg)
|
||||
}
|
||||
|
||||
// LinkChatFiles mocks base method.
|
||||
func (m *MockStore) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -8565,6 +8893,21 @@ func (mr *MockStoreMockRecorder) ListWorkspaceAgentPortShares(ctx, workspaceID a
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListWorkspaceAgentPortShares", reflect.TypeOf((*MockStore)(nil).ListWorkspaceAgentPortShares), ctx, workspaceID)
|
||||
}
|
||||
|
||||
// LockChatAndBumpSnapshotVersion mocks base method.
|
||||
func (m *MockStore) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "LockChatAndBumpSnapshotVersion", ctx, id)
|
||||
ret0, _ := ret[0].(database.Chat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// LockChatAndBumpSnapshotVersion indicates an expected call of LockChatAndBumpSnapshotVersion.
|
||||
func (mr *MockStoreMockRecorder) LockChatAndBumpSnapshotVersion(ctx, id any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LockChatAndBumpSnapshotVersion", reflect.TypeOf((*MockStore)(nil).LockChatAndBumpSnapshotVersion), ctx, id)
|
||||
}
|
||||
|
||||
// MarkAllInboxNotificationsAsRead mocks base method.
|
||||
func (m *MockStore) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -8757,6 +9100,21 @@ func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToFront(ctx, arg any) *
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToFront", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToFront), ctx, arg)
|
||||
}
|
||||
|
||||
// ReorderChatQueuedMessageToHead mocks base method.
|
||||
func (m *MockStore) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ReorderChatQueuedMessageToHead", ctx, arg)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ReorderChatQueuedMessageToHead indicates an expected call of ReorderChatQueuedMessageToHead.
|
||||
func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToHead(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToHead", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToHead), ctx, arg)
|
||||
}
|
||||
|
||||
// ResolveUserChatSpendLimit mocks base method.
|
||||
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -9103,6 +9461,21 @@ func (mr *MockStoreMockRecorder) UpdateChatDebugStep(ctx, arg any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatDebugStep", reflect.TypeOf((*MockStore)(nil).UpdateChatDebugStep), ctx, arg)
|
||||
}
|
||||
|
||||
// UpdateChatExecutionState mocks base method.
|
||||
func (m *MockStore) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpdateChatExecutionState", ctx, arg)
|
||||
ret0, _ := ret[0].(database.Chat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpdateChatExecutionState indicates an expected call of UpdateChatExecutionState.
|
||||
func (mr *MockStoreMockRecorder) UpdateChatExecutionState(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatExecutionState", reflect.TypeOf((*MockStore)(nil).UpdateChatExecutionState), ctx, arg)
|
||||
}
|
||||
|
||||
// UpdateChatHeartbeats mocks base method.
|
||||
func (m *MockStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -9266,6 +9639,21 @@ func (mr *MockStoreMockRecorder) UpdateChatPlanModeByID(ctx, arg any) *gomock.Ca
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatPlanModeByID", reflect.TypeOf((*MockStore)(nil).UpdateChatPlanModeByID), ctx, arg)
|
||||
}
|
||||
|
||||
// UpdateChatRetryState mocks base method.
|
||||
func (m *MockStore) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpdateChatRetryState", ctx, arg)
|
||||
ret0, _ := ret[0].(database.Chat)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpdateChatRetryState indicates an expected call of UpdateChatRetryState.
|
||||
func (mr *MockStoreMockRecorder) UpdateChatRetryState(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatRetryState", reflect.TypeOf((*MockStore)(nil).UpdateChatRetryState), ctx, arg)
|
||||
}
|
||||
|
||||
// UpdateChatStatus mocks base method.
|
||||
func (m *MockStore) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -10990,6 +11378,20 @@ func (mr *MockStoreMockRecorder) UpsertChatGeneralModelOverride(ctx, value any)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatGeneralModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatGeneralModelOverride), ctx, value)
|
||||
}
|
||||
|
||||
// UpsertChatHeartbeat mocks base method.
|
||||
func (m *MockStore) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertChatHeartbeat", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// UpsertChatHeartbeat indicates an expected call of UpsertChatHeartbeat.
|
||||
func (mr *MockStoreMockRecorder) UpsertChatHeartbeat(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatHeartbeat", reflect.TypeOf((*MockStore)(nil).UpsertChatHeartbeat), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertChatIncludeDefaultSystemPrompt mocks base method.
|
||||
func (m *MockStore) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
@@ -1,18 +1,12 @@
|
||||
package dbpurge
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/dustin/go-humanize"
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
@@ -21,9 +15,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
"github.com/coder/coder/v2/coderd/notifications"
|
||||
"github.com/coder/coder/v2/coderd/pproflabel"
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
@@ -52,18 +44,8 @@ const (
|
||||
// Chat debug run deletions can cascade into steps with large JSONB
|
||||
// payloads, so they use the same conservative batch size.
|
||||
chatDebugRunsBatchSize = 1000
|
||||
// chatAutoArchiveDigestMaxChats bounds how many chat titles a
|
||||
// single digest body lists. Past the cap, surplus titles are
|
||||
// summarized as "...and N more". 25 is a readable email-friendly
|
||||
// length; the cap is unrelated to chatAutoArchiveBatchSize, which
|
||||
// bounds work per tick.
|
||||
chatAutoArchiveDigestMaxChats = 25
|
||||
)
|
||||
|
||||
// defaultChatAutoArchiveBatchSize bounds how many root chats one
|
||||
// tick will archive by default.
|
||||
const defaultChatAutoArchiveBatchSize int32 = 1000
|
||||
|
||||
type Option func(*instance)
|
||||
|
||||
// WithClock overrides the clock used by the purger. Defaults to
|
||||
@@ -72,34 +54,12 @@ func WithClock(clk quartz.Clock) Option {
|
||||
return func(i *instance) { i.clk = clk }
|
||||
}
|
||||
|
||||
// WithChatAutoArchiveBatchSize overrides how many root chats a
|
||||
// single tick will auto-archive. Defaults to
|
||||
// defaultChatAutoArchiveBatchSize (1000).
|
||||
func WithChatAutoArchiveBatchSize(n int32) Option {
|
||||
return func(i *instance) { i.chatAutoArchiveBatchSize = n }
|
||||
}
|
||||
|
||||
// WithNotificationsEnqueuer sets the enqueuer used for digest
|
||||
// notifications. Defaults to notifications.NewNoopEnqueuer(). Panics
|
||||
// if e is nil: a nil enqueuer would NPE on the first dispatch tick,
|
||||
// and failing fast at option-apply time surfaces the misuse at
|
||||
// startup rather than minutes later.
|
||||
func WithNotificationsEnqueuer(e notifications.Enqueuer) Option {
|
||||
if e == nil {
|
||||
panic("developer error: WithNotificationsEnqueuer called with nil enqueuer")
|
||||
}
|
||||
return func(i *instance) { i.enqueuer = e }
|
||||
}
|
||||
|
||||
// New creates a new periodically purging database instance.
|
||||
// Callers must Close the returned instance.
|
||||
//
|
||||
// The auditor pointer is loaded on each dispatch tick so runtime
|
||||
// entitlement changes (e.g. toggling the audit-log feature) take
|
||||
// effect without restarting the process. Notifications enqueuer
|
||||
// defaults to no-op. Use WithNotificationsEnqueuer to pass a real
|
||||
// one.
|
||||
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, auditor *atomic.Pointer[audit.Auditor], opts ...Option) io.Closer {
|
||||
// The auditor pointer is accepted for compatibility with other background
|
||||
// services. Dbpurge does not emit audit logs directly.
|
||||
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, _ *atomic.Pointer[audit.Auditor], opts ...Option) io.Closer {
|
||||
closed := make(chan struct{})
|
||||
|
||||
ctx, cancelFunc := context.WithCancel(ctx)
|
||||
@@ -123,26 +83,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
|
||||
}, []string{"record_type"})
|
||||
reg.MustRegister(recordsPurged)
|
||||
|
||||
chatAutoArchiveRecords := prometheus.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: "coderd",
|
||||
Subsystem: "chat_auto_archive",
|
||||
Name: "records_archived_total",
|
||||
Help: "Total number of chats archived by the auto-archive job (counting both roots and cascaded children).",
|
||||
})
|
||||
reg.MustRegister(chatAutoArchiveRecords)
|
||||
|
||||
inst := &instance{
|
||||
cancel: cancelFunc,
|
||||
closed: closed,
|
||||
logger: logger,
|
||||
vals: vals,
|
||||
clk: quartz.NewReal(),
|
||||
auditor: auditor,
|
||||
enqueuer: notifications.NewNoopEnqueuer(),
|
||||
iterationDuration: iterationDuration,
|
||||
recordsPurged: recordsPurged,
|
||||
chatAutoArchiveRecords: chatAutoArchiveRecords,
|
||||
chatAutoArchiveBatchSize: defaultChatAutoArchiveBatchSize,
|
||||
cancel: cancelFunc,
|
||||
closed: closed,
|
||||
logger: logger,
|
||||
vals: vals,
|
||||
clk: quartz.NewReal(),
|
||||
iterationDuration: iterationDuration,
|
||||
recordsPurged: recordsPurged,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(inst)
|
||||
@@ -185,19 +133,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
|
||||
func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.Time) error {
|
||||
// Read chat configs outside the tx so a corrupt value can't
|
||||
// poison subsequent queries. On config read errors, log and stash
|
||||
// the error, then run unrelated purges best-effort. Retention and
|
||||
// auto-archive errors skip only the conversation purge and
|
||||
// auto-archive work. Debug retention errors skip only the debug
|
||||
// purge. purgeTick returns chatConfigErr after the tx so the failed
|
||||
// iteration is operator-visible via metric and logs.
|
||||
// the error, then run unrelated purges best-effort. Retention
|
||||
// errors skip only the conversation purge. Debug retention errors
|
||||
// skip only the debug purge. purgeTick returns chatConfigErr after
|
||||
// the tx so the failed iteration is operator-visible via metric and
|
||||
// logs.
|
||||
chatRetentionDays, chatRetentionErr := db.GetChatRetentionDays(ctx)
|
||||
if chatRetentionErr != nil {
|
||||
i.logger.Error(ctx, "failed to read chat retention config: skipping chat purge and auto-archive this tick", slog.Error(chatRetentionErr))
|
||||
}
|
||||
|
||||
chatAutoArchiveDays, chatAutoArchiveErr := db.GetChatAutoArchiveDays(ctx, codersdk.DefaultChatAutoArchiveDays)
|
||||
if chatAutoArchiveErr != nil {
|
||||
i.logger.Error(ctx, "failed to read chat auto-archive config: skipping chat purge and auto-archive this tick", slog.Error(chatAutoArchiveErr))
|
||||
i.logger.Error(ctx, "failed to read chat retention config: skipping chat purge this tick", slog.Error(chatRetentionErr))
|
||||
}
|
||||
|
||||
chatDebugRetentionDays, chatDebugRetentionErr := db.GetChatDebugRetentionDays(ctx, codersdk.DefaultChatDebugRetentionDays)
|
||||
@@ -205,11 +148,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
i.logger.Error(ctx, "failed to read chat debug retention config: skipping chat debug purge this tick", slog.Error(chatDebugRetentionErr))
|
||||
}
|
||||
|
||||
chatRetentionConfigErr := errors.Join(chatRetentionErr, chatAutoArchiveErr)
|
||||
chatConfigErr := errors.Join(chatRetentionConfigErr, chatDebugRetentionErr)
|
||||
|
||||
// Populated inside the tx; dispatched post-commit.
|
||||
var archivedChats []database.AutoArchiveInactiveChatsRow
|
||||
chatConfigErr := errors.Join(chatRetentionErr, chatDebugRetentionErr)
|
||||
|
||||
// Start a transaction to grab advisory lock, we don't want to run
|
||||
// multiple purges at the same time (multiple replicas).
|
||||
@@ -316,8 +255,8 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
}
|
||||
|
||||
var purgedChats, purgedChatFiles, purgedChatDebugRuns int64
|
||||
if chatRetentionConfigErr == nil {
|
||||
purgedChats, purgedChatFiles, archivedChats, err = i.purgeChatsInTx(ctx, tx, start, chatRetentionDays, chatAutoArchiveDays)
|
||||
if chatRetentionErr == nil {
|
||||
purgedChats, purgedChatFiles, err = i.purgeChatsInTx(ctx, tx, start, chatRetentionDays)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("failed to purge chats: %w", err)
|
||||
}
|
||||
@@ -345,7 +284,6 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
slog.F("chats", purgedChats),
|
||||
slog.F("chat_files", purgedChatFiles),
|
||||
slog.F("chat_debug_runs", purgedChatDebugRuns),
|
||||
slog.F("auto_archived_chats", len(archivedChats)),
|
||||
slog.F("duration", i.clk.Since(start)),
|
||||
)
|
||||
|
||||
@@ -379,35 +317,17 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
return xerrors.Errorf("chat config read failed this tick: %w", chatConfigErr)
|
||||
}
|
||||
|
||||
// Dispatch audits and digests post-commit. Detached context for audit
|
||||
// so that ticker cancellation cannot truncate the audit trail.
|
||||
// Notification enqueue uses the cancellable parent context to avoid
|
||||
// stalling shutdown.
|
||||
// Owners with more eligible chats than batch size will get a
|
||||
// notification per tick until their backlog drains.
|
||||
// If this is deemed too noisy, users can disable the
|
||||
// "Chats Auto-Archived" template from their notification preferences.
|
||||
if len(archivedChats) > 0 {
|
||||
i.chatAutoArchiveRecords.Add(float64(len(archivedChats)))
|
||||
auditCtx := context.WithoutCancel(ctx)
|
||||
i.dispatchChatAutoArchive(auditCtx, ctx, start, chatAutoArchiveDays, chatRetentionDays, archivedChats)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type instance struct {
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
logger slog.Logger
|
||||
vals *codersdk.DeploymentValues
|
||||
clk quartz.Clock
|
||||
auditor *atomic.Pointer[audit.Auditor]
|
||||
enqueuer notifications.Enqueuer
|
||||
iterationDuration *prometheus.HistogramVec
|
||||
recordsPurged *prometheus.CounterVec
|
||||
chatAutoArchiveRecords prometheus.Counter
|
||||
chatAutoArchiveBatchSize int32
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
logger slog.Logger
|
||||
vals *codersdk.DeploymentValues
|
||||
clk quartz.Clock
|
||||
iterationDuration *prometheus.HistogramVec
|
||||
recordsPurged *prometheus.CounterVec
|
||||
}
|
||||
|
||||
func (i *instance) Close() error {
|
||||
@@ -416,73 +336,8 @@ func (i *instance) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// chatFromAutoArchiveRow reshapes the query row into a database.Chat for
|
||||
// audit.Auditable[database.Chat].
|
||||
func chatFromAutoArchiveRow(logger slog.Logger, r database.AutoArchiveInactiveChatsRow) database.Chat {
|
||||
var labels database.StringMap
|
||||
// sqlc's StringMap override doesn't reach CTE-aliased columns, so Labels
|
||||
// arrives as raw JSON bytes. StringMap.Scan handles []byte and nil.
|
||||
if err := labels.Scan([]byte(r.Labels)); err != nil {
|
||||
logger.Warn(context.Background(), "failed to parse chat labels from auto-archive row",
|
||||
slog.F("chat_id", r.ID),
|
||||
slog.F("raw_labels", string(r.Labels)),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
|
||||
var userACL database.ChatACL
|
||||
if err := userACL.Scan([]byte(r.UserACL)); err != nil {
|
||||
logger.Warn(context.Background(), "failed to parse chat user ACL from auto-archive row",
|
||||
slog.F("chat_id", r.ID),
|
||||
slog.F("raw_user_acl", string(r.UserACL)),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
|
||||
var groupACL database.ChatACL
|
||||
if err := groupACL.Scan([]byte(r.GroupACL)); err != nil {
|
||||
logger.Warn(context.Background(), "failed to parse chat group ACL from auto-archive row",
|
||||
slog.F("chat_id", r.ID),
|
||||
slog.F("raw_group_acl", string(r.GroupACL)),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
|
||||
return database.Chat{
|
||||
ID: r.ID,
|
||||
OwnerID: r.OwnerID,
|
||||
OrganizationID: r.OrganizationID,
|
||||
WorkspaceID: r.WorkspaceID,
|
||||
BuildID: r.BuildID,
|
||||
AgentID: r.AgentID,
|
||||
Title: r.Title,
|
||||
Status: r.Status,
|
||||
WorkerID: r.WorkerID,
|
||||
StartedAt: r.StartedAt,
|
||||
HeartbeatAt: r.HeartbeatAt,
|
||||
CreatedAt: r.CreatedAt,
|
||||
UpdatedAt: r.UpdatedAt,
|
||||
ParentChatID: r.ParentChatID,
|
||||
RootChatID: r.RootChatID,
|
||||
LastModelConfigID: r.LastModelConfigID,
|
||||
Archived: r.Archived,
|
||||
LastError: r.LastError,
|
||||
Mode: r.Mode,
|
||||
MCPServerIDs: r.MCPServerIDs,
|
||||
Labels: labels,
|
||||
UserACL: userACL,
|
||||
GroupACL: groupACL,
|
||||
PinOrder: r.PinOrder,
|
||||
LastReadMessageID: r.LastReadMessageID,
|
||||
LastInjectedContext: r.LastInjectedContext,
|
||||
DynamicTools: r.DynamicTools,
|
||||
PlanMode: r.PlanMode,
|
||||
ClientType: r.ClientType,
|
||||
}
|
||||
}
|
||||
|
||||
// purgeChatsInTx MUST BE CALLED WITH A TRANSACTION
|
||||
func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start time.Time, chatRetentionDays, chatAutoArchiveDays int32) (purgedChats, purgedChatFiles int64, archivedChats []database.AutoArchiveInactiveChatsRow, err error) {
|
||||
func (*instance) purgeChatsInTx(ctx context.Context, tx database.Store, start time.Time, chatRetentionDays int32) (purgedChats, purgedChatFiles int64, err error) {
|
||||
// Delete old archived chats first, then orphaned files
|
||||
// (cascade clears chat_file_links but not chat_files).
|
||||
if chatRetentionDays > 0 {
|
||||
@@ -492,7 +347,7 @@ func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start
|
||||
LimitCount: chatsBatchSize,
|
||||
})
|
||||
if err != nil {
|
||||
return 0, 0, nil, xerrors.Errorf("failed to delete old chats: %w", err)
|
||||
return 0, 0, xerrors.Errorf("failed to delete old chats: %w", err)
|
||||
}
|
||||
|
||||
purgedChatFiles, err = tx.DeleteOldChatFiles(ctx, database.DeleteOldChatFilesParams{
|
||||
@@ -500,149 +355,9 @@ func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start
|
||||
LimitCount: chatFilesBatchSize,
|
||||
})
|
||||
if err != nil {
|
||||
return 0, 0, nil, xerrors.Errorf("failed to delete old chat files: %w", err)
|
||||
return 0, 0, xerrors.Errorf("failed to delete old chat files: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Auto-archive runs after the delete pass so newly
|
||||
// archived chats aren't eligible for deletion this tick.
|
||||
// Eligibility uses UTC day boundaries: a chat is archived on the
|
||||
// start of the UTC day after its inactivity period has elapsed.
|
||||
if chatAutoArchiveDays > 0 {
|
||||
today := dbtime.StartOfDay(start)
|
||||
archiveCutoff := today.Add(-time.Duration(chatAutoArchiveDays) * 24 * time.Hour)
|
||||
archivedChats, err = tx.AutoArchiveInactiveChats(ctx, database.AutoArchiveInactiveChatsParams{
|
||||
ArchiveCutoff: archiveCutoff,
|
||||
LimitCount: i.chatAutoArchiveBatchSize,
|
||||
})
|
||||
if err != nil {
|
||||
return 0, 0, nil, xerrors.Errorf("failed to auto-archive inactive chats: %w", err)
|
||||
}
|
||||
}
|
||||
return purgedChats, purgedChatFiles, archivedChats, nil
|
||||
}
|
||||
|
||||
// dispatchChatAutoArchive audits every archived root chat and enqueues one
|
||||
// notification per owner covering the roots archived in this tick. Children
|
||||
// inherit their root's archival decision and are skipped for audit, matching
|
||||
// the manual archive path (patchChat audits the root only). Enqueue is
|
||||
// per-tick: owners whose backlog spans multiple ticks receive multiple
|
||||
// notifications; notification_messages dedupe does not collapse them because
|
||||
// each tick's payload differs.
|
||||
//
|
||||
// auditCtx is detached from the ticker so audits always complete. enqueueCtx
|
||||
// is the cancellable parent: on shutdown we abandon any remaining digests
|
||||
// rather than blocking Close.
|
||||
func (i *instance) dispatchChatAutoArchive(auditCtx, enqueueCtx context.Context, tickStart time.Time, autoArchiveDays, retentionDays int32, archived []database.AutoArchiveInactiveChatsRow) {
|
||||
// Children inherit their root's archival decision and are skipped
|
||||
// for both audit and digest. Partition once so the two loops
|
||||
// cannot drift apart if the cascade shape ever changes.
|
||||
roots := slice.Filter(archived, func(r database.AutoArchiveInactiveChatsRow) bool {
|
||||
return !r.ParentChatID.Valid
|
||||
})
|
||||
|
||||
auditor := *i.auditor.Load()
|
||||
for _, row := range roots {
|
||||
after := chatFromAutoArchiveRow(i.logger, row)
|
||||
before := after
|
||||
before.Archived = false
|
||||
audit.BackgroundAudit(auditCtx, &audit.BackgroundAuditParams[database.Chat]{
|
||||
Audit: auditor,
|
||||
Log: i.logger,
|
||||
UserID: row.OwnerID,
|
||||
OrganizationID: row.OrganizationID,
|
||||
Action: database.AuditActionWrite,
|
||||
Old: before,
|
||||
New: after,
|
||||
Status: http.StatusOK,
|
||||
AdditionalFields: audit.BackgroundTaskFieldsBytes(auditCtx, i.logger, audit.BackgroundSubsystemChatAutoArchive),
|
||||
})
|
||||
}
|
||||
|
||||
// Group archived roots by owner. Inline because this is the
|
||||
// only call site and the loop body is self-explanatory.
|
||||
rootsByOwner := make(map[uuid.UUID][]database.AutoArchiveInactiveChatsRow, len(roots))
|
||||
for _, row := range roots {
|
||||
rootsByOwner[row.OwnerID] = append(rootsByOwner[row.OwnerID], row)
|
||||
}
|
||||
|
||||
// Sort owner IDs so shutdown abandons a deterministic tail of the dispatch list.
|
||||
ownerIDs := make([]uuid.UUID, 0, len(rootsByOwner))
|
||||
for id := range rootsByOwner {
|
||||
ownerIDs = append(ownerIDs, id)
|
||||
}
|
||||
slices.SortFunc(ownerIDs, func(a, b uuid.UUID) int {
|
||||
return cmp.Compare(a.String(), b.String())
|
||||
})
|
||||
|
||||
dispatched := 0
|
||||
for _, ownerID := range ownerIDs {
|
||||
// Check between iterations so shutdown unblocks promptly. A
|
||||
// hung in-flight enqueue is unblocked by enqueueCtx propagating
|
||||
// cancellation into the DB call. Skipped owners are not
|
||||
// re-notified on the next tick because AutoArchiveInactiveChats
|
||||
// only returns rows with archived = false; we accept that
|
||||
// tradeoff over hanging shutdown.
|
||||
if err := enqueueCtx.Err(); err != nil {
|
||||
i.logger.Warn(enqueueCtx, "chat auto-archive digest dispatch canceled",
|
||||
slog.F("remaining_owners", len(ownerIDs)-dispatched),
|
||||
slog.Error(err))
|
||||
return
|
||||
}
|
||||
dispatched++
|
||||
|
||||
ownerRoots := rootsByOwner[ownerID]
|
||||
data := buildDigestData(ownerRoots, autoArchiveDays, retentionDays, tickStart)
|
||||
|
||||
// nolint:gocritic // Background digest runs as the notifier subject.
|
||||
if _, err := i.enqueuer.EnqueueWithData(
|
||||
dbauthz.AsNotifier(enqueueCtx),
|
||||
ownerID,
|
||||
notifications.TemplateChatAutoArchiveDigest,
|
||||
map[string]string{},
|
||||
data,
|
||||
string(audit.BackgroundSubsystemChatAutoArchive),
|
||||
); err != nil {
|
||||
i.logger.Warn(enqueueCtx, "failed to enqueue chat auto-archive digest",
|
||||
slog.F("owner_id", ownerID),
|
||||
slog.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// buildDigestData builds the notification payload; shape mirrors the
|
||||
// golden fixtures in coderd/notifications/testdata. Truncation keeps
|
||||
// the oldest archived roots (created_at ASC from the query) to
|
||||
// preserve index-driven ordering; revisit if the digest becomes the
|
||||
// primary surface for reviewing archived chats.
|
||||
func buildDigestData(rows []database.AutoArchiveInactiveChatsRow, autoArchiveDays, retentionDays int32, tickStart time.Time) map[string]any {
|
||||
// Cap titles; overflow surfaces as "...and N more" via the template.
|
||||
overflow := 0
|
||||
if len(rows) > chatAutoArchiveDigestMaxChats {
|
||||
overflow = len(rows) - chatAutoArchiveDigestMaxChats
|
||||
rows = rows[:chatAutoArchiveDigestMaxChats]
|
||||
}
|
||||
|
||||
chats := make([]map[string]any, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
chats = append(chats, map[string]any{
|
||||
"title": r.Title,
|
||||
"last_activity_humanized": humanize.RelTime(r.LastActivityAt, tickStart, "ago", "from now"),
|
||||
})
|
||||
}
|
||||
|
||||
// Stringify the int32 config values: the template's
|
||||
// {{if eq .Data.retention_days "0"}} branch requires both
|
||||
// operands to share a type, and Go templates do not coerce
|
||||
// numeric ↔ string. Storing a raw int here would silently
|
||||
// take the deletion-warning branch on every notification.
|
||||
data := map[string]any{
|
||||
"auto_archive_days": strconv.Itoa(int(autoArchiveDays)),
|
||||
"retention_days": strconv.Itoa(int(retentionDays)),
|
||||
"archived_chats": chats,
|
||||
}
|
||||
if overflow > 0 {
|
||||
data["additional_archived_count"] = strconv.Itoa(overflow)
|
||||
}
|
||||
return data
|
||||
return purgedChats, purgedChatFiles, nil
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Generated
+190
-3
@@ -341,7 +341,8 @@ CREATE TYPE chat_status AS ENUM (
|
||||
'paused',
|
||||
'completed',
|
||||
'error',
|
||||
'requires_action'
|
||||
'requires_action',
|
||||
'interrupting'
|
||||
);
|
||||
|
||||
CREATE TYPE connection_status AS ENUM (
|
||||
@@ -716,6 +717,29 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION bump_chat_queue_version_on_queued_message_change() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
changed_chat_id uuid;
|
||||
BEGIN
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
changed_chat_id = OLD.chat_id;
|
||||
ELSE
|
||||
changed_chat_id = NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
UPDATE chats
|
||||
SET queue_version = snapshot_version
|
||||
WHERE id = changed_chat_id;
|
||||
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
RETURN OLD;
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION check_workspace_agent_name_unique() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
@@ -1293,6 +1317,103 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION set_chat_message_revision_before() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
chat_snapshot_version bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'UPDATE' THEN
|
||||
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
|
||||
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
|
||||
END IF;
|
||||
|
||||
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF OLD IS NOT DISTINCT FROM NEW THEN
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
SELECT snapshot_version INTO chat_snapshot_version
|
||||
FROM chats WHERE id = NEW.chat_id;
|
||||
|
||||
IF chat_snapshot_version IS NULL THEN
|
||||
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
NEW.revision = chat_snapshot_version;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION sync_chat_retry_state() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
IF OLD.retry_state_version IS DISTINCT FROM NEW.retry_state_version THEN
|
||||
RAISE EXCEPTION 'chats.retry_state_version must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF NEW.generation_attempt IS DISTINCT FROM OLD.generation_attempt THEN
|
||||
NEW.retry_state = NULL;
|
||||
END IF;
|
||||
|
||||
IF NEW.retry_state IS DISTINCT FROM OLD.retry_state THEN
|
||||
NEW.retry_state_version = NEW.snapshot_version;
|
||||
END IF;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION update_chat_history_after_message_insert() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT chat_id FROM chat_message_history_new_rows
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION update_chat_history_after_message_update() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT n.chat_id
|
||||
FROM chat_message_history_new_rows n
|
||||
JOIN chat_message_history_old_rows o ON o.id = n.id
|
||||
WHERE o IS DISTINCT FROM n
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TABLE ai_gateway_keys (
|
||||
id uuid NOT NULL,
|
||||
created_at timestamp with time zone NOT NULL,
|
||||
@@ -1670,6 +1791,14 @@ CREATE TABLE chat_files (
|
||||
data bytea NOT NULL
|
||||
);
|
||||
|
||||
CREATE UNLOGGED TABLE chat_heartbeats (
|
||||
chat_id uuid NOT NULL,
|
||||
runner_id uuid NOT NULL,
|
||||
heartbeat_at timestamp with time zone NOT NULL
|
||||
);
|
||||
|
||||
COMMENT ON TABLE chat_heartbeats IS 'Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats.';
|
||||
|
||||
CREATE TABLE chat_messages (
|
||||
id bigint NOT NULL,
|
||||
chat_id uuid NOT NULL,
|
||||
@@ -1692,7 +1821,8 @@ CREATE TABLE chat_messages (
|
||||
runtime_ms bigint,
|
||||
deleted boolean DEFAULT false NOT NULL,
|
||||
provider_response_id text,
|
||||
api_key_id text
|
||||
api_key_id text,
|
||||
revision bigint NOT NULL
|
||||
);
|
||||
|
||||
CREATE SEQUENCE chat_messages_id_seq
|
||||
@@ -1726,13 +1856,22 @@ CREATE TABLE chat_model_configs (
|
||||
CONSTRAINT chat_model_configs_context_limit_check CHECK ((context_limit > 0))
|
||||
);
|
||||
|
||||
CREATE SEQUENCE chat_queued_messages_position_seq
|
||||
START WITH 1
|
||||
INCREMENT BY 1
|
||||
NO MINVALUE
|
||||
NO MAXVALUE
|
||||
CACHE 1;
|
||||
|
||||
CREATE TABLE chat_queued_messages (
|
||||
id bigint NOT NULL,
|
||||
chat_id uuid NOT NULL,
|
||||
content jsonb NOT NULL,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
model_config_id uuid,
|
||||
api_key_id text
|
||||
api_key_id text,
|
||||
"position" bigint DEFAULT nextval('chat_queued_messages_position_seq'::regclass) NOT NULL,
|
||||
created_by uuid NOT NULL
|
||||
);
|
||||
|
||||
CREATE SEQUENCE chat_queued_messages_id_seq
|
||||
@@ -1797,6 +1936,14 @@ CREATE TABLE chats (
|
||||
last_turn_summary text,
|
||||
user_acl jsonb DEFAULT '{}'::jsonb NOT NULL,
|
||||
group_acl jsonb DEFAULT '{}'::jsonb NOT NULL,
|
||||
snapshot_version bigint DEFAULT 1 NOT NULL,
|
||||
history_version bigint DEFAULT 0 NOT NULL,
|
||||
queue_version bigint DEFAULT 0 NOT NULL,
|
||||
generation_attempt bigint DEFAULT 0 NOT NULL,
|
||||
retry_state jsonb,
|
||||
retry_state_version bigint DEFAULT 0 NOT NULL,
|
||||
runner_id uuid,
|
||||
requires_action_deadline_at timestamp with time zone,
|
||||
CONSTRAINT chat_acl_only_on_root_chats CHECK ((((parent_chat_id IS NULL) AND (root_chat_id IS NULL)) OR ((user_acl = '{}'::jsonb) AND (group_acl = '{}'::jsonb)))),
|
||||
CONSTRAINT chat_group_acl_not_null_jsonb CHECK (((group_acl IS NOT NULL) AND (jsonb_typeof(group_acl) = 'object'::text))),
|
||||
CONSTRAINT chat_user_acl_not_null_jsonb CHECK (((user_acl IS NOT NULL) AND (jsonb_typeof(user_acl) = 'object'::text))),
|
||||
@@ -1804,6 +1951,12 @@ CREATE TABLE chats (
|
||||
CONSTRAINT chats_pin_order_parent_check CHECK (((pin_order = 0) OR (parent_chat_id IS NULL)))
|
||||
);
|
||||
|
||||
COMMENT ON COLUMN chats.snapshot_version IS 'Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet.';
|
||||
|
||||
COMMENT ON COLUMN chats.history_version IS 'Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version.';
|
||||
|
||||
COMMENT ON COLUMN chats.queue_version IS 'Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version.';
|
||||
|
||||
CREATE TABLE users (
|
||||
id uuid NOT NULL,
|
||||
email text NOT NULL,
|
||||
@@ -1884,6 +2037,14 @@ CREATE VIEW chats_expanded AS
|
||||
c.plan_mode,
|
||||
c.client_type,
|
||||
c.last_turn_summary,
|
||||
c.snapshot_version,
|
||||
c.history_version,
|
||||
c.queue_version,
|
||||
c.generation_attempt,
|
||||
c.retry_state,
|
||||
c.retry_state_version,
|
||||
c.runner_id,
|
||||
c.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, c.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, c.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -3848,6 +4009,9 @@ ALTER TABLE ONLY chat_file_links
|
||||
ALTER TABLE ONLY chat_files
|
||||
ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id);
|
||||
|
||||
ALTER TABLE ONLY chat_heartbeats
|
||||
ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id);
|
||||
|
||||
ALTER TABLE ONLY chat_messages
|
||||
ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id);
|
||||
|
||||
@@ -4190,6 +4354,8 @@ CREATE INDEX api_keys_last_used_idx ON api_keys USING btree (last_used DESC);
|
||||
|
||||
COMMENT ON INDEX api_keys_last_used_idx IS 'Index for optimizing api_keys queries filtering by last_used';
|
||||
|
||||
CREATE INDEX chat_heartbeats_heartbeat_at_idx ON chat_heartbeats USING btree (heartbeat_at);
|
||||
|
||||
CREATE INDEX idx_agent_stats_created_at ON workspace_agent_stats USING btree (created_at);
|
||||
|
||||
CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_id);
|
||||
@@ -4320,6 +4486,8 @@ CREATE INDEX idx_chats_pending ON chats USING btree (status) WHERE (status = 'pe
|
||||
|
||||
CREATE INDEX idx_chats_root_chat_id ON chats USING btree (root_chat_id);
|
||||
|
||||
CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (status, updated_at, id) WHERE (archived = false);
|
||||
|
||||
CREATE INDEX idx_chats_workspace ON chats USING btree (workspace_id);
|
||||
|
||||
CREATE INDEX idx_connection_logs_connect_time_desc ON connection_logs USING btree (connect_time DESC);
|
||||
@@ -4550,6 +4718,12 @@ COMMENT ON TRIGGER remove_organization_member_custom_role ON custom_roles IS 'Wh
|
||||
|
||||
CREATE TRIGGER trigger_aggregate_usage_event AFTER INSERT ON usage_events FOR EACH ROW EXECUTE FUNCTION aggregate_usage_event();
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_delete AFTER DELETE ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_insert AFTER INSERT ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_update AFTER UPDATE OF content, model_config_id, "position", created_by ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
CREATE TRIGGER trigger_delete_group_members_on_org_member_delete BEFORE DELETE ON organization_members FOR EACH ROW EXECUTE FUNCTION delete_group_members_on_org_member_delete();
|
||||
|
||||
CREATE TRIGGER trigger_delete_oauth2_provider_app_token AFTER DELETE ON oauth2_provider_app_tokens FOR EACH ROW EXECUTE FUNCTION delete_deleted_oauth2_provider_app_token_api_key();
|
||||
@@ -4566,6 +4740,16 @@ CREATE TRIGGER trigger_insert_organization_system_roles AFTER INSERT ON organiza
|
||||
|
||||
CREATE TRIGGER trigger_nullify_next_start_at_on_workspace_autostart_modificati AFTER UPDATE ON workspaces FOR EACH ROW EXECUTE FUNCTION nullify_next_start_at_on_workspace_autostart_modification();
|
||||
|
||||
CREATE TRIGGER trigger_set_chat_message_revision_on_insert BEFORE INSERT ON chat_messages FOR EACH ROW EXECUTE FUNCTION set_chat_message_revision_before();
|
||||
|
||||
CREATE TRIGGER trigger_set_chat_message_revision_on_update BEFORE UPDATE ON chat_messages FOR EACH ROW EXECUTE FUNCTION set_chat_message_revision_before();
|
||||
|
||||
CREATE TRIGGER trigger_sync_chat_retry_state BEFORE UPDATE OF retry_state, retry_state_version, generation_attempt ON chats FOR EACH ROW EXECUTE FUNCTION sync_chat_retry_state();
|
||||
|
||||
CREATE TRIGGER trigger_update_chat_history_after_message_insert AFTER INSERT ON chat_messages REFERENCING NEW TABLE AS chat_message_history_new_rows FOR EACH STATEMENT EXECUTE FUNCTION update_chat_history_after_message_insert();
|
||||
|
||||
CREATE TRIGGER trigger_update_chat_history_after_message_update AFTER UPDATE ON chat_messages REFERENCING OLD TABLE AS chat_message_history_old_rows NEW TABLE AS chat_message_history_new_rows FOR EACH STATEMENT EXECUTE FUNCTION update_chat_history_after_message_update();
|
||||
|
||||
CREATE TRIGGER trigger_update_users AFTER INSERT OR UPDATE ON users FOR EACH ROW WHEN ((new.deleted = true)) EXECUTE FUNCTION delete_deleted_user_resources();
|
||||
|
||||
CREATE TRIGGER trigger_upsert_user_links BEFORE INSERT OR UPDATE ON user_links FOR EACH ROW EXECUTE FUNCTION insert_user_links_fail_if_user_deleted();
|
||||
@@ -4636,6 +4820,9 @@ ALTER TABLE ONLY chat_files
|
||||
ALTER TABLE ONLY chat_files
|
||||
ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY chat_heartbeats
|
||||
ADD CONSTRAINT chat_heartbeats_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY chat_messages
|
||||
ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
|
||||
|
||||
+1
@@ -22,6 +22,7 @@ const (
|
||||
ForeignKeyChatFileLinksFileID ForeignKeyConstraint = "chat_file_links_file_id_fkey" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_file_id_fkey FOREIGN KEY (file_id) REFERENCES chat_files(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatFilesOrganizationID ForeignKeyConstraint = "chat_files_organization_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatFilesOwnerID ForeignKeyConstraint = "chat_files_owner_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatHeartbeatsChatID ForeignKeyConstraint = "chat_heartbeats_chat_id_fkey" // ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatMessagesAPIKeyID ForeignKeyConstraint = "chat_messages_api_key_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
ForeignKeyChatMessagesChatID ForeignKeyConstraint = "chat_messages_chat_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatMessagesModelConfigID ForeignKeyConstraint = "chat_messages_model_config_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_model_config_id_fkey FOREIGN KEY (model_config_id) REFERENCES chat_model_configs(id);
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
-- Rollback for the chatd core state machine foundation migration.
|
||||
|
||||
-- 1. Recreate chats_expanded without the new chat fields. We must drop
|
||||
-- the view first because the subsequent column drops would fail with
|
||||
-- "view depends on column".
|
||||
DROP VIEW IF EXISTS chats_expanded;
|
||||
|
||||
-- 2. Drop the worker acquisition candidates index.
|
||||
DROP INDEX IF EXISTS idx_chats_worker_acquisition_candidates;
|
||||
|
||||
-- 3. Drop the retry state trigger and function.
|
||||
DROP TRIGGER IF EXISTS trigger_sync_chat_retry_state ON chats;
|
||||
DROP FUNCTION IF EXISTS sync_chat_retry_state();
|
||||
|
||||
-- 4. Drop the queue version triggers and function.
|
||||
DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_delete ON chat_queued_messages;
|
||||
DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_update ON chat_queued_messages;
|
||||
DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_insert ON chat_queued_messages;
|
||||
DROP FUNCTION IF EXISTS bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
-- 5. Drop the message revision triggers and functions.
|
||||
DROP TRIGGER IF EXISTS trigger_update_chat_history_after_message_update ON chat_messages;
|
||||
DROP TRIGGER IF EXISTS trigger_update_chat_history_after_message_insert ON chat_messages;
|
||||
DROP TRIGGER IF EXISTS trigger_set_chat_message_revision_on_update ON chat_messages;
|
||||
DROP TRIGGER IF EXISTS trigger_set_chat_message_revision_on_insert ON chat_messages;
|
||||
DROP FUNCTION IF EXISTS update_chat_history_after_message_update();
|
||||
DROP FUNCTION IF EXISTS update_chat_history_after_message_insert();
|
||||
-- The pre-split function name is kept here for backward compatibility
|
||||
-- with environments that may have applied an earlier draft of the up
|
||||
-- migration. DROP FUNCTION IF EXISTS is a no-op if the function is
|
||||
-- absent.
|
||||
DROP FUNCTION IF EXISTS update_chat_history_after_message_changes();
|
||||
DROP FUNCTION IF EXISTS set_chat_message_revision_before();
|
||||
DROP FUNCTION IF EXISTS set_chat_message_revision();
|
||||
|
||||
-- 6. Drop chat_heartbeats (and its index by association).
|
||||
DROP TABLE IF EXISTS chat_heartbeats;
|
||||
|
||||
-- 7. Drop chat_queued_messages.position and its default sequence, plus
|
||||
-- created_by.
|
||||
ALTER TABLE chat_queued_messages
|
||||
ALTER COLUMN position DROP DEFAULT;
|
||||
ALTER TABLE chat_queued_messages
|
||||
DROP COLUMN IF EXISTS position,
|
||||
DROP COLUMN IF EXISTS created_by;
|
||||
DROP SEQUENCE IF EXISTS chat_queued_messages_position_seq;
|
||||
|
||||
-- 8. Drop chat_messages.revision.
|
||||
ALTER TABLE chat_messages
|
||||
DROP COLUMN IF EXISTS revision;
|
||||
|
||||
-- 9. Drop the new chats columns.
|
||||
ALTER TABLE chats
|
||||
DROP COLUMN IF EXISTS snapshot_version,
|
||||
DROP COLUMN IF EXISTS history_version,
|
||||
DROP COLUMN IF EXISTS queue_version,
|
||||
DROP COLUMN IF EXISTS generation_attempt,
|
||||
DROP COLUMN IF EXISTS retry_state,
|
||||
DROP COLUMN IF EXISTS retry_state_version,
|
||||
DROP COLUMN IF EXISTS runner_id,
|
||||
DROP COLUMN IF EXISTS requires_action_deadline_at;
|
||||
|
||||
-- 10. Recreate chats_expanded with the pre-migration field list.
|
||||
CREATE VIEW chats_expanded AS
|
||||
SELECT
|
||||
c.id,
|
||||
c.owner_id,
|
||||
c.workspace_id,
|
||||
c.title,
|
||||
c.status,
|
||||
c.worker_id,
|
||||
c.started_at,
|
||||
c.heartbeat_at,
|
||||
c.created_at,
|
||||
c.updated_at,
|
||||
c.parent_chat_id,
|
||||
c.root_chat_id,
|
||||
c.last_model_config_id,
|
||||
c.archived,
|
||||
c.last_error,
|
||||
c.mode,
|
||||
c.mcp_server_ids,
|
||||
c.labels,
|
||||
c.build_id,
|
||||
c.agent_id,
|
||||
c.pin_order,
|
||||
c.last_read_message_id,
|
||||
c.last_injected_context,
|
||||
c.dynamic_tools,
|
||||
c.organization_id,
|
||||
c.plan_mode,
|
||||
c.client_type,
|
||||
c.last_turn_summary,
|
||||
COALESCE(root.user_acl, c.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, c.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM
|
||||
chats c
|
||||
LEFT JOIN chats root ON root.id = COALESCE(c.root_chat_id, c.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = c.owner_id;
|
||||
|
||||
-- 11. The `interrupting` chat_status enum value is intentionally left
|
||||
-- in place. Postgres does not support dropping a single enum value
|
||||
-- without recreating the entire type, which would require rewriting
|
||||
-- every chat row and is unsafe inside a transactional rollback.
|
||||
@@ -0,0 +1,358 @@
|
||||
-- Adds the core chat state-machine storage model.
|
||||
-- Adds new versioning fields to chats, a revision column to chat_messages,
|
||||
-- positional ordering and creator tracking to chat_queued_messages, an
|
||||
-- unlogged chat_heartbeats table for ownership leases, and Postgres
|
||||
-- triggers that keep history/queue versioning consistent.
|
||||
|
||||
-- 1. Add `interrupting` to the chat_status enum.
|
||||
ALTER TYPE chat_status ADD VALUE IF NOT EXISTS 'interrupting';
|
||||
|
||||
-- 2. Add new versioning, ownership, retry, and pending-action fields to chats.
|
||||
ALTER TABLE chats
|
||||
ADD COLUMN snapshot_version bigint NOT NULL DEFAULT 1,
|
||||
ADD COLUMN history_version bigint NOT NULL DEFAULT 0,
|
||||
ADD COLUMN queue_version bigint NOT NULL DEFAULT 0,
|
||||
ADD COLUMN generation_attempt bigint NOT NULL DEFAULT 0,
|
||||
ADD COLUMN retry_state jsonb,
|
||||
ADD COLUMN retry_state_version bigint NOT NULL DEFAULT 0,
|
||||
ADD COLUMN runner_id uuid,
|
||||
ADD COLUMN requires_action_deadline_at timestamp with time zone;
|
||||
|
||||
COMMENT ON COLUMN chats.snapshot_version IS
|
||||
'Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet.';
|
||||
COMMENT ON COLUMN chats.history_version IS
|
||||
'Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version.';
|
||||
COMMENT ON COLUMN chats.queue_version IS
|
||||
'Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version.';
|
||||
|
||||
-- 3. Add `revision` to chat_messages. Adding the column as NOT NULL with
|
||||
-- a constant default backfills existing rows through catalog metadata
|
||||
-- only, so the highest-volume table is neither rewritten nor scanned for
|
||||
-- NOT NULL validation while under ACCESS EXCLUSIVE. The default is
|
||||
-- dropped immediately because the BEFORE INSERT trigger below rejects
|
||||
-- inserts that pre-assign revision and assigns it from
|
||||
-- chats.snapshot_version instead.
|
||||
ALTER TABLE chat_messages
|
||||
ADD COLUMN revision bigint NOT NULL DEFAULT 1;
|
||||
ALTER TABLE chat_messages
|
||||
ALTER COLUMN revision DROP DEFAULT;
|
||||
|
||||
-- 4. Backfill chats.history_version = 1 for chats that already have at
|
||||
-- least one message. We avoid recursive trigger fire by performing the
|
||||
-- backfill before the triggers are created.
|
||||
UPDATE chats
|
||||
SET history_version = 1
|
||||
WHERE EXISTS (
|
||||
SELECT 1 FROM chat_messages WHERE chat_messages.chat_id = chats.id
|
||||
);
|
||||
|
||||
-- 5. Add `position` and `created_by` to chat_queued_messages.
|
||||
ALTER TABLE chat_queued_messages
|
||||
ADD COLUMN position bigint,
|
||||
ADD COLUMN created_by uuid;
|
||||
|
||||
-- 6. Backfill chat_queued_messages.position per chat using row_number(),
|
||||
-- ordering by created_at and breaking ties by id.
|
||||
WITH ordered AS (
|
||||
SELECT
|
||||
id,
|
||||
row_number() OVER (
|
||||
PARTITION BY chat_id
|
||||
ORDER BY created_at, id
|
||||
) AS rn
|
||||
FROM chat_queued_messages
|
||||
)
|
||||
UPDATE chat_queued_messages
|
||||
SET position = ordered.rn
|
||||
FROM ordered
|
||||
WHERE chat_queued_messages.id = ordered.id;
|
||||
|
||||
-- 7. Backfill chat_queued_messages.created_by from chats.owner_id.
|
||||
UPDATE chat_queued_messages
|
||||
SET created_by = chats.owner_id
|
||||
FROM chats
|
||||
WHERE chat_queued_messages.chat_id = chats.id
|
||||
AND chat_queued_messages.created_by IS NULL;
|
||||
|
||||
-- 8. Enforce NOT NULL on chat_queued_messages.position and
|
||||
-- created_by. Legacy queued-message inserts are updated to populate
|
||||
-- created_by from the chat owner when no explicit creator exists.
|
||||
ALTER TABLE chat_queued_messages
|
||||
ALTER COLUMN position SET NOT NULL,
|
||||
ALTER COLUMN created_by SET NOT NULL;
|
||||
|
||||
-- 9. Default sequence for new queued-message positions.
|
||||
-- A global sequence is acceptable because ordering only needs to be
|
||||
-- stable within a chat.
|
||||
CREATE SEQUENCE IF NOT EXISTS chat_queued_messages_position_seq AS bigint START WITH 1;
|
||||
SELECT setval(
|
||||
'chat_queued_messages_position_seq',
|
||||
GREATEST((SELECT COALESCE(MAX(position), 0) FROM chat_queued_messages), 1)
|
||||
);
|
||||
ALTER TABLE chat_queued_messages
|
||||
ALTER COLUMN position SET DEFAULT nextval('chat_queued_messages_position_seq');
|
||||
|
||||
-- 10. Backfill chats.queue_version = 1 for chats that already have queued
|
||||
-- messages. Same trigger-avoidance reasoning as for history_version.
|
||||
UPDATE chats
|
||||
SET queue_version = 1
|
||||
WHERE EXISTS (
|
||||
SELECT 1 FROM chat_queued_messages WHERE chat_queued_messages.chat_id = chats.id
|
||||
);
|
||||
|
||||
-- 11. chat_heartbeats: unlogged table for ownership leases. Keyed by
|
||||
-- (chat_id, runner_id) so a single chat can briefly have entries from
|
||||
-- multiple runners during failover.
|
||||
CREATE UNLOGGED TABLE IF NOT EXISTS chat_heartbeats (
|
||||
chat_id uuid NOT NULL REFERENCES chats(id) ON DELETE CASCADE,
|
||||
runner_id uuid NOT NULL,
|
||||
heartbeat_at timestamp with time zone NOT NULL,
|
||||
PRIMARY KEY (chat_id, runner_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE chat_heartbeats IS
|
||||
'Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats.';
|
||||
|
||||
CREATE INDEX IF NOT EXISTS chat_heartbeats_heartbeat_at_idx
|
||||
ON chat_heartbeats (heartbeat_at);
|
||||
|
||||
-- 12. Message revision trigger.
|
||||
-- The BEFORE-trigger only assigns NEW.revision from chats.snapshot_version
|
||||
-- and validates immutability. The chats.history_version /
|
||||
-- generation_attempt update is performed by an AFTER STATEMENT trigger
|
||||
-- so it doesn't conflict with CTE updates on the chats row in the same
|
||||
-- command (the legacy InsertChatMessages query updates last_model_config_id
|
||||
-- in a CTE on chats and then inserts messages).
|
||||
CREATE FUNCTION set_chat_message_revision_before()
|
||||
RETURNS trigger AS $$
|
||||
DECLARE
|
||||
chat_snapshot_version bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'UPDATE' THEN
|
||||
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
|
||||
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
|
||||
END IF;
|
||||
|
||||
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF OLD IS NOT DISTINCT FROM NEW THEN
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
SELECT snapshot_version INTO chat_snapshot_version
|
||||
FROM chats WHERE id = NEW.chat_id;
|
||||
|
||||
IF chat_snapshot_version IS NULL THEN
|
||||
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
NEW.revision = chat_snapshot_version;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- AFTER STATEMENT trigger functions. Use the transition tables to
|
||||
-- update chats.history_version / generation_attempt once per chat per
|
||||
-- command. Running AFTER row inserts/updates complete lets a CTE
|
||||
-- update on the same chats row in the same command finalize before
|
||||
-- this trigger needs to update it.
|
||||
--
|
||||
-- The INSERT and UPDATE variants are split so the UPDATE variant can
|
||||
-- reference both the OLD and NEW transition tables and skip rows that
|
||||
-- did not actually change. Without that filter, a no-op UPDATE on a
|
||||
-- chat_messages row (one whose OLD IS NOT DISTINCT FROM NEW) would
|
||||
-- still advance chats.history_version whenever the chat's snapshot
|
||||
-- had previously been bumped.
|
||||
CREATE FUNCTION update_chat_history_after_message_insert()
|
||||
RETURNS trigger AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT chat_id FROM chat_message_history_new_rows
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE FUNCTION update_chat_history_after_message_update()
|
||||
RETURNS trigger AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT n.chat_id
|
||||
FROM chat_message_history_new_rows n
|
||||
JOIN chat_message_history_old_rows o ON o.id = n.id
|
||||
WHERE o IS DISTINCT FROM n
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE TRIGGER trigger_set_chat_message_revision_on_insert
|
||||
BEFORE INSERT ON chat_messages
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION set_chat_message_revision_before();
|
||||
|
||||
CREATE TRIGGER trigger_set_chat_message_revision_on_update
|
||||
BEFORE UPDATE ON chat_messages
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION set_chat_message_revision_before();
|
||||
|
||||
CREATE TRIGGER trigger_update_chat_history_after_message_insert
|
||||
AFTER INSERT ON chat_messages
|
||||
REFERENCING NEW TABLE AS chat_message_history_new_rows
|
||||
FOR EACH STATEMENT
|
||||
EXECUTE FUNCTION update_chat_history_after_message_insert();
|
||||
|
||||
CREATE TRIGGER trigger_update_chat_history_after_message_update
|
||||
AFTER UPDATE ON chat_messages
|
||||
REFERENCING OLD TABLE AS chat_message_history_old_rows NEW TABLE AS chat_message_history_new_rows
|
||||
FOR EACH STATEMENT
|
||||
EXECUTE FUNCTION update_chat_history_after_message_update();
|
||||
|
||||
-- 13. Queue version trigger function.
|
||||
CREATE FUNCTION bump_chat_queue_version_on_queued_message_change()
|
||||
RETURNS trigger AS $$
|
||||
DECLARE
|
||||
changed_chat_id uuid;
|
||||
BEGIN
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
changed_chat_id = OLD.chat_id;
|
||||
ELSE
|
||||
changed_chat_id = NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
UPDATE chats
|
||||
SET queue_version = snapshot_version
|
||||
WHERE id = changed_chat_id;
|
||||
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
RETURN OLD;
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_insert
|
||||
AFTER INSERT ON chat_queued_messages
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_update
|
||||
AFTER UPDATE OF content, model_config_id, position, created_by
|
||||
ON chat_queued_messages
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_delete
|
||||
AFTER DELETE ON chat_queued_messages
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change();
|
||||
|
||||
-- 14. Retry state trigger function.
|
||||
CREATE FUNCTION sync_chat_retry_state()
|
||||
RETURNS trigger AS $$
|
||||
BEGIN
|
||||
IF OLD.retry_state_version IS DISTINCT FROM NEW.retry_state_version THEN
|
||||
RAISE EXCEPTION 'chats.retry_state_version must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF NEW.generation_attempt IS DISTINCT FROM OLD.generation_attempt THEN
|
||||
NEW.retry_state = NULL;
|
||||
END IF;
|
||||
|
||||
IF NEW.retry_state IS DISTINCT FROM OLD.retry_state THEN
|
||||
NEW.retry_state_version = NEW.snapshot_version;
|
||||
END IF;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE TRIGGER trigger_sync_chat_retry_state
|
||||
BEFORE UPDATE OF retry_state, retry_state_version, generation_attempt
|
||||
ON chats
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION sync_chat_retry_state();
|
||||
|
||||
-- 15. Index for the chat worker acquisition scan, which runs every 30
|
||||
-- seconds per replica plus on every worker wake. Leading on status lets
|
||||
-- the scan touch only rows in the worker-runnable status set instead of
|
||||
-- sequentially scanning the ever-growing chats table. The status set is
|
||||
-- intentionally not part of the index predicate: 'interrupting' is added
|
||||
-- to chat_status above, and Postgres forbids using a new enum value in
|
||||
-- the same transaction, which all migrations share.
|
||||
CREATE INDEX idx_chats_worker_acquisition_candidates ON chats
|
||||
USING btree (status, updated_at, id)
|
||||
WHERE archived = false;
|
||||
|
||||
-- 16. Refresh chats_expanded to include the new chat fields. Drop and
|
||||
-- recreate so column ordering is stable.
|
||||
DROP VIEW IF EXISTS chats_expanded;
|
||||
CREATE VIEW chats_expanded AS
|
||||
SELECT
|
||||
c.id,
|
||||
c.owner_id,
|
||||
c.workspace_id,
|
||||
c.title,
|
||||
c.status,
|
||||
c.worker_id,
|
||||
c.started_at,
|
||||
c.heartbeat_at,
|
||||
c.created_at,
|
||||
c.updated_at,
|
||||
c.parent_chat_id,
|
||||
c.root_chat_id,
|
||||
c.last_model_config_id,
|
||||
c.archived,
|
||||
c.last_error,
|
||||
c.mode,
|
||||
c.mcp_server_ids,
|
||||
c.labels,
|
||||
c.build_id,
|
||||
c.agent_id,
|
||||
c.pin_order,
|
||||
c.last_read_message_id,
|
||||
c.last_injected_context,
|
||||
c.dynamic_tools,
|
||||
c.organization_id,
|
||||
c.plan_mode,
|
||||
c.client_type,
|
||||
c.last_turn_summary,
|
||||
c.snapshot_version,
|
||||
c.history_version,
|
||||
c.queue_version,
|
||||
c.generation_attempt,
|
||||
c.retry_state,
|
||||
c.retry_state_version,
|
||||
c.runner_id,
|
||||
c.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, c.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, c.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM
|
||||
chats c
|
||||
LEFT JOIN chats root ON root.id = COALESCE(c.root_chat_id, c.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = c.owner_id;
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
-- Fixture coverage for the chat_heartbeats table introduced in
|
||||
-- migration 000500. The earlier chat fixtures already insert at least
|
||||
-- one row into chats; we attach a heartbeat for the first such chat so
|
||||
-- migration tests see a non-empty chat_heartbeats table without
|
||||
-- hard-coding a specific chat ID.
|
||||
INSERT INTO chat_heartbeats (
|
||||
chat_id,
|
||||
runner_id,
|
||||
heartbeat_at
|
||||
)
|
||||
SELECT
|
||||
chats.id,
|
||||
'00000000-0000-0000-0000-0000000fea51'::uuid,
|
||||
'2024-01-01 00:00:00+00'
|
||||
FROM chats
|
||||
ORDER BY created_at, id
|
||||
LIMIT 1;
|
||||
@@ -824,6 +824,14 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams,
|
||||
&i.Chat.PlanMode,
|
||||
&i.Chat.ClientType,
|
||||
&i.Chat.LastTurnSummary,
|
||||
&i.Chat.SnapshotVersion,
|
||||
&i.Chat.HistoryVersion,
|
||||
&i.Chat.QueueVersion,
|
||||
&i.Chat.GenerationAttempt,
|
||||
&i.Chat.RetryState,
|
||||
&i.Chat.RetryStateVersion,
|
||||
&i.Chat.RunnerID,
|
||||
&i.Chat.RequiresActionDeadlineAt,
|
||||
&i.Chat.UserACL,
|
||||
&i.Chat.GroupACL,
|
||||
&i.Chat.OwnerUsername,
|
||||
@@ -891,6 +899,14 @@ func (q *sqlQuerier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID
|
||||
&i.PlanMode,
|
||||
&i.ClientType,
|
||||
&i.LastTurnSummary,
|
||||
&i.SnapshotVersion,
|
||||
&i.HistoryVersion,
|
||||
&i.QueueVersion,
|
||||
&i.GenerationAttempt,
|
||||
&i.RetryState,
|
||||
&i.RetryStateVersion,
|
||||
&i.RunnerID,
|
||||
&i.RequiresActionDeadlineAt,
|
||||
&i.UserACL,
|
||||
&i.GroupACL,
|
||||
&i.OwnerUsername,
|
||||
|
||||
Generated
+65
-33
@@ -1567,6 +1567,7 @@ const (
|
||||
ChatStatusCompleted ChatStatus = "completed"
|
||||
ChatStatusError ChatStatus = "error"
|
||||
ChatStatusRequiresAction ChatStatus = "requires_action"
|
||||
ChatStatusInterrupting ChatStatus = "interrupting"
|
||||
)
|
||||
|
||||
func (e *ChatStatus) Scan(src interface{}) error {
|
||||
@@ -1612,7 +1613,8 @@ func (e ChatStatus) Valid() bool {
|
||||
ChatStatusPaused,
|
||||
ChatStatusCompleted,
|
||||
ChatStatusError,
|
||||
ChatStatusRequiresAction:
|
||||
ChatStatusRequiresAction,
|
||||
ChatStatusInterrupting:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
@@ -1627,6 +1629,7 @@ func AllChatStatusValues() []ChatStatus {
|
||||
ChatStatusCompleted,
|
||||
ChatStatusError,
|
||||
ChatStatusRequiresAction,
|
||||
ChatStatusInterrupting,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4607,38 +4610,46 @@ type BoundaryUsageStat struct {
|
||||
}
|
||||
|
||||
type Chat struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
|
||||
WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"`
|
||||
Title string `db:"title" json:"title"`
|
||||
Status ChatStatus `db:"status" json:"status"`
|
||||
WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"`
|
||||
StartedAt sql.NullTime `db:"started_at" json:"started_at"`
|
||||
HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
||||
ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"`
|
||||
RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"`
|
||||
LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"`
|
||||
Archived bool `db:"archived" json:"archived"`
|
||||
LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"`
|
||||
Mode NullChatMode `db:"mode" json:"mode"`
|
||||
MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"`
|
||||
Labels StringMap `db:"labels" json:"labels"`
|
||||
BuildID uuid.NullUUID `db:"build_id" json:"build_id"`
|
||||
AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"`
|
||||
PinOrder int32 `db:"pin_order" json:"pin_order"`
|
||||
LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"`
|
||||
LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"`
|
||||
DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"`
|
||||
ClientType ChatClientType `db:"client_type" json:"client_type"`
|
||||
LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"`
|
||||
UserACL ChatACL `db:"user_acl" json:"user_acl"`
|
||||
GroupACL ChatACL `db:"group_acl" json:"group_acl"`
|
||||
OwnerUsername string `db:"owner_username" json:"owner_username"`
|
||||
OwnerName string `db:"owner_name" json:"owner_name"`
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
|
||||
WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"`
|
||||
Title string `db:"title" json:"title"`
|
||||
Status ChatStatus `db:"status" json:"status"`
|
||||
WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"`
|
||||
StartedAt sql.NullTime `db:"started_at" json:"started_at"`
|
||||
HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
||||
ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"`
|
||||
RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"`
|
||||
LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"`
|
||||
Archived bool `db:"archived" json:"archived"`
|
||||
LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"`
|
||||
Mode NullChatMode `db:"mode" json:"mode"`
|
||||
MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"`
|
||||
Labels StringMap `db:"labels" json:"labels"`
|
||||
BuildID uuid.NullUUID `db:"build_id" json:"build_id"`
|
||||
AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"`
|
||||
PinOrder int32 `db:"pin_order" json:"pin_order"`
|
||||
LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"`
|
||||
LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"`
|
||||
DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"`
|
||||
ClientType ChatClientType `db:"client_type" json:"client_type"`
|
||||
LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"`
|
||||
SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"`
|
||||
HistoryVersion int64 `db:"history_version" json:"history_version"`
|
||||
QueueVersion int64 `db:"queue_version" json:"queue_version"`
|
||||
GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"`
|
||||
RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"`
|
||||
RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"`
|
||||
RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"`
|
||||
RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"`
|
||||
UserACL ChatACL `db:"user_acl" json:"user_acl"`
|
||||
GroupACL ChatACL `db:"group_acl" json:"group_acl"`
|
||||
OwnerUsername string `db:"owner_username" json:"owner_username"`
|
||||
OwnerName string `db:"owner_name" json:"owner_name"`
|
||||
}
|
||||
|
||||
type ChatDebugRun struct {
|
||||
@@ -4720,6 +4731,13 @@ type ChatFileLink struct {
|
||||
FileID uuid.UUID `db:"file_id" json:"file_id"`
|
||||
}
|
||||
|
||||
// Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats.
|
||||
type ChatHeartbeat struct {
|
||||
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
|
||||
RunnerID uuid.UUID `db:"runner_id" json:"runner_id"`
|
||||
HeartbeatAt time.Time `db:"heartbeat_at" json:"heartbeat_at"`
|
||||
}
|
||||
|
||||
type ChatMessage struct {
|
||||
ID int64 `db:"id" json:"id"`
|
||||
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
|
||||
@@ -4743,6 +4761,7 @@ type ChatMessage struct {
|
||||
Deleted bool `db:"deleted" json:"deleted"`
|
||||
ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
Revision int64 `db:"revision" json:"revision"`
|
||||
}
|
||||
|
||||
type ChatModelConfig struct {
|
||||
@@ -4771,6 +4790,8 @@ type ChatQueuedMessage struct {
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
Position int64 `db:"position" json:"position"`
|
||||
CreatedBy uuid.UUID `db:"created_by" json:"created_by"`
|
||||
}
|
||||
|
||||
type ChatTable struct {
|
||||
@@ -4804,6 +4825,17 @@ type ChatTable struct {
|
||||
LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"`
|
||||
UserACL ChatACL `db:"user_acl" json:"user_acl"`
|
||||
GroupACL ChatACL `db:"group_acl" json:"group_acl"`
|
||||
// Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet.
|
||||
SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"`
|
||||
// Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version.
|
||||
HistoryVersion int64 `db:"history_version" json:"history_version"`
|
||||
// Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version.
|
||||
QueueVersion int64 `db:"queue_version" json:"queue_version"`
|
||||
GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"`
|
||||
RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"`
|
||||
RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"`
|
||||
RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"`
|
||||
RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"`
|
||||
}
|
||||
|
||||
type ChatUsageLimitConfig struct {
|
||||
|
||||
Generated
+85
-5
@@ -77,9 +77,12 @@ type sqlcQuerier interface {
|
||||
// enforces that non-deleted rows always have a provider ID.
|
||||
BackfillChatModelConfigProvider(ctx context.Context, arg BackfillChatModelConfigProviderParams) (sql.Result, error)
|
||||
BackoffChatDiffStatus(ctx context.Context, arg BackoffChatDiffStatusParams) error
|
||||
// Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs.
|
||||
BatchDeleteChatHeartbeats(ctx context.Context, arg BatchDeleteChatHeartbeatsParams) (int64, error)
|
||||
BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg BatchUpdateWorkspaceAgentMetadataParams) error
|
||||
BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg BatchUpdateWorkspaceLastUsedAtParams) error
|
||||
BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg BatchUpdateWorkspaceNextStartAtParams) error
|
||||
BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUpsertChatHeartbeatsParams) error
|
||||
BatchUpsertConnectionLogs(ctx context.Context, arg BatchUpsertConnectionLogsParams) error
|
||||
BulkMarkNotificationMessagesFailed(ctx context.Context, arg BulkMarkNotificationMessagesFailedParams) (int64, error)
|
||||
BulkMarkNotificationMessagesSent(ctx context.Context, arg BulkMarkNotificationMessagesSentParams) (int64, error)
|
||||
@@ -94,6 +97,9 @@ type sqlcQuerier interface {
|
||||
ClearChatMessageProviderResponseIDsByChatID(ctx context.Context, chatID uuid.UUID) error
|
||||
CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error)
|
||||
CountAuditLogs(ctx context.Context, arg CountAuditLogsParams) (int64, error)
|
||||
// Cheap queue-length check used by ChatMachine.Update when deciding
|
||||
// whether the chat is in a "1" sub-state.
|
||||
CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error)
|
||||
CountConnectionLogs(ctx context.Context, arg CountConnectionLogsParams) (int64, error)
|
||||
// Counts enabled, non-deleted model configs that lack both input and
|
||||
// output pricing in their JSONB options.cost configuration.
|
||||
@@ -111,7 +117,11 @@ type sqlcQuerier interface {
|
||||
DeleteAIProviderKey(ctx context.Context, id uuid.UUID) error
|
||||
DeleteAPIKeyByID(ctx context.Context, id string) error
|
||||
DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error
|
||||
// Deletes all heartbeat rows for the chat. Used during ownership
|
||||
// transitions that abandon a lease.
|
||||
DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error
|
||||
DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error
|
||||
DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error)
|
||||
DeleteAllTailnetTunnels(ctx context.Context, arg DeleteAllTailnetTunnelsParams) ([]DeleteAllTailnetTunnelsRow, error)
|
||||
// Deletes all existing webpush subscriptions.
|
||||
// This should be called when the VAPID keypair is regenerated, as the old
|
||||
@@ -133,6 +143,10 @@ type sqlcQuerier interface {
|
||||
DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error
|
||||
DeleteChatModelConfigsByProvider(ctx context.Context, provider string) error
|
||||
DeleteChatQueuedMessage(ctx context.Context, arg DeleteChatQueuedMessageParams) error
|
||||
// Deletes a queued message, scoped to the parent chat. Returns the
|
||||
// number of affected rows so callers can detect missing rows without
|
||||
// a follow-up read.
|
||||
DeleteChatQueuedMessageReturningCount(ctx context.Context, arg DeleteChatQueuedMessageReturningCountParams) (int64, error)
|
||||
DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error
|
||||
DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error
|
||||
DeleteCryptoKey(ctx context.Context, arg DeleteCryptoKeyParams) (CryptoKey, error)
|
||||
@@ -201,6 +215,7 @@ type sqlcQuerier interface {
|
||||
DeleteProvisionerKey(ctx context.Context, id uuid.UUID) error
|
||||
DeleteReplicasUpdatedBefore(ctx context.Context, updatedAt time.Time) error
|
||||
DeleteRuntimeConfig(ctx context.Context, key string) error
|
||||
DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error)
|
||||
DeleteTailnetPeer(ctx context.Context, arg DeleteTailnetPeerParams) (DeleteTailnetPeerRow, error)
|
||||
DeleteTailnetTunnel(ctx context.Context, arg DeleteTailnetTunnelParams) (DeleteTailnetTunnelRow, error)
|
||||
DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error)
|
||||
@@ -323,6 +338,10 @@ type sqlcQuerier interface {
|
||||
// This function returns roles for authorization purposes. Implied member roles
|
||||
// are included.
|
||||
GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (GetAuthorizationUserRolesRow, error)
|
||||
// Returns read-only root chat candidates for state-machine-backed
|
||||
// auto-archive. Activity is computed across the root family. The query
|
||||
// limits roots, not total family members.
|
||||
GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg GetAutoArchiveInactiveChatCandidatesParams) ([]GetAutoArchiveInactiveChatCandidatesRow, error)
|
||||
GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (BoundaryLog, error)
|
||||
GetBoundarySessionByID(ctx context.Context, id uuid.UUID) (BoundarySession, error)
|
||||
GetChatACLByID(ctx context.Context, id uuid.UUID) (GetChatACLByIDRow, error)
|
||||
@@ -334,6 +353,7 @@ type sqlcQuerier interface {
|
||||
// Auto-archive window in days. 0 disables.
|
||||
GetChatAutoArchiveDays(ctx context.Context, defaultAutoArchiveDays int32) (int32, error)
|
||||
GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error)
|
||||
GetChatByIDForShare(ctx context.Context, id uuid.UUID) (Chat, error)
|
||||
GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Chat, error)
|
||||
GetChatComputerUseProvider(ctx context.Context) (string, error)
|
||||
// Per-root-chat cost breakdown for a single user within a date range.
|
||||
@@ -371,6 +391,10 @@ type sqlcQuerier interface {
|
||||
GetChatDiffStatusSummary(ctx context.Context) (GetChatDiffStatusSummaryRow, error)
|
||||
GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds []uuid.UUID) ([]ChatDiffStatus, error)
|
||||
GetChatExploreModelOverride(ctx context.Context) (string, error)
|
||||
// Returns the chat IDs of every chat in a family (root + all children)
|
||||
// in deterministic order. The id parameter must be the root id; the
|
||||
// query does not walk up from a child.
|
||||
GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error)
|
||||
GetChatFileByID(ctx context.Context, id uuid.UUID) (ChatFile, error)
|
||||
// GetChatFileMetadataByChatID returns lightweight file metadata for
|
||||
// all files linked to a chat. The data column is excluded to avoid
|
||||
@@ -378,6 +402,7 @@ type sqlcQuerier interface {
|
||||
GetChatFileMetadataByChatID(ctx context.Context, chatID uuid.UUID) ([]GetChatFileMetadataByChatIDRow, error)
|
||||
GetChatFilesByIDs(ctx context.Context, ids []uuid.UUID) ([]ChatFile, error)
|
||||
GetChatGeneralModelOverride(ctx context.Context) (string, error)
|
||||
GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatParams) (ChatHeartbeat, error)
|
||||
// GetChatIncludeDefaultSystemPrompt preserves the legacy default
|
||||
// for deployments created before the explicit include-default toggle.
|
||||
// When the toggle is unset, a non-empty custom prompt implies false;
|
||||
@@ -391,6 +416,7 @@ type sqlcQuerier interface {
|
||||
GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error)
|
||||
GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error)
|
||||
GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error)
|
||||
GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error)
|
||||
GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error)
|
||||
GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error)
|
||||
GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error)
|
||||
@@ -400,12 +426,18 @@ type sqlcQuerier interface {
|
||||
// personal chat model overrides. It defaults to false when unset.
|
||||
GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error)
|
||||
GetChatPlanModeInstructions(ctx context.Context) (string, error)
|
||||
GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error)
|
||||
// Returns the queue head (lowest position, then lowest id).
|
||||
GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error)
|
||||
GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error)
|
||||
// Returns queued messages in state-machine order (position ASC, id ASC).
|
||||
GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error)
|
||||
// Returns the chat retention period in days. Chats archived longer
|
||||
// than this and orphaned chat files older than this are purged by
|
||||
// dbpurge. Returns 30 (days) when no value has been configured.
|
||||
// A value of 0 disables chat purging entirely.
|
||||
GetChatRetentionDays(ctx context.Context) (int32, error)
|
||||
GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]GetChatStreamSyncRowsRow, error)
|
||||
GetChatSystemPrompt(ctx context.Context) (string, error)
|
||||
// GetChatSystemPromptConfig returns both chat system prompt settings in a
|
||||
// single read to avoid torn reads between separate site-config lookups.
|
||||
@@ -430,11 +462,24 @@ type sqlcQuerier interface {
|
||||
// jsonb_array_elements never raises "cannot extract elements from a
|
||||
// scalar". Backed by idx_chat_messages_user_prompts.
|
||||
GetChatUserPromptsByChatID(ctx context.Context, arg GetChatUserPromptsByChatIDParams) ([]GetChatUserPromptsByChatIDRow, error)
|
||||
// Returns chats that workers may try to acquire. Candidates must be:
|
||||
// - in a worker-runnable execution status;
|
||||
// - unarchived; and
|
||||
// - missing ownership, carrying inconsistent ownership, or lacking a
|
||||
// fresh heartbeat for the assigned runner.
|
||||
//
|
||||
// Missing ownership is worker_id IS NULL. Inconsistent ownership is
|
||||
// runner_id IS NULL while worker_id is set. Stale ownership is no
|
||||
// heartbeat row for (chat_id, runner_id), or one older than
|
||||
// @stale_seconds by database time. Candidates are ordered by oldest
|
||||
// updated_at first so workers drain stale runnable chats predictably.
|
||||
GetChatWorkerAcquisitionCandidates(ctx context.Context, arg GetChatWorkerAcquisitionCandidatesParams) ([]GetChatWorkerAcquisitionCandidatesRow, error)
|
||||
// Returns the global TTL for chat workspaces as a Go duration string.
|
||||
// Returns "0s" (disabled) when no value has been configured.
|
||||
GetChatWorkspaceTTL(ctx context.Context) (string, error)
|
||||
GetChats(ctx context.Context, arg GetChatsParams) ([]GetChatsRow, error)
|
||||
GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) ([]Chat, error)
|
||||
GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]Chat, error)
|
||||
GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]Chat, error)
|
||||
// Retrieves chats updated after the given timestamp for telemetry
|
||||
// snapshot collection. Uses updated_at so that long-running chats
|
||||
@@ -451,6 +496,10 @@ type sqlcQuerier interface {
|
||||
GetCryptoKeysByFeature(ctx context.Context, feature CryptoKeyFeature) ([]CryptoKey, error)
|
||||
GetDBCryptKeys(ctx context.Context) ([]DBCryptKey, error)
|
||||
GetDERPMeshKey(ctx context.Context) (string, error)
|
||||
// Returns the current database timestamp. Used so transitions that
|
||||
// record deadlines or heartbeats rely on a clock that is consistent
|
||||
// with the database rather than the caller's local clock.
|
||||
GetDatabaseNow(ctx context.Context) (time.Time, error)
|
||||
GetDefaultChatModelConfig(ctx context.Context) (ChatModelConfig, error)
|
||||
GetDefaultOrganization(ctx context.Context) (Organization, error)
|
||||
GetDefaultProxyConfig(ctx context.Context) (GetDefaultProxyConfigRow, error)
|
||||
@@ -931,6 +980,8 @@ type sqlcQuerier interface {
|
||||
GetWorkspacesByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceTable, error)
|
||||
GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]GetWorkspacesEligibleForTransitionRow, error)
|
||||
GetWorkspacesForWorkspaceMetrics(ctx context.Context) ([]GetWorkspacesForWorkspaceMetricsRow, error)
|
||||
// Increments generation_attempt and returns the resulting value.
|
||||
IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error)
|
||||
InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error)
|
||||
InsertAIBridgeModelThought(ctx context.Context, arg InsertAIBridgeModelThoughtParams) (AIBridgeModelThought, error)
|
||||
InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error)
|
||||
@@ -961,7 +1012,14 @@ type sqlcQuerier interface {
|
||||
InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error)
|
||||
InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error)
|
||||
InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error)
|
||||
// Legacy queue insertion path. When no caller-supplied creator exists,
|
||||
// preserve the created_by invariant by attributing the queued row to the
|
||||
// chat owner.
|
||||
InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error)
|
||||
// Inserts a queued message that carries a position (from the default
|
||||
// sequence) and an explicit created_by reference. Use this when the
|
||||
// queued-message creator differs from the chat owner.
|
||||
InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error)
|
||||
InsertCryptoKey(ctx context.Context, arg InsertCryptoKeyParams) (CryptoKey, error)
|
||||
InsertCustomRole(ctx context.Context, arg InsertCustomRoleParams) (CustomRole, error)
|
||||
InsertDBCryptKey(ctx context.Context, arg InsertDBCryptKeyParams) error
|
||||
@@ -1041,6 +1099,11 @@ type sqlcQuerier interface {
|
||||
InsertWorkspaceProxy(ctx context.Context, arg InsertWorkspaceProxyParams) (WorkspaceProxy, error)
|
||||
InsertWorkspaceResource(ctx context.Context, arg InsertWorkspaceResourceParams) (WorkspaceResource, error)
|
||||
InsertWorkspaceResourceMetadata(ctx context.Context, arg InsertWorkspaceResourceMetadataParams) ([]WorkspaceResourceMetadatum, error)
|
||||
// Returns true when there is no heartbeat row for (chat_id, runner_id)
|
||||
// or the existing row is older than @stale_seconds seconds by database
|
||||
// time. chatstate calls this in a single query so the staleness check
|
||||
// is atomic and does not depend on the caller's local clock.
|
||||
IsChatHeartbeatStale(ctx context.Context, arg IsChatHeartbeatStaleParams) (bool, error)
|
||||
// LinkChatFiles inserts file associations into the chat_file_links
|
||||
// join table with deduplication (ON CONFLICT DO NOTHING). The INSERT
|
||||
// is conditional: it only proceeds when the total number of links
|
||||
@@ -1091,6 +1154,11 @@ type sqlcQuerier interface {
|
||||
ListUserSecretsWithValues(ctx context.Context, userID uuid.UUID) ([]UserSecret, error)
|
||||
ListUserSkillMetadataByUserID(ctx context.Context, userID uuid.UUID) ([]ListUserSkillMetadataByUserIDRow, error)
|
||||
ListWorkspaceAgentPortShares(ctx context.Context, workspaceID uuid.UUID) ([]WorkspaceAgentPortShare, error)
|
||||
// Locks the chat row with FOR UPDATE and atomically increments its
|
||||
// snapshot_version, returning the post-bump chat. This is the single
|
||||
// entry point ChatMachine.Update uses to acquire the row lock and
|
||||
// allocate a new snapshot version in one round trip.
|
||||
LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (Chat, error)
|
||||
MarkAllInboxNotificationsAsRead(ctx context.Context, arg MarkAllInboxNotificationsAsReadParams) error
|
||||
OIDCClaimFieldValues(ctx context.Context, arg OIDCClaimFieldValuesParams) ([]string, error)
|
||||
// OIDCClaimFields returns a list of distinct keys in the the merged_claims fields.
|
||||
@@ -1115,6 +1183,9 @@ type sqlcQuerier interface {
|
||||
// Mutates only created_at on the target row; ids are unchanged so
|
||||
// consumers can keep tracking queued messages by id.
|
||||
ReorderChatQueuedMessageToFront(ctx context.Context, arg ReorderChatQueuedMessageToFrontParams) (int64, error)
|
||||
// Sets the target queued message's position to one less than the
|
||||
// current minimum position for that chat, moving it to the head.
|
||||
ReorderChatQueuedMessageToHead(ctx context.Context, arg ReorderChatQueuedMessageToHeadParams) (int64, error)
|
||||
// Resolves the effective spend limit for a user using the hierarchy:
|
||||
// 1. Individual user override (highest priority, applies globally across
|
||||
// all organizations since it lives on the users table)
|
||||
@@ -1216,6 +1287,11 @@ type sqlcQuerier interface {
|
||||
// parameter keeps updated_at under the caller's clock, matching
|
||||
// the injectable quartz.Clock used by FinalizeStale sweeps.
|
||||
UpdateChatDebugStep(ctx context.Context, arg UpdateChatDebugStepParams) (ChatDebugStep, error)
|
||||
// Atomically updates the execution-state-managed fields on a chat:
|
||||
// status, archived, last_error, ownership identifiers, and the
|
||||
// requires-action deadline. Callers compose this with transition
|
||||
// mutations inside a single ChatMachine.Update transaction.
|
||||
UpdateChatExecutionState(ctx context.Context, arg UpdateChatExecutionStateParams) (Chat, error)
|
||||
// Bumps the heartbeat timestamp for the given set of chat IDs,
|
||||
// provided they are still running and owned by the specified
|
||||
// worker. Returns the IDs that were actually updated so the
|
||||
@@ -1234,11 +1310,9 @@ type sqlcQuerier interface {
|
||||
// Updates the cached last completed turn summary for sidebar display.
|
||||
// Empty or whitespace-only summaries are stored as NULL here so direct
|
||||
// query callers cannot accidentally persist blank sidebar text.
|
||||
// This intentionally preserves updated_at. The staleness guard relies on
|
||||
// every new-turn query, such as UpdateChatStatus and AcquireChats, bumping
|
||||
// updated_at. Future chat-field updates that do not bump updated_at can let
|
||||
// stale summaries persist. If this query ever bumps updated_at, later
|
||||
// goroutine summary writes will be rejected as stale.
|
||||
// This intentionally preserves updated_at. The staleness guard uses
|
||||
// history_version so worker lifecycle transitions that do not change the
|
||||
// active message history cannot reject final turn summary writes.
|
||||
// Two summary workers using the same freshness marker are last-write-wins.
|
||||
UpdateChatLastTurnSummary(ctx context.Context, arg UpdateChatLastTurnSummaryParams) (int64, error)
|
||||
UpdateChatMCPServerIDs(ctx context.Context, arg UpdateChatMCPServerIDsParams) (Chat, error)
|
||||
@@ -1246,6 +1320,9 @@ type sqlcQuerier interface {
|
||||
UpdateChatModelConfig(ctx context.Context, arg UpdateChatModelConfigParams) (ChatModelConfig, error)
|
||||
UpdateChatPinOrder(ctx context.Context, arg UpdateChatPinOrderParams) error
|
||||
UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatPlanModeByIDParams) (Chat, error)
|
||||
// Stores the client-visible retry payload. retry_state_version is
|
||||
// assigned by trigger from the current snapshot_version.
|
||||
UpdateChatRetryState(ctx context.Context, arg UpdateChatRetryStateParams) (Chat, error)
|
||||
UpdateChatStatus(ctx context.Context, arg UpdateChatStatusParams) (Chat, error)
|
||||
UpdateChatStatusPreserveUpdatedAt(ctx context.Context, arg UpdateChatStatusPreserveUpdatedAtParams) (Chat, error)
|
||||
UpdateChatTitleByID(ctx context.Context, arg UpdateChatTitleByIDParams) (Chat, error)
|
||||
@@ -1396,6 +1473,9 @@ type sqlcQuerier interface {
|
||||
UpsertChatDiffStatusReference(ctx context.Context, arg UpsertChatDiffStatusReferenceParams) (ChatDiffStatus, error)
|
||||
UpsertChatExploreModelOverride(ctx context.Context, value string) error
|
||||
UpsertChatGeneralModelOverride(ctx context.Context, value string) error
|
||||
// Upserts a heartbeat row for the (chat_id, runner_id) lease. Uses
|
||||
// database time so callers do not depend on a local clock.
|
||||
UpsertChatHeartbeat(ctx context.Context, arg UpsertChatHeartbeatParams) error
|
||||
UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error
|
||||
// UpsertChatPersonalModelOverridesEnabled updates whether users may configure
|
||||
// personal chat model overrides.
|
||||
|
||||
@@ -11063,10 +11063,12 @@ func TestInsertChatMessages(t *testing.T) {
|
||||
|
||||
insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) {
|
||||
t.Helper()
|
||||
apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: userID})
|
||||
|
||||
_, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||
ChatID: chatID,
|
||||
CreatedBy: []uuid.UUID{userID},
|
||||
APIKeyID: []string{apiKey.ID},
|
||||
ModelConfigID: []uuid.UUID{modelConfigID},
|
||||
Role: []database.ChatMessageRole{database.ChatMessageRoleUser},
|
||||
ContentVersion: []int16{chatprompt.CurrentContentVersion},
|
||||
@@ -11125,10 +11127,12 @@ func TestInsertChatMessages(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
store, ctx, user, chat, _, modelConfigA := setupChat(t)
|
||||
apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: user.ID})
|
||||
|
||||
msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: []uuid.UUID{user.ID, uuid.Nil, uuid.Nil},
|
||||
APIKeyID: []string{apiKey.ID, "", ""},
|
||||
ModelConfigID: []uuid.UUID{modelConfigA.ID, modelConfigA.ID, modelConfigA.ID},
|
||||
Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool},
|
||||
ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion},
|
||||
@@ -12744,9 +12748,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
affected, err := db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedUpdatedAt: chat.UpdatedAt,
|
||||
LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true},
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, affected)
|
||||
@@ -12757,9 +12761,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) {
|
||||
require.Equal(t, chat.UpdatedAt, fetched.UpdatedAt)
|
||||
|
||||
affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedUpdatedAt: chat.UpdatedAt,
|
||||
LastTurnSummary: sql.NullString{String: " \n\t ", Valid: true},
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: " \n\t ", Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, affected)
|
||||
@@ -12770,9 +12774,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) {
|
||||
require.Equal(t, chat.UpdatedAt, fetched.UpdatedAt)
|
||||
|
||||
affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedUpdatedAt: chat.UpdatedAt,
|
||||
LastTurnSummary: sql.NullString{String: "fresh summary", Valid: true},
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: "fresh summary", Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, affected)
|
||||
@@ -12786,17 +12790,54 @@ func TestUpdateChatLastTurnSummary(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedUpdatedAt: chat.UpdatedAt,
|
||||
LastTurnSummary: sql.NullString{String: "stale summary", Valid: true},
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: "still fresh summary", Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, affected)
|
||||
|
||||
fetched, err = db.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, sql.NullString{String: "still fresh summary", Valid: true}, fetched.LastTurnSummary)
|
||||
require.Equal(t, advancedUpdatedAt, fetched.UpdatedAt)
|
||||
|
||||
_, err = db.LockChatAndBumpSnapshotVersion(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: []uuid.UUID{owner.ID},
|
||||
ModelConfigID: []uuid.UUID{modelCfg.ID},
|
||||
Role: []database.ChatMessageRole{database.ChatMessageRoleUser},
|
||||
Content: []string{`[{"type":"text","text":"new request"}]`},
|
||||
ContentVersion: []int16{chatprompt.CurrentContentVersion},
|
||||
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
|
||||
InputTokens: []int64{0},
|
||||
OutputTokens: []int64{0},
|
||||
TotalTokens: []int64{0},
|
||||
ReasoningTokens: []int64{0},
|
||||
CacheCreationTokens: []int64{0},
|
||||
CacheReadTokens: []int64{0},
|
||||
ContextLimit: []int64{0},
|
||||
Compressed: []bool{false},
|
||||
TotalCostMicros: []int64{0},
|
||||
RuntimeMs: []int64{0},
|
||||
ProviderResponseID: []string{""},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{
|
||||
ID: chat.ID,
|
||||
ExpectedHistoryVersion: chat.HistoryVersion,
|
||||
LastTurnSummary: sql.NullString{String: "stale summary", Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, affected)
|
||||
|
||||
fetched, err = db.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, sql.NullString{String: "fresh summary", Valid: true}, fetched.LastTurnSummary)
|
||||
require.Equal(t, advancedUpdatedAt, fetched.UpdatedAt)
|
||||
require.Equal(t, sql.NullString{String: "still fresh summary", Valid: true}, fetched.LastTurnSummary)
|
||||
require.NotEqual(t, chat.HistoryVersion, fetched.HistoryVersion)
|
||||
}
|
||||
|
||||
func TestDeleteChatDebugDataAfterMessageIDIncludesTriggeredRuns(t *testing.T) {
|
||||
|
||||
Generated
+1825
-111
File diff suppressed because it is too large
Load Diff
@@ -35,6 +35,14 @@ chats_expanded AS (
|
||||
updated_chats.plan_mode,
|
||||
updated_chats.client_type,
|
||||
updated_chats.last_turn_summary,
|
||||
updated_chats.snapshot_version,
|
||||
updated_chats.history_version,
|
||||
updated_chats.queue_version,
|
||||
updated_chats.generation_attempt,
|
||||
updated_chats.retry_state,
|
||||
updated_chats.retry_state_version,
|
||||
updated_chats.runner_id,
|
||||
updated_chats.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -90,6 +98,14 @@ chats_expanded AS (
|
||||
updated_chats.plan_mode,
|
||||
updated_chats.client_type,
|
||||
updated_chats.last_turn_summary,
|
||||
updated_chats.snapshot_version,
|
||||
updated_chats.history_version,
|
||||
updated_chats.queue_version,
|
||||
updated_chats.generation_attempt,
|
||||
updated_chats.retry_state,
|
||||
updated_chats.retry_state_version,
|
||||
updated_chats.runner_id,
|
||||
updated_chats.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -293,6 +309,15 @@ SELECT *
|
||||
FROM chats_expanded
|
||||
WHERE id = @id::uuid;
|
||||
|
||||
-- name: GetChatFamilyIDsByRootID :many
|
||||
-- Returns the chat IDs of every chat in a family (root + all children)
|
||||
-- in deterministic order. The id parameter must be the root id; the
|
||||
-- query does not walk up from a child.
|
||||
SELECT id
|
||||
FROM chats
|
||||
WHERE id = @id::uuid OR root_chat_id = @id::uuid
|
||||
ORDER BY (id = @id::uuid) DESC, created_at ASC, id ASC;
|
||||
|
||||
-- name: GetChatACLByID :one
|
||||
SELECT
|
||||
user_acl AS users,
|
||||
@@ -333,6 +358,18 @@ WHERE
|
||||
ORDER BY
|
||||
created_at ASC;
|
||||
|
||||
-- name: GetChatMessagesByRevisionForStream :many
|
||||
SELECT
|
||||
*
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
chat_id = @chat_id::uuid
|
||||
AND revision > @after_revision::bigint
|
||||
AND visibility IN ('user', 'both')
|
||||
ORDER BY
|
||||
created_at ASC, id ASC;
|
||||
|
||||
-- name: GetChatMessagesByChatIDAscPaginated :many
|
||||
SELECT
|
||||
*
|
||||
@@ -723,6 +760,14 @@ chats_expanded AS (
|
||||
inserted_chat.plan_mode,
|
||||
inserted_chat.client_type,
|
||||
inserted_chat.last_turn_summary,
|
||||
inserted_chat.snapshot_version,
|
||||
inserted_chat.history_version,
|
||||
inserted_chat.queue_version,
|
||||
inserted_chat.generation_attempt,
|
||||
inserted_chat.retry_state,
|
||||
inserted_chat.retry_state_version,
|
||||
inserted_chat.runner_id,
|
||||
inserted_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, inserted_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, inserted_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -860,6 +905,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -915,6 +968,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -968,6 +1029,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1021,6 +1090,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1074,6 +1151,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1126,6 +1211,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1178,6 +1271,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1232,6 +1333,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1248,11 +1357,9 @@ FROM chats_expanded;
|
||||
-- Updates the cached last completed turn summary for sidebar display.
|
||||
-- Empty or whitespace-only summaries are stored as NULL here so direct
|
||||
-- query callers cannot accidentally persist blank sidebar text.
|
||||
-- This intentionally preserves updated_at. The staleness guard relies on
|
||||
-- every new-turn query, such as UpdateChatStatus and AcquireChats, bumping
|
||||
-- updated_at. Future chat-field updates that do not bump updated_at can let
|
||||
-- stale summaries persist. If this query ever bumps updated_at, later
|
||||
-- goroutine summary writes will be rejected as stale.
|
||||
-- This intentionally preserves updated_at. The staleness guard uses
|
||||
-- history_version so worker lifecycle transitions that do not change the
|
||||
-- active message history cannot reject final turn summary writes.
|
||||
-- Two summary workers using the same freshness marker are last-write-wins.
|
||||
UPDATE chats
|
||||
SET
|
||||
@@ -1261,7 +1368,7 @@ SET
|
||||
), '')
|
||||
WHERE
|
||||
id = @id::uuid
|
||||
AND updated_at = @expected_updated_at::timestamptz;
|
||||
AND history_version = @expected_history_version::bigint;
|
||||
|
||||
-- name: UpdateChatMCPServerIDs :one
|
||||
WITH updated_chat AS (
|
||||
@@ -1304,6 +1411,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1413,6 +1528,14 @@ chats_expanded AS (
|
||||
acquired_chats.plan_mode,
|
||||
acquired_chats.client_type,
|
||||
acquired_chats.last_turn_summary,
|
||||
acquired_chats.snapshot_version,
|
||||
acquired_chats.history_version,
|
||||
acquired_chats.queue_version,
|
||||
acquired_chats.generation_attempt,
|
||||
acquired_chats.retry_state,
|
||||
acquired_chats.retry_state_version,
|
||||
acquired_chats.runner_id,
|
||||
acquired_chats.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, acquired_chats.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, acquired_chats.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1470,6 +1593,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1527,6 +1658,14 @@ chats_expanded AS (
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1694,13 +1833,18 @@ RETURNING
|
||||
*;
|
||||
|
||||
-- name: InsertChatQueuedMessage :one
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id)
|
||||
VALUES (
|
||||
@chat_id,
|
||||
@content,
|
||||
-- Legacy queue insertion path. When no caller-supplied creator exists,
|
||||
-- preserve the created_by invariant by attributing the queued row to the
|
||||
-- chat owner.
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by)
|
||||
SELECT
|
||||
@chat_id::uuid,
|
||||
@content::jsonb,
|
||||
sqlc.narg('model_config_id')::uuid,
|
||||
sqlc.narg('api_key_id')::text
|
||||
)
|
||||
sqlc.narg('api_key_id')::text,
|
||||
chats.owner_id
|
||||
FROM chats
|
||||
WHERE chats.id = @chat_id::uuid
|
||||
RETURNING *;
|
||||
|
||||
-- name: GetChatQueuedMessages :many
|
||||
@@ -1786,6 +1930,14 @@ chats_expanded AS (
|
||||
locked_chat.plan_mode,
|
||||
locked_chat.client_type,
|
||||
locked_chat.last_turn_summary,
|
||||
locked_chat.snapshot_version,
|
||||
locked_chat.history_version,
|
||||
locked_chat.queue_version,
|
||||
locked_chat.generation_attempt,
|
||||
locked_chat.retry_state,
|
||||
locked_chat.retry_state_version,
|
||||
locked_chat.runner_id,
|
||||
locked_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, locked_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, locked_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
@@ -1798,6 +1950,63 @@ chats_expanded AS (
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: GetChatByIDForShare :one
|
||||
WITH shared_chat AS (
|
||||
SELECT *
|
||||
FROM chats
|
||||
WHERE id = @id::uuid
|
||||
FOR SHARE
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
shared_chat.id,
|
||||
shared_chat.owner_id,
|
||||
shared_chat.workspace_id,
|
||||
shared_chat.title,
|
||||
shared_chat.status,
|
||||
shared_chat.worker_id,
|
||||
shared_chat.started_at,
|
||||
shared_chat.heartbeat_at,
|
||||
shared_chat.created_at,
|
||||
shared_chat.updated_at,
|
||||
shared_chat.parent_chat_id,
|
||||
shared_chat.root_chat_id,
|
||||
shared_chat.last_model_config_id,
|
||||
shared_chat.archived,
|
||||
shared_chat.last_error,
|
||||
shared_chat.mode,
|
||||
shared_chat.mcp_server_ids,
|
||||
shared_chat.labels,
|
||||
shared_chat.build_id,
|
||||
shared_chat.agent_id,
|
||||
shared_chat.pin_order,
|
||||
shared_chat.last_read_message_id,
|
||||
shared_chat.last_injected_context,
|
||||
shared_chat.dynamic_tools,
|
||||
shared_chat.organization_id,
|
||||
shared_chat.plan_mode,
|
||||
shared_chat.client_type,
|
||||
shared_chat.last_turn_summary,
|
||||
shared_chat.snapshot_version,
|
||||
shared_chat.history_version,
|
||||
shared_chat.queue_version,
|
||||
shared_chat.generation_attempt,
|
||||
shared_chat.retry_state,
|
||||
shared_chat.retry_state_version,
|
||||
shared_chat.runner_id,
|
||||
shared_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, shared_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, shared_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM
|
||||
shared_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(shared_chat.root_chat_id, shared_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = shared_chat.owner_id
|
||||
)
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: GetChatsByChatFileID :many
|
||||
SELECT
|
||||
*
|
||||
@@ -2317,6 +2526,420 @@ WHERE chat_id = @chat_id::uuid
|
||||
AND deleted = false
|
||||
AND content::jsonb @> '[{"type": "context-file"}]';
|
||||
|
||||
-- name: GetChatWorkerAcquisitionCandidates :many
|
||||
-- Returns chats that workers may try to acquire. Candidates must be:
|
||||
-- - in a worker-runnable execution status;
|
||||
-- - unarchived; and
|
||||
-- - missing ownership, carrying inconsistent ownership, or lacking a
|
||||
-- fresh heartbeat for the assigned runner.
|
||||
--
|
||||
-- Missing ownership is worker_id IS NULL. Inconsistent ownership is
|
||||
-- runner_id IS NULL while worker_id is set. Stale ownership is no
|
||||
-- heartbeat row for (chat_id, runner_id), or one older than
|
||||
-- @stale_seconds by database time. Candidates are ordered by oldest
|
||||
-- updated_at first so workers drain stale runnable chats predictably.
|
||||
SELECT
|
||||
chats_expanded.*,
|
||||
chat_heartbeats.heartbeat_at AS current_heartbeat_at,
|
||||
NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_heartbeats current_lease
|
||||
WHERE current_lease.chat_id = chats_expanded.id
|
||||
AND current_lease.runner_id = chats_expanded.runner_id
|
||||
AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int)
|
||||
) AS heartbeat_stale
|
||||
FROM chats_expanded
|
||||
LEFT JOIN chat_heartbeats
|
||||
ON chat_heartbeats.chat_id = chats_expanded.id
|
||||
AND chat_heartbeats.runner_id = chats_expanded.runner_id
|
||||
WHERE
|
||||
chats_expanded.status IN ('running'::chat_status, 'interrupting'::chat_status, 'requires_action'::chat_status)
|
||||
AND chats_expanded.archived = false
|
||||
AND (
|
||||
chats_expanded.worker_id IS NULL
|
||||
OR chats_expanded.runner_id IS NULL
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_heartbeats current_lease
|
||||
WHERE current_lease.chat_id = chats_expanded.id
|
||||
AND current_lease.runner_id = chats_expanded.runner_id
|
||||
AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int)
|
||||
)
|
||||
)
|
||||
ORDER BY chats_expanded.updated_at ASC, chats_expanded.id ASC
|
||||
LIMIT @limit_count::int;
|
||||
|
||||
-- name: GetChatsByIDsForRunnerSync :many
|
||||
SELECT *
|
||||
FROM chats_expanded
|
||||
WHERE id = ANY(@ids::uuid[])
|
||||
ORDER BY id ASC;
|
||||
|
||||
-- name: BatchUpsertChatHeartbeats :exec
|
||||
INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at)
|
||||
SELECT chat_ids.chat_id, runner_ids.runner_id, NOW()
|
||||
FROM unnest(@chat_ids::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord)
|
||||
JOIN unnest(@runner_ids::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord)
|
||||
ON CONFLICT (chat_id, runner_id) DO UPDATE
|
||||
SET heartbeat_at = EXCLUDED.heartbeat_at;
|
||||
|
||||
-- name: DeleteStaleChatHeartbeats :execrows
|
||||
DELETE FROM chat_heartbeats
|
||||
WHERE heartbeat_at < NOW() - (INTERVAL '1 second' * @stale_seconds::int);
|
||||
|
||||
-- name: GetAutoArchiveInactiveChatCandidates :many
|
||||
-- Returns read-only root chat candidates for state-machine-backed
|
||||
-- auto-archive. Activity is computed across the root family. The query
|
||||
-- limits roots, not total family members.
|
||||
SELECT
|
||||
chats_expanded.*,
|
||||
COALESCE(activity.last_activity_at, chats_expanded.created_at)::timestamptz AS last_activity_at
|
||||
FROM chats_expanded
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT MAX(chat_messages.created_at) AS last_activity_at
|
||||
FROM chat_messages
|
||||
JOIN chats family_chat ON family_chat.id = chat_messages.chat_id
|
||||
WHERE (family_chat.id = chats_expanded.id OR family_chat.root_chat_id = chats_expanded.id)
|
||||
AND chat_messages.deleted = false
|
||||
) activity ON TRUE
|
||||
WHERE
|
||||
chats_expanded.archived = false
|
||||
AND chats_expanded.pin_order = 0
|
||||
AND chats_expanded.parent_chat_id IS NULL
|
||||
AND chats_expanded.created_at < @archive_cutoff::timestamptz
|
||||
AND chats_expanded.status NOT IN (
|
||||
'running'::chat_status,
|
||||
'interrupting'::chat_status,
|
||||
'pending'::chat_status,
|
||||
'paused'::chat_status,
|
||||
'requires_action'::chat_status
|
||||
)
|
||||
AND COALESCE(activity.last_activity_at, chats_expanded.created_at) < @archive_cutoff::timestamptz
|
||||
ORDER BY chats_expanded.created_at ASC
|
||||
LIMIT @limit_count::int;
|
||||
|
||||
|
||||
-- name: LockChatAndBumpSnapshotVersion :one
|
||||
-- Locks the chat row with FOR UPDATE and atomically increments its
|
||||
-- snapshot_version, returning the post-bump chat. This is the single
|
||||
-- entry point ChatMachine.Update uses to acquire the row lock and
|
||||
-- allocate a new snapshot version in one round trip.
|
||||
WITH bumped_chat AS (
|
||||
UPDATE chats
|
||||
SET snapshot_version = snapshot_version + 1
|
||||
WHERE id = (
|
||||
SELECT id FROM chats
|
||||
WHERE id = @id::uuid
|
||||
FOR UPDATE
|
||||
)
|
||||
RETURNING *
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
bumped_chat.id,
|
||||
bumped_chat.owner_id,
|
||||
bumped_chat.workspace_id,
|
||||
bumped_chat.title,
|
||||
bumped_chat.status,
|
||||
bumped_chat.worker_id,
|
||||
bumped_chat.started_at,
|
||||
bumped_chat.heartbeat_at,
|
||||
bumped_chat.created_at,
|
||||
bumped_chat.updated_at,
|
||||
bumped_chat.parent_chat_id,
|
||||
bumped_chat.root_chat_id,
|
||||
bumped_chat.last_model_config_id,
|
||||
bumped_chat.archived,
|
||||
bumped_chat.last_error,
|
||||
bumped_chat.mode,
|
||||
bumped_chat.mcp_server_ids,
|
||||
bumped_chat.labels,
|
||||
bumped_chat.build_id,
|
||||
bumped_chat.agent_id,
|
||||
bumped_chat.pin_order,
|
||||
bumped_chat.last_read_message_id,
|
||||
bumped_chat.last_injected_context,
|
||||
bumped_chat.dynamic_tools,
|
||||
bumped_chat.organization_id,
|
||||
bumped_chat.plan_mode,
|
||||
bumped_chat.client_type,
|
||||
bumped_chat.last_turn_summary,
|
||||
bumped_chat.snapshot_version,
|
||||
bumped_chat.history_version,
|
||||
bumped_chat.queue_version,
|
||||
bumped_chat.generation_attempt,
|
||||
bumped_chat.retry_state,
|
||||
bumped_chat.retry_state_version,
|
||||
bumped_chat.runner_id,
|
||||
bumped_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, bumped_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, bumped_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM bumped_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(bumped_chat.root_chat_id, bumped_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = bumped_chat.owner_id
|
||||
)
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: UpdateChatExecutionState :one
|
||||
-- Atomically updates the execution-state-managed fields on a chat:
|
||||
-- status, archived, last_error, ownership identifiers, and the
|
||||
-- requires-action deadline. Callers compose this with transition
|
||||
-- mutations inside a single ChatMachine.Update transaction.
|
||||
WITH updated_chat AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
status = @status::chat_status,
|
||||
archived = @archived::boolean,
|
||||
worker_id = sqlc.narg('worker_id')::uuid,
|
||||
runner_id = sqlc.narg('runner_id')::uuid,
|
||||
last_error = sqlc.narg('last_error')::jsonb,
|
||||
requires_action_deadline_at = sqlc.narg('requires_action_deadline_at')::timestamptz,
|
||||
pin_order = CASE WHEN @archived::boolean THEN 0 ELSE pin_order END,
|
||||
updated_at = NOW()
|
||||
WHERE id = @id::uuid
|
||||
RETURNING *
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
updated_chat.id,
|
||||
updated_chat.owner_id,
|
||||
updated_chat.workspace_id,
|
||||
updated_chat.title,
|
||||
updated_chat.status,
|
||||
updated_chat.worker_id,
|
||||
updated_chat.started_at,
|
||||
updated_chat.heartbeat_at,
|
||||
updated_chat.created_at,
|
||||
updated_chat.updated_at,
|
||||
updated_chat.parent_chat_id,
|
||||
updated_chat.root_chat_id,
|
||||
updated_chat.last_model_config_id,
|
||||
updated_chat.archived,
|
||||
updated_chat.last_error,
|
||||
updated_chat.mode,
|
||||
updated_chat.mcp_server_ids,
|
||||
updated_chat.labels,
|
||||
updated_chat.build_id,
|
||||
updated_chat.agent_id,
|
||||
updated_chat.pin_order,
|
||||
updated_chat.last_read_message_id,
|
||||
updated_chat.last_injected_context,
|
||||
updated_chat.dynamic_tools,
|
||||
updated_chat.organization_id,
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM updated_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = updated_chat.owner_id
|
||||
)
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: UpdateChatRetryState :one
|
||||
-- Stores the client-visible retry payload. retry_state_version is
|
||||
-- assigned by trigger from the current snapshot_version.
|
||||
WITH updated_chat AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
retry_state = @retry_state::jsonb,
|
||||
updated_at = NOW()
|
||||
WHERE id = @id::uuid
|
||||
RETURNING *
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
updated_chat.id,
|
||||
updated_chat.owner_id,
|
||||
updated_chat.workspace_id,
|
||||
updated_chat.title,
|
||||
updated_chat.status,
|
||||
updated_chat.worker_id,
|
||||
updated_chat.started_at,
|
||||
updated_chat.heartbeat_at,
|
||||
updated_chat.created_at,
|
||||
updated_chat.updated_at,
|
||||
updated_chat.parent_chat_id,
|
||||
updated_chat.root_chat_id,
|
||||
updated_chat.last_model_config_id,
|
||||
updated_chat.archived,
|
||||
updated_chat.last_error,
|
||||
updated_chat.mode,
|
||||
updated_chat.mcp_server_ids,
|
||||
updated_chat.labels,
|
||||
updated_chat.build_id,
|
||||
updated_chat.agent_id,
|
||||
updated_chat.pin_order,
|
||||
updated_chat.last_read_message_id,
|
||||
updated_chat.last_injected_context,
|
||||
updated_chat.dynamic_tools,
|
||||
updated_chat.organization_id,
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name
|
||||
FROM updated_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = updated_chat.owner_id
|
||||
)
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: IncrementChatGenerationAttempt :one
|
||||
-- Increments generation_attempt and returns the resulting value.
|
||||
UPDATE chats
|
||||
SET generation_attempt = generation_attempt + 1, updated_at = NOW()
|
||||
WHERE id = @id::uuid
|
||||
RETURNING generation_attempt;
|
||||
|
||||
-- name: GetDatabaseNow :one
|
||||
-- Returns the current database timestamp. Used so transitions that
|
||||
-- record deadlines or heartbeats rely on a clock that is consistent
|
||||
-- with the database rather than the caller's local clock.
|
||||
SELECT NOW()::timestamptz AS now;
|
||||
|
||||
-- name: InsertChatQueuedMessageWithCreator :one
|
||||
-- Inserts a queued message that carries a position (from the default
|
||||
-- sequence) and an explicit created_by reference. Use this when the
|
||||
-- queued-message creator differs from the chat owner.
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by)
|
||||
VALUES (
|
||||
@chat_id::uuid,
|
||||
@content::jsonb,
|
||||
sqlc.narg('model_config_id')::uuid,
|
||||
sqlc.narg('api_key_id')::text,
|
||||
@created_by::uuid
|
||||
)
|
||||
RETURNING *;
|
||||
|
||||
-- name: GetChatQueuedMessagesByPosition :many
|
||||
-- Returns queued messages in state-machine order (position ASC, id ASC).
|
||||
SELECT * FROM chat_queued_messages
|
||||
WHERE chat_id = @chat_id::uuid
|
||||
ORDER BY position ASC, id ASC;
|
||||
|
||||
-- name: CountChatQueuedMessages :one
|
||||
-- Cheap queue-length check used by ChatMachine.Update when deciding
|
||||
-- whether the chat is in a "1" sub-state.
|
||||
SELECT COUNT(*)::bigint AS count
|
||||
FROM chat_queued_messages
|
||||
WHERE chat_id = @chat_id::uuid;
|
||||
|
||||
-- name: GetChatQueuedMessageHead :one
|
||||
-- Returns the queue head (lowest position, then lowest id).
|
||||
SELECT * FROM chat_queued_messages
|
||||
WHERE chat_id = @chat_id::uuid
|
||||
ORDER BY position ASC, id ASC
|
||||
LIMIT 1;
|
||||
|
||||
-- name: GetChatQueuedMessageByID :one
|
||||
SELECT * FROM chat_queued_messages
|
||||
WHERE id = @id::bigint AND chat_id = @chat_id::uuid;
|
||||
|
||||
-- name: DeleteChatQueuedMessageReturningCount :execrows
|
||||
-- Deletes a queued message, scoped to the parent chat. Returns the
|
||||
-- number of affected rows so callers can detect missing rows without
|
||||
-- a follow-up read.
|
||||
DELETE FROM chat_queued_messages
|
||||
WHERE id = @id::bigint AND chat_id = @chat_id::uuid;
|
||||
|
||||
-- name: DeleteAllChatQueuedMessagesReturningCount :execrows
|
||||
DELETE FROM chat_queued_messages
|
||||
WHERE chat_id = @chat_id::uuid;
|
||||
|
||||
-- name: ReorderChatQueuedMessageToHead :execrows
|
||||
-- Sets the target queued message's position to one less than the
|
||||
-- current minimum position for that chat, moving it to the head.
|
||||
UPDATE chat_queued_messages AS target
|
||||
SET position = COALESCE(
|
||||
(SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = @chat_id::uuid),
|
||||
0
|
||||
) - 1
|
||||
WHERE target.id = @id::bigint
|
||||
AND target.chat_id = @chat_id::uuid
|
||||
AND target.position > COALESCE(
|
||||
(SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = @chat_id::uuid),
|
||||
target.position
|
||||
);
|
||||
|
||||
-- name: UpsertChatHeartbeat :exec
|
||||
-- Upserts a heartbeat row for the (chat_id, runner_id) lease. Uses
|
||||
-- database time so callers do not depend on a local clock.
|
||||
INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at)
|
||||
VALUES (@chat_id::uuid, @runner_id::uuid, NOW())
|
||||
ON CONFLICT (chat_id, runner_id) DO UPDATE
|
||||
SET heartbeat_at = EXCLUDED.heartbeat_at;
|
||||
|
||||
-- name: GetChatHeartbeat :one
|
||||
SELECT * FROM chat_heartbeats
|
||||
WHERE chat_id = @chat_id::uuid AND runner_id = @runner_id::uuid;
|
||||
|
||||
-- name: IsChatHeartbeatStale :one
|
||||
-- Returns true when there is no heartbeat row for (chat_id, runner_id)
|
||||
-- or the existing row is older than @stale_seconds seconds by database
|
||||
-- time. chatstate calls this in a single query so the staleness check
|
||||
-- is atomic and does not depend on the caller's local clock.
|
||||
SELECT NOT EXISTS (
|
||||
SELECT 1 FROM chat_heartbeats
|
||||
WHERE chat_id = @chat_id::uuid
|
||||
AND runner_id = @runner_id::uuid
|
||||
AND heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int)
|
||||
) AS stale;
|
||||
|
||||
-- name: BatchDeleteChatHeartbeats :execrows
|
||||
-- Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs.
|
||||
DELETE FROM chat_heartbeats
|
||||
USING unnest(@chat_ids::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord)
|
||||
JOIN unnest(@runner_ids::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord)
|
||||
WHERE chat_heartbeats.chat_id = chat_ids.chat_id
|
||||
AND chat_heartbeats.runner_id = runner_ids.runner_id;
|
||||
|
||||
-- name: DeleteAllChatHeartbeats :exec
|
||||
-- Deletes all heartbeat rows for the chat. Used during ownership
|
||||
-- transitions that abandon a lease.
|
||||
DELETE FROM chat_heartbeats WHERE chat_id = @chat_id::uuid;
|
||||
|
||||
|
||||
-- name: GetChatStreamSyncRows :many
|
||||
SELECT
|
||||
id,
|
||||
snapshot_version,
|
||||
history_version,
|
||||
queue_version,
|
||||
retry_state_version,
|
||||
generation_attempt,
|
||||
status,
|
||||
worker_id
|
||||
FROM chats
|
||||
WHERE id = ANY(@ids::uuid[])
|
||||
ORDER BY id ASC;
|
||||
|
||||
-- name: AutoArchiveInactiveChats :many
|
||||
-- Archives inactive root chats (pinned and already-archived chats skipped),
|
||||
-- cascading to children via root_chat_id. Limits apply to roots, not total
|
||||
|
||||
Generated
+1
@@ -26,6 +26,7 @@ const (
|
||||
UniqueChatDiffStatusesPkey UniqueConstraint = "chat_diff_statuses_pkey" // ALTER TABLE ONLY chat_diff_statuses ADD CONSTRAINT chat_diff_statuses_pkey PRIMARY KEY (chat_id);
|
||||
UniqueChatFileLinksChatIDFileIDKey UniqueConstraint = "chat_file_links_chat_id_file_id_key" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_chat_id_file_id_key UNIQUE (chat_id, file_id);
|
||||
UniqueChatFilesPkey UniqueConstraint = "chat_files_pkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id);
|
||||
UniqueChatHeartbeatsPkey UniqueConstraint = "chat_heartbeats_pkey" // ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id);
|
||||
UniqueChatMessagesPkey UniqueConstraint = "chat_messages_pkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id);
|
||||
UniqueChatModelConfigsPkey UniqueConstraint = "chat_model_configs_pkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_pkey PRIMARY KEY (id);
|
||||
UniqueChatQueuedMessagesPkey UniqueConstraint = "chat_queued_messages_pkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_pkey PRIMARY KEY (id);
|
||||
|
||||
Reference in New Issue
Block a user