Agent-System/internal/squad/process_bridge.go

package squad

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"chat-agent-system/internal/agent"
)

// 过程流常量：节流窗口 / 单帧上限 / 工具文本截断（设计文档 F.1）。
const (
	processFlushInterval = 200 * time.Millisecond
	processMaxFrameBytes = 8 * 1024
	processToolTrimBytes = 4 * 1024
)

// detailCategory 过程事件类别（与 stream.DetailPayload.Category 枚举对齐）。
const (
	detailCategoryAgentThought = "agent_thought"
	detailCategoryToolCall     = "tool_call"
)

// deltaStream 一个聚合项的增量缓冲（节流冲刷）。
// 只被 bridge 主路径使用；timer 回调仅触发 flush 任务。
type deltaStream struct {
	mu     sync.Mutex
	itemID string
	buf    strings.Builder
	timer  *time.Timer
	done   bool
}

func (d *deltaStream) append(text string) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.done || text == "" {
		return
	}
	d.buf.WriteString(text)
}

// takeIfReady 返回并清空缓冲；由外部节拍调用（非定时器自转）。
func (d *deltaStream) take() string {
	d.mu.Lock()
	defer d.mu.Unlock()
	out := d.buf.String()
	d.buf.Reset()
	return out
}

func (d *deltaStream) finish() {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.timer != nil {
		d.timer.Stop()
		d.timer = nil
	}
	d.done = true
}

// processBridge 把 agent.ProcessEvent 翻译为 stream.EventDetail 增量帧。
// 每个 step（小队成员的一次执行）独立持有一个 bridge；无跨 step 共享状态，并行安全。
type processBridge struct {
	sink      ProgressSink
	step      StepSpec
	friendly  string // 角色显示名
	flushTick *time.Ticker
	stopCh    chan struct{}
	stopOnce  sync.Once
	wg        sync.WaitGroup

	mu          sync.Mutex
	thoughtSeq  int
	toolSeq     int
	thought     *thoughtBlock // 当前思考块（nil = 尚未开始）
	currentTool *toolBlock
	emitted     bool // 是否已发出过事件（Stop 只在发过时才有意义）
}

// thoughtBlock 一个 message 周期内的思考聚合项。
type thoughtBlock struct {
	itemID string
	stream *deltaStream
	title  string
}

// toolBlock 一次工具调用的聚合项。
type toolBlock struct {
	itemID  string
	name    string
	started time.Time
	input   strings.Builder
}

// newProcessBridge 创建过程桥。返回 nil 表示该 step 无流式需求（sink 缺失），
// 调用方对 nil bridge 调用任何方法都是安全 no-op。
func newProcessBridge(sink ProgressSink, step StepSpec, friendlyRole string) *processBridge {
	if sink == nil {
		return nil
	}
	b := &processBridge{
		sink:      sink,
		step:     step,
		friendly: friendlyRole,
		stopCh:   make(chan struct{}),
	}
	// 定时冲刷循环：所有 delta 缓冲统一按 flushInterval 节拍冲刷。
	b.flushTick = time.NewTicker(processFlushInterval)
	b.wg.Add(1)
	go func() {
		defer b.wg.Done()
		for {
			select {
			case <-b.stopCh:
				return
			case <-b.flushTick.C:
				b.flushPending()
			}
		}
	}()
	return b
}

// Stop 结束 bridge：冲刷残余、定稿当前项、停止节拍循环。幂等（可安全多次调用）。
func (b *processBridge) Stop() {
	if b == nil {
		return
	}
	b.stopOnce.Do(func() {
		b.flushPending()
		b.mu.Lock()
		th := b.thought
		tl := b.currentTool
		b.mu.Unlock()
		if th != nil {
			b.finishThought(th)
		}
		if tl != nil {
			b.finishTool(tl, "succeeded", "")
		}
		close(b.stopCh)
		b.flushTick.Stop()
		b.wg.Wait()
	})
}

// Handle 把一个 agent.ProcessEvent 翻译为 detail 事件流。
// 调用方保证事件到达顺序与 agentx 循环一致（同 step 内串行）。
func (b *processBridge) Handle(ev agent.ProcessEvent) {
	if b == nil {
		return
	}
	select {
	case <-b.stopCh:
		return // Stop 后忽略（防御：Execute 返回后的残余事件）
	default:
	}
	switch ev.Type {
	case "reasoning_delta":
		b.onReasoningDelta(ev.Text)
	case "text_delta":
		// P0 决策（设计文档 C.4）：任务正文不发 detail——终态走既有 step_output
		// 通道，避免与最终 deliverable 重复推送大文本。
	case "tool_call_delta":
		b.onToolCallDelta(ev)
	case "tool_execution_start":
		b.onToolStart(ev.Tool)
	case "tool_execution_end":
		b.onToolEnd(ev.Tool, ev.Detail)
	case "message_end":
		b.onMessageEnd()
	}
}

func (b *processBridge) onReasoningDelta(text string) {
	b.mu.Lock()
	if b.thought == nil {
		b.thoughtSeq++
		b.thought = &thoughtBlock{
			itemID: fmt.Sprintf("thought#%s#%d", b.step.ID, b.thoughtSeq),
			stream: &deltaStream{},
			title:  fmt.Sprintf("%s思考过程 %d", b.friendly, b.thoughtSeq),
		}
		b.emitted = true
	}
	th := b.thought
	b.mu.Unlock()
	th.stream.append(text)
}

func (b *processBridge) onToolCallDelta(ev agent.ProcessEvent) {
	b.mu.Lock()
	t := b.currentTool
	if t == nil {
		// 模型在 tool_execution_start 前就吐参数增量：只预建内部块（不发 calling 帧，
		// 正式的 calling 帧由随后的 tool_execution_start 发出）。
		b.toolSeq++
		t = &toolBlock{
			itemID:  fmt.Sprintf("tool#%s#%d", b.step.ID, b.toolSeq),
			name:    ev.Tool,
			started: time.Now(),
		}
		b.currentTool = t
		b.emitted = true
	}
	b.mu.Unlock()
	t.input.WriteString(ev.Detail)
}

func (b *processBridge) onToolStart(name string) {
	b.mu.Lock()
	var cur *toolBlock
	if b.currentTool != nil {
		// 已由 tool_call_delta 预建（罕见时序）：复用同一聚合项，只补名与时间。
		cur = b.currentTool
		if cur.name == "" {
			cur.name = name
		}
		cur.started = time.Now()
	} else {
		b.toolSeq++
		cur = &toolBlock{
			itemID:  fmt.Sprintf("tool#%s#%d", b.step.ID, b.toolSeq),
			name:    name,
			started: time.Now(),
		}
		b.currentTool = cur
	}
	b.emitted = true
	itemID, toolName := cur.itemID, cur.name
	b.mu.Unlock()
	// tool_execution_start 是正式的"开始调用"信号：无论是否预建，都要发 calling 帧。
	b.emitDetail(&DetailEvent{
		Category:    detailCategoryToolCall,
		StepID:      b.step.ID,
		Role:        b.step.Role,
		ItemID:      itemID,
		Tool:        &ToolEvent{Name: toolName, Status: "calling"},
		Title:       fmt.Sprintf("%s调用工具 %s", b.friendly, toolName),
		Collapsible: true,
	})
}

func (b *processBridge) onToolEnd(name, result string) {
	b.mu.Lock()
	t := b.currentTool
	b.currentTool = nil
	b.mu.Unlock()
	if t == nil {
		return
	}
	status := "succeeded"
	if strings.Contains(strings.ToLower(result), "error") || strings.Contains(strings.ToLower(result), "failed") {
		status = "failed"
	}
	duration := time.Since(t.started).Milliseconds()
	info := &ToolEvent{
		Name:       name,
		Status:     status,
		Input:      trimProcessText(t.input.String()),
		Output:     trimProcessText(result),
		DurationMs: duration,
	}
	b.emitDetail(&DetailEvent{
		Category:    detailCategoryToolCall,
		StepID:      b.step.ID,
		Role:        b.step.Role,
		ItemID:      t.itemID,
		Tool:        info,
		Collapsible: true,
	})
	// 工具项即时定稿（状态帧 + done 标记合并语义：终帧即最后状态）
	b.emitDetail(&DetailEvent{
		Category:    detailCategoryToolCall,
		StepID:      b.step.ID,
		Role:        b.step.Role,
		ItemID:      t.itemID,
		Tool:        info,
		Done:        true,
		Collapsible: true,
	})
}

func (b *processBridge) onMessageEnd() {
	b.mu.Lock()
	th := b.thought
	b.thought = nil
	b.mu.Unlock()
	if th != nil {
		b.finishThought(th)
	}
}

// flushPending 冲刷当前思考块缓冲（节拍或 Stop 调用）。
func (b *processBridge) flushPending() {
	if b == nil {
		return
	}
	b.mu.Lock()
	th := b.thought
	b.mu.Unlock()
	if th == nil {
		return
	}
	text := th.stream.take()
	if text == "" {
		return
	}
	if len(text) > processMaxFrameBytes {
		text = text[:processMaxFrameBytes] + "\n…(截断)"
	}
	b.emitDetail(&DetailEvent{
		Category:    detailCategoryAgentThought,
		StepID:      b.step.ID,
		Role:        b.step.Role,
		ItemID:      th.itemID,
		Delta:       text,
		Collapsible: true,
	})
}

func (b *processBridge) finishThought(th *thoughtBlock) {
	th.stream.finish()
	// 冲残余
	if text := th.stream.take(); text != "" {
		b.emitDetail(&DetailEvent{
			Category:    detailCategoryAgentThought,
			StepID:      b.step.ID,
			Role:        b.step.Role,
			ItemID:      th.itemID,
			Delta:       text,
			Collapsible: true,
		})
	}
	b.emitDetail(&DetailEvent{
		Category:    detailCategoryAgentThought,
		StepID:      b.step.ID,
		Role:        b.step.Role,
		ItemID:      th.itemID,
		Title:       th.title,
		Done:        true,
		Collapsible: true,
	})
}

func (b *processBridge) finishTool(t *toolBlock, status, result string) {
	b.emitDetail(&DetailEvent{
		Category: detailCategoryToolCall,
		StepID:   b.step.ID,
		Role:     b.step.Role,
		ItemID:   t.itemID,
		Tool: &ToolEvent{
			Name:       t.name,
			Status:     status,
			Input:      trimProcessText(t.input.String()),
			Output:     trimProcessText(result),
			DurationMs: time.Since(t.started).Milliseconds(),
		},
		Done:        true,
		Collapsible: true,
	})
}

func (b *processBridge) emitDetail(p *DetailEvent) {
	if b == nil || b.sink == nil {
		return
	}
	b.sink.EmitDetail(*p)
}

// trimProcessText 工具入参/结果截断（防超大帧；保留首尾）。
func trimProcessText(s string) string {
	if len(s) <= processToolTrimBytes {
		return s
	}
	headN := processToolTrimBytes / 2
	tailN := processToolTrimBytes / 4
	head := s[:headN]
	tail := s[len(s)-tailN:]
	return head + "\n…(中间省略 " + fmt.Sprint(len(s)-headN-tailN) + " 字节)…\n" + tail
}

// AttachProcessBridge 给 agent.Request 装配过程桥（若 sink 缺失或无流式需求则原样返回）。
// 返回的 stop 函数在 Execute 返回后必须调用一次（定稿与冲刷残余）。
func AttachProcessBridge(sink ProgressSink, step StepSpec, req *agent.Request) (stop func()) {
	if sink == nil || req == nil {
		return func() {}
	}
	bridge := newProcessBridge(sink, step, FriendlyRoleName(step.Role))
	if bridge == nil {
		return func() {}
	}
	prev := req.OnEvent
	if prev == nil {
		req.OnEvent = bridge.Handle
	} else {
		req.OnEvent = func(ev agent.ProcessEvent) {
			prev(ev)
			bridge.Handle(ev)
		}
	}
	return bridge.Stop
}