fix: save initial hire budgets within lifecycle transaction

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/agent-permissions-routes.test.ts

@@ -2383,7 +2383,8 @@ describe("agent permission routes", () => {
     const app = createApp({ type: "board", userId: "operator-user", source: "local_implicit", companyIds: [companyId] });
     const response = await requestApp(app, url => request(url).post(`/api/companies/${companyId}/${route}`).send({ name: "New employee", role: "engineer", adapterType: "process", budgetMonthlyCents: 5000 }));
     expect(response.status).toBe(201);
-    expect(mockBudgetService.upsertPolicy).toHaveBeenCalledWith(companyId, expect.objectContaining({ scopeType: "agent", scopeId: agentId, amount: 5000 }), "operator-user");
+    expect(mockAgentService.create).toHaveBeenCalledWith(companyId, expect.objectContaining({ budgetMonthlyCents: 5000 }), expect.objectContaining({ createdByUserId: "operator-user" }));
+    expect(mockBudgetService.upsertPolicy).not.toHaveBeenCalled();
   });
   it("rejects heartbeat cancellation outside the caller company scope", async () => {
     mockHeartbeatService.getRun.mockResolvedValue({

server/src/modules/agent-lifecycle/__tests__/lifecycle.test.ts

@@ -298,6 +298,35 @@ const support = await getEmbeddedPostgresTestSupport();
     expect((await current(setupId)).lifecycleHolds).toContain("setup");
   });
 
+  it("saves the immediate hire budget and actual operator audit in its transaction", async () => {
+    const agent = await createAgentLifecycle(db).requestHire(companyId, { name: "Capped hire", adapterType: "process", budgetMonthlyCents: 5000 }, { createdByUserId: "operator-user" });
+    expect(await db.select().from(budgetPolicies).where(eq(budgetPolicies.scopeId, agent.id))).toMatchObject([{ amount: 5000 }]);
+    const audit = await db.select().from(activityLog).where(and(eq(activityLog.companyId, companyId), eq(activityLog.action, "budget.policy_upserted")));
+    expect(audit).toMatchObject([{ actorType: "user", actorId: "operator-user", agentId: null }]);
+  });
+
+  it("does not activate a proposed hire budget before human approval", async () => {
+    const agent = await createAgentLifecycle(db).requestHire(companyId, { name: "Proposed cap", status: "pending_approval", budgetMonthlyCents: 5000 }, { createdByUserId: "operator-user" });
+    expect(agent.budgetMonthlyCents).toBe(5000);
+    expect(await db.select().from(budgetPolicies).where(eq(budgetPolicies.scopeId, agent.id))).toHaveLength(0);
+  });
+
+  it("rolls back the hire, grants, policy and audit together when its initial budget write fails", async () => {
+    const effects = createAgentLifecycleEffects();
+    const setAgentBudget = vi.fn(async (...args: Parameters<typeof effects.setAgentBudget>) => {
+      await effects.setAgentBudget(...args);
+      throw new Error("Initial budget failed");
+    });
+    const lifecycle = createLifecycleCommands(db, { ...effects, setAgentBudget });
+    await expect(lifecycle.requestHire(companyId, { name: "Budget rollback", adapterType: "process", budgetMonthlyCents: 5000 }, { createdByUserId: "operator-user" }))
+      .rejects.toThrow("Initial budget failed");
+    expect(setAgentBudget).toHaveBeenCalledWith(expect.anything(), expect.any(Array), companyId, expect.any(String), 5000, "operator-user");
+    expect(await db.select().from(agents).where(eq(agents.companyId, companyId))).toHaveLength(0);
+    expect(await db.select().from(principalPermissionGrants).where(eq(principalPermissionGrants.companyId, companyId))).toHaveLength(0);
+    expect(await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, companyId))).toHaveLength(0);
+    expect(await db.select().from(activityLog).where(eq(activityLog.companyId, companyId))).toHaveLength(0);
+  });
+
   it("rolls back a hire and its grants when hire initialization fails", async () => {
     const effects = createAgentLifecycleEffects();
     const lifecycle = createLifecycleCommands(db, { ...effects,

server/src/modules/agent-lifecycle/adapters/records.ts

@@ -13,7 +13,7 @@ export function agentRecords(db: Db, effects: LifecycleEffects) {
     getById,
     create: async (companyId: string, data: CreateAgentData, options?: CreateAgentOptions) => {
       if (Object.keys(data).some(key => key.startsWith("lifecycle"))) throw conflict("Lifecycle fields belong to the lifecycle module");
-      return effects.transaction(db, companyId, async (txDb) => {
+      return effects.transaction(db, companyId, async (txDb, publications) => {
         const tx = txDb;
         const setupId = options?.setupAgentId ?? (options?.setupOnly ? data.id : undefined);
         const [setup] = setupId ? await tx.select().from(agents).where(eq(agents.id, setupId)).for("update") : [];
@@ -58,6 +58,10 @@ export function agentRecords(db: Db, effects: LifecycleEffects) {
           .returning()
           .then((rows) => rows[0]);
         await effects.initializeHire(txDb, created, options);
+        if (!options?.setupOnly && status !== "pending_approval" && status !== "terminated"
+          && created.budgetMonthlyCents > 0 && options?.createdByUserId) {
+          await effects.setAgentBudget(txDb, publications, companyId, created.id, created.budgetMonthlyCents, options.createdByUserId);
+        }
         if (!options?.setupOnly && created.status !== "pending_approval" && created.status !== "terminated") {
           await effects.recordCreation(txDb, companyId, created.id);
         }

server/src/routes/agents.ts

@@ -4601,14 +4601,7 @@ export function agentRoutes(
 
       let approval: Awaited<ReturnType<typeof approvalsSvc.getById>> | null = null;
       const actor = getActorInfo(req);
-      if (!requiresApproval && agent.budgetMonthlyCents > 0) {
-        await budgets.upsertPolicy(companyId, {
-          scopeType: "agent",
-          scopeId: agent.id,
-          amount: agent.budgetMonthlyCents,
-          windowKind: "calendar_month_utc",
-        }, actor.actorId);
-      }
+
 
       if (requiresApproval) {
         const requestedAdapterType = normalizedHireInput.adapterType ?? agent.adapterType;
@@ -4918,18 +4911,7 @@ export function agentRoutes(
     );
     await builtInAgentService(db).ensureCompanyDefaultAgentGrants(companyId);
 
-    if (agent.budgetMonthlyCents > 0) {
-      await budgets.upsertPolicy(
-        companyId,
-        {
-          scopeType: "agent",
-          scopeId: agent.id,
-          amount: agent.budgetMonthlyCents,
-          windowKind: "calendar_month_utc",
-        },
-        actor.actorType === "user" ? actor.actorId : null,
-      );
-    }
+
 
     res.status(201).json({ ...redactAgentRowForResponse(agent), ...(configurationWarnings.length ? { configurationWarnings } : {}) });
   });