Agent-System/internal/ctxmgr/message_update.go

package ctxmgr

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

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

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

	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
			WHERE id = $3
		`, content, metaJSON, 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
				}
				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...)
}