441 lines
16 KiB
TypeScript
441 lines
16 KiB
TypeScript
import { EventEmitter } from "node:events";
|
|
import { existsSync, mkdtempSync, readFileSync, readdirSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { dirname, join } from "node:path";
|
|
import { afterEach, expect, test, vi } from "vitest";
|
|
import { loadConfig } from "../src/config.js";
|
|
import { createPiProviderSmoke } from "../src/pi/provider-smoke.js";
|
|
import type { RuntimeModel, RuntimeModelCatalog } from "../src/models/runtime-model-catalog.js";
|
|
|
|
afterEach(() => vi.unstubAllEnvs());
|
|
|
|
function rpcChild(onCommand: (command: any, emit: (message: unknown) => void) => void) {
|
|
const child: any = new EventEmitter();
|
|
child.stdout = new EventEmitter();
|
|
child.stderr = { resume: vi.fn() };
|
|
child.kill = vi.fn();
|
|
const emit = (message: unknown) => queueMicrotask(() => {
|
|
child.stdout.emit("data", `${JSON.stringify(message)}\n`);
|
|
});
|
|
child.stdin = {
|
|
write: (data: unknown) => {
|
|
onCommand(JSON.parse(String(data)), emit);
|
|
return true;
|
|
},
|
|
};
|
|
return child;
|
|
}
|
|
|
|
function successfulProviderChild() {
|
|
return rpcChild((command, emit) => {
|
|
if (command.type === "set_model" || command.type === "set_thinking_level") {
|
|
emit({ type: "response", id: command.id, success: true });
|
|
}
|
|
if (command.type === "prompt") {
|
|
emit({
|
|
type: "message_end",
|
|
message: { role: "assistant", stopReason: "stop", content: "must-not-be-returned" },
|
|
});
|
|
emit({
|
|
type: "agent_end",
|
|
messages: [{ role: "assistant", stopReason: "stop", content: "must-not-be-returned" }],
|
|
});
|
|
}
|
|
});
|
|
}
|
|
|
|
const MANAGED_CONFIG_ERROR = "Pi provider/model configuration is invalid";
|
|
|
|
test("provider smoke resolves the selected catalog credential from the secret bundle", async () => {
|
|
const root = mkdtempSync(join(tmpdir(), "thothii-smoke-catalog-credential-"));
|
|
const secret = join(root, "thothii.secrets");
|
|
writeFileSync(secret, "ZAI_API_KEY=catalog-secret\nTHT_MODEL_API_KEY=legacy-secret\n", { mode: 0o600 });
|
|
const model: RuntimeModel = {
|
|
id: "openai/test-model",
|
|
provider: "openai",
|
|
model: "test-model",
|
|
label: "Test model",
|
|
upstreamModel: "test-model",
|
|
authentication: { mode: "secret_env", apiKeyEnv: "ZAI_API_KEY" },
|
|
sessionAdapter: { mode: "pi_builtin" },
|
|
session: { reasoning: false },
|
|
};
|
|
const modelCatalog: RuntimeModelCatalog = {
|
|
defaultSession: model.id,
|
|
defaultMetadataGeneration: null,
|
|
embedding: null,
|
|
sessionModels: () => [model],
|
|
metadataModels: () => [],
|
|
hasSession: (id) => id === model.id,
|
|
};
|
|
let spawnEnv: NodeJS.ProcessEnv | undefined;
|
|
const smoke = createPiProviderSmoke(loadConfig({ THT_SECRETS_FILE: secret }), {
|
|
modelCatalog,
|
|
authProviders: () => new Set(),
|
|
readModelsStore: () => undefined,
|
|
spawnFn: (_command, _args, options) => {
|
|
spawnEnv = options.env;
|
|
return successfulProviderChild();
|
|
},
|
|
});
|
|
try {
|
|
await expect(smoke({
|
|
provider: "openai", model: "test-model", reasoning: "medium", timeoutMs: 750,
|
|
})).resolves.toBeUndefined();
|
|
expect(spawnEnv?.ZAI_API_KEY).toBe("catalog-secret");
|
|
expect(spawnEnv).not.toHaveProperty("OPENAI_API_KEY");
|
|
expect(spawnEnv).not.toHaveProperty("THT_MODEL_API_KEY");
|
|
} finally {
|
|
rmSync(root, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
// Catches an isolated smoke agent that copies auth.json but drops the selected custom
|
|
// provider/model from models.json, causing set_model to fail before the real request.
|
|
test("provider smoke reaches the selected custom provider from an isolated models.json", async () => {
|
|
let isolatedAgentDir: string | undefined;
|
|
let providerRequests = 0;
|
|
const child = rpcChild((command, emit) => {
|
|
if (command.type === "set_model") {
|
|
const models = JSON.parse(readFileSync(`${isolatedAgentDir}/models.json`, "utf8"));
|
|
const selectedProvider = models.providers?.[command.provider];
|
|
const selectedModel = selectedProvider?.models?.find(
|
|
(candidate: { id?: unknown }) => candidate.id === command.modelId,
|
|
);
|
|
emit({ type: "response", id: command.id, success: Boolean(selectedModel) });
|
|
}
|
|
if (command.type === "set_thinking_level") {
|
|
emit({ type: "response", id: command.id, success: true });
|
|
}
|
|
if (command.type === "prompt") {
|
|
providerRequests++;
|
|
emit({
|
|
type: "message_end",
|
|
message: { role: "assistant", stopReason: "stop", content: "must-not-be-returned" },
|
|
});
|
|
emit({
|
|
type: "agent_end",
|
|
messages: [{ role: "assistant", stopReason: "stop", content: "must-not-be-returned" }],
|
|
});
|
|
}
|
|
});
|
|
const smoke = createPiProviderSmoke(loadConfig({ PI_BIN: "/usr/local/bin/pi" }), {
|
|
spawnFn: (_command, _args, options) => {
|
|
isolatedAgentDir = options.env.PI_CODING_AGENT_DIR;
|
|
expect(readdirSync(isolatedAgentDir)).toEqual(["auth.json", "models.json"]);
|
|
expect(JSON.parse(readFileSync(`${isolatedAgentDir}/models.json`, "utf8"))).toEqual({
|
|
providers: {
|
|
"custom-openai": {
|
|
baseUrl: "https://selected.invalid/v1",
|
|
apiKey: "$CUSTOM_OPENAI_API_KEY",
|
|
api: "openai-completions",
|
|
headers: { "X-Literal-Bang": "$!literal-value" },
|
|
models: [{ id: "selected-model", name: "Selected model", reasoning: true }],
|
|
},
|
|
},
|
|
});
|
|
expect(JSON.parse(readFileSync(`${isolatedAgentDir}/auth.json`, "utf8"))).toEqual({
|
|
"custom-openai": { type: "api_key", key: "${CUSTOM_OPENAI_API_KEY}" },
|
|
});
|
|
return child;
|
|
},
|
|
authProviders: () => new Set(["custom-openai"]),
|
|
readAuthStore: () => JSON.stringify({
|
|
"custom-openai": { type: "api_key", key: "${CUSTOM_OPENAI_API_KEY}" },
|
|
unrelated: { type: "api_key", key: "!must-not-run-or-enter-isolated-context" },
|
|
}),
|
|
readModelsStore: () => JSON.stringify({
|
|
providers: {
|
|
"custom-openai": {
|
|
baseUrl: "https://selected.invalid/v1",
|
|
apiKey: "$CUSTOM_OPENAI_API_KEY",
|
|
api: "openai-completions",
|
|
headers: { "X-Literal-Bang": "$!literal-value" },
|
|
models: [
|
|
{ id: "selected-model", name: "Selected model", reasoning: true },
|
|
{ id: "unrelated-model", name: "Must not enter isolated context" },
|
|
],
|
|
},
|
|
unrelated: {
|
|
baseUrl: "https://unrelated.invalid/v1",
|
|
api: "openai-completions",
|
|
apiKey: "!must-not-run",
|
|
models: [{ id: "unrelated-model" }],
|
|
},
|
|
},
|
|
}),
|
|
});
|
|
|
|
await expect(smoke({
|
|
provider: "custom-openai", model: "selected-model", reasoning: "medium", timeoutMs: 750,
|
|
})).resolves.toBeUndefined();
|
|
expect(providerRequests).toBe(1);
|
|
expect(child.kill).toHaveBeenCalledOnce();
|
|
expect(isolatedAgentDir && existsSync(dirname(isolatedAgentDir))).toBe(false);
|
|
});
|
|
|
|
const selectedExecutableConfigCases: Array<{
|
|
name: string;
|
|
selectedAuth?: unknown;
|
|
selectedProvider: Record<string, unknown>;
|
|
}> = [
|
|
{
|
|
name: "selected auth credential",
|
|
selectedAuth: {
|
|
type: "api_key",
|
|
key: "!sensitive-credential-command /private/credential-path",
|
|
},
|
|
selectedProvider: {},
|
|
},
|
|
{
|
|
name: "selected provider apiKey",
|
|
selectedProvider: {
|
|
apiKey: "!sensitive-api-key-command /private/key-path",
|
|
},
|
|
},
|
|
{
|
|
name: "nested selected-provider headers",
|
|
selectedProvider: {
|
|
headers: { Authorization: "!sensitive-header-command /private/header-path" },
|
|
},
|
|
},
|
|
{
|
|
name: "selected model object",
|
|
selectedProvider: {
|
|
models: [{
|
|
id: "selected-model",
|
|
name: "!sensitive-model-command /private/model-path",
|
|
}],
|
|
},
|
|
},
|
|
{
|
|
name: "selected model override",
|
|
selectedProvider: {
|
|
modelOverrides: {
|
|
"selected-model": {
|
|
headers: { "X-Override": "!sensitive-override-command /private/override-path" },
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "array nested in selected provider configuration",
|
|
selectedProvider: {
|
|
compat: {
|
|
nested: ["literal", { value: "!sensitive-array-command /private/array-path" }],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
|
|
// Catches a defense that validates only known top-level fields or waits until after the isolated
|
|
// Pi process starts. Every selected value crossing into auth.json/models.json must be declarative.
|
|
test.each(selectedExecutableConfigCases)(
|
|
"provider smoke rejects executable config in $name before isolated Pi spawn",
|
|
async ({ selectedAuth, selectedProvider }) => {
|
|
const child = successfulProviderChild();
|
|
const spawnFn = vi.fn(() => child);
|
|
const smoke = createPiProviderSmoke(loadConfig({}), {
|
|
spawnFn,
|
|
authProviders: () => new Set(["custom-openai"]),
|
|
readAuthStore: () => JSON.stringify({
|
|
"custom-openai": selectedAuth ?? { type: "api_key", key: "test-only" },
|
|
unrelated: { type: "api_key", key: "!must-not-contaminate-selected-provider" },
|
|
}),
|
|
readModelsStore: () => JSON.stringify({
|
|
providers: {
|
|
"custom-openai": {
|
|
baseUrl: "https://selected.invalid/v1",
|
|
api: "openai-completions",
|
|
models: [{ id: "selected-model", name: "Selected model" }],
|
|
...selectedProvider,
|
|
},
|
|
unrelated: {
|
|
apiKey: "!must-not-contaminate-selected-provider",
|
|
models: [{ id: "unrelated-model" }],
|
|
},
|
|
},
|
|
}),
|
|
});
|
|
|
|
let caught: unknown;
|
|
try {
|
|
await smoke({
|
|
provider: "custom-openai", model: "selected-model", reasoning: "medium", timeoutMs: 750,
|
|
});
|
|
} catch (error) {
|
|
caught = error;
|
|
}
|
|
expect(caught).toBeInstanceOf(Error);
|
|
expect((caught as Error).message).toBe(MANAGED_CONFIG_ERROR);
|
|
expect(String(caught)).not.toMatch(/sensitive|private|command|path/i);
|
|
expect(spawnFn).not.toHaveBeenCalled();
|
|
},
|
|
);
|
|
|
|
// Catches a provider smoke process that runs from the trusted harness or leaves Pi tools,
|
|
// extensions, skills, context files, templates, themes, or session persistence enabled.
|
|
test("provider smoke makes one configured request from an isolated no-capability Pi process", async () => {
|
|
vi.stubEnv("THT_DATA_ROOT", "/mounted-workflow-state");
|
|
vi.stubEnv("THT_SESSION", "mounted-session-id");
|
|
vi.stubEnv("THT_AUTHOR", "mounted-author");
|
|
vi.stubEnv("THT_CONFIG", "/mounted-workflow-state/config.yaml");
|
|
vi.stubEnv("PI_CODING_AGENT_DIR", "/home/thoth/.pi/agent");
|
|
vi.stubEnv("PI_CODING_AGENT_SESSION_DIR", "/mounted-session-state");
|
|
const commands: any[] = [];
|
|
const spawns: any[][] = [];
|
|
const child = rpcChild((command, emit) => {
|
|
commands.push(command);
|
|
if (command.type === "set_model" || command.type === "set_thinking_level") {
|
|
emit({ type: "response", id: command.id, success: true });
|
|
}
|
|
if (command.type === "prompt") {
|
|
emit({
|
|
type: "message_end",
|
|
message: { role: "assistant", stopReason: "stop", content: "raw-provider-output" },
|
|
});
|
|
emit({
|
|
type: "agent_end",
|
|
messages: [{ role: "assistant", stopReason: "stop", content: "raw-provider-output" }],
|
|
});
|
|
}
|
|
});
|
|
const smoke = createPiProviderSmoke(loadConfig({
|
|
THT_HARNESS_DIR: "/app/harness",
|
|
PI_BIN: "/usr/local/bin/pi",
|
|
THT_DATA_ROOT: "/mounted-workflow-state",
|
|
}), {
|
|
spawnFn: (...args) => {
|
|
spawns.push(args);
|
|
expect(args[2].cwd).not.toBe("/app/harness");
|
|
expect(readdirSync(args[2].cwd)).toEqual([]);
|
|
expect(args[2].env.PI_CODING_AGENT_DIR).not.toBe("/home/thoth/.pi/agent");
|
|
expect(readdirSync(args[2].env.PI_CODING_AGENT_DIR)).toEqual(["auth.json"]);
|
|
expect(JSON.parse(readFileSync(`${args[2].env.PI_CODING_AGENT_DIR}/auth.json`, "utf8")))
|
|
.toEqual({ zai: { type: "api_key", key: "test-only" } });
|
|
return child;
|
|
},
|
|
authProviders: () => new Set(["zai"]),
|
|
readAuthStore: () => JSON.stringify({
|
|
zai: { type: "api_key", key: "test-only" },
|
|
deepseek: { type: "api_key", key: "must-not-enter-isolated-context" },
|
|
}),
|
|
readModelsStore: () => undefined,
|
|
});
|
|
|
|
await expect(smoke({
|
|
provider: "zai", model: "glm-5.2", reasoning: "medium", timeoutMs: 750,
|
|
})).resolves.toBeUndefined();
|
|
expect(spawns).toHaveLength(1);
|
|
expect(spawns[0][0]).toBe("/usr/local/bin/pi");
|
|
expect(spawns[0][1]).toEqual([
|
|
"--mode", "rpc",
|
|
"--no-session",
|
|
"--no-tools",
|
|
"--no-extensions",
|
|
"--no-skills",
|
|
"--no-prompt-templates",
|
|
"--no-themes",
|
|
"--no-context-files",
|
|
"--no-approve",
|
|
]);
|
|
expect(spawns[0][2].env).not.toHaveProperty("THT_DATA_ROOT");
|
|
expect(spawns[0][2].env).not.toHaveProperty("THT_SESSION");
|
|
expect(spawns[0][2].env).not.toHaveProperty("THT_AUTHOR");
|
|
expect(spawns[0][2].env).not.toHaveProperty("THT_CONFIG");
|
|
expect(spawns[0][2].env).not.toHaveProperty("PI_CODING_AGENT_SESSION_DIR");
|
|
expect(existsSync(dirname(spawns[0][2].cwd))).toBe(false);
|
|
expect(commands.map(({ id: _id, ...command }) => command)).toEqual([
|
|
{ type: "set_model", provider: "zai", modelId: "glm-5.2" },
|
|
{ type: "set_thinking_level", level: "medium" },
|
|
{ type: "prompt", message: expect.stringMatching(/health check/i) },
|
|
]);
|
|
expect(child.kill).toHaveBeenCalledOnce();
|
|
});
|
|
|
|
const unexpectedToolEvents = [
|
|
{
|
|
name: "streamed tool call",
|
|
event: {
|
|
type: "message_update",
|
|
assistantMessageEvent: { type: "toolcall_start", contentIndex: 0 },
|
|
},
|
|
},
|
|
{
|
|
name: "tool execution",
|
|
event: { type: "tool_execution_start", toolCallId: "tool-1", toolName: "read" },
|
|
},
|
|
{
|
|
name: "completed message tool call",
|
|
event: {
|
|
type: "message_end",
|
|
message: { role: "assistant", stopReason: "toolUse", content: [{ type: "toolCall" }] },
|
|
},
|
|
},
|
|
{
|
|
name: "turn tool result",
|
|
event: { type: "turn_end", toolResults: [{ role: "toolResult" }] },
|
|
},
|
|
];
|
|
|
|
// Catches Pi/provider regressions that surface a tool capability despite the fixed no-tools argv;
|
|
// accepting agent_end after any such event could hide a mounted-state read or mutation.
|
|
test.each(unexpectedToolEvents)("provider smoke fails closed on an unexpected $name event", async ({ event }) => {
|
|
const child = rpcChild((command, emit) => {
|
|
if (command.type === "set_model" || command.type === "set_thinking_level") {
|
|
emit({ type: "response", id: command.id, success: true });
|
|
}
|
|
if (command.type === "prompt") {
|
|
emit(event);
|
|
emit({ type: "agent_end", messages: [] });
|
|
}
|
|
});
|
|
const smoke = createPiProviderSmoke(loadConfig({}), {
|
|
spawnFn: () => child,
|
|
authProviders: () => new Set(["zai"]),
|
|
readAuthStore: () => '{"zai":{"type":"api_key","key":"test-only"}}',
|
|
readModelsStore: () => undefined,
|
|
});
|
|
|
|
await expect(smoke({
|
|
provider: "zai", model: "glm-5.2", reasoning: "medium", timeoutMs: 750,
|
|
})).rejects.toThrow("Pi provider smoke check failed");
|
|
expect(child.kill).toHaveBeenCalledOnce();
|
|
});
|
|
|
|
// Catches provider errors that are accepted as a successful health check or returned with raw
|
|
// credential/output diagnostics.
|
|
test("provider smoke rejects a failed model turn with a stable non-secret error", async () => {
|
|
const child = rpcChild((command, emit) => {
|
|
if (command.type === "set_model" || command.type === "set_thinking_level") {
|
|
emit({ type: "response", id: command.id, success: true });
|
|
}
|
|
if (command.type === "prompt") {
|
|
emit({
|
|
type: "message_end",
|
|
message: {
|
|
role: "assistant", stopReason: "error",
|
|
errorMessage: '401 {"token":"raw-provider-secret"}',
|
|
},
|
|
});
|
|
emit({ type: "agent_end", messages: [] });
|
|
}
|
|
});
|
|
const smoke = createPiProviderSmoke(loadConfig({}), {
|
|
spawnFn: () => child,
|
|
authProviders: () => new Set(["zai"]),
|
|
readAuthStore: () => '{"zai":{"type":"api_key","key":"test-only"}}',
|
|
readModelsStore: () => undefined,
|
|
});
|
|
|
|
let caught: unknown;
|
|
try {
|
|
await smoke({ provider: "zai", model: "glm-5.2", reasoning: "medium", timeoutMs: 750 });
|
|
} catch (error) {
|
|
caught = error;
|
|
}
|
|
expect(caught).toBeInstanceOf(Error);
|
|
expect((caught as Error).message).toBe("Pi provider smoke check failed");
|
|
expect(String(caught)).not.toContain("raw-provider-secret");
|
|
});
|