agent-relay-server 0.77.0 → 0.78.0

package/src/quota-advisory.ts

@@ -0,0 +1,300 @@
+import { isRecord } from "agent-relay-sdk";
+import { getDb } from "./db/connection.ts";
+import { emitRelayEvent } from "./events";
+import { routeNotification } from "./notification-router";
+import { DEFAULT_QUOTA_THRESHOLD_CONFIG, validateQuotaThresholdConfig } from "./quota-threshold-config";
+import { parseJson } from "./utils";
+import { setProviderQuotaAdvisoryHandler, type ProviderQuotaAdvisoryInput } from "./db/provider-quotas.ts";
+import type { QuotaState } from "./types";
+
+interface QuotaAdvisoryWindowState {
+  lastQuotaAdvisoryBand?: number;
+  lastQuotaAdvisoryAt?: number;
+  lastRemainingPercent?: number;
+  lastResetsAt?: number;
+}
+
+type QuotaAdvisoryState = Record<string, QuotaAdvisoryWindowState>;
+type QuotaAdvisoryStage = "warn" | "critical" | "reset";
+
+type AgentRecipientRow = {
+  id: string;
+  name: string;
+  kind?: string | null;
+  machine?: string | null;
+  rig?: string | null;
+  ready?: number | null;
+  status: string;
+  meta?: string | null;
+};
+
+const RESET_ROUNDING_PERCENT = 99;
+
+export function installQuotaAdvisoryHandler(): void {
+  setProviderQuotaAdvisoryHandler(handleProviderQuotaAdvisory);
+}
+
+function handleProviderQuotaAdvisory(input: ProviderQuotaAdvisoryInput): string {
+  const config = effectiveQuotaNotificationConfig().quotaThreshold;
+  const state = parseJson<QuotaAdvisoryState>(input.previousAdvisoryState ?? "{}", {});
+  const nextState: QuotaAdvisoryState = { ...state };
+
+  for (const window of input.quota.windows) {
+    const previousWindow = input.previousQuota?.windows.find((candidate) => candidate.name === window.name);
+    const remaining = quotaRemaining(window.utilization);
+    const remainingPercent = percent(remaining);
+    const previous = nextState[window.name];
+    const reset = previous?.lastQuotaAdvisoryBand !== undefined && isQuotaReset(remaining, previousWindow, window, config.resetRemaining);
+    if (reset) {
+      notifyQuotaAdvisory({
+        provider: input.provider,
+        accountKey: input.accountKey,
+        windowName: window.name,
+        stage: "reset",
+        band: 100,
+        remainingPercent,
+        resetsAt: window.resetsAt,
+        recipients: quotaAdvisoryRecipients(input.provider, input.accountKey, input.sourceAgentIds, input.now),
+        headroomHint: crossProviderHeadroomHint(input.provider),
+        now: input.now,
+      });
+      delete nextState[window.name];
+      continue;
+    }
+
+    if (remaining > config.warnRemaining) {
+      if (previous) delete nextState[window.name];
+      continue;
+    }
+
+    const band = quotaRemainingBand(remaining, config.bandSize);
+    if (previous?.lastQuotaAdvisoryBand === band) continue;
+    const stage: QuotaAdvisoryStage = remaining <= config.criticalRemaining ? "critical" : "warn";
+    notifyQuotaAdvisory({
+      provider: input.provider,
+      accountKey: input.accountKey,
+      windowName: window.name,
+      stage,
+      band,
+      remainingPercent,
+      resetsAt: window.resetsAt,
+      recipients: quotaAdvisoryRecipients(input.provider, input.accountKey, input.sourceAgentIds, input.now),
+      headroomHint: crossProviderHeadroomHint(input.provider),
+      now: input.now,
+    });
+    nextState[window.name] = {
+      lastQuotaAdvisoryBand: band,
+      lastQuotaAdvisoryAt: input.now,
+      lastRemainingPercent: remainingPercent,
+      ...(window.resetsAt !== undefined ? { lastResetsAt: window.resetsAt } : {}),
+    };
+  }
+
+  return JSON.stringify(nextState);
+}
+
+function notifyQuotaAdvisory(input: {
+  provider: string;
+  accountKey: string;
+  windowName: string;
+  stage: QuotaAdvisoryStage;
+  band: number;
+  remainingPercent: number;
+  resetsAt?: number;
+  recipients: string[];
+  headroomHint: string;
+  now: number;
+}): void {
+  const payload = {
+    kind: "agent.quota_threshold",
+    provider: input.provider,
+    accountKey: input.accountKey,
+    window: input.windowName,
+    stage: input.stage,
+    band: input.band,
+    remainingPercent: input.remainingPercent,
+    ...(input.resetsAt !== undefined ? { resetsAt: input.resetsAt } : {}),
+    recipients: input.recipients,
+  };
+
+  emitRelayEvent({
+    type: "agent.quota_threshold",
+    source: "server",
+    subject: `${input.provider}:${input.accountKey}:${input.windowName}`,
+    data: payload,
+  });
+
+  const notifications = effectiveQuotaNotificationConfig();
+  if (!notifications.enabled || !notifications.quotaThreshold.enabled) return;
+
+  routeNotification({
+    type: "agent.quota_threshold",
+    scope: {},
+    recipients: input.recipients,
+    message: {
+      subject: input.stage === "reset"
+        ? `${windowLabel(input.windowName)} ${providerLabel(input.provider)} quota reset`
+        : `${windowLabel(input.windowName)} ${providerLabel(input.provider)} quota ${input.stage === "critical" ? "critical" : "warning"}`,
+      body: quotaAdvisoryBody(input),
+      payload,
+      replyExpected: false,
+    },
+    emittedAt: input.now,
+  });
+}
+
+function effectiveQuotaNotificationConfig(): { enabled: boolean; quotaThreshold: ReturnType<typeof validateQuotaThresholdConfig> } {
+  const row = getDb().query("SELECT value FROM config WHERE namespace = 'notifications' AND key = 'default'").get() as { value?: string | null } | undefined;
+  const value = parseJson<Record<string, unknown>>(row?.value ?? "{}", {});
+  return {
+    enabled: value.enabled === undefined ? true : value.enabled === true,
+    quotaThreshold: validateQuotaThresholdConfig({ ...DEFAULT_QUOTA_THRESHOLD_CONFIG, ...(isRecord(value.quotaThreshold) ? value.quotaThreshold : {}) }),
+  };
+}
+
+function quotaAdvisoryBody(input: {
+  provider: string;
+  windowName: string;
+  stage: QuotaAdvisoryStage;
+  band: number;
+  remainingPercent: number;
+  resetsAt?: number;
+  headroomHint: string;
+  now: number;
+}): string {
+  const provider = providerLabel(input.provider);
+  const window = windowLabel(input.windowName);
+  if (input.stage === "reset" || input.remainingPercent >= RESET_ROUNDING_PERCENT) {
+    return `${window} ${provider} quota reset to 100%. ${input.headroomHint}`;
+  }
+  const reset = input.resetsAt === undefined ? "reset time unknown" : `resets in ${formatDuration(input.resetsAt - input.now)}`;
+  return `${window} ${provider} quota below ${input.band}% remaining - ${input.remainingPercent}% remaining, ${reset}. ${input.headroomHint}`;
+}
+
+function quotaAdvisoryRecipients(provider: string, accountKey: string, sourceAgentIds: string[], now: number): string[] {
+  const sourceHosts = agentHosts(sourceAgentIds);
+  const accountHost = hostFromAccountKey(accountKey);
+  const candidates = getDb().query(`
+    SELECT id, name, kind, machine, rig, ready, status, meta
+    FROM agents
+    WHERE kind = 'provider'
+      AND ready = 1
+      AND status NOT IN ('offline', 'stale')
+  `).all() as AgentRecipientRow[];
+  return candidates
+    .filter((agent) => isSpawnCapableAgent(agent, now))
+    .filter((agent) => isBoundToQuotaAccount(agent, provider, accountKey, accountHost, sourceHosts, sourceAgentIds))
+    .map((agent) => agent.id);
+}
+
+function isSpawnCapableAgent(agent: AgentRecipientRow, now: number): boolean {
+  const meta = parseJson<Record<string, unknown>>(agent.meta ?? "{}", {});
+  const auth = isRecord(meta.auth) ? meta.auth : undefined;
+  const jti = typeof auth?.jti === "string" ? auth.jti : undefined;
+  if (!jti) return false;
+  const row = getDb().query("SELECT scope, revoked_at, expires_at FROM tokens WHERE jti = ?").get(jti) as { scope: string; revoked_at?: number | null; expires_at?: number | null } | undefined;
+  if (!row || row.revoked_at) return false;
+  const nowSeconds = Math.floor(now / 1000);
+  if (row.expires_at !== null && row.expires_at !== undefined && row.expires_at <= nowSeconds) return false;
+  const scopes = parseJson<string[]>(row.scope, []);
+  return scopes.includes("command:spawn") || scopes.includes("command:*") || scopes.includes("admin:*") || scopes.includes("*");
+}
+
+function isBoundToQuotaAccount(agent: AgentRecipientRow, provider: string, accountKey: string, accountHost: string | undefined, sourceHosts: Set<string>, sourceAgentIds: string[]): boolean {
+  if (sourceAgentIds.includes(agent.id)) return true;
+  const meta = parseJson<Record<string, unknown>>(agent.meta ?? "{}", {});
+  if (stringMeta(meta.provider)?.toLowerCase() === provider && accountKeyMatches(meta, accountKey)) return true;
+  if (accountKeyMatches(meta, accountKey)) return true;
+  const host = stringMeta(agent.machine) ?? stringMeta(agent.rig);
+  if (!host) return false;
+  return host === accountHost || sourceHosts.has(host);
+}
+
+function accountKeyMatches(meta: Record<string, unknown>, accountKey: string): boolean {
+  return [stringMeta(meta.accountKey), stringMeta(meta.accountId), stringMeta(meta.orgId)].includes(accountKey);
+}
+
+function agentHosts(agentIds: string[]): Set<string> {
+  const hosts = new Set<string>();
+  if (agentIds.length === 0) return hosts;
+  const rows = getDb().query(`
+    SELECT machine, rig FROM agents
+    WHERE id IN (${agentIds.map(() => "?").join(",")})
+  `).all(...agentIds) as Array<{ machine?: string | null; rig?: string | null }>;
+  for (const row of rows) {
+    const host = stringMeta(row.machine) ?? stringMeta(row.rig);
+    if (host) hosts.add(host);
+  }
+  return hosts;
+}
+
+function crossProviderHeadroomHint(provider: string): string {
+  const rows = getDb().query(`
+    SELECT provider, quota_state
+    FROM provider_quotas
+    WHERE provider <> ? AND quota_state IS NOT NULL
+  `).all(provider) as Array<{ provider: string; quota_state: string }>;
+  const options = rows.flatMap((row) => {
+    const quota = parseJson<QuotaState>(row.quota_state, { windows: [], source: "probe", updatedAt: 0 });
+    return quota.windows.map((window) => ({
+      provider: row.provider,
+      window: window.name,
+      remainingPercent: percent(quotaRemaining(window.utilization)),
+    }));
+  }).sort((a, b) => b.remainingPercent - a.remainingPercent);
+  const best = options[0];
+  if (!best) return "No separate provider headroom is currently recorded.";
+  return `${providerLabel(best.provider)} pool separate, ${windowLabel(best.window)} at ${best.remainingPercent}% remaining - route new work there.`;
+}
+
+function isQuotaReset(
+  remaining: number,
+  previousWindow: QuotaState["windows"][number] | undefined,
+  window: QuotaState["windows"][number],
+  resetRemaining: number,
+): boolean {
+  if (remaining >= resetRemaining) return true;
+  if (!previousWindow?.resetsAt || !window.resetsAt || window.resetsAt <= previousWindow.resetsAt) return false;
+  return remaining - quotaRemaining(previousWindow.utilization) >= 0.5 && remaining >= 0.8;
+}
+
+function quotaRemainingBand(remaining: number, bandSize: number): number {
+  return Math.min(100, Math.max(0, percent(Math.ceil(remaining / bandSize) * bandSize)));
+}
+
+function quotaRemaining(utilization: number): number {
+  return Math.max(0, Math.min(1, 1 - utilization));
+}
+
+function percent(value: number): number {
+  return Math.round(value * 100);
+}
+
+function hostFromAccountKey(accountKey: string): string | undefined {
+  if (!accountKey.startsWith("host:")) return undefined;
+  const host = accountKey.slice("host:".length);
+  return host && !host.includes(":") ? host : undefined;
+}
+
+function stringMeta(value: unknown): string | undefined {
+  return typeof value === "string" && value.trim() ? value.trim() : undefined;
+}
+
+function providerLabel(provider: string): string {
+  return provider.slice(0, 1).toUpperCase() + provider.slice(1);
+}
+
+function windowLabel(value: string): string {
+  if (value === "five_hour") return "5h";
+  if (value === "seven_day") return "7d";
+  return value.replace(/_/g, " ");
+}
+
+function formatDuration(ms: number): string {
+  if (!Number.isFinite(ms) || ms <= 0) return "now";
+  const minutes = Math.max(1, Math.round(ms / 60_000));
+  if (minutes < 60) return `${minutes}m`;
+  const hours = Math.round(minutes / 60);
+  if (hours < 48) return `${hours}h`;
+  return `${Math.round(hours / 24)}d`;
+}

package/src/quota-threshold-config.ts

@@ -0,0 +1,46 @@
+import { isRecord } from "agent-relay-sdk";
+import { ValidationError } from "./db/connection.ts";
+import type { NotificationsConfig } from "./types";
+
+export const DEFAULT_QUOTA_THRESHOLD_CONFIG: NotificationsConfig["quotaThreshold"] = {
+  enabled: true,
+  warnRemaining: 0.25,
+  criticalRemaining: 0.1,
+  resetRemaining: 0.99,
+  bandSize: 0.05,
+};
+
+export function validateQuotaThresholdConfig(value: unknown): NotificationsConfig["quotaThreshold"] {
+  const input = isRecord(value) ? value : {};
+  const warnRemaining = cleanFraction(input.warnRemaining, "quotaThreshold.warnRemaining", DEFAULT_QUOTA_THRESHOLD_CONFIG.warnRemaining);
+  const criticalRemaining = cleanFraction(input.criticalRemaining, "quotaThreshold.criticalRemaining", DEFAULT_QUOTA_THRESHOLD_CONFIG.criticalRemaining);
+  const resetRemaining = cleanFraction(input.resetRemaining, "quotaThreshold.resetRemaining", DEFAULT_QUOTA_THRESHOLD_CONFIG.resetRemaining);
+  const bandSize = cleanFraction(input.bandSize, "quotaThreshold.bandSize", DEFAULT_QUOTA_THRESHOLD_CONFIG.bandSize);
+  if (bandSize <= 0) throw new ValidationError("quotaThreshold.bandSize must be greater than 0");
+  if (criticalRemaining >= warnRemaining) {
+    throw new ValidationError("quotaThreshold.criticalRemaining must be less than quotaThreshold.warnRemaining");
+  }
+  if (resetRemaining <= warnRemaining) {
+    throw new ValidationError("quotaThreshold.resetRemaining must be greater than quotaThreshold.warnRemaining");
+  }
+  return {
+    enabled: input.enabled === undefined ? DEFAULT_QUOTA_THRESHOLD_CONFIG.enabled : cleanBoolean(input.enabled, "quotaThreshold.enabled"),
+    warnRemaining,
+    criticalRemaining,
+    resetRemaining,
+    bandSize,
+  };
+}
+
+function cleanBoolean(value: unknown, field: string): boolean {
+  if (typeof value !== "boolean") throw new ValidationError(`${field} must be a boolean`);
+  return value;
+}
+
+function cleanFraction(value: unknown, field: string, fallback: number): number {
+  if (value === undefined || value === null) return fallback;
+  if (typeof value !== "number" || !Number.isFinite(value) || value < 0 || value > 1) {
+    throw new ValidationError(`${field} must be a number between 0 and 1`);
+  }
+  return value;
+}

package/src/routes/index.ts

@@ -5,7 +5,7 @@ import { deleteAgentTerminalSession, postAgentAction, postAgentPermissionDecisio
 import { deleteArtifactById, getArtifactById, getArtifactContent, getArtifacts, headArtifactContent, postArtifact } from "./artifacts";
 import { deleteAutomationById, getAutomationById, getAutomationRuns, getAutomations, patchAutomation, postAutomation, postAutomationRun } from "./automations";
 import { deleteCommandById, deleteProviderModelRoute, getProviderConfigRoute, getProviderConfigsRoute, getProvidersRoute, patchProviderModelRoute, postProviderConfigTestRoute, putProviderConfigRoute, putProviderModelRoute } from "./provider-config";
-import { getProviderQuotasRoute, postProviderQuotaRoute } from "./provider-quotas";
+import { getProviderQuotasRoute, postProviderQuotaLeaseAcquireRoute, postProviderQuotaLeaseReleaseRoute, postProviderQuotaRoute } from "./provider-quotas";
 import { deleteConfigKey, getConfigKey, getConfigKeyHistory, getConfigNamespace, putConfigKey } from "./config";
 import { deleteInboxDraftRoute, getChatHistoryImportsRoute, getInboxStateRoute, patchInboxThreadState, postChatHistoryImportRoute, putInboxDraft } from "./inbox";
 import { deleteMemoryRoute, getMemories, getMemoryBrokerInfo, getMemoryById, getMemoryStats, patchMemory, postMemory, postMemoryInject, postMemorySearch } from "./memory";
@@ -90,6 +90,8 @@ const routes: Route[] = [
   route("POST", "/api/route", postRouteAdvice),
   route("GET", "/api/provider-quotas", getProviderQuotasRoute),
   route("POST", "/api/provider-quotas", postProviderQuotaRoute),
+  route("POST", "/api/orchestrators/:id/provider-quota-leases/acquire", postProviderQuotaLeaseAcquireRoute),
+  route("POST", "/api/orchestrators/:id/provider-quota-leases/release", postProviderQuotaLeaseReleaseRoute),
   route("GET", "/api/agents/spawn/directories", getHostDirectories),
   route("GET", "/api/agents/:id/chat", getAgentChatHistory),
   route("GET", "/api/agents/:id/context-snapshots", getAgentContextSnapshots),

package/src/routes/provider-quotas.ts

@@ -1,8 +1,8 @@
 import { isRecord } from "agent-relay-sdk";
-import { listProviderQuotas, upsertProviderQuota } from "../db";
+import { acquireProviderQuotaLease, getOrchestrator, listProviderQuotas, releaseProviderQuotaLease, upsertProviderQuota } from "../db";
 import { quotaStateFromUnknown, providerQuotaErrorFromUnknown } from "../quota";
 import { cleanString } from "../validation";
-import { error, json, parseBody, parseQueryInt, type Handler } from "./_shared";
+import { authorizeRoute, error, json, parseBody, parseQueryInt, type Handler } from "./_shared";
 import type { ProviderQuotaUpdateInput } from "../types";
 
 function cleanProviderQuotaUpdate(body: unknown): ProviderQuotaUpdateInput | string {
@@ -46,3 +46,57 @@ export const postProviderQuotaRoute: Handler = async (req) => {
   if (typeof input === "string") return error(input);
   return json(upsertProviderQuota(input), 201);
 };
+
+function cleanLeaseBody(body: unknown): {
+  provider: string;
+  accountKey: string;
+  leaseToken?: string;
+  ttlMs: number;
+} | string {
+  if (!isRecord(body)) return "JSON object body required";
+  const provider = cleanString(body.provider, "provider", { required: true, max: 80 })!;
+  const accountKey = cleanString(body.accountKey, "accountKey", { required: true, max: 240 })!;
+  const leaseToken = cleanString(body.leaseToken, "leaseToken", { max: 120 });
+  const ttlMs = typeof body.ttlMs === "number" && Number.isFinite(body.ttlMs)
+    ? Math.max(5_000, Math.min(10 * 60_000, Math.round(body.ttlMs)))
+    : 90_000;
+  return { provider, accountKey, ...(leaseToken ? { leaseToken } : {}), ttlMs };
+}
+
+export const postProviderQuotaLeaseAcquireRoute: Handler = async (req, params) => {
+  const orch = getOrchestrator(params.id!);
+  if (!orch) return error("orchestrator not found", 404);
+  const denied = authorizeRoute(req, { scope: "command:write", resource: { orchestratorId: orch.id } });
+  if (denied) return denied;
+  const parsed = await parseBody<unknown>(req);
+  if (!parsed.ok) return error(parsed.error, parsed.status);
+  const input = cleanLeaseBody(parsed.body);
+  if (typeof input === "string") return error(input);
+  return json(acquireProviderQuotaLease({
+    provider: input.provider,
+    accountKey: input.accountKey,
+    orchestratorId: orch.id,
+    leaseToken: input.leaseToken,
+    ttlMs: input.ttlMs,
+  }));
+};
+
+export const postProviderQuotaLeaseReleaseRoute: Handler = async (req, params) => {
+  const orch = getOrchestrator(params.id!);
+  if (!orch) return error("orchestrator not found", 404);
+  const denied = authorizeRoute(req, { scope: "command:write", resource: { orchestratorId: orch.id } });
+  if (denied) return denied;
+  const parsed = await parseBody<unknown>(req);
+  if (!parsed.ok) return error(parsed.error, parsed.status);
+  const input = cleanLeaseBody(parsed.body);
+  if (typeof input === "string") return error(input);
+  if (!input.leaseToken) return error("leaseToken required");
+  return json({
+    released: releaseProviderQuotaLease({
+      provider: input.provider,
+      accountKey: input.accountKey,
+      orchestratorId: orch.id,
+      leaseToken: input.leaseToken,
+    }),
+  });
+};

package/src/token-db.ts

@@ -67,7 +67,7 @@ const BUILT_IN_PROFILES: Array<Omit<TokenProfile, "createdAt" | "updatedAt">> =
     name: "Orchestrator Host",
     description: "Host supervisor access for spawn, command, log, and terminal operations.",
     role: "orchestrator",
-    scope: ["agent:read", "agent:write", "command:read", "command:write", "task:read", "task:write", "logs:read", "terminal:attach"],
+    scope: ["agent:read", "agent:write", "command:read", "command:write", "task:read", "task:write", "logs:read", "terminal:attach", "quota:write"],
     constraints: { logsRead: true, terminalAttach: true },
     ttlSeconds: 30 * 24 * 60 * 60,
     builtIn: true,