第十四章：Event/Streaming 架构与 AG-UI 协议

package streaming

import "time"

// Event 是事件流中的基本单位
type Event struct {
    ID        string      `json:"id"`
    Type      EventType   `json:"type"`
    Timestamp time.Time   `json:"timestamp"`
    RunID     string      `json:"run_id"`
    NodeID    string      `json:"node_id,omitempty"`
    Payload   interface{} `json:"payload"`
}

// EventType 枚举所有事件类型
type EventType string

const (
    EventRunStarted       EventType = "run.started"
    EventRunCompleted     EventType = "run.completed"
    EventRunFailed        EventType = "run.failed"
    EventMessageDelta     EventType = "message.delta"
    EventMessageCompleted EventType = "message.completed"
    EventToolCallStarted  EventType = "tool_call.started"
    EventToolCallCompleted EventType = "tool_call.completed"
    EventToolResult       EventType = "tool.result"
    EventStateSnapshot    EventType = "state.snapshot"
    EventStateDelta       EventType = "state.delta"
    EventThoughtDelta     EventType = "thought.delta"
)

// EventChannel 是事件的传输管道
type EventChannel struct {
    ch     chan Event
    done   chan struct{}
    closed bool
}

// NewEventChannel 创建带缓冲的事件通道
func NewEventChannel(bufferSize int) *EventChannel {
    return &EventChannel{
        ch:   make(chan Event, bufferSize),
        done: make(chan struct{}),
    }
}

// Emit 发送一个事件到通道
func (ec *EventChannel) Emit(evt Event) error {
    select {
    case ec.ch <- evt:
        return nil
    case <-ec.done:
        return ErrChannelClosed
    }
}

// Subscribe 返回只读事件流
func (ec *EventChannel) Subscribe() <-chan Event {
    return ec.ch
}

// Close 关闭事件通道，通知所有订阅者
func (ec *EventChannel) Close() {
    if !ec.closed {
        ec.closed = true
        close(ec.done)
        close(ec.ch)
    }
}

package agent

import (
    "context"
    "fmt"
    "time"

    "github.com/google/uuid"
)

// Run 执行 Agent 并通过 EventChannel 流式输出
func (a *Agent) Run(ctx context.Context, input string, ec *EventChannel) error {
    runID := uuid.New().String()

    // 1. 发射运行开始事件
    ec.Emit(Event{
        ID:        uuid.New().String(),
        Type:      EventRunStarted,
        Timestamp: time.Now(),
        RunID:     runID,
        Payload:   map[string]string{"input": input},
    })

    // 2. 调用 LLM（流式）
    llmStream, err := a.llm.StreamChat(ctx, a.buildMessages(input))
    if err != nil {
        ec.Emit(Event{
            ID:      uuid.New().String(),
            Type:    EventRunFailed,
            RunID:   runID,
            Payload: map[string]string{"error": err.Error()},
        })
        return err
    }

    // 3. 逐 token 转发 LLM 输出为事件
    var fullContent string
    for chunk := range llmStream {
        if chunk.Content != "" {
            fullContent += chunk.Content
            ec.Emit(Event{
                ID:      uuid.New().String(),
                Type:    EventMessageDelta,
                RunID:   runID,
                Payload: map[string]string{"delta": chunk.Content},
            })
        }

        // 4. 检测工具调用
        if chunk.ToolCall != nil {
            ec.Emit(Event{
                ID:      uuid.New().String(),
                Type:    EventToolCallStarted,
                RunID:   runID,
                Payload: chunk.ToolCall,
            })

            // 执行工具
            result, toolErr := a.executeTool(ctx, chunk.ToolCall)

            ec.Emit(Event{
                ID:      uuid.New().String(),
                Type:    EventToolResult,
                RunID:   runID,
                Payload: map[string]interface{}{
                    "tool":   chunk.ToolCall.Name,
                    "result": result,
                    "error":  toolErr,
                },
            })
        }
    }

    // 5. 发射完成事件
    ec.Emit(Event{
        ID:      uuid.New().String(),
        Type:    EventMessageCompleted,
        RunID:   runID,
        Payload: map[string]string{"content": fullContent},
    })

    ec.Emit(Event{
        ID:      uuid.New().String(),
        Type:    EventRunCompleted,
        RunID:   runID,
        Payload: map[string]interface{}{"tokens_used": llmStream.Usage()},
    })

    return nil
}

package handler

import (
    "encoding/json"
    "fmt"
    "net/http"
)

// StreamHandler 将 EventChannel 桥接到 HTTP SSE
func StreamHandler(w http.ResponseWriter, r *http.Request) {
    // 设置 SSE 响应头
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 缓冲

    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming not supported", http.StatusInternalServerError)
        return
    }

    // 创建事件通道
    ec := NewEventChannel(64)
    defer ec.Close()

    // 后台启动 Agent 执行
    go func() {
        input := r.URL.Query().Get("input")
        agent := GetAgent(r.Context())
        if err := agent.Run(r.Context(), input, ec); err != nil {
            // Agent 执行失败，错误已通过事件通道发送
        }
        ec.Close() // Agent 完成，关闭通道
    }()

    // 主循环：从 channel 读事件，写 SSE
    for evt := range ec.Subscribe() {
        data, _ := json.Marshal(evt)
        fmt.Fprintf(w, "id: %s\n", evt.ID)
        fmt.Fprintf(w, "event: %s\n", evt.Type)
        fmt.Fprintf(w, "data: %s\n\n", data)
        flusher.Flush() // 立即推送，不等缓冲区满
    }
}

package streaming

import "sync"

// FanIn 将多个 EventChannel 合并为一个
// 类似"多车道合并为一车道"
func FanIn(channels ...*EventChannel) *EventChannel {
    merged := NewEventChannel(128)
    var wg sync.WaitGroup

    for _, ch := range channels {
        wg.Add(1)
        go func(c *EventChannel) {
            defer wg.Done()
            for evt := range c.Subscribe() {
                merged.Emit(evt)
            }
        }(ch)
    }

    // 所有子 channel 关闭后，关闭合并 channel
    go func() {
        wg.Wait()
        merged.Close()
    }()

    return merged
}

package observability

import (
    "context"
    "time"

    "go.opentelemetry.io/otel/trace"
)

// EventTracer 从事件流构建 OpenTelemetry Trace
type EventTracer struct {
    tracer   trace.Tracer
    spans    map[string]trace.Span  // eventID -> span
    rootSpan trace.Span
}

func NewEventTracer(tracer trace.Tracer) *EventTracer {
    return &EventTracer{
        tracer: tracer,
        spans:  make(map[string]trace.Span),
    }
}

// ProcessEvent 处理每个事件，自动构建 Span 树
func (et *EventTracer) ProcessEvent(ctx context.Context, evt Event) {
    switch evt.Type {
    case EventRunStarted:
        // 创建根 Span
        ctx, span := et.tracer.Start(ctx, "agent.run",
            trace.WithAttributes(
                attribute.String("run_id", evt.RunID),
            ),
        )
        et.rootSpan = span
        et.spans[evt.ID] = span

    case EventToolCallStarted:
        // 在根 Span 下创建子 Span
        payload := evt.Payload.(map[string]string)
        _, span := et.tracer.Start(ctx, "tool."+payload["name"],
            trace.WithAttributes(
                attribute.String("tool.name", payload["name"]),
            ),
        )
        et.spans[evt.ID] = span

    case EventToolResult:
        // 结束工具 Span
        if span, ok := et.spans[evt.ID]; ok {
            span.End()
        }

    case EventRunCompleted:
        // 结束根 Span，附加统计信息
        payload := evt.Payload.(map[string]interface{})
        et.rootSpan.SetAttributes(
            attribute.Int("tokens.total", payload["tokens_used"].(int)),
        )
        et.rootSpan.End()

    case EventRunFailed:
        // 记录错误并结束
        payload := evt.Payload.(map[string]string)
        et.rootSpan.RecordError(fmt.Errorf(payload["error"]))
        et.rootSpan.SetStatus(codes.Error, payload["error"])
        et.rootSpan.End()
    }
}

package metrics

import (
    "time"

    "go.opentelemetry.io/otel/metric"
)

// EventMetricsCollector 从事件流中提取关键指标
type EventMetricsCollector struct {
    // 计数器
    runCounter     metric.Int64Counter
    toolCallCounter metric.Int64Counter
    errorCounter   metric.Int64Counter

    // 直方图（延迟分布）
    runDuration     metric.Float64Histogram
    toolDuration    metric.Float64Histogram
    firstTokenDelay metric.Float64Histogram

    // 临时状态
    runStartTimes  map[string]time.Time
    toolStartTimes map[string]time.Time
}

func (mc *EventMetricsCollector) ProcessEvent(evt Event) {
    switch evt.Type {
    case EventRunStarted:
        mc.runCounter.Add(context.Background(), 1)
        mc.runStartTimes[evt.RunID] = evt.Timestamp

    case EventMessageDelta:
        // 计算 First Token 延迟
        if startTime, ok := mc.runStartTimes[evt.RunID]; ok {
            delay := evt.Timestamp.Sub(startTime).Seconds()
            mc.firstTokenDelay.Record(context.Background(), delay)
            delete(mc.runStartTimes, evt.RunID) // 只记第一个 token
        }

    case EventToolCallStarted:
        mc.toolCallCounter.Add(context.Background(), 1)
        mc.toolStartTimes[evt.ID] = evt.Timestamp

    case EventToolResult:
        if startTime, ok := mc.toolStartTimes[evt.ID]; ok {
            duration := evt.Timestamp.Sub(startTime).Seconds()
            mc.toolDuration.Record(context.Background(), duration)
        }

    case EventRunCompleted:
        if startTime, ok := mc.runStartTimes[evt.RunID]; ok {
            duration := evt.Timestamp.Sub(startTime).Seconds()
            mc.runDuration.Record(context.Background(), duration)
        }

    case EventRunFailed:
        mc.errorCounter.Add(context.Background(), 1)
    }
}

// 计算 Streaming Gap 分布
type GapAnalyzer struct {
    lastEventTime time.Time
    gaps          []time.Duration
}

func (ga *GapAnalyzer) ProcessEvent(evt Event) {
    if !ga.lastEventTime.IsZero() {
        gap := evt.Timestamp.Sub(ga.lastEventTime)
        ga.gaps = append(ga.gaps, gap)

        // 告警：如果某个 gap > 5s，可能有问题
        if gap > 5*time.Second {
            log.Warn("Large streaming gap detected",
                "gap", gap,
                "after_event", evt.Type,
                "run_id", evt.RunID,
            )
        }
    }
    ga.lastEventTime = evt.Timestamp
}

package debug

import (
    "encoding/json"
    "os"
)

// EventRecorder 录制事件流，用于事后回放
type EventRecorder struct {
    file *os.File
}

func NewEventRecorder(path string) (*EventRecorder, error) {
    f, err := os.Create(path)
    if err != nil {
        return nil, err
    }
    return &EventRecorder{file: f}, nil
}

// Record 逐行写入事件（JSONL 格式）
func (er *EventRecorder) Record(evt Event) error {
    data, err := json.Marshal(evt)
    if err != nil {
        return err
    }
    _, err = er.file.Write(append(data, '\n'))
    return err
}

// Replay 从文件回放事件流
func Replay(path string, handler func(Event)) error {
    f, err := os.Open(path)
    if err != nil {
        return err
    }
    defer f.Close()

    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        var evt Event
        if err := json.Unmarshal(scanner.Bytes(), &evt); err != nil {
            continue
        }
        handler(evt)
    }
    return scanner.Err()
}

// buffer 大小的选择
// 太小（如 1）：频繁阻塞，降低吞吐
// 太大（如 10000）：内存浪费，延迟感知不准
// 经验值：64-256 对大多数 Agent 场景够用
ec := NewEventChannel(128)

// 带丢弃策略的事件发送
func (ec *EventChannel) EmitWithDrop(evt Event) bool {
    select {
    case ec.ch <- evt:
        return true // 成功发送
    default:
        // buffer 满了，丢弃这个事件（但记录日志）
        log.Warn("Event dropped due to backpressure",
            "type", evt.Type, "run_id", evt.RunID)
        return false
    }
}

// 服务端：支持从指定位置续传
func StreamHandler(w http.ResponseWriter, r *http.Request) {
    // 检查客户端是否在续传
    lastEventID := r.Header.Get("Last-Event-ID")

    // 如果有 lastEventID，从事件日志中找到断点
    if lastEventID != "" {
        // 从存储中获取该 ID 之后的所有事件
        missedEvents := eventStore.GetAfter(lastEventID)
        for _, evt := range missedEvents {
            writeSSEEvent(w, evt)
        }
        flusher.Flush()
    }

    // 继续正常流式输出后续事件
    // ...
}

// 多租户事件路由
type EventRouter struct {
    channels map[string]*EventChannel // tenantID -> channel
    mu       sync.RWMutex
}

func (er *EventRouter) GetChannel(tenantID string) *EventChannel {
    er.mu.RLock()
    ch, ok := er.channels[tenantID]
    er.mu.RUnlock()

    if !ok {
        er.mu.Lock()
        ch = NewEventChannel(128)
        er.channels[tenantID] = ch
        er.mu.Unlock()
    }
    return ch
}

// 确保事件只发给正确的租户
func (er *EventRouter) Route(evt Event, tenantID string) {
    ch := er.GetChannel(tenantID)
    ch.Emit(evt)
}

// 策略1：预分配 buffer，避免反复 GC
var bufPool = sync.Pool{
    New: func() interface{} {
        return bytes.NewBuffer(make([]byte, 0, 512))
    },
}

func marshalEvent(evt Event) []byte {
    buf := bufPool.Get().(*bytes.Buffer)
    defer bufPool.Put(buf)
    buf.Reset()

    // 使用 json.NewEncoder 避免额外内存分配
    enc := json.NewEncoder(buf)
    enc.Encode(evt)
    return buf.Bytes()
}

// 策略2：对高频事件使用手动序列化
func marshalMessageDelta(delta string) []byte {
    // 跳过 reflect，直接拼 JSON
    return []byte(`{"type":"message.delta","payload":{"delta":"` +
        jsonEscape(delta) + `"}}`)
}

func (a *Agent) Run(ctx context.Context, input string, ec *EventChannel) error {
    runID := uuid.New().String()
    ec.Emit(Event{Type: EventRunStarted, RunID: runID})

    // 监听 context 取消
    go func() {
        <-ctx.Done()
        // 发送取消事件，然后关闭
        ec.Emit(Event{
            Type:    EventRunFailed,
            RunID:   runID,
            Payload: map[string]string{"error": "cancelled by user"},
        })
        ec.Close()
    }()

    // ... 正常执行逻辑
    // 所有 LLM 调用和工具调用都传入 ctx
    // 当 ctx 被 cancel 时，它们会自行中断
}