mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-13 15:55:46 +01:00
285 lines
21 KiB
TypeScript
285 lines
21 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
import { mkdtemp, realpath, rm } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { and, eq, sql } from "drizzle-orm";
|
|
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
|
|
import { agents, companies, companyMemberships, createDb, environmentLeases, environments, heartbeatRuns, plugins } from "@paperclipai/db";
|
|
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
|
|
import { aiConnectionService } from "../services/ai-connections.js";
|
|
import { prepareManagedAiRuntime } from "../services/ai-connection-runtime.js";
|
|
import { registerPluginTaskCredentialHandoff, settlePluginTaskCredentialHandoff, syncPluginTaskCredentialHandoff } from "../services/plugin-task-credential-writeback.js";
|
|
import { executeEnvironmentTask, settleDispatchedEnvironmentTask } from "../services/environment-task-runtime.js";
|
|
import { secretService } from "../services/secrets.js";
|
|
import { syncManagedAiCredentials } from "../services/managed-ai-credential-writeback.js";
|
|
import { appendNativeCredentialUpdate } from "../services/native-runtime/native-credential-updates.js";
|
|
import type { PrpEvent } from "../vendor/paperclip-runner/index.js";
|
|
import type { PluginWorkerManager } from "../services/plugin-worker-manager.js";
|
|
|
|
vi.mock("../services/subscription-refresh.js", () => ({ refreshSubscriptionConnection: async () => {} }));
|
|
const support = await getEmbeddedPostgresTestSupport();
|
|
const suite = support.supported ? describe : describe.skip;
|
|
const auth = (token: string, year = 2020) => JSON.stringify({ tokens: { account_id: "fixture-account", access_token: "fixture-access", refresh_token: token }, last_refresh: `${year}-01-01T00:00:00Z` });
|
|
|
|
suite("durable plugin task credential handoff", () => {
|
|
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>, db: ReturnType<typeof createDb>, home: string;
|
|
beforeAll(async () => {
|
|
home = await realpath(await mkdtemp(join(tmpdir(), "plugin-task-secret-")));
|
|
vi.stubEnv("PAPERCLIP_HOME", home); vi.stubEnv("PAPERCLIP_INSTANCE_ID", "plugin-task-fixture");
|
|
database = await startEmbeddedPostgresTestDatabase("plugin-task-credential-"); db = createDb(database.connectionString);
|
|
}, 90_000);
|
|
afterAll(async () => { await database?.cleanup(); vi.unstubAllEnvs(); if (home) await rm(home, { recursive: true, force: true }); });
|
|
async function fixture(dispatched = true, provider: "codex" | "grok" = "codex") {
|
|
const selectedProvider = provider === "codex" ? "openai" as const : "xai" as const;
|
|
const credential = (token: string, year = 2020) => provider === "codex" ? auth(token, year) : JSON.stringify({
|
|
"https://issuer.example.test::11111111-1111-4111-8111-111111111111": {
|
|
key: "fixture-access", refresh_token: token, expires_at: `${year}-01-01T00:00:00Z`,
|
|
},
|
|
});
|
|
const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(), pluginId = randomUUID(), environmentId = randomUUID(), leaseId = randomUUID();
|
|
const owner = `fixture-${randomUUID()}`;
|
|
await db.insert(companies).values({ id: companyId, name: "Task credential fixture", issuePrefix: `C${companyId.slice(0, 6)}` });
|
|
await db.insert(agents).values({ id: agentId, companyId, name: "Fixture", adapterType: "paperclip_runner" });
|
|
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
|
const service = aiConnectionService(db);
|
|
await service.save(companyId, owner, { provider: selectedProvider, method: "subscription", ownership: "personal", name: "Fixture account", loginSessionId: "fixture", agentIds: [], allAgents: true }, credential("fixture-original"));
|
|
const [selection] = await service.quotaAccounts(companyId, owner, true);
|
|
const managed = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "paperclip_runner",
|
|
binding: { mode: "responsible_user", provider: selectedProvider, method: "subscription" },
|
|
config: provider === "codex" ? { provider: "codex" } : { provider: "acpx", acpxAgent: "grok" } });
|
|
expect(managed.credentialWriteback).toBeDefined();
|
|
await db.insert(environments).values({ id: environmentId, name: `Fixture ${environmentId}`, driver: "plugin", status: "active", config: { pluginKey: pluginId, driverKey: "tasks", driverConfig: {} } });
|
|
await db.insert(plugins).values({ id: pluginId, pluginKey: pluginId, packageName: "fixture-tasks", version: "1.0.0", status: "ready", manifestJson: {
|
|
capabilities: ["environment.drivers.register"], environmentDrivers: [{ driverKey: "tasks", supportsTasks: true }],
|
|
} as never });
|
|
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "manual", status: "running" });
|
|
await db.insert(environmentLeases).values({ id: leaseId, companyId, environmentId, heartbeatRunId: runId, provider: `plugin:${pluginId}:tasks`, providerLeaseId: "fixture-task", status: "active",
|
|
metadata: { driver: "plugin", pluginId, driverKey: "tasks", runnerTaskLease: true, runnerTaskDispatched: dispatched, runnerTaskExecutionStopped: false } });
|
|
await registerPluginTaskCredentialHandoff(db, { companyId, leaseId }, managed.sessionIdentity, managed.credentialWriteback!);
|
|
// Recovery must not depend on the original controller or its private home.
|
|
await rm(managed.home!, { recursive: true, force: true });
|
|
const calls: string[] = [];
|
|
let failExport = false, exported: string | null = auth("fixture-refreshed", 2021);
|
|
const workers = { getWorker: () => ({ supportedMethods: ["environmentTask"] }),
|
|
call: async (_plugin: string, _method: string, params: any) => {
|
|
calls.push(params.operation.kind);
|
|
expect(params.companyId).toBe(companyId); expect(params.taskId).toBe("fixture-task");
|
|
if (params.operation.kind === "status") return { kind: "status", taskId: "fixture-task", phase: "cancelled", executionStopped: true };
|
|
if (params.operation.kind === "credentials") {
|
|
expect(params.operation.accountIdentity).toBe(managed.sessionIdentity);
|
|
if (failExport) throw new Error("fixture export unavailable");
|
|
return { kind: "credentials", taskId: "fixture-task", credential: exported === null ? null : { type: "codex-auth-json", authJson: exported } };
|
|
}
|
|
return { kind: "accepted", taskId: "fixture-task" };
|
|
} } as unknown as PluginWorkerManager;
|
|
const lease = () => db.select().from(environmentLeases).where(and(eq(environmentLeases.id, leaseId), eq(environmentLeases.companyId, companyId))).then(rows => rows[0]!);
|
|
return { accountCredential: credential, companyId, agentId, runId, environmentId, leaseId, owner, selection: selection!, managed, workers, calls, lease,
|
|
synchronize: () => syncPluginTaskCredentialHandoff(db, { companyId, leaseId, expectedTaskId: "fixture-task" }, async accountIdentity => {
|
|
const result = await executeEnvironmentTask(db, workers, { companyId, leaseId, expectedTaskId: "fixture-task", operation: { kind: "credentials", accountIdentity } });
|
|
if (result.kind !== "credentials") throw new Error("Unexpected credential result");
|
|
return result.credential;
|
|
}),
|
|
exportCredential(value: string | null) { exported = value; },
|
|
failExport(value: boolean) { failExport = value; }, credential: () => service.credential(selection!),
|
|
settle: (executionStopObserved = false) => settleDispatchedEnvironmentTask(db, workers, { companyId, id: leaseId }, { executionStopObserved }) };
|
|
}
|
|
it.each(["codex", "grok"] as const)("writes managed %s refreshes without any plugin or invocation home", async provider => {
|
|
const f = await fixture(true, provider);
|
|
const read = vi.fn(async () => f.accountCredential("fixture-refreshed", 2021));
|
|
expect(await syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider, read)).toBe("rotated");
|
|
expect(await f.credential()).toBe(f.accountCredential("fixture-refreshed", 2021));
|
|
expect(await syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider, read)).toBe("unchanged");
|
|
await secretService(db).rotate(f.managed.credentialWriteback!.secretId, { value: f.accountCredential("fixture-reconnected", 2022) }, { userId: f.owner });
|
|
read.mockClear();
|
|
expect(await syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider, read)).toBe("binding_changed");
|
|
expect(read).not.toHaveBeenCalled();
|
|
expect(await f.credential()).toBe(f.accountCredential("fixture-reconnected", 2022));
|
|
expect(f.calls).toEqual([]);
|
|
});
|
|
it.each((["codex", "grok"] as const).flatMap(provider => ["{partial", "{}"].map(snapshot => [provider, snapshot] as const)))(
|
|
"rejects unusable %s snapshots for retry: %s", async (provider, snapshot) => {
|
|
const f = await fixture(true, provider);
|
|
await expect(syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider, async () => snapshot))
|
|
.rejects.toThrow("snapshot invalid");
|
|
expect(await f.credential()).toBe(f.accountCredential("fixture-original"));
|
|
expect(await syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider,
|
|
async () => f.accountCredential("fixture-refreshed", 2021))).toBe("rotated");
|
|
},
|
|
);
|
|
it.each(["codex", "grok"] as const)("rechecks %s authority after the exporter yields", async provider => {
|
|
const f = await fixture(true, provider);
|
|
const result = await syncManagedAiCredentials(db, f.companyId, f.managed.credentialWriteback!, provider, async () => {
|
|
await secretService(db).rotate(f.managed.credentialWriteback!.secretId, { value: f.accountCredential("fixture-reconnected", 2022) }, { userId: f.owner });
|
|
return f.accountCredential("fixture-refreshed", 2021);
|
|
});
|
|
expect(result).toBe("binding_changed");
|
|
expect(await f.credential()).toBe(f.accountCredential("fixture-reconnected", 2022));
|
|
});
|
|
it("replays a validated PRP refresh into the tenant secret after export failure", async () => {
|
|
const f = await fixture(), normalizedSessionId = randomUUID(), runnerInstanceId = randomUUID();
|
|
await db.update(heartbeatRuns).set({ status: "succeeded", runtimeMode: "native", nativeSessionId: normalizedSessionId, runnerInstanceId })
|
|
.where(eq(heartbeatRuns.id, f.runId));
|
|
await db.update(environmentLeases).set({ status: "retained" }).where(eq(environmentLeases.id, f.leaseId));
|
|
const event: PrpEvent = {
|
|
schema: "paperclip.prp.event.v3", schemaVersion: 3, eventType: "auth.credentials.updated", priority: 0,
|
|
sourceKind: "runner", sourceInstanceId: runnerInstanceId, sourceEventId: "live-refresh", sourceSeq: 1,
|
|
runId: f.runId, normalizedSessionId, emittedAt: "2020-01-01T00:00:00Z",
|
|
payload: { schema: "paperclip.auth.credentials.updated.v1", provider: "codex", bindingId: "harness-auth", revision: 1, detection: "credential_store" },
|
|
};
|
|
const notification = { companyId: f.companyId, agentId: f.agentId, runnerInstanceId, normalizedSessionId,
|
|
event, synchronizeCredentials: f.synchronize };
|
|
f.failExport(true);
|
|
await expect(appendNativeCredentialUpdate(db, notification)).rejects.toThrow();
|
|
expect(await f.credential()).toBe(auth("fixture-original"));
|
|
f.failExport(false);
|
|
await appendNativeCredentialUpdate(db, notification);
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect(f.calls).toEqual(["credentials", "credentials"]);
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
});
|
|
it("writes repeated live refreshes without stopping or completing the warm task", async () => {
|
|
const f = await fixture();
|
|
expect(await f.synchronize()).toBe("rotated");
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect(await f.synchronize()).toBe("unchanged");
|
|
f.exportCredential(auth("fixture-next-refresh", 2022));
|
|
expect(await f.synchronize()).toBe("rotated");
|
|
expect(await f.credential()).toBe(auth("fixture-next-refresh", 2022));
|
|
expect(f.calls).toEqual(["credentials", "credentials", "credentials"]);
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
expect(await f.settle()).toBe("completed");
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete", outcome: "unchanged" });
|
|
});
|
|
it("keeps live export failures replayable and rejects missing snapshots", async () => {
|
|
const f = await fixture(); f.failExport(true);
|
|
await expect(f.synchronize()).rejects.toThrow();
|
|
f.failExport(false); f.exportCredential(null);
|
|
await expect(f.synchronize()).rejects.toThrow("snapshot unavailable");
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
expect(await f.credential()).toBe(auth("fixture-original"));
|
|
f.exportCredential(auth("fixture-refreshed", 2021));
|
|
expect(await f.synchronize()).toBe("rotated");
|
|
expect(f.calls).not.toContain("complete");
|
|
});
|
|
it("fences live writeback after explicit reconnect and physical task replacement", async () => {
|
|
const f = await fixture();
|
|
await secretService(db).rotate(f.managed.credentialWriteback!.secretId, { value: auth("fixture-reconnected", 2022) }, { userId: f.owner });
|
|
expect(await f.synchronize()).toBe("binding_changed");
|
|
expect(f.calls).toEqual([]);
|
|
expect(await f.credential()).toBe(auth("fixture-reconnected", 2022));
|
|
await db.update(environmentLeases).set({ providerLeaseId: "replacement-task" }).where(eq(environmentLeases.id, f.leaseId));
|
|
await expect(f.synchronize()).rejects.toThrow("lease unavailable");
|
|
expect(f.calls).toEqual([]);
|
|
});
|
|
it("records a null retirement export without treating it as a secret rotation", async () => {
|
|
const f = await fixture(); f.exportCredential(null);
|
|
expect(await f.settle()).toBe("completed");
|
|
expect(f.calls).toContain("credentials");
|
|
expect(await f.credential()).toBe(auth("fixture-original"));
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete", outcome: "no_credentials" });
|
|
});
|
|
it("exports a dispatched stopped task even after the dispatch flag was cleared", async () => {
|
|
const f = await fixture();
|
|
const metadata = { ...(await f.lease()).metadata, runnerTaskDispatched: false, runnerTaskExecutionStopped: true,
|
|
runnerTaskIdentity: { runnerId: "fixture-runner" } };
|
|
await db.update(environmentLeases).set({ metadata }).where(eq(environmentLeases.id, f.leaseId));
|
|
expect(await f.settle()).toBe("stopped");
|
|
expect(f.calls).toEqual(["credentials"]);
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete", outcome: "rotated" });
|
|
});
|
|
it("persists refreshed credentials before completing a stopped task, after controller files disappear", async () => {
|
|
const f = await fixture();
|
|
expect(await f.settle()).toBe("completed");
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect(f.calls.indexOf("credentials")).toBeLessThan(f.calls.indexOf("complete"));
|
|
const saved = await f.lease();
|
|
expect(saved.metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" });
|
|
expect(JSON.stringify(saved.metadata)).not.toContain("fixture-refreshed");
|
|
expect(JSON.stringify(saved.metadata)).not.toContain(f.managed.home!);
|
|
});
|
|
it("writes back credentials before completing a task whose stop was already observed, without stopping it again", async () => {
|
|
const f = await fixture();
|
|
expect(await f.settle(true)).toBe("completed");
|
|
expect(f.calls).toEqual(["status", "credentials", "complete"]);
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" });
|
|
});
|
|
it("leaves an observed-stop task for cleanup when credential writeback fails", async () => {
|
|
const f = await fixture(); f.failExport(true);
|
|
await expect(f.settle(true)).rejects.toThrow("termination is unconfirmed");
|
|
expect(f.calls).toEqual(["status", "credentials"]);
|
|
expect(await f.lease()).toMatchObject({ status: "pending_cleanup", cleanupStatus: "failed" });
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
f.failExport(false);
|
|
expect(await f.settle()).toBe("completed");
|
|
expect(f.calls.slice(2)).toEqual(["stop", "status", "credentials", "complete"]);
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
});
|
|
it("keeps a failed transfer pending and recovers through a fresh cleanup after run and environment deletion", async () => {
|
|
const f = await fixture(); f.failExport(true);
|
|
await expect(f.settle()).rejects.toThrow("termination is unconfirmed");
|
|
expect(f.calls).not.toContain("complete");
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
expect(await f.credential()).toBe(auth("fixture-original"));
|
|
await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId));
|
|
await db.delete(environments).where(eq(environments.id, f.environmentId));
|
|
f.failExport(false);
|
|
expect(await f.settle()).toBe("completed");
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect(f.calls.at(-1)).toBe("complete");
|
|
});
|
|
it("retains the Runner credential after a database save failure until a fresh receiver commits", async () => {
|
|
const f = await fixture();
|
|
await db.execute(sql.raw(`create function reject_runner_credential_save() returns trigger language plpgsql as $$ begin
|
|
if new.secret_id = '${f.managed.credentialWriteback!.secretId}'::uuid then raise exception 'fixture receiver save unavailable'; end if; return new; end $$;
|
|
create constraint trigger reject_runner_credential_save after insert on company_secret_versions
|
|
deferrable initially deferred for each row execute function reject_runner_credential_save();`));
|
|
try {
|
|
await expect(f.settle()).rejects.toThrow();
|
|
expect(f.calls).toContain("credentials");
|
|
expect(f.calls).not.toContain("complete");
|
|
expect(await f.credential()).toBe(auth("fixture-original"));
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
} finally { await db.execute(sql`drop trigger reject_runner_credential_save on company_secret_versions; drop function reject_runner_credential_save()`); }
|
|
const receiver = createDb(database.connectionString);
|
|
expect(await settleDispatchedEnvironmentTask(receiver, f.workers, { companyId: f.companyId, id: f.leaseId })).toBe("completed");
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect(f.calls.filter(call => call === "credentials")).toHaveLength(2);
|
|
expect(f.calls.at(-1)).toBe("complete");
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" });
|
|
});
|
|
it("does not overwrite an explicit account replacement", async () => {
|
|
const f = await fixture();
|
|
await secretService(db).rotate(f.managed.credentialWriteback!.secretId, { value: auth("fixture-reconnected", 2022) }, { userId: f.owner });
|
|
expect(await f.settle()).toBe("completed");
|
|
expect(f.calls).not.toContain("credentials");
|
|
expect(await f.credential()).toBe(auth("fixture-reconnected", 2022));
|
|
});
|
|
it("allows concurrent cleanup to complete the same durable handoff", async () => {
|
|
const f = await fixture();
|
|
expect(await Promise.all([f.settle(), f.settle()])).not.toContain(false);
|
|
expect(await f.credential()).toBe(auth("fixture-refreshed", 2021));
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" });
|
|
});
|
|
it("does not mark a replacement task's identical credential handoff complete", async () => {
|
|
const f = await fixture(false), prior = await f.lease();
|
|
await db.update(environmentLeases).set({ providerLeaseId: "replacement-task" }).where(eq(environmentLeases.id, f.leaseId));
|
|
await expect(settlePluginTaskCredentialHandoff(db, { companyId: f.companyId, leaseId: f.leaseId, expectedTaskId: "fixture-task" },
|
|
prior.metadata!, async () => { throw new Error("Undispatched task must not export credentials"); })).rejects.toThrow("handoff changed");
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" });
|
|
});
|
|
|
|
it("retires an undispatched task without exporting credentials and blocks a late submit", async () => {
|
|
const f = await fixture(false);
|
|
expect(await f.settle()).toBe(false);
|
|
expect(f.calls).toEqual([]);
|
|
expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" });
|
|
await expect(executeEnvironmentTask(db, f.workers, { companyId: f.companyId, leaseId: f.leaseId, operation: {
|
|
kind: "submit", projectIds: [], bootstrapTicket: "fixture-ticket", runner: { harness: "codex", protocolMin: 1, protocolMax: 3,
|
|
runnerId: "runner", leaseId: f.leaseId, runId: f.runId, sessionId: "session", turnId: "turn", itemId: "item" },
|
|
} })).rejects.toThrow("not active");
|
|
expect(f.calls).toEqual([]);
|
|
});
|
|
});
|