FranklinWH-Toolbox/service/utils/agentserver/sse.go

package agentserver

import (
	"bytes"
	"encoding/json"
)

// Resource 会话下的子资源（run / 交互 / 文件），代理据此校验不带会话 ID 的请求
type Resource struct {
	Id        string
	SessionId string
}

// ResourcesFromEvent 从一条事件中提取子资源 ID：
//   - run.started 等带 runId 的事件 → runId
//   - interaction.required → payload.interactionId
//   - file.changed → payload.fileId
func ResourcesFromEvent(e AgentEvent) []Resource {
	if e.SessionId == "" {
		return nil
	}
	var out []Resource
	if e.Type == "run.started" && e.RunId != nil && *e.RunId != "" {
		out = append(out, Resource{Id: *e.RunId, SessionId: e.SessionId})
	}
	switch e.Type {
	case "interaction.required":
		var p struct {
			InteractionId string `json:"interactionId"`
		}
		if json.Unmarshal(e.Payload, &p) == nil && p.InteractionId != "" {
			out = append(out, Resource{Id: p.InteractionId, SessionId: e.SessionId})
		}
	case "file.changed":
		var p struct {
			FileId string `json:"fileId"`
		}
		if json.Unmarshal(e.Payload, &p) == nil && p.FileId != "" {
			out = append(out, Resource{Id: p.FileId, SessionId: e.SessionId})
		}
	}
	return out
}

// SSEScanner 增量解析 SSE 字节流，对每个完整事件的 data 回调一次。
// 只用于旁路提取资源 ID，不影响原样转发；单个事件超过 maxEventBytes 时丢弃该事件。
type SSEScanner struct {
	buf    []byte
	onData func(data []byte)
}

const maxEventBytes = 4 << 20

// NewSSEScanner 创建解析器
func NewSSEScanner(onData func(data []byte)) *SSEScanner {
	return &SSEScanner{onData: onData}
}

// Write 输入一段原始字节
func (s *SSEScanner) Write(p []byte) {
	s.buf = append(s.buf, p...)
	for {
		idx, sepLen := eventBoundary(s.buf)
		if idx < 0 {
			break
		}
		s.emit(s.buf[:idx])
		s.buf = s.buf[idx+sepLen:]
	}
	if len(s.buf) > maxEventBytes {
		s.buf = s.buf[:0]
	}
}

// eventBoundary 查找事件分隔（空行），兼容 \n\n 与 \r\n\r\n
func eventBoundary(b []byte) (int, int) {
	i := bytes.Index(b, []byte("\n\n"))
	j := bytes.Index(b, []byte("\r\n\r\n"))
	switch {
	case i < 0 && j < 0:
		return -1, 0
	case j < 0 || (i >= 0 && i < j):
		return i, 2
	default:
		return j, 4
	}
}

func (s *SSEScanner) emit(block []byte) {
	var data []byte
	for _, line := range bytes.Split(block, []byte("\n")) {
		line = bytes.TrimSuffix(line, []byte("\r"))
		if bytes.HasPrefix(line, []byte("data:")) {
			v := bytes.TrimPrefix(line, []byte("data:"))
			v = bytes.TrimPrefix(v, []byte(" "))
			if data != nil {
				data = append(data, '\n')
			}
			data = append(data, v...)
		}
	}
	if data != nil && s.onData != nil {
		s.onData(data)
	}
}

// ResourcesFromSSEData 解析一条 SSE data（事件 JSON）
func ResourcesFromSSEData(data []byte) []Resource {
	var e AgentEvent
	if err := json.Unmarshal(data, &e); err != nil {
		return nil
	}
	return ResourcesFromEvent(e)
}

// ResourcesFromResponse 从普通 JSON 响应中提取子资源 ID，sessionId 为路径中的会话
func ResourcesFromResponse(action Action, sessionId string, body []byte) []Resource {
	var out []Resource
	switch action {
	case ActionSendMessage:
		var r RunDto
		if json.Unmarshal(body, &r) == nil && r.Id != "" {
			out = append(out, Resource{Id: r.Id, SessionId: sessionId})
		}
	case ActionUploadFiles, ActionListFiles:
		var fs []FileDto
		if json.Unmarshal(body, &fs) == nil {
			for _, f := range fs {
				if f.Id != "" {
					out = append(out, Resource{Id: f.Id, SessionId: sessionId})
				}
			}
		}
	case ActionEvents:
		var page EventsPage
		if json.Unmarshal(body, &page) == nil {
			for _, e := range page.Items {
				out = append(out, ResourcesFromEvent(e)...)
			}
		}
	case ActionSessionDetail:
		var d SessionDetailDto
		if json.Unmarshal(body, &d) == nil {
			if d.ActiveRunId != nil && *d.ActiveRunId != "" {
				out = append(out, Resource{Id: *d.ActiveRunId, SessionId: sessionId})
			}
			for _, it := range d.PendingInteractions {
				if it.InteractionId != "" {
					out = append(out, Resource{Id: it.InteractionId, SessionId: sessionId})
				}
			}
		}
	case ActionInteractions:
		var items []InteractionRefDto
		if json.Unmarshal(body, &items) == nil {
			for _, it := range items {
				if it.InteractionId != "" {
					out = append(out, Resource{Id: it.InteractionId, SessionId: sessionId})
				}
			}
		}
	}
	// 只保留与路径会话一致的资源，防止伪造的事件写入他人会话的索引
	filtered := out[:0]
	for _, r := range out {
		if r.SessionId == sessionId {
			filtered = append(filtered, r)
		}
	}
	return filtered
}