285 lines
13 KiB
Go
285 lines
13 KiB
Go
|
package vikingmemory
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"crypto/sha256"
|
||
|
|
"encoding/hex"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"data-terminal-cloud/internal/callmemory"
|
||
|
|
"data-terminal-cloud/internal/outbox"
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"github.com/jackc/pgx/v5"
|
||
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Managed struct {
|
||
|
|
pool *pgxpool.Pool
|
||
|
|
client *Client
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewManaged(p *pgxpool.Pool, c *Client) *Managed { return &Managed{p, c} }
|
||
|
|
|
||
|
|
// Upload only ended calls with the currently authorized, retained transcript.
|
||
|
|
// A late final subtitle creates another revision; the stable remote session ID
|
||
|
|
// replaces that session instead of accumulating duplicate uploads on retry.
|
||
|
|
func (m *Managed) SyncNext(ctx context.Context) error {
|
||
|
|
_, err := m.pool.Exec(ctx, `INSERT INTO call_memory_sync(session_id,transcript_revision)
|
||
|
|
SELECT id,transcript_revision FROM call_session c WHERE budget_snapshot->>'protocol'='DIRECT-CHAT-1'
|
||
|
|
AND ended_at IS NOT NULL AND transcript_state IN ('PARTIAL','COMPLETE') AND transcript_revision>0 AND memory_mode<>'OFF'
|
||
|
|
AND managed_call_memory_authorized(subject_id) ON CONFLICT DO NOTHING`)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
conn, err := m.pool.Acquire(ctx)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
defer conn.Release()
|
||
|
|
var sid, tenant, subject, state string
|
||
|
|
var revision int64
|
||
|
|
var uploaded *time.Time
|
||
|
|
err = conn.QueryRow(ctx, `SELECT j.session_id::text,c.tenant_id::text,c.subject_id::text,j.transcript_revision,j.state,j.uploaded_at
|
||
|
|
FROM call_memory_sync j JOIN call_session c ON c.id=j.session_id
|
||
|
|
WHERE (j.state<>'SKIPPED' OR (j.uploaded_at IS NOT NULL AND NOT managed_call_memory_authorized(c.subject_id))) AND j.next_attempt_at<=now() ORDER BY j.next_attempt_at LIMIT 1`).Scan(&sid, &tenant, &subject, &revision, &state, &uploaded)
|
||
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
user := UserID(tenant, subject)
|
||
|
|
var locked bool
|
||
|
|
if err = conn.QueryRow(ctx, `SELECT pg_try_advisory_lock(hashtextextended($1,0))`, user).Scan(&locked); err != nil || !locked {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
defer func() {
|
||
|
|
c, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
if _, e := conn.Exec(c, `SELECT pg_advisory_unlock(hashtextextended($1,0))`, user); e != nil {
|
||
|
|
conn.Conn().Close(c)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
// Re-read after the lock: another worker may already have handled this job.
|
||
|
|
err = conn.QueryRow(ctx, `SELECT j.state,j.uploaded_at FROM call_memory_sync j JOIN call_session c ON c.id=j.session_id WHERE j.session_id=$1::uuid AND j.transcript_revision=$2 AND (j.state<>'SKIPPED' OR (j.uploaded_at IS NOT NULL AND NOT managed_call_memory_authorized(c.subject_id))) AND j.next_attempt_at<=now()`, sid, revision).Scan(&state, &uploaded)
|
||
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
var allowed, suppressed bool
|
||
|
|
var current int64
|
||
|
|
err = conn.QueryRow(ctx, `SELECT managed_call_memory_authorized(c.subject_id),c.transcript_revision,
|
||
|
|
EXISTS(SELECT 1 FROM memory_source_suppression s CROSS JOIN LATERAL jsonb_array_elements(s.source_refs) r WHERE s.subject_id=c.subject_id AND r->>'session_id'=c.id::text)
|
||
|
|
FROM call_session c WHERE id=$1::uuid`, sid).Scan(&allowed, ¤t, &suppressed)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
skip := func() error {
|
||
|
|
_, e := conn.Exec(ctx, `UPDATE call_memory_sync SET state='SKIPPED',last_error=NULL,next_attempt_at=now()+interval '5 minutes' WHERE session_id=$1::uuid AND transcript_revision=$2`, sid, revision)
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if !allowed {
|
||
|
|
// A withdrawal must also remove managed provider data; do not merely hide it.
|
||
|
|
if err = m.client.deleteManagedUser(ctx, user); err != nil {
|
||
|
|
return m.retry(ctx, conn, sid, revision, "VIKING_CLEANUP_FAILED")
|
||
|
|
}
|
||
|
|
return skip()
|
||
|
|
}
|
||
|
|
if revision != current || suppressed {
|
||
|
|
return skip()
|
||
|
|
}
|
||
|
|
remote := SessionID(sid, 0)
|
||
|
|
if uploaded == nil {
|
||
|
|
rows, e := conn.Query(ctx, `SELECT speaker,text_content,source_at FROM call_utterance WHERE session_id=$1::uuid ORDER BY source_at,utterance_id`, sid)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
messages := []Message{}
|
||
|
|
hasUser := false
|
||
|
|
hasAssistant := false
|
||
|
|
for rows.Next() {
|
||
|
|
var speaker, text string
|
||
|
|
var at time.Time
|
||
|
|
if e = rows.Scan(&speaker, &text, &at); e != nil {
|
||
|
|
rows.Close()
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if strings.TrimSpace(text) == "" || text == "[deleted]" {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
role, id := "", user
|
||
|
|
if speaker == "USER" {
|
||
|
|
role = "user"
|
||
|
|
hasUser = true
|
||
|
|
} else if speaker == "ASSISTANT" {
|
||
|
|
role = "assistant"
|
||
|
|
id = ManagedAssistantID
|
||
|
|
hasAssistant = true
|
||
|
|
}
|
||
|
|
if role != "" {
|
||
|
|
messages = append(messages, Message{Role: role, RoleID: id, Content: text, Time: at.UnixMilli()})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
e = rows.Err()
|
||
|
|
rows.Close()
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if !hasUser || !hasAssistant {
|
||
|
|
return skip()
|
||
|
|
}
|
||
|
|
if len(messages) > 1000 {
|
||
|
|
return m.retry(ctx, conn, sid, revision, "TRANSCRIPT_TOO_LARGE")
|
||
|
|
}
|
||
|
|
if e = m.client.AddManagedSession(ctx, remote, user, messages, messages[0].Time); e != nil {
|
||
|
|
return m.retry(ctx, conn, sid, revision, "VIKING_UPLOAD_FAILED")
|
||
|
|
}
|
||
|
|
_, err = conn.Exec(ctx, `UPDATE call_memory_sync SET state='UPLOADED',uploaded_at=now(),attempts=attempts+1,last_error=NULL,next_attempt_at=now()+interval '10 seconds' WHERE session_id=$1::uuid AND transcript_revision=$2`, sid, revision)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
items, err := m.client.ManagedEvents(ctx, user, remote)
|
||
|
|
if err != nil {
|
||
|
|
return m.retry(ctx, conn, sid, revision, "VIKING_RETRIEVAL_FAILED")
|
||
|
|
}
|
||
|
|
if err = m.importEvents(ctx, conn, sid, tenant, subject, revision, items); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
// Extraction is asynchronous. Keep polling, including after an empty result;
|
||
|
|
// an accepted upload is not proof that extraction completed or found a fact.
|
||
|
|
delay := 15
|
||
|
|
if time.Since(*uploaded) > 10*time.Minute {
|
||
|
|
delay = 300
|
||
|
|
}
|
||
|
|
if len(items) > 0 {
|
||
|
|
state = "SYNCED"
|
||
|
|
} else {
|
||
|
|
state = "UPLOADED"
|
||
|
|
}
|
||
|
|
_, err = conn.Exec(ctx, `UPDATE call_memory_sync SET state=$3,last_error=NULL,next_attempt_at=now()+make_interval(secs=>$4) WHERE session_id=$1::uuid AND transcript_revision=$2`, sid, revision, state, delay)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
func (m *Managed) retry(ctx context.Context, conn *pgxpool.Conn, sid string, revision int64, code string) error {
|
||
|
|
_, e := conn.Exec(ctx, `UPDATE call_memory_sync SET attempts=attempts+1,last_error=$3,next_attempt_at=now()+make_interval(secs=>LEAST(300,15*(attempts+1))) WHERE session_id=$1::uuid AND transcript_revision=$2`, sid, revision, code)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
return errors.New(code)
|
||
|
|
}
|
||
|
|
func (m *Managed) importEvents(ctx context.Context, conn *pgxpool.Conn, sid, tenant, subject string, revision int64, items []Item) error {
|
||
|
|
tx, e := conn.Begin(ctx)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
defer tx.Rollback(context.Background())
|
||
|
|
var allowed bool
|
||
|
|
e = tx.QueryRow(ctx, `SELECT managed_call_memory_authorized(us.subject_id) AND c.transcript_revision=$3 AND NOT EXISTS(SELECT 1 FROM memory_source_suppression s CROSS JOIN LATERAL jsonb_array_elements(s.source_refs) r WHERE s.subject_id=us.subject_id AND r->>'session_id'=$2) FROM memory_user_state us JOIN call_session c ON c.subject_id=us.subject_id WHERE us.subject_id=$1::uuid AND c.id=$2::uuid FOR UPDATE OF us`, subject, sid, revision).Scan(&allowed)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if !allowed {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
for _, item := range items {
|
||
|
|
sum := sha256.Sum256([]byte(item.Content))
|
||
|
|
digest := hex.EncodeToString(sum[:])
|
||
|
|
id := uuid.NewSHA1(uuid.NameSpaceOID, []byte("viking-managed:"+m.client.options.ResourceID+":"+subject+":"+item.ID)).String()
|
||
|
|
var localVersion, importedVersion int64
|
||
|
|
var status, policy, previousDigest string
|
||
|
|
e = tx.QueryRow(ctx, `SELECT cm.version,cm.status,cm.proactive_policy,o.imported_version,o.provider_digest FROM curated_memory cm JOIN managed_memory_origin o ON o.memory_id=cm.id WHERE cm.id=$1::uuid FOR UPDATE OF cm`, id).Scan(&localVersion, &status, &policy, &importedVersion, &previousDigest)
|
||
|
|
exists := e == nil
|
||
|
|
if e != nil && !errors.Is(e, pgx.ErrNoRows) {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
// Only untouched automatic records may follow a provider correction.
|
||
|
|
// Never overwrite a user's edit, disable, or deletion.
|
||
|
|
if exists && (localVersion != importedVersion || status != "ACTIVE" || policy != "ALLOWED" || previousDigest == digest) {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var refs []byte
|
||
|
|
if e = tx.QueryRow(ctx, `SELECT jsonb_build_array(jsonb_build_object('type','provider_memory','session_id',$1::text,'provider_id',$2::text)) || COALESCE(jsonb_agg(jsonb_build_object('type','user_turn','session_id',session_id::text,'utterance_id',utterance_id,'source_revision',source_revision)),'[]'::jsonb) FROM call_utterance WHERE session_id=$1::uuid`, sid, item.ID).Scan(&refs); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if !exists {
|
||
|
|
if _, e = tx.Exec(ctx, `INSERT INTO curated_memory(id,public_id,tenant_id,subject_id,status,memory_type,proactive_policy,sensitivity,entity_key,statement_basis,source_time) VALUES($1::uuid,$1::uuid,$2::uuid,$3::uuid,'ACTIVE','EVENT','ALLOWED','normal',$4,'MODEL_EXTRACTED',now())`, id, tenant, subject, "viking:"+item.ID); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
}
|
||
|
|
vid := uuid.NewString()
|
||
|
|
if _, e = tx.Exec(ctx, `INSERT INTO memory_version(id,memory_id,version_no,content,source_refs,change_reason,content_digest) VALUES($1::uuid,$2::uuid,$6,$3,$4::jsonb,'viking-session-extraction',$5)`, vid, id, item.Content, refs, digest, localVersion+1); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if _, e = tx.Exec(ctx, `UPDATE curated_memory SET current_version_id=$2::uuid,version=$3,updated_at=now() WHERE id=$1::uuid`, id, vid, localVersion+1); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if _, e = tx.Exec(ctx, `INSERT INTO managed_memory_origin(memory_id,subject_id,session_id,provider_id,provider_digest,imported_version) VALUES($1::uuid,$2::uuid,$3::uuid,$4,$5,$6) ON CONFLICT(memory_id) DO UPDATE SET provider_digest=EXCLUDED.provider_digest,imported_version=EXCLUDED.imported_version`, id, subject, sid, item.ID, digest, localVersion+1); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if _, e = tx.Exec(ctx, `INSERT INTO memory_index_entry(tenant_id,subject_id,memory_id,memory_version_id,index_type,status) VALUES($1::uuid,$2::uuid,$3::uuid,$4::uuid,'pgvector','PENDING')`, tenant, subject, id, vid); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
dedupe := "memory-index-v1:" + vid
|
||
|
|
payload, _ := json.Marshal(map[string]string{"memory_version_id": vid})
|
||
|
|
if _, e = outbox.NewStore(m.pool).Enqueue(ctx, tx, outbox.NewEvent{TenantID: &tenant, AggregateType: "curated_memory", AggregateID: id, EventType: "memory.index_requested", DeduplicationKey: &dedupe, Payload: payload}); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return tx.Commit(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
// The platform mirror is the serving source for pure S2S. Native MemoryConfig
|
||
|
|
// stays off, so cloud extraction cannot bypass local edits or deletions.
|
||
|
|
func (m *Managed) Project(ctx context.Context, r callmemory.ProjectionRecord, priorKind, priorID string, reconcile bool) (callmemory.ProjectionResult, error) {
|
||
|
|
result := callmemory.ProjectionResult{Kind: "event", ID: r.ID, Visible: true}
|
||
|
|
var sid, id string
|
||
|
|
err := m.pool.QueryRow(ctx, `SELECT o.session_id::text,o.provider_id FROM managed_memory_origin o JOIN curated_memory cm ON cm.id=o.memory_id WHERE cm.public_id=$1::uuid AND cm.subject_id=$2::uuid AND cm.tenant_id=$3::uuid`, r.ID, r.Subject, r.Tenant).Scan(&sid, &id)
|
||
|
|
// Manually authored platform memories do not require a provider mirror to serve.
|
||
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
||
|
|
var direct bool
|
||
|
|
if e := m.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM call_session WHERE subject_id=$1::uuid AND budget_snapshot->>'protocol'='DIRECT-CHAT-1')`, r.Subject).Scan(&direct); e != nil {
|
||
|
|
return result, e
|
||
|
|
}
|
||
|
|
if !direct {
|
||
|
|
return NewProjectionAdapter(m.pool, m.client).Project(ctx, r, priorKind, priorID, reconcile)
|
||
|
|
}
|
||
|
|
result.Kind = "local"
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return result, err
|
||
|
|
}
|
||
|
|
result.ID = id
|
||
|
|
items, err := m.client.ManagedEvents(ctx, UserID(r.Tenant, r.Subject), SessionID(sid, 0))
|
||
|
|
if err != nil {
|
||
|
|
return result, err
|
||
|
|
}
|
||
|
|
for _, item := range items {
|
||
|
|
if item.ID != id {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
serve := r.Status == "ACTIVE" && r.Policy == "ALLOWED" && (r.Kind != "OPEN_ITEM" || r.ItemStatus == "OPEN")
|
||
|
|
if !serve {
|
||
|
|
err = m.client.call(ctx, "/api/memory/event/delete", map[string]any{"event_id": id}, nil)
|
||
|
|
} else if item.Content != r.Content {
|
||
|
|
// original_messages is returned by search but rejected by event/update.
|
||
|
|
err = m.client.call(ctx, "/api/memory/event/update", map[string]any{"event_id": id, "memory_info": map[string]string{"summary": r.Content}}, nil)
|
||
|
|
} else {
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
result.Visible = false
|
||
|
|
return result, err
|
||
|
|
}
|
||
|
|
// Absence means provider cleanup is complete. Local records remain authoritative.
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
func (c *Client) deleteManagedUser(ctx context.Context, user string) error {
|
||
|
|
if user == "" {
|
||
|
|
return errors.New("managed user required")
|
||
|
|
}
|
||
|
|
return c.call(ctx, "/api/memory/event/batch_delete", map[string]any{"delete_type": "full", "filter": map[string]any{"user_id": []string{user}, "assistant_id": []string{ManagedAssistantID}}}, nil)
|
||
|
|
}
|