diff --git a/deploy/postgres/context_init.sql b/deploy/postgres/context_init.sql index f994e94c..459d2326 100644 --- a/deploy/postgres/context_init.sql +++ b/deploy/postgres/context_init.sql @@ -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, diff --git a/deploy/postgres/migrations/2026-09-11-chat-turns.sql b/deploy/postgres/migrations/2026-09-11-chat-turns.sql new file mode 100644 index 00000000..0dc1fcdd --- /dev/null +++ b/deploy/postgres/migrations/2026-09-11-chat-turns.sql @@ -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); diff --git a/internal/ctxmgr/chat_turns.go b/internal/ctxmgr/chat_turns.go new file mode 100644 index 00000000..f56ddae0 --- /dev/null +++ b/internal/ctxmgr/chat_turns.go @@ -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 +} diff --git a/internal/ctxmgr/chat_turns_test.go b/internal/ctxmgr/chat_turns_test.go new file mode 100644 index 00000000..354ad713 --- /dev/null +++ b/internal/ctxmgr/chat_turns_test.go @@ -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) + } +}