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

package agentserver

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/gofiber/fiber/v2"
	"go.uber.org/zap"
)

// 需要解析的 JSON 响应最大读取字节数（事件列表最多 5000 条）
const maxInspectBytes = 64 << 20

// 透传给上游的请求头（Authorization 会被替换为 API Key，Cookie 等不转发）
var forwardRequestHeaders = []string{"Content-Type", "Accept", "Last-Event-ID"}

// 透传给浏览器的响应头
var forwardResponseHeaders = []string{
	"Content-Type", "Content-Disposition", "Cache-Control",
	"X-Content-Type-Options", "Content-Security-Policy",
}

// Proxy /agent/v1/* → AgentServer /v1/*（technical-design §3.5）
type Proxy struct {
	Client *Client
	Guard  *Guard
	Sink   ResourceSink
	// OnMessage 发送消息成功后回调（刷新会话 lastActive）
	OnMessage func(userId, sessionId string)
	// OnTitle 经代理改了标题后回调（同步归属表）
	OnTitle func(userId, sessionId, title string)
}

// WriteError 以 AgentServer 的错误格式返回，SDK 可以直接解析
func WriteError(c *fiber.Ctx, status int, code, message string) error {
	return c.Status(status).JSON(fiber.Map{"error": fiber.Map{"code": code, "message": message}})
}

// Handle 处理一次代理请求；rest 为 /agent/v1/ 之后的路径
func (p *Proxy) Handle(c *fiber.Ctx, userId, rest string) error {
	method := c.Method()
	path := "/" + strings.Trim(rest, "/")
	rawQuery := string(c.Request().URI().QueryString())
	query, _ := url.ParseQuery(rawQuery)

	route, ok := MatchRoute(method, path, query)
	if !ok {
		return WriteError(c, fiber.StatusForbidden, "forbidden", "不支持的接口")
	}

	sessionId, owned, err := p.Guard.Check(c.UserContext(), userId, route)
	if err != nil {
		zap.L().Warn("[Agent] 归属校验失败", zap.String("path", path), zap.Error(err))
		return WriteError(c, fiber.StatusBadGateway, "upstream_error", "AI 服务暂不可用")
	}
	if !owned {
		return WriteError(c, fiber.StatusNotFound, "not_found", "会话不存在或无权访问")
	}

	// 请求体：PATCH 只保留 title、model
	var body []byte
	var patchedTitle *string
	if route.Action == ActionSessionPatch {
		body, patchedTitle, err = filterSessionPatch(c.Body())
		if err != nil {
			return WriteError(c, fiber.StatusBadRequest, "bad_request", "请求体不是有效的 JSON")
		}
	} else if len(c.Body()) > 0 {
		body = append([]byte(nil), c.Body()...)
	}

	header := http.Header{}
	for _, h := range forwardRequestHeaders {
		if v := c.Get(h); v != "" {
			header.Set(h, v)
		}
	}
	if route.Action == ActionSessionPatch {
		header.Set("Content-Type", "application/json")
	}

	// 上游请求的生命周期：流式/透传响应在处理函数返回后才读取完，由读取方在结束时取消
	ctx, cancel := context.WithCancel(context.Background())
	var reqBody io.Reader
	if body != nil {
		reqBody = bytes.NewReader(body)
	}
	resp, err := p.Client.Forward(ctx, method, "/v1"+path, rawQuery, header, reqBody)
	if err != nil {
		cancel()
		zap.L().Warn("[Agent] 转发失败", zap.String("path", path), zap.Error(err))
		return WriteError(c, fiber.StatusBadGateway, "upstream_error", "AI 服务暂不可用")
	}

	for _, h := range forwardResponseHeaders {
		if v := resp.Header.Get(h); v != "" {
			c.Set(h, v)
		}
	}
	c.Status(resp.StatusCode)

	switch {
	case route.Action == ActionStream && resp.StatusCode == http.StatusOK:
		p.stream(c, resp, cancel, sessionId)
		return nil

	case needsInspect(route.Action) && resp.StatusCode < 300:
		defer cancel()
		defer resp.Body.Close()
		data, err := io.ReadAll(io.LimitReader(resp.Body, maxInspectBytes))
		if err != nil {
			return WriteError(c, fiber.StatusBadGateway, "upstream_error", "读取 AI 服务响应失败")
		}
		if rs := ResourcesFromResponse(route.Action, sessionId, data); len(rs) > 0 && p.Sink != nil {
			p.Sink.Record(rs)
		}
		if route.Action == ActionSendMessage && p.OnMessage != nil {
			p.OnMessage(userId, sessionId)
		}
		if route.Action == ActionSessionPatch && patchedTitle != nil && p.OnTitle != nil {
			p.OnTitle(userId, sessionId, *patchedTitle)
		}
		return c.Send(data)

	default:
		// 原样透传（文件下载等），读取结束时关闭上游连接
		c.Context().SetBodyStream(&cancelOnClose{ReadCloser: resp.Body, cancel: cancel}, int(resp.ContentLength))
		return nil
	}
}

func needsInspect(a Action) bool {
	switch a {
	case ActionEvents, ActionSessionDetail, ActionSessionPatch, ActionSendMessage,
		ActionUploadFiles, ActionListFiles, ActionInteractions:
		return true
	}
	return false
}

// stream SSE 透传：逐块写出并立即 Flush；旁路解析事件记录子资源；客户端断开时取消上游请求
func (p *Proxy) stream(c *fiber.Ctx, resp *http.Response, cancel context.CancelFunc, sessionId string) {
	c.Set("Content-Type", "text/event-stream; charset=utf-8")
	c.Set("Cache-Control", "no-cache, no-transform")
	c.Set("Connection", "keep-alive")
	c.Set("X-Accel-Buffering", "no")

	scanner := NewSSEScanner(func(data []byte) {
		var rs []Resource
		for _, r := range ResourcesFromSSEData(data) {
			if r.SessionId == sessionId {
				rs = append(rs, r)
			}
		}
		if len(rs) > 0 && p.Sink != nil {
			p.Sink.Record(rs)
		}
	})

	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
		defer cancel()
		defer resp.Body.Close()
		buf := make([]byte, 32<<10)
		for {
			n, err := resp.Body.Read(buf)
			if n > 0 {
				chunk := buf[:n]
				scanner.Write(chunk)
				if _, werr := w.Write(chunk); werr != nil {
					return
				}
				if ferr := w.Flush(); ferr != nil {
					// 客户端已断开
					return
				}
			}
			if err != nil {
				return
			}
		}
	})
}

// filterSessionPatch 只保留 title、model 字段
func filterSessionPatch(raw []byte) ([]byte, *string, error) {
	in := map[string]json.RawMessage{}
	if len(bytes.TrimSpace(raw)) > 0 {
		if err := json.Unmarshal(raw, &in); err != nil {
			return nil, nil, err
		}
	}
	out := map[string]json.RawMessage{}
	var title *string
	for _, k := range []string{"title", "model"} {
		if v, ok := in[k]; ok {
			out[k] = v
		}
	}
	if v, ok := out["title"]; ok {
		var t string
		if json.Unmarshal(v, &t) == nil {
			title = &t
		}
	}
	b, err := json.Marshal(out)
	return b, title, err
}

// cancelOnClose 透传响应体读取结束（fasthttp 调用 Close）时取消上游请求
type cancelOnClose struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (r *cancelOnClose) Close() error {
	err := r.ReadCloser.Close()
	r.cancel()
	return err
}