diff --git a/internal/ctxmgr/chat_turns.go b/internal/ctxmgr/chat_turns.go index 2d2c7c84..b31d2527 100644 --- a/internal/ctxmgr/chat_turns.go +++ b/internal/ctxmgr/chat_turns.go @@ -6,7 +6,6 @@ import ( "encoding/hex" "errors" "fmt" - "sync" "time" ) @@ -165,8 +164,25 @@ func (m *ContextManager) acquireChatTurnDB( return &AcquireTurnResult{Record: rec, WonLease: true}, nil } -var memTurnsMu sync.Mutex -var memTurns = map[string]*ChatTurnRecord{} // sessionID|clientTurnID -> record + +const maxMemTurns = 2000 +const memTurnTTL = 2 * time.Hour + +// pruneStaleTurnsLocked 调用方需持有 m.mu 写锁 +func (m *ContextManager) pruneStaleTurnsLocked(ttl time.Duration) int { + if ttl <= 0 { + ttl = memTurnTTL + } + now := time.Now() + cleaned := 0 + for k, rec := range m.turns { + if rec == nil || now.Sub(rec.UpdatedAt) > ttl { + delete(m.turns, k) + cleaned++ + } + } + return cleaned +} func turnMemKey(sessionID, clientTurnID string) string { return sessionID + "|" + clientTurnID @@ -176,11 +192,11 @@ func (m *ContextManager) acquireChatTurnMemory( sessionID, clientTurnID, bodyHash, requestID string, userID, projectID int64, ) (*AcquireTurnResult, error) { - memTurnsMu.Lock() - defer memTurnsMu.Unlock() + m.mu.Lock() + defer m.mu.Unlock() key := turnMemKey(sessionID, clientTurnID) - if existing, ok := memTurns[key]; ok { + if existing, ok := m.turns[key]; ok { if existing.RequestBodySHA256 != bodyHash { return nil, ErrIdempotencyKeyReused } @@ -205,7 +221,7 @@ func (m *ContextManager) acquireChatTurnMemory( now := time.Now() leaseUntil := now.Add(5 * time.Minute) - if existing, ok := memTurns[key]; ok { + if existing, ok := m.turns[key]; ok { existing.Status = TurnStatusProcessing existing.OwnerToken = fmt.Sprintf("mem_%d", now.UnixNano()) existing.LastRequestID = requestID @@ -215,8 +231,12 @@ func (m *ContextManager) acquireChatTurnMemory( return &AcquireTurnResult{Record: existing, WonLease: true}, nil } + if len(m.turns) > maxMemTurns { + m.pruneStaleTurnsLocked(memTurnTTL) + } + rec := &ChatTurnRecord{ - ID: int64(len(memTurns) + 1), + ID: int64(len(m.turns) + 1), SessionID: sessionID, ClientTurnID: clientTurnID, UserID: userID, @@ -231,7 +251,7 @@ func (m *ContextManager) acquireChatTurnMemory( UpdatedAt: now, LeaseExpiresAt: &leaseUntil, } - memTurns[key] = rec + m.turns[key] = rec return &AcquireTurnResult{Record: rec, WonLease: true}, nil } @@ -244,9 +264,9 @@ func (m *ContextManager) UpdateChatTurnUserMessage(ctx context.Context, sessionI `, userMessageID, sessionID, clientTurnID) return err } - memTurnsMu.Lock() - defer memTurnsMu.Unlock() - if rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)]; ok { + m.mu.Lock() + defer m.mu.Unlock() + if rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]; ok { rec.UserMessageID = userMessageID rec.UpdatedAt = time.Now() } @@ -264,9 +284,9 @@ func (m *ContextManager) CompleteChatTurn(ctx context.Context, sessionID, client `, TurnStatusCompleted, assistantMessageID, now, sessionID, clientTurnID) return err } - memTurnsMu.Lock() - defer memTurnsMu.Unlock() - if rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)]; ok { + m.mu.Lock() + defer m.mu.Unlock() + if rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]; ok { rec.Status = TurnStatusCompleted rec.AssistantMessageID = assistantMessageID rec.CompletedAt = &now @@ -284,9 +304,9 @@ func (m *ContextManager) FailChatTurn(ctx context.Context, sessionID, clientTurn `, TurnStatusFailed, sessionID, clientTurnID) return err } - memTurnsMu.Lock() - defer memTurnsMu.Unlock() - if rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)]; ok { + m.mu.Lock() + defer m.mu.Unlock() + if rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]; ok { rec.Status = TurnStatusFailed rec.UpdatedAt = time.Now() } @@ -323,9 +343,9 @@ func (m *ContextManager) GetChatTurn(ctx context.Context, sessionID, clientTurnI rec.LeaseExpiresAt = leaseExpiresAt return &rec, nil } - memTurnsMu.Lock() - defer memTurnsMu.Unlock() - rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)] + m.mu.Lock() + defer m.mu.Unlock() + rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)] if !ok { return nil, ErrChatTurnNotFound } diff --git a/internal/ctxmgr/chat_turns_test.go b/internal/ctxmgr/chat_turns_test.go index 354ad713..8e0dae68 100644 --- a/internal/ctxmgr/chat_turns_test.go +++ b/internal/ctxmgr/chat_turns_test.go @@ -2,12 +2,13 @@ package ctxmgr import ( "context" + "strconv" "sync" "testing" + "time" ) func TestAcquireChatTurn_DuplicateSameBody(t *testing.T) { - memTurns = map[string]*ChatTurnRecord{} mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig()) ctx := context.Background() sessionID := "sess-dup" @@ -32,7 +33,6 @@ func TestAcquireChatTurn_DuplicateSameBody(t *testing.T) { } func TestAcquireChatTurn_ConcurrentDuplicate(t *testing.T) { - memTurns = map[string]*ChatTurnRecord{} mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig()) ctx := context.Background() sessionID := "sess-conc" @@ -65,7 +65,6 @@ func TestAcquireChatTurn_ConcurrentDuplicate(t *testing.T) { } func TestAcquireChatTurn_DifferentBodyRejected(t *testing.T) { - memTurns = map[string]*ChatTurnRecord{} mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig()) ctx := context.Background() sessionID := "sess-reuse" @@ -80,3 +79,32 @@ func TestAcquireChatTurn_DifferentBodyRejected(t *testing.T) { t.Fatalf("expected ErrIdempotencyKeyReused, got %v", err) } } + +func TestAcquireChatTurn_StaleEntriesPruned(t *testing.T) { + mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig()) + ctx := context.Background() + mgr.mu.Lock() + now := time.Now() + for i := 0; i < maxMemTurns+50; i++ { + key := turnMemKey("sess-prune", "old-"+strconv.Itoa(i)) + mgr.turns[key] = &ChatTurnRecord{ + ID: int64(i + 1), + SessionID: "sess-prune", + ClientTurnID: "old-" + strconv.Itoa(i), + Status: TurnStatusCompleted, + UpdatedAt: now.Add(-3 * time.Hour), + } + } + mgr.mu.Unlock() + bodyHash := HashTurnBody("sess-prune", "fresh") + _, err := mgr.AcquireChatTurn(ctx, "sess-prune", "fresh-turn", bodyHash, "req-fresh", 1, 1) + if err != nil { + t.Fatalf("acquire: %v", err) + } + mgr.mu.RLock() + n := len(mgr.turns) + mgr.mu.RUnlock() + if n > maxMemTurns { + t.Fatalf("expected turns pruned under capacity, got %d", n) + } +} diff --git a/internal/ctxmgr/manager.go b/internal/ctxmgr/manager.go index 4e8d79b8..864ec7a3 100644 --- a/internal/ctxmgr/manager.go +++ b/internal/ctxmgr/manager.go @@ -33,6 +33,7 @@ type ContextManager struct { summaries map[string]*ChatSummary // session_id -> rolling summary digests map[string]map[string]*DailyDigest // session_id -> day_key -> DailyDigest readPositions map[string]*ChatReadPosition // session_id:user_id -> read position + turns map[string]*ChatTurnRecord // sessionID|clientTurnID -> 幂等回合(仅内存降级) cache *MultiLevelCache // 多级统一缓存系统 (C1-C5) semanticSummarizer SemanticSummarizer // no-tool summarizer; nil → heuristic default nextActivityID int64 @@ -60,6 +61,7 @@ func NewContextManager(pool *pgxpool.Pool, cas *storage.CAS, decayCfg TimeDecayC summaries: make(map[string]*ChatSummary), digests: make(map[string]map[string]*DailyDigest), readPositions: make(map[string]*ChatReadPosition), + turns: make(map[string]*ChatTurnRecord), cache: NewMultiLevelCache(), nextActivityID: 1, nextMessageID: 1,