fix: harden model catalog projections
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
import { spawnSync } from "node:child_process";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { PostgreSqlContainer } from "@testcontainers/postgresql";
|
||||
import { CamelCasePlugin, Kysely, PostgresDialect, sql } from "kysely";
|
||||
import { Pool } from "pg";
|
||||
@@ -33,7 +34,30 @@ test.skipIf(!dockerAvailable)("PostgreSQL migration enforces one database per wo
|
||||
await upDescriptionGeneration(db);
|
||||
await upSensitiveSuggestionRuns(db);
|
||||
await upAiTokenUsage(db);
|
||||
const historicalDatabaseId = randomUUID();
|
||||
await db.insertInto("workspaceDatabases").values({
|
||||
id: historicalDatabaseId,
|
||||
workspaceId: "migration-history",
|
||||
engine: "postgres",
|
||||
databaseName: "warehouse",
|
||||
schemaName: "public",
|
||||
}).execute();
|
||||
await db.insertInto("descriptionGenerationRuns").values({
|
||||
id: randomUUID(), databaseId: historicalDatabaseId, scope: "all",
|
||||
modelId: "openai-mini", language: "en", status: "completed", total: 1,
|
||||
processed: 1, generated: 1,
|
||||
}).execute();
|
||||
await db.insertInto("sensitiveDataSuggestionRuns").values({
|
||||
id: randomUUID(), databaseId: historicalDatabaseId, scope: "all",
|
||||
modelId: "openai-mini", status: "completed", total: 1,
|
||||
suggestedSensitive: 1,
|
||||
}).execute();
|
||||
await upCanonicalModelIds(db);
|
||||
await expect(db.insertInto("descriptionGenerationRuns").values({
|
||||
id: randomUUID(), databaseId: historicalDatabaseId, scope: "all",
|
||||
modelId: "openai/gpt-5-mini", language: "en", status: "completed", total: 1,
|
||||
processed: 1, generated: 1,
|
||||
}).execute()).resolves.toBeDefined();
|
||||
await sql`CREATE ROLE thothii_catalog_runtime`.execute(db);
|
||||
await upRuntimeSequencePrivileges(db);
|
||||
const sequencePrivilege = await sql<{ allowed: boolean }>`
|
||||
|
||||
@@ -6,6 +6,7 @@ import { join } from "node:path";
|
||||
import path from "node:path";
|
||||
import { createPiModelLister } from "../src/pi/list-models.js";
|
||||
import { loadConfig } from "../src/config.js";
|
||||
import type { RuntimeModel, RuntimeModelCatalog } from "../src/models/runtime-model-catalog.js";
|
||||
|
||||
const FAKE = path.resolve("../harness/tests/fake_pi/fake_pi_rpc.mjs");
|
||||
|
||||
@@ -44,6 +45,36 @@ test("createPiModelLister returns mapped PiModel[] from get_available_models", a
|
||||
}
|
||||
});
|
||||
|
||||
test("catalog listing translates upstream Pi IDs back to canonical model keys", async () => {
|
||||
const script = scriptWith([
|
||||
{ provider: "local", id: "qwen2.5:7b", name: "Upstream label", reasoning: false },
|
||||
]);
|
||||
const model: RuntimeModel = {
|
||||
id: "local/qwen", provider: "local", model: "qwen", label: "Catalog Qwen",
|
||||
upstreamModel: "qwen2.5:7b", endpoint: { baseUrl: "http://ollama:11434/v1" },
|
||||
authentication: { mode: "none" }, sessionAdapter: { mode: "openai_compatible" },
|
||||
session: { reasoning: true, contextWindow: 32768, maxTokens: 8192 },
|
||||
};
|
||||
const modelCatalog: RuntimeModelCatalog = {
|
||||
defaultSession: model.id, defaultMetadataGeneration: null, embedding: null,
|
||||
sessionModels: () => [model], metadataModels: () => [], hasSession: (id) => id === model.id,
|
||||
};
|
||||
try {
|
||||
const lister = createPiModelLister(loadConfig({ THT_HARNESS_DIR: "../harness" }), {
|
||||
...noManagedModels,
|
||||
modelCatalog,
|
||||
loadEnabledModels: enabled("local/qwen2.5:7b"),
|
||||
spawnFn: () => spawn("node", [FAKE, script]) as any,
|
||||
});
|
||||
|
||||
await expect(lister()).resolves.toEqual([{
|
||||
provider: "local", id: "qwen", name: "Catalog Qwen", reasoning: true,
|
||||
}]);
|
||||
} finally {
|
||||
rmSync(path.dirname(script), { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("createPiModelLister caches within ttl (spawns once for two calls)", async () => {
|
||||
const script = scriptWith([{ provider: "zai", id: "glm-5.2", name: "GLM 5.2", reasoning: true }]);
|
||||
try {
|
||||
|
||||
@@ -718,6 +718,54 @@ test.each([["OpenAI", "openai"], ["gemini", "google"]])(
|
||||
},
|
||||
);
|
||||
|
||||
test("set_model translates a canonical catalog key to its upstream Pi model ID", async () => {
|
||||
const root = mkdtempSync(path.join(tmpdir(), "thothii-upstream-model-"));
|
||||
const agentDir = path.join(root, "agent");
|
||||
mkdirSync(agentDir, { mode: 0o700 });
|
||||
writeFileSync(path.join(agentDir, "models.json"), JSON.stringify({
|
||||
providers: {
|
||||
local: {
|
||||
baseUrl: "http://ollama:11434/v1", apiKey: "local",
|
||||
models: [{ id: "qwen2.5:7b" }],
|
||||
},
|
||||
},
|
||||
}), { mode: 0o600 });
|
||||
vi.stubEnv("PI_CODING_AGENT_DIR", agentDir);
|
||||
const child = recordingChild();
|
||||
child.stderr.resume = () => {};
|
||||
child.stdin.write = (data: unknown) => {
|
||||
const request = JSON.parse(String(data));
|
||||
child._writes.push(String(data));
|
||||
if (request.id) {
|
||||
queueMicrotask(() => child.stdout.emit("data", `${JSON.stringify({
|
||||
type: "response", id: request.id, success: true,
|
||||
})}\n`));
|
||||
}
|
||||
return true;
|
||||
};
|
||||
const model: RuntimeModel = {
|
||||
id: "local/qwen", provider: "local", model: "qwen", label: "Qwen",
|
||||
upstreamModel: "qwen2.5:7b", endpoint: { baseUrl: "http://ollama:11434/v1" },
|
||||
authentication: { mode: "none" }, sessionAdapter: { mode: "openai_compatible" },
|
||||
session: { reasoning: false, contextWindow: 32768, maxTokens: 8192 },
|
||||
};
|
||||
const modelCatalog: RuntimeModelCatalog = {
|
||||
defaultSession: model.id, defaultMetadataGeneration: null, embedding: null,
|
||||
sessionModels: () => [model], metadataModels: () => [], hasSession: (id) => id === model.id,
|
||||
};
|
||||
const mgr = new PiProcessManager(loadConfig({ PI_BIN: "/usr/local/bin/pi" }), {
|
||||
modelCatalog, authProviders: () => new Set(), spawnFn: () => child as any,
|
||||
});
|
||||
try {
|
||||
await mgr.spawnFor("upstream-model", { provider: "local", model: "qwen" });
|
||||
expect(child._writes.join("")).toContain('"modelId":"qwen2.5:7b"');
|
||||
} finally {
|
||||
mgr.teardown("upstream-model");
|
||||
vi.unstubAllEnvs();
|
||||
rmSync(root, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test.each(["installation-local", "private-compatible"])(
|
||||
"provider %s configured with a literal apiKey spawns without a managed key",
|
||||
async (provider) => {
|
||||
|
||||
@@ -300,11 +300,22 @@ test("provider smoke makes one configured request from an isolated no-capability
|
||||
});
|
||||
}
|
||||
});
|
||||
const smokeModel: RuntimeModel = {
|
||||
id: "zai/catalog-glm", provider: "zai", model: "catalog-glm", label: "GLM",
|
||||
upstreamModel: "glm-5.2", authentication: { mode: "pi_auth" },
|
||||
sessionAdapter: { mode: "pi_builtin" }, session: { reasoning: true },
|
||||
};
|
||||
const smokeCatalog: RuntimeModelCatalog = {
|
||||
defaultSession: smokeModel.id, defaultMetadataGeneration: null, embedding: null,
|
||||
sessionModels: () => [smokeModel], metadataModels: () => [],
|
||||
hasSession: (id) => id === smokeModel.id,
|
||||
};
|
||||
const smoke = createPiProviderSmoke(loadConfig({
|
||||
THT_HARNESS_DIR: "/app/harness",
|
||||
PI_BIN: "/usr/local/bin/pi",
|
||||
THT_DATA_ROOT: "/mounted-workflow-state",
|
||||
}), {
|
||||
modelCatalog: smokeCatalog,
|
||||
spawnFn: (...args) => {
|
||||
spawns.push(args);
|
||||
expect(args[2].cwd).not.toBe("/app/harness");
|
||||
@@ -324,7 +335,7 @@ test("provider smoke makes one configured request from an isolated no-capability
|
||||
});
|
||||
|
||||
await expect(smoke({
|
||||
provider: "zai", model: "glm-5.2", reasoning: "medium", timeoutMs: 750,
|
||||
provider: "zai", model: "catalog-glm", reasoning: "medium", timeoutMs: 750,
|
||||
})).resolves.toBeUndefined();
|
||||
expect(spawns).toHaveLength(1);
|
||||
expect(spawns[0][0]).toBe("/usr/local/bin/pi");
|
||||
|
||||
@@ -73,6 +73,7 @@ test("derives the internal Qdrant and Ollama runtime shape from workspace v4 plu
|
||||
vector: { engine: "qdrant", base_url: "http://qdrant:6333", collection: "psd-clinical" },
|
||||
embeddings: {
|
||||
provider: "ollama_internal", base_url: "http://embedding:11434",
|
||||
id: "ollama/qwen3-embedding:0.6b",
|
||||
model: "qwen3-embedding:0.6b", dimensions: 1024,
|
||||
},
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user