infraweaver 0.1.0

package/src/utils/patchWorkflowRunFields.ts

@@ -0,0 +1,173 @@
+import type { AgentUsage } from "#app/agents/shared";
+import type { CloudToolContext } from "#app/mcp/capabilityContext";
+import { log } from "#app/utils/cli";
+import { cloudRequest } from "#app/utils/cloudClient";
+import { parseRetryAfter, retry } from "#app/utils/retry";
+
+/**
+ * Artifact tracking fields — one-off PATCHes from MCP tools as GitHub entities
+ * are created during the run. Strings only (GraphQL node IDs).
+ * Keep in sync with `STRING_FIELDS` in `app/api/workflow-run/[runId]/route.ts`.
+ */
+export type WorkflowRunArtifactPatchKey =
+  | "prNodeId"
+  | "issueNodeId"
+  | "reviewNodeId"
+  | "planCommentNodeId"
+  | "summarySnapshot";
+
+/**
+ * Usage fields — aggregated across all agent calls and PATCHed once at
+ * end-of-run. Token counts are Int4 on the DB side (ample for any realistic
+ * run); `costUsd` is a Decimal populated by provider-reported dollar amounts.
+ * Keep in sync with `INT_FIELDS` + `DECIMAL_FIELDS` in the server route.
+ */
+export type WorkflowRunUsagePatchKey =
+  | "inputTokens"
+  | "outputTokens"
+  | "cacheReadTokens"
+  | "cacheWriteTokens"
+  | "costUsd";
+
+export type WorkflowRunPatch = Partial<
+  Record<WorkflowRunArtifactPatchKey, string>
+> &
+  Partial<Record<WorkflowRunUsagePatchKey, number>>;
+
+const STRING_KEYS: WorkflowRunArtifactPatchKey[] = [
+  "prNodeId",
+  "issueNodeId",
+  "reviewNodeId",
+  "planCommentNodeId",
+  "summarySnapshot",
+];
+
+const NUMBER_KEYS: WorkflowRunUsagePatchKey[] = [
+  "inputTokens",
+  "outputTokens",
+  "cacheReadTokens",
+  "cacheWriteTokens",
+  "costUsd",
+];
+
+/** PATCH workflow-run fields (Infraweaver JWT, not GitHub). */
+export async function patchWorkflowRunFields(
+  ctx: CloudToolContext,
+  fields: WorkflowRunPatch,
+): Promise<void> {
+  if (ctx.runId === undefined || !ctx.apiToken) return;
+  const body: Record<string, string | number> = {};
+  for (const key of STRING_KEYS) {
+    const value = fields[key];
+    if (typeof value === "string" && value.length > 0) {
+      body[key] = value;
+    }
+  }
+  for (const key of NUMBER_KEYS) {
+    const value = fields[key];
+    if (typeof value === "number" && Number.isFinite(value) && value >= 0) {
+      body[key] = value;
+    }
+  }
+  if (Object.keys(body).length === 0) return;
+  try {
+    await retry(
+      async () => {
+        const response = await cloudRequest({
+          path: `/api/workflow-run/${ctx.runId}`,
+          method: "PATCH",
+          credential: ctx.apiToken,
+          json: body,
+          timeoutMs: 10_000,
+        });
+        // Capture Retry-After before draining, so a 429/503 floors the backoff to the
+        // server's hint instead of the local jitter.
+        const retryAfterMs = response.ok
+          ? undefined
+          : parseRetryAfter(response.headers.get("retry-after"));
+        // drain the (unused) body on BOTH branches so undici/keep-alive returns the
+        // socket to the pool promptly instead of pinning it until GC.
+        await response.body?.cancel();
+        if (!response.ok) {
+          const err: Error & { status?: number; retryAfterMs?: number } =
+            new Error(`PATCH workflow-run: ${response.status}`);
+          // Carry the status so retry() matches it by code (not just the message regex),
+          // and the server's Retry-After so it can honour it.
+          err.status = response.status;
+          if (retryAfterMs !== undefined) err.retryAfterMs = retryAfterMs;
+          throw err;
+        }
+      },
+      {
+        maxAttempts: 3,
+        delayMs: 2000,
+        label: "patchWorkflowRunFields",
+      },
+    );
+  } catch (error) {
+    log.warning(`patchWorkflowRunFields exhausted retries: ${error}`);
+  }
+}
+
+/**
+ * Postgres INTEGER / Prisma Int4 is signed 32-bit. Aggregated usage won't
+ * realistically hit this in a single run (2.1B tokens ≈ $6000+ of input on
+ * Claude Opus), but clamping here keeps the wire payload self-consistent:
+ * the server rejects out-of-range INT fields individually, so without a
+ * client-side clamp a single overflow would write a partial row where
+ * some columns land and others silently don't.
+ */
+const INT4_MAX = 2_147_483_647;
+
+function clampInt(value: number, field: WorkflowRunUsagePatchKey): number {
+  if (value > INT4_MAX) {
+    log.warning(
+      `aggregateUsage: ${field}=${value} exceeds INT4_MAX (${INT4_MAX}) — clamping so the rest of the usage row still persists.`,
+    );
+    return INT4_MAX;
+  }
+  return value;
+}
+
+/**
+ * Sum per-agent usage entries into a single WorkflowRunPatch payload.
+ * Returns an empty object when there's nothing to report, which causes
+ * `patchWorkflowRunFields` to no-op — safe to call unconditionally from
+ * end-of-run paths. Zero-valued fields are dropped so the DB only stores
+ * positive sums (and NULL means "not reported").
+ *
+ * Token sums are clamped to INT4_MAX to guarantee the payload the server
+ * sees is always self-consistent across all numeric columns.
+ */
+export function aggregateUsage(entries: AgentUsage[]): WorkflowRunPatch {
+  if (entries.length === 0) return {};
+
+  const sum = entries.reduce(
+    (acc, e) => ({
+      inputTokens: acc.inputTokens + e.inputTokens,
+      outputTokens: acc.outputTokens + e.outputTokens,
+      cacheReadTokens: acc.cacheReadTokens + (e.cacheReadTokens ?? 0),
+      cacheWriteTokens: acc.cacheWriteTokens + (e.cacheWriteTokens ?? 0),
+      costUsd: acc.costUsd + (e.costUsd ?? 0),
+    }),
+    {
+      inputTokens: 0,
+      outputTokens: 0,
+      cacheReadTokens: 0,
+      cacheWriteTokens: 0,
+      costUsd: 0,
+    },
+  );
+
+  const out: WorkflowRunPatch = {};
+  if (sum.inputTokens > 0)
+    out.inputTokens = clampInt(sum.inputTokens, "inputTokens");
+  if (sum.outputTokens > 0)
+    out.outputTokens = clampInt(sum.outputTokens, "outputTokens");
+  if (sum.cacheReadTokens > 0)
+    out.cacheReadTokens = clampInt(sum.cacheReadTokens, "cacheReadTokens");
+  if (sum.cacheWriteTokens > 0)
+    out.cacheWriteTokens = clampInt(sum.cacheWriteTokens, "cacheWriteTokens");
+  if (sum.costUsd > 0) out.costUsd = sum.costUsd;
+  return out;
+}