ModelRouter/internal/core/channels_zcode_test.go

package core

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"llmsproxy/internal/lua"
)

// zcodeHarness：ZCode 两段式服务端中介授权流的测试台。
//
// 参考语义（zcode-login.ts 逐行）：
//
//	① POST /api/v1/oauth/cli/init  → {flow_id, authorize_url, poll_interval_sec}
//	   Authorization 是**本地造的会话密钥**（不是用户凭据）
//	② 用户在浏览器打开 authorize_url
//	③ GET  /api/v1/oauth/cli/poll/{flow_id} → {code:0, data:{status,...}}
//	   判据：4xx（除 408/429）= 终态失败；5xx/网络错误 = 继续等；
//	         code!=0 或 data 缺失 = pending；status=pending = 继续
//
// 此前标 flow="manual"（理由"要解 captcha"），但该授权流**完全不涉及 captcha**
// —— captcha 属于另一条已废弃的旁路（读取本机官方客户端数据）。
type zcodeHarness struct {
	t       *testing.T
	ps      *lua.Plugins
	ts      *httptest.Server
	polls   int32
	readyAt int32
	gotInit map[string]interface{}
	gotAuth string
}

func newZcodeHarness(t *testing.T) *zcodeHarness {
	t.Helper()
	h := &zcodeHarness{t: t, readyAt: 2}
	h.ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch {
		case r.Method == "POST" && r.URL.Path == "/api/v1/oauth/cli/init":
			body, _ := io.ReadAll(r.Body)
			_ = json.Unmarshal(body, &h.gotInit)
			h.gotAuth = r.Header.Get("Authorization")
			fmt.Fprint(w, `{"code":0,"data":{"flow_id":"FLOW-1","poll_token":"PT",
			  "authorize_url":"https://bigmodel.cn/login?appId=zcode&redirect=x",
			  "expires_at":9999999999,"poll_interval_sec":2}}`)
		case r.Method == "GET" && strings.HasPrefix(r.URL.Path, "/api/v1/oauth/cli/poll/"):
			if !strings.HasSuffix(r.URL.Path, "FLOW-1") {
				fmt.Fprint(w, `{"code":1,"msg":"unknown flow"}`)
				return
			}
			n := atomic.AddInt32(&h.polls, 1)
			if n < atomic.LoadInt32(&h.readyAt) {
				fmt.Fprint(w, `{"code":0,"data":{"status":"pending"}}`)
				return
			}
			fmt.Fprint(w, `{"code":0,"data":{"status":"ready",
			  "token":"ZCODE-JWT-1",
			  "user":{"user_id":"u7","name":"zcode-dev"},
			  "bigmodel":{"access_token":"BM-AT-1","refresh_token":"BM-RT-1"}}}`)
		default:
			w.WriteHeader(404)
		}
	}))
	t.Cleanup(h.ts.Close)

	c := kernelTestCore(t)
	h.ps = c.Plugins()
	src, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	if err := h.ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	_ = h.ps.SetPluginKernelPermission("channels", []string{"add_source", "list_sources", "status"}, false)
	if _, err := h.ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "zcode",
		"initpoll": map[string]interface{}{
			"init_url": h.ts.URL + "/api/v1/oauth/cli/init",
			"poll_url": h.ts.URL + "/api/v1/oauth/cli/poll/{flow_id}",
		},
	}); err != nil {
		t.Fatalf("set_endpoints: %v", err)
	}
	return h
}

func (h *zcodeHarness) begin() (map[string]interface{}, string) {
	out, err := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"zcode"}`, Role: "admin",
	})
	if err != nil {
		h.t.Fatal(err)
	}
	j := out.JSON.(map[string]interface{})
	return j, j["session"].(map[string]interface{})["id"].(string)
}

func (h *zcodeHarness) pollOnce(sid string) map[string]interface{} {
	out, err := h.ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
		Method: "GET", Query: map[string]string{"id": sid}, Role: "admin",
	})
	if err != nil {
		h.t.Fatal(err)
	}
	return out.JSON.(map[string]interface{})
}

// TestZcodeInitPollBehavior 逐行锁定参考的 ZCode 授权流语义。
func TestZcodeInitPollBehavior(t *testing.T) {
	h := newZcodeHarness(t)

	t.Run("行1 init带本地会话密钥与provider", func(t *testing.T) {
		atomic.StoreInt32(&h.polls, 0)
		j, _ := h.begin()
		if h.gotInit["provider"] != "bigmodel" {
			t.Fatalf("init body 应带 provider=bigmodel, 实得 %v", h.gotInit)
		}
		// Authorization 是我们自己生成的会话密钥（Bearer 后非空），不是用户凭据
		if !strings.HasPrefix(h.gotAuth, "Bearer ") || len(h.gotAuth) < 20 {
			t.Fatalf("init 应带本地会话密钥作 Bearer, 实得 %q", h.gotAuth)
		}
		if u, _ := j["url"].(string); !strings.Contains(u, "bigmodel.cn") {
			t.Fatalf("应给出 authorize_url, 实得 %q", u)
		}
	})

	t.Run("行2 pending状态保持pending", func(t *testing.T) {
		atomic.StoreInt32(&h.polls, 0)
		atomic.StoreInt32(&h.readyAt, 99)
		_, sid := h.begin()
		for i := 0; i < 3; i++ {
			if st := h.pollOnce(sid); st["state"] != "pending" {
				t.Fatalf("status=pending 应保持 pending, 实得 %v", st["state"])
			}
		}
	})

	t.Run("行3 ready即完成_凭据含zcodeJwt与第三方token", func(t *testing.T) {
		atomic.StoreInt32(&h.polls, 0)
		atomic.StoreInt32(&h.readyAt, 2)
		_, sid := h.begin()
		if st := h.pollOnce(sid); st["state"] != "pending" {
			t.Fatalf("首拍应 pending: %v", st)
		}
		st := h.pollOnce(sid)
		if st["state"] != "completed" {
			t.Fatalf("ready 应 completed, 实得 %v", st)
		}
		v, ok, _ := h.ps.GetSecret("channels", "channels:zcode")
		if !ok {
			t.Fatal("凭据未落库")
		}
		for _, want := range []string{"ZCODE-JWT-1", "BM-AT-1", "zcode-dev"} {
			if !strings.Contains(v, want) {
				t.Fatalf("凭据应含 %q, 实得 %s", want, v)
			}
		}
	})

	t.Run("行4 池路径_凭据落账号ref", func(t *testing.T) {
		atomic.StoreInt32(&h.polls, 0)
		atomic.StoreInt32(&h.readyAt, 1)
		add, err := h.ps.InvokeHandler("channels", "account_add", &lua.HandlerRequest{
			Method: "POST", Body: `{"platform":"zcode"}`, Role: "admin",
		})
		if err != nil {
			t.Fatal(err)
		}
		aid := add.JSON.(map[string]interface{})["account_id"].(string)
		out, _ := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
			Method: "POST", Body: fmt.Sprintf(`{"platform":"zcode","account_id":%q}`, aid), Role: "admin",
		})
		sid := out.JSON.(map[string]interface{})["session"].(map[string]interface{})["id"].(string)
		if st := h.pollOnce(sid); st["state"] != "completed" {
			t.Fatalf("池路径应完成: %v", st)
		}
		if _, ok, _ := h.ps.GetSecret("channels", "channels:zcode:"+aid); !ok {
			t.Fatalf("池路径凭据未落到 channels:zcode:%s", aid)
		}
	})
}