469 lines
14 KiB
Go
469 lines
14 KiB
Go
package ctxmgr
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"errors"
|
||
"fmt"
|
||
"log/slog"
|
||
"time"
|
||
)
|
||
|
||
var (
|
||
ErrIdempotencyKeyReused = errors.New("idempotency key reused with different body")
|
||
ErrChatTurnNotFound = errors.New("chat turn not found")
|
||
)
|
||
|
||
// ChatTurnStatus 回合状态
|
||
type ChatTurnStatus string
|
||
|
||
const (
|
||
TurnStatusAccepted ChatTurnStatus = "accepted"
|
||
TurnStatusProcessing ChatTurnStatus = "processing"
|
||
TurnStatusCompleted ChatTurnStatus = "completed"
|
||
TurnStatusFailed ChatTurnStatus = "failed"
|
||
TurnStatusCancelled ChatTurnStatus = "cancelled"
|
||
)
|
||
|
||
// ChatTurnRecord 回合记录
|
||
type ChatTurnRecord struct {
|
||
ID int64
|
||
SessionID string
|
||
ClientTurnID string
|
||
UserID int64
|
||
ProjectID int64
|
||
UserMessageID int64
|
||
AssistantMessageID int64
|
||
RequestBodySHA256 string
|
||
Status ChatTurnStatus
|
||
OwnerToken string
|
||
FirstRequestID string
|
||
LastRequestID string
|
||
AttemptCount int
|
||
CreatedAt time.Time
|
||
UpdatedAt time.Time
|
||
CompletedAt *time.Time
|
||
LeaseExpiresAt *time.Time
|
||
}
|
||
|
||
// AcquireTurnResult 取得执行权结果
|
||
type AcquireTurnResult struct {
|
||
Record *ChatTurnRecord
|
||
WonLease bool // true = 本请求赢得执行权，应跑 Agent
|
||
}
|
||
|
||
func HashTurnBody(sessionID, message string) string {
|
||
h := sha256.Sum256([]byte(sessionID + "\n" + message))
|
||
return hex.EncodeToString(h[:])
|
||
}
|
||
|
||
// HashTurnBodyWithKey 在消息哈希上并入附件引用，避免同文不同附件被当成同一回合。
|
||
func HashTurnBodyWithKey(sessionID, message, extra string) string {
|
||
if extra == "" {
|
||
return HashTurnBody(sessionID, message)
|
||
}
|
||
return HashTurnBody(sessionID, message+"\n"+extra)
|
||
}
|
||
|
||
// AcquireChatTurn 尝试取得回合执行权（INSERT ON CONFLICT）
|
||
func (m *ContextManager) AcquireChatTurn(
|
||
ctx context.Context,
|
||
sessionID, clientTurnID, bodyHash, requestID string,
|
||
userID, projectID int64,
|
||
) (*AcquireTurnResult, error) {
|
||
if sessionID == "" || clientTurnID == "" {
|
||
return nil, errors.New("session_id and client_turn_id required")
|
||
}
|
||
if bodyHash == "" {
|
||
bodyHash = HashTurnBody(sessionID, clientTurnID)
|
||
}
|
||
|
||
if m.pool != nil {
|
||
return m.acquireChatTurnDB(ctx, sessionID, clientTurnID, bodyHash, requestID, userID, projectID)
|
||
}
|
||
return m.acquireChatTurnMemory(sessionID, clientTurnID, bodyHash, requestID, userID, projectID)
|
||
}
|
||
|
||
func (m *ContextManager) acquireChatTurnDB(
|
||
ctx context.Context,
|
||
sessionID, clientTurnID, bodyHash, requestID string,
|
||
userID, projectID int64,
|
||
) (*AcquireTurnResult, error) {
|
||
ownerToken := fmt.Sprintf("own_%d", time.Now().UnixNano())
|
||
leaseUntil := time.Now().Add(5 * time.Minute)
|
||
|
||
var existingID int64
|
||
var existingHash string
|
||
var existingStatus string
|
||
var leaseExpiresAt *time.Time
|
||
err := m.pool.QueryRow(ctx, `
|
||
SELECT id, request_body_sha256, status, lease_expires_at
|
||
FROM chat_turns
|
||
WHERE session_id = $1 AND client_turn_id = $2
|
||
`, sessionID, clientTurnID).Scan(&existingID, &existingHash, &existingStatus, &leaseExpiresAt)
|
||
if err == nil {
|
||
if existingHash != "" && existingHash != bodyHash {
|
||
return nil, ErrIdempotencyKeyReused
|
||
}
|
||
if existingStatus == string(TurnStatusProcessing) && leaseExpiresAt != nil && time.Now().After(*leaseExpiresAt) {
|
||
_, _ = m.pool.Exec(ctx, `
|
||
UPDATE chat_turns SET status = $1, updated_at = now()
|
||
WHERE id = $2
|
||
`, TurnStatusFailed, existingID)
|
||
} else {
|
||
rec, recErr := m.GetChatTurn(ctx, sessionID, clientTurnID)
|
||
if recErr != nil {
|
||
return nil, recErr
|
||
}
|
||
if rec.Status != TurnStatusProcessing || (rec.LeaseExpiresAt != nil && time.Now().Before(*rec.LeaseExpiresAt)) {
|
||
_, _ = m.pool.Exec(ctx, `
|
||
UPDATE chat_turns
|
||
SET attempt_count = attempt_count + 1,
|
||
last_request_id = $1,
|
||
updated_at = now()
|
||
WHERE id = $2
|
||
`, requestID, existingID)
|
||
rec.AttemptCount++
|
||
rec.LastRequestID = requestID
|
||
return &AcquireTurnResult{Record: rec, WonLease: false}, nil
|
||
}
|
||
}
|
||
}
|
||
|
||
var turnID int64
|
||
err = m.pool.QueryRow(ctx, `
|
||
INSERT INTO chat_turns (
|
||
session_id, client_turn_id, user_id, project_id,
|
||
request_body_sha256, status, owner_token, lease_expires_at,
|
||
first_request_id, last_request_id, attempt_count
|
||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $9, 1)
|
||
ON CONFLICT (session_id, client_turn_id) DO NOTHING
|
||
RETURNING id
|
||
`, sessionID, clientTurnID, userID, projectID, bodyHash, TurnStatusProcessing, ownerToken, leaseUntil, requestID).Scan(&turnID)
|
||
if err != nil {
|
||
// conflict race: re-read
|
||
rec, recErr := m.GetChatTurn(ctx, sessionID, clientTurnID)
|
||
if recErr != nil {
|
||
return nil, recErr
|
||
}
|
||
if rec.RequestBodySHA256 != bodyHash {
|
||
return nil, ErrIdempotencyKeyReused
|
||
}
|
||
return &AcquireTurnResult{Record: rec, WonLease: false}, nil
|
||
}
|
||
|
||
leaseCopy := leaseUntil
|
||
rec := &ChatTurnRecord{
|
||
ID: turnID,
|
||
SessionID: sessionID,
|
||
ClientTurnID: clientTurnID,
|
||
UserID: userID,
|
||
ProjectID: projectID,
|
||
RequestBodySHA256: bodyHash,
|
||
Status: TurnStatusProcessing,
|
||
OwnerToken: ownerToken,
|
||
FirstRequestID: requestID,
|
||
LastRequestID: requestID,
|
||
AttemptCount: 1,
|
||
CreatedAt: time.Now(),
|
||
UpdatedAt: time.Now(),
|
||
LeaseExpiresAt: &leaseCopy,
|
||
}
|
||
return &AcquireTurnResult{Record: rec, WonLease: true}, nil
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
func (m *ContextManager) acquireChatTurnMemory(
|
||
sessionID, clientTurnID, bodyHash, requestID string,
|
||
userID, projectID int64,
|
||
) (*AcquireTurnResult, error) {
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
|
||
key := turnMemKey(sessionID, clientTurnID)
|
||
if existing, ok := m.turns[key]; ok {
|
||
if existing.RequestBodySHA256 != bodyHash {
|
||
return nil, ErrIdempotencyKeyReused
|
||
}
|
||
leaseExpired := existing.Status == TurnStatusProcessing &&
|
||
existing.LeaseExpiresAt != nil &&
|
||
time.Now().After(*existing.LeaseExpiresAt)
|
||
if leaseExpired {
|
||
existing.Status = TurnStatusFailed
|
||
existing.UpdatedAt = time.Now()
|
||
} else if existing.Status == TurnStatusProcessing {
|
||
existing.AttemptCount++
|
||
existing.LastRequestID = requestID
|
||
existing.UpdatedAt = time.Now()
|
||
return &AcquireTurnResult{Record: existing, WonLease: false}, nil
|
||
} else {
|
||
existing.AttemptCount++
|
||
existing.LastRequestID = requestID
|
||
existing.UpdatedAt = time.Now()
|
||
return &AcquireTurnResult{Record: existing, WonLease: false}, nil
|
||
}
|
||
}
|
||
|
||
now := time.Now()
|
||
leaseUntil := now.Add(5 * time.Minute)
|
||
if existing, ok := m.turns[key]; ok {
|
||
existing.Status = TurnStatusProcessing
|
||
existing.OwnerToken = fmt.Sprintf("mem_%d", now.UnixNano())
|
||
existing.LastRequestID = requestID
|
||
existing.AttemptCount++
|
||
existing.UpdatedAt = now
|
||
existing.LeaseExpiresAt = &leaseUntil
|
||
return &AcquireTurnResult{Record: existing, WonLease: true}, nil
|
||
}
|
||
|
||
if len(m.turns) > maxMemTurns {
|
||
m.pruneStaleTurnsLocked(memTurnTTL)
|
||
}
|
||
|
||
rec := &ChatTurnRecord{
|
||
ID: int64(len(m.turns) + 1),
|
||
SessionID: sessionID,
|
||
ClientTurnID: clientTurnID,
|
||
UserID: userID,
|
||
ProjectID: projectID,
|
||
RequestBodySHA256: bodyHash,
|
||
Status: TurnStatusProcessing,
|
||
OwnerToken: fmt.Sprintf("mem_%d", now.UnixNano()),
|
||
FirstRequestID: requestID,
|
||
LastRequestID: requestID,
|
||
AttemptCount: 1,
|
||
CreatedAt: now,
|
||
UpdatedAt: now,
|
||
LeaseExpiresAt: &leaseUntil,
|
||
}
|
||
m.turns[key] = rec
|
||
return &AcquireTurnResult{Record: rec, WonLease: true}, nil
|
||
}
|
||
|
||
// UpdateChatTurnUserMessage 绑定用户消息 ID
|
||
func (m *ContextManager) UpdateChatTurnUserMessage(ctx context.Context, sessionID, clientTurnID string, userMessageID int64) error {
|
||
if m.pool != nil {
|
||
_, err := m.pool.Exec(ctx, `
|
||
UPDATE chat_turns SET user_message_id = $1, updated_at = now()
|
||
WHERE session_id = $2 AND client_turn_id = $3
|
||
`, userMessageID, sessionID, clientTurnID)
|
||
return err
|
||
}
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
if rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]; ok {
|
||
rec.UserMessageID = userMessageID
|
||
rec.UpdatedAt = time.Now()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// CompleteChatTurn 标记回合完成
|
||
func (m *ContextManager) CompleteChatTurn(ctx context.Context, sessionID, clientTurnID string, assistantMessageID int64) error {
|
||
now := time.Now()
|
||
if m.pool != nil {
|
||
_, err := m.pool.Exec(ctx, `
|
||
UPDATE chat_turns
|
||
SET status = $1, assistant_message_id = $2, completed_at = $3, updated_at = $3
|
||
WHERE session_id = $4 AND client_turn_id = $5
|
||
`, TurnStatusCompleted, assistantMessageID, now, sessionID, clientTurnID)
|
||
if err != nil {
|
||
// T20-B5：补偿路径不再静默；结构化日志便于告警与计数。
|
||
slog.ErrorContext(ctx, "chat_turn_complete_failed",
|
||
"session_id", sessionID, "client_turn_id", clientTurnID,
|
||
"assistant_message_id", assistantMessageID, "error", err)
|
||
}
|
||
return err
|
||
}
|
||
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
|
||
rec.UpdatedAt = now
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// FailChatTurn 标记回合失败
|
||
func (m *ContextManager) FailChatTurn(ctx context.Context, sessionID, clientTurnID string) error {
|
||
if m.pool != nil {
|
||
_, err := m.pool.Exec(ctx, `
|
||
UPDATE chat_turns SET status = $1, updated_at = now()
|
||
WHERE session_id = $2 AND client_turn_id = $3
|
||
`, TurnStatusFailed, sessionID, clientTurnID)
|
||
if err != nil {
|
||
// T20-B5：失败补偿自身失败必须可见（否则孤儿 processing 无声累积）。
|
||
slog.ErrorContext(ctx, "chat_turn_fail_compensation_failed",
|
||
"session_id", sessionID, "client_turn_id", clientTurnID, "error", err)
|
||
}
|
||
return err
|
||
}
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
if rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]; ok {
|
||
rec.Status = TurnStatusFailed
|
||
rec.UpdatedAt = time.Now()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// GetChatTurn 查询回合
|
||
func (m *ContextManager) GetChatTurn(ctx context.Context, sessionID, clientTurnID string) (*ChatTurnRecord, error) {
|
||
if m.pool != nil {
|
||
var rec ChatTurnRecord
|
||
var completedAt *time.Time
|
||
var leaseExpiresAt *time.Time
|
||
err := m.pool.QueryRow(ctx, `
|
||
SELECT id, session_id, client_turn_id,
|
||
COALESCE(user_id, 0), COALESCE(project_id, 0),
|
||
COALESCE(user_message_id, 0), COALESCE(assistant_message_id, 0),
|
||
request_body_sha256, status,
|
||
COALESCE(first_request_id, ''), COALESCE(last_request_id, ''),
|
||
attempt_count, created_at, updated_at, completed_at, lease_expires_at
|
||
FROM chat_turns
|
||
WHERE session_id = $1 AND client_turn_id = $2
|
||
`, sessionID, clientTurnID).Scan(
|
||
&rec.ID, &rec.SessionID, &rec.ClientTurnID,
|
||
&rec.UserID, &rec.ProjectID,
|
||
&rec.UserMessageID, &rec.AssistantMessageID,
|
||
&rec.RequestBodySHA256, &rec.Status,
|
||
&rec.FirstRequestID, &rec.LastRequestID,
|
||
&rec.AttemptCount, &rec.CreatedAt, &rec.UpdatedAt, &completedAt, &leaseExpiresAt,
|
||
)
|
||
if err != nil {
|
||
return nil, ErrChatTurnNotFound
|
||
}
|
||
rec.CompletedAt = completedAt
|
||
rec.LeaseExpiresAt = leaseExpiresAt
|
||
return &rec, nil
|
||
}
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
rec, ok := m.turns[turnMemKey(sessionID, clientTurnID)]
|
||
if !ok {
|
||
return nil, ErrChatTurnNotFound
|
||
}
|
||
cp := *rec
|
||
return &cp, nil
|
||
}
|
||
|
||
// GetChatTurnByAssistantMessageID 按会话 + 助手消息定位回合；用于 copy_only 显式 pin。
|
||
func (m *ContextManager) GetChatTurnByAssistantMessageID(ctx context.Context, sessionID string, assistantMessageID int64) (*ChatTurnRecord, error) {
|
||
if sessionID == "" || assistantMessageID <= 0 {
|
||
return nil, ErrChatTurnNotFound
|
||
}
|
||
if m.pool != nil {
|
||
var rec ChatTurnRecord
|
||
var completedAt *time.Time
|
||
var leaseExpiresAt *time.Time
|
||
err := m.pool.QueryRow(ctx, `
|
||
SELECT id, session_id, client_turn_id,
|
||
COALESCE(user_id, 0), COALESCE(project_id, 0),
|
||
COALESCE(user_message_id, 0), COALESCE(assistant_message_id, 0),
|
||
request_body_sha256, status,
|
||
COALESCE(first_request_id, ''), COALESCE(last_request_id, ''),
|
||
attempt_count, created_at, updated_at, completed_at, lease_expires_at
|
||
FROM chat_turns
|
||
WHERE session_id = $1 AND assistant_message_id = $2
|
||
ORDER BY id DESC
|
||
LIMIT 1
|
||
`, sessionID, assistantMessageID).Scan(
|
||
&rec.ID, &rec.SessionID, &rec.ClientTurnID,
|
||
&rec.UserID, &rec.ProjectID,
|
||
&rec.UserMessageID, &rec.AssistantMessageID,
|
||
&rec.RequestBodySHA256, &rec.Status,
|
||
&rec.FirstRequestID, &rec.LastRequestID,
|
||
&rec.AttemptCount, &rec.CreatedAt, &rec.UpdatedAt, &completedAt, &leaseExpiresAt,
|
||
)
|
||
if err != nil {
|
||
return nil, ErrChatTurnNotFound
|
||
}
|
||
rec.CompletedAt = completedAt
|
||
rec.LeaseExpiresAt = leaseExpiresAt
|
||
return &rec, nil
|
||
}
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
for _, rec := range m.turns {
|
||
if rec != nil && rec.SessionID == sessionID && rec.AssistantMessageID == assistantMessageID {
|
||
cp := *rec
|
||
return &cp, nil
|
||
}
|
||
}
|
||
return nil, ErrChatTurnNotFound
|
||
}
|
||
|
||
// SweepExpiredProcessingTurns 批量清理孤儿回合（T20-B5）：
|
||
// status=processing 且 lease_expires_at < now 的回合标记为 failed。
|
||
// 返回清理条数；错误不打断扫描（单条失败仅记日志）。
|
||
// 幂等：重复扫描对已终态回合无影响。
|
||
func (m *ContextManager) SweepExpiredProcessingTurns(ctx context.Context, limit int) (int, error) {
|
||
if limit <= 0 {
|
||
limit = 100
|
||
}
|
||
if m.pool != nil {
|
||
tag, err := m.pool.Exec(ctx, `
|
||
UPDATE chat_turns SET status = $1, updated_at = now()
|
||
WHERE id IN (
|
||
SELECT id FROM chat_turns
|
||
WHERE status = $2 AND lease_expires_at < now()
|
||
ORDER BY lease_expires_at ASC
|
||
LIMIT $3
|
||
)
|
||
`, TurnStatusFailed, TurnStatusProcessing, limit)
|
||
if err != nil {
|
||
slog.ErrorContext(ctx, "chat_turn_sweep_failed", "error", err)
|
||
return 0, err
|
||
}
|
||
n := tag.RowsAffected()
|
||
if n > 0 {
|
||
slog.WarnContext(ctx, "chat_turn_sweep_expired_processing",
|
||
"marked_failed", n, "limit", limit)
|
||
}
|
||
return int(n), nil
|
||
}
|
||
// 内存模式：扫描过期 processing
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
now := time.Now()
|
||
swept := 0
|
||
for key, rec := range m.turns {
|
||
if rec.Status == TurnStatusProcessing &&
|
||
rec.LeaseExpiresAt != nil && rec.LeaseExpiresAt.Before(now) {
|
||
rec.Status = TurnStatusFailed
|
||
rec.UpdatedAt = now
|
||
swept++
|
||
slog.WarnContext(ctx, "chat_turn_sweep_expired_processing", "key", key)
|
||
}
|
||
if swept >= limit {
|
||
break
|
||
}
|
||
}
|
||
return swept, nil
|
||
}
|