@desplega.ai/agent-swarm 1.71.2 → 1.72.0

package/src/tools/task-action.ts

@@ -1,5 +1,10 @@
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import * as z from "zod";
+import { canClaim } from "@/be/budget-admission";
+import {
+  type BudgetRefusalContext,
+  emitBudgetRefusalSideEffects,
+} from "@/be/budget-refusal-notify";
 import {
   acceptTask,
   checkDependencies,
@@ -14,12 +19,13 @@ import {
   moveTaskFromBacklog,
   moveTaskToBacklog,
   reassociateSessionLogs,
+  recordBudgetRefusalNotification,
   rejectTask,
   releaseTask,
   updateTaskClaudeSessionId,
 } from "@/be/db";
 import { createToolRegistrar } from "@/tools/utils";
-import { AgentTaskSchema } from "@/types";
+import { AgentTaskSchema, BudgetRefusalCauseSchema } from "@/types";
 
 const TaskActionSchema = z.enum([
   "create",
@@ -83,6 +89,14 @@ export const registerTaskActionTool = (server: McpServer) => {
         success: z.boolean(),
         message: z.string(),
         task: AgentTaskSchema.optional(),
+        // Phase 3: budget-admission refusal fields. Populated only on
+        // `accept` action when the per-agent or global daily budget is blown.
+        refusalCause: BudgetRefusalCauseSchema.optional(),
+        agentSpend: z.number().optional(),
+        agentBudget: z.number().optional(),
+        globalSpend: z.number().optional(),
+        globalBudget: z.number().optional(),
+        resetAt: z.string().optional(),
       }),
     },
     async (input, requestInfo, _meta) => {
@@ -241,6 +255,60 @@ export const registerTaskActionTool = (server: McpServer) => {
                 message: `Task "${taskId}" has unmet dependencies: ${blockedBy.join(", ")}. Cannot accept until dependencies are completed.`,
               };
             }
+            // Budget admission gate (Phase 3). Same in-transaction placement
+            // as the /api/poll gates so capacity AND budget share atomicity.
+            // Phase 5: record dedup row + capture side-effect context for the
+            // after-commit lead follow-up + workflow event-bus emit.
+            const admission = canClaim(agentId, new Date());
+            if (!admission.allowed) {
+              const causeMsg =
+                admission.cause === "agent"
+                  ? "agent daily budget exceeded"
+                  : "global daily budget exceeded";
+              const utcDate = new Date().toISOString().slice(0, 10);
+              const dedup = recordBudgetRefusalNotification({
+                taskId,
+                date: utcDate,
+                agentId,
+                cause: admission.cause,
+                agentSpendUsd: admission.agentSpend,
+                agentBudgetUsd: admission.agentBudget,
+                globalSpendUsd: admission.globalSpend,
+                globalBudgetUsd: admission.globalBudget,
+              });
+              return {
+                success: false,
+                message: `Refused: ${causeMsg}. Resets at ${admission.resetAt}.`,
+                refusalCause: admission.cause,
+                ...(admission.agentSpend !== undefined && { agentSpend: admission.agentSpend }),
+                ...(admission.agentBudget !== undefined && { agentBudget: admission.agentBudget }),
+                ...(admission.globalSpend !== undefined && { globalSpend: admission.globalSpend }),
+                ...(admission.globalBudget !== undefined && {
+                  globalBudget: admission.globalBudget,
+                }),
+                resetAt: admission.resetAt,
+                refusalSideEffects: {
+                  context: {
+                    task: {
+                      id: existingTask.id,
+                      task: existingTask.task,
+                      slackChannelId: existingTask.slackChannelId,
+                      slackThreadTs: existingTask.slackThreadTs,
+                      slackUserId: existingTask.slackUserId,
+                    },
+                    agentId,
+                    date: utcDate,
+                    cause: admission.cause,
+                    agentSpendUsd: admission.agentSpend,
+                    agentBudgetUsd: admission.agentBudget,
+                    globalSpendUsd: admission.globalSpend,
+                    globalBudgetUsd: admission.globalBudget,
+                    resetAt: admission.resetAt,
+                  } satisfies BudgetRefusalContext,
+                  inserted: dedup.inserted,
+                },
+              };
+            }
             const acceptedTask = acceptTask(taskId, agentId);
             if (!acceptedTask) {
               return { success: false, message: `Failed to accept task "${taskId}".` };
@@ -334,11 +402,33 @@ export const registerTaskActionTool = (server: McpServer) => {
 
       const result = txn();
 
+      // Phase 5: when the accept gate refused, run after-commit side
+      // effects (lead follow-up + workflow bus). The dedup row was recorded
+      // inside the txn; this just consumes the captured context.
+      if (
+        "refusalSideEffects" in result &&
+        result.refusalSideEffects &&
+        typeof result.refusalSideEffects === "object"
+      ) {
+        const sideEffects = result.refusalSideEffects as {
+          context: BudgetRefusalContext;
+          inserted: boolean;
+        };
+        emitBudgetRefusalSideEffects(sideEffects.context, sideEffects.inserted);
+      }
+
+      // Strip the internal-only `refusalSideEffects` field from the wire
+      // response — workers receive only the public refusal envelope.
+      const { refusalSideEffects: _omit, ...publicResult } = result as {
+        refusalSideEffects?: unknown;
+        [key: string]: unknown;
+      };
+
       return {
         content: [{ type: "text", text: result.message }],
         structuredContent: {
           yourAgentId: agentId,
-          ...result,
+          ...publicResult,
         },
       };
     },

package/src/tools/templates.ts

@@ -84,3 +84,32 @@ Decide whether to reassign, retry, or handle the failure. Use \`get-task-details
   ],
   category: "task_lifecycle",
 });
+
+// ============================================================================
+// Budget refusal follow-up (Phase 5: created when an agent is refused due to
+// per-agent or global daily budget exhaustion)
+// ============================================================================
+
+registerTemplate({
+  eventType: "task.budget.refused",
+  header: "",
+  defaultBody: `Budget refusal \u2014 task is blocked.
+
+Cause: {{cause}}
+Agent: {{agent_name}}
+Task: {{task_desc}}
+Spend / budget: {{spend_summary}}
+Resets at: {{reset_at}}
+
+Decide whether to raise the budget, reassign, or wait for the daily reset.
+Use \`get-task-details\` with taskId "{{task_id}}" for full details.`,
+  variables: [
+    { name: "cause", description: "'agent' or 'global'" },
+    { name: "agent_name", description: "Refusing agent name or ID prefix" },
+    { name: "task_desc", description: "Task description (truncated to 200 chars)" },
+    { name: "spend_summary", description: 'Formatted "$X / $Y" pair' },
+    { name: "reset_at", description: "UTC reset time (human readable)" },
+    { name: "task_id", description: "Original task ID" },
+  ],
+  category: "task_lifecycle",
+});

package/src/types.ts

@@ -68,6 +68,32 @@ export const AgentTaskSourceSchema = z.enum([
 ]);
 export type AgentTaskSource = z.infer<typeof AgentTaskSourceSchema>;
 
+// ---------------------------------------------------------------------------
+// Harness Provider
+// ---------------------------------------------------------------------------
+// String identifiers accepted by `HARNESS_PROVIDER` and the
+// `createProviderAdapter` factory in `src/providers/index.ts`. Keep this in
+// sync with the factory's switch and the unknown-provider error message.
+export const ProviderNameSchema = z.enum(["claude", "codex", "pi", "devin", "claude-managed"]);
+export type ProviderName = z.infer<typeof ProviderNameSchema>;
+
+export type DevinProviderMeta = {
+  sessionUrl: string;
+  maxAcuLimit?: number;
+  acuCostUsd?: number;
+};
+
+// These providers do not have metadata yet.
+type NoProviderMeta = Record<string, never>;
+
+export type ProviderMetaMap = {
+  devin: DevinProviderMeta;
+  claude: NoProviderMeta;
+  codex: NoProviderMeta;
+  pi: NoProviderMeta;
+  "claude-managed": NoProviderMeta;
+};
+
 export const AgentTaskSchema = z.object({
   id: z.uuid(),
   agentId: z.uuid().nullable(), // Nullable for unassigned tasks
@@ -169,6 +195,10 @@ export const AgentTaskSchema = z.object({
   // agent-swarm package version at task creation time. Enables benchmarking
   // performance across releases. Nullable for rows created before tracking was added.
   swarmVersion: z.string().optional(),
+
+  // Provider tracking — which harness provider ran this task
+  provider: ProviderNameSchema.optional(),
+  providerMeta: z.record(z.string(), z.unknown()).optional(),
 });
 
 // ============================================================================
@@ -365,6 +395,11 @@ export const AgentLogEventTypeSchema = z.enum([
   "service_registered",
   "service_unregistered",
   "service_status_change",
+  // Phase 6: budget / pricing operator-mutation audit log events
+  "budget.upserted",
+  "budget.deleted",
+  "pricing.inserted",
+  "pricing.deleted",
 ]);
 
 export const AgentLogSchema = z.object({
@@ -396,6 +431,9 @@ export const SessionLogSchema = z.object({
 export type SessionLog = z.infer<typeof SessionLogSchema>;
 
 // Session Cost Types (aggregated cost data per session)
+export const SessionCostSourceSchema = z.enum(["harness", "pricing-table"]);
+export type SessionCostSource = z.infer<typeof SessionCostSourceSchema>;
+
 export const SessionCostSchema = z.object({
   id: z.uuid(),
   sessionId: z.string(),
@@ -410,6 +448,10 @@ export const SessionCostSchema = z.object({
   numTurns: z.number().int().min(1),
   model: z.string(),
   isError: z.boolean().default(false),
+  // Phase 6: where the recorded totalCostUsd came from. New rows write the
+  // actual source ('pricing-table' when the API recomputed Codex USD from DB
+  // pricing rows, 'harness' otherwise). Defaults to 'harness' for back-compat.
+  costSource: SessionCostSourceSchema.default("harness"),
   createdAt: z.iso.datetime(),
 });
 
@@ -1122,3 +1164,77 @@ export const ContextSnapshotSchema = z.object({
 });
 
 export type ContextSnapshot = z.infer<typeof ContextSnapshotSchema>;
+
+// ============================================================================
+// Budgets + Pricing (per-agent daily cost budget — V1)
+// ============================================================================
+//
+// Timestamp convention for these schemas: number = epoch milliseconds (UTC).
+// This is a deliberate divergence from the rest of types.ts (which uses
+// `z.iso.datetime()` strings) so that the price-book "largest
+// effective_from <= now" lookup is a pure integer comparison. Matches the
+// SQL columns in migration 046_budgets_and_pricing.sql verbatim.
+
+export const BudgetScopeSchema = z.enum(["global", "agent"]);
+export type BudgetScope = z.infer<typeof BudgetScopeSchema>;
+
+export const BudgetSchema = z.object({
+  scope: BudgetScopeSchema,
+  scopeId: z.string(), // '' (empty string) for the global row
+  dailyBudgetUsd: z.number().nonnegative(),
+  createdAt: z.number(), // epoch ms
+  lastUpdatedAt: z.number(), // epoch ms
+});
+export type Budget = z.infer<typeof BudgetSchema>;
+
+export const PricingProviderSchema = z.enum(["claude", "codex", "pi"]);
+export type PricingProvider = z.infer<typeof PricingProviderSchema>;
+
+export const PricingTokenClassSchema = z.enum(["input", "cached_input", "output"]);
+export type PricingTokenClass = z.infer<typeof PricingTokenClassSchema>;
+
+export const PricingRowSchema = z.object({
+  provider: PricingProviderSchema,
+  model: z.string(),
+  tokenClass: PricingTokenClassSchema,
+  effectiveFrom: z.number().nonnegative(), // epoch ms; 0 = seed
+  pricePerMillionUsd: z.number().nonnegative(),
+  createdAt: z.number(), // epoch ms
+  lastUpdatedAt: z.number(), // epoch ms
+});
+export type PricingRow = z.infer<typeof PricingRowSchema>;
+
+export const BudgetRefusalCauseSchema = z.enum(["agent", "global"]);
+export type BudgetRefusalCause = z.infer<typeof BudgetRefusalCauseSchema>;
+
+export const BudgetRefusalNotificationSchema = z.object({
+  taskId: z.string(),
+  date: z.string(), // 'YYYY-MM-DD' UTC
+  agentId: z.string(),
+  cause: BudgetRefusalCauseSchema,
+  agentSpendUsd: z.number().nullable().optional(),
+  agentBudgetUsd: z.number().nullable().optional(),
+  globalSpendUsd: z.number().nullable().optional(),
+  globalBudgetUsd: z.number().nullable().optional(),
+  followUpTaskId: z.string().nullable().optional(),
+  createdAt: z.number(), // epoch ms
+});
+export type BudgetRefusalNotification = z.infer<typeof BudgetRefusalNotificationSchema>;
+
+/**
+ * Phase 3 — `budget_refused` is the new variant of the `/api/poll` trigger
+ * envelope returned when an admission gate (`canClaim`) refuses to let the
+ * agent take a task. Older workers receiving this discriminator fall through
+ * to default polling without back-off (degrades gracefully); Phase 4 teaches
+ * the runner to recognize it.
+ */
+export const BudgetRefusedTriggerSchema = z.object({
+  type: z.literal("budget_refused"),
+  cause: BudgetRefusalCauseSchema,
+  agentSpend: z.number().optional(),
+  agentBudget: z.number().optional(),
+  globalSpend: z.number().optional(),
+  globalBudget: z.number().optional(),
+  resetAt: z.string(), // ISO 8601, next UTC midnight
+});
+export type BudgetRefusedTrigger = z.infer<typeof BudgetRefusedTriggerSchema>;

package/src/utils/budget-backoff.ts

@@ -0,0 +1,34 @@
+/**
+ * Phase 4 — exponential back-off for `budget_refused` poll responses.
+ *
+ * The worker poll loop short-circuits on `trigger.type === "budget_refused"`
+ * to avoid busy-looping the API while the agent is over-budget. Each
+ * consecutive refusal doubles the sleep, capped at 5 minutes (per the
+ * scoping decision in the per-agent-daily-cost-budget plan).
+ *
+ * This module is a pure helper so it can be unit-tested in isolation
+ * without standing up the full poll loop.
+ */
+
+/** Hard cap on back-off, per the plan (5 minutes). */
+export const BUDGET_BACKOFF_CAP_MS = 5 * 60 * 1000;
+
+/**
+ * Compute the back-off delay for the Nth consecutive `budget_refused`
+ * trigger. The first refusal sleeps `basePollMs`; each subsequent one
+ * doubles, capped at {@link BUDGET_BACKOFF_CAP_MS}.
+ *
+ * @param consecutiveRefusals - 1-indexed count of consecutive refusals
+ *        (i.e. the Nth refusal in a row, including the current one).
+ *        Must be >= 1; values < 1 fall back to `basePollMs`.
+ * @param basePollMs - Base poll interval (today's `PollIntervalMs`, ~2s).
+ * @returns Sleep duration in milliseconds, capped at 5 minutes.
+ */
+export function computeBudgetBackoffMs(consecutiveRefusals: number, basePollMs: number): number {
+  const n = Math.max(1, Math.floor(consecutiveRefusals));
+  // 2 ** (n-1) grows quickly; cap before multiplying to avoid Infinity for
+  // pathological inputs. JS handles 2 ** 30 fine but cap-first is cleaner.
+  const exponent = Math.min(30, n - 1);
+  const raw = basePollMs * 2 ** exponent;
+  return Math.min(BUDGET_BACKOFF_CAP_MS, raw);
+}

package/src/utils/credentials.ts

@@ -5,6 +5,7 @@ export const CREDENTIAL_POOL_VARS = [
   "OPENROUTER_API_KEY",
   "OPENAI_API_KEY",
   "CODEX_OAUTH",
+  "DEVIN_API_KEY",
 ] as const;
 
 /**
@@ -23,6 +24,7 @@ export const PROVIDER_CREDENTIAL_VARS: Record<string, readonly string[]> = {
   // pi-mono accepts either router or anthropic keys
   pi: ["OPENROUTER_API_KEY", "ANTHROPIC_API_KEY"],
   codex: ["OPENAI_API_KEY", "CODEX_OAUTH"],
+  devin: ["DEVIN_API_KEY"],
 };
 
 /**
@@ -43,6 +45,8 @@ export function deriveProviderFromKeyType(keyType: string): string {
     case "OPENAI_API_KEY":
     case "CODEX_OAUTH":
       return "codex";
+    case "DEVIN_API_KEY":
+      return "devin";
     default:
       return "claude";
   }

package/src/utils/provider-metadata.ts

@@ -0,0 +1,9 @@
+import type { ProviderMetaMap, ProviderName } from "@/types";
+
+export function parseProviderMeta<TProvider extends ProviderName>(
+  rawProvider: TProvider | null | undefined,
+  rawMeta: string | undefined | null,
+): ProviderMetaMap[TProvider] | undefined {
+  if (!rawMeta || !rawProvider) return undefined;
+  return JSON.parse(rawMeta) as ProviderMetaMap[TProvider];
+}