story/internal/vikingmemory/service.go

package vikingmemory

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"data-terminal-cloud/internal/outbox"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const SyncEvent = "conversation.viking_sync_requested"
const eligibleSQL = `SELECT EXISTS(SELECT 1 FROM subject_profile sp WHERE sp.id=$2::uuid AND sp.tenant_id=$1::uuid AND sp.status='active'
 AND NOT EXISTS(SELECT 1 FROM memory_user_state mus WHERE mus.subject_id=sp.id)
 AND NOT EXISTS(SELECT 1 FROM privacy_request pr WHERE pr.tenant_id=sp.tenant_id AND pr.subject_id=sp.id AND pr.request_type='DELETE')
 AND EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='personalized_memory' AND cr.withdrawn_at IS NULL AND (cr.expires_at IS NULL OR cr.expires_at>now()))
 AND (NOT EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='transcript_storage')
 OR EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='transcript_storage' AND cr.withdrawn_at IS NULL AND (cr.expires_at IS NULL OR cr.expires_at>now()))))`

type Service struct {
	pool   *pgxpool.Pool
	client *Client
}

// NativeAccess contains only non-secret values that Volcano needs to query
// this user's memory. The Viking API key is never returned to the device.
type NativeAccess struct {
	SubjectID, UserID, CollectionName, AssistantID string
	MemoryTypes                                    []string
	Limit                                          int
	TransitionWords                                string
}

func NewService(pool *pgxpool.Pool, client *Client) *Service {
	return &Service{pool: pool, client: client}
}
func (s *Service) Eligible(ctx context.Context, tenant, subject string) (bool, error) {
	var ok bool
	err := s.pool.QueryRow(ctx, eligibleSQL, tenant, subject).Scan(&ok)
	return ok, err
}
func (s *Service) Search(ctx context.Context, tenant, subject, query string, limit int) ([]Item, error) {
	ok, err := s.Eligible(ctx, tenant, subject)
	if err != nil || !ok {
		return []Item{}, err
	}
	items, err := s.client.Search(ctx, UserID(tenant, subject), query, limit)
	if err != nil {
		return nil, err
	}
	// Recheck after network I/O so a concurrent withdrawal cannot expose results.
	ok, err = s.Eligible(ctx, tenant, subject)
	if err != nil || !ok {
		return []Item{}, err
	}
	return items, nil
}
func (s *Service) PrepareForDevice(ctx context.Context, tenant, device string) (string, string, error) {
	var subject string
	err := s.pool.QueryRow(ctx, `SELECT b.subject_id::text FROM device_binding b JOIN device d ON d.id=b.device_id AND d.tenant_id=b.tenant_id WHERE b.device_id=$1::uuid AND b.tenant_id=$2::uuid AND b.unbound_at IS NULL AND d.status='registered'`, device, tenant).Scan(&subject)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", "", nil
	}
	if err != nil {
		return "", "", err
	}
	// Empty query retrieves recent memories and the user profile at session start.
	items, err := s.Search(ctx, tenant, subject, "", 4)
	if err != nil {
		return "", "", err
	}
	parts := []string{}
	for _, item := range items {
		parts = append(parts, item.Content)
	}
	return subject, strings.Join(parts, "\n"), nil
}

func (s *Service) NativeAccessForDevice(ctx context.Context, tenant, device string) (*NativeAccess, error) {
	if strings.TrimSpace(s.client.options.CollectionName) == "" {
		return nil, nil
	}
	var subject string
	err := s.pool.QueryRow(ctx, `SELECT b.subject_id::text FROM device_binding b JOIN device d ON d.id=b.device_id AND d.tenant_id=b.tenant_id WHERE b.device_id=$1::uuid AND b.tenant_id=$2::uuid AND b.unbound_at IS NULL AND d.status='registered'`, device, tenant).Scan(&subject)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	ok, err := s.Eligible(ctx, tenant, subject)
	if err != nil || !ok {
		return nil, err
	}
	memoryTypes := []string{s.client.options.EventType}
	if s.client.options.ProfileType != "" && s.client.options.ProfileType != s.client.options.EventType {
		memoryTypes = append(memoryTypes, s.client.options.ProfileType)
	}
	return &NativeAccess{
		SubjectID: subject, UserID: UserID(tenant, subject), CollectionName: s.client.options.CollectionName,
		AssistantID: AssistantID, MemoryTypes: memoryTypes, Limit: 5,
		TransitionWords: "可参考的用户历史记忆如下，请自然承接；若与用户最新说法冲突，以最新说法为准：",
	}, nil
}
func (s *Service) Enqueue(ctx context.Context, event outbox.Event) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(context.Background())
	var tenant string
	var generation int64
	err = tx.QueryRow(ctx, `SELECT c.tenant_id::text,COALESCE(max(ct.sequence_no),0) FROM conversation c LEFT JOIN conversation_turn ct ON ct.conversation_id=c.id WHERE c.id=$1::uuid AND c.status IN ('COMPLETED','FAILED') GROUP BY c.tenant_id`, event.AggregateID).Scan(&tenant, &generation)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil
	}
	if err != nil {
		return err
	}
	if generation == 0 {
		return nil
	}
	payload, _ := json.Marshal(map[string]int64{"generation_version": generation})
	key := fmt.Sprintf("viking-v1:%s:%d", event.AggregateID, generation)
	_, err = outbox.NewStore(s.pool).Enqueue(ctx, tx, outbox.NewEvent{TenantID: &tenant, AggregateType: "conversation", AggregateID: event.AggregateID, EventType: SyncEvent, DeduplicationKey: &key, Payload: payload})
	if err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Service) Handle(ctx context.Context, event outbox.Event) error {
	var tenant, subject, publicID string
	var when time.Time
	var frozen []byte
	err := s.pool.QueryRow(ctx, `SELECT c.tenant_id::text,c.subject_id::text,c.public_id::text,COALESCE(c.started_at,c.created_at),cs.effective_consents FROM conversation c JOIN consent_snapshot cs ON cs.id=c.consent_snapshot_id AND cs.tenant_id=c.tenant_id WHERE c.id=$1::uuid AND c.status IN ('COMPLETED','FAILED')`, event.AggregateID).Scan(&tenant, &subject, &publicID, &when, &frozen)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil
	}
	if err != nil {
		return err
	}
	consents := map[string]bool{}
	if json.Unmarshal(frozen, &consents) != nil || !consents["personalized_memory"] || !consents["transcript_storage"] {
		return nil
	}
	ok, err := s.Eligible(ctx, tenant, subject)
	if err != nil || !ok {
		return err
	}
	// Serialize external writes and erasure per user, including retries from another worker.
	conn, err := s.pool.Acquire(ctx)
	if err != nil {
		return err
	}
	defer conn.Release()
	user := UserID(tenant, subject)
	if _, err = conn.Exec(ctx, `SELECT pg_advisory_lock(hashtextextended($1,0))`, user); err != nil {
		return err
	}
	defer conn.Exec(context.Background(), `SELECT pg_advisory_unlock(hashtextextended($1,0))`, user)
	ok, err = s.Eligible(ctx, tenant, subject)
	if err != nil || !ok {
		return err
	}
	rows, err := conn.Query(ctx, `SELECT speaker,final_text,start_at FROM conversation_turn WHERE conversation_id=$1::uuid ORDER BY sequence_no,id`, event.AggregateID)
	if err != nil {
		return err
	}
	messages := []Message{}
	hasUser, hasAssistant := false, false
	for rows.Next() {
		var speaker, text string
		var at time.Time
		if err = rows.Scan(&speaker, &text, &at); err != nil {
			rows.Close()
			return err
		}
		role, id := "", user
		if speaker == "subject" {
			role = "user"
			hasUser = true
		} else if speaker == "agent" {
			role = "assistant"
			id = AssistantID
			hasAssistant = true
		}
		if role != "" && strings.TrimSpace(text) != "" {
			messages = append(messages, Message{Role: role, Content: text, RoleID: id, Time: at.UnixMilli()})
		}
	}
	rows.Close()
	if err = rows.Err(); err != nil {
		return err
	}
	if !hasUser || !hasAssistant || len(messages) == 0 {
		return nil
	}
	serialized, _ := json.Marshal(messages)
	sum := sha256.Sum256(serialized)
	digest := hex.EncodeToString(sum[:])
	var already bool
	err = conn.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM audit_event WHERE tenant_id=$1::uuid AND resource_id=$2::uuid AND action='memory.viking.session_sync' AND detail->>'digest'=$3 AND detail->>'collection'=$4)`, tenant, publicID, digest, s.client.options.ResourceID).Scan(&already)
	if err != nil || already {
		return err
	}
	for start := 0; start < len(messages); start += 100 {
		end := min(start+100, len(messages))
		if err = s.client.AddSession(ctx, SessionID(publicID, start/100), user, messages[start:end], when.UnixMilli()); err != nil {
			return err
		}
	}
	ok, err = s.Eligible(ctx, tenant, subject)
	if err != nil {
		return err
	}
	if !ok {
		return s.client.DeleteUser(ctx, user)
	}
	detail, _ := json.Marshal(map[string]any{"digest": digest, "collection": s.client.options.ResourceID, "message_count": len(messages), "user_id": user})
	_, err = conn.Exec(ctx, `INSERT INTO audit_event(tenant_id,actor_type,action,resource_type,resource_id,result,detail) VALUES($1::uuid,'system','memory.viking.session_sync','conversation',$2::uuid,'success',$3::jsonb)`, tenant, publicID, detail)
	return err
}

func (s *Service) DeleteSubject(ctx context.Context, tenant, subject string) error {
	if tenant == "" || subject == "" {
		return errors.New("Viking erasure requires tenant and subject")
	}
	conn, err := s.pool.Acquire(ctx)
	if err != nil {
		return err
	}
	defer conn.Release()
	user := UserID(tenant, subject)
	if _, err = conn.Exec(ctx, `SELECT pg_advisory_lock(hashtextextended($1,0))`, user); err != nil {
		return err
	}
	defer conn.Exec(context.Background(), `SELECT pg_advisory_unlock(hashtextextended($1,0))`, user)
	if err := s.client.deleteManagedUser(ctx, user); err != nil {
		return err
	}
	return s.client.DeleteUser(ctx, user)
}