diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 732d4d3608..65910f5d0d 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -361,7 +361,7 @@ func main() { defer engine.Close() // Initialize Redis cache - if err = redis.Init(); err != nil { + if err = redis.Init(ctx); err != nil { common.Fatal("Failed to initialize Redis", zap.Error(err)) } defer redis.Close() @@ -784,7 +784,7 @@ func startServer(ctx context.Context) { agentOpts.stateSerializer, agentOpts.runTracker, ) - agentHandler := handler.NewAgentHandler(agentService, fileService) + agentHandler := handler.NewAgentHandler(ctx, agentService, fileService) // Public chatbot/agentbot endpoints (api/v1/chatbots/..., // api/v1/agentbots/...) and the agent attachment download. diff --git a/internal/admin/handler.go b/internal/admin/handler.go index 91bf959c38..053286f452 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -17,7 +17,6 @@ package admin import ( - "context" "encoding/json" "errors" "fmt" @@ -113,7 +112,7 @@ func (h *Handler) Login(c *gin.Context) { return } - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { common.ErrorWithCode(c, common.CodeServerError, fmt.Sprintf("Failed to get secret key: %s", err.Error())) return @@ -1235,8 +1234,9 @@ func (h *Handler) PingStore(c *gin.Context) { } func (h *Handler) PingCache(c *gin.Context) { + ctx := c.Request.Context() redisClient := redis.Get() - if redisClient.Health() { + if redisClient.Health(ctx) { common.SuccessNoMessage(c, "SUCCESS") } else { common.ErrorWithCode(c, common.CodeServerError, "cache health check failed") @@ -1245,7 +1245,7 @@ func (h *Handler) PingCache(c *gin.Context) { func (h *Handler) PingEngine(c *gin.Context) { docEngine := engine.Get() - ctx := context.Background() + ctx := c.Request.Context() if err := docEngine.Ping(ctx); err != nil { var coded interface { Code() common.ErrorCode diff --git a/internal/admin/service.go b/internal/admin/service.go index 047b0aa01a..c5d6a00152 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -1121,7 +1121,7 @@ func (s *Service) getRedisInfo(ctx context.Context) ServiceStatus { startTime := time.Now() redisClient := redis.Get() - if redisClient.Health() { + if redisClient.Health(ctx) { return newServiceStatus(serviceType, name, "alive", startTime, "") } diff --git a/internal/agent/audio/model_provider_synthesizer.go b/internal/agent/audio/model_provider_synthesizer.go index a1ae930275..b2b7289c82 100644 --- a/internal/agent/audio/model_provider_synthesizer.go +++ b/internal/agent/audio/model_provider_synthesizer.go @@ -108,7 +108,7 @@ func (m *modelProviderSynthesizer) Synthesize(ctx context.Context, req Synthesiz // log and fall through to the model provider. cacheKey := buildTTSCacheKey(tenantID, req) if cacheKey != "" && m.redis != nil { - if cached, _ := m.redis.Get(cacheKey); cached != "" { + if cached, _ := m.redis.Get(ctx, cacheKey); cached != "" { if b, err := hex.DecodeString(cached); err == nil && len(b) > 0 { return &SynthesizeResponse{Audio: b, MediaType: "audio/mpeg"}, nil } @@ -139,7 +139,7 @@ func (m *modelProviderSynthesizer) Synthesize(ctx context.Context, req Synthesiz if cacheKey != "" && m.redis != nil { ttl := ttsCacheTTL() if ttl > 0 { - m.redis.Set(cacheKey, hex.EncodeToString(resp.Audio), ttl) + m.redis.Set(ctx, cacheKey, hex.EncodeToString(resp.Audio), ttl) } } return resp, nil diff --git a/internal/deepdoc/parser/pdf/inference/cache.go b/internal/deepdoc/parser/pdf/inference/cache.go index 4183aa2d18..f683f72c83 100644 --- a/internal/deepdoc/parser/pdf/inference/cache.go +++ b/internal/deepdoc/parser/pdf/inference/cache.go @@ -45,8 +45,8 @@ const cacheKeyPrefix = "ddoc:cache:" // out of the wrapping path. type cacheStore interface { Enabled() bool - GetObj(key string, dest any) bool - SetObj(key string, value any, ttl time.Duration) bool + GetObj(ctx context.Context, key string, dest any) bool + SetObj(ctx context.Context, key string, value any, ttl time.Duration) bool } // redisCacheStore is the production store, backed by the @@ -61,13 +61,13 @@ func (redisCacheStore) Enabled() bool { return redis.IsEnabled() } // GetObj forwards to the package-level Redis singleton. // Errors are logged inside the engine; we treat any failure // (including connectivity) as a cache miss for safety. -func (redisCacheStore) GetObj(key string, dest any) bool { - return redis.Get().GetObj(key, dest) +func (redisCacheStore) GetObj(ctx context.Context, key string, dest any) bool { + return redis.Get().GetObj(ctx, key, dest) } // SetObj forwards to the package-level Redis singleton. -func (redisCacheStore) SetObj(key string, value any, ttl time.Duration) bool { - return redis.Get().SetObj(key, value, ttl) +func (redisCacheStore) SetObj(ctx context.Context, key string, value any, ttl time.Duration) bool { + return redis.Get().SetObj(ctx, key, value, ttl) } // DocAnalyzerCache wraps an inner doctype.DocAnalyzer and @@ -156,7 +156,7 @@ func cacheKeyOrEmpty(method string, img image.Image) string { func (c *DocAnalyzerCache) DLA(ctx context.Context, img image.Image) ([]doctype.DLARegion, error) { if key := cacheKeyOrEmpty("dla", img); key != "" && c.store.Enabled() { var cached []doctype.DLARegion - if c.store.GetObj(key, &cached) { + if c.store.GetObj(ctx, key, &cached) { return cached, nil } } @@ -165,7 +165,7 @@ func (c *DocAnalyzerCache) DLA(ctx context.Context, img image.Image) ([]doctype. return out, err } if key := cacheKeyOrEmpty("dla", img); key != "" && c.store.Enabled() { - c.store.SetObj(key, out, c.ttl) + c.store.SetObj(ctx, key, out, c.ttl) } return out, nil } @@ -174,7 +174,7 @@ func (c *DocAnalyzerCache) DLA(ctx context.Context, img image.Image) ([]doctype. func (c *DocAnalyzerCache) TSR(ctx context.Context, img image.Image) ([]doctype.TSRCell, error) { if key := cacheKeyOrEmpty("tsr", img); key != "" && c.store.Enabled() { var cached []doctype.TSRCell - if c.store.GetObj(key, &cached) { + if c.store.GetObj(ctx, key, &cached) { return cached, nil } } @@ -183,7 +183,7 @@ func (c *DocAnalyzerCache) TSR(ctx context.Context, img image.Image) ([]doctype. return out, err } if key := cacheKeyOrEmpty("tsr", img); key != "" && c.store.Enabled() { - c.store.SetObj(key, out, c.ttl) + c.store.SetObj(ctx, key, out, c.ttl) } return out, nil } @@ -192,7 +192,7 @@ func (c *DocAnalyzerCache) TSR(ctx context.Context, img image.Image) ([]doctype. func (c *DocAnalyzerCache) OCRDetect(ctx context.Context, img image.Image) ([]doctype.OCRBox, error) { if key := cacheKeyOrEmpty("ocr_detect", img); key != "" && c.store.Enabled() { var cached []doctype.OCRBox - if c.store.GetObj(key, &cached) { + if c.store.GetObj(ctx, key, &cached) { return cached, nil } } @@ -201,7 +201,7 @@ func (c *DocAnalyzerCache) OCRDetect(ctx context.Context, img image.Image) ([]do return out, err } if key := cacheKeyOrEmpty("ocr_detect", img); key != "" && c.store.Enabled() { - c.store.SetObj(key, out, c.ttl) + c.store.SetObj(ctx, key, out, c.ttl) } return out, nil } @@ -210,7 +210,7 @@ func (c *DocAnalyzerCache) OCRDetect(ctx context.Context, img image.Image) ([]do func (c *DocAnalyzerCache) OCRRecognize(ctx context.Context, img image.Image) ([]doctype.OCRText, error) { if key := cacheKeyOrEmpty("ocr_recognize", img); key != "" && c.store.Enabled() { var cached []doctype.OCRText - if c.store.GetObj(key, &cached) { + if c.store.GetObj(ctx, key, &cached) { return cached, nil } } @@ -219,7 +219,7 @@ func (c *DocAnalyzerCache) OCRRecognize(ctx context.Context, img image.Image) ([ return out, err } if key := cacheKeyOrEmpty("ocr_recognize", img); key != "" && c.store.Enabled() { - c.store.SetObj(key, out, c.ttl) + c.store.SetObj(ctx, key, out, c.ttl) } return out, nil } diff --git a/internal/deepdoc/parser/pdf/inference/cache_test.go b/internal/deepdoc/parser/pdf/inference/cache_test.go index 2f74b67521..2928c5b8c1 100644 --- a/internal/deepdoc/parser/pdf/inference/cache_test.go +++ b/internal/deepdoc/parser/pdf/inference/cache_test.go @@ -48,7 +48,7 @@ func newFakeStore() *fakeStore { func (s *fakeStore) Enabled() bool { return s.enabled } -func (s *fakeStore) GetObj(key string, dest any) bool { +func (s *fakeStore) GetObj(ctx context.Context, key string, dest any) bool { atomic.AddInt32(&s.gets, 1) s.mu.Lock() defer s.mu.Unlock() @@ -59,7 +59,7 @@ func (s *fakeStore) GetObj(key string, dest any) bool { return json.Unmarshal(raw, dest) == nil } -func (s *fakeStore) SetObj(key string, value any, _ time.Duration) bool { +func (s *fakeStore) SetObj(ctx context.Context, key string, value any, _ time.Duration) bool { atomic.AddInt32(&s.sets, 1) raw, err := json.Marshal(value) if err != nil { diff --git a/internal/engine/redis/redis.go b/internal/engine/redis/redis.go index 3a460083c4..24e1d86c6c 100644 --- a/internal/engine/redis/redis.go +++ b/internal/engine/redis/redis.go @@ -106,7 +106,7 @@ const ( ) // Init InitRedis initializes Redis client -func Init() error { +func Init(ctx context.Context) error { var initErr error once.Do(func() { globalConfig := server.GetConfig() @@ -124,10 +124,10 @@ func Init() error { }) // Test connection - ctx, cancel := context.WithTimeout(context.Background(), server.DefaultConnectTimeout) + redisCtx, cancel := context.WithTimeout(ctx, server.DefaultConnectTimeout) defer cancel() - if err := client.Ping(ctx).Err(); err != nil { + if err := client.Ping(redisCtx).Err(); err != nil { initErr = fmt.Errorf("failed to connect to Redis: %w", err) return } @@ -167,11 +167,10 @@ func IsEnabled() bool { } // Health checks if Redis is healthy -func (r *Client) Health() bool { +func (r *Client) Health(ctx context.Context) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.Ping(ctx).Err(); err != nil { return false } @@ -190,11 +189,10 @@ func (r *Client) Health() bool { } // Info returns Redis server information -func (r *Client) Info() map[string]interface{} { +func (r *Client) Info(ctx context.Context) map[string]interface{} { if r.client == nil { return nil } - ctx := context.Background() infoStr, err := r.client.Info(ctx).Result() if err != nil { common.Warn("Failed to get Redis info", zap.Error(err)) @@ -278,11 +276,10 @@ func (r *Client) IsAlive() bool { } // Exist checks if key exists -func (r *Client) Exist(key string) (bool, error) { +func (r *Client) Exist(ctx context.Context, key string) (bool, error) { if r.client == nil { return false, nil } - ctx := context.Background() exists, err := r.client.Exists(ctx, key).Result() if err != nil { common.Warn("Redis Exist error", zap.String("key", key), zap.Error(err)) @@ -292,11 +289,10 @@ func (r *Client) Exist(key string) (bool, error) { } // Get gets value by key -func (r *Client) Get(key string) (string, error) { +func (r *Client) Get(ctx context.Context, key string) (string, error) { if r.client == nil { return "", nil } - ctx := context.Background() val, err := r.client.Get(ctx, key).Result() if err == redis.Nil { return "", nil @@ -309,11 +305,10 @@ func (r *Client) Get(key string) (string, error) { } // SetObj sets object with JSON serialization -func (r *Client) SetObj(key string, obj interface{}, exp time.Duration) bool { +func (r *Client) SetObj(ctx context.Context, key string, obj interface{}, exp time.Duration) bool { if r.client == nil { return false } - ctx := context.Background() data, err := json.Marshal(obj) if err != nil { common.Warn("Redis SetObj marshal error", zap.String("key", key), zap.Error(err)) @@ -326,12 +321,11 @@ func (r *Client) SetObj(key string, obj interface{}, exp time.Duration) bool { return true } -// GetObj gets and unmarshals object from Redis -func (r *Client) GetObj(key string, dest interface{}) bool { +// GetObj gets and unmarshal object from Redis +func (r *Client) GetObj(ctx context.Context, key string, dest interface{}) bool { if r.client == nil { return false } - ctx := context.Background() data, err := r.client.Get(ctx, key).Result() if err == redis.Nil { return false @@ -348,11 +342,10 @@ func (r *Client) GetObj(key string, dest interface{}) bool { } // Set sets value with expiration -func (r *Client) Set(key string, value string, exp time.Duration) bool { +func (r *Client) Set(ctx context.Context, key string, value string, exp time.Duration) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.Set(ctx, key, value, exp).Err(); err != nil { common.Warn("Redis Set error", zap.String("key", key), zap.Error(err)) return false @@ -361,11 +354,10 @@ func (r *Client) Set(key string, value string, exp time.Duration) bool { } // SetNX sets value only if key does not exist -func (r *Client) SetNX(key string, value string, exp time.Duration) bool { +func (r *Client) SetNX(ctx context.Context, key string, value string, exp time.Duration) bool { if r.client == nil { return false } - ctx := context.Background() ok, err := r.client.SetNX(ctx, key, value, exp).Result() if err != nil { common.Warn("Redis SetNX error", zap.String("key", key), zap.Error(err)) @@ -376,11 +368,10 @@ func (r *Client) SetNX(key string, value string, exp time.Duration) bool { // GetOrCreateKey atomically retrieves an existing key or creates a new one // Uses Redis SETNX command to ensure atomicity across multiple goroutines/processes -func (r *Client) GetOrCreateKey(key string, value string) (string, error) { +func (r *Client) GetOrCreateKey(ctx context.Context, key string, value string) (string, error) { if r.client == nil { return "", nil } - ctx := context.Background() // First, try to get the existing key existingKey, err := r.client.Get(ctx, key).Result() if err == nil { @@ -412,11 +403,10 @@ func (r *Client) GetOrCreateKey(key string, value string) (string, error) { } // SAdd adds member to set -func (r *Client) SAdd(key string, member string) bool { +func (r *Client) SAdd(ctx context.Context, key string, member string) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.SAdd(ctx, key, member).Err(); err != nil { common.Warn("Redis SAdd error", zap.String("key", key), zap.Error(err)) return false @@ -425,11 +415,10 @@ func (r *Client) SAdd(key string, member string) bool { } // SRem removes member from set -func (r *Client) SRem(key string, member string) bool { +func (r *Client) SRem(ctx context.Context, key string, member string) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.SRem(ctx, key, member).Err(); err != nil { common.Warn("Redis SRem error", zap.String("key", key), zap.Error(err)) return false @@ -438,11 +427,10 @@ func (r *Client) SRem(key string, member string) bool { } // SMembers returns all members of a set -func (r *Client) SMembers(key string) ([]string, error) { +func (r *Client) SMembers(ctx context.Context, key string) ([]string, error) { if r.client == nil { return nil, nil } - ctx := context.Background() members, err := r.client.SMembers(ctx, key).Result() if err != nil { common.Warn("Redis SMembers error", zap.String("key", key), zap.Error(err)) @@ -452,11 +440,10 @@ func (r *Client) SMembers(key string) ([]string, error) { } // SIsMember checks if member exists in set -func (r *Client) SIsMember(key string, member string) bool { +func (r *Client) SIsMember(ctx context.Context, key string, member string) bool { if r.client == nil { return false } - ctx := context.Background() ok, err := r.client.SIsMember(ctx, key, member).Result() if err != nil { common.Warn("Redis SIsMember error", zap.String("key", key), zap.Error(err)) @@ -466,11 +453,10 @@ func (r *Client) SIsMember(key string, member string) bool { } // ZAdd adds member with score to sorted set -func (r *Client) ZAdd(key string, member string, score float64) bool { +func (r *Client) ZAdd(ctx context.Context, key string, member string, score float64) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.ZAdd(ctx, key, redis.Z{Score: score, Member: member}).Err(); err != nil { common.Warn("Redis ZAdd error", zap.String("key", key), zap.Error(err)) return false @@ -479,11 +465,10 @@ func (r *Client) ZAdd(key string, member string, score float64) bool { } // ZCount returns count of members with score in range -func (r *Client) ZCount(key string, min, max float64) int64 { +func (r *Client) ZCount(ctx context.Context, key string, min, max float64) int64 { if r.client == nil { return 0 } - ctx := context.Background() count, err := r.client.ZCount(ctx, key, fmt.Sprintf("%f", min), fmt.Sprintf("%f", max)).Result() if err != nil { common.Warn("Redis ZCount error", zap.String("key", key), zap.Error(err)) @@ -493,11 +478,10 @@ func (r *Client) ZCount(key string, min, max float64) int64 { } // ZPopMin pops minimum score members from sorted set -func (r *Client) ZPopMin(key string, count int) ([]redis.Z, error) { +func (r *Client) ZPopMin(ctx context.Context, key string, count int) ([]redis.Z, error) { if r.client == nil { return nil, nil } - ctx := context.Background() members, err := r.client.ZPopMin(ctx, key, int64(count)).Result() if err != nil { common.Warn("Redis ZPopMin error", zap.String("key", key), zap.Error(err)) @@ -507,11 +491,10 @@ func (r *Client) ZPopMin(key string, count int) ([]redis.Z, error) { } // ZRangeByScore returns members with score in range -func (r *Client) ZRangeByScore(key string, min, max float64) ([]string, error) { +func (r *Client) ZRangeByScore(ctx context.Context, key string, min, max float64) ([]string, error) { if r.client == nil { return nil, nil } - ctx := context.Background() members, err := r.client.ZRangeByScore(ctx, key, &redis.ZRangeBy{ Min: fmt.Sprintf("%f", min), Max: fmt.Sprintf("%f", max), @@ -524,11 +507,10 @@ func (r *Client) ZRangeByScore(key string, min, max float64) ([]string, error) { } // ZRemRangeByScore removes members with score in range -func (r *Client) ZRemRangeByScore(key string, min, max float64) int64 { +func (r *Client) ZRemRangeByScore(ctx context.Context, key string, min, max float64) int64 { if r.client == nil { return 0 } - ctx := context.Background() count, err := r.client.ZRemRangeByScore(ctx, key, fmt.Sprintf("%f", min), fmt.Sprintf("%f", max)).Result() if err != nil { common.Warn("Redis ZRemRangeByScore error", zap.String("key", key), zap.Error(err)) @@ -538,11 +520,10 @@ func (r *Client) ZRemRangeByScore(key string, min, max float64) int64 { } // IncrBy increments key by increment -func (r *Client) IncrBy(key string, increment int64) (int64, error) { +func (r *Client) IncrBy(ctx context.Context, key string, increment int64) (int64, error) { if r.client == nil { return 0, nil } - ctx := context.Background() val, err := r.client.IncrBy(ctx, key, increment).Result() if err != nil { common.Warn("Redis IncrBy error", zap.String("key", key), zap.Error(err)) @@ -552,11 +533,10 @@ func (r *Client) IncrBy(key string, increment int64) (int64, error) { } // DecrBy decrements key by decrement -func (r *Client) DecrBy(key string, decrement int64) (int64, error) { +func (r *Client) DecrBy(ctx context.Context, key string, decrement int64) (int64, error) { if r.client == nil { return 0, nil } - ctx := context.Background() val, err := r.client.DecrBy(ctx, key, decrement).Result() if err != nil { common.Warn("Redis DecrBy error", zap.String("key", key), zap.Error(err)) @@ -566,7 +546,7 @@ func (r *Client) DecrBy(key string, decrement int64) (int64, error) { } // GenerateAutoIncrementID generates auto-increment ID -func (r *Client) GenerateAutoIncrementID(keyPrefix string, namespace string, increment int64, ensureMinimum *int64) int64 { +func (r *Client) GenerateAutoIncrementID(ctx context.Context, keyPrefix string, namespace string, increment int64, ensureMinimum *int64) int64 { if r.client == nil { return -1 } @@ -581,7 +561,6 @@ func (r *Client) GenerateAutoIncrementID(keyPrefix string, namespace string, inc } redisKey := fmt.Sprintf("%s:%s", keyPrefix, namespace) - ctx := context.Background() // Check if key exists exists, err := r.client.Exists(ctx, redisKey).Result() @@ -616,11 +595,10 @@ func (r *Client) GenerateAutoIncrementID(keyPrefix string, namespace string, inc } // Transaction sets key with NX flag (transaction-like behavior) -func (r *Client) Transaction(key string, value string, exp time.Duration) bool { +func (r *Client) Transaction(ctx context.Context, key string, value string, exp time.Duration) bool { if r.client == nil { return false } - ctx := context.Background() pipe := r.client.Pipeline() pipe.SetNX(ctx, key, value, exp) _, err := pipe.Exec(ctx) @@ -632,11 +610,10 @@ func (r *Client) Transaction(key string, value string, exp time.Duration) bool { } // QueueProduct produces a message to Redis Stream -func (r *Client) QueueProduct(queue string, message interface{}) bool { +func (r *Client) QueueProduct(ctx context.Context, queue string, message interface{}) bool { if r.client == nil { return false } - ctx := context.Background() for i := 0; i < 3; i++ { data, err := json.Marshal(message) @@ -659,11 +636,10 @@ func (r *Client) QueueProduct(queue string, message interface{}) bool { } // QueueConsumer consumes a message from Redis Stream -func (r *Client) QueueConsumer(queueName, groupName, consumerName string, msgID string) (*Message, error) { +func (r *Client) QueueConsumer(ctx context.Context, queueName, groupName, consumerName string, msgID string) (*Message, error) { if r.client == nil { return nil, nil } - ctx := context.Background() for i := 0; i < 3; i++ { // Create consumer group if not exists @@ -730,11 +706,10 @@ func (r *Client) QueueConsumer(queueName, groupName, consumerName string, msgID } // Ack acknowledges the message -func (m *Message) Ack() bool { +func (m *Message) Ack(ctx context.Context) bool { if m.consumer == nil { return false } - ctx := context.Background() err := m.consumer.XAck(ctx, m.queueName, m.groupName, m.msgID).Err() if err != nil { common.Warn("Message Ack error", zap.Error(err)) @@ -754,11 +729,11 @@ func (m *Message) GetMsgID() string { } // GetPendingMsg gets pending messages -func (r *Client) GetPendingMsg(queue, groupName string) ([]redis.XPendingExt, error) { +func (r *Client) GetPendingMsg(ctx context.Context, queue, groupName string) ([]redis.XPendingExt, error) { if r.client == nil { return nil, nil } - ctx := context.Background() + messages, err := r.client.XPendingExt(ctx, &redis.XPendingExtArgs{ Stream: queue, Group: groupName, @@ -776,11 +751,10 @@ func (r *Client) GetPendingMsg(queue, groupName string) ([]redis.XPendingExt, er } // RequeueMsg re-enqueues a message -func (r *Client) RequeueMsg(queue, groupName, msgID string) { +func (r *Client) RequeueMsg(ctx context.Context, queue, groupName, msgID string) { if r.client == nil { return } - ctx := context.Background() for i := 0; i < 3; i++ { msgs, err := r.client.XRange(ctx, queue, msgID, msgID).Result() @@ -803,11 +777,10 @@ func (r *Client) RequeueMsg(queue, groupName, msgID string) { } // QueueInfo returns queue group info -func (r *Client) QueueInfo(queue, groupName string) (map[string]interface{}, error) { +func (r *Client) QueueInfo(ctx context.Context, queue, groupName string) (map[string]interface{}, error) { if r.client == nil { return nil, nil } - ctx := context.Background() for i := 0; i < 3; i++ { groups, err := r.client.XInfoGroups(ctx, queue).Result() @@ -833,11 +806,10 @@ func (r *Client) QueueInfo(queue, groupName string) (map[string]interface{}, err } // DeleteIfEqual deletes key if its value equals expected value (atomic) -func (r *Client) DeleteIfEqual(key, expectedValue string) bool { +func (r *Client) DeleteIfEqual(ctx context.Context, key, expectedValue string) bool { if r.client == nil { return false } - ctx := context.Background() result, err := r.luaDeleteIfEqual.Run(ctx, r.client, []string{key}, expectedValue).Result() if err != nil { common.Warn("Redis DeleteIfEqual error", zap.Error(err)) @@ -847,11 +819,10 @@ func (r *Client) DeleteIfEqual(key, expectedValue string) bool { } // Delete deletes a key -func (r *Client) Delete(key string) bool { +func (r *Client) Delete(ctx context.Context, key string) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.Del(ctx, key).Err(); err != nil { common.Warn("Redis Delete error", zap.String("key", key), zap.Error(err)) return false @@ -860,11 +831,10 @@ func (r *Client) Delete(key string) bool { } // Expire sets expiration on a key -func (r *Client) Expire(key string, exp time.Duration) bool { +func (r *Client) Expire(ctx context.Context, key string, exp time.Duration) bool { if r.client == nil { return false } - ctx := context.Background() if err := r.client.Expire(ctx, key, exp).Err(); err != nil { common.Warn("Redis Expire error", zap.String("key", key), zap.Error(err)) return false @@ -873,11 +843,10 @@ func (r *Client) Expire(key string, exp time.Duration) bool { } // TTL gets remaining time to live of a key -func (r *Client) TTL(key string) time.Duration { +func (r *Client) TTL(ctx context.Context, key string) time.Duration { if r.client == nil { return -2 } - ctx := context.Background() ttl, err := r.client.TTL(ctx, key).Result() if err != nil { common.Warn("Redis TTL error", zap.String("key", key), zap.Error(err)) @@ -913,13 +882,13 @@ func NewDistributedLock(lockKey string, lockValue string, timeout time.Duration, } // Acquire acquires the lock -func (l *DistributedLock) Acquire() bool { +func (l *DistributedLock) Acquire(ctx context.Context) bool { if l.client == nil { return false } // Delete if stale - l.client.DeleteIfEqual(l.lockKey, l.lockValue) - return l.client.SetNX(l.lockKey, l.lockValue, l.timeout) + l.client.DeleteIfEqual(ctx, l.lockKey, l.lockValue) + return l.client.SetNX(ctx, l.lockKey, l.lockValue, l.timeout) } // SpinAcquire keeps trying to acquire the lock @@ -929,8 +898,8 @@ func (l *DistributedLock) SpinAcquire(ctx context.Context) error { case <-ctx.Done(): return ctx.Err() default: - l.client.DeleteIfEqual(l.lockKey, l.lockValue) - if l.client.SetNX(l.lockKey, l.lockValue, l.timeout) { + l.client.DeleteIfEqual(ctx, l.lockKey, l.lockValue) + if l.client.SetNX(ctx, l.lockKey, l.lockValue, l.timeout) { return nil } time.Sleep(10 * time.Second) @@ -939,11 +908,11 @@ func (l *DistributedLock) SpinAcquire(ctx context.Context) error { } // Release releases the lock -func (l *DistributedLock) Release() bool { +func (l *DistributedLock) Release(ctx context.Context) bool { if l.client == nil { return false } - return l.client.DeleteIfEqual(l.lockKey, l.lockValue) + return l.client.DeleteIfEqual(ctx, l.lockKey, l.lockValue) } // TokenBucket token bucket rate limiter @@ -968,11 +937,10 @@ func NewTokenBucket(key string, capacity, rate float64) *TokenBucket { } // Allow checks if request is allowed -func (tb *TokenBucket) Allow(cost float64) (bool, float64) { +func (tb *TokenBucket) Allow(ctx context.Context, cost float64) (bool, float64) { if tb.client == nil || tb.client.client == nil { return true, 0 } - ctx := context.Background() now := float64(time.Now().Unix()) result, err := tb.client.luaTokenBucket.Run(ctx, tb.client.client, []string{tb.key}, diff --git a/internal/engine/redis/redis_test.go b/internal/engine/redis/redis_test.go index d3fb2006eb..cb77b995b2 100644 --- a/internal/engine/redis/redis_test.go +++ b/internal/engine/redis/redis_test.go @@ -52,7 +52,7 @@ func newStrictTestClient(t *testing.T) (*Client, *miniredis.Miniredis) { // third should be denied. This is the happy-path security gate. func TestEvalTokenBucketStrict_AllowedThenDenied(t *testing.T) { r, _ := newStrictTestClient(t) - ctx := context.Background() + ctx := t.Context() for i := 1; i <= 2; i++ { ok, err := r.EvalTokenBucketStrict(ctx, "tb:webhook", 2, 0.1) @@ -80,7 +80,7 @@ func TestEvalTokenBucketStrict_RedisDownFailsClosed(t *testing.T) { r, mr := newStrictTestClient(t) mr.Close() // break the connection - ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) defer cancel() ok, err := r.EvalTokenBucketStrict(ctx, "tb:webhook", 5, 1) @@ -97,7 +97,7 @@ func TestEvalTokenBucketStrict_RedisDownFailsClosed(t *testing.T) { // (false, error) so the webhook handler can surface 102. func TestEvalTokenBucketStrict_NilClient(t *testing.T) { var r *Client - ok, err := r.EvalTokenBucketStrict(context.Background(), "tb:webhook", 1, 1) + ok, err := r.EvalTokenBucketStrict(t.Context(), "tb:webhook", 1, 1) if err == nil { t.Fatalf("expected error on nil client, got nil") } diff --git a/internal/handler/agent.go b/internal/handler/agent.go index df2301ff5d..c2bdd64f4d 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -129,13 +129,13 @@ func (h *AgentHandler) WithDocumentService(s documentAccessChecker) *AgentHandle // NewAgentHandler create agent handler -func NewAgentHandler(agentService *service.AgentService, fileService *file.FileService) *AgentHandler { +func NewAgentHandler(ctx context.Context, agentService *service.AgentService, fileService *file.FileService) *AgentHandler { return &AgentHandler{ agentService: agentService, chatRunner: agentService, fileService: fileService, loader: agentService, - redisGet: func(key string) (string, error) { return redis.Get().Get(key) }, + redisGet: func(key string) (string, error) { return redis.Get().Get(ctx, key) }, redisStore: redis.Get(), newExecutor: func(taskCtx *task.TaskContext, canvasID string, docBulkSize int) (debugExecutor, error) { return task.NewPipelineExecutor(taskCtx, canvasID, docBulkSize) diff --git a/internal/handler/agent_logs_test.go b/internal/handler/agent_logs_test.go index 793b4fe9f0..3b0bfd1459 100644 --- a/internal/handler/agent_logs_test.go +++ b/internal/handler/agent_logs_test.go @@ -113,6 +113,7 @@ func TestParseAgentLogs(t *testing.T) { func TestGetAgentLogs_E2EViaMiniredis(t *testing.T) { gin.SetMode(gin.TestMode) + ctx := t.Context() db := setupHandlerAgentsTestDB(t) orig := dao.DB dao.DB = db @@ -134,13 +135,13 @@ func TestGetAgentLogs_E2EViaMiniredis(t *testing.T) { `{"component_id":"File","trace":[{"progress":1,"message":"parsed","datetime":"10:00:00","timestamp":1.0,"elapsed_time":0}]},` + `{"component_id":"END","trace":[{"progress":1,"message":"done","datetime":"10:00:01","timestamp":2.0,"elapsed_time":1.0}]}` + `]` - if err := rdb.Set(context.Background(), logKey, arrayPayload, 0).Err(); err != nil { + if err = rdb.Set(ctx, logKey, arrayPayload, 0).Err(); err != nil { t.Fatalf("seed redis: %v", err) } - h := NewAgentHandler(service.NewAgentService(), nil). + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithRedisGetter(func(key string) (string, error) { - return rdb.Get(context.Background(), key).Result() + return rdb.Get(ctx, key).Result() }) run := func(messageID string) map[string]interface{} { @@ -224,6 +225,7 @@ func clientConsidersComplete(arr []map[string]interface{}) bool { // byte-for-byte through JSON round-tripping. func TestGetAgentLogs_EndSignalCompletion(t *testing.T) { gin.SetMode(gin.TestMode) + ctx := t.Context() db := setupHandlerAgentsTestDB(t) orig := dao.DB @@ -246,7 +248,7 @@ func TestGetAgentLogs_EndSignalCompletion(t *testing.T) { `{"component_id":"File","trace":[{"progress":1,"message":"parsed","datetime":"10:00:00","timestamp":1.0,"elapsed_time":0}]},` + `{"component_id":"END","trace":[{"progress":1,"message":"run finished","datetime":"10:00:01","timestamp":2.0,"elapsed_time":1.0}]}` + `]` - if err := rdb.Set(context.Background(), "c1-msg-good-logs", goodPayload, 0).Err(); err != nil { + if err = rdb.Set(ctx, "c1-msg-good-logs", goodPayload, 0).Err(); err != nil { t.Fatalf("seed redis: %v", err) } @@ -257,13 +259,13 @@ func TestGetAgentLogs_EndSignalCompletion(t *testing.T) { `{"component_id":"File","trace":[{"progress":1,"message":"parsed","datetime":"10:00:00","timestamp":1.0,"elapsed_time":0}]},` + `{"component_id":"END","trace":[{"progress":1,"message":"","datetime":"10:00:01","timestamp":2.0,"elapsed_time":1.0}]}` + `]` - if err := rdb.Set(context.Background(), "c1-msg-bad-logs", badPayload, 0).Err(); err != nil { + if err = rdb.Set(ctx, "c1-msg-bad-logs", badPayload, 0).Err(); err != nil { t.Fatalf("seed redis: %v", err) } - h := NewAgentHandler(service.NewAgentService(), nil). + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithRedisGetter(func(key string) (string, error) { - return rdb.Get(context.Background(), key).Result() + return rdb.Get(ctx, key).Result() }) call := func(messageID string) []map[string]interface{} { @@ -310,7 +312,7 @@ type capturedStore struct { data map[string]string } -func (s *capturedStore) Set(key, value string, _ time.Duration) bool { +func (s *capturedStore) Set(ctx context.Context, key, value string, _ time.Duration) bool { s.mu.Lock() defer s.mu.Unlock() if s.data == nil { @@ -320,7 +322,7 @@ func (s *capturedStore) Set(key, value string, _ time.Duration) bool { return true } -func (s *capturedStore) get(key string) (string, bool) { +func (s *capturedStore) Get(ctx context.Context, key string) (string, bool) { s.mu.Lock() defer s.mu.Unlock() v, ok := s.data[key] @@ -463,10 +465,10 @@ func TestRespondWithDebugResult_ErrorCarriesMessageID(t *testing.T) { // before — the failure log was written but unreachable because message_id was // dropped on the error path. func TestRunCanvasPipelineDebug_ErrorStillExposesMessageID(t *testing.T) { - ctx := context.Background() + ctx := t.Context() store := &capturedStore{} - h := NewAgentHandler(service.NewAgentService(), nil). + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithRedisStore(store). WithNewExecutor(func(taskCtx *task.TaskContext, canvasID string, docBulkSize int) (debugExecutor, error) { return &fakeDebugExecutor{ @@ -489,7 +491,7 @@ func TestRunCanvasPipelineDebug_ErrorStillExposesMessageID(t *testing.T) { // The failure log must be written under the composed key. key := "c1-" + result.MessageID + "-logs" - raw, ok := store.get(key) + raw, ok := store.Get(ctx, key) if !ok { t.Fatalf("failure log not written under key %q; store keys=%v", key, keysOf(store)) } @@ -511,10 +513,10 @@ func TestRunCanvasPipelineDebug_ErrorStillExposesMessageID(t *testing.T) { // satisfy the completion predicate (END last, non-empty END message) so the // Log box stops polling. func TestRunCanvasPipelineDebug_WiresMessageIDAndLog(t *testing.T) { - ctx := context.Background() + ctx := t.Context() store := &capturedStore{} - h := NewAgentHandler(service.NewAgentService(), nil). + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithRedisStore(store). WithNewExecutor(func(taskCtx *task.TaskContext, canvasID string, docBulkSize int) (debugExecutor, error) { return &fakeDebugExecutor{ @@ -537,7 +539,7 @@ func TestRunCanvasPipelineDebug_WiresMessageIDAndLog(t *testing.T) { // The log array must be written under the composed key. key := "c1-" + result.MessageID + "-logs" - raw, ok := store.get(key) + raw, ok := store.Get(ctx, key) if !ok { t.Fatalf("log not written under key %q; store keys=%v", key, keysOf(store)) } @@ -570,8 +572,8 @@ type miniredisDebugStore struct { rdb *goredis.Client } -func (s miniredisDebugStore) Set(key, value string, ttl time.Duration) bool { - if err := s.rdb.Set(context.Background(), key, value, ttl).Err(); err != nil { +func (s miniredisDebugStore) Set(ctx context.Context, key, value string, ttl time.Duration) bool { + if err := s.rdb.Set(ctx, key, value, ttl).Err(); err != nil { return false } return true @@ -589,7 +591,7 @@ func (s miniredisDebugStore) Set(key, value string, ttl time.Duration) bool { // seeds Redis directly rather than going through the writer. func TestRunCanvasPipelineDebug_WriteThenReadViaMiniredis(t *testing.T) { gin.SetMode(gin.TestMode) - ctx := context.Background() + ctx := t.Context() db := setupHandlerAgentsTestDB(t) orig := dao.DB @@ -608,7 +610,7 @@ func TestRunCanvasPipelineDebug_WriteThenReadViaMiniredis(t *testing.T) { // Both seams point at the same miniredis: the writer stores via // WithRedisStore and the reader fetches via WithRedisGetter, mirroring // production where both hit one Redis. - h := NewAgentHandler(service.NewAgentService(), nil). + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithRedisStore(miniredisDebugStore{rdb: rdb}). WithRedisGetter(func(key string) (string, error) { return rdb.Get(ctx, key).Result() diff --git a/internal/handler/agent_test.go b/internal/handler/agent_test.go index 8e7501fdb1..da88e819c4 100644 --- a/internal/handler/agent_test.go +++ b/internal/handler/agent_test.go @@ -112,7 +112,8 @@ func TestListAgentVersionsHandler_Success(t *testing.T) { }, }) - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.ListVersions(c) if w.Code != http.StatusOK { @@ -162,7 +163,8 @@ func TestListAgentVersionsHandler_NoPermission(t *testing.T) { // Canvas owned by user-b db.Create(&entity.UserCanvas{ID: "canvas-b", UserID: "user-b", Title: sptr("Not Yours")}) - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.ListVersions(c) var resp map[string]interface{} @@ -192,7 +194,8 @@ func TestListAgentVersionsHandler_CanvasNotFound(t *testing.T) { c.Set("user_id", "user-1") c.Params = gin.Params{{Key: "canvas_id", Value: "non-existent"}} - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.ListVersions(c) var resp map[string]interface{} @@ -244,7 +247,8 @@ func TestGetAgentVersionHandler_Success(t *testing.T) { }, }) - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.GetVersion(c) if w.Code != http.StatusOK { @@ -291,7 +295,8 @@ func TestGetAgentVersionHandler_VersionNotFound(t *testing.T) { Title: sptr("Test Agent"), }) - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.GetVersion(c) var resp map[string]interface{} @@ -658,7 +663,8 @@ func TestAgentChatCompletions_RequiresAgentID(t *testing.T) { c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.AgentChatCompletions(c) if w.Code != http.StatusOK { @@ -686,7 +692,8 @@ func TestAgentChatCompletions_OpenAICompat_EmptyMessages(t *testing.T) { c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.AgentChatCompletions(c) var resp map[string]interface{} @@ -998,7 +1005,8 @@ func TestAgentChatCompletions_OpenAICompat_NonStreamReturnsChoices(t *testing.T) c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.AgentChatCompletions(c) var resp map[string]interface{} @@ -1035,7 +1043,8 @@ func TestRerunAgent_RequiresAllFields(t *testing.T) { c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.RerunAgent(c) var resp map[string]interface{} @@ -1068,7 +1077,8 @@ func TestRerunAgent_AcceptsCompleteRequest(t *testing.T) { c.Set("user_id", "u1") stub := &stubDocService{accessible: true} - h := NewAgentHandler(service.NewAgentService(), nil). + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithDocumentService(stub) h.RerunAgent(c) @@ -1089,7 +1099,8 @@ func TestPromptsReturnsHardcodedFields(t *testing.T) { c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.Prompts(c) var resp map[string]interface{} @@ -1126,7 +1137,8 @@ func TestGetAgentWebhookLogsReturnsEmptyPoll(t *testing.T) { c.Set("user_id", "u1") c.Params = gin.Params{{Key: "canvas_id", Value: "c1"}} - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) h.GetAgentWebhookLogs(c) var resp map[string]interface{} @@ -1169,7 +1181,8 @@ func TestRerunAgent_RejectsInaccessibleDocument(t *testing.T) { // round 5), so the deny-all stub injects cleanly without standing // up the real DocumentService (DB, storage, ...). stub := &stubDocService{accessible: false} - h := NewAgentHandler(service.NewAgentService(), nil). + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil). WithDocumentService(stub) h.RerunAgent(c) @@ -1201,7 +1214,8 @@ func TestRerunAgent_NoDocumentServiceFailsClosed(t *testing.T) { c.Set("user", &entity.User{ID: "u1"}) c.Set("user_id", "u1") - h := NewAgentHandler(service.NewAgentService(), nil) + ctx := t.Context() + h := NewAgentHandler(ctx, service.NewAgentService(), nil) // Note: no WithDocumentService call → documentService is nil. // Production wiring (cmd/server_main.go) always calls // WithDocumentService; a nil here means the handler was diff --git a/internal/handler/agent_webhook.go b/internal/handler/agent_webhook.go index ef2ded7725..53a43fedfa 100644 --- a/internal/handler/agent_webhook.go +++ b/internal/handler/agent_webhook.go @@ -520,7 +520,7 @@ func (h *AgentHandler) runWebhookDetached( zap.String("canvas", cv.ID), zap.Error(err)) if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) + appendWebhookTrace(ctx, cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) } return } @@ -529,7 +529,7 @@ func (h *AgentHandler) runWebhookDetached( ev.SessionID = sessionID } if isTest { - appendWebhookTrace(cv.ID, startTs, ev) + appendWebhookTrace(ctx, cv.ID, startTs, ev) } } } @@ -553,8 +553,8 @@ func (h *AgentHandler) runWebhookSync( events, err := h.loader.RunAgentWithWebhook(ctx, cv.UserID, cv.ID, payload) if err != nil { if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": false})}) + appendWebhookTrace(ctx, cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) + appendWebhookTrace(ctx, cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": false})}) } return webhookSyncResult{status: http.StatusBadRequest, body: gin.H{ "code": 400, @@ -571,7 +571,7 @@ func (h *AgentHandler) runWebhookSync( ev.SessionID = sessionID } if isTest { - appendWebhookTrace(cv.ID, startTs, ev) + appendWebhookTrace(ctx, cv.ID, startTs, ev) } switch ev.Type { case "message": @@ -602,7 +602,7 @@ func (h *AgentHandler) runWebhookSync( } final := strings.Join(contents, "") if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": true})}) + appendWebhookTrace(ctx, cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": true})}) } return webhookSyncResult{status: status, body: gin.H{ "message": final, @@ -629,14 +629,14 @@ func mustJSON(v any) string { // The trace key is `webhook-trace--logs` with a 600 s TTL. // Each event is recorded as {"ts": , "event": , ...}. // Tests use miniredis to verify the key shape. -func appendWebhookTrace(agentID string, startTs time.Time, ev canvas.RunEvent) { +func appendWebhookTrace(ctx context.Context, agentID string, startTs time.Time, ev canvas.RunEvent) { rdb := rediscli.Get() if rdb == nil { return } key := fmt.Sprintf("webhook-trace-%s-logs", agentID) - raw, _ := rdb.Get(key) + raw, _ := rdb.Get(ctx, key) obj := map[string]any{} if raw != "" { _ = json.Unmarshal([]byte(raw), &obj) @@ -677,5 +677,5 @@ func appendWebhookTrace(agentID string, startTs time.Time, ev canvas.RunEvent) { common.Warn("webhook trace marshal failed", zap.Error(err)) return } - rdb.SetObj(key, string(encoded), 600*time.Second) + rdb.SetObj(ctx, key, string(encoded), 600*time.Second) } diff --git a/internal/handler/system.go b/internal/handler/system.go index 6b471ee10d..e0756f58e0 100644 --- a/internal/handler/system.go +++ b/internal/handler/system.go @@ -109,7 +109,8 @@ func (h *SystemHandler) GetStatus(c *gin.Context) { return } - status, err := h.systemService.GetStatus() + ctx := c.Request.Context() + status, err := h.systemService.GetStatus(ctx) if err != nil { jsonInternalError(c, err) return diff --git a/internal/handler/user.go b/internal/handler/user.go index 7ab49662dd..23fe7ca1b8 100644 --- a/internal/handler/user.go +++ b/internal/handler/user.go @@ -99,7 +99,7 @@ func (h *UserHandler) Register(c *gin.Context) { return } - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, false, err.Error()) return @@ -169,7 +169,7 @@ func (h *UserHandler) Login(c *gin.Context) { operationLog.UserID = user.ID // Sign the access_token using itsdangerous (compatible with Python) - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { errMessage := fmt.Sprintf("Failed to get secret key: %s", err.Error()) common.ResponseWithCodeData(c, common.CodeServerError, false, errMessage) @@ -256,7 +256,7 @@ func (h *UserHandler) LoginByEmail(c *gin.Context) { } operationLog.UserID = user.ID - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { errorMessage := fmt.Sprintf("Failed to get secret key: %s", err.Error()) common.ResponseWithCodeData(c, common.CodeServerError, false, errorMessage) @@ -713,7 +713,7 @@ func (h *UserHandler) ForgotResetPassword(c *gin.Context) { return } - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, false, fmt.Sprintf("Failed to get secret key: %s", err.Error())) return diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 2d33629441..2c69fbaa72 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -921,7 +921,7 @@ func (c *ExtractorComponent) runEnableMetadata(ctx context.Context, db *gorm.DB, // Best-effort: a missing Redis client or any cache error falls through to // a live call instead of failing the extraction. var parsed map[string]any - if cached, hit := getMetadataLLMCache(in.llmID, schemaStr, chunkText); hit { + if cached, hit := getMetadataLLMCache(ctx, in.llmID, schemaStr, chunkText); hit { parsed = cached } else { metaTemp := extractorTemperature @@ -956,7 +956,7 @@ func (c *ExtractorComponent) runEnableMetadata(ctx context.Context, db *gorm.DB, return nil } parsed = parsedObj - setMetadataLLMCache(in.llmID, schemaStr, chunkText, parsed) + setMetadataLLMCache(ctx, in.llmID, schemaStr, chunkText, parsed) } // Merge into the chunk metadata map, preserving existing keys. var meta map[string]any @@ -997,17 +997,17 @@ func metadataLLMCacheKey(llmID, schemaJSON, chunkText string) string { // getMetadataLLMCache returns a cached extraction for the given chunk, or // (nil, false) on miss / Redis unavailable / decode error. Best-effort. -func getMetadataLLMCache(llmID, schemaJSON, chunkText string) (map[string]any, bool) { +func getMetadataLLMCache(ctx context.Context, llmID, schemaJSON, chunkText string) (map[string]any, bool) { client := redis.Get() if client == nil { return nil, false } - data, err := client.Get(metadataLLMCacheKey(llmID, schemaJSON, chunkText)) + data, err := client.Get(ctx, metadataLLMCacheKey(llmID, schemaJSON, chunkText)) if err != nil || data == "" { return nil, false } var parsed map[string]any - if err := json.Unmarshal([]byte(data), &parsed); err != nil { + if err = json.Unmarshal([]byte(data), &parsed); err != nil { return nil, false } return parsed, true @@ -1015,7 +1015,7 @@ func getMetadataLLMCache(llmID, schemaJSON, chunkText string) (map[string]any, b // setMetadataLLMCache stores an extraction result for 24h. Best-effort: a // missing Redis client or marshal error is silently ignored. -func setMetadataLLMCache(llmID, schemaJSON, chunkText string, parsed map[string]any) { +func setMetadataLLMCache(ctx context.Context, llmID, schemaJSON, chunkText string, parsed map[string]any) { client := redis.Get() if client == nil { return @@ -1024,7 +1024,7 @@ func setMetadataLLMCache(llmID, schemaJSON, chunkText string, parsed map[string] if err != nil { return } - client.Set(metadataLLMCacheKey(llmID, schemaJSON, chunkText), string(data), metadataLLMCacheTTL) + client.Set(ctx, metadataLLMCacheKey(llmID, schemaJSON, chunkText), string(data), metadataLLMCacheTTL) } // cleanExtractionResult strips `` tags and rejects `**ERROR**` responses, diff --git a/internal/ingestion/component/extractor_tag.go b/internal/ingestion/component/extractor_tag.go index 1b1a221000..69276dc16a 100644 --- a/internal/ingestion/component/extractor_tag.go +++ b/internal/ingestion/component/extractor_tag.go @@ -664,7 +664,7 @@ func llmTagChunk( return } - if cached := getTaggerLLMCache(llmID, text, allTags, topN); cached != nil { + if cached := getTaggerLLMCache(ctx, llmID, text, allTags, topN); cached != nil { chunk[common.TAG_FLD] = cached return } @@ -713,7 +713,7 @@ func llmTagChunk( if len(result) > 0 { chunk[common.TAG_FLD] = result - setTaggerLLMCache(llmID, text, allTags, topN, result) + setTaggerLLMCache(ctx, llmID, text, allTags, topN, result) } } @@ -803,13 +803,13 @@ func taggerCacheKey(llmID, text string, allTags map[string]float64, topN int) st return fmt.Sprintf("tagger:%x", hasher.Sum64()) } -func getTaggerLLMCache(llmID, text string, allTags map[string]float64, topN int) map[string]int { +func getTaggerLLMCache(ctx context.Context, llmID, text string, allTags map[string]float64, topN int) map[string]int { client := redis.Get() if client == nil { return nil } key := taggerCacheKey(llmID, text, allTags, topN) - data, err := client.Get(key) + data, err := client.Get(ctx, key) if err != nil || data == "" { return nil } @@ -820,7 +820,7 @@ func getTaggerLLMCache(llmID, text string, allTags map[string]float64, topN int) return result } -func setTaggerLLMCache(llmID, text string, allTags map[string]float64, topN int, result map[string]int) { +func setTaggerLLMCache(ctx context.Context, llmID, text string, allTags map[string]float64, topN int, result map[string]int) { if result == nil { return } @@ -833,7 +833,7 @@ func setTaggerLLMCache(llmID, text string, allTags map[string]float64, topN int, if err != nil { return } - client.Set(key, string(data), 24*time.Hour) + client.Set(ctx, key, string(data), 24*time.Hour) } func sortedTagNames(allTags map[string]float64) []string { diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index a73aec65ef..dcb1c1a16b 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -460,7 +460,7 @@ func (e *Ingestor) markStopped(ctx context.Context, taskID string) bool { } if rc := redis2.Get(); rc != nil { utility.BestEffort(fmt.Sprintf("clear cancel flag for %s", taskID), func() error { - rc.Delete(fmt.Sprintf("%s-cancel", taskID)) + rc.Delete(ctx, fmt.Sprintf("%s-cancel", taskID)) return nil // Delete returns bool; the bool does not distinguish "not found" from "error" }) } @@ -503,7 +503,7 @@ func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool if rc := redis2.Get(); rc != nil { key := fmt.Sprintf("%s-cancel", task.ID) utility.BestEffort(fmt.Sprintf("clear stale cancel flag for %s", task.ID), func() error { - rc.Delete(key) + rc.Delete(ctx, key) return nil // Delete returns bool; false may mean "key not found" or "error" }) } @@ -665,7 +665,7 @@ func (e *Ingestor) ackOrNack(taskCtx *taskpkg.TaskContext, terminal bool) { func (e *Ingestor) defaultCancelCheck(ctx context.Context, taskID string) bool { rc := redis2.Get() if rc != nil { - if ok, _ := rc.Exist(fmt.Sprintf("%s-cancel", taskID)); ok { + if ok, _ := rc.Exist(ctx, fmt.Sprintf("%s-cancel", taskID)); ok { return true } } diff --git a/internal/ingestion/task/debug_log_sink.go b/internal/ingestion/task/debug_log_sink.go index ce55ae406a..51b3a7dc91 100644 --- a/internal/ingestion/task/debug_log_sink.go +++ b/internal/ingestion/task/debug_log_sink.go @@ -35,7 +35,7 @@ const DebugLogTTL = 30 * time.Minute // ttl)); tests pass a capturing closure. Keeping it to a single method keeps the // sink free of any redis-package import. type DebugLogStore interface { - Set(key, value string, ttl time.Duration) bool + Set(ctx context.Context, key, value string, ttl time.Duration) bool } // funcStore adapts a plain function to DebugLogStore so tests (and callers that @@ -250,7 +250,7 @@ func (s *DebugLogSink) Flush(ctx context.Context, finalErr error) { } } key := s.canvasID + "-" + s.messageID + "-logs" - s.store.Set(key, string(payload), s.ttl) + s.store.Set(ctx, key, string(payload), s.ttl) } // truncateRunes returns s truncated to at most max runes, preserving the original diff --git a/internal/ingestion/task/debug_log_sink_test.go b/internal/ingestion/task/debug_log_sink_test.go index 97388a33d6..3e0c306ffa 100644 --- a/internal/ingestion/task/debug_log_sink_test.go +++ b/internal/ingestion/task/debug_log_sink_test.go @@ -28,9 +28,10 @@ import ( "time" "unicode/utf8" - "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/pipeline" + + "gorm.io/gorm" ) // capturedStore is a fake DebugLogStore that records every write so tests can @@ -40,7 +41,7 @@ type capturedStore struct { data map[string]string } -func (c *capturedStore) Set(key, value string, ttl time.Duration) bool { +func (c *capturedStore) Set(ctx context.Context, key, value string, ttl time.Duration) bool { c.mu.Lock() defer c.mu.Unlock() if c.data == nil { @@ -50,7 +51,7 @@ func (c *capturedStore) Set(key, value string, ttl time.Duration) bool { return true } -func (c *capturedStore) get(key string) string { +func (c *capturedStore) Get(ctx context.Context, key string) string { c.mu.Lock() defer c.mu.Unlock() return c.data[key] @@ -96,20 +97,21 @@ func loadArray(t *testing.T, raw string) []map[string]any { func TestDebugLogSink_RecordsTraceAndEndMarker(t *testing.T) { store := &capturedStore{} sink := NewDebugLogSink("c1", "m1", store) + ctx := t.Context() - sink.OnComponentProgress(context.Background(), pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ Component: "File", Message: "File Started", Phase: phaseEnter, }) - sink.OnComponentProgress(context.Background(), pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ Component: "File", Message: "File Done", Phase: phaseExit, }) - sink.OnComponentProgress(context.Background(), pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ Component: "Chunker", Message: "Chunker Done", Phase: phaseExit, }) - sink.Flush(context.Background(), nil) + sink.Flush(ctx, nil) - raw := store.get("c1-m1-logs") + raw := store.Get(ctx, "c1-m1-logs") if raw == "" { t.Fatalf("expected key c1-m1-logs to be written") } @@ -148,13 +150,14 @@ func TestDebugLogSink_RecordsTraceAndEndMarker(t *testing.T) { func TestDebugLogSink_ErrorPrefix(t *testing.T) { store := &capturedStore{} sink := NewDebugLogSink("c1", "m1", store) + ctx := t.Context() - sink.OnComponentProgress(context.Background(), pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ Component: "Tokenizer", Message: "Tokenizer: boom", Phase: phaseError, }) - sink.Flush(context.Background(), errors.New("boom")) + sink.Flush(ctx, errors.New("boom")) - raw := store.get("c1-m1-logs") + raw := store.Get(ctx, "c1-m1-logs") arr := loadArray(t, raw) // First element is the errored component; its message must carry [ERROR]. @@ -180,9 +183,11 @@ func TestDebugLogSink_ErrorPrefix(t *testing.T) { // the debug log array shape does not use it. func TestDebugLogSink_OnComponentTotalIsNoOp(t *testing.T) { store := &capturedStore{} + ctx := t.Context() + sink := NewDebugLogSink("c1", "m1", store) // Must not panic and must not write. - sink.OnComponentTotal(context.Background(), "task-1", 5) + sink.OnComponentTotal(ctx, "task-1", 5) if len(store.data) != 0 { t.Errorf("OnComponentTotal must not write to the store, wrote %v", store.data) } @@ -195,7 +200,7 @@ func TestDebugLogSink_OnComponentTotalIsNoOp(t *testing.T) { // pipeline and its events land in the persisted debug-log array (with the END // marker), without touching the DB/index persist path. func TestPipelineExecutor_DebugRunWritesLogViaSink(t *testing.T) { - ctx := context.Background() + ctx := t.Context() taskCtx := NewDebugTaskContext("tenant-1", "c1", "doc.pdf", nil) exec, err := NewPipelineExecutor(taskCtx, "c1", 0) if err != nil { @@ -228,7 +233,7 @@ func TestPipelineExecutor_DebugRunWritesLogViaSink(t *testing.T) { // Mirrors runCanvasPipelineDebug: flush (incl. END marker) after the run. sink.Flush(ctx, nil) - raw := store.get("c1-m1-logs") + raw := store.Get(ctx, "c1-m1-logs") if raw == "" { t.Fatalf("expected debug log written to c1-m1-logs") } @@ -280,6 +285,7 @@ func (traceChunkComponent) Invoke(_ context.Context, _ *gorm.DB, _ map[string]an // observation that a debug run's polled log appeared to stay at progress 0 — // per the code, each component's trace must contain [Started, Done]. func TestDebugLogSink_RealPipeline_EachComponentTraceHasStartedAndDone(t *testing.T) { + ctx := t.Context() const ( compA = "trace.RealStubA" compB = "trace.RealStubB" @@ -305,12 +311,12 @@ func TestDebugLogSink_RealPipeline_EachComponentTraceHasStartedAndDone(t *testin if err != nil { t.Fatalf("NewPipelineFromDSL: %v", err) } - if _, err := pipe.Run(context.Background(), map[string]any{"name": "doc-trace"}, nil); err != nil { + if _, err = pipe.Run(ctx, map[string]any{"name": "doc-trace"}, nil); err != nil { t.Fatalf("Run: %v", err) } - sink.Flush(context.Background(), nil) + sink.Flush(ctx, nil) - raw := store.get("c-trace-m-trace-logs") + raw := store.Get(ctx, "c-trace-m-trace-logs") if raw == "" { t.Fatalf("expected debug log written to c-trace-m-trace-logs") } @@ -370,6 +376,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerCarriesDSL(t *testing.T) { compC = "trace.RealStubC" compD = "trace.RealStubD" ) + ctx := t.Context() runtime.MustRegister(compC, runtime.CategoryIngestion, func(_ string, _ map[string]any) (runtime.Component, error) { return traceStubComponent{}, nil }, runtime.Metadata{Version: "1.0.0"}) @@ -391,7 +398,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerCarriesDSL(t *testing.T) { if err != nil { t.Fatalf("NewPipelineFromDSL: %v", err) } - output, err := pipe.Run(context.Background(), map[string]any{"name": "doc-dsl"}, nil) + output, err := pipe.Run(ctx, map[string]any{"name": "doc-dsl"}, nil) if err != nil { t.Fatalf("Run: %v", err) } @@ -402,9 +409,9 @@ func TestDebugLogSink_RealPipeline_EndMarkerCarriesDSL(t *testing.T) { t.Fatalf("BuildDebugResultDSL: %v", err) } sink.SetResult(resultDSL) - sink.Flush(context.Background(), nil) + sink.Flush(ctx, nil) - raw := store.get("c-dsl-m-dsl-logs") + raw := store.Get(ctx, "c-dsl-m-dsl-logs") if raw == "" { t.Fatalf("expected debug log written to c-dsl-m-dsl-logs") } @@ -467,6 +474,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerDSLShowsChunks(t *testing.T) { compC = "trace.RealStubChunks" compD = "trace.RealStubD2" ) + ctx := t.Context() runtime.MustRegister(compC, runtime.CategoryIngestion, func(_ string, _ map[string]any) (runtime.Component, error) { return traceChunkComponent{}, nil }, runtime.Metadata{Version: "1.0.0"}) @@ -488,7 +496,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerDSLShowsChunks(t *testing.T) { if err != nil { t.Fatalf("NewPipelineFromDSL: %v", err) } - output, err := pipe.Run(context.Background(), map[string]any{"name": "doc-chunk"}, nil) + output, err := pipe.Run(ctx, map[string]any{"name": "doc-chunk"}, nil) if err != nil { t.Fatalf("Run: %v", err) } @@ -498,9 +506,9 @@ func TestDebugLogSink_RealPipeline_EndMarkerDSLShowsChunks(t *testing.T) { t.Fatalf("BuildDebugResultDSL: %v", err) } sink.SetResult(resultDSL) - sink.Flush(context.Background(), nil) + sink.Flush(ctx, nil) - raw := store.get("c-chunk-m-chunk-logs") + raw := store.Get(ctx, "c-chunk-m-chunk-logs") if raw == "" { t.Fatalf("expected debug log written") } @@ -514,7 +522,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerDSLShowsChunks(t *testing.T) { var dslDoc map[string]any switch v := endFirst["dsl"].(type) { case string: - if err := json.Unmarshal([]byte(v), &dslDoc); err != nil { + if err = json.Unmarshal([]byte(v), &dslDoc); err != nil { t.Fatalf("END dsl not valid JSON: %v", err) } case map[string]any: @@ -569,7 +577,7 @@ func TestDebugLogSink_RealPipeline_EndMarkerDSLShowsChunks(t *testing.T) { func TestDebugLogSink_ElapsedTimeIsInSeconds(t *testing.T) { store := &capturedStore{} sink := NewDebugLogSink("c-elapsed", "m-elapsed", store) - ctx := context.Background() + ctx := t.Context() t0 := time.Now() sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ @@ -582,7 +590,7 @@ func TestDebugLogSink_ElapsedTimeIsInSeconds(t *testing.T) { }) sink.Flush(ctx, nil) - raw := store.get("c-elapsed-m-elapsed-logs") + raw := store.Get(ctx, "c-elapsed-m-elapsed-logs") if raw == "" { t.Fatalf("expected debug log written to c-elapsed-m-elapsed-logs") } @@ -621,7 +629,7 @@ func TestDebugLogSink_ElapsedTimeIsInSeconds(t *testing.T) { func TestDebugLogSink_TimestampIsInSeconds(t *testing.T) { store := &capturedStore{} sink := NewDebugLogSink("c-ts", "m-ts", store) - ctx := context.Background() + ctx := t.Context() before := time.Now().Unix() sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ @@ -629,7 +637,7 @@ func TestDebugLogSink_TimestampIsInSeconds(t *testing.T) { }) sink.Flush(ctx, nil) - raw := store.get("c-ts-m-ts-logs") + raw := store.Get(ctx, "c-ts-m-ts-logs") if raw == "" { t.Fatalf("expected debug log written to c-ts-m-ts-logs") } @@ -661,7 +669,7 @@ func TestDebugLogSink_TimestampIsInSeconds(t *testing.T) { // Counts are deliberately bounded so the test stays fast: we only need to exceed // each cap by a small margin to prove clamping, not replay millions of events. func TestDebugLogSink_RespectsCaps(t *testing.T) { - ctx := context.Background() + ctx := t.Context() // Component cap: more distinct components than maxLogEntries (each one trace, // one over-long message that must be truncated to the rune cap). @@ -676,7 +684,7 @@ func TestDebugLogSink_RespectsCaps(t *testing.T) { } sinkA.Flush(ctx, nil) - arrA := loadArray(t, storeA.get("c1-mA-logs")) + arrA := loadArray(t, storeA.Get(ctx, "c1-mA-logs")) // Flush always appends the terminal END marker, so count only real components. var compCountA int for _, el := range arrA { @@ -702,7 +710,7 @@ func TestDebugLogSink_RespectsCaps(t *testing.T) { } } if !clientConsidersComplete(arrA) { - t.Errorf("clamped log must still satisfy completion predicate; raw=%s", storeA.get("c1-mA-logs")) + t.Errorf("clamped log must still satisfy completion predicate; raw=%s", storeA.Get(ctx, "c1-mA-logs")) } // Trace cap: a single component with more traces than maxTracePerEntry. @@ -717,17 +725,17 @@ func TestDebugLogSink_RespectsCaps(t *testing.T) { } sinkB.Flush(ctx, nil) - arrB := loadArray(t, storeB.get("c1-mB-logs")) + arrB := loadArray(t, storeB.Get(ctx, "c1-mB-logs")) // Busy (the only real component) + the END completion marker = 2 elements. if len(arrB) != 2 { - t.Fatalf("Busy + END expected, got %d elements: %s", len(arrB), storeB.get("c1-mB-logs")) + t.Fatalf("Busy + END expected, got %d elements: %s", len(arrB), storeB.Get(ctx, "c1-mB-logs")) } busyTrace, _ := arrB[0]["trace"].([]any) if len(busyTrace) > maxTracePerEntry { t.Fatalf("Busy trace len=%d want <= %d", len(busyTrace), maxTracePerEntry) } if !clientConsidersComplete(arrB) { - t.Errorf("trace-capped log must still satisfy completion predicate; raw=%s", storeB.get("c1-mB-logs")) + t.Errorf("trace-capped log must still satisfy completion predicate; raw=%s", storeB.Get(ctx, "c1-mB-logs")) } } @@ -738,7 +746,7 @@ func TestDebugLogSink_RespectsCaps(t *testing.T) { func TestDebugLogSink_PayloadHardCap(t *testing.T) { store := &capturedStore{} sink := NewDebugLogSink("c1", "m1", store) - ctx := context.Background() + ctx := t.Context() for i := 0; i < maxLogEntries; i++ { comp := fmt.Sprintf("Comp-%d", i) @@ -752,7 +760,7 @@ func TestDebugLogSink_PayloadHardCap(t *testing.T) { } sink.Flush(ctx, nil) - raw := store.get("c1-m1-logs") + raw := store.Get(ctx, "c1-m1-logs") if raw == "" { t.Fatalf("expected debug log written to c1-m1-logs") } diff --git a/internal/server/variable.go b/internal/server/variable.go index 953f488eaa..28a89690da 100644 --- a/internal/server/variable.go +++ b/internal/server/variable.go @@ -35,9 +35,9 @@ type Variables struct { // VariableStore interface for persistent storage (e.g., Redis) type VariableStore interface { - Get(key string) (string, error) - Set(key string, value string, exp time.Duration) bool - SetNX(key string, value string, exp time.Duration) bool + Get(ctx context.Context, key string) (string, error) + Set(ctx context.Context, key string, value string, exp time.Duration) bool + SetNX(ctx context.Context, key string, value string, exp time.Duration) bool } var ( @@ -90,7 +90,7 @@ func InitVariables(store VariableStore) error { //} // GetSecretKey returns the current secret key -func GetSecretKey(store VariableStore) (string, error) { +func GetSecretKey(ctx context.Context, store VariableStore) (string, error) { if globalConfig.GetSecretKey() != "" { return globalConfig.GetSecretKey(), nil } @@ -100,7 +100,7 @@ func GetSecretKey(store VariableStore) (string, error) { return "", fmt.Errorf("failed to generate secret key: %w", err) } - secretKey, err := GetOrCreateKey(store, SecretKeyRedisKey, generatedKey) + secretKey, err := GetOrCreateKey(ctx, store, SecretKeyRedisKey, generatedKey) if err != nil { return "", fmt.Errorf("failed to get secret key: %w", err) } @@ -121,7 +121,7 @@ func GetSecretKey(store VariableStore) (string, error) { // - If key exists in store, returns the stored value // - If key doesn't exist, calls createFn to generate value, stores it, and returns it // - Uses SetNX to ensure atomic creation (only one caller succeeds when key doesn't exist) -func GetOrCreateKey(store VariableStore, key string, newValue string) (string, error) { +func GetOrCreateKey(ctx context.Context, store VariableStore, key string, newValue string) (string, error) { if store == nil { err := fmt.Errorf("store is nil") common.Warn("VariableStore is nil, cannot get or create key", zap.String("key", key)) @@ -129,7 +129,7 @@ func GetOrCreateKey(store VariableStore, key string, newValue string) (string, e } // Try to get existing value - value, err := store.Get(key) + value, err := store.Get(ctx, key) if err != nil { common.Warn("Failed to get key from store", zap.String("key", key), zap.Error(err)) return "", err @@ -145,13 +145,13 @@ func GetOrCreateKey(store VariableStore, key string, newValue string) (string, e common.Info("Generating new value for key", zap.String("key", key)) // Try to set with NX (only if not exists) - ensures atomicity - if store.SetNX(key, newValue, SecretKeyTTL) { + if store.SetNX(ctx, key, newValue, SecretKeyTTL) { common.Info("New value stored successfully", zap.String("key", key)) return newValue, nil } // Another process might have set it, try to get again - value, err = store.Get(key) + value, err = store.Get(ctx, key) if err != nil { common.Warn("Failed to get key after SetNX", zap.String("key", key), zap.Error(err)) return newValue, nil // Return our generated value as fallback @@ -168,7 +168,7 @@ func GetOrCreateKey(store VariableStore, key string, newValue string) (string, e // RefreshVariables refreshes all variables from storage // Call this when you want to reload variables from persistent storage -func RefreshVariables(store VariableStore) error { +func RefreshVariables(ctx context.Context, store VariableStore) error { if store == nil { return fmt.Errorf("store is nil") } @@ -181,7 +181,7 @@ func RefreshVariables(store VariableStore) error { } // Refresh SecretKey - secretKey, err := store.Get(SecretKeyRedisKey) + secretKey, err := store.Get(ctx, SecretKeyRedisKey) if err != nil { common.Warn("Failed to refresh secret key from store", zap.Error(err)) return err @@ -215,11 +215,11 @@ func (w *VariableWatcher) Start(interval time.Duration) { w.wg.Go(func() { ticker := time.NewTicker(interval) defer ticker.Stop() - + ctx := context.Background() for { select { case <-ticker.C: - if err := RefreshVariables(w.store); err != nil { + if err := RefreshVariables(ctx, w.store); err != nil { common.Debug("Failed to refresh variables", zap.Error(err)) } case <-w.stopChan: diff --git a/internal/service/bot.go b/internal/service/bot.go index 8213353263..edefe521c5 100644 --- a/internal/service/bot.go +++ b/internal/service/bot.go @@ -187,7 +187,7 @@ func (s *BotService) AgentbotLogs(ctx context.Context, tenantID, agentID, messag if _, err := s.loadCanvas(ctx, tenantID, agentID); err != nil { return nil, common.CodeDataError, err } - payload, err := redis.Get().Get(fmt.Sprintf("%s-%s-logs", agentID, messageID)) + payload, err := redis.Get().Get(ctx, fmt.Sprintf("%s-%s-logs", agentID, messageID)) if err != nil { return nil, common.CodeServerError, errors.New("failed to read agent logs") } @@ -251,7 +251,7 @@ func (s *BotService) persistLock(sessionID string) *sync.Mutex { // loadCanvas is the IDOR guard for agentbot reads. It mirrors the // private loadCanvasForUser helper on AgentService without taking a -// dependency on the agentService pointer (so BotService can be unit- +// dependency on the agentService pointer (so BotService can be // tested with a nil agentService). func (s *BotService) loadCanvas(ctx context.Context, tenantID, agentID string) (*entity.UserCanvas, error) { if agentID == "" { diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 5be3b3195f..4f1e1c9817 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -223,7 +223,7 @@ func (s *ChunkService) cancelAllTasksOfDoc(ctx context.Context, docID string) er if task == nil { continue } - redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 0) + redisClient.Set(ctx, fmt.Sprintf("%s-cancel", task.ID), "x", 0) } return nil diff --git a/internal/service/connector.go b/internal/service/connector.go index 6dd611d3af..210848938f 100644 --- a/internal/service/connector.go +++ b/internal/service/connector.go @@ -455,7 +455,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(ctx context.Context, userID, sour CodeVerifier: codeVerifier, CreatedAt: time.Now().Unix(), } - if ok := redisClient.SetObj(webStateCacheKey(flowID, source), state, webFlowTTL); !ok { + if ok := redisClient.SetObj(ctx, webStateCacheKey(flowID, source), state, webFlowTTL); !ok { return nil, common.CodeServerError, fmt.Errorf("failed to initialize Google OAuth flow. Please verify the uploaded client configuration") } @@ -484,17 +484,17 @@ func (s *ConnectorService) GoogleWebOAuthCallback(ctx context.Context, source, s stateKey := webStateCacheKey(stateID, source) var state googleWebOAuthState - if ok := redisClient.GetObj(stateKey, &state); !ok { + if ok := redisClient.GetObj(ctx, stateKey, &state); !ok { return renderWebOAuthPopup(stateID, false, "Authorization session expired. Please restart from the main window.", source) } if state.ClientConfig == nil { - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(stateID, false, "Authorization session was invalid. Please retry.", source) } if strings.TrimSpace(oauthError) != "" { - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) message := strings.TrimSpace(errorDescription) if message == "" { message = strings.TrimSpace(oauthError) @@ -512,7 +512,7 @@ func (s *ConnectorService) GoogleWebOAuthCallback(ctx context.Context, source, s credentials, err := exchangeGoogleWebOAuthCode(state.ClientConfig, googleOAuthScopesForSource(source), state.RedirectURI, code, state.CodeVerifier) if err != nil { - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(stateID, false, "Failed to exchange tokens with Google. Please retry.", source) } @@ -520,11 +520,11 @@ func (s *ConnectorService) GoogleWebOAuthCallback(ctx context.Context, source, s UserID: state.UserID, Credentials: credentials, } - if ok := redisClient.SetObj(webResultCacheKey(stateID, source), result, webFlowTTL); !ok { - redisClient.Delete(stateKey) + if ok := redisClient.SetObj(ctx, webResultCacheKey(stateID, source), result, webFlowTTL); !ok { + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(stateID, false, "Failed to exchange tokens with Google. Please retry.", source) } - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(stateID, true, "Authorization completed successfully.", source) } @@ -545,7 +545,7 @@ func (s *ConnectorService) PollGoogleWebOAuthResult(ctx context.Context, userID, resultKey := webResultCacheKey(strings.TrimSpace(req.FlowID), source) var result googleWebOAuthResult - if ok := redisClient.GetObj(resultKey, &result); !ok { + if ok := redisClient.GetObj(ctx, resultKey, &result); !ok { return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } @@ -553,7 +553,7 @@ func (s *ConnectorService) PollGoogleWebOAuthResult(ctx context.Context, userID, return nil, common.CodePermissionError, fmt.Errorf("you are not allowed to access this authorization result") } - redisClient.Delete(resultKey) + redisClient.Delete(ctx, resultKey) return &PollGoogleWebOAuthResultResponse{Credentials: result.Credentials}, common.CodeSuccess, nil } @@ -1079,7 +1079,7 @@ func (s *ConnectorService) StartBoxWebOAuth(ctx context.Context, userID string, RedirectURI: redirectURI, CreatedAt: time.Now().Unix(), } - if ok := redisClient.SetObj(webStateCacheKey(flowID, "box"), state, webFlowTTL); !ok { + if ok := redisClient.SetObj(ctx, webStateCacheKey(flowID, "box"), state, webFlowTTL); !ok { return nil, common.CodeServerError, fmt.Errorf("failed to initialize Box OAuth flow. Please verify the client configuration") } @@ -1103,12 +1103,12 @@ func (s *ConnectorService) BoxWebOAuthCallback(ctx context.Context, flowID strin stateKey := webStateCacheKey(flowID, "box") var state boxWebOAuthState - if ok := redisClient.GetObj(stateKey, &state); !ok { + if ok := redisClient.GetObj(ctx, stateKey, &state); !ok { return renderWebOAuthPopup(flowID, false, "Box OAuth session expired or invalid.", "box") } if strings.TrimSpace(oauthError) != "" { - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) message := strings.TrimSpace(errorDescription) if message == "" { message = strings.TrimSpace(oauthError) @@ -1126,7 +1126,7 @@ func (s *ConnectorService) BoxWebOAuthCallback(ctx context.Context, flowID strin token, err := exchangeBoxAuthorizationCode(state.ClientID, state.ClientSecret, state.RedirectURI, code) if err != nil { - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(flowID, false, "Failed to exchange tokens with Box. Please retry.", "box") } @@ -1137,11 +1137,11 @@ func (s *ConnectorService) BoxWebOAuthCallback(ctx context.Context, flowID strin AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, } - if ok := redisClient.SetObj(webResultCacheKey(flowID, "box"), result, webFlowTTL); !ok { - redisClient.Delete(stateKey) + if ok := redisClient.SetObj(ctx, webResultCacheKey(flowID, "box"), result, webFlowTTL); !ok { + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(flowID, false, "Failed to exchange tokens with Box. Please retry.", "box") } - redisClient.Delete(stateKey) + redisClient.Delete(ctx, stateKey) return renderWebOAuthPopup(flowID, true, "Authorization completed successfully.", "box") } @@ -1158,7 +1158,7 @@ func (s *ConnectorService) PollBoxWebOAuthResult(ctx context.Context, userID str resultKey := webResultCacheKey(strings.TrimSpace(req.FlowID), "box") var result boxWebOAuthCredentials - if ok := redisClient.GetObj(resultKey, &result); !ok { + if ok := redisClient.GetObj(ctx, resultKey, &result); !ok { return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } @@ -1166,7 +1166,7 @@ func (s *ConnectorService) PollBoxWebOAuthResult(ctx context.Context, userID str return nil, common.CodePermissionError, fmt.Errorf("you are not allowed to access this authorization result") } - redisClient.Delete(resultKey) + redisClient.Delete(ctx, resultKey) result.UserID = "" return &PollBoxWebOAuthResultResponse{Credentials: result}, common.CodeSuccess, nil } diff --git a/internal/service/dataset/index.go b/internal/service/dataset/index.go index 2c4d584be5..5909e05bd8 100644 --- a/internal/service/dataset/index.go +++ b/internal/service/dataset/index.go @@ -134,9 +134,9 @@ func createDatasetIndexTaskInTx(tx *gorm.DB, task *entity.Task, queueDocID strin return &document, nil } -func enqueueDatasetIndexTask(priority int, queueMessage map[string]interface{}) error { +func enqueueDatasetIndexTask(ctx context.Context, priority int, queueMessage map[string]interface{}) error { redisClient := redisengine.Get() - if redisClient == nil || !redisClient.QueueProduct(datasetIndexQueueName(priority), queueMessage) { + if redisClient == nil || !redisClient.QueueProduct(ctx, datasetIndexQueueName(priority), queueMessage) { return errors.New("can't access Redis. Please check the Redis' status") } return nil @@ -251,12 +251,12 @@ func datasetIndexQueueName(priority int) string { return fmt.Sprintf("%s.%d.common", serverQueueNamePrefix, priority) } -func clearGraphPhaseMarkers(redisClient *redisengine.Client, datasetID string) { +func clearGraphPhaseMarkers(ctx context.Context, redisClient *redisengine.Client, datasetID string) { if redisClient == nil || datasetID == "" { return } for _, phase := range []string{graphPhaseResolutionDone, graphPhaseCommunityDone} { - if !redisClient.Delete(fmt.Sprintf("graphrag:phase:%s:%s", datasetID, phase)) { + if !redisClient.Delete(ctx, fmt.Sprintf("graphrag:phase:%s:%s", datasetID, phase)) { common.Warn("Failed to clear GraphRAG phase marker", zap.String("dataset_id", datasetID), zap.String("phase", phase)) } } @@ -343,7 +343,7 @@ func (d *DatasetService) RunIndex(ctx context.Context, userID, datasetID, indexT return nil, common.CodeDataError, errors.New("internal server error") } - if err = enqueueDatasetIndexTask(0, queueMessage); err != nil { + if err = enqueueDatasetIndexTask(ctx, 0, queueMessage); err != nil { if cleanupErr := cleanupFailedDatasetIndexTask(task.ID, updatedDocument, kb.ID, indexType); cleanupErr != nil { err = errors.Join(err, cleanupErr) } @@ -690,7 +690,7 @@ func (d *DatasetService) DeleteIndex(ctx context.Context, userID, datasetID, ind if taskID != "" { redisClient := redisengine.Get() - if redisClient == nil || !redisClient.Set(fmt.Sprintf("%s-cancel", taskID), "x", 0) { + if redisClient == nil || !redisClient.Set(ctx, fmt.Sprintf("%s-cancel", taskID), "x", 0) { common.Warn("Failed to set dataset index cancellation marker", zap.String("dataset_id", datasetID), zap.String("task_id", taskID)) } if err := dao.DB.Unscoped().Where("id = ?", taskID).Delete(&entity.Task{}).Error; err != nil { @@ -712,7 +712,7 @@ func (d *DatasetService) DeleteIndex(ctx context.Context, userID, datasetID, ind common.Warn("Failed to delete GraphRAG artefacts", zap.String("dataset_id", datasetID), zap.Error(err)) return common.CodeDataError, errors.New("internal server error") } - clearGraphPhaseMarkers(redisengine.Get(), datasetID) + clearGraphPhaseMarkers(ctx, redisengine.Get(), datasetID) common.Info("delete_index: cleared GraphRAG artefacts and phase markers", zap.String("dataset_id", datasetID)) } else if wipe && indexType == "raptor" { if d.docEngine == nil { diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 950b3310f6..1f60bde712 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -254,7 +254,7 @@ func (s *IngestionTaskService) RequestStop(ctx context.Context, taskID string) ( // running worker's pollCancel detects the stop immediately rather // than waiting for the next DB poll (up to 3s). if rc := redis2.Get(); rc != nil { - rc.Set(fmt.Sprintf("%s-cancel", taskID), "x", 1*time.Hour) + rc.Set(ctx, fmt.Sprintf("%s-cancel", taskID), "x", 1*time.Hour) } return task, nil default: diff --git a/internal/service/memory_extractor.go b/internal/service/memory_extractor.go index 3e07192729..ed8489c393 100644 --- a/internal/service/memory_extractor.go +++ b/internal/service/memory_extractor.go @@ -83,7 +83,7 @@ func (s *MemoryMessageService) StartTaskConsumer(ctx context.Context) { if ctx.Err() != nil { return } - msg, err := redisClient.QueueConsumer(queueName, memoryTaskConsumerGroup, consumerName, ">") + msg, err := redisClient.QueueConsumer(ctx, queueName, memoryTaskConsumerGroup, consumerName, ">") if err != nil { common.Error("memory task consumer: consume error", err) select { @@ -99,13 +99,13 @@ func (s *MemoryMessageService) StartTaskConsumer(ctx context.Context) { payload := msg.GetMessage() if taskType, _ := payload["task_type"].(string); taskType != "memory" { common.Warn(fmt.Sprintf("memory task consumer: skip task_type %q", taskType)) - msg.Ack() + msg.Ack(ctx) continue } if err := s.HandleSaveToMemoryTask(ctx, payload); err != nil { common.Error("memory task consumer: handle task failed", err) } - msg.Ack() + msg.Ack(ctx) } } @@ -174,7 +174,7 @@ func (s *MemoryMessageService) saveExtractedToMemory(ctx context.Context, memory now := time.Now().UTC() messages := make([]map[string]any, 0, len(extracted)) for _, item := range extracted { - messages = append(messages, buildExtractedMessage(generateRawMessageID(), sourceID, memoryID, msg, item, now)) + messages = append(messages, buildExtractedMessage(generateRawMessageID(ctx), sourceID, memoryID, msg, item, now)) } if err := s.embedAndSaveMessages(ctx, mem, messages); err != nil { return err diff --git a/internal/service/memory_message_service.go b/internal/service/memory_message_service.go index 24799b8c03..707aa7b73e 100644 --- a/internal/service/memory_message_service.go +++ b/internal/service/memory_message_service.go @@ -148,7 +148,7 @@ func (s *MemoryMessageService) QueueSaveToMemoryTask(ctx context.Context, memory // keeps the same field set as Python:344-386 so the // downstream extractor can consume the row without // schema changes. - rawMessageID := generateRawMessageID() + rawMessageID := generateRawMessageID(ctx) rawMessage := buildRawMessage(rawMessageID, memoryID, msg) if err := s.embedAndSave(ctx, mem, rawMessage); err != nil { @@ -167,7 +167,7 @@ func (s *MemoryMessageService) QueueSaveToMemoryTask(ctx context.Context, memory }) continue } - if err := queueMemoryTask(memoryID, mem.TenantID, rawMessageID, task, msg); err != nil { + if err = queueMemoryTask(ctx, memoryID, mem.TenantID, rawMessageID, task, msg); err != nil { res.Failed = append(res.Failed, MemoryFailure{ MemoryID: memoryID, FailMsg: err.Error(), @@ -179,9 +179,9 @@ func (s *MemoryMessageService) QueueSaveToMemoryTask(ctx context.Context, memory // generateRawMessageID returns the Redis auto-increment id used by the Python // side (`REDIS_CONN.generate_auto_increment_id(namespace="memory")`). -func generateRawMessageID() int64 { +func generateRawMessageID(ctx context.Context) int64 { if redisClient := redisengine.Get(); redisClient != nil { - if id := redisClient.GenerateAutoIncrementID("id_generator", "memory", 1, nil); id > 0 { + if id := redisClient.GenerateAutoIncrementID(ctx, "id_generator", "memory", 1, nil); id > 0 { return id } } @@ -335,7 +335,7 @@ func taskFromRow(row map[string]any) *entity.Task { } } -func queueMemoryTask(memoryID, tenantID string, rawMessageID int64, task map[string]any, msg MemoryMessage) error { +func queueMemoryTask(ctx context.Context, memoryID, tenantID string, rawMessageID int64, task map[string]any, msg MemoryMessage) error { taskID := fmt.Sprint(task["id"]) message := map[string]any{ "id": taskID, @@ -352,7 +352,7 @@ func queueMemoryTask(memoryID, tenantID string, rawMessageID int64, task map[str "agent_response": msg.AgentResponse, }, } - if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(memoryTaskQueueName(0), message) { + if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(ctx, memoryTaskQueueName(0), message) { return errors.New("Can't access Redis.") } return nil diff --git a/internal/service/memory_message_service_test.go b/internal/service/memory_message_service_test.go index 832062a635..011ba0c377 100644 --- a/internal/service/memory_message_service_test.go +++ b/internal/service/memory_message_service_test.go @@ -150,8 +150,9 @@ func TestTaskFromRow_InitializesProgressMessage(t *testing.T) { // values. (Wall-clock based today; the Redis-backed counter will // be added when the project's Redis client lands.) func TestGenerateRawMessageID_Unique(t *testing.T) { - a := generateRawMessageID() - b := generateRawMessageID() + ctx := t.Context() + a := generateRawMessageID(ctx) + b := generateRawMessageID(ctx) if a == b { // Allow a 1-second tie when the clock hasn't ticked. // GenerateRawMessageID uses Unix seconds; two calls diff --git a/internal/service/system.go b/internal/service/system.go index 64112a0fb0..ee76a25e4a 100644 --- a/internal/service/system.go +++ b/internal/service/system.go @@ -102,13 +102,13 @@ type StatusResponse struct { } // GetStatus gets health status for core system dependencies. -func (s *SystemService) GetStatus() (*StatusResponse, error) { +func (s *SystemService) GetStatus(ctx context.Context) (*StatusResponse, error) { return &StatusResponse{ DocEngine: s.getDocEngineStatus(), Storage: s.getStorageStatus(), Database: s.getDatabaseStatus(), - Redis: s.getRedisStatus(), - TaskExecutorHeartbeats: s.getTaskExecutorHeartbeats(), + Redis: s.getRedisStatus(ctx), + TaskExecutorHeartbeats: s.getTaskExecutorHeartbeats(ctx), }, nil } @@ -228,7 +228,7 @@ func (s *SystemService) getDatabaseStatus() ComponentStatus { } } -func (s *SystemService) getRedisStatus() ComponentStatus { +func (s *SystemService) getRedisStatus(ctx context.Context) ComponentStatus { startedAt := time.Now() redisClient := redis.Get() if redisClient == nil { @@ -238,7 +238,7 @@ func (s *SystemService) getRedisStatus() ComponentStatus { "error": "redis not initialized", } } - if !redisClient.Health() { + if !redisClient.Health(ctx) { return ComponentStatus{ "status": "red", "elapsed": elapsedMilliseconds(startedAt), @@ -252,21 +252,21 @@ func (s *SystemService) getRedisStatus() ComponentStatus { } } -func (s *SystemService) getTaskExecutorHeartbeats() map[string][]interface{} { +func (s *SystemService) getTaskExecutorHeartbeats(ctx context.Context) map[string][]interface{} { heartbeatsByExecutor := map[string][]interface{}{} redisClient := redis.Get() if redisClient == nil { return heartbeatsByExecutor } - taskExecutorIDs, err := redisClient.SMembers("TASKEXE") + taskExecutorIDs, err := redisClient.SMembers(ctx, "TASKEXE") if err != nil { return heartbeatsByExecutor } now := float64(time.Now().Unix()) for _, taskExecutorID := range taskExecutorIDs { - rawHeartbeats, err := redisClient.ZRangeByScore(taskExecutorID, now-60*30, now) + rawHeartbeats, err := redisClient.ZRangeByScore(ctx, taskExecutorID, now-60*30, now) if err != nil { continue } @@ -329,7 +329,7 @@ func GetComponentsHealthz(ctx context.Context) (*HealthzResponse, bool) { redisOK, redisMeta := timedHealthCheck(func() error { redisClient := redis.Get() - if redisClient == nil || !redisClient.Health() { + if redisClient == nil || !redisClient.Health(ctx) { return fmt.Errorf("redis is not healthy") } return nil diff --git a/internal/service/tag.go b/internal/service/tag.go index 16d7d7029e..444f1c0eb6 100644 --- a/internal/service/tag.go +++ b/internal/service/tag.go @@ -56,7 +56,7 @@ func getTagsCacheKey(kbIDs []string) string { // GetTagsFromCache retrieves cached tags for given kb_ids // Returns nil if not found (cache miss) -func GetTagsFromCache(kbIDs []string) (map[string]float64, error) { +func GetTagsFromCache(ctx context.Context, kbIDs []string) (map[string]float64, error) { if len(kbIDs) == 0 { return nil, nil } @@ -68,7 +68,7 @@ func GetTagsFromCache(kbIDs []string) (map[string]float64, error) { } key := getTagsCacheKey(kbIDs) - data, err := redisClient.Get(key) + data, err := redisClient.Get(ctx, key) if err != nil || data == "" { // Cache miss or error return nil, nil @@ -84,7 +84,7 @@ func GetTagsFromCache(kbIDs []string) (map[string]float64, error) { } // SetTagsToCache stores tags in cache for given kb_ids with 10 minute expiry -func SetTagsToCache(kbIDs []string, tags map[string]float64) error { +func SetTagsToCache(ctx context.Context, kbIDs []string, tags map[string]float64) error { if len(kbIDs) == 0 || tags == nil { return nil } @@ -102,7 +102,7 @@ func SetTagsToCache(kbIDs []string, tags map[string]float64) error { } // Cache for 10 minutes (600 seconds) - ok := redisClient.Set(key, string(data), 10*time.Minute) + ok := redisClient.Set(ctx, key, string(data), 10*time.Minute) if !ok { common.Warn("Failed to set tags cache") return fmt.Errorf("failed to set tags cache") @@ -290,7 +290,7 @@ func (s *MetadataService) LabelQuestion(ctx context.Context, question string, kb common.Debug("tag_kb_ids found in parser_config", zap.Strings("tag_kb_ids", tagKBIDs)) // Get all tags from cache or compute and cache - allTags, err := GetTagsFromCache(tagKBIDs) + allTags, err := GetTagsFromCache(ctx, tagKBIDs) if err != nil { common.Warn("Failed to get tags from cache", zap.Error(err)) } @@ -302,7 +302,7 @@ func (s *MetadataService) LabelQuestion(ctx context.Context, question string, kb return nil } // Store in cache for future lookups - if err = SetTagsToCache(tagKBIDs, allTags); err != nil { + if err = SetTagsToCache(ctx, tagKBIDs, allTags); err != nil { common.Warn("Failed to set tags cache", zap.Error(err)) } } diff --git a/internal/service/user.go b/internal/service/user.go index 58cbf8edfd..2a8fd8b4ec 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -595,7 +595,7 @@ func defaultUserLanguage() string { // using itsdangerous URLSafeTimedSerializer to get the actual access_token func (s *UserService) GetUserByToken(ctx context.Context, authorization string) (*entity.User, common.ErrorCode, error) { // Get secret key from config - secretKey, err := server.GetSecretKey(redis.Get()) + secretKey, err := server.GetSecretKey(ctx, redis.Get()) if err != nil { return nil, common.CodeUnauthorized, err } @@ -1081,7 +1081,7 @@ func (s *UserService) ForgotIssueCaptcha(ctx context.Context, email string) (cap return "", "", common.CodeServerError, err } captchaID = utility.GenerateToken() - if ok := redis.Get().Set(utility.CaptchaIDRedisKey(captchaID), text, 60*time.Second); !ok { + if ok := redis.Get().Set(ctx, utility.CaptchaIDRedisKey(captchaID), text, 60*time.Second); !ok { return "", "", common.CodeServerError, fmt.Errorf("failed to store captcha") } imageDataURL = utility.RenderCaptchaPNGDataURL(text) @@ -1103,7 +1103,7 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc rc := redis.Get() captchaKey := utility.CaptchaIDRedisKey(captchaID) - stored, _ := rc.Get(captchaKey) + stored, _ := rc.Get(ctx, captchaKey) if stored == "" { return common.CodeNotEffective, fmt.Errorf("invalid or expired captcha") } @@ -1112,7 +1112,7 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc } // One-shot: consume the captcha so a leaked captcha_id cannot be // reused for a stream of OTP requests. - rc.Delete(captchaKey) + rc.Delete(ctx, captchaKey) codeKey, attemptsKey, lastSentKey, lockKey := utility.OTPRedisKeys(email) @@ -1120,12 +1120,12 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc // let a request for a new OTP wipe the lock (deliberate divergence // from the Python implementation, which deletes the lock here and so // allows a locked attacker to clear their own lockout by re-requesting). - if locked, _ := rc.Get(lockKey); locked != "" { + if locked, _ := rc.Get(ctx, lockKey); locked != "" { return common.CodeNotEffective, fmt.Errorf("too many attempts, try later") } // Resend cooldown — refuse if we already sent within the window. - if lastSent, _ := rc.Get(lastSentKey); lastSent != "" { + if lastSent, _ := rc.Get(ctx, lastSentKey); lastSent != "" { ts, parseErr := strconv.ParseInt(lastSent, 10, 64) if parseErr == nil { elapsed := time.Since(time.Unix(ts, 0)) @@ -1150,15 +1150,15 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc // Snapshot the previous OTP-flow state so we can restore it if email // delivery fails — otherwise the user is throttled by lastSentKey // even though they never received the code. - prevCode, _ := rc.Get(codeKey) - prevAttempts, _ := rc.Get(attemptsKey) - prevLastSent, _ := rc.Get(lastSentKey) + prevCode, _ := rc.Get(ctx, codeKey) + prevAttempts, _ := rc.Get(ctx, attemptsKey) + prevLastSent, _ := rc.Get(ctx, lastSentKey) - if !rc.Set(codeKey, utility.EncodeOTPStorageValue(codeHash, salt), utility.OTPTTL) { + if !rc.Set(ctx, codeKey, utility.EncodeOTPStorageValue(codeHash, salt), utility.OTPTTL) { return common.CodeServerError, fmt.Errorf("failed to store otp") } - rc.Set(attemptsKey, "0", utility.OTPTTL) - rc.Set(lastSentKey, now, utility.OTPTTL) + rc.Set(ctx, attemptsKey, "0", utility.OTPTTL) + rc.Set(ctx, lastSentKey, now, utility.OTPTTL) // Note: lockKey is intentionally not cleared here. If the user has // been locked out by a previous verify burst, requesting a new OTP // does not lift the lock — we already refused above. @@ -1170,19 +1170,19 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc // keys we just wrote so the next attempt isn't blocked by the // resend cooldown a failed send just installed. if prevCode != "" { - rc.Set(codeKey, prevCode, utility.OTPTTL) + rc.Set(ctx, codeKey, prevCode, utility.OTPTTL) } else { - rc.Delete(codeKey) + rc.Delete(ctx, codeKey) } if prevAttempts != "" { - rc.Set(attemptsKey, prevAttempts, utility.OTPTTL) + rc.Set(ctx, attemptsKey, prevAttempts, utility.OTPTTL) } else { - rc.Delete(attemptsKey) + rc.Delete(ctx, attemptsKey) } if prevLastSent != "" { - rc.Set(lastSentKey, prevLastSent, utility.OTPTTL) + rc.Set(ctx, lastSentKey, prevLastSent, utility.OTPTTL) } else { - rc.Delete(lastSentKey) + rc.Delete(ctx, lastSentKey) } return common.CodeServerError, fmt.Errorf("failed to send email") } @@ -1203,11 +1203,11 @@ func (s *UserService) ForgotVerifyOTP(ctx context.Context, email, otp string) (c rc := redis.Get() codeKey, attemptsKey, lastSentKey, lockKey := utility.OTPRedisKeys(email) - if locked, _ := rc.Get(lockKey); locked != "" { + if locked, _ := rc.Get(ctx, lockKey); locked != "" { return common.CodeNotEffective, fmt.Errorf("too many attempts, try later") } - stored, _ := rc.Get(codeKey) + stored, _ := rc.Get(ctx, codeKey) if stored == "" { return common.CodeNotEffective, fmt.Errorf("expired otp") } @@ -1219,25 +1219,25 @@ func (s *UserService) ForgotVerifyOTP(ctx context.Context, email, otp string) (c if utility.HashOTPCode(strings.ToUpper(strings.TrimSpace(otp)), salt) != storedHash { // bump attempts; lock on >= limit attempts := 0 - if cur, _ := rc.Get(attemptsKey); cur != "" { + if cur, _ := rc.Get(ctx, attemptsKey); cur != "" { if n, perr := strconv.Atoi(cur); perr == nil { attempts = n } } attempts++ - rc.Set(attemptsKey, strconv.Itoa(attempts), utility.OTPTTL) + rc.Set(ctx, attemptsKey, strconv.Itoa(attempts), utility.OTPTTL) if attempts >= utility.OTPAttemptLimit { - rc.Set(lockKey, strconv.FormatInt(time.Now().Unix(), 10), utility.OTPAttemptLockDuration) + rc.Set(ctx, lockKey, strconv.FormatInt(time.Now().Unix(), 10), utility.OTPAttemptLockDuration) } return common.CodeAuthenticationError, fmt.Errorf("expired otp") } // Success: clear OTP state, mark email verified. - rc.Delete(codeKey) - rc.Delete(attemptsKey) - rc.Delete(lastSentKey) - rc.Delete(lockKey) - if !rc.Set(utility.OTPVerifiedRedisKey(email), "1", utility.OTPTTL) { + rc.Delete(ctx, codeKey) + rc.Delete(ctx, attemptsKey) + rc.Delete(ctx, lastSentKey) + rc.Delete(ctx, lockKey) + if !rc.Set(ctx, utility.OTPVerifiedRedisKey(email), "1", utility.OTPTTL) { return common.CodeServerError, fmt.Errorf("failed to set verification state") } return common.CodeSuccess, nil @@ -1272,7 +1272,7 @@ func (s *UserService) ForgotResetPassword(ctx context.Context, req *ForgotResetP rc := redis.Get() verifiedKey := utility.OTPVerifiedRedisKey(req.Email) - if v, _ := rc.Get(verifiedKey); v != "1" { + if v, _ := rc.Get(ctx, verifiedKey); v != "1" { return nil, common.CodeAuthenticationError, fmt.Errorf("email not verified") } @@ -1310,6 +1310,6 @@ func (s *UserService) ForgotResetPassword(ctx context.Context, req *ForgotResetP return nil, common.CodeServerError, fmt.Errorf("failed to reset password: %w", err) } - rc.Delete(verifiedKey) + rc.Delete(ctx, verifiedKey) return user, common.CodeSuccess, nil }