Agent-System/internal/ctxmgr/chat_turns.go

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
	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[:])
}

// 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
}

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
		}
		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 := memTurns[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
	}

	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,
		LeaseExpiresAt:    &leaseUntil,
	}
	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
		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
	}
	memTurnsMu.Lock()
	defer memTurnsMu.Unlock()
	rec, ok := memTurns[turnMemKey(sessionID, clientTurnID)]
	if !ok {
		return nil, ErrChatTurnNotFound
	}
	cp := *rec
	return &cp, nil
}