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

package agentserver

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"github.com/gofiber/fiber/v2"
)

type lockedSink struct {
	mu  sync.Mutex
	got []Resource
}

func (s *lockedSink) Record(rs []Resource) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.got = append(s.got, rs...)
}

func (s *lockedSink) ids() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []string
	for _, r := range s.got {
		out = append(out, r.Id)
	}
	return out
}

// 假的 AgentServer：记录收到的鉴权头与请求体
func fakeAgentServer(t *testing.T, seen *http.Header, lastBody *string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		*seen = r.Header.Clone()
		b, _ := io.ReadAll(r.Body)
		*lastBody = string(b)
		switch {
		case r.URL.Path == "/v1/sessions/ses_a1/stream":
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			io.WriteString(w, "retry: 2000\n\nid: 1\nevent: run.started\ndata: "+
				`{"seq":1,"type":"run.started","sessionId":"ses_a1","runId":"run_s1","payload":{}}`+"\n\n")
		case r.URL.Path == "/v1/sessions/ses_a1/messages":
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(202)
			io.WriteString(w, `{"id":"run_m1","sessionId":"ses_a1","status":"queued"}`)
		case r.URL.Path == "/v1/sessions/ses_a1":
			w.Header().Set("Content-Type", "application/json")
			io.WriteString(w, `{"id":"ses_a1","title":"x"}`)
		case r.URL.Path == "/v1/files/file_a":
			w.Header().Set("Content-Type", "text/plain")
			w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''a.txt")
			io.WriteString(w, "file-content")
		default:
			w.WriteHeader(404)
			io.WriteString(w, `{"error":{"code":"not_found","message":"x"}}`)
		}
	}))
}

func newTestApp(t *testing.T, upstreamURL string, sink *lockedSink) (*fiber.App, *[]string) {
	var touched []string
	guard := &Guard{
		Store: &fakeStore{
			owners:    map[string]string{"ses_a1": "alice", "ses_b1": "bob"},
			resources: map[string]string{"file_a": "ses_a1"},
		},
		Upstream: &fakeUpstream{},
		Sink:     sink,
	}
	p := &Proxy{
		Client:    NewClient(upstreamURL, "ask_test"),
		Guard:     guard,
		Sink:      sink,
		OnMessage: func(uid, sid string) { touched = append(touched, uid+"@"+sid) },
	}
	app := fiber.New()
	app.All("/agent/v1/*", func(c *fiber.Ctx) error { return p.Handle(c, "alice", c.Params("*")) })
	return app, &touched
}

func TestProxy(t *testing.T) {
	var seen http.Header
	var body string
	up := fakeAgentServer(t, &seen, &body)
	defer up.Close()
	sink := &lockedSink{}
	app, touched := newTestApp(t, up.URL, sink)

	do := func(method, path, reqBody string) (*http.Response, string) {
		req := httptest.NewRequest(method, path, strings.NewReader(reqBody))
		req.Header.Set("Authorization", "Bearer toolbox-token")
		req.Header.Set("Cookie", "secret=1")
		if reqBody != "" {
			req.Header.Set("Content-Type", "application/json")
		}
		resp, err := app.Test(req, 5000)
		if err != nil {
			t.Fatal(err)
		}
		b, _ := io.ReadAll(resp.Body)
		return resp, string(b)
	}

	// 白名单外
	if resp, _ := do("GET", "/agent/v1/sessions", ""); resp.StatusCode != 403 {
		t.Errorf("GET /sessions 应为 403，得到 %d", resp.StatusCode)
	}
	// 他人会话
	if resp, b := do("GET", "/agent/v1/sessions/ses_b1/events", ""); resp.StatusCode != 404 || !strings.Contains(b, "not_found") {
		t.Errorf("他人会话应为 404，得到 %d %s", resp.StatusCode, b)
	}

	// 发送消息：替换鉴权头、不转发 Cookie、记录 runId、刷新活跃时间
	resp, b := do("POST", "/agent/v1/sessions/ses_a1/messages", `{"text":"hi"}`)
	if resp.StatusCode != 202 || !strings.Contains(b, "run_m1") {
		t.Fatalf("发送消息 %d %s", resp.StatusCode, b)
	}
	if seen.Get("Authorization") != "Bearer ask_test" || seen.Get("Cookie") != "" {
		t.Errorf("上游请求头 Authorization=%q Cookie=%q", seen.Get("Authorization"), seen.Get("Cookie"))
	}
	if len(*touched) != 1 || (*touched)[0] != "alice@ses_a1" {
		t.Errorf("OnMessage %v", *touched)
	}

	// PATCH 只保留 title、model
	do("PATCH", "/agent/v1/sessions/ses_a1", `{"title":"t","permissionMode":"auto"}`)
	if body != `{"title":"t"}` {
		t.Errorf("PATCH 转发的请求体 %s", body)
	}

	// 文件下载透传
	resp, b = do("GET", "/agent/v1/files/file_a?inline=1", "")
	if resp.StatusCode != 200 || b != "file-content" || resp.Header.Get("Content-Disposition") == "" {
		t.Errorf("文件下载 %d %q %v", resp.StatusCode, b, resp.Header)
	}

	// SSE 透传并记录 runId
	resp, b = do("GET", "/agent/v1/sessions/ses_a1/stream", "")
	if resp.StatusCode != 200 || !strings.Contains(b, "run.started") || resp.Header.Get("X-Accel-Buffering") != "no" {
		t.Errorf("SSE %d %q", resp.StatusCode, b)
	}

	ids := strings.Join(sink.ids(), ",")
	if !strings.Contains(ids, "run_m1") || !strings.Contains(ids, "run_s1") {
		t.Errorf("记录的子资源 %s", ids)
	}
}