408 lines
12 KiB
Go
408 lines
12 KiB
Go
package ctxmgr
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
ShadowStatusPending = "pending"
|
|
shadowJobMaxAttempt = 3
|
|
shadowLeaseTTL = 2 * time.Minute
|
|
)
|
|
|
|
// SemanticShadowJob is a durable unit of work for advisory semantic generation.
|
|
// It never replaces the served deterministic Narrative (REPLACE stays unwired).
|
|
type SemanticShadowJob struct {
|
|
ID int64 `json:"id,omitempty"`
|
|
SessionID string `json:"session_id"`
|
|
SummaryRevision int `json:"summary_revision"`
|
|
InputDigest string `json:"input_digest"`
|
|
VisibilityDigest string `json:"visibility_digest"`
|
|
WatermarkToMessageID int64 `json:"watermark_to_message_id"`
|
|
PriorRevision int `json:"prior_revision"`
|
|
Messages []MinimizedMessage `json:"messages"`
|
|
PriorNarrative string `json:"prior_narrative,omitempty"`
|
|
PriorFacts map[string]string `json:"prior_facts,omitempty"`
|
|
Status string `json:"status,omitempty"`
|
|
Attempt int `json:"attempt,omitempty"`
|
|
}
|
|
|
|
type shadowJobPayload struct {
|
|
Messages []MinimizedMessage `json:"messages"`
|
|
PriorNarrative string `json:"prior_narrative,omitempty"`
|
|
PriorFacts map[string]string `json:"prior_facts,omitempty"`
|
|
}
|
|
|
|
// EnqueueSemanticShadow schedules background semantic evaluation after the
|
|
// deterministic summary is already persisted and returned to the caller.
|
|
func (m *ContextManager) EnqueueSemanticShadow(ctx context.Context, job SemanticShadowJob) error {
|
|
if m == nil || job.SessionID == "" || job.InputDigest == "" {
|
|
return fmt.Errorf("invalid semantic shadow job")
|
|
}
|
|
if job.Status == "" {
|
|
job.Status = "pending"
|
|
}
|
|
if m.pool == nil {
|
|
m.rememberShadowJob(job)
|
|
m.kickShadowWorker()
|
|
return nil
|
|
}
|
|
payload, err := json.Marshal(shadowJobPayload{
|
|
Messages: job.Messages,
|
|
PriorNarrative: job.PriorNarrative,
|
|
PriorFacts: job.PriorFacts,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = m.pool.Exec(ctx, `
|
|
INSERT INTO chat_summary_shadow_jobs (
|
|
session_id, summary_revision, input_digest, visibility_digest,
|
|
watermark_to_message_id, prior_revision, payload, status, attempt, updated_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,$7,'pending',0,now())
|
|
ON CONFLICT (session_id, summary_revision, input_digest) DO NOTHING
|
|
`, job.SessionID, job.SummaryRevision, job.InputDigest, job.VisibilityDigest,
|
|
job.WatermarkToMessageID, job.PriorRevision, payload)
|
|
if err != nil {
|
|
slog.Warn("semantic shadow job persist failed; using in-process queue", "session_id", job.SessionID, "err", err)
|
|
m.rememberShadowJob(job)
|
|
}
|
|
m.kickShadowWorker()
|
|
return nil
|
|
}
|
|
|
|
func (m *ContextManager) rememberShadowJob(job SemanticShadowJob) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.shadowJobs == nil {
|
|
m.shadowJobs = make([]SemanticShadowJob, 0, 4)
|
|
}
|
|
for _, existing := range m.shadowJobs {
|
|
if existing.SessionID == job.SessionID && existing.SummaryRevision == job.SummaryRevision && existing.InputDigest == job.InputDigest {
|
|
return
|
|
}
|
|
}
|
|
m.shadowJobs = append(m.shadowJobs, job)
|
|
}
|
|
|
|
func (m *ContextManager) kickShadowWorker() {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.shadowWorkerRunning {
|
|
return
|
|
}
|
|
m.shadowWorkerRunning = true
|
|
go m.shadowWorkerLoop()
|
|
}
|
|
|
|
// StartSemanticShadowWorker starts a long-lived reclaim loop for durable jobs.
|
|
func (m *ContextManager) StartSemanticShadowWorker(stop <-chan struct{}) {
|
|
if m == nil || m.pool == nil {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
if m.shadowReclaimStarted {
|
|
m.mu.Unlock()
|
|
return
|
|
}
|
|
m.shadowReclaimStarted = true
|
|
m.mu.Unlock()
|
|
go func() {
|
|
ticker := time.NewTicker(5 * time.Second)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-stop:
|
|
return
|
|
case <-ticker.C:
|
|
m.kickShadowWorker()
|
|
}
|
|
}
|
|
}()
|
|
m.kickShadowWorker()
|
|
}
|
|
|
|
// DrainSemanticShadowTasks processes queued shadow jobs synchronously (tests / recovery).
|
|
func (m *ContextManager) DrainSemanticShadowTasks(ctx context.Context) int {
|
|
if m == nil {
|
|
return 0
|
|
}
|
|
n := 0
|
|
for {
|
|
job, ok := m.claimNextShadowJob(ctx)
|
|
if !ok {
|
|
return n
|
|
}
|
|
m.runShadowJob(ctx, job)
|
|
n++
|
|
}
|
|
}
|
|
|
|
func (m *ContextManager) shadowWorkerLoop() {
|
|
defer func() {
|
|
m.mu.Lock()
|
|
m.shadowWorkerRunning = false
|
|
m.mu.Unlock()
|
|
}()
|
|
ctx := context.Background()
|
|
for {
|
|
job, ok := m.claimNextShadowJob(ctx)
|
|
if !ok {
|
|
return
|
|
}
|
|
m.runShadowJob(ctx, job)
|
|
}
|
|
}
|
|
|
|
func (m *ContextManager) claimNextShadowJob(ctx context.Context) (SemanticShadowJob, bool) {
|
|
if m.pool != nil {
|
|
if job, ok := m.claimDurableShadowJob(ctx); ok {
|
|
return job, true
|
|
}
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if len(m.shadowJobs) == 0 {
|
|
return SemanticShadowJob{}, false
|
|
}
|
|
job := m.shadowJobs[0]
|
|
m.shadowJobs = m.shadowJobs[1:]
|
|
job.Status = "leased"
|
|
job.Attempt++
|
|
return job, true
|
|
}
|
|
|
|
func (m *ContextManager) claimDurableShadowJob(ctx context.Context) (SemanticShadowJob, bool) {
|
|
owner := fmt.Sprintf("shadow-%d", time.Now().UnixNano())
|
|
tx, err := m.pool.Begin(ctx)
|
|
if err != nil {
|
|
return SemanticShadowJob{}, false
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
|
|
var (
|
|
job SemanticShadowJob
|
|
payload []byte
|
|
)
|
|
err = tx.QueryRow(ctx, `
|
|
SELECT id, session_id, summary_revision, input_digest, visibility_digest,
|
|
watermark_to_message_id, prior_revision, payload, attempt
|
|
FROM chat_summary_shadow_jobs
|
|
WHERE status IN ('pending','leased')
|
|
AND (lease_until IS NULL OR lease_until < now())
|
|
ORDER BY updated_at, id
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1
|
|
`).Scan(
|
|
&job.ID, &job.SessionID, &job.SummaryRevision, &job.InputDigest, &job.VisibilityDigest,
|
|
&job.WatermarkToMessageID, &job.PriorRevision, &payload, &job.Attempt,
|
|
)
|
|
if err != nil {
|
|
return SemanticShadowJob{}, false
|
|
}
|
|
job.Attempt++
|
|
var body shadowJobPayload
|
|
_ = json.Unmarshal(payload, &body)
|
|
job.Messages = body.Messages
|
|
job.PriorNarrative = body.PriorNarrative
|
|
job.PriorFacts = body.PriorFacts
|
|
job.Status = "leased"
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE chat_summary_shadow_jobs
|
|
SET status='leased', attempt=$2, lease_owner=$3, lease_until=now()+$4::interval, updated_at=now()
|
|
WHERE id=$1
|
|
`, job.ID, job.Attempt, owner, shadowLeaseTTL.String())
|
|
if err != nil {
|
|
return SemanticShadowJob{}, false
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return SemanticShadowJob{}, false
|
|
}
|
|
return job, true
|
|
}
|
|
|
|
func (m *ContextManager) runShadowJob(ctx context.Context, job SemanticShadowJob) {
|
|
report := &SemanticShadowReport{
|
|
Status: ShadowStatusFallback,
|
|
InputDigest: job.InputDigest,
|
|
ToolPermissions: SemanticToolPermissions,
|
|
GeneratedAt: time.Now(),
|
|
}
|
|
summarizer := m.semanticSummarizerOrDefault()
|
|
res, err := summarizer.Summarize(ctx, SemanticSummaryRequest{
|
|
SessionID: job.SessionID,
|
|
Messages: job.Messages,
|
|
InputDigest: job.InputDigest,
|
|
VisibilityDigest: job.VisibilityDigest,
|
|
PriorNarrative: job.PriorNarrative,
|
|
PriorFacts: job.PriorFacts,
|
|
WatermarkToMessageID: job.WatermarkToMessageID,
|
|
ExpectedRevision: job.SummaryRevision,
|
|
})
|
|
if err != nil || res == nil || strings.TrimSpace(res.Narrative) == "" {
|
|
reason := "empty_or_error"
|
|
if err != nil {
|
|
reason = "model_error"
|
|
}
|
|
report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: %s; deterministic retained at %s", reason, time.Now().UTC().Format(time.RFC3339))
|
|
_ = m.persistShadowProvenance(ctx, job, report, err)
|
|
return
|
|
}
|
|
if verr := ValidateShadowResult(job, res); verr != nil {
|
|
report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: validation_failed (%v); deterministic retained at %s", verr, time.Now().UTC().Format(time.RFC3339))
|
|
_ = m.persistShadowProvenance(ctx, job, report, verr)
|
|
return
|
|
}
|
|
if err := m.validateShadowAgainstCurrent(ctx, job); err != nil {
|
|
report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: stale (%v); deterministic retained at %s", err, time.Now().UTC().Format(time.RFC3339))
|
|
_ = m.persistShadowProvenance(ctx, job, report, err)
|
|
return
|
|
}
|
|
report.Status = ShadowStatusOK
|
|
report.Narrative = res.Narrative
|
|
report.Model = res.Model
|
|
report.Protocol = res.Protocol
|
|
if report.Protocol == "" {
|
|
report.Protocol = semanticProtocolV1
|
|
}
|
|
// Explicitly do NOT swap Narrative even if REPLACE env is true.
|
|
_ = SemanticSummaryReplaceEnabled()
|
|
_ = m.persistShadowProvenance(ctx, job, report, nil)
|
|
}
|
|
|
|
func (m *ContextManager) validateShadowAgainstCurrent(ctx context.Context, job SemanticShadowJob) error {
|
|
current, err := m.GetChatSummary(ctx, job.SessionID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if current == nil {
|
|
return fmt.Errorf("summary missing")
|
|
}
|
|
return ValidateShadowWatermark(job, current)
|
|
}
|
|
|
|
// ValidateShadowWatermark rejects shadow results that no longer match the served summary.
|
|
func ValidateShadowWatermark(job SemanticShadowJob, current *ChatSummary) error {
|
|
if current == nil {
|
|
return fmt.Errorf("summary missing")
|
|
}
|
|
if current.Revision != job.SummaryRevision {
|
|
return fmt.Errorf("revision mismatch want %d got %d", job.SummaryRevision, current.Revision)
|
|
}
|
|
if current.VisibilityDigest != "" && job.VisibilityDigest != "" && current.VisibilityDigest != job.VisibilityDigest {
|
|
return fmt.Errorf("visibility digest mismatch")
|
|
}
|
|
if job.WatermarkToMessageID > 0 && current.ToMessageID > 0 && current.ToMessageID != job.WatermarkToMessageID {
|
|
return fmt.Errorf("watermark mismatch want %d got %d", job.WatermarkToMessageID, current.ToMessageID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ValidateShadowResult checks optional echoed digests/watermarks from the summarizer.
|
|
func ValidateShadowResult(job SemanticShadowJob, res *SemanticSummaryResult) error {
|
|
if res == nil {
|
|
return fmt.Errorf("nil result")
|
|
}
|
|
if strings.TrimSpace(res.Narrative) == "" {
|
|
return fmt.Errorf("empty narrative")
|
|
}
|
|
if res.InputDigest != "" && res.InputDigest != job.InputDigest {
|
|
return fmt.Errorf("input digest mismatch")
|
|
}
|
|
if res.WatermarkToMessageID > 0 && job.WatermarkToMessageID > 0 && res.WatermarkToMessageID != job.WatermarkToMessageID {
|
|
return fmt.Errorf("result watermark mismatch")
|
|
}
|
|
if res.VisibilityDigest != "" && job.VisibilityDigest != "" && res.VisibilityDigest != job.VisibilityDigest {
|
|
return fmt.Errorf("result visibility digest mismatch")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *ContextManager) persistShadowProvenance(ctx context.Context, job SemanticShadowJob, report *SemanticShadowReport, runErr error) error {
|
|
current, err := m.GetChatSummary(ctx, job.SessionID)
|
|
if err != nil {
|
|
m.finishShadowJob(ctx, job, runErr)
|
|
return err
|
|
}
|
|
if current == nil {
|
|
m.finishShadowJob(ctx, job, fmt.Errorf("summary missing"))
|
|
return fmt.Errorf("summary missing")
|
|
}
|
|
if err := ValidateShadowWatermark(job, current); err != nil {
|
|
m.finishShadowJob(ctx, job, err)
|
|
return err
|
|
}
|
|
current.Shadow = report
|
|
if report != nil && report.FreshnessNote != "" {
|
|
current.FreshnessNote = report.FreshnessNote
|
|
} else if report != nil && report.Status == ShadowStatusOK {
|
|
current.FreshnessNote = ""
|
|
}
|
|
if report != nil && report.Model != "" {
|
|
current.ModelVersion = report.Model
|
|
}
|
|
if report != nil && report.Protocol != "" {
|
|
current.ProtocolVersion = report.Protocol
|
|
}
|
|
current.Source = "deterministic"
|
|
saveErr := m.SaveChatSummary(ctx, current)
|
|
finishErr := runErr
|
|
if saveErr != nil && finishErr == nil {
|
|
finishErr = saveErr
|
|
}
|
|
m.finishShadowJob(ctx, job, finishErr)
|
|
return saveErr
|
|
}
|
|
|
|
func (m *ContextManager) finishShadowJob(ctx context.Context, job SemanticShadowJob, runErr error) {
|
|
if m.pool == nil || job.ID == 0 {
|
|
return
|
|
}
|
|
status := "done"
|
|
errText := ""
|
|
if runErr != nil {
|
|
errText = runErr.Error()
|
|
if job.Attempt >= shadowJobMaxAttempt {
|
|
status = "failed"
|
|
} else {
|
|
status = "pending"
|
|
}
|
|
}
|
|
_, _ = m.pool.Exec(ctx, `
|
|
UPDATE chat_summary_shadow_jobs
|
|
SET status=$2, last_error=$3, lease_owner='', lease_until=NULL, updated_at=now()
|
|
WHERE id=$1
|
|
`, job.ID, status, errText)
|
|
}
|
|
|
|
// PendingShadowJobCount returns in-process queued jobs (tests).
|
|
func (m *ContextManager) PendingShadowJobCount() int {
|
|
if m == nil {
|
|
return 0
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
return len(m.shadowJobs)
|
|
}
|
|
|
|
// WaitShadowIdle waits until the in-process worker is idle (tests).
|
|
func (m *ContextManager) WaitShadowIdle(timeout time.Duration) bool {
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
m.mu.RLock()
|
|
running := m.shadowWorkerRunning
|
|
pending := len(m.shadowJobs)
|
|
m.mu.RUnlock()
|
|
if !running && pending == 0 {
|
|
return true
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
return false
|
|
}
|