feat(db): add chat_turns idempotency table
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
WeiCHEN
committed by
weichen
weichen
co-authored by
Cursor
Cursor
parent
62ab734021
commit
c887bc0480
@@ -121,6 +121,29 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_sid_category ON chat_messages(sessi
|
||||
CREATE EXTENSION IF NOT EXISTS pg_trgm;
|
||||
CREATE INDEX IF NOT EXISTS idx_chat_messages_content_trgm ON chat_messages USING GIN (content gin_trgm_ops);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS chat_turns (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
client_turn_id TEXT NOT NULL,
|
||||
user_id BIGINT,
|
||||
project_id BIGINT,
|
||||
user_message_id BIGINT REFERENCES chat_messages(id) ON DELETE SET NULL,
|
||||
assistant_message_id BIGINT REFERENCES chat_messages(id) ON DELETE SET NULL,
|
||||
request_body_sha256 TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'accepted',
|
||||
owner_token TEXT NOT NULL DEFAULT '',
|
||||
lease_expires_at TIMESTAMPTZ,
|
||||
first_request_id TEXT NOT NULL DEFAULT '',
|
||||
last_request_id TEXT NOT NULL DEFAULT '',
|
||||
attempt_count INT NOT NULL DEFAULT 1,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
completed_at TIMESTAMPTZ,
|
||||
CONSTRAINT uq_chat_turns_session_client UNIQUE (session_id, client_turn_id)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_chat_turns_session_status ON chat_turns(session_id, status);
|
||||
CREATE INDEX IF NOT EXISTS idx_chat_turns_client ON chat_turns(client_turn_id);
|
||||
|
||||
-- 4.1 每日智能摘要表 (Daily Digests)
|
||||
CREATE TABLE IF NOT EXISTS chat_daily_digests (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
-- chat_turns: 端到端回合幂等权威表
|
||||
CREATE TABLE IF NOT EXISTS chat_turns (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
client_turn_id TEXT NOT NULL,
|
||||
user_id BIGINT,
|
||||
project_id BIGINT,
|
||||
user_message_id BIGINT REFERENCES chat_messages(id) ON DELETE SET NULL,
|
||||
assistant_message_id BIGINT REFERENCES chat_messages(id) ON DELETE SET NULL,
|
||||
request_body_sha256 TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'accepted',
|
||||
owner_token TEXT NOT NULL DEFAULT '',
|
||||
lease_expires_at TIMESTAMPTZ,
|
||||
first_request_id TEXT NOT NULL DEFAULT '',
|
||||
last_request_id TEXT NOT NULL DEFAULT '',
|
||||
attempt_count INT NOT NULL DEFAULT 1,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
completed_at TIMESTAMPTZ,
|
||||
CONSTRAINT uq_chat_turns_session_client UNIQUE (session_id, client_turn_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_chat_turns_session_status ON chat_turns(session_id, status);
|
||||
CREATE INDEX IF NOT EXISTS idx_chat_turns_client ON chat_turns(client_turn_id);
|
||||
@@ -0,0 +1,294 @@
|
||||
package ctxmgr
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"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
|
||||
}
|
||||
|
||||
// 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[:])
|
||||
}
|
||||
|
||||
// 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
|
||||
err := m.pool.QueryRow(ctx, `
|
||||
SELECT id, request_body_sha256, status
|
||||
FROM chat_turns
|
||||
WHERE session_id = $1 AND client_turn_id = $2
|
||||
`, sessionID, clientTurnID).Scan(&existingID, &existingHash, &existingStatus)
|
||||
if err == nil {
|
||||
if existingHash != "" && existingHash != bodyHash {
|
||||
return nil, ErrIdempotencyKeyReused
|
||||
}
|
||||
rec, recErr := m.GetChatTurn(ctx, sessionID, clientTurnID)
|
||||
if recErr != nil {
|
||||
return nil, recErr
|
||||
}
|
||||
_, _ = 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
|
||||
}
|
||||
|
||||
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(),
|
||||
}
|
||||
return &AcquireTurnResult{Record: rec, WonLease: true}, nil
|
||||
}
|
||||
|
||||
var memTurnsMu sync.Mutex
|
||||
var memTurns = map[string]*ChatTurnRecord{} // sessionID|clientTurnID -> record
|
||||
|
||||
func turnMemKey(sessionID, clientTurnID string) string {
|
||||
return sessionID + "|" + clientTurnID
|
||||
}
|
||||
|
||||
func (m *ContextManager) acquireChatTurnMemory(
|
||||
sessionID, clientTurnID, bodyHash, requestID string,
|
||||
userID, projectID int64,
|
||||
) (*AcquireTurnResult, error) {
|
||||
memTurnsMu.Lock()
|
||||
defer memTurnsMu.Unlock()
|
||||
|
||||
key := turnMemKey(sessionID, clientTurnID)
|
||||
if existing, ok := memTurns[key]; ok {
|
||||
if existing.RequestBodySHA256 != bodyHash {
|
||||
return nil, ErrIdempotencyKeyReused
|
||||
}
|
||||
existing.AttemptCount++
|
||||
existing.LastRequestID = requestID
|
||||
existing.UpdatedAt = time.Now()
|
||||
return &AcquireTurnResult{Record: existing, WonLease: false}, nil
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
rec := &ChatTurnRecord{
|
||||
ID: int64(len(memTurns) + 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,
|
||||
}
|
||||
memTurns[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
|
||||
}
|
||||
memTurnsMu.Lock()
|
||||
defer memTurnsMu.Unlock()
|
||||
if rec, ok := memTurns[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)
|
||||
return err
|
||||
}
|
||||
memTurnsMu.Lock()
|
||||
defer memTurnsMu.Unlock()
|
||||
if rec, ok := memTurns[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)
|
||||
return err
|
||||
}
|
||||
memTurnsMu.Lock()
|
||||
defer memTurnsMu.Unlock()
|
||||
if rec, ok := memTurns[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
|
||||
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
|
||||
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,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, ErrChatTurnNotFound
|
||||
}
|
||||
rec.CompletedAt = completedAt
|
||||
return &rec, nil
|
||||
}
|
||||
memTurnsMu.Lock()
|
||||
defer memTurnsMu.Unlock()
|
||||
rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)]
|
||||
if !ok {
|
||||
return nil, ErrChatTurnNotFound
|
||||
}
|
||||
cp := *rec
|
||||
return &cp, nil
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package ctxmgr
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAcquireChatTurn_DuplicateSameBody(t *testing.T) {
|
||||
memTurns = map[string]*ChatTurnRecord{}
|
||||
mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig())
|
||||
ctx := context.Background()
|
||||
sessionID := "sess-dup"
|
||||
turnID := "turn-abc"
|
||||
bodyHash := HashTurnBody(sessionID, "hello world")
|
||||
|
||||
acq1, err := mgr.AcquireChatTurn(ctx, sessionID, turnID, bodyHash, "req-1", 1, 1)
|
||||
if err != nil || !acq1.WonLease {
|
||||
t.Fatalf("first acquire should win lease: %v %+v", err, acq1)
|
||||
}
|
||||
|
||||
acq2, err := mgr.AcquireChatTurn(ctx, sessionID, turnID, bodyHash, "req-2", 1, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("second acquire err: %v", err)
|
||||
}
|
||||
if acq2.WonLease {
|
||||
t.Fatal("second acquire must not win lease")
|
||||
}
|
||||
if acq2.Record.AttemptCount < 2 {
|
||||
t.Fatalf("attempt_count should increment, got %d", acq2.Record.AttemptCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcquireChatTurn_ConcurrentDuplicate(t *testing.T) {
|
||||
memTurns = map[string]*ChatTurnRecord{}
|
||||
mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig())
|
||||
ctx := context.Background()
|
||||
sessionID := "sess-conc"
|
||||
turnID := "turn-conc"
|
||||
bodyHash := HashTurnBody(sessionID, "concurrent msg")
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wins := 0
|
||||
var mu sync.Mutex
|
||||
for i := 0; i < 5; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
acq, err := mgr.AcquireChatTurn(ctx, sessionID, turnID, bodyHash, "req", 1, 1)
|
||||
if err != nil {
|
||||
t.Errorf("acquire %d: %v", n, err)
|
||||
return
|
||||
}
|
||||
if acq.WonLease {
|
||||
mu.Lock()
|
||||
wins++
|
||||
mu.Unlock()
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
if wins != 1 {
|
||||
t.Fatalf("expected exactly 1 winner, got %d", wins)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcquireChatTurn_DifferentBodyRejected(t *testing.T) {
|
||||
memTurns = map[string]*ChatTurnRecord{}
|
||||
mgr := NewContextManager(nil, nil, DefaultTimeDecayConfig())
|
||||
ctx := context.Background()
|
||||
sessionID := "sess-reuse"
|
||||
turnID := "turn-x"
|
||||
|
||||
_, err := mgr.AcquireChatTurn(ctx, sessionID, turnID, HashTurnBody(sessionID, "msg-a"), "req-1", 1, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = mgr.AcquireChatTurn(ctx, sessionID, turnID, HashTurnBody(sessionID, "msg-b"), "req-2", 1, 1)
|
||||
if err != ErrIdempotencyKeyReused {
|
||||
t.Fatalf("expected ErrIdempotencyKeyReused, got %v", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user