mtfyun/ragflow
1
0 Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-11-10 21:32:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context to redis client (#17689)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai authored and GitHub committed 2026-08-02 22:50:54 +08:00
1 parent b2521ebf51
commit 266837eb33
31 files changed
+302 -308
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cmd/ragflow_server.go
+2 -2
View File
@@ -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.
internal/admin/handler.go
+4 -4
View File
@@ -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
internal/admin/service.go
+1 -1
View File
@@ -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, "")
}
internal/agent/audio/model_provider_synthesizer.go
+2 -2
View File
@@ -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
internal/deepdoc/parser/pdf/inference/cache.go
+14 -14
View File
@@ -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
}
internal/deepdoc/parser/pdf/inference/cache_test.go
+2 -2
View File
@@ -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 {
internal/engine/redis/redis.go
+45 -77
View File
@@ -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},
internal/engine/redis/redis_test.go
+3 -3
View File
@@ -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")
}
internal/handler/agent.go
+2 -2
View File
@@ -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)
internal/handler/agent_logs_test.go
+21 -19
View File
@@ -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()
internal/handler/agent_test.go
+28 -14
View File
@@ -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
internal/handler/agent_webhook.go
+9 -9
View File
@@ -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-<agent_id>-logs` with a 600 s TTL.
// Each event is recorded as {"ts": <float>, "event": <type>, ...}.
// 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)
}
internal/handler/system.go
+2 -1
View File
@@ -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
internal/handler/user.go
+4 -4
View File
@@ -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
internal/ingestion/component/extractor.go
+7 -7
View File
@@ -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 `</think>` tags and rejects `**ERROR**` responses,
internal/ingestion/component/extractor_tag.go
+6 -6
View File
@@ -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 {
internal/ingestion/service/ingestion_service.go
+3 -3
View File
@@ -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
}
}
internal/ingestion/task/debug_log_sink.go
+2 -2
View File
@@ -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
internal/ingestion/task/debug_log_sink_test.go
+44 -36
View File
@@ -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")
}
internal/server/variable.go
+13 -13
View File
@@ -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:
internal/service/bot.go
+2 -2
View File
@@ -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 == "" {
internal/service/chunk_types.go
+1 -1
View File
@@ -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
internal/service/connector.go
+19 -19
View File
@@ -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
}
internal/service/dataset/index.go
+7 -7
View File
@@ -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 {
internal/service/ingestion_task_service.go
+1 -1
View File
@@ -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:
internal/service/memory_extractor.go
+4 -4
View File
@@ -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
internal/service/memory_message_service.go
+6 -6
View File
@@ -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
internal/service/memory_message_service_test.go
+3 -2
View File
@@ -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
internal/service/system.go
+9 -9
View File
@@ -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
internal/service/tag.go
+6 -6
View File
@@ -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))
}
}
internal/service/user.go
+30 -30
View File
@@ -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
}