mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-23 17:33:51 +02:00
Merge accounting lifecycle fences and durable OAuth recovery
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
commit
8f10b0b7c5
20 files changed
+664
-177
No files matched your search
+19
-1
@@ -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.
|
||||
@@ -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" });
|
||||
|
||||
@@ -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 () => {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 };
|
||||
});
|
||||
}
|
||||
@@ -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 })
|
||||
|
||||
@@ -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 }); },
|
||||
};
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user