Wire task runs to provider gateway
This commit is contained in:
ZacharyZhang-NY committed
1 parent
273f69eb05
commit
a27658b56f
19 files changed
+591
-11
No files matched your search
@@ -0,0 +1,206 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { describe, expect, test } from "vitest";
|
||||
import { DatabaseConnection } from "../../src/main/database/connection";
|
||||
import { OrganizationService } from "../../src/main/services/organization-service";
|
||||
import { ProviderService } from "../../src/main/services/provider-service";
|
||||
import { providerAdapters, type ProviderAdapterRegistry } from "../../src/main/services/provider-adapters";
|
||||
import { SnapshotService } from "../../src/main/services/snapshot-service";
|
||||
import { TaskExecutionService } from "../../src/main/services/task-execution-service";
|
||||
import { RISK_LEVEL, TASK_PRIORITY, TASK_STATUS } from "../../src/shared/domain";
|
||||
|
||||
function createConnection(): DatabaseConnection {
|
||||
return new DatabaseConnection(join(tmpdir(), `agentcompany-new-task-execution-${randomUUID()}.db`));
|
||||
}
|
||||
|
||||
function openAiProviderId(connection: DatabaseConnection, orgId: string): string {
|
||||
return (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'openai'")
|
||||
.get(orgId) as { id: string }).id;
|
||||
}
|
||||
|
||||
const testVault = {
|
||||
encryptSecret: (secret: string) => Buffer.from(secret, "utf8"),
|
||||
decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8"),
|
||||
};
|
||||
|
||||
describe("task execution service", () => {
|
||||
test("runs an assigned task through the provider gateway", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
const snapshot = new SnapshotService(connection.db);
|
||||
let capturedPrompt = "";
|
||||
const adapters: ProviderAdapterRegistry = {
|
||||
...providerAdapters,
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
createResponse: async ({ prompt, model }) => {
|
||||
capturedPrompt = prompt;
|
||||
return {
|
||||
text: `executed:${model}`,
|
||||
rawResponseJson: JSON.stringify({ id: "resp-task-1" }),
|
||||
rateLimitState: null,
|
||||
};
|
||||
},
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
const execution = new TaskExecutionService(connection.db, providers);
|
||||
let orgId = "";
|
||||
let providerId = "";
|
||||
let taskId = "";
|
||||
|
||||
connection.transaction(() => {
|
||||
orgId = organization.createOrganization({
|
||||
organizationName: "Acme Operations",
|
||||
employeeName: "Dana Lee",
|
||||
mission: "Coordinate agents with auditable approvals.",
|
||||
goal: "Ship a verified control plane.",
|
||||
});
|
||||
providerId = openAiProviderId(connection, orgId);
|
||||
providers.saveCredential({ orgId, providerId, secret: "sk-live-provider", model: "gpt-5.4" });
|
||||
const agentId = organization.createAgent({
|
||||
orgId,
|
||||
name: "Ada",
|
||||
role: "Execution lead",
|
||||
department: "Engineering",
|
||||
modelProviderId: providerId,
|
||||
modelId: "gpt-5.4",
|
||||
budgetCents: 10000,
|
||||
managerAgentId: null,
|
||||
});
|
||||
taskId = organization.createTask({
|
||||
orgId,
|
||||
title: "Draft launch checklist",
|
||||
description: "Prepare an auditable checklist for release readiness.",
|
||||
status: TASK_STATUS.IN_PROGRESS,
|
||||
assigneeAgentId: agentId,
|
||||
risk: RISK_LEVEL.MEDIUM,
|
||||
projectId: null,
|
||||
goalId: null,
|
||||
workspaceId: null,
|
||||
priority: TASK_PRIORITY.HIGH,
|
||||
requiresFilesystem: false,
|
||||
});
|
||||
});
|
||||
|
||||
const runId = await execution.runTask({ orgId, taskId });
|
||||
const state = snapshot.load();
|
||||
const run = state.taskRuns.find((item) => item.id === runId);
|
||||
const request = state.providerRequests.find((item) => item.id === run?.providerRequestId);
|
||||
|
||||
expect(run).toMatchObject({ taskId, agentId: expect.any(String), providerId, state: "completed", errorMessage: null });
|
||||
expect(request).toMatchObject({ providerId, state: "succeeded", responseText: "executed:gpt-5.4" });
|
||||
expect(capturedPrompt).toContain("Task: Draft launch checklist");
|
||||
expect(capturedPrompt).toContain("Assignee: Ada (Execution lead)");
|
||||
expect(state.auditEvents.some((event) => event.entityType === "task_run" && event.action === "completed")).toBe(true);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
test("marks task runs failed when provider execution fails", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
const snapshot = new SnapshotService(connection.db);
|
||||
const adapters: ProviderAdapterRegistry = {
|
||||
...providerAdapters,
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
createResponse: async () => {
|
||||
throw new Error("provider rejected task");
|
||||
},
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
const execution = new TaskExecutionService(connection.db, providers);
|
||||
let orgId = "";
|
||||
let taskId = "";
|
||||
|
||||
connection.transaction(() => {
|
||||
orgId = organization.createOrganization({
|
||||
organizationName: "Acme Operations",
|
||||
employeeName: "Dana Lee",
|
||||
mission: "Coordinate agents with auditable approvals.",
|
||||
goal: "Ship a verified control plane.",
|
||||
});
|
||||
const providerId = openAiProviderId(connection, orgId);
|
||||
providers.saveCredential({ orgId, providerId, secret: "sk-live-provider", model: "gpt-5.4" });
|
||||
const agentId = organization.createAgent({
|
||||
orgId,
|
||||
name: "Ada",
|
||||
role: "Execution lead",
|
||||
department: "Engineering",
|
||||
modelProviderId: providerId,
|
||||
modelId: "gpt-5.4",
|
||||
budgetCents: 10000,
|
||||
managerAgentId: null,
|
||||
});
|
||||
taskId = organization.createTask({
|
||||
orgId,
|
||||
title: "Draft launch checklist",
|
||||
description: "Prepare an auditable checklist for release readiness.",
|
||||
status: TASK_STATUS.IN_PROGRESS,
|
||||
assigneeAgentId: agentId,
|
||||
risk: RISK_LEVEL.MEDIUM,
|
||||
projectId: null,
|
||||
goalId: null,
|
||||
workspaceId: null,
|
||||
priority: TASK_PRIORITY.HIGH,
|
||||
requiresFilesystem: false,
|
||||
});
|
||||
});
|
||||
|
||||
await expect(execution.runTask({ orgId, taskId })).rejects.toThrow("provider rejected task");
|
||||
const state = snapshot.load();
|
||||
const run = state.taskRuns[0];
|
||||
const request = state.providerRequests.find((item) => item.id === run?.providerRequestId);
|
||||
|
||||
expect(run).toMatchObject({ taskId, state: "failed", errorMessage: "provider rejected task" });
|
||||
expect(request).toMatchObject({ state: "failed", errorMessage: "provider rejected task" });
|
||||
expect(state.auditEvents.some((event) => event.entityType === "task_run" && event.action === "failed")).toBe(true);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
test("requires task assignees to have a model provider", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
const execution = new TaskExecutionService(connection.db);
|
||||
let orgId = "";
|
||||
let taskId = "";
|
||||
|
||||
connection.transaction(() => {
|
||||
orgId = organization.createOrganization({
|
||||
organizationName: "Acme Operations",
|
||||
employeeName: "Dana Lee",
|
||||
mission: "Coordinate agents with auditable approvals.",
|
||||
goal: "Ship a verified control plane.",
|
||||
});
|
||||
const agentId = organization.createAgent({
|
||||
orgId,
|
||||
name: "Ada",
|
||||
role: "Execution lead",
|
||||
department: "Engineering",
|
||||
modelProviderId: null,
|
||||
modelId: null,
|
||||
budgetCents: 10000,
|
||||
managerAgentId: null,
|
||||
});
|
||||
taskId = organization.createTask({
|
||||
orgId,
|
||||
title: "Draft launch checklist",
|
||||
description: "Prepare an auditable checklist for release readiness.",
|
||||
status: TASK_STATUS.IN_PROGRESS,
|
||||
assigneeAgentId: agentId,
|
||||
risk: RISK_LEVEL.MEDIUM,
|
||||
projectId: null,
|
||||
goalId: null,
|
||||
workspaceId: null,
|
||||
priority: TASK_PRIORITY.HIGH,
|
||||
requiresFilesystem: false,
|
||||
});
|
||||
});
|
||||
|
||||
await expect(execution.runTask({ orgId, taskId })).rejects.toThrow("Task assignee requires a model provider.");
|
||||
expect(connection.db.prepare("select count(*) as count from task_runs").get()).toEqual({ count: 0 });
|
||||
connection.close();
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user