diff --git a/internal/chat/repository.go b/internal/chat/repository.go new file mode 100644 index 0000000..7646257 --- /dev/null +++ b/internal/chat/repository.go @@ -0,0 +1,1153 @@ +// Package chat — repository.go 是 sqlc 生成层之上的薄包装: +// - pgtype.* ↔ 领域类型互转 +// - pgx.ErrNoRows → errs.CodeChat*NotFound +// - PG 唯一约束冲突 → 可识别的 AppError +// - WithTx 提供事务作用域 +package chat + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "math/big" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgerrcode" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + chatsqlc "github.com/yan1h/agent-coding-workflow/internal/chat/sqlc" +) + +// Postgres epoch: 2000-01-01 00:00:00 UTC +var pgEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + +// Repository 是 chat 模块对 PG 的全部依赖。 +// Service 单元测试通过手写 fakeRepo 满足该接口。 +type Repository interface { + // ===== Conversation ===== + InsertConversation(ctx context.Context, p InsertConversationParams) (*Conversation, error) + GetConversation(ctx context.Context, id uuid.UUID) (*Conversation, error) + ListConversationsByUser(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error) + SetConversationTitleManual(ctx context.Context, id uuid.UUID, title *string) error + UpdateConversationTitle(ctx context.Context, id uuid.UUID, title *string) error + SoftDeleteConversation(ctx context.Context, id uuid.UUID) error + + // ===== Message ===== + InsertMessage(ctx context.Context, conversationID uuid.UUID, role MessageRole, content string, status MessageStatus) (*Message, error) + GetMessage(ctx context.Context, id int64) (*Message, error) + GetFirstUserMessage(ctx context.Context, conversationID uuid.UUID) (*Message, error) + ListMessagesByConversation(ctx context.Context, conversationID uuid.UUID, beforeID *int64, limit int) ([]Message, error) + ListMessagesAscending(ctx context.Context, conversationID uuid.UUID) ([]Message, error) + UpdateMessageOK(ctx context.Context, p UpdateMessageOKParams) error + UpdateMessageError(ctx context.Context, id int64, errMsg *string) error + UpdateMessageCancelled(ctx context.Context, id int64) error + MarkPendingAsErrorOnStartup(ctx context.Context) error + DeleteMessage(ctx context.Context, id int64) error + + // ===== Attachment ===== + InsertAttachment(ctx context.Context, p InsertAttachmentParams) (*Attachment, error) + GetAttachment(ctx context.Context, id uuid.UUID) (*Attachment, error) + AttachToMessage(ctx context.Context, ids []uuid.UUID, messageID int64) error + ListAttachmentsForMessage(ctx context.Context, messageID int64) ([]Attachment, error) + ListOrphanAttachmentsForUser(ctx context.Context, userID uuid.UUID, ids []uuid.UUID) ([]Attachment, error) + ListExpiredOrphans(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) + ListAttachmentsForSoftDeletedConversations(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) + DeleteAttachment(ctx context.Context, id uuid.UUID) error + + // ===== PromptTemplate ===== + InsertPromptTemplate(ctx context.Context, p InsertTemplateParams) (*PromptTemplate, error) + GetPromptTemplate(ctx context.Context, id uuid.UUID) (*PromptTemplate, error) + ListPromptTemplatesForUser(ctx context.Context, userID uuid.UUID) ([]PromptTemplate, error) + UpdatePromptTemplate(ctx context.Context, id uuid.UUID, name, content string) error + DeletePromptTemplate(ctx context.Context, id uuid.UUID) error + + // ===== LLMEndpoint ===== + InsertLLMEndpoint(ctx context.Context, p InsertEndpointParams) (*LLMEndpoint, error) + GetLLMEndpoint(ctx context.Context, id uuid.UUID) (*LLMEndpoint, error) + ListLLMEndpoints(ctx context.Context) ([]LLMEndpoint, error) + UpdateLLMEndpoint(ctx context.Context, id uuid.UUID, in UpdateEndpointInput) error + DeleteLLMEndpoint(ctx context.Context, id uuid.UUID) error + + // ===== LLMModel ===== + InsertLLMModel(ctx context.Context, p InsertModelParams) (*LLMModel, error) + GetLLMModel(ctx context.Context, id uuid.UUID) (*LLMModel, error) + GetTitleGeneratorModel(ctx context.Context) (*LLMModel, error) + ListLLMModels(ctx context.Context, onlyEnabled bool) ([]LLMModel, error) + UpdateLLMModel(ctx context.Context, id uuid.UUID, in UpdateModelInput) error + DeleteLLMModel(ctx context.Context, id uuid.UUID) error + + // ===== Usage ===== + InsertUsage(ctx context.Context, r UsageRecord) error + SummarizeUsageByUser(ctx context.Context, from, to time.Time) ([]UserUsageRow, error) + SummarizeUsageByModel(ctx context.Context, from, to time.Time) ([]ModelUsageRow, error) + SummarizeUsageByDay(ctx context.Context, from, to time.Time) ([]DayUsageRow, error) + SummarizeUserDaily(ctx context.Context, userID uuid.UUID, from, to time.Time) ([]DayUsageRow, error) + + // ===== Tx ===== + WithTx(ctx context.Context) (Repository, Tx, error) +} + +// Tx 是暴露给调用方的事务句柄;调用方必须在 defer 或 cleanup 中调用 Rollback, +// 成功后调用 Commit。 +type Tx interface { + Commit(ctx context.Context) error + Rollback(ctx context.Context) error +} + +// ===== Repo-layer input structs ===== + +// InsertConversationParams 是 InsertConversation 的参数。 +type InsertConversationParams struct { + ID uuid.UUID + UserID uuid.UUID + ProjectID *uuid.UUID + WorkspaceID *uuid.UUID + IssueID *uuid.UUID + ModelID uuid.UUID + SystemPrompt string + Title *string +} + +// UpdateMessageOKParams 是 UpdateMessageOK 的参数。 +type UpdateMessageOKParams struct { + ID int64 + Content string + Thinking *string + PromptTokens *int32 + CompletionTokens *int32 + ThinkingTokens *int32 +} + +// InsertAttachmentParams 是 InsertAttachment 的参数。 +type InsertAttachmentParams struct { + ID uuid.UUID + UserID uuid.UUID + Filename string + MimeType string + SizeBytes int64 + SHA256 string + StoragePath string +} + +// InsertTemplateParams 是 InsertPromptTemplate 的参数。 +type InsertTemplateParams struct { + ID uuid.UUID + Name string + Content string + Scope TemplateScope + OwnerID *uuid.UUID +} + +// InsertEndpointParams 是 InsertLLMEndpoint 的参数。 +type InsertEndpointParams struct { + ID uuid.UUID + Provider LLMProvider + DisplayName string + BaseURL string + APIKeyEncrypted []byte + CreatedBy uuid.UUID +} + +// InsertModelParams 是 InsertLLMModel 的参数。 +type InsertModelParams struct { + ID uuid.UUID + EndpointID uuid.UUID + ModelID string + DisplayName string + Capabilities ModelCapabilities + ContextWindow int + MaxOutputTokens int + PromptPrice float64 + CompletionPrice float64 + ThinkingPrice float64 + IsTitleGenerator bool + Enabled bool + SortOrder int +} + +// AttachmentRef は 孤立 attachment 清理时使用的轻量引用。 +type AttachmentRef struct { + ID uuid.UUID + StoragePath string +} + +// ===== compile-time assertion ===== +var _ Repository = (*pgRepo)(nil) + +// pgRepo 是 Repository 的 Postgres 实现。 +type pgRepo struct { + pool *pgxpool.Pool + tx pgx.Tx // 非 nil 时处于事务中 + q *chatsqlc.Queries +} + +// NewRepository 用 pgxpool 构造 Repository。 +func NewRepository(pool *pgxpool.Pool) Repository { + return &pgRepo{pool: pool, q: chatsqlc.New(pool)} +} + +// WithTx 开启一个 pgx 事务,返回绑定在该事务的新 Repository 和 Tx 句柄。 +func (r *pgRepo) WithTx(ctx context.Context) (Repository, Tx, error) { + tx, err := r.pool.Begin(ctx) + if err != nil { + return nil, nil, errs.Wrap(err, errs.CodeInternal, "begin tx") + } + return &pgRepo{pool: r.pool, tx: tx, q: chatsqlc.New(tx)}, tx, nil +} + +// ===== pgx.Tx 直接实现 Tx 接口 ===== +// pgx.Tx 已有 Commit/Rollback(ctx) error,满足接口要求。 + +// ===== 类型转换 helpers ===== + +func toPgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} } + +func fromPgUUID(p pgtype.UUID) uuid.UUID { + if !p.Valid { + return uuid.Nil + } + return uuid.UUID(p.Bytes) +} + +func toPgUUIDPtr(u *uuid.UUID) pgtype.UUID { + if u == nil { + return pgtype.UUID{Valid: false} + } + return pgtype.UUID{Bytes: *u, Valid: true} +} + +func fromPgUUIDPtr(p pgtype.UUID) *uuid.UUID { + if !p.Valid { + return nil + } + u := uuid.UUID(p.Bytes) + return &u +} + +func fromPgTime(t pgtype.Timestamptz) time.Time { + if !t.Valid { + return time.Time{} + } + return t.Time +} + +func fromPgTimePtr(t pgtype.Timestamptz) *time.Time { + if !t.Valid { + return nil + } + v := t.Time + return &v +} + +func toPgTimestamptz(t time.Time) pgtype.Timestamptz { + return pgtype.Timestamptz{Time: t, Valid: true} +} + +// numericToFloat64 は pgtype.Numeric → float64 変換。 +// Float64Value() は (Float8, error) を返す; Float8.Float64 を使う。 +func numericToFloat64(n pgtype.Numeric) float64 { + if !n.Valid { + return 0 + } + f8, err := n.Float64Value() + if err != nil { + return 0 + } + return f8.Float64 +} + +// float64ToNumeric は float64 → pgtype.Numeric 変換。 +// pgtype.Numeric は Int (big.Int) + Exp (int32) で表現するため、 +// fmt.Sprintf で文字列変換してから ScanScientific で読む。 +func float64ToNumeric(f float64) pgtype.Numeric { + var n pgtype.Numeric + if err := n.ScanScientific(fmt.Sprintf("%.10e", f)); err != nil { + // fallback: 整数部のみ + bi := new(big.Int) + bi.SetInt64(int64(f)) + return pgtype.Numeric{Int: bi, Exp: 0, Valid: true} + } + return n +} + +// intervalToTime は date_trunc('day',...) が pgtype.Interval として返ってくる +// sqlc 型推論ドリフトへの対処。Postgres は timestamptz を binary で +// 「2000-01-01 エポックからのマイクロ秒」として送るが、pgx が pgtype.Interval +// へスキャンした場合 Microseconds フィールドに入る。 +// Drift 注記: sqlc v1.31.1 は date_trunc の返り型を誤って interval と推論する。 +func intervalToTime(iv pgtype.Interval) time.Time { + if !iv.Valid { + return time.Time{} + } + // Microseconds はここでは「Postgres epoch からのマイクロ秒数」として解釈する。 + return pgEpoch.Add(time.Duration(iv.Microseconds) * time.Microsecond) +} + +// ===== PG エラー変換 ===== + +// translatePgError は PG 固有のエラーをドメイン AppError に変換する。 +// 対象外のエラーはそのまま返す。 +func translatePgError(err error) error { + var pgErr *pgconn.PgError + if !errors.As(err, &pgErr) { + return err + } + if pgErr.Code == pgerrcode.UniqueViolation { + switch pgErr.ConstraintName { + case "messages_one_pending_per_conv": + return errs.Wrap(err, errs.CodeChatConcurrentMessage, "another message is already streaming") + case "llm_models_one_title_generator": + return errs.Wrap(err, errs.CodeConflict, "another model is already the title generator") + } + } + return err +} + +// ===== Row → Domain converters ===== + +func convFromRow(r chatsqlc.Conversation) *Conversation { + return &Conversation{ + ID: fromPgUUID(r.ID), + UserID: fromPgUUID(r.UserID), + ProjectID: fromPgUUIDPtr(r.ProjectID), + WorkspaceID: fromPgUUIDPtr(r.WorkspaceID), + IssueID: fromPgUUIDPtr(r.IssueID), + ModelID: fromPgUUID(r.ModelID), + SystemPrompt: r.SystemPrompt, + Title: r.Title, + TitleGeneratedAt: fromPgTimePtr(r.TitleGeneratedAt), + CreatedAt: fromPgTime(r.CreatedAt), + UpdatedAt: fromPgTime(r.UpdatedAt), + DeletedAt: fromPgTimePtr(r.DeletedAt), + } +} + +func messageFromRow(r chatsqlc.Message) *Message { + var ptrs *int + var ctrs *int + var ttrs *int + if r.PromptTokens != nil { + v := int(*r.PromptTokens) + ptrs = &v + } + if r.CompletionTokens != nil { + v := int(*r.CompletionTokens) + ctrs = &v + } + if r.ThinkingTokens != nil { + v := int(*r.ThinkingTokens) + ttrs = &v + } + return &Message{ + ID: r.ID, + ConversationID: fromPgUUID(r.ConversationID), + Role: MessageRole(r.Role), + Content: r.Content, + Thinking: r.Thinking, + ToolCalls: r.ToolCalls, + Status: MessageStatus(r.Status), + ErrorMessage: r.ErrorMessage, + PromptTokens: ptrs, + CompletionTokens: ctrs, + ThinkingTokens: ttrs, + CreatedAt: fromPgTime(r.CreatedAt), + UpdatedAt: fromPgTime(r.UpdatedAt), + } +} + +func attachmentFromRow(r chatsqlc.MessageAttachment) *Attachment { + return &Attachment{ + ID: fromPgUUID(r.ID), + UserID: fromPgUUID(r.UserID), + MessageID: r.MessageID, + Filename: r.Filename, + MimeType: r.MimeType, + SizeBytes: r.SizeBytes, + SHA256: r.Sha256, + StoragePath: r.StoragePath, + CreatedAt: fromPgTime(r.CreatedAt), + } +} + +func templateFromRow(r chatsqlc.PromptTemplate) *PromptTemplate { + return &PromptTemplate{ + ID: fromPgUUID(r.ID), + Name: r.Name, + Content: r.Content, + Scope: TemplateScope(r.Scope), + OwnerID: fromPgUUIDPtr(r.OwnerID), + CreatedAt: fromPgTime(r.CreatedAt), + UpdatedAt: fromPgTime(r.UpdatedAt), + } +} + +func endpointFromRow(r chatsqlc.LlmEndpoint) *LLMEndpoint { + return &LLMEndpoint{ + ID: fromPgUUID(r.ID), + Provider: LLMProvider(r.Provider), + DisplayName: r.DisplayName, + BaseURL: r.BaseUrl, + APIKeyEncrypted: r.ApiKeyEncrypted, + CreatedBy: fromPgUUID(r.CreatedBy), + CreatedAt: fromPgTime(r.CreatedAt), + UpdatedAt: fromPgTime(r.UpdatedAt), + } +} + +func modelFromRow(r chatsqlc.LlmModel) *LLMModel { + var caps ModelCapabilities + if len(r.Capabilities) > 0 { + _ = json.Unmarshal(r.Capabilities, &caps) + } + return &LLMModel{ + ID: fromPgUUID(r.ID), + EndpointID: fromPgUUID(r.EndpointID), + ModelID: r.ModelID, + DisplayName: r.DisplayName, + Capabilities: caps, + ContextWindow: int(r.ContextWindow), + MaxOutputTokens: int(r.MaxOutputTokens), + PromptPricePerMillionUSD: numericToFloat64(r.PromptPricePerMillionUsd), + CompletionPricePerMillionUSD: numericToFloat64(r.CompletionPricePerMillionUsd), + ThinkingPricePerMillionUSD: numericToFloat64(r.ThinkingPricePerMillionUsd), + IsTitleGenerator: r.IsTitleGenerator, + Enabled: r.Enabled, + SortOrder: int(r.SortOrder), + CreatedAt: fromPgTime(r.CreatedAt), + UpdatedAt: fromPgTime(r.UpdatedAt), + } +} + +func usageFromRow(r chatsqlc.LlmUsage) *UsageRecord { + return &UsageRecord{ + ID: r.ID, + ConversationID: fromPgUUID(r.ConversationID), + MessageID: r.MessageID, + UserID: fromPgUUID(r.UserID), + EndpointID: fromPgUUID(r.EndpointID), + ModelID: fromPgUUID(r.ModelID), + PromptTokens: int(r.PromptTokens), + CompletionTokens: int(r.CompletionTokens), + ThinkingTokens: int(r.ThinkingTokens), + CostUSD: numericToFloat64(r.CostUsd), + CreatedAt: fromPgTime(r.CreatedAt), + } +} + +func userUsageRowFromRow(r chatsqlc.SummarizeUsageByUserRow) UserUsageRow { + return UserUsageRow{ + UserID: fromPgUUID(r.UserID), + PromptTokens: int(r.PromptTokens), + CompletionTokens: int(r.CompletionTokens), + ThinkingTokens: int(r.ThinkingTokens), + CostUSD: numericToFloat64(r.CostUsd), + } +} + +func modelUsageRowFromRow(r chatsqlc.SummarizeUsageByModelRow) ModelUsageRow { + return ModelUsageRow{ + ModelID: fromPgUUID(r.ModelID), + PromptTokens: int(r.PromptTokens), + CompletionTokens: int(r.CompletionTokens), + ThinkingTokens: int(r.ThinkingTokens), + CostUSD: numericToFloat64(r.CostUsd), + } +} + +func dayUsageRowFromRow(r chatsqlc.SummarizeUsageByDayRow) DayUsageRow { + return DayUsageRow{ + // Drift: sqlc 把 date_trunc('day', timestamptz) 推断为 pgtype.Interval。 + // pgx 在 binary 模式下把 timestamptz 值的 "Postgres epoch + microseconds" + // 送入 Interval.Microseconds 字段。用 pgEpoch + Duration 还原。 + Day: intervalToTime(r.Day), + PromptTokens: int(r.PromptTokens), + CompletionTokens: int(r.CompletionTokens), + ThinkingTokens: int(r.ThinkingTokens), + CostUSD: numericToFloat64(r.CostUsd), + } +} + +// ===== Conversation ===== + +func (r *pgRepo) InsertConversation(ctx context.Context, p InsertConversationParams) (*Conversation, error) { + row, err := r.q.InsertConversation(ctx, chatsqlc.InsertConversationParams{ + ID: toPgUUID(p.ID), + UserID: toPgUUID(p.UserID), + ProjectID: toPgUUIDPtr(p.ProjectID), + WorkspaceID: toPgUUIDPtr(p.WorkspaceID), + IssueID: toPgUUIDPtr(p.IssueID), + ModelID: toPgUUID(p.ModelID), + SystemPrompt: p.SystemPrompt, + Title: p.Title, + }) + if err != nil { + return nil, errs.Wrap(translatePgError(err), errs.CodeInternal, "insert conversation") + } + return convFromRow(row), nil +} + +func (r *pgRepo) GetConversation(ctx context.Context, id uuid.UUID) (*Conversation, error) { + row, err := r.q.GetConversation(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatConversationNotFound, "conversation not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get conversation") + } + return convFromRow(row), nil +} + +func (r *pgRepo) ListConversationsByUser(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error) { + limit := f.Limit + if limit <= 0 { + limit = 20 + } + rows, err := r.q.ListConversationsByUser(ctx, chatsqlc.ListConversationsByUserParams{ + UserID: toPgUUID(userID), + Column2: toPgUUIDPtr(f.ProjectID), + Column3: toPgUUIDPtr(f.WorkspaceID), + Column4: toPgUUIDPtr(f.IssueID), + Column5: f.TitleQuery, + Limit: int32(limit), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list conversations") + } + out := make([]Conversation, 0, len(rows)) + for _, row := range rows { + out = append(out, *convFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) SetConversationTitleManual(ctx context.Context, id uuid.UUID, title *string) error { + if err := r.q.SetConversationTitleManual(ctx, chatsqlc.SetConversationTitleManualParams{ + ID: toPgUUID(id), + Title: title, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "set conversation title manual") + } + return nil +} + +func (r *pgRepo) UpdateConversationTitle(ctx context.Context, id uuid.UUID, title *string) error { + if err := r.q.UpdateConversationTitle(ctx, chatsqlc.UpdateConversationTitleParams{ + ID: toPgUUID(id), + Title: title, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update conversation title") + } + return nil +} + +func (r *pgRepo) SoftDeleteConversation(ctx context.Context, id uuid.UUID) error { + if err := r.q.SoftDeleteConversation(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "soft delete conversation") + } + return nil +} + +// ===== Message ===== + +func (r *pgRepo) InsertMessage(ctx context.Context, conversationID uuid.UUID, role MessageRole, content string, status MessageStatus) (*Message, error) { + row, err := r.q.InsertMessage(ctx, chatsqlc.InsertMessageParams{ + ConversationID: toPgUUID(conversationID), + Role: string(role), + Content: content, + Status: string(status), + }) + if err != nil { + return nil, translatePgError(err) + } + return messageFromRow(row), nil +} + +func (r *pgRepo) GetMessage(ctx context.Context, id int64) (*Message, error) { + row, err := r.q.GetMessage(ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatMessageNotFound, "message not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get message") + } + return messageFromRow(row), nil +} + +func (r *pgRepo) GetFirstUserMessage(ctx context.Context, conversationID uuid.UUID) (*Message, error) { + row, err := r.q.GetFirstUserMessage(ctx, toPgUUID(conversationID)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatMessageNotFound, "no user message found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get first user message") + } + return messageFromRow(row), nil +} + +func (r *pgRepo) ListMessagesByConversation(ctx context.Context, conversationID uuid.UUID, beforeID *int64, limit int) ([]Message, error) { + var before int64 + if beforeID != nil { + before = *beforeID + } + if limit <= 0 { + limit = 50 + } + rows, err := r.q.ListMessagesByConversation(ctx, chatsqlc.ListMessagesByConversationParams{ + ConversationID: toPgUUID(conversationID), + Column2: before, + Limit: int32(limit), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list messages") + } + out := make([]Message, 0, len(rows)) + for _, row := range rows { + out = append(out, *messageFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) ListMessagesAscending(ctx context.Context, conversationID uuid.UUID) ([]Message, error) { + rows, err := r.q.ListMessagesAscending(ctx, toPgUUID(conversationID)) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list messages ascending") + } + out := make([]Message, 0, len(rows)) + for _, row := range rows { + out = append(out, *messageFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) UpdateMessageOK(ctx context.Context, p UpdateMessageOKParams) error { + if err := r.q.UpdateMessageOK(ctx, chatsqlc.UpdateMessageOKParams{ + ID: p.ID, + Content: p.Content, + Thinking: p.Thinking, + PromptTokens: p.PromptTokens, + CompletionTokens: p.CompletionTokens, + ThinkingTokens: p.ThinkingTokens, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update message ok") + } + return nil +} + +func (r *pgRepo) UpdateMessageError(ctx context.Context, id int64, errMsg *string) error { + if err := r.q.UpdateMessageError(ctx, chatsqlc.UpdateMessageErrorParams{ + ID: id, + ErrorMessage: errMsg, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update message error") + } + return nil +} + +func (r *pgRepo) UpdateMessageCancelled(ctx context.Context, id int64) error { + if err := r.q.UpdateMessageCancelled(ctx, id); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update message cancelled") + } + return nil +} + +func (r *pgRepo) MarkPendingAsErrorOnStartup(ctx context.Context) error { + if err := r.q.MarkPendingAsErrorOnStartup(ctx); err != nil { + return errs.Wrap(err, errs.CodeInternal, "mark pending as error on startup") + } + return nil +} + +func (r *pgRepo) DeleteMessage(ctx context.Context, id int64) error { + if err := r.q.DeleteMessage(ctx, id); err != nil { + return errs.Wrap(err, errs.CodeInternal, "delete message") + } + return nil +} + +// ===== Attachment ===== + +func (r *pgRepo) InsertAttachment(ctx context.Context, p InsertAttachmentParams) (*Attachment, error) { + row, err := r.q.InsertAttachment(ctx, chatsqlc.InsertAttachmentParams{ + ID: toPgUUID(p.ID), + UserID: toPgUUID(p.UserID), + Filename: p.Filename, + MimeType: p.MimeType, + SizeBytes: p.SizeBytes, + Sha256: p.SHA256, + StoragePath: p.StoragePath, + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "insert attachment") + } + return attachmentFromRow(row), nil +} + +func (r *pgRepo) GetAttachment(ctx context.Context, id uuid.UUID) (*Attachment, error) { + row, err := r.q.GetAttachment(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeNotFound, "attachment not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get attachment") + } + return attachmentFromRow(row), nil +} + +func (r *pgRepo) AttachToMessage(ctx context.Context, ids []uuid.UUID, messageID int64) error { + pgIDs := make([]pgtype.UUID, len(ids)) + for i, id := range ids { + pgIDs[i] = toPgUUID(id) + } + if err := r.q.AttachToMessage(ctx, chatsqlc.AttachToMessageParams{ + Column1: pgIDs, + MessageID: &messageID, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "attach to message") + } + return nil +} + +func (r *pgRepo) ListAttachmentsForMessage(ctx context.Context, messageID int64) ([]Attachment, error) { + rows, err := r.q.ListAttachmentsForMessage(ctx, &messageID) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list attachments for message") + } + out := make([]Attachment, 0, len(rows)) + for _, row := range rows { + out = append(out, *attachmentFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) ListOrphanAttachmentsForUser(ctx context.Context, userID uuid.UUID, ids []uuid.UUID) ([]Attachment, error) { + pgIDs := make([]pgtype.UUID, len(ids)) + for i, id := range ids { + pgIDs[i] = toPgUUID(id) + } + rows, err := r.q.ListOrphanAttachmentsForUser(ctx, chatsqlc.ListOrphanAttachmentsForUserParams{ + UserID: toPgUUID(userID), + Column2: pgIDs, + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list orphan attachments") + } + out := make([]Attachment, 0, len(rows)) + for _, row := range rows { + out = append(out, *attachmentFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) ListExpiredOrphans(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) { + rows, err := r.q.ListExpiredOrphans(ctx, pgtype.Interval{ + Microseconds: int64(olderThan / time.Microsecond), + Valid: true, + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list expired orphans") + } + out := make([]AttachmentRef, 0, len(rows)) + for _, row := range rows { + out = append(out, AttachmentRef{ + ID: fromPgUUID(row.ID), + StoragePath: row.StoragePath, + }) + } + return out, nil +} + +func (r *pgRepo) ListAttachmentsForSoftDeletedConversations(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) { + rows, err := r.q.ListAttachmentsForSoftDeletedConversations(ctx, pgtype.Interval{ + Microseconds: int64(olderThan / time.Microsecond), + Valid: true, + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list attachments for soft deleted") + } + out := make([]AttachmentRef, 0, len(rows)) + for _, row := range rows { + out = append(out, AttachmentRef{ + ID: fromPgUUID(row.ID), + StoragePath: row.StoragePath, + }) + } + return out, nil +} + +func (r *pgRepo) DeleteAttachment(ctx context.Context, id uuid.UUID) error { + if err := r.q.DeleteAttachment(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "delete attachment") + } + return nil +} + +// ===== PromptTemplate ===== + +func (r *pgRepo) InsertPromptTemplate(ctx context.Context, p InsertTemplateParams) (*PromptTemplate, error) { + row, err := r.q.InsertPromptTemplate(ctx, chatsqlc.InsertPromptTemplateParams{ + ID: toPgUUID(p.ID), + Name: p.Name, + Content: p.Content, + Scope: string(p.Scope), + OwnerID: toPgUUIDPtr(p.OwnerID), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "insert template") + } + return templateFromRow(row), nil +} + +func (r *pgRepo) GetPromptTemplate(ctx context.Context, id uuid.UUID) (*PromptTemplate, error) { + row, err := r.q.GetPromptTemplate(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatTemplateNotFound, "template not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get template") + } + return templateFromRow(row), nil +} + +func (r *pgRepo) ListPromptTemplatesForUser(ctx context.Context, userID uuid.UUID) ([]PromptTemplate, error) { + rows, err := r.q.ListPromptTemplatesForUser(ctx, toPgUUID(userID)) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list templates") + } + out := make([]PromptTemplate, 0, len(rows)) + for _, row := range rows { + out = append(out, *templateFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) UpdatePromptTemplate(ctx context.Context, id uuid.UUID, name, content string) error { + if err := r.q.UpdatePromptTemplate(ctx, chatsqlc.UpdatePromptTemplateParams{ + ID: toPgUUID(id), + Name: name, + Content: content, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update template") + } + return nil +} + +func (r *pgRepo) DeletePromptTemplate(ctx context.Context, id uuid.UUID) error { + if err := r.q.DeletePromptTemplate(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "delete template") + } + return nil +} + +// ===== LLMEndpoint ===== + +func (r *pgRepo) InsertLLMEndpoint(ctx context.Context, p InsertEndpointParams) (*LLMEndpoint, error) { + row, err := r.q.InsertLLMEndpoint(ctx, chatsqlc.InsertLLMEndpointParams{ + ID: toPgUUID(p.ID), + Provider: string(p.Provider), + DisplayName: p.DisplayName, + BaseUrl: p.BaseURL, + ApiKeyEncrypted: p.APIKeyEncrypted, + CreatedBy: toPgUUID(p.CreatedBy), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "insert endpoint") + } + return endpointFromRow(row), nil +} + +func (r *pgRepo) GetLLMEndpoint(ctx context.Context, id uuid.UUID) (*LLMEndpoint, error) { + row, err := r.q.GetLLMEndpoint(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatEndpointNotFound, "endpoint not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get endpoint") + } + return endpointFromRow(row), nil +} + +func (r *pgRepo) ListLLMEndpoints(ctx context.Context) ([]LLMEndpoint, error) { + rows, err := r.q.ListLLMEndpoints(ctx) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list endpoints") + } + out := make([]LLMEndpoint, 0, len(rows)) + for _, row := range rows { + out = append(out, *endpointFromRow(row)) + } + return out, nil +} + +// UpdateLLMEndpoint は UpdateEndpointInput に基づいて既存エンドポイントを更新する。 +// nil フィールドは変更なし(先に Get して現在値でマージする)。 +func (r *pgRepo) UpdateLLMEndpoint(ctx context.Context, id uuid.UUID, in UpdateEndpointInput) error { + existing, err := r.GetLLMEndpoint(ctx, id) + if err != nil { + return err + } + displayName := existing.DisplayName + if in.DisplayName != nil { + displayName = *in.DisplayName + } + baseURL := existing.BaseURL + if in.BaseURL != nil { + baseURL = *in.BaseURL + } + apiKey := existing.APIKeyEncrypted + if in.APIKey != nil { + // in.APIKey は明文;service 層で暗号化後にここに渡す想定だが、 + // repository 層では []byte として受け取る。UpdateEndpointInput.APIKey は + // *string(明文)なので、service 層で変換済みのものが渡る。 + // ここでは文字列をバイト列に変換するだけ。 + apiKey = []byte(*in.APIKey) + } + if err := r.q.UpdateLLMEndpoint(ctx, chatsqlc.UpdateLLMEndpointParams{ + ID: toPgUUID(id), + DisplayName: displayName, + BaseUrl: baseURL, + ApiKeyEncrypted: apiKey, + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "update endpoint") + } + return nil +} + +func (r *pgRepo) DeleteLLMEndpoint(ctx context.Context, id uuid.UUID) error { + if err := r.q.DeleteLLMEndpoint(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "delete endpoint") + } + return nil +} + +// ===== LLMModel ===== + +func (r *pgRepo) InsertLLMModel(ctx context.Context, p InsertModelParams) (*LLMModel, error) { + capsBytes, err := json.Marshal(p.Capabilities) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "marshal capabilities") + } + row, err := r.q.InsertLLMModel(ctx, chatsqlc.InsertLLMModelParams{ + ID: toPgUUID(p.ID), + EndpointID: toPgUUID(p.EndpointID), + ModelID: p.ModelID, + DisplayName: p.DisplayName, + Capabilities: capsBytes, + ContextWindow: int32(p.ContextWindow), + MaxOutputTokens: int32(p.MaxOutputTokens), + PromptPricePerMillionUsd: float64ToNumeric(p.PromptPrice), + CompletionPricePerMillionUsd: float64ToNumeric(p.CompletionPrice), + ThinkingPricePerMillionUsd: float64ToNumeric(p.ThinkingPrice), + IsTitleGenerator: p.IsTitleGenerator, + Enabled: p.Enabled, + SortOrder: int32(p.SortOrder), + }) + if err != nil { + return nil, translatePgError(err) + } + return modelFromRow(row), nil +} + +func (r *pgRepo) GetLLMModel(ctx context.Context, id uuid.UUID) (*LLMModel, error) { + row, err := r.q.GetLLMModel(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatModelNotFound, "model not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get model") + } + return modelFromRow(row), nil +} + +func (r *pgRepo) GetTitleGeneratorModel(ctx context.Context) (*LLMModel, error) { + row, err := r.q.GetTitleGeneratorModel(ctx) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeChatModelNotFound, "no title generator model found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "get title generator model") + } + return modelFromRow(row), nil +} + +func (r *pgRepo) ListLLMModels(ctx context.Context, onlyEnabled bool) ([]LLMModel, error) { + rows, err := r.q.ListLLMModels(ctx, onlyEnabled) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list models") + } + out := make([]LLMModel, 0, len(rows)) + for _, row := range rows { + out = append(out, *modelFromRow(row)) + } + return out, nil +} + +// UpdateLLMModel は UpdateModelInput に基づいて既存モデルを更新する。 +// nil フィールドは変更なし(先に Get して現在値でマージする)。 +func (r *pgRepo) UpdateLLMModel(ctx context.Context, id uuid.UUID, in UpdateModelInput) error { + existing, err := r.GetLLMModel(ctx, id) + if err != nil { + return err + } + + displayName := existing.DisplayName + if in.DisplayName != nil { + displayName = *in.DisplayName + } + caps := existing.Capabilities + if in.Capabilities != nil { + caps = *in.Capabilities + } + contextWindow := existing.ContextWindow + if in.ContextWindow != nil { + contextWindow = *in.ContextWindow + } + maxOutputTokens := existing.MaxOutputTokens + if in.MaxOutputTokens != nil { + maxOutputTokens = *in.MaxOutputTokens + } + promptPrice := existing.PromptPricePerMillionUSD + if in.PromptPrice != nil { + promptPrice = *in.PromptPrice + } + completionPrice := existing.CompletionPricePerMillionUSD + if in.CompletionPrice != nil { + completionPrice = *in.CompletionPrice + } + thinkingPrice := existing.ThinkingPricePerMillionUSD + if in.ThinkingPrice != nil { + thinkingPrice = *in.ThinkingPrice + } + isTitleGenerator := existing.IsTitleGenerator + if in.IsTitleGenerator != nil { + isTitleGenerator = *in.IsTitleGenerator + } + enabled := existing.Enabled + if in.Enabled != nil { + enabled = *in.Enabled + } + sortOrder := existing.SortOrder + if in.SortOrder != nil { + sortOrder = *in.SortOrder + } + + capsBytes, err := json.Marshal(caps) + if err != nil { + return errs.Wrap(err, errs.CodeInternal, "marshal capabilities") + } + + if err := r.q.UpdateLLMModel(ctx, chatsqlc.UpdateLLMModelParams{ + ID: toPgUUID(id), + DisplayName: displayName, + Capabilities: capsBytes, + ContextWindow: int32(contextWindow), + MaxOutputTokens: int32(maxOutputTokens), + PromptPricePerMillionUsd: float64ToNumeric(promptPrice), + CompletionPricePerMillionUsd: float64ToNumeric(completionPrice), + ThinkingPricePerMillionUsd: float64ToNumeric(thinkingPrice), + IsTitleGenerator: isTitleGenerator, + Enabled: enabled, + SortOrder: int32(sortOrder), + }); err != nil { + return translatePgError(err) + } + return nil +} + +func (r *pgRepo) DeleteLLMModel(ctx context.Context, id uuid.UUID) error { + if err := r.q.DeleteLLMModel(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "delete model") + } + return nil +} + +// ===== Usage ===== + +func (r *pgRepo) InsertUsage(ctx context.Context, rec UsageRecord) error { + if err := r.q.InsertUsage(ctx, chatsqlc.InsertUsageParams{ + ConversationID: toPgUUID(rec.ConversationID), + MessageID: rec.MessageID, + UserID: toPgUUID(rec.UserID), + EndpointID: toPgUUID(rec.EndpointID), + ModelID: toPgUUID(rec.ModelID), + PromptTokens: int32(rec.PromptTokens), + CompletionTokens: int32(rec.CompletionTokens), + ThinkingTokens: int32(rec.ThinkingTokens), + CostUsd: float64ToNumeric(rec.CostUSD), + }); err != nil { + return errs.Wrap(err, errs.CodeInternal, "insert usage") + } + return nil +} + +func (r *pgRepo) SummarizeUsageByUser(ctx context.Context, from, to time.Time) ([]UserUsageRow, error) { + rows, err := r.q.SummarizeUsageByUser(ctx, chatsqlc.SummarizeUsageByUserParams{ + CreatedAt: toPgTimestamptz(from), + CreatedAt_2: toPgTimestamptz(to), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by user") + } + out := make([]UserUsageRow, 0, len(rows)) + for _, row := range rows { + out = append(out, userUsageRowFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) SummarizeUsageByModel(ctx context.Context, from, to time.Time) ([]ModelUsageRow, error) { + rows, err := r.q.SummarizeUsageByModel(ctx, chatsqlc.SummarizeUsageByModelParams{ + CreatedAt: toPgTimestamptz(from), + CreatedAt_2: toPgTimestamptz(to), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by model") + } + out := make([]ModelUsageRow, 0, len(rows)) + for _, row := range rows { + out = append(out, modelUsageRowFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) SummarizeUsageByDay(ctx context.Context, from, to time.Time) ([]DayUsageRow, error) { + rows, err := r.q.SummarizeUsageByDay(ctx, chatsqlc.SummarizeUsageByDayParams{ + CreatedAt: toPgTimestamptz(from), + CreatedAt_2: toPgTimestamptz(to), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by day") + } + out := make([]DayUsageRow, 0, len(rows)) + for _, row := range rows { + out = append(out, dayUsageRowFromRow(row)) + } + return out, nil +} + +func (r *pgRepo) SummarizeUserDaily(ctx context.Context, userID uuid.UUID, from, to time.Time) ([]DayUsageRow, error) { + rows, err := r.q.SummarizeUserDaily(ctx, chatsqlc.SummarizeUserDailyParams{ + UserID: toPgUUID(userID), + CreatedAt: toPgTimestamptz(from), + CreatedAt_2: toPgTimestamptz(to), + }) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "summarize user daily") + } + out := make([]DayUsageRow, 0, len(rows)) + for _, row := range rows { + out = append(out, DayUsageRow{ + Day: intervalToTime(row.Day), + PromptTokens: int(row.PromptTokens), + CompletionTokens: int(row.CompletionTokens), + ThinkingTokens: int(row.ThinkingTokens), + CostUSD: numericToFloat64(row.CostUsd), + }) + } + return out, nil +} diff --git a/internal/chat/repository_test.go b/internal/chat/repository_test.go new file mode 100644 index 0000000..ff9e7e6 --- /dev/null +++ b/internal/chat/repository_test.go @@ -0,0 +1,339 @@ +//go:build integration + +package chat + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" + tcpg "github.com/testcontainers/testcontainers-go/modules/postgres" + + "github.com/yan1h/agent-coding-workflow/internal/infra/db" + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" +) + +// setupRepo 起一个临时 postgres 容器、跑迁移、注入 user / llm_endpoint / llm_model +// 行作外键基础,返回 Repository 与 pool、uid、eid、mid。 +func setupRepo(t *testing.T) (Repository, *pgxpool.Pool, uuid.UUID, uuid.UUID, uuid.UUID) { + t.Helper() + ctx := context.Background() + container, err := tcpg.Run(ctx, "postgres:16-alpine", + tcpg.WithDatabase("acw"), tcpg.WithUsername("acw"), tcpg.WithPassword("acw")) + require.NoError(t, err) + t.Cleanup(func() { _ = container.Terminate(ctx) }) + + dsn, err := container.ConnectionString(ctx, "sslmode=disable") + require.NoError(t, err) + require.NoError(t, db.Migrate(ctx, dsn, "file://../../migrations")) + pool, err := db.NewPool(ctx, dsn) + require.NoError(t, err) + t.Cleanup(pool.Close) + + uid := uuid.New() + _, err = pool.Exec(ctx, `INSERT INTO users (id,email,password_hash,display_name,is_admin,created_at,updated_at) + VALUES ($1,'a@b.c','x','admin',true,now(),now())`, uid) + require.NoError(t, err) + eid := uuid.New() + _, err = pool.Exec(ctx, `INSERT INTO llm_endpoints (id,provider,display_name,base_url,api_key_encrypted,created_by) + VALUES ($1,'anthropic','test','https://example.com','\x',$2)`, eid, uid) + require.NoError(t, err) + mid := uuid.New() + _, err = pool.Exec(ctx, `INSERT INTO llm_models (id,endpoint_id,model_id,display_name,capabilities,context_window,max_output_tokens) + VALUES ($1,$2,'claude','Claude','{}',200000,8192)`, mid, eid) + require.NoError(t, err) + + return NewRepository(pool), pool, uid, eid, mid +} + +// ===== 1. Conversation CRUD ===== + +func TestRepo_Conversation_CRUD(t *testing.T) { + t.Parallel() + repo, _, uid, _, mid := setupRepo(t) + ctx := context.Background() + + cid := uuid.New() + conv, err := repo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "be helpful", + }) + require.NoError(t, err) + require.Equal(t, cid, conv.ID) + require.Equal(t, "be helpful", conv.SystemPrompt) + + got, err := repo.GetConversation(ctx, cid) + require.NoError(t, err) + require.Equal(t, cid, got.ID) + + require.NoError(t, repo.SoftDeleteConversation(ctx, cid)) + + _, err = repo.GetConversation(ctx, cid) + var ae *errs.AppError + require.ErrorAs(t, err, &ae) + require.Equal(t, errs.CodeChatConversationNotFound, ae.Code) +} + +// ===== 2. Messages: pending unique index ===== + +func TestRepo_Messages_PendingUniqueIndex(t *testing.T) { + t.Parallel() + repo, _, uid, _, mid := setupRepo(t) + ctx := context.Background() + + cid := uuid.New() + _, err := repo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", + }) + require.NoError(t, err) + + // first pending assistant message — should succeed + _, err = repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) + require.NoError(t, err) + + // second pending assistant message in same conv — must fail with CodeChatConcurrentMessage + _, err = repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) + require.Error(t, err) + var ae *errs.AppError + require.ErrorAs(t, err, &ae) + require.Equal(t, errs.CodeChatConcurrentMessage, ae.Code) + require.True(t, strings.Contains(ae.Message, "another message is already streaming")) +} + +// ===== 3. Message status transitions ===== + +func TestRepo_Messages_StatusTransitions(t *testing.T) { + t.Parallel() + repo, _, uid, _, mid := setupRepo(t) + ctx := context.Background() + + cid := uuid.New() + _, err := repo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", + }) + require.NoError(t, err) + + msg, err := repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) + require.NoError(t, err) + require.Equal(t, StatusPending, msg.Status) + + pt := int32(10) + ct := int32(20) + require.NoError(t, repo.UpdateMessageOK(ctx, UpdateMessageOKParams{ + ID: msg.ID, + Content: "hello", + PromptTokens: &pt, + CompletionTokens: &ct, + })) + + updated, err := repo.GetMessage(ctx, msg.ID) + require.NoError(t, err) + require.Equal(t, StatusOK, updated.Status) + require.Equal(t, "hello", updated.Content) + require.NotNil(t, updated.PromptTokens) + require.Equal(t, 10, *updated.PromptTokens) +} + +// ===== 4. Attachments orphan lifecycle ===== + +func TestRepo_Attachments_OrphanLifecycle(t *testing.T) { + t.Parallel() + repo, _, uid, _, mid := setupRepo(t) + ctx := context.Background() + + aid := uuid.New() + att, err := repo.InsertAttachment(ctx, InsertAttachmentParams{ + ID: aid, UserID: uid, + Filename: "file.txt", MimeType: "text/plain", + SizeBytes: 42, SHA256: "abc", StoragePath: "/tmp/file", + }) + require.NoError(t, err) + require.Nil(t, att.MessageID, "newly inserted attachment must be orphan") + + // create conversation + message to attach to + cid := uuid.New() + _, err = repo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", + }) + require.NoError(t, err) + msg, err := repo.InsertMessage(ctx, cid, RoleUser, "hello", StatusOK) + require.NoError(t, err) + + require.NoError(t, repo.AttachToMessage(ctx, []uuid.UUID{aid}, msg.ID)) + + atts, err := repo.ListAttachmentsForMessage(ctx, msg.ID) + require.NoError(t, err) + require.Len(t, atts, 1) + require.Equal(t, aid, atts[0].ID) + require.NotNil(t, atts[0].MessageID) + require.Equal(t, msg.ID, *atts[0].MessageID) +} + +// ===== 5. Templates scope filtering ===== + +func TestRepo_Templates_ScopeFiltering(t *testing.T) { + t.Parallel() + repo, pool, uid, _, _ := setupRepo(t) + ctx := context.Background() + + // seed user B + uidB := uuid.New() + _, err := pool.Exec(ctx, `INSERT INTO users (id,email,password_hash,display_name,is_admin,created_at,updated_at) + VALUES ($1,'b@b.c','x','B',false,now(),now())`, uidB) + require.NoError(t, err) + + // system template (no owner) + _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ + ID: uuid.New(), Name: "sys", Content: "system", Scope: TemplateScopeSystem, + }) + require.NoError(t, err) + + // user A template + _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ + ID: uuid.New(), Name: "tmpl-a", Content: "for a", Scope: TemplateScopeUser, OwnerID: &uid, + }) + require.NoError(t, err) + + // user B template + _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ + ID: uuid.New(), Name: "tmpl-b", Content: "for b", Scope: TemplateScopeUser, OwnerID: &uidB, + }) + require.NoError(t, err) + + templates, err := repo.ListPromptTemplatesForUser(ctx, uid) + require.NoError(t, err) + require.Len(t, templates, 2, "user A should see system + own template only") + names := make([]string, len(templates)) + for i, tmpl := range templates { + names[i] = tmpl.Name + } + require.Contains(t, names, "sys") + require.Contains(t, names, "tmpl-a") + require.NotContains(t, names, "tmpl-b") +} + +// ===== 6. Endpoints CRUD ===== + +func TestRepo_Endpoints_CRUD(t *testing.T) { + t.Parallel() + repo, _, uid, _, _ := setupRepo(t) + ctx := context.Background() + + eid := uuid.New() + ep, err := repo.InsertLLMEndpoint(ctx, InsertEndpointParams{ + ID: eid, Provider: ProviderOpenAI, DisplayName: "OpenAI", + BaseURL: "https://api.openai.com", APIKeyEncrypted: []byte("key"), CreatedBy: uid, + }) + require.NoError(t, err) + require.Equal(t, eid, ep.ID) + require.Equal(t, ProviderOpenAI, ep.Provider) + + newName := "OpenAI Updated" + require.NoError(t, repo.UpdateLLMEndpoint(ctx, eid, UpdateEndpointInput{ + DisplayName: &newName, + })) + got, err := repo.GetLLMEndpoint(ctx, eid) + require.NoError(t, err) + require.Equal(t, "OpenAI Updated", got.DisplayName) + + require.NoError(t, repo.DeleteLLMEndpoint(ctx, eid)) + _, err = repo.GetLLMEndpoint(ctx, eid) + var ae *errs.AppError + require.ErrorAs(t, err, &ae) + require.Equal(t, errs.CodeChatEndpointNotFound, ae.Code) +} + +// ===== 7. Models: title generator unique constraint ===== + +func TestRepo_Models_TitleGeneratorUnique(t *testing.T) { + t.Parallel() + repo, _, _, eid, _ := setupRepo(t) + ctx := context.Background() + + // first title-generator model + _, err := repo.InsertLLMModel(ctx, InsertModelParams{ + ID: uuid.New(), EndpointID: eid, ModelID: "m1", DisplayName: "M1", + Capabilities: ModelCapabilities{}, ContextWindow: 4096, MaxOutputTokens: 1024, + IsTitleGenerator: true, Enabled: true, + }) + require.NoError(t, err) + + // second title-generator model — must fail with CodeConflict + _, err = repo.InsertLLMModel(ctx, InsertModelParams{ + ID: uuid.New(), EndpointID: eid, ModelID: "m2", DisplayName: "M2", + Capabilities: ModelCapabilities{}, ContextWindow: 4096, MaxOutputTokens: 1024, + IsTitleGenerator: true, Enabled: true, + }) + require.Error(t, err) + var ae *errs.AppError + require.ErrorAs(t, err, &ae) + require.Equal(t, errs.CodeConflict, ae.Code) +} + +// ===== 8. Usage summarize by day ===== + +func TestRepo_Usage_SummarizeByDay(t *testing.T) { + t.Parallel() + repo, _, uid, eid, mid := setupRepo(t) + ctx := context.Background() + + cid := uuid.New() + _, err := repo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", + }) + require.NoError(t, err) + msg, err := repo.InsertMessage(ctx, cid, RoleUser, "hi", StatusOK) + require.NoError(t, err) + + // insert 3 usage records + for i := 0; i < 3; i++ { + require.NoError(t, repo.InsertUsage(ctx, UsageRecord{ + ConversationID: cid, MessageID: msg.ID, UserID: uid, + EndpointID: eid, ModelID: mid, + PromptTokens: 10, CompletionTokens: 20, ThinkingTokens: 0, + CostUSD: 0.001, + })) + } + + from := time.Now().UTC().Truncate(24 * time.Hour) + to := from.Add(24 * time.Hour) + rows, err := repo.SummarizeUsageByDay(ctx, from, to) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, 30, rows[0].PromptTokens) + require.Equal(t, 60, rows[0].CompletionTokens) +} + +// ===== 9. Tx rollback undoes insert ===== + +func TestRepo_Tx_RollbackUndoesInsert(t *testing.T) { + t.Parallel() + repo, _, uid, _, mid := setupRepo(t) + ctx := context.Background() + + txRepo, tx, err := repo.WithTx(ctx) + require.NoError(t, err) + + cid := uuid.New() + _, err = txRepo.InsertConversation(ctx, InsertConversationParams{ + ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "rollback me", + }) + require.NoError(t, err) + + // can read within same tx + got, err := txRepo.GetConversation(ctx, cid) + require.NoError(t, err) + require.Equal(t, cid, got.ID) + + // rollback + require.NoError(t, tx.Rollback(ctx)) + + // after rollback, reading via original repo should return not found + _, err = repo.GetConversation(ctx, cid) + var ae *errs.AppError + require.ErrorAs(t, err, &ae) + require.Equal(t, errs.CodeChatConversationNotFound, ae.Code) +}