Merge accounting lifecycle fences and durable OAuth recovery

Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
Devin FoleyandPaperclip committed 2026-10-03 23:45:16 -07:00
commit 8f10b0b7c5
20 files changed
+664 -177

No files matched your search

+19 -1
View File
@@ -177,7 +177,25 @@ The mutation gate retains partial results and per-run reports/logs under
Failed native turns persist observed run-delta usage before semantic-result
finalization. A matching terminal failure, cancellation, or interruption closes
that receipt; missing terminal evidence and session-only totals remain pending.
that usage snapshot; ledger acknowledgement still waits until the native coordinator
has a result or a terminal failure. A retryable failed heartbeat keeps its
reservation and can resume. Replacing its recorder restores the native run's
cumulative snapshot; recovery does not add the same tokens twice, and an empty
final result cannot erase earlier observed usage. Missing terminal evidence and
session-only totals remain pending.
Budget admission completes before the final dispatch ownership check. A rejected
handoff releases a newly created hold using explicit evidence that the adapter
was never entered. Process adapter configuration and spawn failures likewise
return pre-provider proof; failures after a child exists cannot claim it.
Scheduled retry gates acquire company, issue, then run locks, in the same order
as accounting and attribution writes. Delayed budget cancellation rechecks the
policy version and current blocking state when it writes durable stop intent,
under the company admission lock. External shutdown happens after that lock is
released so the provider can save its final receipt. Raising a budget before the
stop is claimed preserves newly admitted work, including previously queued rows;
a run with an already claimed stop cannot enter provider work after the grant.
For multiple provider attempts, known prices remain in the exact spend total
when another attempt is unpriced. The aggregate stays marked unpriced, so the
configured unknown-price policy still controls whether new work may start.
+12
View File
@@ -540,3 +540,15 @@ the private runtime home for retry; cleanup deletes it only after a successful
transaction or an intentional freshness/authorization discard. These retained
homes are recovery evidence, not a background replay queue; persistent database
or lock failures still require operator intervention.
Quota OAuth replacement tokens are encrypted with the instance secrets master
key and fsynced under `<instance-root>/quota-credential-recovery/<company-id>/`
before vault, grant, or activity writes. Storage and encryption are checked before
exchanging a single-use refresh token. Failed saves retry without another OAuth
exchange; persistent failures keep the encrypted record for the next quota poll,
including after a restart. The journal is removed only after database commit (or
when an authorized newer credential makes it obsolete). Replay checks the grant,
connection, secret, and original credential fingerprint and cannot reactivate a
revoked grant or overwrite a reconnect. Back up this directory and the instance
secrets key with the instance data. Loss of durable storage while receiving a
provider token can still require reconnecting the account.
@@ -176,3 +176,32 @@ describe("incremental protocol accounting", () => {
expect(second.usage).toMatchObject(Object.fromEntries(Object.entries(fixture.tokens!).map(([key, count]) => [key, count * multiplier])));
});
});
describe("process adapter bootstrap accounting proof", () => {
const context = (config: Record<string, unknown>): AdapterExecutionContext => ({
runId: `bootstrap-${crypto.randomUUID()}`,
agent: { id: "agent", companyId: "company", name: "Process", adapterType: "process", adapterConfig: {} },
runtime: { sessionId: null, sessionParams: null, sessionDisplayId: null, taskKey: null },
config, context: {}, onLog: async () => {},
});
it.each([{}, { command: "/nonexistent-paperclip-accounting-command" },
{ command: process.execPath, cwd: "/nonexistent-paperclip-accounting-directory" }])(
"proves no provider work for a failure before spawn: %j", async (config) => {
const { execute } = await import("../adapters/process/execute.js");
const onSpawn = vi.fn();
const result = await execute({ ...context(config), onSpawn });
expect(result).toMatchObject({ exitCode: 1, executionRecovery: { kind: "bootstrap", providerWorkStarted: false } });
expect(result.errorMessage).toBeTruthy();
expect(onSpawn).not.toHaveBeenCalled();
},
);
it("does not invent bootstrap proof when a spawned process fails", async () => {
const { execute } = await import("../adapters/process/execute.js");
const onSpawn = vi.fn(async () => { throw new Error("Metadata persistence failed"); });
const result = await execute({ ...context({ command: process.execPath, args: ["-e", "process.exit(7)"] }), onSpawn });
expect(onSpawn).toHaveBeenCalledOnce();
expect(result.exitCode).toBe(7);
expect(result.executionRecovery).toBeUndefined();
});
});
@@ -11,6 +11,7 @@ import { and, eq, sql } from "drizzle-orm";
import { createDb, companies, agents, heartbeatRuns, companyMemberships, connectionGrants, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests, companySecrets, companySecretVersions, userSecretDefinitions, principalPermissionGrants } from "@paperclipai/db";
import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres";
import { fetchCompanyQuotaWindows } from "../services/quota-windows.js";
import { quotaCredentialRecoveryPath, quotaCredentialRecovery, quotaCredentialHash } from "../services/quota-credential-recovery.js";
import { aiConnectionService } from "../services/ai-connections.js";
import * as codexAdapter from "@paperclipai/adapter-codex-local/server";
import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
@@ -104,6 +105,92 @@ describe("managed AI connections", () => {
} finally { request.mockRestore(); }
});
it.each(["rotation", "grant", "activity", "commit"] as const)("recovers exchanged OAuth tokens after %s persistence failure without another exchange", async (failure) => {
const owner = `quota-rollback-${failure}`;
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const original = JSON.stringify({ tokens: { access_token: `expired-${failure}`, refresh_token: "single-use", account_id: owner } });
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
const [row] = await service.quotaAccounts(companyId, owner);
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
const table = failure === "grant" ? "connection_grants" : failure === "activity" ? "activity_log" : "company_secret_versions";
const condition = failure === "grant" ? `new.id = '${row.grant.id}'::uuid` : failure === "activity" ? `new.action = 'ai_connection.credential_refreshed'` : `new.secret_id = '${ref.secretId}'::uuid`;
await db.execute(sql.raw(`create function quota_save_failure() returns trigger language plpgsql as $$ begin
if ${condition} then raise exception 'injected quota save failure'; end if; return new; end $$;
create ${failure === "commit" ? "constraint " : ""}trigger quota_save_failure ${failure === "commit" ? "after" : "before"} ${failure === "grant" ? "update" : "insert"} on ${table}
${failure === "commit" ? "deferrable initially deferred" : ""} for each row execute function quota_save_failure();`));
const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` }));
const file = quotaCredentialRecoveryPath(companyId, row.grant.id);
let triggerDropped = false;
try {
await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow();
expect(request).toHaveBeenCalledTimes(1);
expect(await service.credential(row)).toBe(original);
const encrypted = await readFile(file, "utf8");
expect(encrypted).not.toContain(`fresh-${failure}`);
expect(encrypted).not.toContain(`replacement-${failure}`);
await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`));
triggerDropped = true;
// A new service instance recovers from disk, without an in-memory cache.
const recovered = aiConnectionService(db);
await recovered.refreshQuotaCredential(row, original, new AbortController().signal);
expect(JSON.parse(await recovered.credential(row))).toMatchObject({ tokens: { access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` } });
expect(request).toHaveBeenCalledTimes(1);
await expect(access(file)).rejects.toThrow();
expect((await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).map(version => version.version).sort()).toEqual([1, 2]);
} finally {
request.mockRestore();
if (!triggerDropped) await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`));
}
});
it.each(["rollback", "lost_commit_reply"] as const)("automatically saves one exchanged token across %s", async (failure) => {
const owner = `quota-once-${failure}`;
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: owner } });
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
const [row] = await service.quotaAccounts(companyId, owner);
const transaction = db.transaction.bind(db);
const interrupted = vi.spyOn(db, "transaction").mockImplementationOnce(async (work, config) => {
await transaction(async (tx) => {
await work(tx);
if (failure === "rollback") throw new Error("Injected rollback after writes");
}, config);
throw new Error("Injected lost commit response");
});
const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: "fresh-once", refresh_token: "replacement-once" }));
try {
await service.refreshQuotaCredential(row, original, new AbortController().signal);
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-once", refresh_token: "replacement-once" } });
expect(request).toHaveBeenCalledTimes(1);
await expect(access(quotaCredentialRecoveryPath(companyId, row.grant.id))).rejects.toThrow();
} finally { interrupted.mockRestore(); request.mockRestore(); }
});
it.each(["reconnected", "revoked"] as const)("does not apply a pending refresh over a %s identity", async (change) => {
const owner = `quota-replay-${change}`;
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: owner } });
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
const [row] = await service.quotaAccounts(companyId, owner);
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
const journal = quotaCredentialRecovery(companyId, row.grant.id);
await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId, baseHash: quotaCredentialHash(original), value: "old-pending-replacement" });
const request = vi.spyOn(globalThis, "fetch");
try {
if (change === "reconnected") {
await secretService(db).rotate(ref.secretId, { value: "new-authorized-identity" });
expect(await service.refreshQuotaCredential(row, original, new AbortController().signal)).toBe("new-authorized-identity");
expect(await service.credential(row)).toBe("new-authorized-identity");
expect(await journal.read()).toBeNull();
} else {
await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, row.grant.id));
await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow("credentials_unavailable");
expect(await journal.read()).not.toBeNull();
}
expect(request).not.toHaveBeenCalled();
} finally { request.mockRestore(); await journal.clear(); }
});
it("saves exchanged tokens when the quota deadline expires during response-body reading", async () => {
const owner = "quota-body-deadline-owner";
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
+2 -2
View File
@@ -188,7 +188,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
expect(agentAfterHardStop).toMatchObject({ status: "paused", pauseReason: "budget" });
expect(agentAfterHardStop?.pausedAt).toBeInstanceOf(Date);
expect(cancelWorkForScope).toHaveBeenCalledTimes(2);
expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "agent", scopeId: agentId, createdBefore: expect.any(Date) });
expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "agent", scopeId: agentId, createdBefore: expect.any(Date), enforcement: { policyId: expect.any(String), version: expect.any(Number) } });
const block = await service.getInvocationBlock(companyId, agentId);
expect(block).toEqual({
@@ -273,7 +273,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
.from(projects);
expect(projectAfterHardStop?.pauseReason).toBe("budget");
expect(projectAfterHardStop?.pausedAt).toBeInstanceOf(Date);
expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "project", scopeId: projectId, createdBefore: expect.any(Date) });
expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "project", scopeId: projectId, createdBefore: expect.any(Date), enforcement: { policyId: expect.any(String), version: expect.any(Number) } });
const overviewWhileBlocked = await service.overview(companyId);
expect(overviewWhileBlocked.pausedProjectCount).toBe(1);
@@ -4,7 +4,7 @@ import os from "node:os";
import path from "node:path";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import { and, eq, sql } from "drizzle-orm";
import { accountingRuntimeBaselines, agents, agentRuntimeState, billingInvoiceLines, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, projects, runUsageReceipts } from "@paperclipai/db";
import { accountingRuntimeBaselines, agents, agentRuntimeState, billingInvoiceLines, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, nativeRunFinalizations, issues, projects, runUsageReceipts } from "@paperclipai/db";
import { accountingIntegrityService } from "../services/accounting-integrity.js";
import { billingReconciliationService } from "../services/billing-reconciliation.js";
import { budgetService } from "../services/budgets.js";
@@ -67,6 +67,35 @@ const support = await getEmbeddedPostgresTestSupport();
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId,f.run.id)))[0].state).toBe("released");
expect(await db.select().from(costEvents).where(eq(costEvents.companyId,f.company.id))).toEqual([]);
});
it.each(["retryable_failure", "observed"])("keeps native %s receipts open through same-run recovery", async (phase) => {
const f = await fixture();
const [issue] = await db.insert(issues).values({ companyId: f.company.id, title: "Recoverable native work" }).returning();
await db.update(heartbeatRuns).set({ runtimeMode: "native", nativeIssueId: issue.id }).where(eq(heartbeatRuns.id, f.run.id));
await reserveRunBudget(db, f.company.id, f.run.id, null);
await db.insert(nativeRunFinalizations).values({ companyId: f.company.id, runId: f.run.id, issueId: issue.id, phase });
const input = { companyId: f.company.id, runId: f.run.id, adapterType: "paperclip_runner" };
const spool = path.join(directory, randomUUID());
const recorder = await createRunUsageRecorder(db, input, spool);
await recorder.capture({ complete: true, usageBasis: "per_run", usage: { inputTokens: 100, outputTokens: 10 }, costUsd: 0.1 });
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.run.id));
expect(await accountRunCost(db, f.run.id)).toBe(false);
await reconcileRunCosts(db);
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0].costAccountedAt).toBeNull();
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, f.run.id)))[0].state).toBe("held");
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, f.run.id));
const resumed = await createRunUsageRecorder(db, input, spool);
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0].usageJson?.accountingReceiptReady).toBe(false);
const empty = await resumed.complete({ exitCode: 0, signal: null, timedOut: false });
expect(empty).toMatchObject({ complete: false, usage: { inputTokens: 100, outputTokens: 10 } });
// The provider's native delta is cumulative for this run across recovery.
await resumed.capture({ complete: true, usageBasis: "per_run", usage: { inputTokens: 150, outputTokens: 15 }, costUsd: 0.15 });
await db.update(nativeRunFinalizations).set({ phase: "terminal_failure" }).where(eq(nativeRunFinalizations.runId, f.run.id));
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.run.id));
expect(await accountRunCost(db, f.run.id)).toBe(true);
expect(await accountRunCost(db, f.run.id)).toBe(false);
expect((await db.select().from(costEvents).where(eq(costEvents.heartbeatRunId, f.run.id)))[0]).toMatchObject({ inputTokens: 150, outputTokens: 15, costCents: 15 });
});
it("reports over-limit advisory budgets without blocking, and tolerates legacy active zero limits", async () => {
const f = await fixture(), budgets = budgetService(db);
const policy = await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 1, hardStopEnabled: false }, "board");
@@ -75,7 +75,7 @@ const support = await getEmbeddedPostgresTestSupport();
const original = await reserveRunBudget(db, f.company.id, run.id, null);
await expect(reserveRunBudget(db, f.company.id, run.id, null)).rejects.toThrow("already reserved");
await expect(reserveRunBudget(db, f.company.id, run.id, null, {}, "stale-owner")).rejects.toThrow("already reserved");
expect(await reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).toEqual(original);
expect(await reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).toEqual({ ...original, reused: true });
expect((await accountingIntegrityService(db).health(f.company.id)).heldReservationCents).toBe("6.0000000");
await db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }).where(eq(nativeRunFinalizations.runId, run.id));
await expect(reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).rejects.toThrow("already reserved");
@@ -2,9 +2,11 @@ import { upsertBudgetPolicySchema } from "@paperclipai/shared";
import { randomUUID } from "node:crypto";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import { and, eq, sql } from "drizzle-orm";
import { activityLog, agentRuntimeState, agents, approvals, budgetIncidents, budgetPolicies, companies, costEvents, createDb, financeEvents, goals, heartbeatRuns, issues, projects } from "@paperclipai/db";
import { costService } from "../services/costs.js";
import { budgetService } from "../services/budgets.js";
import { activityLog, agentRuntimeState, agentWakeupRequests, agents, approvals, budgetIncidents, budgetPolicies, companies, costEvents, createDb, financeEvents, goals, heartbeatRuns, issues, projects } from "@paperclipai/db";
import { reserveRunBudget } from "../services/budget-reservations.js";
import { createRunDispatch } from "../modules/run-dispatch/index.js";
import { costService, createCostEventInTransaction } from "../services/costs.js";
import { budgetService, withCurrentBudgetEnforcement, type BudgetEnforcementScope } from "../services/budgets.js";
import { accountRunCost, reconcileRunCosts } from "../services/run-cost-accounting.js";
import { financeService } from "../services/finance.js";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
@@ -57,6 +59,31 @@ databaseDescribe("cost accounting reliability (PostgreSQL)", () => {
});
}, 5000);
it("takes the company lock before retry issue/run locks during an attributed cost write", async () => {
const f = await fixture();
await db.update(issues).set({ status: "in_progress", assigneeAgentId: f.agent.id }).where(eq(issues.id, f.issue.id));
await db.update(heartbeatRuns).set({ status: "scheduled_retry", scheduledRetryAt: new Date(0), scheduledRetryReason: "workspace_busy",
contextSnapshot: { issueId: f.issue.id, projectId: f.project.id },
}).where(eq(heartbeatRuns.id, f.run.id));
let promotion: ReturnType<ReturnType<typeof createRunDispatch>["promoteScheduledRetry"]> | undefined;
await db.transaction(async (tx) => {
await tx.select().from(companies).where(eq(companies.id, f.company.id)).for("no key update");
promotion = createRunDispatch(db).promoteScheduledRetry({ companyId: f.company.id, runId: f.run.id });
// Observe the actual blocked company-lock query before checking the issue;
// no timing guess or timeout is used as proof of lock ordering.
await vi.waitFor(async () => {
const result = await db.execute(sql`select pid from pg_stat_activity where datname = current_database()
and pid <> pg_backend_pid() and wait_event_type = 'Lock' and query like '%companies%for no key update%'`);
expect(result.length).toBeGreaterThan(0);
}, { timeout: 3000 });
await tx.execute(sql`select id from issues where id = ${f.issue.id} for key share nowait`);
await createCostEventInTransaction(tx as unknown as ReturnType<typeof createDb>, f.company.id,
{ ...f.receipt, issueId: f.issue.id, projectId: f.project.id }, []);
});
expect((await promotion)?.outcome).toBe("promoted");
expect((await costService(db).summary(f.company.id)).spendCents).toBe(100);
}, 10000);
it("deduplicates concurrent retries and rejects changed receipt content", async () => {
const f = await fixture();
const receipt = { ...f.receipt, idempotencyKey: "provider-receipt-1" };
@@ -419,4 +446,57 @@ databaseDescribe("cost accounting reliability (PostgreSQL)", () => {
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, newRun.id)))[0].status).toBe("scheduled_retry");
});
it.each(["agent", "company", "project"] as const)("rechecks delayed %s enforcement after a grant admits an old queued run", async (scopeType) => {
const f = await fixture();
const { heartbeatService } = await import("../services/heartbeat.js");
const heartbeat = heartbeatService(db);
const budgets = budgetService(db);
const scopeId = scopeType === "agent" ? f.agent.id : scopeType === "company" ? f.company.id : f.project.id;
const policyInput = { scopeType, scopeId, amount: 1 };
await db.update(heartbeatRuns).set({ status: "queued", contextSnapshot: { projectId: f.project.id } }).where(eq(heartbeatRuns.id, f.run.id));
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: f.company.id, agentId: f.agent.id, source: "automation", payload: { projectId: f.project.id } }).returning();
await budgets.upsertPolicy(f.company.id, policyInput, "board");
await costService(db).createEvent(f.company.id, { ...f.receipt, projectId: f.project.id, costCents: 2 });
let release!: () => void;
const paused = new Promise<void>(resolve => { release = resolve; });
let snapshotReady!: (scope: BudgetEnforcementScope) => void;
const snapshot = new Promise<BudgetEnforcementScope>(resolve => { snapshotReady = resolve; });
const delivery = budgetService(db, { cancelWorkForScope: async (scope) => {
snapshotReady(scope); await paused; await heartbeat.cancelBudgetScopeWork(scope);
} }).deliverPendingEnforcement(f.company.id);
try {
expect((await snapshot).enforcement).toMatchObject({ policyId: expect.any(String), version: expect.any(Number) });
await budgets.upsertPolicy(f.company.id, { ...policyInput, amount: 100 }, "board");
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, f.run.id));
await reserveRunBudget(db, f.company.id, f.run.id, f.project.id);
} finally { release(); }
await delivery;
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0]).toMatchObject({ status: "running", resultJson: null });
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
});
it("keeps a durably claimed budget stop fenced after a later grant", async () => {
const f = await fixture(), budgets = budgetService(db);
await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 1 }, "board");
await costService(db).createEvent(f.company.id, { ...f.receipt, costCents: 2 });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, f.company.id));
const scope: BudgetEnforcementScope = { companyId: f.company.id, scopeType: "company", scopeId: f.company.id,
enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
expect(await withCurrentBudgetEnforcement(db, scope, async (tx) => {
await tx.update(heartbeatRuns).set({ status: "running", resultJson: { cancellation: { reason: "Budget stop" } } }).where(eq(heartbeatRuns.id, f.run.id));
return true;
})).toBe(true);
await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 100 }, "board");
await expect(reserveRunBudget(db, f.company.id, f.run.id, null)).rejects.toThrow("cancellation already requested");
const effect = vi.fn(async () => true);
for (const invalid of [
scope,
{ ...scope, enforcement: { ...scope.enforcement!, policyId: randomUUID() } },
{ ...scope, scopeId: randomUUID() },
{ ...scope, scopeType: "agent" as const },
{ ...scope, enforcement: { ...scope.enforcement!, version: policy.enforcementVersion + 1 } },
]) expect(await withCurrentBudgetEnforcement(db, invalid, effect)).toBeNull();
expect(effect).not.toHaveBeenCalled();
});
});
@@ -7553,6 +7553,19 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
},
);
it.each([{}, { command: "/nonexistent-paperclip-bootstrap-command" }])("releases holds for real process bootstrap failures: %j", async (config) => {
const { execute } = await import("../adapters/process/execute.js");
mockAdapterExecute.mockImplementationOnce(execute);
const { runId, agentId } = await seedRunFixture({ runtimeMode: "legacy", agentStatus: "idle", runStatus: "queued", includeIssue: false });
await db.update(agents).set({ adapterType: "process", adapterConfig: config }).where(eq(agents.id, agentId));
const heartbeat = heartbeatService(db);
await heartbeat.resumeQueuedRuns();
await heartbeat.drainActiveRunExecutions();
expect(await heartbeat.getRun(runId)).toMatchObject({ status: "failed", costAccountingPending: false, costAccountedAt: expect.any(Date) });
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("released");
expect(await db.select().from(costEvents).where(eq(costEvents.heartbeatRunId, runId))).toEqual([]);
});
it("releases the dispatch reservation after a proven pre-provider adapter failure", async () => {
mockAdapterExecute.mockResolvedValueOnce({
exitCode: 1, signal: null, timedOut: false,
@@ -6,6 +6,7 @@ import {
agentWakeupRequests,
companies,
costEvents,
budgetReservations,
createDb,
documentRevisions,
documents,
@@ -847,6 +848,9 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => {
afterContinuationDispatchCheck = async ({ runId: guardedRunId, issueId: guardedIssueId }) => {
expect(guardedRunId).toBe(runId);
expect(guardedIssueId).toBe(issueId);
// Budget admission must finish before this last ownership check. Putting
// an awaited reservation inside dispatch lets work escape the row lock.
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("held");
ordering.push("validated");
parkPromise = Promise.resolve(
db
@@ -890,6 +894,8 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => {
expect(issue?.status).toBe("backlog");
expect(ordering).toEqual(["validated", "parked"]);
expect(countExecuteCallsForRun(runId)).toBe(0);
await heartbeat.drainActiveRunExecutions();
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("released");
});
it("rate-limits skipped generic timer wakes by advancing the timer baseline", async () => {
+31 -15
View File
@@ -43,6 +43,7 @@ import {
releaseRuntimeServicesForRun,
UnresolvedWorkspaceBaseRefError,
resetRuntimeServicesForTests,
resetRuntimeServicePortReservationsForTests,
MANAGED_RUNTIME_PUBLIC_URL_ENV,
resolveManagedPaperclipRuntimePublicOrigin,
resolveRuntimeProvisionCommand,
@@ -454,6 +455,9 @@ afterEach(async () => {
delete process.env.PAPERCLIP_WORKTREES_DIR;
delete process.env.DATABASE_URL;
await resetRuntimeServicesForTests();
// Registry reset does not clear the process-local allocation claims. A
// failed-start fixture must not reserve another test's ephemeral port.
resetRuntimeServicePortReservationsForTests();
});
describe("sanitizeRuntimeServiceBaseEnv", () => {
@@ -4099,6 +4103,7 @@ describe("realizeExecutionWorkspace", () => {
"utf8",
);
process.env.PAPERCLIP_WORKTREES_DIR = worktreesDir;
const canonicalInstanceRoot = await fs.realpath(instanceRoot);
await cleanupExecutionWorkspaceArtifacts({
workspace: {
@@ -4133,7 +4138,7 @@ describe("realizeExecutionWorkspace", () => {
expect(operations[0]?.command).toBe("printf 'cleanup ok\\n'");
expect(operations[1]?.metadata).toMatchObject({
cleanupAction: "remove_worktree_instance",
instanceRoot,
instanceRoot: canonicalInstanceRoot,
});
expect(operations[2]?.metadata).toMatchObject({
cleanupAction: "worktree_remove",
@@ -6759,7 +6764,8 @@ describeEmbeddedPostgres("workspace runtime service control persistence", () =>
status: "starting",
healthStatus: "unknown",
});
expect(startingRow.providerRef).toMatch(/^\d+$/);
// Starting is persisted before the asynchronous spawn metadata write.
// The child marker alone does not prove that write has committed.
expect(startingRow.port).toEqual(expect.any(Number));
const services = await startPromise;
@@ -6772,6 +6778,7 @@ describeEmbeddedPostgres("workspace runtime service control persistence", () =>
const runningRow = await waitForPersistedStatus("running");
expect(runningRow.id).toBe(startingRow.id);
expect(runningRow.providerRef).toMatch(/^\d+$/);
await expect(fetch(services[0]!.url!)).resolves.toMatchObject({ ok: true });
const runtimeProvisionOperations = await db
.select()
@@ -7643,16 +7650,25 @@ describeEmbeddedPostgres("workspace runtime startup reconciliation", () => {
process.env.PAPERCLIP_HOME = paperclipHome;
process.env.PAPERCLIP_INSTANCE_ID = `runtime-desired-reconcile-${randomUUID()}`;
const reservePort = async () => {
const probe = net.createServer();
await new Promise<void>((resolve) => probe.listen(0, "127.0.0.1", resolve));
const address = probe.address();
const port = typeof address === "object" && address ? address.port : null;
await new Promise<void>((resolve, reject) => {
probe.close((error) => error ? reject(error) : resolve());
});
if (!port) throw new Error("Failed to reserve runtime reconciliation test port");
return port;
const reservePorts = async () => {
const probes = [net.createServer(), net.createServer()];
try {
// Keep both sockets open until both allocations finish. Closing the
// first probe early lets the kernel return that same port again.
for (const probe of probes) await new Promise<void>((resolve, reject) => {
probe.once("error", reject);
probe.listen(0, "127.0.0.1", resolve);
});
return probes.map(probe => {
const address = probe.address();
if (!address || typeof address !== "object") throw new Error("Failed to reserve runtime reconciliation test port");
return address.port;
});
} finally {
await Promise.all(probes.filter(probe => probe.listening).map(probe => new Promise<void>((resolve, reject) => {
probe.close(error => error ? reject(error) : resolve());
})));
}
};
const isLoopbackPortFree = async (port: number) => {
const probe = net.createServer();
@@ -7673,8 +7689,8 @@ describeEmbeddedPostgres("workspace runtime startup reconciliation", () => {
}
throw new Error(`Port ${port} did not become free in time`);
};
const stoppedPort = await reservePort();
const livePort = await reservePort();
const [stoppedPort, livePort] = await reservePorts();
expect(stoppedPort).not.toBe(livePort);
const companyId = randomUUID();
const projectId = randomUUID();
const projectWorkspaceId = randomUUID();
@@ -9728,7 +9744,7 @@ describe("realizeExecutionWorkspace with an exact existing branch", () => {
const workspace = await realizeExistingBranch(repoRoot, "feature/legacy-checkout");
expect(workspace.cwd).toBe(path.resolve(legacyPath));
expect(workspace.cwd).toBe(await fs.realpath(legacyPath));
expect(workspace.branchName).toBe("feature/legacy-checkout");
expect(workspace.created).toBe(false);
expect(await readGit(workspace.cwd, ["rev-parse", "HEAD"])).toBe(branchTip);
+18
View File
@@ -15,6 +15,24 @@ import {
} from "../utils.js";
export async function execute(ctx: AdapterExecutionContext): Promise<AdapterExecutionResult> {
let spawned = false;
try {
return await executeProcess({ ...ctx, onSpawn: async (meta) => {
// Set this before metadata persistence: a failed callback cannot prove
// that the already-created process did no billable work.
spawned = true;
await ctx.onSpawn?.(meta);
} });
} catch (error) {
if (spawned) throw error;
return { exitCode: 1, signal: null, timedOut: false,
errorMessage: error instanceof Error ? error.message : String(error),
executionRecovery: { kind: "bootstrap", providerWorkStarted: false },
};
}
}
async function executeProcess(ctx: AdapterExecutionContext): Promise<AdapterExecutionResult> {
const { runId, agent, config, onLog, onMeta, authToken } = ctx;
const command = asString(config.command, "");
if (!command) throw new Error("Process adapter missing command");
@@ -15,9 +15,11 @@ import {
issues,
} from "@paperclipai/db";
import { ISSUE_DISPOSITION_REPAIR_RETRY_REASON } from "@paperclipai/shared";
import { withAccountingTransaction } from "../../../services/accounting-transaction.js";
import type { ActivityPublication } from "../../../services/activity-log.js";
import { parseObject } from "../../../adapters/utils.js";
import { evaluateAgentInvokabilityFromDb } from "../../../services/agent-invokability.js";
import { budgetService } from "../../../services/budgets.js";
import { budgetService, budgetServiceInTransaction } from "../../../services/budgets.js";
import { isCompletedOnboardingHandoffWake } from "../../../services/chat-completion-delivery.js";
import { isHeartbeatWakeOnDemandEnabled } from "../../../services/heartbeat-policy.js";
import { collectDispositionRepairSourceState } from "../../../services/recovery/disposition-repair.js";
@@ -190,7 +192,7 @@ export function createPostgresRunDispatchAdapter(
async function withIssueThenRunLocks<T>(
input: { runId: string; companyId: string },
onMissing: () => T,
operation: (tx: Db, run: HeartbeatRun) => Promise<T>,
operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise<T>,
): Promise<T> {
// Queue editing and claiming already use issue -> wake -> run. Read the
// immutable issue reference without a lock first, then acquire issue ->
@@ -206,8 +208,9 @@ export function createPostgresRunDispatchAdapter(
const hintedIssueId = readNonEmptyString(parseObject(hint.contextSnapshot).issueId);
const result: { kind: "missing" | "retry" } | { kind: "value"; value: T } =
await db.transaction(async (tx) => {
const typedTx = tx as unknown as Db;
await withAccountingTransaction(db, input.companyId, async (typedTx, publications) => {
// Budget gates mutate policy state. Acquire company before issue/run,
// matching receipt writers (including their attribution FK locks).
if (hintedIssueId) {
await typedTx
.select({ id: issues.id })
@@ -231,7 +234,7 @@ export function createPostgresRunDispatchAdapter(
const lockedIssueId = readNonEmptyString(parseObject(run.contextSnapshot).issueId);
if (lockedIssueId !== hintedIssueId) return { kind: "retry" as const };
return { kind: "value" as const, value: await operation(typedTx, run) };
return { kind: "value" as const, value: await operation(typedTx, run, publications) };
});
if (result.kind === "missing") return onMissing();
@@ -247,12 +250,13 @@ export function createPostgresRunDispatchAdapter(
input: LoadGateFactsInput,
now: Date,
tx?: unknown,
publications?: ActivityPublication[],
): Promise<LoadGateFactsResult> {
// Semantic adapter operations pass their transaction here so the fact
// read and the state transition share one unit of work. This helper is
// deliberately not exposed through the module's public API.
const dbOrTx = (tx as Db | undefined) ?? db;
const budgetsForRead = tx ? budgetService(dbOrTx) : budgets;
const budgetsForRead = tx ? budgetServiceInTransaction(dbOrTx, publications) : budgets;
const treeControlForRead = tx ? issueTreeControlService(dbOrTx) : treeControlSvc;
const issuesSvcForRead = tx ? issueService(dbOrTx) : issuesSvc;
@@ -747,7 +751,7 @@ export function createPostgresRunDispatchAdapter(
): Promise<PromoteScheduledRetryOutcome> {
const now = input.now;
const promoteLockedRun = async (tx: Db, run: HeartbeatRun) => {
const promoteLockedRun = async (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => {
if (
run.status !== "scheduled_retry" ||
!run.scheduledRetryAt ||
@@ -768,6 +772,7 @@ export function createPostgresRunDispatchAdapter(
},
now,
tx,
publications,
);
if (!factsResult.agentFound) {
@@ -1010,7 +1015,7 @@ export function createPostgresRunDispatchAdapter(
input: DispatchResolvedInteractionInput<T>,
): Promise<DispatchResolvedInteractionOutcome<T>> {
const dispatchLockedRun = async (tx: Db, run: HeartbeatRun) => {
if (run.status !== input.expectedStatus) {
if (run.status !== input.expectedStatus || run.resultJson?.cancellation || run.resultJson?.startupCancellation) {
return { dispatched: false as const, cancellation: { outcome: "lost_race" as const } };
}
const { issueId, decision: initialDecision } = await decideCurrentRunStaleness(
+90 -58
View File
@@ -34,6 +34,8 @@ import {
type AiConnectionLoginIntent,
} from "@paperclipai/shared";
import { forbidden, notFound, unprocessable } from "../errors.js";
import { quotaCredentialRecovery, quotaCredentialHash } from "./quota-credential-recovery.js";
import { logger } from "../middleware/logger.js";
import { logActivity } from "./activity-log.js";
import { secretService } from "./secrets.js";
import { probeAiConnectionUsage } from "./ai-connection-usage.js";
@@ -860,66 +862,96 @@ export function aiConnectionService(db: Db) {
}
/** Refresh only the selected vaulted identity, serializing rotation and reconnect. */
async function refreshQuotaCredential(row: Parameters<typeof credential>[0], failedValue: string, signal: AbortSignal) {
return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => db.transaction(async (tx) => {
// Do not wait on a transaction that might itself be waiting for the file
// lock. All database locks must be available before consuming OAuth tokens.
const companyId = row.connection.companyId;
const [grant] = await tx.select().from(connectionGrants).where(and(
eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId),
eq(connectionGrants.connectionId, row.connection.id),
)).for("update", { noWait: true });
if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
if (!ref) throw new Error("credentials_unavailable");
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId),
)).for("update", { noWait: true });
const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant });
signal.throwIfAborted();
// Another poll, an agent, or reconnect already supplied a newer credential.
if (current !== failedValue) return current;
// An active provider may rotate this single-use token itself. Defer polling
// until its credential write-back instead of invalidating its live copy.
const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns)
.innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and(
eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]),
or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`,
sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")),
)).limit(1);
if (active) throw new Error("provider_unavailable");
const auth = JSON.parse(current);
const refreshToken = auth.tokens?.refresh_token;
if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required");
// Once exchange starts, its lifetime is independent of the dashboard's
// shorter deadline. A consumed single-use token must still be read/saved.
signal.throwIfAborted();
const refreshSignal = AbortSignal.timeout(60_000);
// Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs).
const response = await fetch("https://auth.openai.com/oauth/token", {
method: "POST", redirect: "error", signal: refreshSignal,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }),
return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => {
const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id);
let savingReplacement = false;
const save = () => db.transaction(async (tx) => {
// Do not wait on a transaction that might itself be waiting for the file
// lock. All database locks must be available before consuming OAuth tokens.
const companyId = row.connection.companyId;
const [grant] = await tx.select().from(connectionGrants).where(and(
eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId),
eq(connectionGrants.connectionId, row.connection.id),
)).for("update", { noWait: true });
if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
if (!ref) throw new Error("credentials_unavailable");
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId),
)).for("update", { noWait: true });
const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant });
signal.throwIfAborted();
// Another poll, an agent, or reconnect already supplied a newer credential.
if (current !== failedValue) return current;
const pending = await recovery.read();
if (pending && (pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current))) {
throw new Error("credentials_unavailable");
}
// An active provider may rotate this single-use token itself. Defer polling
// until its credential write-back instead of invalidating its live copy.
const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns)
.innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and(
eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]),
or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`,
sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")),
)).limit(1);
if (active && !pending) throw new Error("provider_unavailable");
let refreshed = pending?.value;
if (!refreshed) {
await recovery.prepare();
const auth = JSON.parse(current);
const refreshToken = auth.tokens?.refresh_token;
if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required");
// Once exchange starts, its lifetime is independent of the dashboard's
// shorter deadline. A consumed single-use token must still be read/saved.
signal.throwIfAborted();
const refreshSignal = AbortSignal.timeout(60_000);
// Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs).
const response = await fetch("https://auth.openai.com/oauth/token", {
method: "POST", redirect: "error", signal: refreshSignal,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }),
});
if (!response.ok) {
await response.body?.cancel();
throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable");
}
const tokens = await response.json() as Record<string, unknown>;
if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable");
refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens,
access_token: tokens.access_token,
...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}),
...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}),
}, last_refresh: new Date().toISOString() });
// This fsynced encrypted journal survives transaction rollback and host
// restart. Never exchange the consumed token again after a failed save.
await recovery.write({ companyId, grantId: grant.id, connectionId: row.connection.id,
secretId: ref.secretId, baseHash: quotaCredentialHash(current), value: refreshed });
}
savingReplacement = true;
// Persist the rotated token even if the quota deadline expires after the
// successful exchange; losing it would strand all subsequent executions.
await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId });
await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id));
await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota",
action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
details: { provider: "openai", grantId: grant.id } });
return refreshed;
});
if (!response.ok) {
await response.body?.cancel();
throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable");
for (let attempt = 0; ; attempt++) {
try {
const value = await save();
// Clear only after commit. A committed-but-lost response is detected
// by the current credential comparison on the next attempt.
await recovery.clear().catch((error) => logger.warn({ err: error, grantId: row.grant.id }, "Saved quota credential recovery cleanup pending"));
return value;
} catch (error) {
if (!savingReplacement || attempt >= 2) throw error;
// Saving retries have no provider side effect. Persistent failures
// retain the journal for the next quota poll, including after restart.
}
}
const tokens = await response.json() as Record<string, unknown>;
if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable");
const refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens,
access_token: tokens.access_token,
...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}),
...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}),
}, last_refresh: new Date().toISOString() });
// Persist the rotated token even if the quota deadline expires after the
// successful exchange; losing it would strand all subsequent executions.
await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId });
await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id));
await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota",
action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
details: { provider: "openai", grantId: grant.id } });
return refreshed;
}));
});
}
// A quota read uses the same credential audience as execution. Operator status
// alone never grants access to another member's personal subscription.
+3 -2
View File
@@ -12,6 +12,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
return withAccountingTransaction(db, companyId, async (tx, publications) => {
const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.companyId, companyId))).for("update");
if (!run) throw notFound("Run not found");
if (run.resultJson?.cancellation || run.resultJson?.startupCancellation) throw conflict("Run cancellation already requested");
if (run.costAccountedAt || !["queued", "running"].includes(run.status)) throw conflict("Run can no longer start provider work");
const [existing] = await tx.select().from(budgetReservations).where(and(eq(budgetReservations.companyId, companyId), eq(budgetReservations.runId, runId)));
if (existing) {
@@ -22,7 +23,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
.from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, companyId),
eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner),
sql`${nativeRunFinalizations.leaseExpiresAt} > now()`)).for("update") : [];
if (owner) return existing;
if (owner) return { ...existing, reused: true };
throw conflict("Provider dispatch has already reserved this run");
}
const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId });
@@ -50,6 +51,6 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
await tx.update(heartbeatRuns).set({ costAccountingPending: true,
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptReady: false, ledgerScope: { ...(run.usageJson?.ledgerScope as object ?? {}), ...ledgerScope, projectId } })}::jsonb`,
}).where(eq(heartbeatRuns.id, runId));
return reservation;
return { ...reservation, reused: false };
});
}
+19 -1
View File
@@ -45,6 +45,8 @@ export type BudgetEnforcementScope = {
scopeId: string;
/** Snapshot taken under the company lock, before a later budget grant can admit new work. */
createdBefore?: Date;
/** Revalidated at the durable cancellation write, under the admission lock. */
enforcement?: { policyId: string; version: number };
};
export type BudgetServiceHooks = {
@@ -939,6 +941,22 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
};
}
/** Fence the durable stop intent against policy edits and admission. External
* process shutdown must happen after this transaction, without holding locks
* that the executor needs to persist its final usage receipt. */
export async function withCurrentBudgetEnforcement<T>(db: Db, scope: BudgetEnforcementScope, work: (tx: Db) => Promise<T>): Promise<T | null> {
return withAccountingTransaction(db, scope.companyId, async (tx) => {
if (scope.enforcement) {
const [policy] = await tx.select().from(budgetPolicies).where(and(
eq(budgetPolicies.id, scope.enforcement.policyId), eq(budgetPolicies.companyId, scope.companyId),
));
if (!policy || policy.scopeType !== scope.scopeType || policy.scopeId !== scope.scopeId
|| policy.enforcementVersion !== scope.enforcement.version || !(await policyBlocks(tx, policy))) return null;
}
return work(tx);
});
}
/** Cancellation is an at-least-once external effect. A failed delivery never
* rolls back committed spend, and its version remains pending for recovery. */
export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks, companyId?: string) {
@@ -952,7 +970,7 @@ export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks
const cancellation = await withAccountingTransaction(db, policy.companyId, async (tx) => {
const [current] = await tx.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
if (!current || current.enforcementVersion !== policy.enforcementVersion || !(await policyBlocks(tx, current))) return null;
return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId, createdBefore: new Date() };
return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId, createdBefore: new Date(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
});
if (cancellation) await hooks.cancelWorkForScope(cancellation);
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion })
+124 -82
View File
@@ -340,7 +340,7 @@ import {
import { reportRunFailure } from "./run-failure-report.js";
import { collectRunFailureSecretValues, type RunFailureReportOptions } from "./run-failure-diagnostics.js";
import { companySkillService } from "./company-skills.js";
import { budgetService, type BudgetEnforcementScope } from "./budgets.js";
import { budgetService, withCurrentBudgetEnforcement, type BudgetEnforcementScope } from "./budgets.js";
import { secretService, type MissingRuntimeBinding } from "./secrets.js";
import {
resolveDefaultAgentWorkspaceDir,
@@ -22688,7 +22688,7 @@ export function heartbeatService(
},
);
}
const dispatchResolvedInteractionContinuationWithAtomicGate = async <T>(
const dispatchResolvedInteractionContinuationAfterAdmission = async <T>(
dispatch: (markDispatchStarted: () => void) => Promise<T>,
): Promise<
{ dispatched: true; resultPromise: Promise<T> } | { dispatched: false }
@@ -22754,6 +22754,30 @@ export function heartbeatService(
}
return { dispatched: false };
};
const dispatchResolvedInteractionContinuationWithAtomicGate = async <T>(
dispatch: (markDispatchStarted: () => void) => Promise<T>,
) => {
// Admission can wait for accounting locks. Finish that wait before the
// final ownership gate, whose callback must enter the adapter directly.
const reservation = await reserveRunBudget(db, run.companyId, run.id,
readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner);
let entered = false;
try {
return await dispatchResolvedInteractionContinuationAfterAdmission((markDispatchStarted) => {
entered = true;
return dispatch(markDispatchStarted);
});
} finally {
if (!entered && !reservation.reused) {
await db.update(heartbeatRuns).set({ costAccountingPending: true,
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || '{"accountingProviderWorkStarted":false}'::jsonb`,
}).where(and(eq(heartbeatRuns.id, run.id), isNull(heartbeatRuns.costAccountedAt)));
await accountRunCost(db, run.id, budgetHooks).catch((err) => {
logger.warn({ err, runId: run.id }, "Undispatched reservation release remains pending");
});
}
}
};
if (!executionTarget || executionTarget.kind === "local") {
try {
runScratch = await prepareHeartbeatRunScratch({
@@ -24656,8 +24680,7 @@ export function heartbeatService(
);
const guardedDispatch =
await dispatchResolvedInteractionContinuationWithAtomicGate(
async (markDispatchStarted) => {
await reserveRunBudget(db, run.companyId, run.id, readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner);
(markDispatchStarted) => {
return executePaperclipNativeSession({
db,
execution: nativeExecution,
@@ -24884,8 +24907,7 @@ export function heartbeatService(
}
const guardedDispatch =
await dispatchResolvedInteractionContinuationWithAtomicGate(
async (markDispatchStarted) => {
await reserveRunBudget(db, run.companyId, run.id, readNonEmptyString(runLedgerScope.projectId), runLedgerScope);
(markDispatchStarted) => {
legacyAdapterEntered = true;
return adapter.execute({
getFreshSessionHandoff,
@@ -29331,6 +29353,7 @@ export function heartbeatService(
async function listProjectScopedWakeupIds(
companyId: string,
projectId: string,
database: Db = db,
) {
const wakeIssueId = sql<
string | null
@@ -29339,7 +29362,7 @@ export function heartbeatService(
string | null
>`coalesce(${agentWakeupRequests.payload} ->> 'projectId', ${issues.projectId}::text)`;
const rows = await db
const rows = await database
.selectDistinctOn([agentWakeupRequests.id], {
id: agentWakeupRequests.id,
})
@@ -29369,68 +29392,72 @@ export function heartbeatService(
async function cancelPendingWakeupsForBudgetScope(
scope: BudgetEnforcementScope,
) {
const now = new Date();
let wakeupIds: string[] = [];
return withCurrentBudgetEnforcement(db, scope, async (tx) => {
const now = new Date();
let wakeupIds: string[] = [];
if (scope.scopeType === "company") {
wakeupIds = await db
.select({ id: agentWakeupRequests.id })
.from(agentWakeupRequests)
.where(
and(
eq(agentWakeupRequests.companyId, scope.companyId),
inArray(agentWakeupRequests.status, [
"queued",
"deferred_issue_execution",
]),
sql`${agentWakeupRequests.runId} is null`,
),
)
.then((rows) => rows.map((row) => row.id));
} else if (scope.scopeType === "agent") {
wakeupIds = await db
.select({ id: agentWakeupRequests.id })
.from(agentWakeupRequests)
.where(
and(
eq(agentWakeupRequests.companyId, scope.companyId),
eq(agentWakeupRequests.agentId, scope.scopeId),
inArray(agentWakeupRequests.status, [
"queued",
"deferred_issue_execution",
]),
sql`${agentWakeupRequests.runId} is null`,
),
)
.then((rows) => rows.map((row) => row.id));
} else {
wakeupIds = await listProjectScopedWakeupIds(
scope.companyId,
scope.scopeId,
);
}
if (scope.scopeType === "company") {
wakeupIds = await tx
.select({ id: agentWakeupRequests.id })
.from(agentWakeupRequests)
.where(
and(
eq(agentWakeupRequests.companyId, scope.companyId),
inArray(agentWakeupRequests.status, [
"queued",
"deferred_issue_execution",
]),
sql`${agentWakeupRequests.runId} is null`,
),
)
.then((rows) => rows.map((row) => row.id));
} else if (scope.scopeType === "agent") {
wakeupIds = await tx
.select({ id: agentWakeupRequests.id })
.from(agentWakeupRequests)
.where(
and(
eq(agentWakeupRequests.companyId, scope.companyId),
eq(agentWakeupRequests.agentId, scope.scopeId),
inArray(agentWakeupRequests.status, [
"queued",
"deferred_issue_execution",
]),
sql`${agentWakeupRequests.runId} is null`,
),
)
.then((rows) => rows.map((row) => row.id));
} else {
wakeupIds = await listProjectScopedWakeupIds(
scope.companyId,
scope.scopeId,
tx,
);
}
if (wakeupIds.length === 0) return 0;
if (wakeupIds.length === 0) return 0;
await db
.update(agentWakeupRequests)
.set({
status: "cancelled",
finishedAt: now,
error: "Cancelled due to budget pause",
updatedAt: now,
})
.where(and(
inArray(agentWakeupRequests.id, wakeupIds),
inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]),
isNull(agentWakeupRequests.runId),
scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined,
));
await tx
.update(agentWakeupRequests)
.set({
status: "cancelled",
finishedAt: now,
error: "Cancelled due to budget pause",
updatedAt: now,
})
.where(and(
inArray(agentWakeupRequests.id, wakeupIds),
inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]),
isNull(agentWakeupRequests.runId),
scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined,
));
return wakeupIds.length;
return wakeupIds.length;
});
}
type CancelRunOptions = {
budgetEnforcement?: BudgetEnforcementScope;
errorCode?: string;
resultJson?: Record<string, unknown>;
eventMessage?: string;
@@ -29505,25 +29532,40 @@ export function heartbeatService(
// Established legacy processes must still be stopped if the database is
// unavailable. Only native or not-yet-dispatched preparation needs this
// additional durable fence before its existing cancellation path.
if (run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)) {
const [fenced] = await db.update(heartbeatRuns).set({
// Record handoff intent before native cancellation can finalize and release
// the run. Only its audited stop acknowledgement suppresses recovery.
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) ||
${JSON.stringify(options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true
? { reassignmentStopRequested: true } : {})}::jsonb ||
${JSON.stringify({ cancellation })}::jsonb ||
jsonb_build_object('startupCancellation', jsonb_build_object(
'requestedAt', ${new Date().toISOString()}::text,
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
and ${heartbeatRuns.runtimeModeResolvedAt} is null
and ${heartbeatRuns.executionStage} = 'preparing'
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
))`,
}).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status,
pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES],
))).returning();
if (!fenced) return getRun(runId);
if (options.budgetEnforcement || run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)) {
const fence = async (tx: Db) => {
const [fenced] = await tx.update(heartbeatRuns).set({
// Record handoff intent before native cancellation can finalize and release
// the run. Only its audited stop acknowledgement suppresses recovery.
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) ||
${JSON.stringify(options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true
? { reassignmentStopRequested: true } : {})}::jsonb ||
${JSON.stringify({ cancellation })}::jsonb ||
jsonb_build_object('startupCancellation', jsonb_build_object(
'requestedAt', ${new Date().toISOString()}::text,
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
and ${heartbeatRuns.runtimeModeResolvedAt} is null
and ${heartbeatRuns.executionStage} = 'preparing'
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
))`,
}).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status,
pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES],
))).returning();
return fenced ?? null;
};
let fenced: typeof run | null;
try {
fenced = options.budgetEnforcement
? await withCurrentBudgetEnforcement(db, options.budgetEnforcement, fence)
: await fence(db);
} catch (error) {
stopOwnership?.release();
throw error;
}
if (!fenced) {
stopOwnership?.release();
return getRun(runId);
}
run = fenced;
}
const resultJson = { ...(agent
@@ -29903,7 +29945,7 @@ export function heartbeatService(
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]),
)).then((rows) => rows.map((row) => row.id));
for (const runId of runIds) await cancelRunInternal(runId, "Cancelled due to budget pause");
for (const runId of runIds) await cancelRunInternal(runId, "Cancelled due to budget pause", { budgetEnforcement: scope });
await cancelPendingWakeupsForBudgetScope(scope);
}
@@ -0,0 +1,59 @@
import { createHash, randomUUID } from "node:crypto";
import { promises as fs } from "node:fs";
import path from "node:path";
import { z } from "zod";
import { resolvePaperclipInstanceRoot } from "../home-paths.js";
import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js";
const recoverySchema = z.object({
companyId: z.string().uuid(), grantId: z.string().uuid(), connectionId: z.string().uuid(), secretId: z.string().uuid(),
baseHash: z.string().length(64), value: z.string().min(1),
});
type Recovery = z.infer<typeof recoverySchema>;
export const quotaCredentialHash = (value: string) => createHash("sha256").update(value).digest("hex");
export const quotaCredentialRecoveryPath = (companyId: string, grantId: string) =>
path.join(resolvePaperclipInstanceRoot(), "quota-credential-recovery", companyId, `${grantId}.json`);
async function syncDirectory(directory: string) {
if (process.platform === "win32") return;
const handle = await fs.open(directory, "r");
try { await handle.sync(); } finally { await handle.close(); }
}
/** Called only while holding the company credential mutation lock. The local
* encrypted provider uses the instance secrets master key, including when the
* destination vault is remote. No plaintext token enters this journal. */
export function quotaCredentialRecovery(companyId: string, grantId: string) {
const file = quotaCredentialRecoveryPath(companyId, grantId);
const directory = path.dirname(file);
async function writeEncrypted(target: string, value: string) {
const encrypted = await localEncryptedProvider.createVersion({ value });
await fs.mkdir(directory, { recursive: true, mode: 0o700 });
const temporary = `${target}.${randomUUID()}.tmp`;
const handle = await fs.open(temporary, "wx", 0o600);
try {
try { await handle.writeFile(JSON.stringify(encrypted.material)); await handle.sync(); }
finally { await handle.close(); }
await fs.rename(temporary, target);
await syncDirectory(directory);
} catch (error) { await fs.rm(temporary, { force: true }); throw error; }
}
return {
// Verify encryption, storage, and fsync before exchanging a single-use token.
prepare: async () => {
const probe = `${file}.${randomUUID()}.probe`;
try { await writeEncrypted(probe, "credential-recovery-probe"); }
finally { await fs.rm(probe, { force: true }); }
},
read: async (): Promise<Recovery | null> => {
let material: string;
try { material = await fs.readFile(file, "utf8"); }
catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; throw error; }
const pending = recoverySchema.parse(JSON.parse(await localEncryptedProvider.resolveVersion({ material: JSON.parse(material), externalRef: null })));
if (pending.companyId !== companyId || pending.grantId !== grantId) throw new Error("credentials_unavailable");
return pending;
},
write: (input: Recovery) => writeEncrypted(file, JSON.stringify(recoverySchema.parse(input))),
clear: async () => { await fs.rm(file, { force: true }); },
};
}
+13 -1
View File
@@ -1,6 +1,6 @@
import { usdToUnits, unitsToCents } from "@paperclipai/shared";
import { and, asc, eq, inArray, sql } from "drizzle-orm";
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, issues, projects, type Db } from "@paperclipai/db";
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, type Db } from "@paperclipai/db";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { createCostEventInTransaction } from "./costs.js";
import { budgetService, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
@@ -25,6 +25,13 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
}
const preProviderFailure = object(object(run.resultJson).executionRecovery).providerWorkStarted === false
|| object(run.usageJson).accountingProviderWorkStarted === false;
// A failed native turn is not necessarily the end of this heartbeat: the
// coordinator can resume the same run. Only its closed provider boundary
// authorizes acknowledgement; finalization may still run after a result.
if (run.runtimeMode === "native" && !preProviderFailure) {
const [coordinator] = await tx.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id));
if (coordinator && !coordinator.resultId && coordinator.phase !== "terminal_failure") return false;
}
// A stop can mark the run terminal before the provider's shutdown receipt
// arrives. Keep the debt pending instead of acknowledging an invented zero
// and permanently discarding that late receipt.
@@ -118,6 +125,11 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) {
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.costAccountingPending, true), inArray(heartbeatRuns.status, terminalStatuses),
sql`(${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false'
or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false'
or not exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
and ${heartbeatRuns.runtimeMode} = 'native' and ${nativeRunFinalizations.resultId} is null
and ${nativeRunFinalizations.phase} <> 'terminal_failure'))`,
sql`(${heartbeatRuns.usageJson}->>'accountingReceiptReady' is distinct from 'false' or ${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false' or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false')`,
)).orderBy(asc(heartbeatRuns.updatedAt), asc(heartbeatRuns.id)).limit(100);
let accounted = 0;
+12 -2
View File
@@ -121,12 +121,22 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
const probe = await fs.open(path.join(directory, `${sourceId}.probe`), "wx", 0o600);
try { await probe.sync(); } finally { await probe.close(); await fs.rm(path.join(directory, `${sourceId}.probe`)); }
await syncDirectory(directory);
const bound = await db.update(heartbeatRuns).set({ usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0 })}::jsonb` })
.where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt))).returning({ id: heartbeatRuns.id });
const bound = await db.update(heartbeatRuns).set({ usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0, accountingReceiptReady: false })}::jsonb` })
.where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt))).returning({ id: heartbeatRuns.id, runtimeMode: heartbeatRuns.runtimeMode, usageJson: heartbeatRuns.usageJson });
if (bound.length !== 1) throw conflict("Run cannot accept a new usage recorder");
let sequence = 0;
let currentAttempt = "default";
const attempts = new Map<string, AdapterUsageCheckpoint>();
// Native run deltas retain their baseline across same-run recovery. Restore
// the previous snapshot as the same attempt, never add it a second time.
// An empty final result must not erase the last durable observation.
const priorReceiptId = bound[0].usageJson?.accountingReceiptId;
if (bound[0].runtimeMode === "native" && typeof priorReceiptId === "string") {
const [previous] = await db.select().from(runUsageReceipts).where(and(
eq(runUsageReceipts.id, priorReceiptId), eq(runUsageReceipts.runId, input.runId), eq(runUsageReceipts.companyId, input.companyId),
));
if (previous) attempts.set("default", envelopeSchema.parse(previous.receiptJson).receipt as AdapterUsageCheckpoint);
}
let chain = Promise.resolve();
let captureFailed = false;
async function capture(raw: AdapterUsageCheckpoint) {