diff --git a/build.sh b/build.sh index 67b93270a3..a07b1fbc55 100755 --- a/build.sh +++ b/build.sh @@ -155,7 +155,7 @@ check_office_oxide_deps() { if [ ! -f "$lib_path" ] || [ ! -f "$header_path" ]; then echo -e "${RED}Error: office_oxide native library not found${NC}" echo " Expected: ${lib_path}" - echo " Run: uv run python3 ragflow_deps/download_deps.py" + echo " Run: uv run python3 ragflow_deps/download_go_deps.py" echo " Or manually download: https://github.com/yfedoseev/office_oxide/releases/download/v${OFFICE_OXIDE_VERSION}/native-linux-x86_64.tar.gz" exit 1 fi @@ -171,7 +171,7 @@ check_office_oxide_deps() { echo " Required: v${OFFICE_OXIDE_VERSION}; found: ${found_version:-unknown}" echo " A stale lib silently loses PPT97 (.ppt) slide content. Refresh:" echo " rm -rf ~/ragflow-native-libs/office_oxide ragflow_deps/office_oxide-linux-x86_64.tar.gz" - echo " uv run python3 ragflow_deps/download_deps.py" + echo " uv run python3 ragflow_deps/download_go_deps.py" exit 1 fi @@ -191,7 +191,7 @@ check_pdfium_deps() { echo " pdfium (static) not found" echo " Expected: ${lib_path}" - echo " Run: uv run python3 ragflow_deps/download_deps.py" + echo " Run: uv run python3 ragflow_deps/download_go_deps.py" echo " Or: curl -fsSL https://github.com/kognitos/pdfium-static/releases/download/chromium%2F${PDFIUM_STATIC_VERSION}/pdfium-linux-x64-static.tgz | tar xz -C ${PDFIUM_STATIC_PREFIX}" return 1 } @@ -228,7 +228,7 @@ check_pdf_oxide_deps() { echo " pdf_oxide (static) not found" echo " Expected: ${lib_path}" - echo " Run: uv run python3 ragflow_deps/download_deps.py" + echo " Run: uv run python3 ragflow_deps/download_go_deps.py" echo " Or: curl -fsSL https://github.com/yfedoseev/pdf_oxide/releases/download/v${PDF_OXIDE_VERSION}/pdf_oxide-go-ffi-linux-amd64.tar.gz | tar xz -C ${PDF_OXIDE_PREFIX}" return 1 } @@ -510,7 +510,7 @@ DEPENDENCIES: - cmake >= 4.0 - go >= 1.24 - g++ with C++17/23 support - - office_oxide native library (download with: uv run python3 ragflow_deps/download_deps.py) + - office_oxide native library (download with: uv run python3 ragflow_deps/download_go_deps.py) - lld (Linux only): sudo apt install lld-20 && sudo ln -s /usr/bin/ld.lld-20 /usr/bin/ld.lld - pcre2 development files - Debian/Ubuntu: libpcre2-dev diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 4af6bf56ee..dc8e020f1d 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -212,6 +212,9 @@ func printHelp(args *serverArgs) { } func main() { + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGUSR2) + defer cancel() + arguments, err := parseArgs() if err != nil { fmt.Printf("Failed to parse arguments: %v\n", err) @@ -340,7 +343,7 @@ func main() { server.PrintAll() // Initialize database - if err = dao.InitDB(arguments.migrateDB); err != nil { + if err = dao.InitDB(ctx, arguments.migrateDB); err != nil { common.Fatal("Failed to initialize database", zap.Error(err)) } @@ -371,7 +374,6 @@ func main() { common.Warn("Failed to initialize server variables from Redis, using defaults", zap.String("error", err.Error())) } - ctx, cancel := context.WithCancel(context.Background()) if err = server.StartServer(ctx, cancel, serverName); err != nil { common.Error("Failed to start EE server", err) os.Exit(1) @@ -384,7 +386,7 @@ func main() { switch *arguments.mode { case "api": - if err = runAPI(arguments); err != nil { + if err = runAPI(ctx, arguments); err != nil { fmt.Printf("Failed to start API server: %v\n", err) os.Exit(1) } @@ -478,11 +480,11 @@ func runAdmin(args *serverArgs) error { common.Info("Shutting down RAGFlow HTTP server...") // Create context with timeout for graceful shutdown - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() + quitCtx, quitCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer quitCancel() // Shutdown HTTP server - if err := srv.Shutdown(ctx); err != nil { + if err := srv.Shutdown(quitCtx); err != nil { common.Fatal("Server forced to shutdown", zap.Error(err)) } @@ -646,7 +648,7 @@ func runSyncer(args *serverArgs) error { return nil } -func runAPI(args *serverArgs) error { +func runAPI(ctx context.Context, args *serverArgs) error { // Initialize admin status (default: unavailable=1) local.InitAdminStatus(1, "admin server not connected") @@ -666,14 +668,14 @@ func runAPI(args *serverArgs) error { } config := server.GetConfig() - startServer(config) + startServer(ctx, config) common.Info("Server exited") return nil } -func startServer(config *server.Config) { +func startServer(ctx context.Context, config *server.Config) { // Set Gin mode if config.Server.Mode == "release" { @@ -744,7 +746,7 @@ func startServer(config *server.Config) { return handler.MCPListDatasets(datasetsService, userID, page, pageSize, orderBy, desc) }, func(userID string, page, pageSize int, orderBy string, desc bool) ([]map[string]interface{}, int64, error) { - return handler.MCPListChats(chatService, userID, page, pageSize, orderBy, desc) + return handler.MCPListChats(ctx, chatService, userID, page, pageSize, orderBy, desc) }, func(userID string, req mcp.RetrievalRequest) (string, error) { return handler.MCPRetrieval(datasetsService, userID, req) diff --git a/internal/admin/handler.go b/internal/admin/handler.go index d0835febdd..ff7ffaa78d 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -432,7 +432,8 @@ func (h *Handler) ListUserAPITokens(c *gin.Context) { return } - apiKeys, err := h.service.ListUserAPITokens(username) + ctx := c.Request.Context() + apiKeys, err := h.service.ListUserAPITokens(ctx, username) if err != nil { common.ErrorWithCode(c, common.CodeServerError, err.Error()) return @@ -448,7 +449,8 @@ func (h *Handler) GenerateUserAPIToken(c *gin.Context) { return } - apiKey, err := h.service.GenerateUserAPIToken(username) + ctx := c.Request.Context() + apiKey, err := h.service.GenerateUserAPIToken(ctx, username) if err != nil { common.ErrorWithCode(c, common.CodeServerError, err.Error()) return @@ -471,7 +473,8 @@ func (h *Handler) DeleteUserAPIToken(c *gin.Context) { return } - if err = h.service.DeleteUserAPIToken(username, key); err != nil { + ctx := c.Request.Context() + if err = h.service.DeleteUserAPIToken(ctx, username, key); err != nil { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } diff --git a/internal/admin/service.go b/internal/admin/service.go index c3ded9888c..07a7e23123 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -17,6 +17,7 @@ package admin import ( + "context" "crypto/rand" "crypto/tls" "encoding/base64" @@ -840,7 +841,7 @@ func (s *Service) GetUserAgents(username string) ([]map[string]interface{}, erro // API Key methods // ListUserAPITokens get user API keys -func (s *Service) ListUserAPITokens(username string) ([]map[string]interface{}, error) { +func (s *Service) ListUserAPITokens(ctx context.Context, username string) ([]map[string]interface{}, error) { // 1. Get user details user, err := s.userDAO.GetByEmail(username) if err != nil { @@ -856,7 +857,7 @@ func (s *Service) ListUserAPITokens(username string) ([]map[string]interface{}, tenantID := userTenants[0].TenantID // 3. Get API tokens by tenant ID - tokens, err := s.apiTokenDAO.GetByTenantID(tenantID) + tokens, err := s.apiTokenDAO.GetByTenantID(ctx, tenantID) if err != nil { return nil, fmt.Errorf("failed to get API tokens: %w", err) } @@ -881,7 +882,7 @@ func (s *Service) ListUserAPITokens(username string) ([]map[string]interface{}, } // GenerateUserAPIToken generate API key for user -func (s *Service) GenerateUserAPIToken(username string) (map[string]interface{}, error) { +func (s *Service) GenerateUserAPIToken(ctx context.Context, username string) (map[string]interface{}, error) { // 1. Get user details user, err := s.userDAO.GetByEmail(username) if err != nil { @@ -907,7 +908,7 @@ func (s *Service) GenerateUserAPIToken(username string) (map[string]interface{}, } // 4. Save API token - if err = s.apiTokenDAO.Create(apiToken); err != nil { + if err = s.apiTokenDAO.Create(ctx, apiToken); err != nil { return nil, fmt.Errorf("failed to generate API key: %w", err) } @@ -923,7 +924,7 @@ func (s *Service) GenerateUserAPIToken(username string) (map[string]interface{}, } // DeleteUserAPIToken delete user API key -func (s *Service) DeleteUserAPIToken(username, key string) error { +func (s *Service) DeleteUserAPIToken(ctx context.Context, username, key string) error { // 1. Get user details user, err := s.userDAO.GetByEmail(username) if err != nil { @@ -939,7 +940,7 @@ func (s *Service) DeleteUserAPIToken(username, key string) error { tenantID := userTenants[0].TenantID // 3. Delete API token - rowsAffected, err := s.apiTokenDAO.DeleteByTenantIDAndToken(tenantID, key) + rowsAffected, err := s.apiTokenDAO.DeleteByTenantIDAndToken(ctx, tenantID, key) if err != nil { return fmt.Errorf("failed to delete API key: %w", err) } diff --git a/internal/dao/api_token.go b/internal/dao/api_token.go index 54a1eb8bc2..837235e428 100644 --- a/internal/dao/api_token.go +++ b/internal/dao/api_token.go @@ -17,6 +17,7 @@ package dao import ( + "context" "errors" "ragflow/internal/entity" @@ -31,27 +32,27 @@ func NewAPITokenDAO() *APITokenDAO { } // Create creates a new API token -func (dao *APITokenDAO) Create(apiToken *entity.APIToken) error { - return DB.Create(apiToken).Error +func (dao *APITokenDAO) Create(ctx context.Context, apiToken *entity.APIToken) error { + return DB.WithContext(ctx).Create(apiToken).Error } // GetByTenantID gets API tokens by tenant ID -func (dao *APITokenDAO) GetByTenantID(tenantID string) ([]*entity.APIToken, error) { +func (dao *APITokenDAO) GetByTenantID(ctx context.Context, tenantID string) ([]*entity.APIToken, error) { var tokens []*entity.APIToken - err := DB.Where("tenant_id = ?", tenantID).Find(&tokens).Error + err := DB.WithContext(ctx).Where("tenant_id = ?", tenantID).Find(&tokens).Error return tokens, err } // DeleteByTenantID deletes all API tokens by tenant ID (hard delete) -func (dao *APITokenDAO) DeleteByTenantID(tenantID string) (int64, error) { - result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.APIToken{}) +func (dao *APITokenDAO) DeleteByTenantID(ctx context.Context, tenantID string) (int64, error) { + result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } // GetByToken gets API token by access key -func (dao *APITokenDAO) GetUserByAPIToken(token string) (*entity.APIToken, error) { +func (dao *APITokenDAO) GetUserByAPIToken(ctx context.Context, token string) (*entity.APIToken, error) { var apiToken entity.APIToken - err := DB.Where("token = ?", token).First(&apiToken).Error + err := DB.WithContext(ctx).Where("token = ?", token).First(&apiToken).Error if err != nil { return nil, err } @@ -60,24 +61,24 @@ func (dao *APITokenDAO) GetUserByAPIToken(token string) (*entity.APIToken, error // GetByBeta gets API tokens by beta key (SDK/bot authorization token). // Mirrors Python's APIToken.query(beta=token), which returns a list. -func (dao *APITokenDAO) GetByBeta(beta string) ([]*entity.APIToken, error) { +func (dao *APITokenDAO) GetByBeta(ctx context.Context, beta string) ([]*entity.APIToken, error) { var tokens []*entity.APIToken - err := DB.Where("beta = ?", beta).Find(&tokens).Error + err := DB.WithContext(ctx).Where("beta = ?", beta).Find(&tokens).Error return tokens, err } // DeleteByDialogIDs deletes API tokens by dialog IDs (hard delete) -func (dao *APITokenDAO) DeleteByDialogIDs(dialogIDs []string) (int64, error) { +func (dao *APITokenDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.APIToken{}) + result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } // DeleteByTenantIDAndToken deletes a specific API token by tenant ID and token value -func (dao *APITokenDAO) DeleteByTenantIDAndToken(tenantID, token string) (int64, error) { - result := DB.Unscoped().Where("tenant_id = ? AND token = ?", tenantID, token).Delete(&entity.APIToken{}) +func (dao *APITokenDAO) DeleteByTenantIDAndToken(ctx context.Context, tenantID, token string) (int64, error) { + result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ? AND token = ?", tenantID, token).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } @@ -105,11 +106,11 @@ type ConversationStatsRow struct { // DAO does not assign defaults because session creation paths in the // Python agent API generate a uuid + tenant timestamp and rely on the // round-trip shape being byte-identical. -func (dao *API4ConversationDAO) Create(conv *entity.API4Conversation) error { +func (dao *API4ConversationDAO) Create(ctx context.Context, conv *entity.API4Conversation) error { if conv == nil { return errors.New("api4 conversation: nil row") } - return DB.Create(conv).Error + return DB.WithContext(ctx).Create(conv).Error } // Update writes back an existing api_4_conversation row. The bot @@ -117,18 +118,18 @@ func (dao *API4ConversationDAO) Create(conv *entity.API4Conversation) error { // turn so multi-turn chatbot sessions carry prior history into the next // LLM call. Matches the Python conversation_service.update pattern at // api/db/services/conversation_service.py:236 (async_iframe_completion). -func (dao *API4ConversationDAO) Update(conv *entity.API4Conversation) error { +func (dao *API4ConversationDAO) Update(ctx context.Context, conv *entity.API4Conversation) error { if conv == nil { return errors.New("api4 conversation: nil row") } if conv.ID == "" { return errors.New("api4 conversation: empty id") } - return DB.Save(conv).Error + return DB.WithContext(ctx).Save(conv).Error } // Stats returns daily conversation aggregates for a tenant. -func (dao *API4ConversationDAO) Stats(tenantID, fromDate, toDate string, source *string) ([]ConversationStatsRow, error) { +func (dao *API4ConversationDAO) Stats(ctx context.Context, tenantID, fromDate, toDate string, source *string) ([]ConversationStatsRow, error) { var rows []ConversationStatsRow dateExpr := "DATE_FORMAT(a.create_date, '%Y-%m-%d 00:00:00')" db := DB.Table("api_4_conversation AS a"). @@ -156,9 +157,9 @@ func (dao *API4ConversationDAO) Stats(tenantID, fromDate, toDate string, source return rows, err } -func (dao *API4ConversationDAO) GetBySessionID(sessionID, agentID string) (*entity.API4Conversation, error) { +func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, sessionID, agentID string) (*entity.API4Conversation, error) { var result entity.API4Conversation - tx := DB.Where("id = ? AND dialog_id = ?", sessionID, agentID).Find(&result) + tx := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Find(&result) if tx.Error != nil { return nil, tx.Error } @@ -169,23 +170,23 @@ func (dao *API4ConversationDAO) GetBySessionID(sessionID, agentID string) (*enti } // ListIDsByAgentID lists conversation IDs for one agent. -func (dao *API4ConversationDAO) ListIDsByAgentID(agentID string) ([]string, error) { +func (dao *API4ConversationDAO) ListIDsByAgentID(ctx context.Context, agentID string) ([]string, error) { var ids []string - err := DB.Model(&entity.API4Conversation{}).Where("dialog_id = ?", agentID).Pluck("id", &ids).Error + err := DB.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", agentID).Pluck("id", &ids).Error return ids, err } // DeleteBySessionIDAndAgentID deletes API4Conversations by sessionID and agentID -func (dao *API4ConversationDAO) DeleteBySessionIDAndAgentID(sessionID, agentID string) (int64, error) { - result := DB.Where("id = ? AND dialog_id = ?", sessionID, agentID).Delete(&entity.API4Conversation{}) +func (dao *API4ConversationDAO) DeleteBySessionIDAndAgentID(ctx context.Context, sessionID, agentID string) (int64, error) { + result := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Delete(&entity.API4Conversation{}) return result.RowsAffected, result.Error } // DeleteByDialogIDs deletes API4Conversations by dialog IDs (hard delete) -func (dao *API4ConversationDAO) DeleteByDialogIDs(dialogIDs []string) (int64, error) { +func (dao *API4ConversationDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.API4Conversation{}) + result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.API4Conversation{}) return result.RowsAffected, result.Error } diff --git a/internal/dao/api_token_beta_test.go b/internal/dao/api_token_beta_test.go index 77b3322c87..eb6264ef26 100644 --- a/internal/dao/api_token_beta_test.go +++ b/internal/dao/api_token_beta_test.go @@ -55,7 +55,8 @@ func TestAPITokenDAOGetByBeta(t *testing.T) { t.Fatalf("failed to create api token: %v", err) } - got, err := NewAPITokenDAO().GetByBeta(beta) + ctx := t.Context() + got, err := NewAPITokenDAO().GetByBeta(ctx, beta) if err != nil { t.Fatalf("GetByBeta failed: %v", err) } diff --git a/internal/dao/api_token_test.go b/internal/dao/api_token_test.go index 87eeef5cfa..78bebf72c4 100644 --- a/internal/dao/api_token_test.go +++ b/internal/dao/api_token_test.go @@ -63,7 +63,8 @@ func TestAPI4ConversationDAOGetBySessionID(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") - session, err := NewAPI4ConversationDAO().GetBySessionID("session-1", "agent-1") + ctx := t.Context() + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "session-1", "agent-1") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -84,7 +85,8 @@ func TestAPI4ConversationDAOGetBySessionIDWrongAgent(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") - session, err := NewAPI4ConversationDAO().GetBySessionID("session-1", "agent-2") + ctx := t.Context() + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "session-1", "agent-2") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -99,7 +101,8 @@ func TestAPI4ConversationDAOGetBySessionIDNoRows(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") - session, err := NewAPI4ConversationDAO().GetBySessionID("missing-session", "agent-1") + ctx := t.Context() + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "missing-session", "agent-1") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -116,7 +119,8 @@ func TestAPI4ConversationDAOListIDsByAgentID(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-2", "agent-1") createAPI4ConversationForDAOTest(t, "session-other", "agent-2") - ids, err := NewAPI4ConversationDAO().ListIDsByAgentID("agent-1") + ctx := t.Context() + ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, "agent-1") if err != nil { t.Fatalf("ListIDsByAgentID failed: %v", err) } @@ -139,7 +143,8 @@ func TestAPI4ConversationDAOListIDsByAgentIDNoRows(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") - ids, err := NewAPI4ConversationDAO().ListIDsByAgentID("agent-2") + ctx := t.Context() + ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, "agent-2") if err != nil { t.Fatalf("ListIDsByAgentID failed: %v", err) } diff --git a/internal/dao/canvas_template.go b/internal/dao/canvas_template.go index bd407bc9a5..1c0a4ef907 100644 --- a/internal/dao/canvas_template.go +++ b/internal/dao/canvas_template.go @@ -17,6 +17,7 @@ package dao import ( + "context" "ragflow/internal/entity" ) @@ -31,9 +32,9 @@ func NewCanvasTemplateDAO() *CanvasTemplateDAO { // GetAll returns every row in canvas_template ordered by create_time desc, so // templates appear newest first in the UI. Mirrors the Python // CanvasTemplateService.get_all() behaviour. -func (dao *CanvasTemplateDAO) GetAll() ([]*entity.CanvasTemplate, error) { +func (dao *CanvasTemplateDAO) GetAll(ctx context.Context) ([]*entity.CanvasTemplate, error) { var templates []*entity.CanvasTemplate - if err := DB.Order("create_time desc").Find(&templates).Error; err != nil { + if err := DB.WithContext(ctx).Order("create_time desc").Find(&templates).Error; err != nil { return nil, err } return templates, nil diff --git a/internal/dao/canvas_template_seed.go b/internal/dao/canvas_template_seed.go index debf7fe62f..9a3b48c65a 100644 --- a/internal/dao/canvas_template_seed.go +++ b/internal/dao/canvas_template_seed.go @@ -18,6 +18,7 @@ package dao import ( "bytes" + "context" "encoding/json" "fmt" "os" @@ -35,8 +36,8 @@ import ( // SeedCanvasTemplates seeds the canvas_template table from the built-in // agent/templates/*.json and internal/ingestion/pipeline/template/*.json files. -func SeedCanvasTemplates() error { - if err := addColumnIfNotExists(DB, "canvas_template", "parser_ids", "LONGTEXT NULL"); err != nil { +func SeedCanvasTemplates(ctx context.Context) error { + if err := addColumnIfNotExists(ctx, DB, "canvas_template", "parser_ids", "LONGTEXT NULL"); err != nil { return fmt.Errorf("failed to ensure canvas_template.parser_ids column: %w", err) } diff --git a/internal/dao/canvas_template_seed_test.go b/internal/dao/canvas_template_seed_test.go index 1d1efffae6..d4e90a53ae 100644 --- a/internal/dao/canvas_template_seed_test.go +++ b/internal/dao/canvas_template_seed_test.go @@ -21,9 +21,10 @@ import ( "path/filepath" "testing" + "ragflow/internal/entity" + "github.com/glebarez/sqlite" "gorm.io/gorm" - "ragflow/internal/entity" ) func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { @@ -31,7 +32,7 @@ func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.CanvasTemplate{}); err != nil { + if err = db.AutoMigrate(&entity.CanvasTemplate{}); err != nil { t.Fatalf("migrate canvas templates: %v", err) } @@ -39,7 +40,7 @@ func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { writeTemplate := func(name, id, title string) { t.Helper() body := `{"id":"` + id + `","title":{"en":"` + title + `"},"description":{},"dsl":{}}` - if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600); err != nil { + if err = os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600); err != nil { t.Fatalf("write template: %v", err) } } @@ -52,7 +53,7 @@ func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { if err != nil { t.Fatalf("read templates: %v", err) } - if _, err := seedCanvasTemplates(db, dir, entries); err != nil { + if _, err = seedCanvasTemplates(db, dir, entries); err != nil { t.Fatalf("seed templates: %v", err) } } @@ -65,7 +66,7 @@ func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { seed() var templates []entity.CanvasTemplate - if err := db.Find(&templates).Error; err != nil { + if err = db.Find(&templates).Error; err != nil { t.Fatalf("read seeded templates: %v", err) } if len(templates) != 1 { diff --git a/internal/dao/chat.go b/internal/dao/chat.go index 9303accc40..a0369088a7 100644 --- a/internal/dao/chat.go +++ b/internal/dao/chat.go @@ -17,6 +17,7 @@ package dao import ( + "context" "fmt" "strings" "time" @@ -35,10 +36,10 @@ func NewChatDAO() *ChatDAO { } // ListByTenantID list chats by tenant ID -func (dao *ChatDAO) ListByTenantID(tenantID string, status string) ([]*entity.Chat, error) { +func (dao *ChatDAO) ListByTenantID(ctx context.Context, tenantID string, status string) ([]*entity.Chat, error) { var chats []*entity.Chat - query := DB.Model(&entity.Chat{}). + query := DB.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ?", tenantID) if status != "" { @@ -54,12 +55,12 @@ func (dao *ChatDAO) ListByTenantID(tenantID string, status string) ([]*entity.Ch } // ListByTenantIDs list chats by tenant IDs with pagination and filtering -func (dao *ChatDAO) ListByTenantIDs(tenantIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { +func (dao *ChatDAO) ListByTenantIDs(ctx context.Context, tenantIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { var chats []*entity.ChatListItem var total int64 // Build query with join to user table for nickname and avatar - query := DB.Model(&entity.Chat{}). + query := DB.WithContext(ctx).Model(&entity.Chat{}). Select(` dialog.*, user.nickname, @@ -106,11 +107,11 @@ func (dao *ChatDAO) ListByTenantIDs(tenantIDs []string, userID string, page, pag } // ListByOwnerIDs list chats by owner IDs with filtering (manual pagination) -func (dao *ChatDAO) ListByOwnerIDs(ownerIDs []string, userID string, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { +func (dao *ChatDAO) ListByOwnerIDs(ctx context.Context, ownerIDs []string, userID string, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { var chats []*entity.ChatListItem // Build query with join to user table - query := DB.Model(&entity.Chat{}). + query := DB.WithContext(ctx).Model(&entity.Chat{}). Select(` dialog.*, user.nickname, @@ -145,9 +146,9 @@ func (dao *ChatDAO) ListByOwnerIDs(ownerIDs []string, userID string, orderby str } // GetByID gets chat by ID -func (dao *ChatDAO) GetByID(id string) (*entity.Chat, error) { +func (dao *ChatDAO) GetByID(ctx context.Context, id string) (*entity.Chat, error) { var chat entity.Chat - err := DB.Where("id = ?", id).First(&chat).Error + err := DB.WithContext(ctx).Where("id = ?", id).First(&chat).Error if err != nil { return nil, err } @@ -155,9 +156,9 @@ func (dao *ChatDAO) GetByID(id string) (*entity.Chat, error) { } // GetByIDAndStatus gets chat by ID and status -func (dao *ChatDAO) GetByIDAndStatus(id string, status string) (*entity.Chat, error) { +func (dao *ChatDAO) GetByIDAndStatus(ctx context.Context, id string, status string) (*entity.Chat, error) { var chat entity.Chat - err := DB.Where("id = ? AND status = ?", id, status).First(&chat).Error + err := DB.WithContext(ctx).Where("id = ? AND status = ?", id, status).First(&chat).Error if err != nil { return nil, err } @@ -165,32 +166,32 @@ func (dao *ChatDAO) GetByIDAndStatus(id string, status string) (*entity.Chat, er } // GetExistingNames gets existing dialog names for a tenant -func (dao *ChatDAO) GetExistingNames(tenantID string, status string) ([]string, error) { +func (dao *ChatDAO) GetExistingNames(ctx context.Context, tenantID string, status string) ([]string, error) { var names []string - err := DB.Model(&entity.Chat{}). + err := DB.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ? AND status = ?", tenantID, status). Pluck("name", &names).Error return names, err } // ExistsByNameTenantStatus checks whether a chat with the given name exists. -func (dao *ChatDAO) ExistsByNameTenantStatus(name, tenantID, status string) (bool, error) { +func (dao *ChatDAO) ExistsByNameTenantStatus(ctx context.Context, name, tenantID, status string) (bool, error) { var count int64 - err := DB.Model(&entity.Chat{}). + err := DB.WithContext(ctx).Model(&entity.Chat{}). Where("name = ? AND tenant_id = ? AND status = ?", name, tenantID, status). Count(&count).Error return count > 0, err } // Create creates a new chat/dialog -func (dao *ChatDAO) Create(chat *entity.Chat) error { +func (dao *ChatDAO) Create(ctx context.Context, chat *entity.Chat) error { // Select("*") forces GORM to persist explicit zero values (e.g. similarity_threshold=0, // vector_similarity_weight=0, top_n=0) instead of substituting the column defaults. - return DB.Select("*").Create(chat).Error + return DB.WithContext(ctx).Select("*").Create(chat).Error } // UpdateByID updates a chat by ID -func (dao *ChatDAO) UpdateByID(id string, updates map[string]interface{}) error { +func (dao *ChatDAO) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { if updates == nil { updates = make(map[string]interface{}) } @@ -199,7 +200,7 @@ func (dao *ChatDAO) UpdateByID(id string, updates map[string]interface{}) error updates["update_time"] = now.UnixMilli() updates["update_date"] = now.Truncate(time.Second) - result := DB.Session(&gorm.Session{SkipHooks: true}).Model(&entity.Chat{}).Where("id = ?", id).Updates(updates) + result := DB.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.Chat{}).Where("id = ?", id).Updates(updates) if result.Error != nil { return result.Error } @@ -216,13 +217,13 @@ func (dao *ChatDAO) UpdateByID(id string, updates map[string]interface{}) error } // UpdateManyByID updates multiple chats by ID (batch update) -func (dao *ChatDAO) UpdateManyByID(updates []map[string]interface{}) error { +func (dao *ChatDAO) UpdateManyByID(ctx context.Context, updates []map[string]interface{}) error { if len(updates) == 0 { return nil } // Use transaction for batch update - tx := DB.Begin() + tx := DB.WithContext(ctx).Begin() if tx.Error != nil { return tx.Error } @@ -252,15 +253,15 @@ func (dao *ChatDAO) UpdateManyByID(updates []map[string]interface{}) error { } // DeleteByTenantID deletes all chats by tenant ID (hard delete) -func (dao *ChatDAO) DeleteByTenantID(tenantID string) (int64, error) { - result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Chat{}) +func (dao *ChatDAO) DeleteByTenantID(ctx context.Context, tenantID string) (int64, error) { + result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Chat{}) return result.RowsAffected, result.Error } // GetAllDialogIDsByTenantID gets all dialog IDs by tenant ID -func (dao *ChatDAO) GetAllDialogIDsByTenantID(tenantID string) ([]string, error) { +func (dao *ChatDAO) GetAllDialogIDsByTenantID(ctx context.Context, tenantID string) ([]string, error) { var dialogIDs []string - err := DB.Model(&entity.Chat{}). + err := DB.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ?", tenantID). Pluck("id", &dialogIDs).Error return dialogIDs, err @@ -269,8 +270,8 @@ func (dao *ChatDAO) GetAllDialogIDsByTenantID(tenantID string) ([]string, error) // QueryByTenantIDAndID checks if a chat exists with given tenant_id and id // Reference: Python DialogService.query(tenant_id=tenant.tenant_id, id=chat_id, status=StatusEnum.VALID.value) // Used for permission verification in get_chat API -func (dao *ChatDAO) QueryByTenantIDAndID(tenantID string, chatID string, status string) ([]*entity.Chat, error) { +func (dao *ChatDAO) QueryByTenantIDAndID(ctx context.Context, tenantID string, chatID string, status string) ([]*entity.Chat, error) { var chats []*entity.Chat - err := DB.Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, status).Find(&chats).Error + err := DB.WithContext(ctx).Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, status).Find(&chats).Error return chats, err } diff --git a/internal/dao/chat_channel.go b/internal/dao/chat_channel.go index 484a31b121..2fe6da29ab 100644 --- a/internal/dao/chat_channel.go +++ b/internal/dao/chat_channel.go @@ -1,6 +1,9 @@ package dao -import "ragflow/internal/entity" +import ( + "context" + "ragflow/internal/entity" +) type ChatChannelDAO struct{} @@ -8,22 +11,22 @@ func NewChatChannel() *ChatChannelDAO { return &ChatChannelDAO{} } -func (dao *ChatChannelDAO) Create(channel *entity.ChatChannel) error { - return DB.Create(channel).Error +func (dao *ChatChannelDAO) Create(ctx context.Context, channel *entity.ChatChannel) error { + return DB.WithContext(ctx).Create(channel).Error } -func (dao *ChatChannelDAO) GetByIDOnly(id string) (*entity.ChatChannel, error) { +func (dao *ChatChannelDAO) GetByIDOnly(ctx context.Context, id string) (*entity.ChatChannel, error) { var channel entity.ChatChannel - err := DB.Where("id = ?", id).First(&channel).Error + err := DB.WithContext(ctx).Where("id = ?", id).First(&channel).Error if err != nil { return nil, err } return &channel, err } -func (dao *ChatChannelDAO) GetByID(id string, tenantID string) (*entity.ChatChannel, error) { +func (dao *ChatChannelDAO) GetByID(ctx context.Context, id string, tenantID string) (*entity.ChatChannel, error) { var channel entity.ChatChannel - err := DB.Where("id = ? AND tenant_id = ?", id, tenantID).First(&channel).Error + err := DB.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&channel).Error if err != nil { return nil, err } @@ -31,20 +34,20 @@ func (dao *ChatChannelDAO) GetByID(id string, tenantID string) (*entity.ChatChan } // UpdateByID Update a single record by ID -func (dao *ChatChannelDAO) UpdateByID(id string, tenantID string, updates map[string]any) error { - return DB.Model(&entity.ChatChannel{}).Where("id = ? AND tenant_id = ?", id, tenantID).Updates(updates).Error +func (dao *ChatChannelDAO) UpdateByID(ctx context.Context, id string, tenantID string, updates map[string]any) error { + return DB.WithContext(ctx).Model(&entity.ChatChannel{}).Where("id = ? AND tenant_id = ?", id, tenantID).Updates(updates).Error } // DeleteByID Delete a single record by ID -func (dao *ChatChannelDAO) DeleteByID(id string, tenantID string) error { - return DB.Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.ChatChannel{}).Error +func (dao *ChatChannelDAO) DeleteByID(ctx context.Context, id string, tenantID string) error { + return DB.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.ChatChannel{}).Error } // ListByTenantID List a single record by TenantID -func (dao *ChatChannelDAO) ListByTenantID(tenantID string) ([]*entity.ChatChannelListResponse, error) { +func (dao *ChatChannelDAO) ListByTenantID(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) { results := make([]*entity.ChatChannelListResponse, 0) - err := DB.Table("chat_channel"). + err := DB.WithContext(ctx).Table("chat_channel"). Select("chat_channel.id, chat_channel.name, chat_channel.channel, chat_channel.chat_id, chat_channel.status, dialog.name as dialog_name"). Joins("LEFT JOIN dialog ON dialog.id = chat_channel.chat_id"). Where("chat_channel.tenant_id = ?", tenantID). diff --git a/internal/dao/chat_channel_test.go b/internal/dao/chat_channel_test.go index 14fe46d6f9..074afa2fe3 100644 --- a/internal/dao/chat_channel_test.go +++ b/internal/dao/chat_channel_test.go @@ -68,13 +68,15 @@ func TestChatChannelDAO_CRUD(t *testing.T) { Status: statusActive, } - err := dao.Create(cc) + ctx := t.Context() + + err := dao.Create(ctx, cc) if err != nil { t.Fatalf("failed to create chat channel: %v", err) } // 2. Test GetByID - res, err := dao.GetByID("chan-1", "tenant-1") + res, err := dao.GetByID(ctx, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get chat channel: %v", err) } @@ -90,7 +92,7 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // 2b. Test tenant isolation for GetByID - _, err = dao.GetByID("chan-1", "tenant-2") + _, err = dao.GetByID(ctx, "chan-1", "tenant-2") if err == nil { t.Fatalf("expected error (not found) when getting with wrong tenant, got nil") } @@ -100,12 +102,12 @@ func TestChatChannelDAO_CRUD(t *testing.T) { "name": "Updated WeCom Bot", } // Try updating with wrong tenant - err = dao.UpdateByID("chan-1", "tenant-2", updates) + err = dao.UpdateByID(ctx, "chan-1", "tenant-2", updates) if err != nil { t.Fatalf("failed to run UpdateByID with wrong tenant: %v", err) } // Verify it was NOT updated (should still be "Test WeCom Bot" since wrong tenant was used) - res, err = dao.GetByID("chan-1", "tenant-1") + res, err = dao.GetByID(ctx, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get chat channel: %v", err) } @@ -114,12 +116,12 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // Update with correct tenant - err = dao.UpdateByID("chan-1", "tenant-1", updates) + err = dao.UpdateByID(ctx, "chan-1", "tenant-1", updates) if err != nil { t.Fatalf("failed to update chat channel: %v", err) } - res, err = dao.GetByID("chan-1", "tenant-1") + res, err = dao.GetByID(ctx, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get updated chat channel: %v", err) } @@ -128,23 +130,23 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // 3b. Test DeleteByID with wrong tenant (should not delete) - err = dao.DeleteByID("chan-1", "tenant-2") + err = dao.DeleteByID(ctx, "chan-1", "tenant-2") if err != nil { t.Fatalf("failed to delete with wrong tenant: %v", err) } // Verify it still exists for tenant-1 - _, err = dao.GetByID("chan-1", "tenant-1") + _, err = dao.GetByID(ctx, "chan-1", "tenant-1") if err != nil { t.Fatalf("expected chat channel to still exist for tenant-1, got error: %v", err) } // 4. Test DeleteByID - err = dao.DeleteByID("chan-1", "tenant-1") + err = dao.DeleteByID(ctx, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to delete chat channel: %v", err) } - _, err = dao.GetByID("chan-1", "tenant-1") + _, err = dao.GetByID(ctx, "chan-1", "tenant-1") if err == nil { t.Fatalf("expected record not found error, got nil") } @@ -209,7 +211,8 @@ func TestChatChannelDAO_ListByTenantID(t *testing.T) { } // Perform query - list, err := dao.ListByTenantID("tenant-1") + ctx := t.Context() + list, err := dao.ListByTenantID(ctx, "tenant-1") if err != nil { t.Fatalf("ListByTenantID failed: %v", err) } diff --git a/internal/dao/chat_session.go b/internal/dao/chat_session.go index 3913c5623e..53cee49621 100644 --- a/internal/dao/chat_session.go +++ b/internal/dao/chat_session.go @@ -17,6 +17,7 @@ package dao import ( + "context" "strconv" "strings" "time" @@ -51,9 +52,9 @@ func NewChatSessionDAO() *ChatSessionDAO { } // GetByID gets chat session by ID -func (dao *ChatSessionDAO) GetByID(id string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetByID(ctx context.Context, id string) (*entity.ChatSession, error) { var conv entity.ChatSession - err := DB.Where("id = ?", id).First(&conv).Error + err := DB.WithContext(ctx).Where("id = ?", id).First(&conv).Error if err != nil { return nil, err } @@ -61,9 +62,9 @@ func (dao *ChatSessionDAO) GetByID(id string) (*entity.ChatSession, error) { } // GetBySessionIDAndChatID gets a chat session by session ID and chat ID. -func (dao *ChatSessionDAO) GetBySessionIDAndChatID(sessionID, chatID string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) { var conv entity.ChatSession - err := DB.Where("id = ? AND dialog_id = ?", sessionID, chatID).First(&conv).Error + err := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, chatID).First(&conv).Error if err != nil { return nil, err } @@ -71,12 +72,12 @@ func (dao *ChatSessionDAO) GetBySessionIDAndChatID(sessionID, chatID string) (*e } // Create creates a new chat session -func (dao *ChatSessionDAO) Create(conv *entity.ChatSession) error { - return DB.Create(conv).Error +func (dao *ChatSessionDAO) Create(ctx context.Context, conv *entity.ChatSession) error { + return DB.WithContext(ctx).Create(conv).Error } // UpdateByID updates a chat session by ID -func (dao *ChatSessionDAO) UpdateByID(id string, updates map[string]interface{}) error { +func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { if updates == nil { updates = make(map[string]interface{}) } @@ -85,13 +86,13 @@ func (dao *ChatSessionDAO) UpdateByID(id string, updates map[string]interface{}) updates["update_time"] = now.UnixMilli() updates["update_date"] = now.Truncate(time.Second) - result := DB.Session(&gorm.Session{SkipHooks: true}).Model(&entity.ChatSession{}).Where("id = ?", id).Updates(updates) + result := DB.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.ChatSession{}).Where("id = ?", id).Updates(updates) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { var count int64 - if err := DB.Model(&entity.ChatSession{}).Where("id = ?", id).Count(&count).Error; err != nil { + if err := DB.WithContext(ctx).Model(&entity.ChatSession{}).Where("id = ?", id).Count(&count).Error; err != nil { return err } if count == 0 { @@ -102,23 +103,23 @@ func (dao *ChatSessionDAO) UpdateByID(id string, updates map[string]interface{}) } // DeleteByID deletes a chat session by ID (hard delete) -func (dao *ChatSessionDAO) DeleteByID(id string) error { - return DB.Where("id = ?", id).Delete(&entity.ChatSession{}).Error +func (dao *ChatSessionDAO) DeleteByID(ctx context.Context, id string) error { + return DB.WithContext(ctx).Where("id = ?", id).Delete(&entity.ChatSession{}).Error } // ListByChatID lists chat sessions by chat ID -func (dao *ChatSessionDAO) ListByChatID(chatID string) ([]*entity.ChatSession, error) { +func (dao *ChatSessionDAO) ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) { var convs []*entity.ChatSession - err := DB.Where("dialog_id = ?", chatID). + err := DB.WithContext(ctx).Where("dialog_id = ?", chatID). Order("create_time DESC"). Find(&convs).Error return convs, err } // CheckDialogExists checks if a dialog exists with given tenant_id and dialog_id -func (dao *ChatSessionDAO) CheckDialogExists(tenantID, chatID string) (bool, error) { +func (dao *ChatSessionDAO) CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) { var count int64 - err := DB.Model(&entity.Chat{}). + err := DB.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, common.StatusDialogValid). Count(&count).Error if err != nil { @@ -128,9 +129,9 @@ func (dao *ChatSessionDAO) CheckDialogExists(tenantID, chatID string) (bool, err } // GetDialogByID gets dialog by ID -func (dao *ChatSessionDAO) GetDialogByID(chatID string) (*entity.Chat, error) { +func (dao *ChatSessionDAO) GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) { var dialog entity.Chat - err := DB.Where("id = ? AND status = ?", chatID, common.StatusDialogValid).First(&dialog).Error + err := DB.WithContext(ctx).Where("id = ? AND status = ?", chatID, common.StatusDialogValid).First(&dialog).Error if err != nil { return nil, err } @@ -138,17 +139,17 @@ func (dao *ChatSessionDAO) GetDialogByID(chatID string) (*entity.Chat, error) { } // DeleteByDialogIDs deletes chat sessions by dialog IDs (hard delete) -func (dao *ChatSessionDAO) DeleteByDialogIDs(dialogIDs []string) (int64, error) { +func (dao *ChatSessionDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.ChatSession{}) + result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.ChatSession{}) return result.RowsAffected, result.Error } -func (dao *ChatSessionDAO) ListAgentSessionNames(agentID, expUserID string) ([]map[string]interface{}, error) { +func (dao *ChatSessionDAO) ListAgentSessionNames(ctx context.Context, agentID, expUserID string) ([]map[string]interface{}, error) { var rows []map[string]interface{} - err := DB.Model(&entity.API4Conversation{}). + err := DB.WithContext(ctx).Model(&entity.API4Conversation{}). Select("id", "name"). Where("dialog_id = ? AND exp_user_id = ?", agentID, expUserID). Order("create_date DESC"). @@ -183,8 +184,8 @@ func normalizeAgentSessionOrderBy(orderBy string) string { } } -func (dao *ChatSessionDAO) ListAgentSessions(params ListAgentSessionsParams) (int64, []*entity.API4Conversation, error) { - query := DB.Model(&entity.API4Conversation{}).Where("dialog_id = ?", params.AgentID) +func (dao *ChatSessionDAO) ListAgentSessions(ctx context.Context, params ListAgentSessionsParams) (int64, []*entity.API4Conversation, error) { + query := DB.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", params.AgentID) if !params.IncludeDSL { query = query.Omit("dsl") } diff --git a/internal/dao/chat_session_test.go b/internal/dao/chat_session_test.go index 077a5d8dc0..114626e424 100644 --- a/internal/dao/chat_session_test.go +++ b/internal/dao/chat_session_test.go @@ -96,11 +96,12 @@ func TestChatSessionDAOUpdateByIDRefreshesTimestampsOnEmptyUpdate(t *testing.T) oldUpdateTime := int64(1000) createChatSessionForDAOTest(t, db, "session-1", "chat-1", "same", oldUpdateTime) - if err := NewChatSessionDAO().UpdateByID("session-1", map[string]interface{}{}); err != nil { + ctx := t.Context() + if err := NewChatSessionDAO().UpdateByID(ctx, "session-1", map[string]interface{}{}); err != nil { t.Fatalf("UpdateByID failed: %v", err) } - session, err := NewChatSessionDAO().GetByID("session-1") + session, err := NewChatSessionDAO().GetByID(ctx, "session-1") if err != nil { t.Fatalf("GetByID failed: %v", err) } @@ -118,7 +119,8 @@ func TestChatSessionDAOUpdateByIDSameValueSucceeds(t *testing.T) { createChatSessionForDAOTest(t, db, "session-1", "chat-1", "same", 1000) - if err := NewChatSessionDAO().UpdateByID("session-1", map[string]interface{}{"name": "same"}); err != nil { + ctx := t.Context() + if err := NewChatSessionDAO().UpdateByID(ctx, "session-1", map[string]interface{}{"name": "same"}); err != nil { t.Fatalf("UpdateByID failed: %v", err) } } @@ -127,7 +129,8 @@ func TestChatSessionDAOUpdateByIDMissingSession(t *testing.T) { db := setupChatSessionDAOTestDB(t) pushDB(t, db) - err := NewChatSessionDAO().UpdateByID("missing", nil) + ctx := t.Context() + err := NewChatSessionDAO().UpdateByID(ctx, "missing", nil) if !errors.Is(err, gorm.ErrRecordNotFound) { t.Fatalf("expected ErrRecordNotFound, got %v", err) } @@ -142,7 +145,8 @@ func TestChatSessionDAOListAgentSessionsOrdersByUpdateTimeDesc(t *testing.T) { createAgentSessionForDAOTest(t, db, "session-middle", "agent-1", "user-1", 2000) createAgentSessionForDAOTest(t, db, "session-other-agent", "agent-2", "user-1", 9999) - total, sessions, err := NewChatSessionDAO().ListAgentSessions(ListAgentSessionsParams{ + ctx := t.Context() + total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, ListAgentSessionsParams{ AgentID: "agent-1", Page: 1, PageSize: 10, @@ -180,7 +184,8 @@ func TestChatSessionDAOListAgentSessionsFiltersAndPaginates(t *testing.T) { createAgentSessionForDAOTest(t, db, "session-3", "agent-1", "user-1", 3000) createAgentSessionForDAOTest(t, db, "session-other-user", "agent-1", "user-2", 4000) - total, sessions, err := NewChatSessionDAO().ListAgentSessions(ListAgentSessionsParams{ + ctx := t.Context() + total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, ListAgentSessionsParams{ AgentID: "agent-1", UserID: "user-1", Page: 2, diff --git a/internal/dao/database.go b/internal/dao/database.go index 427f5960f6..b5dc6569cc 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -17,6 +17,7 @@ package dao import ( + "context" "fmt" "os" "path/filepath" @@ -65,7 +66,7 @@ type LLMFactoriesFile struct { } // InitDB initialize database connection -func InitDB(migrateDB bool) error { +func InitDB(ctx context.Context, migrateDB bool) error { cfg := server.GetConfig() dbCfg := cfg.Database @@ -167,7 +168,7 @@ func InitDB(migrateDB bool) error { } // Run manual migrations for complex schema changes - if err = RunMigrations(DB); err != nil { + if err = RunMigrations(ctx, DB); err != nil { return fmt.Errorf("failed to run manual migrations: %w", err) } common.Info("Database schema migrated successfully") @@ -175,7 +176,7 @@ func InitDB(migrateDB bool) error { // Seed built-in agent templates so the Go backend can serve the // "create agent from template" catalogue without relying on Python-side // initialization. - if err = SeedCanvasTemplates(); err != nil { + if err = SeedCanvasTemplates(ctx); err != nil { common.Warn("Failed to seed canvas templates", zap.Error(err)) } diff --git a/internal/dao/migration.go b/internal/dao/migration.go index dc774a8a81..cc5a890e4a 100644 --- a/internal/dao/migration.go +++ b/internal/dao/migration.go @@ -17,6 +17,7 @@ package dao import ( + "context" "fmt" "ragflow/internal/common" "ragflow/internal/entity" @@ -28,42 +29,42 @@ import ( // RunMigrations runs all manual database migrations // These are migrations that cannot be handled by AutoMigrate alone -func RunMigrations(db *gorm.DB) error { +func RunMigrations(ctx context.Context, db *gorm.DB) error { // Check if tenant_llm table has composite primary key and migrate to ID primary key - if err := migrateTenantLLMPrimaryKey(db); err != nil { + if err := migrateTenantLLMPrimaryKey(ctx, db); err != nil { return fmt.Errorf("failed to migrate tenant_llm primary key: %w", err) } // Rename columns (correct typos) - if err := renameColumnIfExists(db, "task", "process_duation", "process_duration"); err != nil { + if err := renameColumnIfExists(ctx, db, "task", "process_duation", "process_duration"); err != nil { return fmt.Errorf("failed to rename task.process_duation: %w", err) } - if err := renameColumnIfExists(db, "document", "process_duation", "process_duration"); err != nil { + if err := renameColumnIfExists(ctx, db, "document", "process_duation", "process_duration"); err != nil { return fmt.Errorf("failed to rename document.process_duation: %w", err) } // Add unique index on user.email - if err := migrateAddUniqueEmail(db); err != nil { + if err := migrateAddUniqueEmail(ctx, db); err != nil { return fmt.Errorf("failed to add unique index on user.email: %w", err) } // Add unique index on ingestion_task.document_id - if err := migrateIngestionTaskDocumentIDUnique(db); err != nil { + if err := migrateIngestionTaskDocumentIDUnique(ctx, db); err != nil { return fmt.Errorf("failed to add unique index on ingestion_task.document_id: %w", err) } // Modify column types that AutoMigrate may not handle correctly - if err := modifyColumnTypes(db); err != nil { + if err := modifyColumnTypes(ctx, db); err != nil { return fmt.Errorf("failed to modify column types: %w", err) } // Add case-insensitive unique constraint on knowledgebase (tenant_id, name) - if err := migrateKnowledgebaseNameUnique(db); err != nil { + if err := migrateKnowledgebaseNameUnique(ctx, db); err != nil { return fmt.Errorf("failed to add unique index on knowledgebase (tenant_id, name): %w", err) } // Add unique constraint on user_canvas (user_id, canvas_category, title) - if err := migrateUserCanvasTitleUnique(db); err != nil { + if err := migrateUserCanvasTitleUnique(ctx, db); err != nil { return fmt.Errorf("failed to add unique index on user_canvas (user_id, canvas_category, title): %w", err) } @@ -73,15 +74,15 @@ func RunMigrations(db *gorm.DB) error { // migrateTenantLLMPrimaryKey migrates tenant_llm from composite primary key to ID primary key // This corresponds to Python's update_tenant_llm_to_id_primary_key function -func migrateTenantLLMPrimaryKey(db *gorm.DB) error { +func migrateTenantLLMPrimaryKey(ctx context.Context, db *gorm.DB) error { // Check if tenant_llm table exists - if !db.Migrator().HasTable("tenant_llm") { + if !db.WithContext(ctx).Migrator().HasTable("tenant_llm") { return nil } // Check if 'id' column already exists using raw SQL var idColumnExists int64 - err := db.Raw(` + err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = 'tenant_llm' AND COLUMN_NAME = 'id' `).Scan(&idColumnExists).Error @@ -92,7 +93,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { if idColumnExists > 0 { // Check if id is already a primary key with auto_increment var count int64 - err = db.Raw(` + err = db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = 'tenant_llm' AND COLUMN_NAME = 'id' @@ -110,7 +111,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { common.Info("Migrating tenant_llm to use ID primary key...") // Start transaction - return db.Transaction(func(tx *gorm.DB) error { + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { // Check for temp_id column and drop it if exists var tempIdExists int64 tx.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS @@ -159,14 +160,14 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { } // migrateAddUniqueEmail adds unique index on user.email -func migrateAddUniqueEmail(db *gorm.DB) error { - if !db.Migrator().HasTable("user") { +func migrateAddUniqueEmail(ctx context.Context, db *gorm.DB) error { + if !db.WithContext(ctx).Migrator().HasTable("user") { return nil } // Check if unique index already exists using raw SQL var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'user' AND INDEX_NAME = 'idx_user_email_unique'`).Scan(&count) if count > 0 { return nil @@ -174,7 +175,7 @@ func migrateAddUniqueEmail(db *gorm.DB) error { // Check if there's a duplicate email issue first var duplicateCount int64 - err := db.Raw(` + err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM ( SELECT email FROM user GROUP BY email HAVING COUNT(*) > 1 ) AS duplicates @@ -189,7 +190,7 @@ func migrateAddUniqueEmail(db *gorm.DB) error { } common.Info("Adding unique index on user.email...") - if err = db.Exec(`ALTER TABLE user ADD UNIQUE INDEX idx_user_email_unique (email)`).Error; err != nil { + if err = db.WithContext(ctx).Exec(`ALTER TABLE user ADD UNIQUE INDEX idx_user_email_unique (email)`).Error; err != nil { // Check if error is MySQL duplicate index error (Error 1061) errStr := err.Error() @@ -203,15 +204,15 @@ func migrateAddUniqueEmail(db *gorm.DB) error { return nil } -func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { - if !db.Migrator().HasTable("ingestion_task") { +func migrateIngestionTaskDocumentIDUnique(ctx context.Context, db *gorm.DB) error { + if !db.WithContext(ctx).Migrator().HasTable("ingestion_task") { return nil } const indexName = "idx_ingestion_task_document_id" var uniqueCount int64 - if err := db.Raw(` + if err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'ingestion_task' AND INDEX_NAME = ? @@ -225,7 +226,7 @@ func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { } var duplicateCount int64 - if err := db.Raw(` + if err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM ( SELECT document_id FROM ingestion_task GROUP BY document_id HAVING COUNT(*) > 1 ) AS duplicates @@ -238,7 +239,7 @@ func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { } var existingIndexCount int64 - if err := db.Raw(` + if err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'ingestion_task' AND INDEX_NAME = ? @@ -246,12 +247,12 @@ func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { return err } if existingIndexCount > 0 { - if err := db.Exec(`ALTER TABLE ingestion_task DROP INDEX ` + indexName).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE ingestion_task DROP INDEX ` + indexName).Error; err != nil { return fmt.Errorf("failed to drop existing index %s: %w", indexName, err) } } - if err := db.Exec(`ALTER TABLE ingestion_task ADD UNIQUE INDEX ` + indexName + ` (document_id)`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE ingestion_task ADD UNIQUE INDEX ` + indexName + ` (document_id)`).Error; err != nil { errStr := err.Error() if strings.Contains(errStr, "Error 1061") && strings.Contains(errStr, "Duplicate key name") { common.Info("Index already exists, skipping", zap.String("error", errStr)) @@ -270,20 +271,20 @@ func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { // check-then-write path in CreateDataset/UpdateDataset against concurrent // duplicate inserts, and the resulting duplicate-key error is mapped back to the // "already exists" domain error at the service layer. -func migrateKnowledgebaseNameUnique(db *gorm.DB) error { - if !db.Migrator().HasTable("knowledgebase") { +func migrateKnowledgebaseNameUnique(ctx context.Context, db *gorm.DB) error { + if !db.WithContext(ctx).Migrator().HasTable("knowledgebase") { return nil } // Add the generated column if it does not exist yet. var colExists int64 - if err := db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + if err := db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = 'knowledgebase' AND COLUMN_NAME = 'name_ci'`).Scan(&colExists).Error; err != nil { return err } if colExists == 0 { common.Info("Adding generated column name_ci to knowledgebase...") - if err := db.Exec(`ALTER TABLE knowledgebase + if err := db.WithContext(ctx).Exec(`ALTER TABLE knowledgebase ADD COLUMN name_ci VARCHAR(128) GENERATED ALWAYS AS ( CASE WHEN status = '1' THEN LOWER(name) ELSE NULL END ) VIRTUAL`).Error; err != nil { @@ -300,7 +301,7 @@ func migrateKnowledgebaseNameUnique(db *gorm.DB) error { // Check whether the unique index already exists. var idxExists int64 - if err := db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + if err := db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = 'knowledgebase' AND INDEX_NAME = ?`, indexName).Scan(&idxExists).Error; err != nil { return err } @@ -310,7 +311,7 @@ func migrateKnowledgebaseNameUnique(db *gorm.DB) error { // Check for duplicate valid names before adding the index. var duplicateCount int64 - if err := db.Raw(` + if err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM ( SELECT tenant_id, name_ci FROM knowledgebase WHERE name_ci IS NOT NULL @@ -324,7 +325,7 @@ func migrateKnowledgebaseNameUnique(db *gorm.DB) error { } common.Info("Adding unique index on knowledgebase (tenant_id, name_ci)...") - if err := db.Exec("ALTER TABLE knowledgebase ADD UNIQUE INDEX " + indexName + " (tenant_id, name_ci)").Error; err != nil { + if err := db.WithContext(ctx).Exec("ALTER TABLE knowledgebase ADD UNIQUE INDEX " + indexName + " (tenant_id, name_ci)").Error; err != nil { errStr := err.Error() if strings.Contains(errStr, "Error 1061") && strings.Contains(errStr, "Duplicate key name") { common.Info("Index already exists, skipping", zap.String("error", errStr)) @@ -336,15 +337,15 @@ func migrateKnowledgebaseNameUnique(db *gorm.DB) error { return nil } -func migrateUserCanvasTitleUnique(db *gorm.DB) error { - if !db.Migrator().HasTable("user_canvas") { +func migrateUserCanvasTitleUnique(ctx context.Context, db *gorm.DB) error { + if !db.WithContext(ctx).Migrator().HasTable("user_canvas") { return nil } const indexName = "idx_user_canvas_user_category_title" var idxExists int64 - if err := db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + if err := db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = 'user_canvas' AND INDEX_NAME = ?`, indexName).Scan(&idxExists).Error; err != nil { return err } @@ -353,7 +354,7 @@ func migrateUserCanvasTitleUnique(db *gorm.DB) error { } var duplicateCount int64 - if err := db.Raw(` + if err := db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM ( SELECT user_id, canvas_category, title FROM user_canvas WHERE title IS NOT NULL @@ -367,7 +368,7 @@ func migrateUserCanvasTitleUnique(db *gorm.DB) error { } common.Info("Adding unique index on user_canvas (user_id, canvas_category, title)...") - if err := db.Exec("ALTER TABLE user_canvas ADD UNIQUE INDEX " + indexName + " (user_id, canvas_category, title)").Error; err != nil { + if err := db.WithContext(ctx).Exec("ALTER TABLE user_canvas ADD UNIQUE INDEX " + indexName + " (user_id, canvas_category, title)").Error; err != nil { errStr := err.Error() if strings.Contains(errStr, "Error 1061") && strings.Contains(errStr, "Duplicate key name") { common.Info("Index already exists, skipping", zap.String("error", errStr)) @@ -380,67 +381,67 @@ func migrateUserCanvasTitleUnique(db *gorm.DB) error { } // modifyColumnTypes modifies column types that need explicit ALTER statements -func modifyColumnTypes(db *gorm.DB) error { +func modifyColumnTypes(ctx context.Context, db *gorm.DB) error { columnExists := func(table, column string) bool { var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND COLUMN_NAME = ?`, table, column).Scan(&count) return count > 0 } // dialog.top_k: ensure its INTEGER with default 1024 - if db.Migrator().HasTable("dialog") && columnExists("dialog", "top_k") { - if err := db.Exec(`ALTER TABLE dialog MODIFY COLUMN top_k BIGINT NOT NULL DEFAULT 1024`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("dialog") && columnExists("dialog", "top_k") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE dialog MODIFY COLUMN top_k BIGINT NOT NULL DEFAULT 1024`).Error; err != nil { common.Warn("Failed to modify dialog.top_k", zap.Error(err)) } } // tenant_llm.api_key: ensure it's TEXT type - if db.Migrator().HasTable("tenant_llm") && columnExists("tenant_llm", "api_key") { - if err := db.Exec(`ALTER TABLE tenant_llm MODIFY COLUMN api_key VARCHAR(8192)`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("tenant_llm") && columnExists("tenant_llm", "api_key") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE tenant_llm MODIFY COLUMN api_key VARCHAR(8192)`).Error; err != nil { common.Warn("Failed to modify tenant_llm.api_key", zap.Error(err)) } } // api_token.dialog_id: ensure it's varchar(32) - if db.Migrator().HasTable("api_token") && columnExists("api_token", "dialog_id") { - if err := db.Exec(`ALTER TABLE api_token MODIFY COLUMN dialog_id VARCHAR(32)`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("api_token") && columnExists("api_token", "dialog_id") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE api_token MODIFY COLUMN dialog_id VARCHAR(32)`).Error; err != nil { common.Warn("Failed to modify api_token.dialog_id", zap.Error(err)) } } // canvas_template.title and description: ensure they're LONGTEXT type (same as Python JSONField) // Note: Python's JSONField uses null=True with application-level default, not database DEFAULT - if db.Migrator().HasTable("canvas_template") { + if db.WithContext(ctx).Migrator().HasTable("canvas_template") { if columnExists("canvas_template", "title") { - if err := db.Exec(`ALTER TABLE canvas_template MODIFY COLUMN title LONGTEXT NULL`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE canvas_template MODIFY COLUMN title LONGTEXT NULL`).Error; err != nil { common.Warn("Failed to modify canvas_template.title", zap.Error(err)) } } if columnExists("canvas_template", "description") { - if err := db.Exec(`ALTER TABLE canvas_template MODIFY COLUMN description LONGTEXT NULL`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE canvas_template MODIFY COLUMN description LONGTEXT NULL`).Error; err != nil { common.Warn("Failed to modify canvas_template.description", zap.Error(err)) } } } // system_settings.value: ensure it's LONGTEXT - if db.Migrator().HasTable("system_settings") && columnExists("system_settings", "value") { - if err := db.Exec(`ALTER TABLE system_settings MODIFY COLUMN value LONGTEXT NOT NULL`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("system_settings") && columnExists("system_settings", "value") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE system_settings MODIFY COLUMN value LONGTEXT NOT NULL`).Error; err != nil { common.Warn("Failed to modify system_settings.value", zap.Error(err)) } } // knowledgebase.raptor_task_finish_at: ensure it's DateTime - if db.Migrator().HasTable("knowledgebase") && columnExists("knowledgebase", "raptor_task_finish_at") { - if err := db.Exec(`ALTER TABLE knowledgebase MODIFY COLUMN raptor_task_finish_at DATETIME`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("knowledgebase") && columnExists("knowledgebase", "raptor_task_finish_at") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE knowledgebase MODIFY COLUMN raptor_task_finish_at DATETIME`).Error; err != nil { common.Warn("Failed to modify knowledgebase.raptor_task_finish_at", zap.Error(err)) } } // knowledgebase.mindmap_task_finish_at: ensure it's DateTime - if db.Migrator().HasTable("knowledgebase") && columnExists("knowledgebase", "mindmap_task_finish_at") { - if err := db.Exec(`ALTER TABLE knowledgebase MODIFY COLUMN mindmap_task_finish_at DATETIME`).Error; err != nil { + if db.WithContext(ctx).Migrator().HasTable("knowledgebase") && columnExists("knowledgebase", "mindmap_task_finish_at") { + if err := db.WithContext(ctx).Exec(`ALTER TABLE knowledgebase MODIFY COLUMN mindmap_task_finish_at DATETIME`).Error; err != nil { common.Warn("Failed to modify knowledgebase.mindmap_task_finish_at", zap.Error(err)) } } @@ -449,15 +450,15 @@ func modifyColumnTypes(db *gorm.DB) error { } // renameColumnIfExists renames a column if it exists and the new column doesn't exist -func renameColumnIfExists(db *gorm.DB, tableName, oldName, newName string) error { - if !db.Migrator().HasTable(tableName) { +func renameColumnIfExists(ctx context.Context, db *gorm.DB, tableName, oldName, newName string) error { + if !db.WithContext(ctx).Migrator().HasTable(tableName) { return nil } // Helper to check if column exists columnExists := func(column string) bool { var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND COLUMN_NAME = ?`, tableName, column).Scan(&count) return count > 0 } @@ -474,25 +475,25 @@ func renameColumnIfExists(db *gorm.DB, tableName, oldName, newName string) error zap.String("table", tableName), zap.String("oldColumn", oldName), zap.String("newColumn", newName)) - return db.Migrator().DropColumn(tableName, oldName) + return db.WithContext(ctx).Migrator().DropColumn(tableName, oldName) } common.Info("Renaming column", zap.String("table", tableName), zap.String("oldColumn", oldName), zap.String("newColumn", newName)) - return db.Migrator().RenameColumn(tableName, oldName, newName) + return db.WithContext(ctx).Migrator().RenameColumn(tableName, oldName, newName) } // addColumnIfNotExists adds a column if it doesn't exist -func addColumnIfNotExists(db *gorm.DB, tableName, columnName, columnDef string) error { - if !db.Migrator().HasTable(tableName) { +func addColumnIfNotExists(ctx context.Context, db *gorm.DB, tableName, columnName, columnDef string) error { + if !db.WithContext(ctx).Migrator().HasTable(tableName) { return nil } // Check if column exists using raw SQL var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND COLUMN_NAME = ?`, tableName, columnName).Scan(&count) if count > 0 { return nil @@ -502,13 +503,13 @@ func addColumnIfNotExists(db *gorm.DB, tableName, columnName, columnDef string) zap.String("table", tableName), zap.String("column", columnName)) sql := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, columnName, columnDef) - return db.Exec(sql).Error + return db.WithContext(ctx).Exec(sql).Error } // migrateSkillSearchTables creates skill search related tables -func migrateSkillSearchTables(db *gorm.DB) error { +func migrateSkillSearchTables(ctx context.Context, db *gorm.DB) error { // Create skill_search_configs table only - if !db.Migrator().HasTable("skill_search_configs") { + if !db.WithContext(ctx).Migrator().HasTable("skill_search_configs") { common.Info("Creating skill_search_configs table...") sql := ` CREATE TABLE IF NOT EXISTS skill_search_configs ( @@ -533,50 +534,50 @@ func migrateSkillSearchTables(db *gorm.DB) error { UNIQUE INDEX idx_tenant_space_embd (tenant_id, space_id, embd_id) ) ` - if err := db.Exec(sql).Error; err != nil { + if err := db.WithContext(ctx).Exec(sql).Error; err != nil { common.Warn("Failed to create skill_search_configs table with MySQL dialect, trying generic", zap.Error(err)) - if err := db.AutoMigrate(&entity.SkillSearchConfig{}); err != nil { + if err = db.WithContext(ctx).AutoMigrate(&entity.SkillSearchConfig{}); err != nil { return err } // AutoMigrate doesn't create unique indexes, so create them explicitly common.Info("Creating unique indexes for skill_search_configs...") - if err := db.Exec(`ALTER TABLE skill_search_configs ADD UNIQUE INDEX idx_tenant_space_embd (tenant_id, space_id, embd_id)`).Error; err != nil { + if err = db.WithContext(ctx).Exec(`ALTER TABLE skill_search_configs ADD UNIQUE INDEX idx_tenant_space_embd (tenant_id, space_id, embd_id)`).Error; err != nil { return fmt.Errorf("failed to create unique index idx_tenant_space_embd: %w", err) } } } else { // Add space_id for existing installations. - if err := addColumnIfNotExists(db, "skill_search_configs", "space_id", "VARCHAR(128) NOT NULL DEFAULT 'default'"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_search_configs", "space_id", "VARCHAR(128) NOT NULL DEFAULT 'default'"); err != nil { return fmt.Errorf("failed to add space_id column to skill_search_configs: %w", err) } - if err := addColumnIfNotExists(db, "skill_search_configs", "create_date", "DATETIME"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_search_configs", "create_date", "DATETIME"); err != nil { return fmt.Errorf("failed to add create_date column to skill_search_configs: %w", err) } - if err := addColumnIfNotExists(db, "skill_search_configs", "update_date", "DATETIME"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_search_configs", "update_date", "DATETIME"); err != nil { return fmt.Errorf("failed to add update_date column to skill_search_configs: %w", err) } - if err := db.Exec(`ALTER TABLE skill_search_configs MODIFY COLUMN update_time BIGINT`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE skill_search_configs MODIFY COLUMN update_time BIGINT`).Error; err != nil { common.Warn("Failed to modify skill_search_configs.update_time", zap.Error(err)) } // Drop legacy unique index (tenant_id, embd_id) to allow per-space configs. var legacyIndexExists int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_search_configs' AND INDEX_NAME = 'idx_tenant_embd'`).Scan(&legacyIndexExists) if legacyIndexExists > 0 { common.Info("Dropping legacy unique index idx_tenant_embd from skill_search_configs...") - if err := db.Exec(`ALTER TABLE skill_search_configs DROP INDEX idx_tenant_embd`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE skill_search_configs DROP INDEX idx_tenant_embd`).Error; err != nil { return fmt.Errorf("failed to drop legacy unique index idx_tenant_embd: %w", err) } } // Table exists, check if unique index exists var indexExists int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.WithContext(ctx).Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_search_configs' AND INDEX_NAME = 'idx_tenant_space_embd'`).Scan(&indexExists) if indexExists == 0 { common.Info("Adding unique index idx_tenant_space_embd to skill_search_configs...") - if err := db.Exec(`ALTER TABLE skill_search_configs + if err := db.WithContext(ctx).Exec(`ALTER TABLE skill_search_configs ADD UNIQUE INDEX idx_tenant_space_embd (tenant_id, space_id, embd_id)`).Error; err != nil { return fmt.Errorf("failed to add unique index idx_tenant_space_embd: %w", err) } @@ -587,8 +588,8 @@ func migrateSkillSearchTables(db *gorm.DB) error { } // migrateSkillSpaceTables creates skill space related tables -func migrateSkillSpaceTables(db *gorm.DB) error { - if !db.Migrator().HasTable("skill_spaces") { +func migrateSkillSpaceTables(ctx context.Context, db *gorm.DB) error { + if !db.WithContext(ctx).Migrator().HasTable("skill_spaces") { common.Info("Creating skill_spaces table...") sql := ` CREATE TABLE IF NOT EXISTS skill_spaces ( @@ -609,34 +610,34 @@ func migrateSkillSpaceTables(db *gorm.DB) error { UNIQUE INDEX idx_tenant_name_status (tenant_id, name, status) ) ` - if err := db.Exec(sql).Error; err != nil { + if err := db.WithContext(ctx).Exec(sql).Error; err != nil { common.Warn("Failed to create skill_spaces table with MySQL dialect, trying generic", zap.Error(err)) // Try with AutoMigrate as fallback - if err := db.AutoMigrate(&entity.SkillSpace{}); err != nil { + if err = db.WithContext(ctx).AutoMigrate(&entity.SkillSpace{}); err != nil { return err } // AutoMigrate doesn't create unique indexes, so create them explicitly common.Info("Creating unique indexes for skill_spaces...") - if err := db.Exec(`ALTER TABLE skill_spaces ADD UNIQUE INDEX idx_tenant_name_status (tenant_id, name, status)`).Error; err != nil { + if err = db.WithContext(ctx).Exec(`ALTER TABLE skill_spaces ADD UNIQUE INDEX idx_tenant_name_status (tenant_id, name, status)`).Error; err != nil { return fmt.Errorf("failed to create unique index idx_tenant_name_status: %w", err) } } } else { // Migrate existing table: add status column first, then update index - if err := addColumnIfNotExists(db, "skill_spaces", "status", "VARCHAR(1) NOT NULL DEFAULT '1'"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_spaces", "status", "VARCHAR(1) NOT NULL DEFAULT '1'"); err != nil { return fmt.Errorf("failed to add status column to skill_spaces: %w", err) } - if err := addColumnIfNotExists(db, "skill_spaces", "create_date", "DATETIME"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_spaces", "create_date", "DATETIME"); err != nil { return fmt.Errorf("failed to add create_date column to skill_spaces: %w", err) } - if err := addColumnIfNotExists(db, "skill_spaces", "update_date", "DATETIME"); err != nil { + if err := addColumnIfNotExists(ctx, db, "skill_spaces", "update_date", "DATETIME"); err != nil { return fmt.Errorf("failed to add update_date column to skill_spaces: %w", err) } - if err := db.Exec(`ALTER TABLE skill_spaces MODIFY COLUMN update_time BIGINT`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`ALTER TABLE skill_spaces MODIFY COLUMN update_time BIGINT`).Error; err != nil { common.Warn("Failed to modify skill_spaces.update_time", zap.Error(err)) } // Migrate index after status column exists - if err := migrateSkillSpaceIndex(db); err != nil { + if err := migrateSkillSpaceIndex(ctx, db); err != nil { return fmt.Errorf("failed to migrate skill_space index: %w", err) } } @@ -645,31 +646,31 @@ func migrateSkillSpaceTables(db *gorm.DB) error { } // migrateSkillSpaceIndex migrates the unique index to include status -func migrateSkillSpaceIndex(db *gorm.DB) error { +func migrateSkillSpaceIndex(ctx context.Context, db *gorm.DB) error { // Check if old index exists and drop it var oldIndexExists int64 - db.Raw(` + db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_spaces' AND INDEX_NAME = 'idx_tenant_name' `).Scan(&oldIndexExists) if oldIndexExists > 0 { common.Info("Dropping old idx_tenant_name index from skill_spaces...") - if err := db.Exec(`DROP INDEX idx_tenant_name ON skill_spaces`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`DROP INDEX idx_tenant_name ON skill_spaces`).Error; err != nil { return fmt.Errorf("failed to drop old index idx_tenant_name: %w", err) } } // Check if new index exists var newIndexExists int64 - db.Raw(` + db.WithContext(ctx).Raw(` SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_spaces' AND INDEX_NAME = 'idx_tenant_name_status' `).Scan(&newIndexExists) if newIndexExists == 0 { common.Info("Creating new idx_tenant_name_status index on skill_spaces...") - if err := db.Exec(`CREATE UNIQUE INDEX idx_tenant_name_status ON skill_spaces(tenant_id, name, status)`).Error; err != nil { + if err := db.WithContext(ctx).Exec(`CREATE UNIQUE INDEX idx_tenant_name_status ON skill_spaces(tenant_id, name, status)`).Error; err != nil { return fmt.Errorf("failed to create unique index idx_tenant_name_status: %w", err) } } diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 793c6a248e..5e2c82dd8c 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -359,8 +359,8 @@ func (h *AgentHandler) ListTemplates(c *gin.Context) { common.ErrorWithCode(c, errorCode, errorMessage) return } - - templates, err := h.agentService.ListTemplates() + ctx := c.Request.Context() + templates, err := h.agentService.ListTemplates(ctx) if err != nil { jsonInternalError(c, err) return @@ -629,7 +629,8 @@ func (h *AgentHandler) ListAgentTemplates(c *gin.Context) { common.ResponseWithCodeData(c, code, nil, msg) return } - rows, err := h.agentService.ListTemplates() + ctx := c.Request.Context() + rows, err := h.agentService.ListTemplates(ctx) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return @@ -711,8 +712,8 @@ func (h *AgentHandler) ListAgentSessions(c *gin.Context) { sessionID := c.Query("id") expUserID := c.Query("user_id") includeDSL := c.Query("dsl") == "true" - - resp, code, err := h.agentService.ListAgentSessions(user.ID, user.ID, canvasID, service.ListAgentSessionsRequest{ + ctx := c.Request.Context() + resp, code, err := h.agentService.ListAgentSessions(ctx, user.ID, user.ID, canvasID, service.ListAgentSessionsRequest{ Page: page, PageSize: pageSize, Keywords: keywords, @@ -749,7 +750,8 @@ func (h *AgentHandler) CreateAgentSession(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "Invalid request: "+err.Error()) return } - row, code, err := h.agentService.CreateAgentSession(&service.CreateAgentSessionRequest{ + ctx := c.Request.Context() + row, code, err := h.agentService.CreateAgentSession(ctx, &service.CreateAgentSessionRequest{ UserID: user.ID, AgentID: canvasID, Name: body.Name, @@ -772,7 +774,8 @@ func (h *AgentHandler) GetAgentSession(c *gin.Context) { } canvasID := c.Param("canvas_id") sessionID := c.Param("session_id") - row, code, err := h.agentService.GetAgentSession(user.ID, canvasID, sessionID) + ctx := c.Request.Context() + row, code, err := h.agentService.GetAgentSession(ctx, user.ID, canvasID, sessionID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -791,10 +794,11 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) { common.ResponseWithCodeData(c, code, nil, msg) return } + ctx := c.Request.Context() canvasID := c.Param("canvas_id") sessionID := c.Param("session_id") if sessionID != "" { - ok, code, err := h.agentService.DeleteAgentSessionItem(user.ID, canvasID, sessionID) + ok, code, err := h.agentService.DeleteAgentSessionItem(ctx, user.ID, canvasID, sessionID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -812,7 +816,7 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) { } } } - result, code, err := h.agentService.DeleteAgentSessions(user.ID, canvasID, ids, deleteAll) + result, code, err := h.agentService.DeleteAgentSessions(ctx, user.ID, canvasID, ids, deleteAll) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/api_token.go b/internal/handler/api_token.go index a41240f05f..d7f158f2e6 100644 --- a/internal/handler/api_token.go +++ b/internal/handler/api_token.go @@ -50,9 +50,10 @@ func (h *SystemHandler) ListAPIKeys(c *gin.Context) { } tenantID := tenants[0].TenantID + ctx := c.Request.Context() // Get keys for the tenant - keys, err := h.systemService.ListAPIKeys(tenantID) + keys, err := h.systemService.ListAPIKeys(ctx, tenantID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to list keys") return @@ -92,8 +93,10 @@ func (h *SystemHandler) CreateKey(c *gin.Context) { return } + ctx := c.Request.Context() + // Create key - key, err := h.systemService.CreateAPIKey(tenantID, &req) + key, err := h.systemService.CreateAPIKey(ctx, tenantID, &req) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to create key") return @@ -133,8 +136,10 @@ func (h *SystemHandler) DeleteKey(c *gin.Context) { return } + ctx := c.Request.Context() + // Delete key - if err = h.systemService.DeleteAPIKey(tenantID, key); err != nil { + if err = h.systemService.DeleteAPIKey(ctx, tenantID, key); err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to delete key") return } diff --git a/internal/handler/chat.go b/internal/handler/chat.go index c0a6680505..3bf711ad9c 100644 --- a/internal/handler/chat.go +++ b/internal/handler/chat.go @@ -103,9 +103,10 @@ func (h *ChatHandler) ListChats(c *gin.Context) { } ownerIDs := getOwnerIDs(c) + ctx := c.Request.Context() // List chats - default to valid status "1" (same as Python StatusEnum.VALID.value) - result, err := h.chatService.ListChats(userID, "1", keywords, page, pageSize, orderby, desc, ownerIDs) + result, err := h.chatService.ListChats(ctx, userID, "1", keywords, page, pageSize, orderby, desc, ownerIDs) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error()) return @@ -141,7 +142,8 @@ func (h *ChatHandler) Create(c *gin.Context) { req = map[string]interface{}{} } - result, code, err := h.chatService.Create(user.ID, req) + ctx := c.Request.Context() + result, code, err := h.chatService.Create(ctx, user.ID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -232,7 +234,8 @@ func (h *ChatHandler) DeleteChat(c *gin.Context) { return } - if err := h.chatService.DeleteChat(userID, chatID); err != nil { + ctx := c.Request.Context() + if err := h.chatService.DeleteChat(ctx, userID, chatID); err != nil { if err.Error() == "no authorization" { common.ResponseWithCodeData(c, common.CodeAuthenticationError, false, "No authorization.") return @@ -263,9 +266,10 @@ func (h *ChatHandler) BulkDeleteChats(c *gin.Context) { return } + ctx := c.Request.Context() if len(req.IDs) == 0 && !req.DeleteAll { if req.ChatID != "" { - if err := h.chatService.DeleteChat(userID, req.ChatID); err != nil { + if err := h.chatService.DeleteChat(ctx, userID, req.ChatID); err != nil { if err.Error() == "no authorization" { common.ResponseWithCodeData(c, common.CodeAuthenticationError, false, "No authorization.") return @@ -282,7 +286,7 @@ func (h *ChatHandler) BulkDeleteChats(c *gin.Context) { return } - result, err := h.chatService.BulkDeleteChats(userID, &req) + result, err := h.chatService.BulkDeleteChats(ctx, userID, &req) if err != nil { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return @@ -327,7 +331,8 @@ func (h *ChatHandler) GetChat(c *gin.Context) { } // Get chat detail with permission check - chat, err := h.chatService.GetChat(userID, chatID) + ctx := c.Request.Context() + chat, err := h.chatService.GetChat(ctx, userID, chatID) if err != nil { errMsg := err.Error() // Check if it's an authorization error @@ -404,14 +409,13 @@ func (h *ChatHandler) updateChatByMethod(c *gin.Context, patch bool) { return } - var ( - result map[string]interface{} - err error - ) + var result map[string]interface{} + var err error + ctx := c.Request.Context() if patch { - result, err = h.chatService.PatchChat(user.ID, chatID, req) + result, err = h.chatService.PatchChat(ctx, user.ID, chatID, req) } else { - result, err = h.chatService.UpdateChat(user.ID, chatID, req) + result, err = h.chatService.UpdateChat(ctx, user.ID, chatID, req) } if err != nil { if err.Error() == "no authorization" { diff --git a/internal/handler/chat_channel.go b/internal/handler/chat_channel.go index ec4838f681..351b4cc65c 100644 --- a/internal/handler/chat_channel.go +++ b/internal/handler/chat_channel.go @@ -15,6 +15,7 @@ package handler import ( + "context" "strings" "github.com/gin-gonic/gin" @@ -25,11 +26,11 @@ import ( ) type ChatChannelService interface { - CreateChatChannel(tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) - List(tenantID string) ([]*entity.ChatChannelListResponse, error) - GetChatChannel(userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) - UpdateChatChannel(userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) - DeleteChatChannel(userID, channelID string) (bool, common.ErrorCode, error) + CreateChatChannel(ctx context.Context, userID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) + List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) + GetChatChannel(ctx context.Context, userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) + UpdateChatChannel(ctx context.Context, userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) + DeleteChatChannel(ctx context.Context, userID, channelID string) (bool, common.ErrorCode, error) } type ChatChannelHandler struct { @@ -66,7 +67,10 @@ func (h *ChatChannelHandler) CreateChatChannel(c *gin.Context) { return } + ctx := c.Request.Context() + row, err := h.chatChannelService.CreateChatChannel( + ctx, user.ID, req.Name, req.Channel, @@ -88,7 +92,8 @@ func (h *ChatChannelHandler) ListChatChannel(c *gin.Context) { return } - rows, err := h.chatChannelService.List(user.ID) + ctx := c.Request.Context() + rows, err := h.chatChannelService.List(ctx, user.ID) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return @@ -116,7 +121,8 @@ func (h *ChatChannelHandler) GetChatChannel(c *gin.Context) { return } - channel, code, err := h.chatChannelService.GetChatChannel(userID, channelID) + ctx := c.Request.Context() + channel, code, err := h.chatChannelService.GetChatChannel(ctx, userID, channelID) if code != common.CodeSuccess || err != nil { writeChatChannelError(c, code, chatChannelErrMsg(code, err)) return @@ -151,7 +157,8 @@ func (h *ChatChannelHandler) UpdateChatChannel(c *gin.Context) { return } - result, code, err := h.chatChannelService.UpdateChatChannel(userID, channelID, unwrapChatChannelPayload(request)) + ctx := c.Request.Context() + result, code, err := h.chatChannelService.UpdateChatChannel(ctx, userID, channelID, unwrapChatChannelPayload(request)) if code != common.CodeSuccess || err != nil { writeChatChannelError(c, code, chatChannelErrMsg(code, err)) return @@ -180,7 +187,8 @@ func (h *ChatChannelHandler) DeleteChatChannel(c *gin.Context) { return } - result, code, err := h.chatChannelService.DeleteChatChannel(userID, channelID) + ctx := c.Request.Context() + result, code, err := h.chatChannelService.DeleteChatChannel(ctx, userID, channelID) if code != common.CodeSuccess || err != nil { writeChatChannelError(c, code, chatChannelErrMsg(code, err)) return diff --git a/internal/handler/chat_channel_test.go b/internal/handler/chat_channel_test.go index da68cb4a9a..c925fdd34b 100644 --- a/internal/handler/chat_channel_test.go +++ b/internal/handler/chat_channel_test.go @@ -1,6 +1,7 @@ package handler import ( + "context" "encoding/json" "errors" "net/http" @@ -22,35 +23,35 @@ type fakeChatChannelService struct { deleteFn func(userID, channelID string) (bool, common.ErrorCode, error) } -func (f fakeChatChannelService) CreateChatChannel(tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) { +func (f fakeChatChannelService) CreateChatChannel(ctx context.Context, tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) { if f.createFn == nil { return nil, errors.New("unexpected CreateChatChannel call") } return f.createFn(tenantID, name, channelType, config, chatID) } -func (f fakeChatChannelService) List(tenantID string) ([]*entity.ChatChannelListResponse, error) { +func (f fakeChatChannelService) List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) { if f.listFn == nil { return nil, errors.New("unexpected List call") } return f.listFn(tenantID) } -func (f fakeChatChannelService) GetChatChannel(userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) { +func (f fakeChatChannelService) GetChatChannel(ctx context.Context, userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) { if f.getFn == nil { return nil, common.CodeServerError, errors.New("unexpected GetChatChannel call") } return f.getFn(userID, channelID) } -func (f fakeChatChannelService) UpdateChatChannel(userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) { +func (f fakeChatChannelService) UpdateChatChannel(ctx context.Context, userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) { if f.updateFn == nil { return nil, common.CodeServerError, errors.New("unexpected UpdateChatChannel call") } return f.updateFn(userID, channelID, req) } -func (f fakeChatChannelService) DeleteChatChannel(userID, channelID string) (bool, common.ErrorCode, error) { +func (f fakeChatChannelService) DeleteChatChannel(ctx context.Context, userID, channelID string) (bool, common.ErrorCode, error) { if f.deleteFn == nil { return false, common.CodeServerError, errors.New("unexpected DeleteChatChannel call") } diff --git a/internal/handler/chat_session.go b/internal/handler/chat_session.go index f24185753b..5b2a1ac91a 100644 --- a/internal/handler/chat_session.go +++ b/internal/handler/chat_session.go @@ -68,7 +68,8 @@ func (h *ChatSessionHandler) ListChatSessions(c *gin.Context) { } // Call service to list chat sessions - result, err := h.chatSessionService.ListChatSessions(userID, chatID) + ctx := c.Request.Context() + result, err := h.chatSessionService.ListChatSessions(ctx, userID, chatID) if err != nil { // Check if it's an authorization error if err.Error() == "Only owner of dialog authorized for this operation" { @@ -247,7 +248,8 @@ func (h *ChatSessionHandler) GetSession(c *gin.Context) { userID := user.ID chatID, sessionID := c.Param("chat_id"), c.Param("session_id") - result, code, err := h.chatSessionService.GetSession(userID, chatID, sessionID) + ctx := c.Request.Context() + result, code, err := h.chatSessionService.GetSession(ctx, userID, chatID, sessionID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -288,7 +290,8 @@ func (h *ChatSessionHandler) CreateSession(c *gin.Context) { req = map[string]interface{}{} } - result, code, err := h.chatSessionService.CreateSession(userID, chatID, req) + ctx := c.Request.Context() + result, code, err := h.chatSessionService.CreateSession(ctx, userID, chatID, req) if err != nil { if code == common.CodeAuthenticationError { common.ResponseWithCodeData(c, code, false, err.Error()) @@ -334,7 +337,8 @@ func (h *ChatSessionHandler) DeleteSessions(c *gin.Context) { req = map[string]interface{}{} } - result, message, code, err := h.chatSessionService.DeleteSessions(userID, chatID, req) + ctx := c.Request.Context() + result, message, code, err := h.chatSessionService.DeleteSessions(ctx, userID, chatID, req) if err != nil { if code == common.CodeAuthenticationError { common.ResponseWithCodeData(c, code, false, err.Error()) @@ -372,7 +376,8 @@ func (h *ChatSessionHandler) UpdateSession(c *gin.Context) { return } - result, code, err := h.chatSessionService.UpdateSession(userID, chatID, sessionID, req) + ctx := c.Request.Context() + result, code, err := h.chatSessionService.UpdateSession(ctx, userID, chatID, sessionID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -389,7 +394,8 @@ func (h *ChatSessionHandler) DeleteSessionMessage(c *gin.Context) { userID := user.ID chatID, sessionID, msgID := c.Param("chat_id"), c.Param("session_id"), c.Param("msg_id") - result, code, err := h.chatSessionService.DeleteSessionMessage(userID, chatID, sessionID, msgID) + ctx := c.Request.Context() + result, code, err := h.chatSessionService.DeleteSessionMessage(ctx, userID, chatID, sessionID, msgID) if err != nil { if code == common.CodeAuthenticationError { common.ResponseWithCodeData(c, code, false, err.Error()) diff --git a/internal/handler/mcp_server.go b/internal/handler/mcp_server.go index b490f70caf..10da9c86ef 100644 --- a/internal/handler/mcp_server.go +++ b/internal/handler/mcp_server.go @@ -17,6 +17,7 @@ package handler import ( + "context" "encoding/json" "fmt" "io" @@ -123,8 +124,8 @@ func MCPListDatasets(ds *dataset.DatasetService, userID string, page, pageSize i // MCPListChats wraps ChatService.ListChats for the MCP tool handler, // converting the typed response into a generic []map[string]interface{}. -func MCPListChats(cs *service.ChatService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { - resp, err := cs.ListChats(userID, "1", "", page, pageSize, orderby, desc, nil) +func MCPListChats(ctx context.Context, chatService *service.ChatService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + resp, err := chatService.ListChats(ctx, userID, "1", "", page, pageSize, orderby, desc, nil) if err != nil { return nil, 0, err } diff --git a/internal/handler/stats.go b/internal/handler/stats.go index a02f9f17a1..c90cd4dcdd 100644 --- a/internal/handler/stats.go +++ b/internal/handler/stats.go @@ -57,8 +57,10 @@ func (h *StatsHandler) GetStats(c *gin.Context) { agentSource := "agent" source = &agentSource } + ctx := c.Request.Context() - stats, err := h.statsService.GetStats(user.ID, fromDate, toDate, source) + // Get stats + stats, err := h.statsService.GetStats(ctx, user.ID, fromDate, toDate, source) if err != nil { code := common.CodeExceptionError if errors.Is(err, service.ErrTenantNotFound) { diff --git a/internal/service/agent.go b/internal/service/agent.go index ca5db0f568..b582383533 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -355,8 +355,8 @@ func NewAgentServiceWithOptions( // ListTemplates returns every canvas template. Mirrors Python // agent_api.list_agent_template, which iterates CanvasTemplateService.get_all() // and serialises each row. -func (s *AgentService) ListTemplates() ([]*entity.CanvasTemplate, error) { - return s.canvasTemplateDAO.GetAll() +func (s *AgentService) ListTemplates(ctx context.Context) ([]*entity.CanvasTemplate, error) { + return s.canvasTemplateDAO.GetAll(ctx) } // AgentItem is one entry in the list response. @@ -1020,7 +1020,7 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID dsl = normalisedDSLForRun(versionRow) } if sessionID != "" && s.api4ConversationDAO != nil { - session, sessionErr := s.api4ConversationDAO.GetBySessionID(sessionID, canvasID) + session, sessionErr := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, canvasID) if sessionErr != nil { return nil, fmt.Errorf("RunAgent: load session %q: %w: %w", sessionID, sessionErr, ErrAgentStorageError) } @@ -1032,7 +1032,7 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID } } if newSession && len(dsl) > 0 { - if err := s.createAgentRunSession(sessionID, userID, canvasID, dsl, versionRow); err != nil { + if err = s.createAgentRunSession(ctx, sessionID, userID, canvasID, dsl, versionRow); err != nil { return nil, fmt.Errorf("RunAgent: create session %q: %w: %w", sessionID, err, ErrAgentStorageError) } } @@ -1426,7 +1426,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if answer != "" { appendAssistantHistory(state, partialAssistantOutput(answer, downloads)) } - if persistErr := s.persistAgentRunSession(canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, answer != ""); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, answer != ""); persistErr != nil { return nil, fmt.Errorf("persist interrupted agent session: %w: %w", persistErr, ErrAgentStorageError) } if answer != "" { @@ -1443,7 +1443,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } if shouldTreatAsCompletedLoopRun(err, answer) { appendAssistantHistory(state, assistantOutput) - if persistErr := s.persistAgentRunSession(canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } @@ -1477,7 +1477,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // Emit message + message_end (mirrors Python's ans dict). appendAssistantHistory(state, assistantOutput) - if persistErr := s.persistAgentRunSession(canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } @@ -1524,6 +1524,7 @@ func beginLayoutRecognize(c *canvas.Canvas) string { } func (s *AgentService) createAgentRunSession( + ctx context.Context, sessionID, userID, agentID string, runDSL map[string]any, versionRow *entity.UserCanvasVersion, @@ -1544,7 +1545,7 @@ func (s *AgentService) createAgentRunSession( if versionRow != nil { session.VersionTitle = versionRow.Title } - return s.api4ConversationDAO.Create(session) + return s.api4ConversationDAO.Create(ctx, session) } // runIDFor builds the per-run CanvasState identifier: canvasID @@ -1582,6 +1583,7 @@ func emptyDownloadValue(value any) bool { } func (s *AgentService) persistAgentRunSession( + ctx context.Context, agentID, userID, sessionID, messageID string, userInput any, answer string, @@ -1593,7 +1595,7 @@ func (s *AgentService) persistAgentRunSession( if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil { return nil } - session, err := s.api4ConversationDAO.GetBySessionID(sessionID, agentID) + session, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) if err != nil { common.Warn("agent run: load session for update failed", zap.String("agent_id", agentID), zap.String("session_id", sessionID), zap.Error(err)) return nil @@ -1620,7 +1622,7 @@ func (s *AgentService) persistAgentRunSession( if state != nil { session.DSL = buildPersistedAgentDSL(runDSL, state) } - return s.api4ConversationDAO.Update(session) + return s.api4ConversationDAO.Update(ctx, session) } func buildPersistedAgentDSL(runDSL map[string]any, state *canvas.CanvasState) entity.JSONMap { diff --git a/internal/service/agent_sessions.go b/internal/service/agent_sessions.go index c29cb9b7db..6b029e2f7d 100644 --- a/internal/service/agent_sessions.go +++ b/internal/service/agent_sessions.go @@ -17,6 +17,7 @@ package service import ( + "context" "encoding/json" "errors" "fmt" @@ -311,7 +312,7 @@ func checkDuplicateSessionIDs(ids []string) ([]string, []string) { } // ListAgentSessions returns paginated agent sessions visible to the caller. -func (s *AgentService) ListAgentSessions(userID, tenantID, agentID string, req ListAgentSessionsRequest) (*ListAgentSessionsResponse, common.ErrorCode, error) { +func (s *AgentService) ListAgentSessions(ctx context.Context, userID, tenantID, agentID string, req ListAgentSessionsRequest) (*ListAgentSessionsResponse, common.ErrorCode, error) { if agentID == "" { return nil, common.CodeArgumentError, errors.New("agent_id is required") } @@ -335,7 +336,7 @@ func (s *AgentService) ListAgentSessions(userID, tenantID, agentID string, req L sessionDAO := dao.NewChatSessionDAO() if req.ExpUserID != "" { - rows, err := sessionDAO.ListAgentSessionNames(agentID, req.ExpUserID) + rows, err := sessionDAO.ListAgentSessionNames(ctx, agentID, req.ExpUserID) if err != nil { return nil, common.CodeServerError, err } @@ -351,7 +352,7 @@ func (s *AgentService) ListAgentSessions(userID, tenantID, agentID string, req L return nil, common.CodeArgumentError, err } - total, sessions, err := sessionDAO.ListAgentSessions(dao.ListAgentSessionsParams{ + total, sessions, err := sessionDAO.ListAgentSessions(ctx, dao.ListAgentSessionsParams{ AgentID: agentID, Page: req.Page, PageSize: req.PageSize, @@ -377,7 +378,7 @@ func (s *AgentService) ListAgentSessions(userID, tenantID, agentID string, req L } // GetAgentSession fetches a single conversation belonging to agentID. -func (s *AgentService) GetAgentSession(userID, agentID, sessionID string) (*entity.API4Conversation, common.ErrorCode, error) { +func (s *AgentService) GetAgentSession(ctx context.Context, userID, agentID, sessionID string) (*entity.API4Conversation, common.ErrorCode, error) { if sessionID == "" { return nil, common.CodeArgumentError, fmt.Errorf("session_id is required") } @@ -392,7 +393,7 @@ func (s *AgentService) GetAgentSession(userID, agentID, sessionID string) (*enti return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") } - data, err := s.api4ConversationDAO.GetBySessionID(sessionID, agentID) + data, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to fetch session: %w", err) } @@ -403,7 +404,7 @@ func (s *AgentService) GetAgentSession(userID, agentID, sessionID string) (*enti } // DeleteAgentSessionItem removes one conversation if it belongs to agentID. -func (s *AgentService) DeleteAgentSessionItem(userID, agentID, sessionID string) (bool, common.ErrorCode, error) { +func (s *AgentService) DeleteAgentSessionItem(ctx context.Context, userID, agentID, sessionID string) (bool, common.ErrorCode, error) { if sessionID == "" { return false, common.CodeArgumentError, errors.New("session_id is required") } @@ -418,7 +419,7 @@ func (s *AgentService) DeleteAgentSessionItem(userID, agentID, sessionID string) return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") } - row, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(sessionID, agentID) + row, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, sessionID, agentID) if err != nil { return false, common.CodeServerError, err } @@ -431,7 +432,7 @@ func (s *AgentService) DeleteAgentSessionItem(userID, agentID, sessionID string) // DeleteAgentSessions removes multiple conversations owned by agentID. // When ids is empty and deleteAll is true, every session under agentID is // removed. -func (s *AgentService) DeleteAgentSessions(userID, agentID string, ids []string, deleteAll bool) (*DeleteAgentSessionsResult, common.ErrorCode, error) { +func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID string, ids []string, deleteAll bool) (*DeleteAgentSessionsResult, common.ErrorCode, error) { if agentID == "" { return nil, common.CodeArgumentError, errors.New("agent_id is required") } @@ -450,7 +451,7 @@ func (s *AgentService) DeleteAgentSessions(userID, agentID string, ids []string, return &DeleteAgentSessionsResult{}, common.CodeSuccess, nil } - ids, err = s.api4ConversationDAO.ListIDsByAgentID(agentID) + ids, err = s.api4ConversationDAO.ListIDsByAgentID(ctx, agentID) if err != nil { return nil, common.CodeServerError, err } @@ -470,7 +471,7 @@ func (s *AgentService) DeleteAgentSessions(userID, agentID string, ids []string, continue } - conv, err := s.api4ConversationDAO.GetBySessionID(sessionID, agentID) + conv, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) if err != nil { return nil, common.CodeServerError, err } @@ -479,7 +480,7 @@ func (s *AgentService) DeleteAgentSessions(userID, agentID string, ids []string, continue } - if _, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(sessionID, agentID); err != nil { + if _, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, sessionID, agentID); err != nil { return nil, common.CodeServerError, err } successCount++ @@ -686,7 +687,7 @@ type CreateAgentSessionRequest struct { // - update_time : unix-millis // - create_date : local-time.Truncate(time.Second) // - update_date : local-time.Truncate(time.Second) -func (s *AgentService) CreateAgentSession(req *CreateAgentSessionRequest) (*entity.API4Conversation, common.ErrorCode, error) { +func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentSessionRequest) (*entity.API4Conversation, common.ErrorCode, error) { if req == nil { return nil, common.CodeArgumentError, errors.New("create agent session: nil request") } @@ -751,7 +752,7 @@ func (s *AgentService) CreateAgentSession(req *CreateAgentSessionRequest) (*enti Source: sourcePtr, DSL: dsl, } - if err := s.api4ConversationDAO.Create(row); err != nil { + if err := s.api4ConversationDAO.Create(ctx, row); err != nil { return nil, common.CodeServerError, fmt.Errorf("create agent session: %w", err) } return row, common.CodeSuccess, nil diff --git a/internal/service/agent_test.go b/internal/service/agent_test.go index 50f96b4f16..5240d668a7 100644 --- a/internal/service/agent_test.go +++ b/internal/service/agent_test.go @@ -874,7 +874,8 @@ func TestListAgentSessionsServiceSuccess(t *testing.T) { createAgentSessionTestConversation(t, "session-new", "canvas-1", "user-1", 3000) createAgentSessionTestConversation(t, "session-other-agent", "canvas-other", "user-1", 9999) - resp, code, err := NewAgentService().ListAgentSessions("user-1", "user-1", "canvas-1", ListAgentSessionsRequest{ + ctx := t.Context() + resp, code, err := NewAgentService().ListAgentSessions(ctx, "user-1", "user-1", "canvas-1", ListAgentSessionsRequest{ Page: 1, PageSize: 10, OrderBy: "update_time", @@ -913,7 +914,8 @@ func TestListAgentSessionsServiceDenied(t *testing.T) { createAgentSessionTestCanvas(t, "canvas-1", "user-2") - resp, code, err := NewAgentService().ListAgentSessions("user-1", "user-1", "canvas-1", ListAgentSessionsRequest{}) + ctx := t.Context() + resp, code, err := NewAgentService().ListAgentSessions(ctx, "user-1", "user-1", "canvas-1", ListAgentSessionsRequest{}) if err == nil { t.Fatal("expected permission error") } @@ -931,7 +933,8 @@ func TestGetAgentSessionServiceSuccess(t *testing.T) { createAgentSessionTestCanvas(t, "canvas-1", "user-1") createAgentSessionTestConversation(t, "session-1", "canvas-1", "user-1", 1000) - session, code, err := NewAgentService().GetAgentSession("user-1", "canvas-1", "session-1") + ctx := t.Context() + session, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "canvas-1", "session-1") if err != nil { t.Fatalf("GetAgentSession failed: %v", err) } @@ -955,7 +958,8 @@ func TestGetAgentSessionServiceNotFoundWhenSessionBelongsToAnotherAgent(t *testi createAgentSessionTestCanvas(t, "canvas-1", "user-1") createAgentSessionTestConversation(t, "session-other", "canvas-other", "user-1", 1000) - session, code, err := NewAgentService().GetAgentSession("user-1", "canvas-1", "session-other") + ctx := t.Context() + session, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "canvas-1", "session-other") if err == nil { t.Fatal("expected not found error") } @@ -974,7 +978,8 @@ func TestDeleteAgentSessionItemServiceDeletesMatchingSession(t *testing.T) { createAgentSessionTestConversation(t, "session-1", "canvas-1", "user-1", 1000) createAgentSessionTestConversation(t, "session-other", "canvas-other", "user-1", 2000) - deleted, code, err := NewAgentService().DeleteAgentSessionItem("user-1", "canvas-1", "session-1") + ctx := t.Context() + deleted, code, err := NewAgentService().DeleteAgentSessionItem(ctx, "user-1", "canvas-1", "session-1") if err != nil { t.Fatalf("DeleteAgentSessionItem failed: %v", err) } @@ -986,14 +991,14 @@ func TestDeleteAgentSessionItemServiceDeletesMatchingSession(t *testing.T) { } var count int64 - if err := dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { t.Fatalf("failed to count deleted session: %v", err) } if count != 0 { t.Fatalf("expected session-1 to be deleted, count=%d", count) } - if err := dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-other").Count(&count).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-other").Count(&count).Error; err != nil { t.Fatalf("failed to count other session: %v", err) } if count != 1 { @@ -1007,7 +1012,8 @@ func TestDeleteAgentSessionItemServiceNoopForSessionFromAnotherAgent(t *testing. createAgentSessionTestCanvas(t, "canvas-1", "user-1") createAgentSessionTestConversation(t, "session-other", "canvas-other", "user-1", 1000) - deleted, code, err := NewAgentService().DeleteAgentSessionItem("user-1", "canvas-1", "session-other") + ctx := t.Context() + deleted, code, err := NewAgentService().DeleteAgentSessionItem(ctx, "user-1", "canvas-1", "session-other") if err != nil { t.Fatalf("DeleteAgentSessionItem failed: %v", err) } @@ -1035,7 +1041,8 @@ func TestDeleteAgentSessionsServiceDeleteAll(t *testing.T) { createAgentSessionTestConversation(t, "session-2", "canvas-1", "user-1", 2000) createAgentSessionTestConversation(t, "session-other", "canvas-other", "user-1", 3000) - result, code, err := NewAgentService().DeleteAgentSessions("user-1", "canvas-1", nil, true) + ctx := t.Context() + result, code, err := NewAgentService().DeleteAgentSessions(ctx, "user-1", "canvas-1", nil, true) if err != nil { t.Fatalf("DeleteAgentSessions failed: %v", err) } @@ -1050,7 +1057,7 @@ func TestDeleteAgentSessionsServiceDeleteAll(t *testing.T) { } var ownCount int64 - if err := dao.DB.Model(&entity.API4Conversation{}).Where("dialog_id = ?", "canvas-1").Count(&ownCount).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("dialog_id = ?", "canvas-1").Count(&ownCount).Error; err != nil { t.Fatalf("failed to count own sessions: %v", err) } if ownCount != 0 { @@ -1058,7 +1065,7 @@ func TestDeleteAgentSessionsServiceDeleteAll(t *testing.T) { } var otherCount int64 - if err := dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-other").Count(&otherCount).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-other").Count(&otherCount).Error; err != nil { t.Fatalf("failed to count other session: %v", err) } if otherCount != 1 { @@ -1072,7 +1079,8 @@ func TestDeleteAgentSessionsServiceDuplicateIDsPartial(t *testing.T) { createAgentSessionTestCanvas(t, "canvas-1", "user-1") createAgentSessionTestConversation(t, "session-1", "canvas-1", "user-1", 1000) - result, code, err := NewAgentService().DeleteAgentSessions("user-1", "canvas-1", []string{"session-1", "session-1"}, false) + ctx := t.Context() + result, code, err := NewAgentService().DeleteAgentSessions(ctx, "user-1", "canvas-1", []string{"session-1", "session-1"}, false) if err != nil { t.Fatalf("DeleteAgentSessions failed: %v", err) } @@ -1090,7 +1098,7 @@ func TestDeleteAgentSessionsServiceDuplicateIDsPartial(t *testing.T) { } var count int64 - if err := dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { t.Fatalf("failed to count deleted session: %v", err) } if count != 0 { @@ -1103,7 +1111,8 @@ func TestDeleteAgentSessionsServiceMissingSessionError(t *testing.T) { createAgentSessionTestCanvas(t, "canvas-1", "user-1") - result, code, err := NewAgentService().DeleteAgentSessions("user-1", "canvas-1", []string{"missing-session"}, false) + ctx := t.Context() + result, code, err := NewAgentService().DeleteAgentSessions(ctx, "user-1", "canvas-1", []string{"missing-session"}, false) if err == nil { t.Fatal("expected missing session error") } @@ -1124,7 +1133,8 @@ func TestDeleteAgentSessionsServiceRequiresOwner(t *testing.T) { createAgentSessionTestCanvas(t, "canvas-1", "user-2") createAgentSessionTestConversation(t, "session-1", "canvas-1", "user-1", 1000) - result, code, err := NewAgentService().DeleteAgentSessions("user-1", "canvas-1", []string{"session-1"}, false) + ctx := t.Context() + result, code, err := NewAgentService().DeleteAgentSessions(ctx, "user-1", "canvas-1", []string{"session-1"}, false) if err == nil { t.Fatal("expected owner error") } @@ -1136,7 +1146,7 @@ func TestDeleteAgentSessionsServiceRequiresOwner(t *testing.T) { } var count int64 - if err := dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { + if err = dao.DB.Model(&entity.API4Conversation{}).Where("id = ?", "session-1").Count(&count).Error; err != nil { t.Fatalf("failed to count session: %v", err) } if count != 1 { @@ -1841,7 +1851,8 @@ func TestGetAgentSession_RejectsIDOR(t *testing.T) { createAgentSessionTestCanvas(t, "agent-2", "user-1") createAgentSessionTestConversation(t, "session-1", "agent-1", "user-1", 1000) - data, code, err := NewAgentService().GetAgentSession("user-1", "agent-2", "session-1") + ctx := t.Context() + data, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "agent-2", "session-1") if err == nil { t.Fatal("expected non-nil error for cross-agent session access") } @@ -1863,7 +1874,8 @@ func TestGetAgentSession_SuccessWhenAgentMatches(t *testing.T) { createAgentSessionTestCanvas(t, "agent-1", "user-1") createAgentSessionTestConversation(t, "session-1", "agent-1", "user-1", 1000) - data, code, err := NewAgentService().GetAgentSession("user-1", "agent-1", "session-1") + ctx := t.Context() + data, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "agent-1", "session-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -1892,7 +1904,8 @@ func TestDeleteAgentSessionItem_RejectsIDOR(t *testing.T) { createAgentSessionTestCanvas(t, "agent-2", "user-1") createAgentSessionTestConversation(t, "session-1", "agent-1", "user-1", 1000) - deleted, code, err := NewAgentService().DeleteAgentSessionItem("user-1", "agent-2", "session-1") + ctx := t.Context() + deleted, code, err := NewAgentService().DeleteAgentSessionItem(ctx, "user-1", "agent-2", "session-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -1904,7 +1917,7 @@ func TestDeleteAgentSessionItem_RejectsIDOR(t *testing.T) { } // The session must still exist — the cross-agent delete was a no-op. - verify, _, err := NewAgentService().GetAgentSession("user-1", "agent-1", "session-1") + verify, _, err := NewAgentService().GetAgentSession(ctx, "user-1", "agent-1", "session-1") if err != nil { t.Fatalf("session should still exist for the legitimate owner: %v", err) } diff --git a/internal/service/api_token.go b/internal/service/api_token.go index 77a4164c8d..18ca3f532a 100644 --- a/internal/service/api_token.go +++ b/internal/service/api_token.go @@ -17,6 +17,7 @@ package service import ( + "context" "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/utility" @@ -37,9 +38,9 @@ type APIKeyResponse struct { } // ListAPIKeys list all API keys for a tenant -func (s *SystemService) ListAPIKeys(tenantID string) ([]*APIKeyResponse, error) { +func (s *SystemService) ListAPIKeys(ctx context.Context, tenantID string) ([]*APIKeyResponse, error) { APITokenDAO := dao.NewAPITokenDAO() - keys, err := APITokenDAO.GetByTenantID(tenantID) + keys, err := APITokenDAO.GetByTenantID(ctx, tenantID) if err != nil { return nil, err } @@ -49,7 +50,7 @@ func (s *SystemService) ListAPIKeys(tenantID string) ([]*APIKeyResponse, error) beta := key.Beta if beta == nil || *beta == "" { generatedBeta := utility.GenerateBetaAPIToken() - if err = dao.DB.Model(&entity.APIToken{}). + if err = dao.DB.WithContext(ctx).Model(&entity.APIToken{}). Where("tenant_id = ? AND token = ?", tenantID, key.Token). Updates(map[string]interface{}{ "beta": generatedBeta, @@ -82,7 +83,7 @@ type CreateAPIKeyRequest struct { } // CreateAPIKey creates a new API key for a tenant -func (s *SystemService) CreateAPIKey(tenantID string, req *CreateAPIKeyRequest) (*APIKeyResponse, error) { +func (s *SystemService) CreateAPIKey(ctx context.Context, tenantID string, req *CreateAPIKeyRequest) (*APIKeyResponse, error) { APITokenDAO := dao.NewAPITokenDAO() // Generate key and beta values @@ -97,7 +98,7 @@ func (s *SystemService) CreateAPIKey(tenantID string, req *CreateAPIKeyRequest) Beta: &betaAPIKey, } - if err := APITokenDAO.Create(APIKeyData); err != nil { + if err := APITokenDAO.Create(ctx, APIKeyData); err != nil { return nil, err } @@ -114,8 +115,8 @@ func (s *SystemService) CreateAPIKey(tenantID string, req *CreateAPIKeyRequest) }, nil } -func (s *SystemService) DeleteAPIKey(tenantID, key string) error { +func (s *SystemService) DeleteAPIKey(ctx context.Context, tenantID, key string) error { APITokenDAO := dao.NewAPITokenDAO() - _, err := APITokenDAO.DeleteByTenantIDAndToken(tenantID, key) + _, err := APITokenDAO.DeleteByTenantIDAndToken(ctx, tenantID, key) return err } diff --git a/internal/service/bot.go b/internal/service/bot.go index bf45035ce0..adabed4a84 100644 --- a/internal/service/bot.go +++ b/internal/service/bot.go @@ -83,7 +83,7 @@ func NewBotService(agentSvc *AgentService, llmSvc *LLMService) *BotService { func (s *BotService) ChatbotInfo(ctx context.Context, tenantID, dialogID string) ( title, avatar, prologue, llmID string, hasTavilyKey bool, ec common.ErrorCode, err error, ) { - dialog, err := s.chatDAO.GetDialogByID(dialogID) + dialog, err := s.chatDAO.GetDialogByID(ctx, dialogID) if err != nil { return "", "", "", "", false, common.CodeDataError, err } diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index c181dc5bb9..758289f379 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -302,7 +302,7 @@ func (s *BotService) ChatbotCompletion( // ChatSessionDAO.GetDialogByID already filters by status = "1" // so a returned row is valid; we still nil-check defensively // before dereferencing for symmetry with the session path. - dialog, err := s.chatDAO.GetDialogByID(dialogID) + dialog, err := s.chatDAO.GetDialogByID(ctx, dialogID) if err != nil || dialog == nil || dialog.TenantID != tenantID || dialog.Status == nil || *dialog.Status != common.StatusDialogValid { @@ -351,7 +351,7 @@ func (s *BotService) ChatbotCompletion( UserID: tenantID, Message: seedMsg, } - if err = s.api4ConversationDAO.Create(session); err != nil { + if err = s.api4ConversationDAO.Create(ctx, session); err != nil { return nil, common.CodeServerError, err } @@ -375,7 +375,7 @@ func (s *BotService) ChatbotCompletion( return out, common.CodeSuccess, nil } - session, err := s.api4ConversationDAO.GetBySessionID(req.SessionID, dialogID) + session, err := s.api4ConversationDAO.GetBySessionID(ctx, req.SessionID, dialogID) if err != nil { return nil, common.CodeServerError, err } @@ -508,7 +508,7 @@ func (s *BotService) ChatbotCompletion( // pipeline-level error ("**ERROR**" answer) nothing is // persisted, matching the python exception path. if !errored { - if pErr := s.persistChatbotTurn(session, req.Question, fullAnswer, messageID, finalRef); pErr != nil { + if pErr := s.persistChatbotTurn(ctx, session, req.Question, fullAnswer, messageID, finalRef); pErr != nil { common.Error("bot: ChatbotCompletion session update failed", pErr, zap.String("dialog_id", dialogID), @@ -570,7 +570,7 @@ func parseChatbotTurns(raw json.RawMessage) []map[string]any { // turn in its history. Mirrors python // API4ConversationService.append_message. func (s *BotService) persistChatbotTurn( - session *entity.API4Conversation, question, answer, messageID string, reference map[string]any, + ctx context.Context, session *entity.API4Conversation, question, answer, messageID string, reference map[string]any, ) error { // Serialise the read-modify-write per session and re-read the row // inside the lock: the caller's session was loaded before the @@ -580,7 +580,7 @@ func (s *BotService) persistChatbotTurn( lock := s.persistLock(session.ID) lock.Lock() defer lock.Unlock() - fresh, err := s.api4ConversationDAO.GetBySessionID(session.ID, session.DialogID) + fresh, err := s.api4ConversationDAO.GetBySessionID(ctx, session.ID, session.DialogID) if err != nil { return err } @@ -632,7 +632,7 @@ func (s *BotService) persistChatbotTurn( } session.Reference = rawRef - return s.api4ConversationDAO.Update(session) + return s.api4ConversationDAO.Update(ctx, session) } // normalizeBotBoolFlag coerces the JSON-encoded reasoning / internet diff --git a/internal/service/bot_completion_history_test.go b/internal/service/bot_completion_history_test.go index cb0bae88dc..a8db570656 100644 --- a/internal/service/bot_completion_history_test.go +++ b/internal/service/bot_completion_history_test.go @@ -335,7 +335,8 @@ func TestBotService_ChatbotCompletion_NewSessionSkipsLLM(t *testing.T) { // The persisted session must hold exactly the prologue turn — // no empty user message may be written. - row, derr := dao.NewAPI4ConversationDAO().GetBySessionID(got[0].SessionID, "dlg-1") + ctx := t.Context() + row, derr := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, got[0].SessionID, "dlg-1") if derr != nil || row == nil { t.Fatalf("session not persisted: row=%v err=%v", row, derr) } @@ -399,7 +400,8 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { UserID: "tenant-1", Message: seed, } - if err := dao.NewAPI4ConversationDAO().Create(sess); err != nil { + ctx := t.Context() + if err := dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { t.Fatalf("seed session: %v", err) } @@ -408,11 +410,11 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { "chunks": []any{map[string]any{"chunk_id": "c1"}}, "doc_aggs": []any{}, } - if err := svc.persistChatbotTurn(sess, "What is Go?", "A language.", "msg-p1", ref); err != nil { + if err := svc.persistChatbotTurn(ctx, sess, "What is Go?", "A language.", "msg-p1", ref); err != nil { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID("sess-p1", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p1", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -449,17 +451,18 @@ func TestPersistChatbotTurn_NilReferenceDefaultsToEmpty(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) + ctx := t.Context() sess := &entity.API4Conversation{ID: "sess-p2", DialogID: "dlg-p1", UserID: "tenant-1"} - if err := dao.NewAPI4ConversationDAO().Create(sess); err != nil { + if err := dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { t.Fatalf("seed session: %v", err) } svc := NewBotService(nil, nil) - if err := svc.persistChatbotTurn(sess, "q", "a", "msg-p2", nil); err != nil { + if err := svc.persistChatbotTurn(ctx, sess, "q", "a", "msg-p2", nil); err != nil { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID("sess-p2", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p2", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -494,8 +497,9 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { } sqlDB.SetMaxOpenConns(1) + ctx := t.Context() sess := &entity.API4Conversation{ID: "sess-p3", DialogID: "dlg-p1", UserID: "tenant-1"} - if err := dao.NewAPI4ConversationDAO().Create(sess); err != nil { + if err = dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { t.Fatalf("seed session: %v", err) } @@ -507,14 +511,14 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { wg.Add(1) go func(i int, q, a string) { defer wg.Done() - if err := svc.persistChatbotTurn(sess, q, a, fmt.Sprintf("msg-p3-%d", i), nil); err != nil { + if err := svc.persistChatbotTurn(ctx, sess, q, a, fmt.Sprintf("msg-p3-%d", i), nil); err != nil { t.Errorf("persistChatbotTurn %d: %v", i, err) } }(i, turn[0], turn[1]) } wg.Wait() - row, err := dao.NewAPI4ConversationDAO().GetBySessionID("sess-p3", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p3", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -522,7 +526,7 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { t.Fatalf("want both exchanges persisted (4 turns), got %d: %+v", len(turns), turns) } var refs []map[string]any - if err := json.Unmarshal(row.Reference, &refs); err != nil { + if err = json.Unmarshal(row.Reference, &refs); err != nil { t.Fatalf("reference decode: %v", err) } if len(refs) != 2 { diff --git a/internal/service/chat.go b/internal/service/chat.go index a4d2613054..d479ee2f61 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -17,6 +17,7 @@ package service import ( + "context" "encoding/json" "errors" "fmt" @@ -78,13 +79,14 @@ type ListChatsResponse struct { } // ListChats list chats for a user -func (s *ChatService) ListChats(userID, status, keywords string, page, pageSize int, orderBy string, desc bool, ownerIDs []string) (*ListChatsResponse, error) { +func (s *ChatService) ListChats(ctx context.Context, userID, status, keywords string, page, pageSize int, orderBy string, desc bool, ownerIDs []string) (*ListChatsResponse, error) { var chats []*entity.ChatListItem var total int64 var err error if len(ownerIDs) == 0 { chats, total, err = s.chatDAO.ListByTenantIDs( + ctx, nil, userID, page, @@ -97,7 +99,8 @@ func (s *ChatService) ListChats(userID, status, keywords string, page, pageSize return nil, err } } else { - filterOwnerIDs, err := s.filterAccessibleChatOwnerIDs(userID, ownerIDs) + var filterOwnerIDs []string + filterOwnerIDs, err = s.filterAccessibleChatOwnerIDs(userID, ownerIDs) if err != nil { return nil, err } @@ -108,7 +111,7 @@ func (s *ChatService) ListChats(userID, status, keywords string, page, pageSize }, nil } - chats, total, err = s.chatDAO.ListByOwnerIDs(filterOwnerIDs, userID, orderBy, desc, keywords) + chats, total, err = s.chatDAO.ListByOwnerIDs(ctx, filterOwnerIDs, userID, orderBy, desc, keywords) if err != nil { return nil, err } @@ -203,7 +206,7 @@ type CreateChatRequest struct { TenantID *string `json:"tenant_id"` } -func (s *ChatService) Create(userID string, req map[string]interface{}) (map[string]interface{}, common.ErrorCode, error) { +func (s *ChatService) Create(ctx context.Context, userID string, req map[string]interface{}) (map[string]interface{}, common.ErrorCode, error) { tenant, err := s.tenantDAO.GetByID(userID) if err != nil { return nil, common.CodeDataError, errors.New("tenant not found") @@ -317,7 +320,7 @@ func (s *ChatService) Create(userID string, req map[string]interface{}) (map[str applyCreatePromptDefaults(req) filterCreateChatPersistedFields(req) - exists, err := s.chatDAO.ExistsByNameTenantStatus(name, userID, string(entity.StatusValid)) + exists, err := s.chatDAO.ExistsByNameTenantStatus(ctx, name, userID, string(entity.StatusValid)) if err != nil { return nil, common.CodeServerError, err } @@ -326,11 +329,11 @@ func (s *ChatService) Create(userID string, req map[string]interface{}) (map[str } chat := buildCreateChatEntity(req, userID) - if err = s.chatDAO.Create(chat); err != nil { + if err = s.chatDAO.Create(ctx, chat); err != nil { return nil, common.CodeDataError, errors.New("failed to create chat") } - chat, err = s.chatDAO.GetByID(chat.ID) + chat, err = s.chatDAO.GetByID(ctx, chat.ID) if err != nil { return nil, common.CodeDataError, errors.New("failed to retrieve created chat") } @@ -781,8 +784,8 @@ func (s *ChatService) splitModelNameAndFactory(embeddingModelID string) string { return embeddingModelID } -func (s *ChatService) getOwnedValidChat(userID, chatID string) (*entity.Chat, error) { - chat, err := s.chatDAO.GetByIDAndStatus(chatID, string(entity.StatusValid)) +func (s *ChatService) getOwnedValidChat(ctx context.Context, userID, chatID string) (*entity.Chat, error) { + chat, err := s.chatDAO.GetByIDAndStatus(ctx, chatID, string(entity.StatusValid)) if err != nil { return nil, errors.New("no authorization") } @@ -830,17 +833,17 @@ var defaultRerankModels = map[string]struct{}{ } // UpdateChat mirrors PUT /api/v1/chats/ in the Python REST API. -func (s *ChatService) UpdateChat(userID, chatID string, req map[string]interface{}) (map[string]interface{}, error) { - return s.updateChatREST(userID, chatID, req, false) +func (s *ChatService) UpdateChat(ctx context.Context, userID, chatID string, req map[string]interface{}) (map[string]interface{}, error) { + return s.updateChatREST(ctx, userID, chatID, req, false) } // PatchChat mirrors PATCH /api/v1/chats/ in the Python REST API. -func (s *ChatService) PatchChat(userID, chatID string, req map[string]interface{}) (map[string]interface{}, error) { - return s.updateChatREST(userID, chatID, req, true) +func (s *ChatService) PatchChat(ctx context.Context, userID, chatID string, req map[string]interface{}) (map[string]interface{}, error) { + return s.updateChatREST(ctx, userID, chatID, req, true) } -func (s *ChatService) updateChatREST(userID, chatID string, req map[string]interface{}, patch bool) (map[string]interface{}, error) { - currentChat, err := s.getOwnedValidChat(userID, chatID) +func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, req map[string]interface{}, patch bool) (map[string]interface{}, error) { + currentChat, err := s.getOwnedValidChat(ctx, userID, chatID) if err != nil { return nil, err } @@ -948,7 +951,7 @@ func (s *ChatService) updateChatREST(userID, chatID string, req map[string]inter currentName = *currentChat.Name } if strings.ToLower(name) != strings.ToLower(currentName) { - existingNames, err := s.chatDAO.GetExistingNames(userID, string(entity.StatusValid)) + existingNames, err := s.chatDAO.GetExistingNames(ctx, userID, string(entity.StatusValid)) if err != nil { return nil, err } @@ -961,7 +964,7 @@ func (s *ChatService) updateChatREST(userID, chatID string, req map[string]inter } if len(updates) > 0 { - if err = s.chatDAO.UpdateByID(chatID, updates); err != nil { + if err = s.chatDAO.UpdateByID(ctx, chatID, updates); err != nil { if patch { return nil, errors.New("failed to update chat") } @@ -969,7 +972,7 @@ func (s *ChatService) updateChatREST(userID, chatID string, req map[string]inter } } - updatedChat, err := s.chatDAO.GetByID(chatID) + updatedChat, err := s.chatDAO.GetByID(ctx, chatID) if err != nil { return nil, errors.New("failed to retrieve updated chat") } @@ -1150,11 +1153,11 @@ func (s *ChatService) buildRESTChatResponse(chat *entity.Chat) map[string]interf } // DeleteChat soft deletes a single chat owned by the current user. -func (s *ChatService) DeleteChat(userID, chatID string) error { - if _, err := s.getOwnedValidChat(userID, chatID); err != nil { +func (s *ChatService) DeleteChat(ctx context.Context, userID, chatID string) error { + if _, err := s.getOwnedValidChat(ctx, userID, chatID); err != nil { return err } - if err := s.chatDAO.UpdateByID(chatID, map[string]interface{}{ + if err := s.chatDAO.UpdateByID(ctx, chatID, map[string]interface{}{ "status": string(entity.StatusInvalid), }); err != nil { return fmt.Errorf("failed to delete chat %s", chatID) @@ -1195,10 +1198,10 @@ func checkDuplicateChatIDs(ids []string) ([]string, []string) { } // BulkDeleteChats soft deletes chats owned by the current user with partial success semantics. -func (s *ChatService) BulkDeleteChats(userID string, req *BulkDeleteChatsRequest) (map[string]interface{}, error) { +func (s *ChatService) BulkDeleteChats(ctx context.Context, userID string, req *BulkDeleteChatsRequest) (map[string]interface{}, error) { ids := req.IDs if len(ids) == 0 && req.DeleteAll { - chats, err := s.chatDAO.ListByTenantID(userID, string(entity.StatusValid)) + chats, err := s.chatDAO.ListByTenantID(ctx, userID, string(entity.StatusValid)) if err != nil { return nil, err } @@ -1216,11 +1219,11 @@ func (s *ChatService) BulkDeleteChats(userID string, req *BulkDeleteChatsRequest successCount := 0 for _, chatID := range uniqueIDs { - if _, err := s.getOwnedValidChat(userID, chatID); err != nil { + if _, err := s.getOwnedValidChat(ctx, userID, chatID); err != nil { errorsList = append(errorsList, fmt.Sprintf("Chat(%s) not found.", chatID)) continue } - if err := s.chatDAO.UpdateByID(chatID, map[string]interface{}{ + if err := s.chatDAO.UpdateByID(ctx, chatID, map[string]interface{}{ "status": string(entity.StatusInvalid), }); err != nil { errorsList = append(errorsList, fmt.Sprintf("Failed to delete chat %s", chatID)) @@ -1261,7 +1264,7 @@ type GetChatResponse struct { } // GetChat gets chat detail by ID with permission check -func (s *ChatService) GetChat(userID string, chatID string) (*GetChatResponse, error) { +func (s *ChatService) GetChat(ctx context.Context, userID string, chatID string) (*GetChatResponse, error) { // Step 1: Get user tenants (same as Python UserTenantService.query(user_id=current_user.id)) tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -1272,7 +1275,7 @@ func (s *ChatService) GetChat(userID string, chatID string) (*GetChatResponse, e // Python: for tenant in tenants: if DialogService.query(tenant_id=tenant.tenant_id, id=chat_id, status=StatusEnum.VALID.value): break hasPermission := false for _, tenant := range tenants { - chats, err := s.chatDAO.QueryByTenantIDAndID(tenant.TenantID, chatID, "1") + chats, err := s.chatDAO.QueryByTenantIDAndID(ctx, tenant.TenantID, chatID, "1") if err != nil { continue // Try next tenant } @@ -1287,7 +1290,7 @@ func (s *ChatService) GetChat(userID string, chatID string) (*GetChatResponse, e } // Step 3: Get chat detail (same as Python DialogService.get_by_id(chat_id)) - chat, err := s.chatDAO.GetByID(chatID) + chat, err := s.chatDAO.GetByID(ctx, chatID) if err != nil { return nil, fmt.Errorf("chat not found") } diff --git a/internal/service/chat_channel.go b/internal/service/chat_channel.go index f047e300f8..b6b2773775 100644 --- a/internal/service/chat_channel.go +++ b/internal/service/chat_channel.go @@ -15,6 +15,7 @@ package service import ( + "context" "errors" "fmt" "ragflow/internal/utility" @@ -38,7 +39,7 @@ func NewChatChannelService() *ChatChannelService { } } -func (s *ChatChannelService) Insert(channel *entity.ChatChannel) error { +func (s *ChatChannelService) Insert(ctx context.Context, channel *entity.ChatChannel) error { if channel == nil { return errors.New("channel is nil") } @@ -48,23 +49,23 @@ func (s *ChatChannelService) Insert(channel *entity.ChatChannel) error { if channel.Status == 0 { channel.Status = 1 } - return s.chatChannelDAO.Create(channel) + return s.chatChannelDAO.Create(ctx, channel) } -func (s *ChatChannelService) GetByID(id string) (*entity.ChatChannel, error) { +func (s *ChatChannelService) GetByID(ctx context.Context, id string) (*entity.ChatChannel, error) { if id == "" { return nil, errors.New("id is empty") } - return s.chatChannelDAO.GetByIDOnly(id) + return s.chatChannelDAO.GetByIDOnly(ctx, id) } -func (s *ChatChannelService) List(tenantID string) ([]*entity.ChatChannelListResponse, error) { - return s.chatChannelDAO.ListByTenantID(tenantID) +func (s *ChatChannelService) List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) { + return s.chatChannelDAO.ListByTenantID(ctx, tenantID) } -func (s *ChatChannelService) CreateChatChannel(tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) { +func (s *ChatChannelService) CreateChatChannel(ctx context.Context, tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) { if chatID != nil && *chatID != "" { - dialog, err := s.chatDAO.GetByID(*chatID) + dialog, err := s.chatDAO.GetByID(ctx, *chatID) if err != nil { if dao.IsNotFoundErr(err) { return nil, errors.New("Can't find this chat assistant!") @@ -85,19 +86,19 @@ func (s *ChatChannelService) CreateChatChannel(tenantID, name, channelType strin Status: 1, } - if err := s.Insert(row); err != nil { + if err := s.Insert(ctx, row); err != nil { return nil, err } - created, err := s.GetByID(row.ID) + created, err := s.GetByID(ctx, row.ID) if err != nil { return nil, fmt.Errorf("failed to load created chat channel: %w", err) } return created, nil } -func (s *ChatChannelService) accessible(userID, channelID string) (*entity.ChatChannel, bool, error) { - channel, err := s.chatChannelDAO.GetByIDOnly(channelID) +func (s *ChatChannelService) accessible(ctx context.Context, userID, channelID string) (*entity.ChatChannel, bool, error) { + channel, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) if err != nil { if dao.IsNotFoundErr(err) { return nil, false, nil @@ -122,8 +123,8 @@ func (s *ChatChannelService) accessible(userID, channelID string) (*entity.ChatC return channel, false, nil } -func (s *ChatChannelService) GetChatChannel(userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) { - _, ok, err := s.accessible(userID, channelID) +func (s *ChatChannelService) GetChatChannel(ctx context.Context, userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) { + _, ok, err := s.accessible(ctx, userID, channelID) if err != nil { return nil, common.CodeServerError, err } @@ -131,7 +132,7 @@ func (s *ChatChannelService) GetChatChannel(userID, channelID string) (*entity.C return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - channel, err := s.chatChannelDAO.GetByIDOnly(channelID) + channel, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Can't find this chat channel!") @@ -141,8 +142,8 @@ func (s *ChatChannelService) GetChatChannel(userID, channelID string) (*entity.C return channel, common.CodeSuccess, nil } -func (s *ChatChannelService) UpdateChatChannel(userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) { - channel, ok, err := s.accessible(userID, channelID) +func (s *ChatChannelService) UpdateChatChannel(ctx context.Context, userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) { + channel, ok, err := s.accessible(ctx, userID, channelID) if err != nil { return nil, common.CodeServerError, err } @@ -184,7 +185,7 @@ func (s *ChatChannelService) UpdateChatChannel(userID, channelID string, req map return nil, common.CodeDataError, errors.New("chat_id must be string or null") } if chatID != "" { - dialog, err := s.chatDAO.GetByID(chatID) + dialog, err := s.chatDAO.GetByID(ctx, chatID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Can't find this chat assistant!") @@ -200,12 +201,12 @@ func (s *ChatChannelService) UpdateChatChannel(userID, channelID string, req map } if len(updates) > 0 { - if err := s.chatChannelDAO.UpdateByID(channelID, channel.TenantID, updates); err != nil { + if err := s.chatChannelDAO.UpdateByID(ctx, channelID, channel.TenantID, updates); err != nil { return nil, common.CodeDataError, err } } - updated, err := s.chatChannelDAO.GetByIDOnly(channelID) + updated, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Can't find this chat channel!") @@ -215,8 +216,8 @@ func (s *ChatChannelService) UpdateChatChannel(userID, channelID string, req map return updated, common.CodeSuccess, nil } -func (s *ChatChannelService) DeleteChatChannel(userID, channelID string) (bool, common.ErrorCode, error) { - channel, ok, err := s.accessible(userID, channelID) +func (s *ChatChannelService) DeleteChatChannel(ctx context.Context, userID, channelID string) (bool, common.ErrorCode, error) { + channel, ok, err := s.accessible(ctx, userID, channelID) if err != nil { return false, common.CodeServerError, err } @@ -227,7 +228,7 @@ func (s *ChatChannelService) DeleteChatChannel(userID, channelID string) (bool, return false, common.CodeAuthenticationError, errors.New("No authorization.") } - if err := s.chatChannelDAO.DeleteByID(channelID, channel.TenantID); err != nil { + if err = s.chatChannelDAO.DeleteByID(ctx, channelID, channel.TenantID); err != nil { return false, common.CodeDataError, err } return true, common.CodeSuccess, nil diff --git a/internal/service/chat_channel_test.go b/internal/service/chat_channel_test.go index cc241c8745..1117fa7867 100644 --- a/internal/service/chat_channel_test.go +++ b/internal/service/chat_channel_test.go @@ -87,7 +87,9 @@ func TestChatChannelServiceCreateAndList(t *testing.T) { svc := NewChatChannelService() chatID := "dialog-1" + ctx := t.Context() channel, err := svc.CreateChatChannel( + ctx, "tenant-1", "bot-a", "dingtalk", @@ -104,7 +106,7 @@ func TestChatChannelServiceCreateAndList(t *testing.T) { t.Fatalf("unexpected created channel: %+v", channel) } - list, err := svc.List("tenant-1") + list, err := svc.List(ctx, "tenant-1") if err != nil { t.Fatalf("List failed: %v", err) } @@ -128,7 +130,8 @@ func TestChatChannelServiceGetChatChannelAllowsJoinedTenant(t *testing.T) { }) createServiceTestMembership(t, db, "user-2", "tenant-1") - channel, code, err := NewChatChannelService().GetChatChannel("user-2", "cc-1") + ctx := t.Context() + channel, code, err := NewChatChannelService().GetChatChannel(ctx, "user-2", "cc-1") if err != nil { t.Fatalf("GetChatChannel failed: %v", err) } @@ -155,7 +158,8 @@ func TestChatChannelServiceUpdateChatChannelSuccess(t *testing.T) { Status: 1, }) - updated, code, err := NewChatChannelService().UpdateChatChannel("tenant-1", "cc-1", map[string]interface{}{ + ctx := t.Context() + updated, code, err := NewChatChannelService().UpdateChatChannel(ctx, "tenant-1", "cc-1", map[string]interface{}{ "name": "bot-b", "config": map[string]interface{}{"token": "new"}, "chat_id": "dialog-2", @@ -189,7 +193,8 @@ func TestChatChannelServiceUpdateChatChannelRejectsCrossTenantDialog(t *testing. Status: 1, }) - _, code, err := NewChatChannelService().UpdateChatChannel("tenant-1", "cc-1", map[string]interface{}{ + ctx := t.Context() + _, code, err := NewChatChannelService().UpdateChatChannel(ctx, "tenant-1", "cc-1", map[string]interface{}{ "chat_id": "dialog-2", }) if code != common.CodeAuthenticationError { @@ -213,7 +218,8 @@ func TestChatChannelServiceDeleteChatChannel(t *testing.T) { svc := NewChatChannelService() - deleted, code, err := svc.DeleteChatChannel("user-2", "cc-1") + ctx := t.Context() + deleted, code, err := svc.DeleteChatChannel(ctx, "user-2", "cc-1") if code != common.CodeAuthenticationError { t.Fatalf("expected authentication error, got %v", code) } @@ -224,7 +230,7 @@ func TestChatChannelServiceDeleteChatChannel(t *testing.T) { t.Fatal("expected delete to be rejected") } - deleted, code, err = svc.DeleteChatChannel("tenant-1", "cc-1") + deleted, code, err = svc.DeleteChatChannel(ctx, "tenant-1", "cc-1") if err != nil { t.Fatalf("DeleteChatChannel failed: %v", err) } @@ -232,7 +238,7 @@ func TestChatChannelServiceDeleteChatChannel(t *testing.T) { t.Fatalf("unexpected delete result: deleted=%v code=%v err=%v", deleted, code, err) } - if _, err := svc.GetByID("cc-1"); err == nil { + if _, err = svc.GetByID(ctx, "cc-1"); err == nil { t.Fatal("expected deleted record to be gone") } } diff --git a/internal/service/chat_delete_test.go b/internal/service/chat_delete_test.go index 6a0b70c1fc..6b6a8dbac9 100644 --- a/internal/service/chat_delete_test.go +++ b/internal/service/chat_delete_test.go @@ -60,12 +60,13 @@ func TestChatServiceDeleteChatRejectsNonOwner(t *testing.T) { createChatDeleteServiceTestChat(t, db, "chat-1", "tenant-1") svc := NewChatService() - err := svc.DeleteChat("user-1", "chat-1") + ctx := t.Context() + err := svc.DeleteChat(ctx, "user-1", "chat-1") if err == nil || err.Error() != "no authorization" { t.Fatalf("expected no authorization, got %v", err) } - chat, getErr := svc.chatDAO.GetByID("chat-1") + chat, getErr := svc.chatDAO.GetByID(ctx, "chat-1") if getErr != nil { t.Fatalf("failed to fetch chat: %v", getErr) } @@ -81,7 +82,8 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) createChatDeleteServiceTestChat(t, db, "chat-3", "tenant-2") svc := NewChatService() - result, err := svc.BulkDeleteChats("user-1", &BulkDeleteChatsRequest{DeleteAll: true}) + ctx := t.Context() + result, err := svc.BulkDeleteChats(ctx, "user-1", &BulkDeleteChatsRequest{DeleteAll: true}) if err != nil { t.Fatalf("BulkDeleteChats failed: %v", err) } @@ -90,7 +92,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected success_count 2, got %+v", result["success_count"]) } - owned1, err := svc.chatDAO.GetByID("chat-1") + owned1, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch chat-1: %v", err) } @@ -98,7 +100,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected chat-1 invalid, got %+v", owned1.Status) } - owned2, err := svc.chatDAO.GetByID("chat-2") + owned2, err := svc.chatDAO.GetByID(ctx, "chat-2") if err != nil { t.Fatalf("failed to fetch chat-2: %v", err) } @@ -106,7 +108,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected chat-2 invalid, got %+v", owned2.Status) } - other, err := svc.chatDAO.GetByID("chat-3") + other, err := svc.chatDAO.GetByID(ctx, "chat-3") if err != nil { t.Fatalf("failed to fetch chat-3: %v", err) } @@ -121,7 +123,8 @@ func TestChatServiceBulkDeleteChatsReturnsPartialSuccessErrors(t *testing.T) { createChatDeleteServiceTestChat(t, db, "chat-2", "tenant-2") svc := NewChatService() - result, err := svc.BulkDeleteChats("user-1", &BulkDeleteChatsRequest{ + ctx := t.Context() + result, err := svc.BulkDeleteChats(ctx, "user-1", &BulkDeleteChatsRequest{ IDs: []string{"chat-1", "chat-2", "chat-1"}, }) if err != nil { diff --git a/internal/service/chat_list_test.go b/internal/service/chat_list_test.go index 62022d159c..f6a9eb78ce 100644 --- a/internal/service/chat_list_test.go +++ b/internal/service/chat_list_test.go @@ -82,7 +82,8 @@ func TestChatServiceListChatsDefaultReturnsAllWithCorrectTotal(t *testing.T) { createChatListTestChat(t, db, "chat-3", "user-1", "list_test_2") svc := NewChatService() - result, err := svc.ListChats("user-1", "1", "", 0, 0, "create_time", true, nil) + ctx := t.Context() + result, err := svc.ListChats(ctx, "user-1", "1", "", 0, 0, "create_time", true, nil) if err != nil { t.Fatalf("ListChats failed: %v", err) } @@ -123,7 +124,8 @@ func TestChatServiceListChatsFiltersByOwnerIDs(t *testing.T) { createChatListTestChat(t, db, "chat-team", "tenant-2", "team_chat") svc := NewChatService() - result, err := svc.ListChats("user-1", "1", "", 0, 0, "create_time", true, []string{"tenant-2"}) + ctx := t.Context() + result, err := svc.ListChats(ctx, "user-1", "1", "", 0, 0, "create_time", true, []string{"tenant-2"}) if err != nil { t.Fatalf("ListChats failed: %v", err) } @@ -146,7 +148,8 @@ func TestChatServiceListChatsKeywordFiltersCorrectly(t *testing.T) { svc := NewChatService() - exactResult, err := svc.ListChats("user-1", "1", "list_keyword_1", 0, 0, "create_time", true, nil) + ctx := t.Context() + exactResult, err := svc.ListChats(ctx, "user-1", "1", "list_keyword_1", 0, 0, "create_time", true, nil) if err != nil { t.Fatalf("ListChats keyword exact failed: %v", err) } @@ -157,7 +160,7 @@ func TestChatServiceListChatsKeywordFiltersCorrectly(t *testing.T) { t.Fatalf("expected chat name 'list_keyword_1', got %+v", exactResult.Chats[0].Name) } - unknownResult, err := svc.ListChats("user-1", "1", "unknown_keyword", 0, 0, "create_time", true, nil) + unknownResult, err := svc.ListChats(ctx, "user-1", "1", "unknown_keyword", 0, 0, "create_time", true, nil) if err != nil { t.Fatalf("ListChats unknown keyword failed: %v", err) } @@ -165,7 +168,7 @@ func TestChatServiceListChatsKeywordFiltersCorrectly(t *testing.T) { t.Fatalf("expected 0 chats for unknown keyword, got %d", len(unknownResult.Chats)) } - partialResult, err := svc.ListChats("user-1", "1", "list_keyword", 0, 0, "create_time", true, nil) + partialResult, err := svc.ListChats(ctx, "user-1", "1", "list_keyword", 0, 0, "create_time", true, nil) if err != nil { t.Fatalf("ListChats partial keyword failed: %v", err) } @@ -182,7 +185,8 @@ func TestChatServiceListChatsPagination(t *testing.T) { svc := NewChatService() - page1, err := svc.ListChats("user-1", "1", "", 1, 2, "create_time", true, nil) + ctx := t.Context() + page1, err := svc.ListChats(ctx, "user-1", "1", "", 1, 2, "create_time", true, nil) if err != nil { t.Fatalf("ListChats page 1 failed: %v", err) } @@ -193,7 +197,7 @@ func TestChatServiceListChatsPagination(t *testing.T) { t.Fatalf("expected total=5, got %d", page1.Total) } - page3, err := svc.ListChats("user-1", "1", "", 3, 2, "create_time", true, nil) + page3, err := svc.ListChats(ctx, "user-1", "1", "", 3, 2, "create_time", true, nil) if err != nil { t.Fatalf("ListChats page 3 failed: %v", err) } @@ -211,7 +215,8 @@ func TestChatServiceListChatsExcludesDeletedChats(t *testing.T) { db.Model(&entity.Chat{}).Where("id = ?", "chat-2").Update("status", invalidStatus) svc := NewChatService() - result, err := svc.ListChats("user-1", "1", "", 0, 0, "create_time", true, nil) + ctx := t.Context() + result, err := svc.ListChats(ctx, "user-1", "1", "", 0, 0, "create_time", true, nil) if err != nil { t.Fatalf("ListChats failed: %v", err) } diff --git a/internal/service/chat_rest_update_test.go b/internal/service/chat_rest_update_test.go index b9084cdcbf..6d8c9693a4 100644 --- a/internal/service/chat_rest_update_test.go +++ b/internal/service/chat_rest_update_test.go @@ -125,7 +125,8 @@ func TestChatServiceCreateDefaultsMetaDataFilter(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - resp, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + resp, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "created chat", }) if err != nil { @@ -140,7 +141,7 @@ func TestChatServiceCreateDefaultsMetaDataFilter(t *testing.T) { if !ok || chatID == "" { t.Fatalf("expected created chat id, got %+v", resp["id"]) } - chat, err := svc.chatDAO.GetByID(chatID) + chat, err := svc.chatDAO.GetByID(ctx, chatID) if err != nil { t.Fatalf("failed to fetch created chat: %v", err) } @@ -154,7 +155,8 @@ func TestChatServiceCreateAcceptsNilMetaDataFilter(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - resp, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + resp, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "created chat", "meta_data_filter": nil, }) @@ -171,7 +173,8 @@ func TestChatServiceCreateRejectsInvalidMetaDataFilter(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - _, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + _, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "created chat", "meta_data_filter": []interface{}{"invalid"}, }) @@ -188,7 +191,8 @@ func TestChatServicePatchChatMergesPromptConfigAndLLMSetting(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - resp, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + resp, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{ "prompt_config": map[string]interface{}{"quote": false}, "llm_setting": map[string]interface{}{"temperature": float64(0.2)}, }) @@ -202,7 +206,7 @@ func TestChatServicePatchChatMergesPromptConfigAndLLMSetting(t *testing.T) { t.Fatalf("response should expose dataset_ids: %+v", resp) } - chat, err := svc.chatDAO.GetByID("chat-1") + chat, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch updated chat: %v", err) } @@ -225,7 +229,8 @@ func TestChatServiceUpdateChatRejectsTenantID(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{ "tenant_id": "tenant-2", }) if err == nil || err.Error() != "`tenant_id` must not be provided." { @@ -238,14 +243,15 @@ func TestChatServiceUpdateChatRejectsInvalidLLMSetting(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{ "llm_setting": "invalid", }) if err == nil || err.Error() != "`llm_setting` should be an object" { t.Fatalf("expected llm_setting error, got %v", err) } - chat, err := svc.chatDAO.GetByID("chat-1") + chat, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -259,7 +265,8 @@ func TestChatServiceUpdateChatAcceptsMetaDataFilterObject(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{ "name": "chat-chat-1", "meta_data_filter": map[string]interface{}{ "method": "disabled", @@ -270,7 +277,7 @@ func TestChatServiceUpdateChatAcceptsMetaDataFilterObject(t *testing.T) { t.Fatalf("UpdateChat failed: %v", err) } - chat, err := svc.chatDAO.GetByID("chat-1") + chat, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -284,7 +291,8 @@ func TestChatServiceUpdateChatBackfillsNilMetaDataFilter(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - resp, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + resp, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{ "name": "chat-chat-1", }) if err != nil { @@ -292,7 +300,7 @@ func TestChatServiceUpdateChatBackfillsNilMetaDataFilter(t *testing.T) { } assertEmptyMetaDataFilter(t, resp["meta_data_filter"]) - chat, err := svc.chatDAO.GetByID("chat-1") + chat, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -307,7 +315,8 @@ func TestChatServicePatchChatIgnoresTenantIDAndUpdatesName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{ + ctx := t.Context() + _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{ "tenant_id": "tenant-2", "name": " renamed chat ", }) @@ -315,7 +324,7 @@ func TestChatServicePatchChatIgnoresTenantIDAndUpdatesName(t *testing.T) { t.Fatalf("PatchChat failed: %v", err) } - chat, err := svc.chatDAO.GetByID("chat-1") + chat, err := svc.chatDAO.GetByID(ctx, "chat-1") if err != nil { t.Fatalf("failed to fetch updated chat: %v", err) } @@ -331,7 +340,8 @@ func TestChatServiceCreateValidatesName(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - _, code, err := svc.Create("user-1", map[string]interface{}{"name": " "}) + ctx := t.Context() + _, code, err := svc.Create(ctx, "user-1", map[string]interface{}{"name": " "}) if err == nil { t.Fatal("expected name validation error") } @@ -348,7 +358,8 @@ func TestChatServiceCreateRejectsDuplicateName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, code, err := svc.Create("user-1", map[string]interface{}{"name": "chat-chat-1"}) + ctx := t.Context() + _, code, err := svc.Create(ctx, "user-1", map[string]interface{}{"name": "chat-chat-1"}) if err == nil { t.Fatal("expected duplicate name error") } @@ -365,7 +376,8 @@ func TestChatServiceUpdateValidatesName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{"name": " "}) + ctx := t.Context() + _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": " "}) if err == nil { t.Fatal("expected name validation error") } @@ -380,7 +392,8 @@ func TestChatServiceUpdateRejectsDuplicateName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-2", "user-1") svc := NewChatService() - _, err := svc.UpdateChat("user-1", "chat-1", map[string]interface{}{"name": "chat-chat-2"}) + ctx := t.Context() + _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": "chat-chat-2"}) if err == nil { t.Fatal("expected duplicate name error") } @@ -394,7 +407,8 @@ func TestChatServicePatchChatRejectsEmptyName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{"name": ""}) + ctx := t.Context() + _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": ""}) if err == nil || err.Error() != "`name` cannot be empty." { t.Fatalf("expected empty name error, got %v", err) } @@ -405,7 +419,8 @@ func TestChatServicePatchChatRejectsNonStringName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() - _, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{"name": 123}) + ctx := t.Context() + _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": 123}) if err == nil || err.Error() != "Chat name must be a string." { t.Fatalf("expected non-string name error, got %v", err) } @@ -416,8 +431,9 @@ func TestChatServicePatchChatRejectsTooLongName(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-1", "user-1") svc := NewChatService() + ctx := t.Context() longName := strings.Repeat("a", 256) - _, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{"name": longName}) + _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": longName}) if err == nil { t.Fatal("expected too long name error") } @@ -432,7 +448,8 @@ func TestChatServiceUpdateRejectsDuplicateNameCaseInsensitive(t *testing.T) { createChatRESTUpdateServiceTestChat(t, db, "chat-2", "user-1") svc := NewChatService() - _, err := svc.PatchChat("user-1", "chat-1", map[string]interface{}{"name": "CHAT-CHAT-2"}) + ctx := t.Context() + _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": "CHAT-CHAT-2"}) if err == nil || err.Error() != "Duplicated chat name." { t.Fatalf("expected case-insensitive duplicate name error, got %v", err) } @@ -442,7 +459,8 @@ func TestChatServiceCreateRejectsTenantID(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - _, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + _, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "valid chat", "tenant_id": "other-tenant", }) @@ -458,7 +476,8 @@ func TestChatServiceCreateRejectsInvalidPromptConfig(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - _, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + _, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "valid chat", "prompt_config": "invalid", }) @@ -474,7 +493,8 @@ func TestChatServiceCreatePromptDefaultsContract(t *testing.T) { setupChatRESTUpdateServiceTestDB(t) svc := NewChatService() - resp, code, err := svc.Create("user-1", map[string]interface{}{ + ctx := t.Context() + resp, code, err := svc.Create(ctx, "user-1", map[string]interface{}{ "name": "prompt defaults chat", }) if err != nil { diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 8afedc3f15..3dfc363c02 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -41,14 +41,14 @@ import ( // Interfaces for testability — satisfied by the concrete DAO/pipeline types. type chatSessionStore interface { - GetByID(id string) (*entity.ChatSession, error) - GetBySessionIDAndChatID(sessionID, chatID string) (*entity.ChatSession, error) - Create(conv *entity.ChatSession) error - UpdateByID(id string, updates map[string]interface{}) error - DeleteByID(id string) error - ListByChatID(chatID string) ([]*entity.ChatSession, error) - GetDialogByID(chatID string) (*entity.Chat, error) - CheckDialogExists(tenantID, chatID string) (bool, error) + GetByID(ctx context.Context, id string) (*entity.ChatSession, error) + GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) + Create(ctx context.Context, conv *entity.ChatSession) error + UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error + DeleteByID(ctx context.Context, id string) error + ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) + GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) + CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) } type userTenantStore interface { @@ -116,7 +116,7 @@ type SetChatSessionResponse struct { // SetChatSession creates or updates a chat session. // Kept as a compatibility entrypoint for older chat-session callers. -func (s *ChatSessionService) SetChatSession(userID string, req *SetChatSessionRequest) (*SetChatSessionResponse, error) { +func (s *ChatSessionService) SetChatSession(ctx context.Context, userID string, req *SetChatSessionRequest) (*SetChatSessionResponse, error) { name := req.Name if name == "" { name = "New chat session" @@ -130,19 +130,19 @@ func (s *ChatSessionService) SetChatSession(userID string, req *SetChatSessionRe "name": name, "user_id": userID, } - if err := s.chatSessionDAO.UpdateByID(req.SessionID, updates); err != nil { - return nil, errors.New("Chat session not found") + if err := s.chatSessionDAO.UpdateByID(ctx, req.SessionID, updates); err != nil { + return nil, errors.New("chat session not found") } - session, err := s.chatSessionDAO.GetByID(req.SessionID) + session, err := s.chatSessionDAO.GetByID(ctx, req.SessionID) if err != nil { - return nil, errors.New("Fail to update a chat session") + return nil, errors.New("fail to update a chat session") } return &SetChatSessionResponse{ChatSession: session}, nil } - dialog, err := s.chatSessionDAO.GetDialogByID(req.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, req.DialogID) if err != nil { - return nil, errors.New("Dialog not found") + return nil, errors.New("dialog not found") } prologue := "Hi! I'm your assistant. What can I do for you?" @@ -167,8 +167,8 @@ func (s *ChatSessionService) SetChatSession(userID string, req *SetChatSessionRe UserID: &userID, Reference: referenceJSON, } - if err := s.chatSessionDAO.Create(session); err != nil { - return nil, errors.New("Fail to create a chat session") + if err = s.chatSessionDAO.Create(ctx, session); err != nil { + return nil, errors.New("fail to create a chat session") } return &SetChatSessionResponse{ChatSession: session}, nil @@ -176,7 +176,7 @@ func (s *ChatSessionService) SetChatSession(userID string, req *SetChatSessionRe // RemoveChatSessions removes chat sessions. // Kept as a compatibility entrypoint for older chat-session callers. -func (s *ChatSessionService) RemoveChatSessions(userID string, chatSessions []string) error { +func (s *ChatSessionService) RemoveChatSessions(ctx context.Context, userID string, chatSessions []string) error { tenantIDs, err := s.userTenantDAO.GetTenantIDsByUserID(userID) if err != nil { return err @@ -189,14 +189,14 @@ func (s *ChatSessionService) RemoveChatSessions(userID string, chatSessions []st tenantIDSet[userID] = true for _, convID := range chatSessions { - session, err := s.chatSessionDAO.GetByID(convID) + session, err := s.chatSessionDAO.GetByID(ctx, convID) if err != nil { return fmt.Errorf("Chat session not found: %s", convID) } isOwner := false for tenantID := range tenantIDSet { - exists, err := s.chatSessionDAO.CheckDialogExists(tenantID, session.DialogID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, session.DialogID) if err != nil { return err } @@ -209,7 +209,7 @@ func (s *ChatSessionService) RemoveChatSessions(userID string, chatSessions []st return errors.New("Only owner of chat session authorized for this operation") } - if err := s.chatSessionDAO.DeleteByID(convID); err != nil { + if err = s.chatSessionDAO.DeleteByID(ctx, convID); err != nil { return err } } @@ -242,7 +242,7 @@ type ChatSessionPayload struct { } // ListChatSessions lists chat sessions for a dialog -func (s *ChatSessionService) ListChatSessions(userID string, chatID string) (*ListChatSessionsResponse, error) { +func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID string, chatID string) (*ListChatSessionsResponse, error) { // Get user's tenants tenantIDs, err := s.userTenantDAO.GetTenantIDsByUserID(userID) if err != nil { @@ -253,7 +253,7 @@ func (s *ChatSessionService) ListChatSessions(userID string, chatID string) (*Li isOwner := false for _, tenantID := range tenantIDs { var exists bool - exists, err = s.chatSessionDAO.CheckDialogExists(tenantID, chatID) + exists, err = s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) if err != nil { return nil, err } @@ -266,7 +266,7 @@ func (s *ChatSessionService) ListChatSessions(userID string, chatID string) (*Li // Also check with userID as tenant if !isOwner { var exists bool - exists, err = s.chatSessionDAO.CheckDialogExists(userID, chatID) + exists, err = s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) if err != nil { return nil, err } @@ -278,7 +278,7 @@ func (s *ChatSessionService) ListChatSessions(userID string, chatID string) (*Li } // List chat sessions - sessions, err := s.chatSessionDAO.ListByChatID(chatID) + sessions, err := s.chatSessionDAO.ListByChatID(ctx, chatID) if err != nil { return nil, err } @@ -287,8 +287,8 @@ func (s *ChatSessionService) ListChatSessions(userID string, chatID string) (*Li } // GetSession returns one chat session after ownership validation. -func (s *ChatSessionService) GetSession(userID, chatID, sessionID string) (*ChatSessionPayload, common.ErrorCode, error) { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, sessionID string) (*ChatSessionPayload, common.ErrorCode, error) { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -296,7 +296,7 @@ func (s *ChatSessionService) GetSession(userID, chatID, sessionID string) (*Chat return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, sessionID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") @@ -307,7 +307,7 @@ func (s *ChatSessionService) GetSession(userID, chatID, sessionID string) (*Chat return nil, common.CodeDataError, errors.New("Session does not belong to this chat!") } - dialog, err := s.chatSessionDAO.GetDialogByID(chatID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, chatID) if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err } @@ -316,8 +316,8 @@ func (s *ChatSessionService) GetSession(userID, chatID, sessionID string) (*Chat } // CreateSession create a session in a dialog -func (s *ChatSessionService) CreateSession(userID, chatID string, req map[string]interface{}) (*ChatSessionPayload, common.ErrorCode, error) { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) CreateSession(ctx context.Context, userID, chatID string, req map[string]interface{}) (*ChatSessionPayload, common.ErrorCode, error) { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -325,7 +325,7 @@ func (s *ChatSessionService) CreateSession(userID, chatID string, req map[string return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - dialog, err := s.chatSessionDAO.GetDialogByID(chatID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, chatID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Chat not found!") @@ -370,11 +370,11 @@ func (s *ChatSessionService) CreateSession(userID, chatID string, req map[string Reference: referenceJSON, } - if err := s.chatSessionDAO.Create(conv); err != nil { + if err = s.chatSessionDAO.Create(ctx, conv); err != nil { return nil, common.CodeDataError, errors.New("Fail to create a session!") } - session, err := s.chatSessionDAO.GetByID(conv.ID) + session, err := s.chatSessionDAO.GetByID(ctx, conv.ID) if err != nil { return nil, common.CodeDataError, errors.New("Fail to create a session!") } @@ -382,8 +382,8 @@ func (s *ChatSessionService) CreateSession(userID, chatID string, req map[string } // DeleteSessions delete a session in a dialog -func (s *ChatSessionService) DeleteSessions(userID, chatID string, req map[string]interface{}) (interface{}, string, common.ErrorCode, error) { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) DeleteSessions(ctx context.Context, userID, chatID string, req map[string]interface{}) (interface{}, string, common.ErrorCode, error) { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return nil, "", common.CodeServerError, err } @@ -399,7 +399,7 @@ func (s *ChatSessionService) DeleteSessions(userID, chatID string, req map[strin if !hasIDs || len(sessionIDs) == 0 { deleteAll, _ := req["delete_all"].(bool) if deleteAll { - sessions, err := s.chatSessionDAO.ListByChatID(chatID) + sessions, err := s.chatSessionDAO.ListByChatID(ctx, chatID) if err != nil { return nil, "", common.CodeServerError, err } @@ -420,7 +420,7 @@ func (s *ChatSessionService) DeleteSessions(userID, chatID string, req map[strin successCount := 0 for _, sid := range uniqueIDs { - session, err := s.chatSessionDAO.GetBySessionIDAndChatID(sid, chatID) + session, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, sid, chatID) if err != nil { errorsList = append(errorsList, fmt.Sprintf("The chat doesn't own the session %s", sid)) continue @@ -428,7 +428,7 @@ func (s *ChatSessionService) DeleteSessions(userID, chatID string, req map[strin s.removeSessionUploadFiles(userID, session) - if err := s.chatSessionDAO.DeleteByID(sid); err != nil { + if err = s.chatSessionDAO.DeleteByID(ctx, sid); err != nil { return nil, "", common.CodeServerError, err } @@ -540,8 +540,8 @@ func checkDuplicateChatSessionIDs(ids []string) ([]string, []string) { } // UpdateSession updates one chat session after Python-style field validation. -func (s *ChatSessionService) UpdateSession(userID, chatID, sessionID string, req map[string]interface{}) (*ChatSessionPayload, common.ErrorCode, error) { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) UpdateSession(ctx context.Context, userID, chatID, sessionID string, req map[string]interface{}) (*ChatSessionPayload, common.ErrorCode, error) { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -552,7 +552,7 @@ func (s *ChatSessionService) UpdateSession(userID, chatID, sessionID string, req return nil, common.CodeArgumentError, errors.New("Request body cannot be empty") } - if _, err := s.chatSessionDAO.GetBySessionIDAndChatID(sessionID, chatID); err != nil { + if _, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, sessionID, chatID); err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") } @@ -591,14 +591,14 @@ func (s *ChatSessionService) UpdateSession(userID, chatID, sessionID string, req } } - if err := s.chatSessionDAO.UpdateByID(sessionID, updateFields); err != nil { + if err = s.chatSessionDAO.UpdateByID(ctx, sessionID, updateFields); err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") } return nil, common.CodeServerError, err } - session, err := s.chatSessionDAO.GetByID(sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, sessionID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Fail to update a session!") @@ -609,8 +609,8 @@ func (s *ChatSessionService) UpdateSession(userID, chatID, sessionID string, req return s.buildSessionPayload(session, nil, false), common.CodeSuccess, nil } -func (s *ChatSessionService) DeleteSessionMessage(userID, chatID, sessionID, msgID string) (*ChatSessionPayload, common.ErrorCode, error) { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, chatID, sessionID, msgID string) (*ChatSessionPayload, common.ErrorCode, error) { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -618,7 +618,7 @@ func (s *ChatSessionService) DeleteSessionMessage(userID, chatID, sessionID, msg return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, sessionID) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -663,7 +663,7 @@ func (s *ChatSessionService) DeleteSessionMessage(userID, chatID, sessionID, msg if err != nil { return nil, common.CodeServerError, err } - if err := s.chatSessionDAO.UpdateByID(session.ID, map[string]interface{}{ + if err = s.chatSessionDAO.UpdateByID(ctx, session.ID, map[string]interface{}{ "message": messageRaw, "reference": referenceRaw, }); err != nil { @@ -685,7 +685,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return nil, common.CodeServerError, err } for _, tenantID := range tenantIDs { - exists, err := s.chatSessionDAO.CheckDialogExists(tenantID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -695,7 +695,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, } } if ownerTenantID == "" { - exists, err := s.chatSessionDAO.CheckDialogExists(userID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -708,7 +708,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, sessionID) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -779,7 +779,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, if err != nil { return nil, common.CodeServerError, err } - if err := s.chatSessionDAO.UpdateByID(session.ID, map[string]interface{}{"message": messageRaw}); err != nil { + if err = s.chatSessionDAO.UpdateByID(ctx, session.ID, map[string]interface{}{"message": messageRaw}); err != nil { return nil, common.CodeServerError, err } session.Message = messageRaw @@ -1112,14 +1112,14 @@ func floatValue(value interface{}) (float64, bool) { } } -func (s *ChatSessionService) ensureOwnedChat(userID, chatID string) (bool, error) { +func (s *ChatSessionService) ensureOwnedChat(ctx context.Context, userID, chatID string) (bool, error) { tenantIDs, err := s.userTenantDAO.GetTenantIDsByUserID(userID) if err != nil { return false, err } for _, tenantID := range tenantIDs { - exists, err := s.chatSessionDAO.CheckDialogExists(tenantID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) if err != nil { return false, err } @@ -1128,7 +1128,7 @@ func (s *ChatSessionService) ensureOwnedChat(userID, chatID string) (bool, error } } - exists, err := s.chatSessionDAO.CheckDialogExists(userID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) if err != nil { return false, err } @@ -1265,7 +1265,7 @@ func isChatSessionNotFound(err error) bool { // Completion performs chat completion with full RAG support via ChatPipelineService. // Kept as a compatibility entrypoint for callers that still use the pre-ChatCompletions API. -func (s *ChatSessionService) Completion(userID string, conversationID string, messages []map[string]interface{}, llmID string, chatModelConfig map[string]interface{}, messageID string) (map[string]interface{}, error) { +func (s *ChatSessionService) Completion(ctx context.Context, userID string, conversationID string, messages []map[string]interface{}, llmID string, chatModelConfig map[string]interface{}, messageID string) (map[string]interface{}, error) { if len(messages) == 0 { return nil, errors.New("messages cannot be empty") } @@ -1274,12 +1274,12 @@ func (s *ChatSessionService) Completion(userID string, conversationID string, me return nil, errors.New("the last content of this conversation is not from user") } - session, err := s.chatSessionDAO.GetByID(conversationID) + session, err := s.chatSessionDAO.GetByID(ctx, conversationID) if err != nil { return nil, errors.New("Conversation not found") } - dialog, err := s.chatSessionDAO.GetDialogByID(session.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, session.DialogID) if err != nil { return nil, errors.New("Dialog not found") } @@ -1338,7 +1338,7 @@ func (s *ChatSessionService) Completion(userID string, conversationID string, me "id": messageID, "created_at": float64(time.Now().Unix()), }) - s.updateSessionMessages(session, sessionMessages, reference) + s.updateSessionMessages(ctx, session, sessionMessages, reference) } return result, nil @@ -1361,16 +1361,16 @@ func (s *ChatSessionService) CompletionStream(ctx context.Context, userID string return errors.New("the last content of this conversation is not from user") } - session, err := s.chatSessionDAO.GetByID(conversationID) + session, err := s.chatSessionDAO.GetByID(ctx, conversationID) if err != nil { streamChan <- fmt.Sprintf("data: %s\n\n", `{"code": 500, "message": "Conversation not found", "data": {"answer": "**ERROR**: Conversation not found", "reference": []}}`) - return errors.New("Conversation not found") + return errors.New("conversation not found") } - dialog, err := s.chatSessionDAO.GetDialogByID(session.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, session.DialogID) if err != nil { streamChan <- fmt.Sprintf("data: %s\n\n", `{"code": 500, "message": "Dialog not found", "data": {"answer": "**ERROR**: Dialog not found", "reference": []}}`) - return errors.New("Dialog not found") + return errors.New("dialog not found") } sessionMessages := s.buildSessionMessages(session, messages) @@ -1436,7 +1436,7 @@ func (s *ChatSessionService) CompletionStream(ctx context.Context, userID string "id": messageID, "created_at": float64(time.Now().Unix()), }) - s.updateSessionMessages(session, sessionMessages, reference) + s.updateSessionMessages(ctx, session, sessionMessages, reference) } return nil @@ -1491,15 +1491,15 @@ func (s *ChatSessionService) ChatCompletions( var dialog *entity.Chat var session *entity.ChatSession if chatID != "" { - if err := s.checkDialogOwnership(userID, chatID); err != nil { + if err = s.checkDialogOwnership(ctx, userID, chatID); err != nil { return fail(err) } - dialog, err = s.chatSessionDAO.GetDialogByID(chatID) + dialog, err = s.chatSessionDAO.GetDialogByID(ctx, chatID) if err != nil { return fail(errors.New("Chat not found!")) } if sessionID != "" { - session, err = s.chatSessionDAO.GetByID(sessionID) + session, err = s.chatSessionDAO.GetByID(ctx, sessionID) if err != nil { return fail(errors.New("Session not found!")) } @@ -1507,7 +1507,7 @@ func (s *ChatSessionService) ChatCompletions( return fail(errors.New("Session does not belong to this chat!")) } } else { - session, err = s.createSessionForCompletion(chatID, dialog, userID) + session, err = s.createSessionForCompletion(ctx, chatID, dialog, userID) if err != nil { return fail(err) } @@ -1541,14 +1541,14 @@ func (s *ChatSessionService) ChatCompletions( if llmID != "" { hasKey, err := s.checkTenantLLMAPIKey(dialog.TenantID, llmID) if err != nil || !hasKey { - return fail(fmt.Errorf("Cannot use specified model %s", llmID)) + return fail(fmt.Errorf("cannot use specified model %s", llmID)) } dialog.LLMID = llmID dialog.LLMSetting = genConfig } else if dialog.LLMID == "" { tenant, err := dao.NewTenantDAO().GetByID(dialog.TenantID) if err != nil || tenant.LLMID == "" { - return fail(errors.New("No default chat model for tenant.")) + return fail(errors.New("no default chat model for tenant")) } dialog.LLMID = tenant.LLMID if dialog.LLMSetting == nil { @@ -1672,7 +1672,7 @@ func (s *ChatSessionService) ChatCompletions( // Persist session state (matches Python's update_by_id after loop) if session != nil { - s.updateSessionMessages(session, s.getSessionMessagesAsSlice(session), reference) + s.updateSessionMessages(ctx, session, s.getSessionMessagesAsSlice(session), reference) } } else { var answer strings.Builder @@ -1700,7 +1700,7 @@ func (s *ChatSessionService) ChatCompletions( if chatID != "" { result["chat_id"] = chatID } - s.updateSessionMessages(session, s.getSessionMessagesAsSlice(session), reference) + s.updateSessionMessages(ctx, session, s.getSessionMessagesAsSlice(session), reference) return sanitizeJSONFloats(result).(map[string]interface{}), nil } ans["id"] = messageID @@ -1752,11 +1752,11 @@ func (s *ChatSessionService) normalizeCompletionMessages( } if len(requestMsg) == 0 { - return nil, nil, "", errors.New("`messages` must contain a user message.") + return nil, nil, "", errors.New("`messages` must contain a user message") } lastRole, _ := requestMsg[len(requestMsg)-1]["role"].(string) if lastRole != "user" { - return nil, nil, "", errors.New("The last content of this conversation is not from user.") + return nil, nil, "", errors.New("the last content of this conversation is not from user") } // Generate message ID if missing — matches Python's get_uuid() in _normalize_completion_messages. @@ -1777,8 +1777,8 @@ func (s *ChatSessionService) normalizeCompletionMessages( } // checkDialogOwnership checks if the user owns the dialog. -func (s *ChatSessionService) checkDialogOwnership(userID, chatID string) error { - ok, err := s.ensureOwnedChat(userID, chatID) +func (s *ChatSessionService) checkDialogOwnership(ctx context.Context, userID, chatID string) error { + ok, err := s.ensureOwnedChat(ctx, userID, chatID) if err != nil { return err } @@ -1804,7 +1804,7 @@ func (s *ChatSessionService) buildDefaultCompletionDialog(tenantID string) *enti } } -func (s *ChatSessionService) createSessionForCompletion(chatID string, dialog *entity.Chat, userID string) (*entity.ChatSession, error) { +func (s *ChatSessionService) createSessionForCompletion(ctx context.Context, chatID string, dialog *entity.Chat, userID string) (*entity.ChatSession, error) { newID := utility.GenerateUUID() name := "New session" @@ -1829,7 +1829,7 @@ func (s *ChatSessionService) createSessionForCompletion(chatID string, dialog *e UserID: &userID, Reference: refJSON, } - if err := s.chatSessionDAO.Create(session); err != nil { + if err := s.chatSessionDAO.Create(ctx, session); err != nil { return nil, err } return session, nil @@ -2124,7 +2124,7 @@ func (s *ChatSessionService) structureAnswer(session *entity.ChatSession, answer } } -func (s *ChatSessionService) updateSessionMessages(session *entity.ChatSession, messages []map[string]interface{}, reference []interface{}) { +func (s *ChatSessionService) updateSessionMessages(ctx context.Context, session *entity.ChatSession, messages []map[string]interface{}, reference []interface{}) { messagesJSON, err := json.Marshal(messages) if err != nil { common.Warn("updateSessionMessages: failed to marshal messages", zap.Error(err)) @@ -2140,7 +2140,7 @@ func (s *ChatSessionService) updateSessionMessages(session *entity.ChatSession, "message": messagesJSON, "reference": referenceJSON, } - if err := s.chatSessionDAO.UpdateByID(session.ID, updates); err != nil { + if err := s.chatSessionDAO.UpdateByID(ctx, session.ID, updates); err != nil { common.Warn("updateSessionMessages: DAO update failed", zap.Error(err)) return } diff --git a/internal/service/chat_session_contract_test.go b/internal/service/chat_session_contract_test.go index 47b47ea93c..f1978bea15 100644 --- a/internal/service/chat_session_contract_test.go +++ b/internal/service/chat_session_contract_test.go @@ -41,7 +41,8 @@ func TestCreateSession_Success(t *testing.T) { pipeline: &fakePipeline{}, } - resp, code, err := svc.CreateSession("user-1", "chat-1", map[string]interface{}{"name": "valid"}) + ctx := t.Context() + resp, code, err := svc.CreateSession(ctx, "user-1", "chat-1", map[string]interface{}{"name": "valid"}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -70,8 +71,9 @@ func TestCreateSession_RejectsEmptyOrNonStringName(t *testing.T) { pipeline: &fakePipeline{}, } + ctx := t.Context() for _, name := range []interface{}{"", " ", 1} { - _, code, err := svc.CreateSession("user-1", "chat-1", map[string]interface{}{"name": name}) + _, code, err := svc.CreateSession(ctx, "user-1", "chat-1", map[string]interface{}{"name": name}) if err == nil || err.Error() != "`name` can not be empty." { t.Fatalf("name=%#v err=%v", name, err) } @@ -92,8 +94,9 @@ func TestCreateSession_TruncatesLongName(t *testing.T) { pipeline: &fakePipeline{}, } + ctx := t.Context() longName := strings.Repeat("a", 300) - resp, code, err := svc.CreateSession("user-1", "chat-1", map[string]interface{}{"name": longName}) + resp, code, err := svc.CreateSession(ctx, "user-1", "chat-1", map[string]interface{}{"name": longName}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -113,7 +116,8 @@ func TestCreateSession_NotOwner(t *testing.T) { pipeline: &fakePipeline{}, } - _, code, err := svc.CreateSession("user-1", "chat-1", map[string]interface{}{"name": "x"}) + ctx := t.Context() + _, code, err := svc.CreateSession(ctx, "user-1", "chat-1", map[string]interface{}{"name": "x"}) if err == nil || err.Error() != "No authorization." { t.Fatalf("err=%v", err) } @@ -139,7 +143,8 @@ func TestDeleteSessions_SuccessByIDs(t *testing.T) { pipeline: &fakePipeline{}, } - result, message, code, err := svc.DeleteSessions("user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"s1"}}) + ctx := t.Context() + result, message, code, err := svc.DeleteSessions(ctx, "user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"s1"}}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -168,13 +173,14 @@ func TestDeleteSessions_DeleteAllAndInvalidID(t *testing.T) { } // Empty payload -> success with empty result map. - if _, _, code, err := svc.DeleteSessions("user-1", "chat-1", map[string]interface{}{}); err != nil || code != common.CodeSuccess { + ctx := t.Context() + if _, _, code, err := svc.DeleteSessions(ctx, "user-1", "chat-1", map[string]interface{}{}); err != nil || code != common.CodeSuccess { t.Fatalf("empty payload: code=%v err=%v", code, err) } // delete_all removes every session for the chat. store.sessions["s1"] = &entity.ChatSession{ID: "s1", DialogID: "chat-1"} - if _, _, code, err := svc.DeleteSessions("user-1", "chat-1", map[string]interface{}{"delete_all": true}); err != nil || code != common.CodeSuccess { + if _, _, code, err := svc.DeleteSessions(ctx, "user-1", "chat-1", map[string]interface{}{"delete_all": true}); err != nil || code != common.CodeSuccess { t.Fatalf("delete_all: code=%v err=%v", code, err) } if len(store.sessions) != 0 { @@ -183,7 +189,7 @@ func TestDeleteSessions_DeleteAllAndInvalidID(t *testing.T) { // Unknown id -> DataError reporting the unowned session. store.sessions["s1"] = &entity.ChatSession{ID: "s1", DialogID: "chat-1"} - _, _, code, err := svc.DeleteSessions("user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"missing"}}) + _, _, code, err := svc.DeleteSessions(ctx, "user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"missing"}}) if err == nil || !strings.Contains(err.Error(), "The chat doesn't own the session missing") { t.Fatalf("invalid id: err=%v", err) } @@ -200,7 +206,8 @@ func TestDeleteSessions_NotOwner(t *testing.T) { pipeline: &fakePipeline{}, } - _, _, code, err := svc.DeleteSessions("user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"s1"}}) + ctx := t.Context() + _, _, code, err := svc.DeleteSessions(ctx, "user-1", "chat-1", map[string]interface{}{"ids": []interface{}{"s1"}}) if err == nil || err.Error() != "No authorization." { t.Fatalf("err=%v", err) } diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index fea16bbc0a..4698b746f6 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -10,11 +10,12 @@ import ( "sync" "testing" - "gorm.io/gorm" "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" + + "gorm.io/gorm" ) // --------------------------------------------------------------------------- @@ -46,7 +47,7 @@ func newFakeSessionStore() *fakeSessionStore { } } -func (f *fakeSessionStore) GetByID(id string) (*entity.ChatSession, error) { +func (f *fakeSessionStore) GetByID(ctx context.Context, id string) (*entity.ChatSession, error) { if f.getByIDErr != nil { return nil, f.getByIDErr } @@ -57,8 +58,8 @@ func (f *fakeSessionStore) GetByID(id string) (*entity.ChatSession, error) { return s, nil } -func (f *fakeSessionStore) GetBySessionIDAndChatID(sessionID, chatID string) (*entity.ChatSession, error) { - s, err := f.GetByID(sessionID) +func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) { + s, err := f.GetByID(ctx, sessionID) if err != nil { return nil, err } @@ -68,7 +69,7 @@ func (f *fakeSessionStore) GetBySessionIDAndChatID(sessionID, chatID string) (*e return s, nil } -func (f *fakeSessionStore) Create(conv *entity.ChatSession) error { +func (f *fakeSessionStore) Create(ctx context.Context, conv *entity.ChatSession) error { f.mu.Lock() defer f.mu.Unlock() if f.createErr != nil { @@ -79,7 +80,7 @@ func (f *fakeSessionStore) Create(conv *entity.ChatSession) error { return nil } -func (f *fakeSessionStore) UpdateByID(id string, updates map[string]interface{}) error { +func (f *fakeSessionStore) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { f.mu.Lock() defer f.mu.Unlock() if f.updateByIDErr != nil { @@ -112,12 +113,12 @@ func (f *fakeSessionStore) UpdateByID(id string, updates map[string]interface{}) return nil } -func (f *fakeSessionStore) DeleteByID(id string) error { +func (f *fakeSessionStore) DeleteByID(ctx context.Context, id string) error { delete(f.sessions, id) return nil } -func (f *fakeSessionStore) ListByChatID(chatID string) ([]*entity.ChatSession, error) { +func (f *fakeSessionStore) ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) { var result []*entity.ChatSession for _, s := range f.sessions { if s.DialogID == chatID { @@ -127,7 +128,7 @@ func (f *fakeSessionStore) ListByChatID(chatID string) ([]*entity.ChatSession, e return result, nil } -func (f *fakeSessionStore) GetDialogByID(chatID string) (*entity.Chat, error) { +func (f *fakeSessionStore) GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) { if f.getDialogErr != nil { return nil, f.getDialogErr } @@ -138,7 +139,7 @@ func (f *fakeSessionStore) GetDialogByID(chatID string) (*entity.Chat, error) { return d, nil } -func (f *fakeSessionStore) CheckDialogExists(tenantID, chatID string) (bool, error) { +func (f *fakeSessionStore) CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) { key := tenantID + "|" + chatID return f.dialogExists[key], nil } @@ -289,7 +290,8 @@ func TestListChatSessions_Success(t *testing.T) { pipeline: &fakePipeline{}, } - resp, err := svc.ListChatSessions("user-1", "chat-1") + ctx := t.Context() + resp, err := svc.ListChatSessions(ctx, "user-1", "chat-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -307,7 +309,8 @@ func TestListChatSessions_NotOwner(t *testing.T) { pipeline: &fakePipeline{}, } - _, err := svc.ListChatSessions("user-1", "chat-1") + ctx := t.Context() + _, err := svc.ListChatSessions(ctx, "user-1", "chat-1") if err == nil || !strings.Contains(err.Error(), "only owner") { t.Fatalf("expected 'only owner' error, got %v", err) } @@ -340,7 +343,8 @@ func TestGetSession_Success(t *testing.T) { pipeline: &fakePipeline{}, } - resp, code, err := svc.GetSession("user-1", "chat-1", "session-1") + ctx := t.Context() + resp, code, err := svc.GetSession(ctx, "user-1", "chat-1", "session-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -382,7 +386,8 @@ func TestGetSession_NotOwner(t *testing.T) { pipeline: &fakePipeline{}, } - _, code, err := svc.GetSession("user-1", "chat-1", "session-1") + ctx := t.Context() + _, code, err := svc.GetSession(ctx, "user-1", "chat-1", "session-1") if err == nil || err.Error() != "No authorization." { t.Fatalf("err=%v", err) } @@ -402,7 +407,8 @@ func TestGetSession_WrongChat(t *testing.T) { pipeline: &fakePipeline{}, } - _, code, err := svc.GetSession("user-1", "chat-1", "session-1") + ctx := t.Context() + _, code, err := svc.GetSession(ctx, "user-1", "chat-1", "session-1") if err == nil || err.Error() != "Session does not belong to this chat!" { t.Fatalf("err=%v", err) } @@ -428,7 +434,8 @@ func TestUpdateSession_Success(t *testing.T) { } longName := " " + strings.Repeat("x", 260) + " " - resp, code, err := svc.UpdateSession("user-1", "chat-1", "session-1", map[string]interface{}{ + ctx := t.Context() + resp, code, err := svc.UpdateSession(ctx, "user-1", "chat-1", "session-1", map[string]interface{}{ "name": longName, "user_id": "spoof", "chat_id": "spoof-chat", @@ -482,7 +489,8 @@ func TestUpdateSession_ValidationErrors(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - _, code, err := svc.UpdateSession("user-1", "chat-1", "session-1", tc.req) + ctx := t.Context() + _, code, err := svc.UpdateSession(ctx, "user-1", "chat-1", "session-1", tc.req) if err == nil || err.Error() != tc.message { t.Fatalf("err=%v", err) } @@ -503,7 +511,8 @@ func TestUpdateSession_NotFound(t *testing.T) { pipeline: &fakePipeline{}, } - _, code, err := svc.UpdateSession("user-1", "chat-1", "missing", map[string]interface{}{"name": "renamed"}) + ctx := t.Context() + _, code, err := svc.UpdateSession(ctx, "user-1", "chat-1", "missing", map[string]interface{}{"name": "renamed"}) if err == nil || err.Error() != "Session not found!" { t.Fatalf("err=%v", err) } @@ -541,7 +550,8 @@ func TestDeleteSessionMessage_RemovesMessagePairAndReference(t *testing.T) { pipeline: &fakePipeline{}, } - resp, code, err := svc.DeleteSessionMessage("user-1", "chat-1", "session-1", "msg-1") + ctx := t.Context() + resp, code, err := svc.DeleteSessionMessage(ctx, "user-1", "chat-1", "session-1", "msg-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -834,7 +844,8 @@ func TestCompletion_Success(t *testing.T) { pipeline: pipeline, } - result, err := svc.Completion("user-1", "session-1", []map[string]interface{}{ + ctx := t.Context() + result, err := svc.Completion(ctx, "user-1", "session-1", []map[string]interface{}{ {"role": "user", "content": "hi"}, }, "", nil, "msg-1") if err != nil { @@ -1104,7 +1115,8 @@ func TestCompletion_EmptyMessages(t *testing.T) { pipeline: &fakePipeline{}, } - _, err := svc.Completion("user-1", "session-1", nil, "", nil, "msg-1") + ctx := t.Context() + _, err := svc.Completion(ctx, "user-1", "session-1", nil, "", nil, "msg-1") if err == nil || err.Error() != "messages cannot be empty" { t.Fatalf("expected 'messages cannot be empty', got %v", err) } @@ -1117,7 +1129,8 @@ func TestCompletion_LastMessageNotFromUser(t *testing.T) { pipeline: &fakePipeline{}, } - _, err := svc.Completion("user-1", "session-1", []map[string]interface{}{ + ctx := t.Context() + _, err := svc.Completion(ctx, "user-1", "session-1", []map[string]interface{}{ {"role": "assistant", "content": "hello"}, }, "", nil, "msg-1") if err == nil || !strings.Contains(err.Error(), "not from user") { @@ -1134,7 +1147,8 @@ func TestCompletion_ConversationNotFound(t *testing.T) { pipeline: &fakePipeline{}, } - _, err := svc.Completion("user-1", "missing", []map[string]interface{}{ + ctx := t.Context() + _, err := svc.Completion(ctx, "user-1", "missing", []map[string]interface{}{ {"role": "user", "content": "hi"}, }, "", nil, "msg-1") if err == nil || err.Error() != "Conversation not found" { @@ -1156,7 +1170,8 @@ func TestCompletion_DialogNotFound(t *testing.T) { pipeline: &fakePipeline{}, } - _, err := svc.Completion("user-1", "session-1", []map[string]interface{}{ + ctx := t.Context() + _, err := svc.Completion(ctx, "user-1", "session-1", []map[string]interface{}{ {"role": "user", "content": "hi"}, }, "", nil, "msg-1") if err == nil || err.Error() != "Dialog not found" { @@ -1182,7 +1197,8 @@ func TestCompletion_PipelineError(t *testing.T) { pipeline: &fakePipeline{err: errors.New("model unavailable")}, } - _, err := svc.Completion("user-1", "session-1", []map[string]interface{}{ + ctx := t.Context() + _, err := svc.Completion(ctx, "user-1", "session-1", []map[string]interface{}{ {"role": "user", "content": "hi"}, }, "", nil, "msg-1") if err == nil || err.Error() != "model unavailable" { diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index b156dc6774..78bfca4c41 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -217,7 +217,8 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID } } - dialogResp, err := s.chatSvc.GetChat(userID, chatID) + ctx := c.Request.Context() + dialogResp, err := s.chatSvc.GetChat(ctx, userID, chatID) if err != nil { s.writeDataError(c, err.Error()) return @@ -263,7 +264,6 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID completionID := fmt.Sprintf("chatcmpl-%s", openaiReq.ChatID) - ctx := c.Request.Context() lfClient := LangfuseClientFromTenant(ctx, dialog.TenantID, userID, openaiReq.ChatID, openaiReq.Model) if lfClient != nil { ctx = context.WithValue(ctx, langfuseCtxKey, lfClient) diff --git a/internal/service/stats.go b/internal/service/stats.go index 9553c5864a..e9e08eb347 100644 --- a/internal/service/stats.go +++ b/internal/service/stats.go @@ -17,6 +17,7 @@ package service import ( + "context" "errors" "ragflow/internal/dao" @@ -50,13 +51,13 @@ type StatsResponse struct { } // GetStats returns daily API conversation statistics for the first tenant of a user. -func (s *StatsService) GetStats(userID, fromDate, toDate string, source *string) (*StatsResponse, error) { +func (s *StatsService) GetStats(ctx context.Context, userID, fromDate, toDate string, source *string) (*StatsResponse, error) { tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil || len(tenants) == 0 { return nil, ErrTenantNotFound } - rows, err := dao.NewAPI4ConversationDAO().Stats(tenants[0].TenantID, fromDate, toDate, source) + rows, err := dao.NewAPI4ConversationDAO().Stats(ctx, tenants[0].TenantID, fromDate, toDate, source) if err != nil { return nil, err } diff --git a/internal/service/user.go b/internal/service/user.go index 649d567af0..d49ffa5e33 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -1020,7 +1020,7 @@ func (s *UserService) GetUserByAPIToken(ctx context.Context, authorization strin // Query API token from database apiTokenDAO := dao.NewAPITokenDAO() - userToken, err := apiTokenDAO.GetUserByAPIToken(token) + userToken, err := apiTokenDAO.GetUserByAPIToken(ctx, token) if err != nil { return nil, common.CodeUnauthorized, fmt.Errorf("invalid access token") } @@ -1066,7 +1066,7 @@ func (s *UserService) GetAPITokenByBeta(ctx context.Context, authorization strin return nil, fmt.Errorf("invalid authorization format") } apiTokenDAO := dao.NewAPITokenDAO() - tokens, err := apiTokenDAO.GetByBeta(token) + tokens, err := apiTokenDAO.GetByBeta(ctx, token) if err != nil { return nil, err } @@ -1101,7 +1101,7 @@ func (s *UserService) GetUserByBetaAPIToken(ctx context.Context, authorization s } apiTokenDAO := dao.NewAPITokenDAO() - userTokens, err := apiTokenDAO.GetByBeta(token) + userTokens, err := apiTokenDAO.GetByBeta(ctx, token) if err != nil || len(userTokens) == 0 { return nil, common.CodeUnauthorized, fmt.Errorf("invalid beta access token") } diff --git a/ragflow_deps/download_go_deps.py b/ragflow_deps/download_go_deps.py index 82985fddf6..3d43018a87 100644 --- a/ragflow_deps/download_go_deps.py +++ b/ragflow_deps/download_go_deps.py @@ -22,7 +22,7 @@ # # Typical workflow: # -# uv run python3 ragflow_deps/download_deps.py # download +# uv run python3 ragflow_deps/download_go_deps.py # download # cd ragflow_deps # docker build -f Dockerfile -t infiniflow/ragflow_deps . #