Files
UHD_Media_Tool/apps/server/test/pi/upstreams.test.ts
T
rootandClaude Opus 5.5 0765264a47 feat(pi): fetch and refresh upstream model lists, searchable model picker
- POST /pi/upstreams/{id}/models/refresh calls the provider's model list
  endpoint (OpenAI-shaped /models, Anthropic /v1/models, Google
  /v1beta/models, OpenRouter /api/v1/models; Radius via pi-ai refresh),
  normalizes {id, displayName, contextWindow, supportsVision,
  supportsTools, pricing} and stores it in pi_upstream.models (0044)
- GET .../models returns {source, fetchedAt, models[]}: remote merged
  with the static catalog (remote first); failures (401/network/
  unsupported) fall back to static with a one-line error
- POST /pi/models/probe for trying credentials before creating
- wizard and edit drawer: fetch button, searchable dropdown with
  vision/tool tags, manual id, source line; default upstream needs a model
- tests: mock /models, 401, network failure; smoke step; runbook

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-09-29 01:26:07 +08:00

401 lines
19 KiB
TypeScript

import { mkdtempSync, rmSync } from "node:fs";
import { createServer, type Server } from "node:http";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { FakeOAuthControl } from "@umt/agent";
import { type DB, sql } from "@umt/db";
import { Kysely, PostgresDialect } from "kysely";
import pg from "pg";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { createAgentModule } from "../../src/modules/agent/index.ts";
import { drive } from "../../src/modules/agent/runner.ts";
import { ensureSetupToken } from "../../src/modules/auth-settings/auth-service.ts";
import { createSecretStore, createSettingsReader } from "../../src/modules/auth-settings/provided.ts";
import { createPiModule } from "../../src/modules/pi/index.ts";
import { refreshPiTokens } from "../../src/modules/pi/refresh.ts";
import { dbCredentialStore, setDefaultUpstream } from "../../src/modules/pi/store.ts";
import type { ServerContext } from "../../src/modules/types.ts";
import {
Browser,
bearer,
createHarness,
type Harness,
MASTER_KEY,
PG_URL,
pgAvailable,
} from "../auth-settings/harness.ts";
import { DEV_KEY, json } from "../helpers.ts";
// W7-A Pi 上游:端点鉴权、CRUD 与凭据不出 API、OAuth 状态机(fake provider 模拟 browser / device,不联网)、
// 同 provider 登录串行、端口占用提示、刷新任务的轮换与并发写(两个进程同时刷新只刷新一次)、runAgent 用指定上游。
const hasPg = await pgAvailable();
const PASSWORD = "correct horse battery";
describe.skipIf(!hasPg)("Pi 上游(真实 PG)", () => {
let h: Harness;
let admin: Browser;
const control = new FakeOAuthControl();
const agent = createAgentModule({
fakeScripts: [{ match: /.*/, steps: [], finalText: "fake 完成" }],
});
const sessionDir = mkdtempSync(join(tmpdir(), "umt-pi-"));
let mock: Server;
let mockUrl = "";
const seenAuth: string[] = [];
beforeAll(async () => {
const pi = createPiModule({ fakeOAuth: control, refreshIntervalMs: 0, oauth: { completeWaitMs: 2000 } });
h = await createHarness({ extraModules: [pi, agent] });
agent.registerCapabilities?.(h.ctx.capabilities, h.ctx);
const token = await ensureSetupToken(h.ctx);
admin = new Browser(h.app);
const r = await admin.req("POST", "/auth/setup", { setupToken: token, username: "admin", password: PASSWORD });
expect(r.status).toBe(200);
// OpenAI 兼容的最小流式端点:记录 Authorization,回一句话
mock = createServer((req, res) => {
seenAuth.push(String(req.headers.authorization ?? ""));
req.resume();
res.writeHead(200, { "content-type": "text/event-stream" });
const chunk = (o: unknown) => res.write(`data: ${JSON.stringify(o)}\n\n`);
chunk({
id: "c1",
object: "chat.completion.chunk",
choices: [{ index: 0, delta: { role: "assistant", content: "上游回复" }, finish_reason: null }],
});
chunk({
id: "c1",
object: "chat.completion.chunk",
choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5 },
});
res.end("data: [DONE]\n\n");
});
await new Promise<void>((r) => mock.listen(0, "127.0.0.1", r));
mockUrl = `http://127.0.0.1:${(mock.address() as { port: number }).port}/v1`;
}, 120_000);
afterAll(async () => {
await new Promise((r) => mock?.close(r));
for (const m of [agent]) await m.stop?.();
await h?.close();
rmSync(sessionDir, { recursive: true, force: true });
});
const req = async (method: string, path: string, body?: unknown) => {
const res = await admin.req(method, path, body);
return { status: res.status, body: await json(res) };
};
const create = async (body: Record<string, unknown>) => {
const r = await req("POST", "/pi/upstreams", body);
expect(r.status, JSON.stringify(r.body)).toBe(200);
return r.body.data;
};
it("鉴权:未登录 401;Bearer admin 只能读(写/登录要控制台会话);read key 403", async () => {
expect((await h.app.request("/api/v1/pi/providers")).status).toBe(401);
const key = bearer(h.app, DEV_KEY);
expect((await key("GET", "/pi/providers?models=false")).status).toBe(200);
expect((await key("GET", "/pi/upstreams")).status).toBe(200);
expect((await key("POST", "/pi/upstreams", { name: "x", provider: "fake", authKind: "none" })).status).toBe(403);
const readKey = (await json(await admin.req("POST", "/api-keys", { name: "reader", scopes: ["read"] }))).data
.plaintextKey as string;
expect((await bearer(h.app, readKey)("GET", "/pi/upstreams")).status).toBe(403);
const cat = await json(await key("GET", "/pi/providers?models=false"));
expect(cat.data.providers).toHaveLength(43);
const anthropic = cat.data.providers.find((p: { id: string }) => p.id === "anthropic");
expect(anthropic.termsNote).toMatch(/封禁/);
});
it("CRUD:密钥只写、只回掩码;第一条自动成为默认;校验 provider/authKind/baseUrl;重名 409", async () => {
const a = await create({
name: "OpenAI 主号",
provider: "openai",
authKind: "api_key",
apiKey: "sk-live-AAAA-secret-1234",
});
expect(a).toMatchObject({ status: "ok", hasCredential: true, credentialMask: "sk-l…1234", isDefault: true });
expect(JSON.stringify(a)).not.toContain("secret");
const b = await create({
name: "OpenAI 副号",
provider: "openai",
authKind: "api_key",
apiKey: "sk-live-BBBB-secret-5678",
baseUrl: "https://proxy.local/v1",
makeDefault: true,
});
expect(b.isDefault).toBe(true);
const list = (await req("GET", "/pi/upstreams")).body.data;
expect(list.defaultUpstreamId).toBe(b.id);
expect(list.items.filter((u: { isDefault: boolean }) => u.isDefault)).toHaveLength(1);
expect(JSON.stringify(list)).not.toMatch(/secret-\d/);
expect((await h.ctx.settings.get("pi")).defaultUpstreamId).toBe(b.id);
expect(
(await req("POST", "/pi/upstreams", { name: "c", provider: "openai-codex", authKind: "api_key" })).status,
).toBe(422);
expect(
(
await req("POST", "/pi/upstreams", {
name: "c",
provider: "github-copilot",
authKind: "api_key",
baseUrl: "https://x.local",
})
).status,
).toBe(422);
expect(
(await req("POST", "/pi/upstreams", { name: "OpenAI 主号", provider: "openai", authKind: "api_key" })).status,
).toBe(409);
// 清 key → auth_required;测试报 auth
const cleared = (await req("PATCH", `/pi/upstreams/${a.id}`, { apiKey: null })).body.data;
expect(cleared).toMatchObject({ status: "auth_required", hasCredential: false });
const t = (await req("POST", `/pi/upstreams/${a.id}/test`, {})).body.data;
expect(t).toMatchObject({ ok: false, errorKind: "auth" });
// 重新填 key → 测试(只解析凭据 + 校验模型,不联网)通过
await req("PATCH", `/pi/upstreams/${a.id}`, { apiKey: "sk-live-CCCC-secret-9999", defaultModel: "gpt-5" });
const t2 = (await req("POST", `/pi/upstreams/${a.id}/test`, {})).body.data;
expect(t2).toMatchObject({ ok: true, mode: "auth", source: "stored credential", model: "gpt-5" });
expect(t2.upstream.lastVerifiedAt).toBeTruthy();
// 自定义 Base URL:凭证检查不连端点,不改写可用状态,提示用最小补全(QA U2)
await h.db.updateTable("pi_upstream").set({ status: "error" }).where("id", "=", b.id).execute();
const tb = (await req("POST", `/pi/upstreams/${b.id}/test`, {})).body.data;
expect(tb).toMatchObject({ ok: true, mode: "auth", upstream: { status: "error", lastVerifiedAt: null } });
expect(tb.message).toMatch(/最小补全/);
const models = (await req("GET", `/pi/upstreams/${a.id}/models`)).body.data;
expect(models.models.length).toBeGreaterThan(10);
expect(models).toMatchObject({ source: "static", fetchedAt: null, remoteSupported: true });
// 删除默认上游 → 默认清空,secret 一并删除
const secretRef = (
await h.db.selectFrom("pi_upstream").select("secret_ref").where("id", "=", b.id).executeTakeFirstOrThrow()
).secret_ref!;
expect((await req("DELETE", `/pi/upstreams/${b.id}`)).status).toBe(200);
expect((await req("GET", "/pi/upstreams")).body.data.defaultUpstreamId).toBeNull();
expect(await h.db.selectFrom("secret").select("id").where("id", "=", secretRef).executeTakeFirst()).toBeUndefined();
await req("DELETE", `/pi/upstreams/${a.id}`);
expect(h.lines.join("\n")).not.toMatch(/sk-live-\w+-secret/);
});
it("OAuth browser:粘错 state 失败可读;粘对 URL 成功,token 包加密入库", async () => {
const u = await create({ name: "fake 订阅 A", provider: "fake", authKind: "oauth" });
expect(u.status).toBe("auth_required");
const s1 = (await req("POST", `/pi/upstreams/${u.id}/oauth/start`, { mode: "browser" })).body.data;
expect(s1).toMatchObject({ mode: "browser", state: "pending", userCode: null });
expect(s1.authorizeUrl).toMatch(/^https:\/\/fake-oauth\.invalid\/authorize\?state=/);
const bad = (
await req("POST", `/pi/upstreams/${u.id}/oauth/complete`, {
sessionId: s1.sessionId,
pastedUrlOrCode: "http://localhost:1455/auth/callback?code=abc&state=nope",
})
).body.data;
expect(bad).toMatchObject({ state: "failed" });
expect(bad.message).toMatch(/state/);
const s2 = (await req("POST", `/pi/upstreams/${u.id}/oauth/start`, {})).body.data;
const state = new URL(s2.authorizeUrl).searchParams.get("state");
const done = (
await req("POST", `/pi/upstreams/${u.id}/oauth/complete`, {
sessionId: s2.sessionId,
pastedUrlOrCode: `http://localhost:1455/auth/callback?code=xyz&state=${state}`,
})
).body.data;
expect(done).toMatchObject({
state: "succeeded",
upstream: { status: "ok", hasCredential: true, account: "fake-browser-user" },
});
expect(done.upstream.oauthExpiresAt).toBeTruthy();
const row = await h.db
.selectFrom("pi_upstream")
.select("secret_ref")
.where("id", "=", u.id)
.executeTakeFirstOrThrow();
const sec = await h.db
.selectFrom("secret")
.select("ciphertext")
.where("id", "=", row.secret_ref!)
.executeTakeFirstOrThrow();
expect(Buffer.from(sec.ciphertext).toString("utf8")).not.toContain("fake-access");
expect(JSON.stringify(done)).not.toContain("fake-access");
// 过期会话
const gone = (await req("POST", `/pi/upstreams/${u.id}/oauth/complete`, { sessionId: "nope" })).body.data;
expect(gone.state).toBe("expired");
// 测试端点对 OAuth 上游可用
expect((await req("POST", `/pi/upstreams/${u.id}/test`, {})).body.data).toMatchObject({
ok: true,
source: "OAuth",
});
// 登出 → 凭据删除
const out = (await req("POST", `/pi/upstreams/${u.id}/oauth/logout`)).body.data;
expect(out).toMatchObject({ status: "auth_required", hasCredential: false });
await req("DELETE", `/pi/upstreams/${u.id}`);
});
it("OAuth device:服务端轮询,未批准返回 pending,批准后 succeeded;同 provider 登录串行;端口占用提示", async () => {
control.deviceAutoApproveMs = null;
const a = await create({ name: "fake 订阅 D1", provider: "fake", authKind: "oauth" });
const b = await create({ name: "fake 订阅 D2", provider: "fake", authKind: "oauth" });
const s = (await req("POST", `/pi/upstreams/${a.id}/oauth/start`, { mode: "device" })).body.data;
expect(s).toMatchObject({ mode: "device", state: "pending", verificationUrl: "https://fake-oauth.invalid/device" });
expect(s.userCode).toMatch(/^FAKE-/);
// 同 provider 另一个上游的登录要等这一个结束
expect((await req("POST", `/pi/upstreams/${b.id}/oauth/start`, { mode: "device" })).status).toBe(409);
const p = (await req("POST", `/pi/upstreams/${a.id}/oauth/complete`, { sessionId: s.sessionId, waitMs: 100 })).body
.data;
expect(p.state).toBe("pending");
control.approveDevice();
const ok = (await req("POST", `/pi/upstreams/${a.id}/oauth/complete`, { sessionId: s.sessionId, waitMs: 5000 }))
.body.data;
expect(ok).toMatchObject({ state: "succeeded", upstream: { account: "fake-device-user", status: "ok" } });
// 端口占用 → 可读提示
control.browserListenError = "listen EADDRINUSE: address already in use 127.0.0.1:1455";
const f = (await req("POST", `/pi/upstreams/${b.id}/oauth/start`, { mode: "browser" })).body.data;
expect(f.state).toBe("failed");
expect(f.message).toMatch(/端口 1455 被占用/);
control.browserListenError = null;
control.deviceAutoApproveMs = 3000;
await req("DELETE", `/pi/upstreams/${a.id}`);
await req("DELETE", `/pi/upstreams/${b.id}`);
});
/** 第二个“进程”:独立连接池 + 独立 secrets / settings */
function secondProcess(): { ctx: ServerContext; close(): Promise<void> } {
const db = new Kysely<DB>({
dialect: new PostgresDialect({
pool: new pg.Pool({ connectionString: PG_URL, max: 3, options: `-c search_path=${h.schema},public` }),
}),
});
const ctx = {
...h.ctx,
db,
secrets: createSecretStore({ db, masterKey: MASTER_KEY }),
settings: createSettingsReader({ db, logger: h.ctx.logger }),
} as ServerContext;
return { ctx, close: () => db.destroy() };
}
it("刷新任务:两个进程同时刷新只刷新一次,轮换后的 refresh token 写回;失败 → expired + 待处理事项,重新登录后关闭", async () => {
control.deviceAutoApproveMs = 0;
const u = await create({ name: "fake 订阅 R", provider: "fake", authKind: "oauth" });
const s = (await req("POST", `/pi/upstreams/${u.id}/oauth/start`, { mode: "device" })).body.data;
await req("POST", `/pi/upstreams/${u.id}/oauth/complete`, { sessionId: s.sessionId, waitMs: 5000 });
// 让 token 快过期
const store = dbCredentialStore(h.ctx, u.id);
await store.modify(async (cur) => ({ ...(cur as object), expires: Date.now() + 60_000 }) as never);
const before = (await store.read()) as { refresh: string };
control.refreshCalls = 0;
control.refreshDelayMs = 150;
const p2 = secondProcess();
try {
const [r1, r2] = await Promise.all([
refreshPiTokens(h.ctx, { client: { fakeOAuth: control } }),
refreshPiTokens(p2.ctx, { client: { fakeOAuth: control } }),
]);
expect(control.refreshCalls).toBe(1);
expect(r1.failed).toEqual([]);
expect(r2.failed).toEqual([]);
} finally {
await p2.close();
control.refreshDelayMs = 0;
}
const after = (await store.read()) as { refresh: string; expires: number };
expect(after.refresh).not.toBe(before.refresh);
expect(after.expires).toBeGreaterThan(Date.now() + 30 * 60_000);
// 已不在刷新窗口内:再跑一次不刷新
await refreshPiTokens(h.ctx, { client: { fakeOAuth: control } });
expect(control.refreshCalls).toBe(1);
// refresh token 失效
await store.modify(async (cur) => ({ ...(cur as object), expires: Date.now() + 60_000 }) as never);
control.failRefresh = true;
const r = await refreshPiTokens(h.ctx, { client: { fakeOAuth: control } });
control.failRefresh = false;
expect(r.failed).toHaveLength(1);
const view = (await req("GET", `/pi/upstreams/${u.id}`)).body.data;
expect(view).toMatchObject({ status: "expired", hasCredential: true });
expect(view.lastError).toMatch(/invalid_grant/);
const d = await h.db
.selectFrom("decision")
.selectAll()
.where("dedupe_key", "=", `pi_upstream_expired:${u.id}`)
.executeTakeFirstOrThrow();
expect(d).toMatchObject({ status: "open", subject_type: "pi_upstream", subject_id: u.id });
// 失效后不再自动刷新
expect((await refreshPiTokens(h.ctx, { client: { fakeOAuth: control } })).checked).toBe(0);
const s2 = (await req("POST", `/pi/upstreams/${u.id}/oauth/start`, { mode: "device" })).body.data;
const ok = (await req("POST", `/pi/upstreams/${u.id}/oauth/complete`, { sessionId: s2.sessionId, waitMs: 5000 }))
.body.data;
expect(ok.upstream.status).toBe("ok");
const d2 = await h.db.selectFrom("decision").select("status").where("id", "=", d.id).executeTakeFirstOrThrow();
expect(d2.status).toBe("resolved");
control.deviceAutoApproveMs = 3000;
await req("DELETE", `/pi/upstreams/${u.id}`);
});
async function newRun(upstreamId: string | null) {
const r = await h.db
.insertInto("agent_run")
.values({
trigger: "user",
goal: "测试上游选择",
scope: JSON.stringify({ capabilities: ["task.inspect"], allowedKinds: null }),
upstream_id: upstreamId,
})
.returning("id")
.executeTakeFirstOrThrow();
return r.id;
}
it("runAgent:缺省用默认上游,指定 upstreamId 覆盖;key 由服务端解密注入请求;停用上游报可读错误", async () => {
const env = { ...agent.env(h.ctx), opts: { ...agent.env(h.ctx).opts, sessionDir } };
const fake = await create({ name: "默认 fake", provider: "fake", authKind: "none", makeDefault: true });
const custom = await create({
name: "本地兼容端点",
provider: "openai-compatible",
authKind: "api_key",
apiKey: "sk-local-upstream-key",
baseUrl: mockUrl,
defaultModel: "local-model",
thinkingLevel: "off",
});
const r1 = await newRun(null);
const d1 = await drive(env, h.ctx.capabilities, r1, new AbortController().signal);
expect(d1).toMatchObject({ status: "completed", finalText: "fake 完成" });
const row1 = await h.db.selectFrom("agent_run").selectAll().where("id", "=", r1).executeTakeFirstOrThrow();
expect(row1).toMatchObject({ provider: "fake", upstream_id: fake.id });
const r2 = await newRun(custom.id);
const d2 = await drive(env, h.ctx.capabilities, r2, new AbortController().signal);
expect(d2, JSON.stringify(d2)).toMatchObject({ status: "completed", finalText: "上游回复" });
expect(seenAuth.at(-1)).toBe("Bearer sk-local-upstream-key");
const row2 = await h.db.selectFrom("agent_run").selectAll().where("id", "=", r2).executeTakeFirstOrThrow();
expect(row2).toMatchObject({ provider: "openai-compatible", model: "local-model", upstream_id: custom.id });
const events = await sql<{ payload: unknown }>`select payload from agent_run_event where run_id = ${r2}`.execute(
h.db,
);
expect(JSON.stringify(events.rows)).not.toContain("sk-local-upstream-key");
expect(h.lines.join("\n")).not.toContain("sk-local-upstream-key");
await req("PATCH", `/pi/upstreams/${custom.id}`, { enabled: false });
const r3 = await newRun(custom.id);
const d3 = await drive(env, h.ctx.capabilities, r3, new AbortController().signal);
expect(d3).toMatchObject({ status: "failed" });
expect(d3.message).toMatch(/已停用/);
// 没有默认上游:旧 settings.pi 字段不再兜底,直接失败
await setDefaultUpstream(h.ctx, null);
const r4 = await newRun(null);
const d4 = await drive(env, h.ctx.capabilities, r4, new AbortController().signal);
expect(d4).toMatchObject({ status: "failed" });
expect(d4.message).toMatch(/未配置 Agent 上游/);
await setDefaultUpstream(h.ctx, fake.id);
});
});