Agent-System/internal/ctxmgr/message_update.go

package ctxmgr

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"chat-agent-system/internal/app/uow"
	"chat-agent-system/internal/replyexport"

	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

// UpdateMessageMetadata 更新已持久化助手消息的 metadata 与任务状态。
func (m *ContextManager) UpdateMessageMetadata(ctx context.Context, messageID int64, meta map[string]any, taskStatus string) error {
	if messageID <= 0 {
		return errors.New("invalid message_id")
	}
	if meta == nil {
		meta = map[string]any{}
	}
	if taskStatus != "" {
		meta["task_status"] = taskStatus
	}

	if m.pool != nil || hasUoW(ctx) {
		metaJSON, _ := json.Marshal(meta)
		_, err := execChat(ctx, m.pool, `
			UPDATE chat_messages
			SET metadata = COALESCE(metadata, '{}'::jsonb) || $1::jsonb
			WHERE id = $2
		`, metaJSON, messageID)
		if err != nil {
			return fmt.Errorf("update message metadata failed: %w", err)
		}
		return nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	for sid, msgs := range m.messages {
		for i := range msgs {
			if msgs[i].ID == messageID {
				if msgs[i].Metadata == nil {
					msgs[i].Metadata = map[string]any{}
				}
				for k, v := range meta {
					msgs[i].Metadata[k] = v
				}
				m.messages[sid][i] = msgs[i]
				return nil
			}
		}
	}
	return errors.New("message not found")
}

// FinalizeAssistantTaskMessage 替换助手任务占位正文并写入完整 metadata（含 task_status）。
func (m *ContextManager) FinalizeAssistantTaskMessage(ctx context.Context, messageID int64, content string, meta map[string]any, taskStatus string) error {
	if messageID <= 0 {
		return errors.New("invalid message_id")
	}
	if meta == nil {
		meta = map[string]any{}
	}
	if taskStatus != "" {
		meta["task_status"] = taskStatus
	}

	kind, cap := replyexport.EvaluatePolicy(replyexport.PolicyOptions{
		Role:    "assistant",
		Status:  nzStatus(taskStatus),
		Content: content,
	})
	meta["content_kind"] = string(kind)
	meta["export_policy"] = string(cap.Policy)
	meta["content_revision"] = cap.ContentRevision
	meta["content_sha256"] = cap.ContentSha256

	if m.pool != nil || hasUoW(ctx) {
		metaJSON, err := json.Marshal(meta)
		if err != nil {
			return fmt.Errorf("marshal metadata failed: %w", err)
		}
		tag, err := execChat(ctx, m.pool, `
			UPDATE chat_messages
			SET content = $1,
			    metadata = COALESCE(metadata, '{}'::jsonb) || $2::jsonb,
			    export_policy = $3,
			    content_sha256 = $4,
			    content_kind = $5
			WHERE id = $6
		`, content, metaJSON, string(cap.Policy), cap.ContentSha256, string(kind), messageID)
		if err != nil {
			return fmt.Errorf("finalize assistant task message failed: %w", err)
		}
		if tag.RowsAffected() == 0 {
			return errors.New("message not found")
		}
		return nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	for sid, msgs := range m.messages {
		for i := range msgs {
			if msgs[i].ID == messageID {
				msgs[i].Content = content
				if msgs[i].Metadata == nil {
					msgs[i].Metadata = map[string]any{}
				}
				for k, v := range meta {
					msgs[i].Metadata[k] = v
				}
				msgs[i].ExportPolicy = string(cap.Policy)
				msgs[i].ContentSha256 = cap.ContentSha256
				msgs[i].ContentKind = string(kind)
				msgs[i].ContentRevision = cap.ContentRevision
				m.messages[sid][i] = msgs[i]
				return nil
			}
		}
	}
	return errors.New("message not found")
}

// SetInteractiveCardStatus marks a stored chat card resolved or otherwise closed.
func (m *ContextManager) SetInteractiveCardStatus(ctx context.Context, messageID int64, status, summary string) error {
	if m == nil || messageID <= 0 || status == "" {
		return errors.New("invalid card status update")
	}
	if m.pool != nil || hasUoW(ctx) {
		_, err := execChat(ctx, m.pool, `
			UPDATE chat_messages
			SET metadata = jsonb_set(
				jsonb_set(COALESCE(metadata, '{}'::jsonb), '{interactive_card,status}', to_jsonb($2::text), true),
				'{interactive_card,resolutionSummary}', to_jsonb($3::text), true)
			WHERE id = $1
		`, messageID, status, summary)
		if err != nil {
			return fmt.Errorf("set interactive card status: %w", err)
		}
		return nil
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	for sid, msgs := range m.messages {
		for i := range msgs {
			if msgs[i].ID != messageID {
				continue
			}
			card, _ := msgs[i].Metadata["interactive_card"].(map[string]any)
			if card == nil {
				card = map[string]any{}
			}
			card["status"] = status
			card["resolutionSummary"] = summary
			if msgs[i].Metadata == nil {
				msgs[i].Metadata = map[string]any{}
			}
			msgs[i].Metadata["interactive_card"] = card
			m.messages[sid][i] = msgs[i]
			return nil
		}
	}
	return errors.New("message not found")
}

func hasUoW(ctx context.Context) bool {
	_, ok := uow.ScopeFrom(ctx)
	return ok
}

func execChat(ctx context.Context, pool *pgxpool.Pool, q string, args ...any) (pgconn.CommandTag, error) {
	if scope, ok := uow.ScopeFrom(ctx); ok {
		tx, err := uow.AsPG(scope)
		if err != nil {
			return pgconn.CommandTag{}, err
		}
		return tx.Exec(ctx, q, args...)
	}
	if pool == nil {
		return pgconn.CommandTag{}, errors.New("ctxmgr: database required")
	}
	return pool.Exec(ctx, q, args...)
}

func nzStatus(taskStatus string) string {
	if taskStatus == "" {
		return "completed"
	}
	return taskStatus
}