Agent-System/internal/server/vault.go

package server

import (
	"context"
	"fmt"
	"net/http"
	"sync/atomic"
	"time"

	"chat-agent-system/internal/artifactvault"
	"chat-agent-system/internal/config"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/replyexport"
	"chat-agent-system/internal/taskrun"

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

func (s *Server) initArtifactVault() {
	if s == nil || s.cfg == nil {
		return
	}
	env := s.cfg.Env
	if env == "" {
		env = "dev"
	}
	if s.rt != nil {
		s.vaultReg = s.rt.ArtifactVaultRegistry()
	} else {
		s.vaultReg = artifactvault.NewMemoryRegistry()
	}
	s.vaultIngest = &artifactvault.Ingest{Env: env, Reg: s.vaultReg}
	if s.cfg.VaultStateRoot == "" || s.cfg.VaultMasterKey == "" {
		return
	}
	if !s.cfg.FAVPublisherEnabled {
		return
	}
	if err := config.ValidateIndependentVault(s.cfg.VaultStateRoot, s.cfg.DataDir); err != nil {
		return
	}
	ident := s.cfg.VaultStorageIdentity
	if ident == "" {
		ident = "vault-" + env
	}
	store, err := artifactvault.NewFilesystemStore(s.cfg.VaultStateRoot, ident, []byte(s.cfg.VaultMasterKey))
	if err != nil {
		return
	}
	s.vaultStore = store
	s.vaultPub = &artifactvault.Publisher{Reg: s.vaultReg, Store: store}
}

func (s *Server) startVaultPublisher() {
	if s == nil || s.vaultPub == nil || s.cfg == nil || !s.cfg.FAVPublisherEnabled {
		return
	}
	s.vaultStop = make(chan struct{})
	go func() {
		t := time.NewTicker(2 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-s.vaultStop:
				return
			case <-t.C:
				_ = s.vaultPub.Drain(context.Background())
			}
		}
	}()
	atomic.StoreInt32(&s.vaultPublisherStarted, 1)
}

func (s *Server) vaultReady() (identityOK, connected, lagOK bool) {
	if s == nil || s.vaultStore == nil {
		return false, false, true
	}
	idOK, conn, _ := s.vaultStore.Ready()
	lagOK = true
	if s.vaultReg != nil {
		pending, err := s.vaultReg.LoadRequested(context.Background())
		lagOK = err == nil && len(pending) < 100
	}
	return idOK, conn, lagOK
}

func (s *Server) ingestAssistantReply(ctx context.Context, record *taskrun.RunRecord, content, terminalStatus string) {
	if s == nil || s.vaultIngest == nil || record == nil || terminalStatus != "completed" {
		return
	}
	_, cap := replyexport.EvaluatePolicy(replyexport.PolicyOptions{
		Role: "assistant", Status: "completed", Content: content,
	})
	now := time.Now().UTC()
	logical := fmt.Sprintf("msg:%d", record.AssistantMsgID)
	eventID := record.SessionID + ":" + record.ClientTurnID
	if eventID == ":" {
		eventID = record.RunID
	}
	_ = s.vaultIngest.AssistantReply(ctx, record.ProjectID, logical, eventID, record.AssistantMsgID, record.RunID, string(cap.Policy), &now, []byte(content))
}

func (s *Server) wireFinalArtifactHooks() {
	if s == nil || s.vaultIngest == nil {
		return
	}
	if gov := s.gov; gov != nil {
		gov.OnFinalArtifact = func(ctx context.Context, doc *docgov.DocumentRecord, markdown []byte, snapshotID int64) {
			if doc == nil {
				return
			}
			_ = s.vaultIngest.Approved(ctx, doc.ProjectID, doc.DocID, snapshotID, doc.ApprovedAt, markdown)
		}
		gov.OnFinalArtifactTX = func(ctx context.Context, tx pgx.Tx, doc *docgov.DocumentRecord, markdown []byte, snapshotID int64) error {
			if doc == nil {
				return nil
			}
			return s.vaultIngest.ApprovedTX(ctx, tx, doc.ProjectID, doc.DocID, snapshotID, doc.ApprovedAt, markdown)
		}
	}
	if rev := s.rev; rev != nil {
		rev.OnFinalArtifact = func(ctx context.Context, kind, docID string, projectID, snapshotID int64, eventID string, payload []byte, sealedAt *time.Time) {
			switch kind {
			case "agent_commit":
				_ = s.vaultIngest.AgentCommit(ctx, projectID, docID, eventID, snapshotID, payload)
			case "version_sealed":
				_ = s.vaultIngest.SealedVersion(ctx, projectID, docID, eventID, snapshotID, sealedAt, payload)
			}
		}
		rev.OnFinalArtifactTX = func(ctx context.Context, tx pgx.Tx, kind, docID string, projectID, snapshotID int64, eventID string, payload []byte, sealedAt *time.Time) error {
			switch kind {
			case "agent_commit":
				return s.vaultIngest.AgentCommitTX(ctx, tx, projectID, docID, eventID, snapshotID, payload)
			case "version_sealed":
				return s.vaultIngest.SealedVersionTX(ctx, tx, projectID, docID, eventID, snapshotID, sealedAt, payload)
			default:
				return nil
			}
		}
	}
	if fd := s.freeDoc; fd != nil {
		fd.OnFinalArtifact = func(ctx context.Context, projectID int64, docID, eventID string, snapshotID int64, payload []byte) {
			_ = s.vaultIngest.FreeformPublish(ctx, projectID, docID, eventID, snapshotID, payload)
		}
	}
}

func (s *Server) tryVaultDegraded(ctx context.Context, env, logical, projectKey string) ([]byte, string, error) {
	if s == nil || s.vaultPub == nil || s.cfg == nil || !s.cfg.FAVDegradedRead {
		return nil, "", fmt.Errorf("artifactvault: degraded read disabled")
	}
	return s.vaultPub.DegradedPayload(ctx, env, logical, projectKey)
}

func (s *Server) handleVaultStatus(w http.ResponseWriter, r *http.Request) {
	_ = r
	cfg := s.cfg
	degraded := false
	publisher := false
	storeReady := false
	if cfg != nil {
		degraded = cfg.FAVDegradedRead
		publisher = cfg.FAVPublisherEnabled
	}
	if s.vaultStore != nil {
		storeReady = true
	}
	writeJSONResponse(w, http.StatusOK, map[string]any{
		"status":             "ok",
		"consumer":           "fav",
		"consumer_mode":      "off",
		"publisher_enabled":  publisher,
		"publisher_running":  atomic.LoadInt32(&s.vaultPublisherStarted) == 1,
		"degraded_read":      degraded,
		"store_configured":   storeReady,
		"subscriber_default": false,
		"shadow_observe":     true,
		"would_stage":        false,
		"would_publish":      false,
		"live_not_opened":    true,
	})
}