Agent-System-Backend/internal/pgbus/pgbus.go

// Package pgbus 是 PG 档的事件跨实例扇出（阶段 3 / 协作存储 U6 切片一）：
// 落盘事务提交后 NOTIFY agent_events '<conv>|<lastSeq>'；每实例一个 LISTEN 连接，
// 收到通知按会话游标从库分页追赶转发本地 bus（SSE 以 seq 去重，重投无害）。
// 多实例部署契约：run 在其归属实例内存中执行，会话流量需按会话亲和路由
// （网关/客户端择点）；本包只解决“其他实例的订阅者”可见性。
//
// U6（KTD6）修复的两个既有缺口：
//   - 批内遗漏：通知只带末 seq，旧实现 EventAt 只取末条——批内中间事件对
//     实时订阅者不可见。现按会话游标 EventsAfter 分页追赶，末 seq 即水位；
//   - 断线假重连：旧实现对原连接继续 WaitForNotification。现真正重建
//     连接→LISTEN→按全部已知会话追赶，另有周期补查覆盖通知丢失。
package pgbus

import (
	"context"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/bus"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/store"
	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"

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

const channel = "agent_events"

// catchUpPage 追赶分页大小（验收含 1000+ 事件分页可见）。
const catchUpPage = 500

// rescueEvery 周期补查间隔：覆盖 NOTIFY 丢失与监听窗口前后竞争（KTD6）。
const rescueEvery = 2 * time.Second

// Notifier 落盘后通知（任意连接池连接可用）。
type Notifier struct{ st *store.Store }

func NewNotifier(st *store.Store) *Notifier { return &Notifier{st: st} }

// Notify 在 AppendEvents 成功后调用：一批事件一条通知（载荷=会话|末 seq）。
func (n *Notifier) Notify(conversationID string, evs []core.StreamEvent) {
	if len(evs) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	payload := conversationID + "|" + strconv.FormatInt(evs[len(evs)-1].Seq, 10)
	if _, err := n.st.NotifyPG(ctx, channel, payload); err != nil {
		log.Printf("[pgbus] notify %s: %v", conversationID, err)
	}
}

// Listener 每实例一个：LISTEN → 按会话游标追赶 → 本地 bus 转发。
type Listener struct {
	conn *pgx.Conn
	dsn  string
	st   *store.Store
	bus  *bus.Bus
	done chan struct{}

	mu     sync.Mutex
	cursor map[string]int64 // conv → 已转发水位（重启清零：SSE 侧重放兜底）
}

// Start 建立监听连接并开始转发循环；返回后即可用，随 ctx 取消停机。
func Start(ctx context.Context, dsn string, st *store.Store, b *bus.Bus) (*Listener, error) {
	l := &Listener{dsn: dsn, st: st, bus: b, done: make(chan struct{}), cursor: map[string]int64{}}
	if err := l.connect(ctx); err != nil {
		return nil, err
	}
	go l.loop(ctx)
	return l, nil
}

func (l *Listener) connect(ctx context.Context) error {
	conn, err := pgx.Connect(ctx, l.dsn)
	if err != nil {
		return err
	}
	if _, err := conn.Exec(ctx, "LISTEN "+channel); err != nil {
		conn.Close(ctx)
		return err
	}
	l.conn = conn
	return nil
}

func (l *Listener) loop(ctx context.Context) {
	defer close(l.done)
	// 周期补查独立 goroutine：通知丢失场景恰恰没有通知可触发——不能只在
	// 收到通知后顺带补查（KTD6）。
	go l.rescueLoop(ctx)
	for {
		n, err := l.conn.WaitForNotification(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return // 停机
			}
			// 真重连（U6）：旧实现对原连接继续 Wait——断线后永不恢复。
			// 退避重建 → LISTEN → 按已知会话追赶（补断线窗口内落盘的事件）。
			log.Printf("[pgbus] wait: %v（重连中）", err)
			l.reconnect(ctx)
			continue
		}
		conv, seq, ok := parsePayload(n.Payload)
		if !ok {
			continue
		}
		l.catchUp(conv, seq)
	}
}

// rescueLoop 周期补查心跳：随 ctx 或主循环退出（done）停止。
func (l *Listener) rescueLoop(ctx context.Context) {
	t := time.NewTicker(rescueEvery)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-l.done:
			return
		case <-t.C:
			l.rescueKnown()
		}
	}
}

// reconnect 重建监听连接并立即追赶全部已知会话（有界退避重试）。
func (l *Listener) reconnect(ctx context.Context) {
	_ = l.conn.Close(context.Background())
	for attempt := 0; attempt < 3; attempt++ {
		cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
		err := l.connect(cctx)
		cancel()
		if err == nil {
			l.rescueKnown()
			return
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(time.Duration(attempt+1) * time.Second):
		}
	}
}

// rescueKnown 周期补查：不经通知，按各会话游标继续分页追赶（覆盖 NOTIFY
// 丢失、监听生效窗口前后竞争——KTD6「只有落库事件才是重放依据」）。
// 调用方：重连成功后与 loop 的 rescue 心跳（消息稀疏时由 ticker 兜底）。
func (l *Listener) rescueKnown() {
	l.mu.Lock()
	convs := make([]string, 0, len(l.cursor))
	for c := range l.cursor {
		convs = append(convs, c)
	}
	l.mu.Unlock()
	for _, conv := range convs {
		l.catchUp(conv, 1<<62)
	}
}

// catchUp 从会话游标之后分页取回并转发，直到取空或越过水位 toSeq。
// 首见会话从 0 开始（SSE 端按 seq 去重，重投无害；这保证批内中间事件可见）。
func (l *Listener) catchUp(conv string, toSeq int64) {
	l.mu.Lock()
	since := l.cursor[conv]
	l.mu.Unlock()
	for {
		evs, err := l.st.EventsAfter(conv, since, catchUpPage)
		if err != nil {
			log.Printf("[pgbus] catchup %s@%d: %v", conv, since, err)
			return
		}
		stop := false
		for _, ev := range evs {
			if ev.Seq > toSeq {
				stop = true
				break
			}
			l.bus.Publish(ev)
			since = ev.Seq
		}
		l.setCursor(conv, since)
		if stop || len(evs) < catchUpPage {
			return
		}
	}
}

func (l *Listener) setCursor(conv string, seq int64) {
	l.mu.Lock()
	if seq > l.cursor[conv] {
		l.cursor[conv] = seq
	}
	l.mu.Unlock()
}

// parsePayload 拆 '<conv>|<seq>'（会话 ID 不含 '|'）。
func parsePayload(p string) (string, int64, bool) {
	i := strings.LastIndex(p, "|")
	if i < 0 {
		return "", 0, false
	}
	seq, err := strconv.ParseInt(p[i+1:], 10, 64)
	if err != nil {
		return "", 0, false
	}
	return p[:i], seq, true
}

// Close 关闭监听连接。
func (l *Listener) Close(ctx context.Context) {
	l.conn.Close(ctx)
	select {
	case <-l.done:
	case <-time.After(2 * time.Second):
	}
}