@nifrajs/agent 4.0.2 → 5.0.0

package/src/gateway.ts

@@ -27,7 +27,10 @@ const ERROR_CODES: ReadonlySet<string> = new Set(MODEL_GATEWAY_ERROR_CODES)
 const TOKEN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/
 const MAX_ATTEMPTS = 512
 const MAX_TOKENS = 10_000_000
-const MAX_EVIDENCE = 512
+// An attempt records at most four entries plus one fallback; the run adds one terminal entry.
+const MAX_EVIDENCE = MAX_ATTEMPTS * 5 + 1
+// setTimeout fires at once for any delay past 2^31 - 1 ms (about 24.8 days).
+const MAX_TIMER_MS = 2_147_483_647
 
 export interface ModelRoute {
   /** Caller-owned opaque route identifier. It is not a provider or model name. */
@@ -341,44 +344,79 @@ function timedFailure(
   return undefined
 }
 
-async function settleAttempt(
-  gateway: ModelGateway,
-  request: ModelGatewayAttemptRequest,
+type Deadlined<T> =
+  | { readonly settled: true; readonly value: T }
+  | { readonly settled: false; readonly error: ModelGatewayError }
+
+/** Settle `work`, unless the deadline passes or `signal` aborts first. A rejection passes through. */
+async function withinDeadline<T>(
+  work: Promise<T>,
+  signal: AbortSignal,
   deadlineAt: number | undefined,
   now: () => number,
-): Promise<ModelGatewayRawResult> {
-  const before = timedFailure(request.signal, deadlineAt, now)
-  if (before !== undefined) return { ok: false, error: before }
+): Promise<Deadlined<T>> {
+  if (signal.aborted) return { settled: false, error: { code: "cancelled" } }
   const remaining = deadlineAt === undefined ? undefined : Math.max(0, deadlineAt - now())
   let timer: ReturnType<typeof setTimeout> | undefined
   let abort: (() => void) | undefined
   try {
-    const operation = Promise.resolve()
-      .then(() => gateway.complete(request))
-      .then(parseModelGatewayResult, () => ({
-        ok: false as const,
-        error: { code: "internal" as const },
-      }))
-    const races: Promise<ModelGatewayRawResult>[] = [operation]
+    const races: Promise<Deadlined<T>>[] = [
+      work.then((value) => ({ settled: true as const, value })),
+      new Promise((resolve) => {
+        abort = () => resolve({ settled: false, error: { code: "cancelled" } })
+        signal.addEventListener("abort", abort, { once: true })
+      }),
+    ]
     if (remaining !== undefined) {
       races.push(
         new Promise((resolve) => {
-          timer = setTimeout(() => resolve({ ok: false, error: { code: "timeout" } }), remaining)
-        }),
-      )
-    }
-    if (!request.signal.aborted) {
-      races.push(
-        new Promise((resolve) => {
-          abort = () => resolve({ ok: false, error: { code: "cancelled" } })
-          request.signal.addEventListener("abort", abort!, { once: true })
+          let left = remaining
+          const arm = (): void => {
+            const delay = Math.min(left, MAX_TIMER_MS)
+            left -= delay
+            timer = setTimeout(
+              () => (left > 0 ? arm() : resolve({ settled: false, error: { code: "timeout" } })),
+              delay,
+            )
+          }
+          arm()
         }),
       )
     }
     return await Promise.race(races)
   } finally {
     if (timer !== undefined) clearTimeout(timer)
-    if (abort !== undefined) request.signal.removeEventListener("abort", abort)
+    if (abort !== undefined) signal.removeEventListener("abort", abort)
+  }
+}
+
+async function settleAttempt(
+  gateway: ModelGateway,
+  request: Omit<ModelGatewayAttemptRequest, "signal">,
+  signal: AbortSignal,
+  deadlineAt: number | undefined,
+  now: () => number,
+): Promise<ModelGatewayRawResult> {
+  const before = timedFailure(signal, deadlineAt, now)
+  if (before !== undefined) return { ok: false, error: before }
+  // Each attempt has its own signal, aborted when the attempt stops being waited for, so an adapter
+  // that honours it stops working on a request whose result would be discarded.
+  const attempt = new AbortController()
+  const cancel = (): void => attempt.abort()
+  signal.addEventListener("abort", cancel, { once: true })
+  try {
+    const operation = Promise.resolve()
+      .then(() => gateway.complete({ ...request, signal: attempt.signal }))
+      .then(parseModelGatewayResult, () => ({
+        ok: false as const,
+        error: { code: "internal" as const },
+      }))
+    const outcome = await withinDeadline(operation, signal, deadlineAt, now)
+    if (outcome.settled) return outcome.value
+    attempt.abort()
+    return { ok: false, error: outcome.error }
+  } finally {
+    signal.removeEventListener("abort", cancel)
   }
 }
 
@@ -431,7 +469,8 @@ export async function runModelGateway<Input, Output>(
     })
     const raw = await settleAttempt(
       gateway,
-      { routeId: route, input: request.input, signal, envelope },
+      { routeId: route, input: request.input, envelope },
+      signal,
       policy.deadlineAt,
       checked.now,
     )
@@ -461,8 +500,29 @@ export async function runModelGateway<Input, Output>(
         outcome: "passed",
       })
       if (request.parser !== undefined) {
+        const parser = request.parser
+        // The deadline covers parsing too: a success is only returned while it is still in time.
+        const stopped = (error: ModelGatewayError) => {
+          addEvidence(evidence, {
+            kind: "attempt",
+            routeId: route,
+            attempt: attempts,
+            outcome: "failed",
+            code: error.code,
+          })
+          return terminal(error, route, attempts, usage, evidence)
+        }
         try {
-          const output = await request.parser.parse(parsed.output)
+          const parsing = await withinDeadline(
+            Promise.resolve().then(() => parser.parse(parsed.output)),
+            signal,
+            policy.deadlineAt,
+            checked.now,
+          )
+          if (!parsing.settled) return stopped(parsing.error)
+          const late = timedFailure(signal, policy.deadlineAt, checked.now)
+          if (late !== undefined) return stopped(late)
+          const output = parsing.value
           addEvidence(evidence, {
             kind: "parse",
             routeId: route,

package/src/index.ts

@@ -275,6 +275,8 @@ export interface AgentApprovalPort {
     readonly signal: AbortSignal
     /** Wall-clock budget reserved for this approval hop and the following tool execution. */
     readonly deadline?: RequestBudget
+    /** The decision the resuming caller sent, under `resumeApproval: "port"`. Evidence, not a verdict. */
+    readonly callerDecision?: ToolApproval
   }): AgentApprovalResult | PromiseLike<AgentApprovalResult>
 }
 
@@ -454,6 +456,13 @@ export interface AgentPorts {
   readonly idempotency?: ToolIdempotencyStore
   readonly state?: AgentStateStore
   readonly approval?: AgentApprovalPort
+  /**
+   * Who decides a step that suspended for approval when the run resumes. `"caller"` (the default)
+   * takes `resume.approval` as the decision. `"port"` asks `approval` again and hands it the caller's
+   * decision as `callerDecision`; set it when the resuming caller is the party asking for the tool,
+   * as under A2A or a public AG-UI route, so it cannot approve its own call.
+   */
+  readonly resumeApproval?: "caller" | "port"
   readonly telemetry?: AgentTelemetryPort
   /** Receives model deltas when the model port streams. Transient - nothing here is persisted. */
   readonly deltas?: AgentDeltaSink
@@ -621,7 +630,19 @@ export async function turn<
   if (signal.aborted) return suspended(state, { kind: "cancelled", effectId: newId() }, "cancelled")
 
   if (state.status === "suspended") {
-    if (input.resume !== undefined) {
+    if (input.resume !== undefined && state.pending?.tool === undefined) {
+      // A model, budget or cancelled suspension has nothing to replay: echoing its continuation back
+      // retries the model step, as omitting `resume` does. Any other continuation is refused.
+      const pending = state.pending
+      const continuation = input.resume.continuation
+      if (
+        pending === undefined ||
+        continuation.tool !== undefined ||
+        continuation.kind !== pending.kind ||
+        continuation.effectId !== pending.effectId
+      )
+        throw new AgentResumeMismatchError()
+    } else if (input.resume !== undefined) {
       return resumeTool(
         definition,
         state,
@@ -997,13 +1018,22 @@ async function resumeTool<
     )
   // The caller's decision answers only a step that suspended to await one; a step suspended for budget
   // or cancellation asks the approval port as it would have.
+  const callerDecision = resume.approval
+  const port = ports.approval
   const approvalPorts: AgentPorts =
-    resume.approval === undefined || pending.kind !== "approval"
+    callerDecision === undefined || pending.kind !== "approval"
       ? ports
-      : {
-          ...ports,
-          approval: { request: () => ({ status: "approved", approval: resume.approval! }) },
-        }
+      : ports.resumeApproval !== "port"
+        ? {
+            ...ports,
+            approval: { request: () => ({ status: "approved", approval: callerDecision }) },
+          }
+        : port === undefined
+          ? ports
+          : {
+              ...ports,
+              approval: { request: (input) => port.request({ ...input, callerDecision }) },
+            }
   return executeToolStep(
     { ...state, status: "ready" },
     tool,

package/src/mount.ts

@@ -14,10 +14,13 @@
  *   - otherwise -> a single JSON body: the projected run result.
  * - `resume` is `{ continuation, approval? }` (see {@link AgentTurnInput}); the runtime keeps state
  *   token-only, so a suspended tool's input must be replayed by the caller in `continuation.input`.
+ *   Return `resumeApproval: "port"` from `ports` when the caller must not approve its own tool calls.
  * - With an `evidenceLog` configured, `step` frames carry `id: <seq>` and a dropped SSE connection
  *   is resumable: re-POST the same `turnId` with a `Last-Event-ID` header to replay the missed
  *   evidence and rejoin the still-running turn (or receive the stored terminal frame) - the run is
  *   never re-executed. Without the log the header is ignored and every POST starts a run.
+ * - One SSE response emits at most `maxOutputBytes` (default 4 MiB), so a client that stops reading
+ *   cannot make the server queue a whole run; past it the response ends with `event: error`.
  * - The seam performs no authentication or authorization. Wrap it with the app's own route guards, and
  *   scope the store/model returned by `ports` to the caller (RLS born in at the data layer). With an
  *   `evidenceLog`, set `evidenceOwner` too: without it a turn id is a bearer capability, and any
@@ -53,6 +56,8 @@ import {
 
 const DEFAULT_PATH = "/agent"
 const DEFAULT_MAX_BODY_BYTES = 1_000_000
+const DEFAULT_MAX_OUTPUT_BYTES = 4 * 1024 * 1024
+const MAX_OUTPUT_BYTES = 64 * 1024 * 1024
 const PENDING_KINDS: readonly AgentPendingKind[] = ["approval", "budget", "model", "cancelled"]
 const LAST_EVENT_ID_PATTERN = /^\d{1,15}$/
 const TURN_ID_PATTERN = /^[a-zA-Z0-9._:-]{1,128}$/
@@ -80,6 +85,12 @@ export interface MountAgentOptions<
   readonly maxBodyBytes?: number
   /** Prototype-poisoning policy for the framing lane. Default `"reject"`. */
   readonly protoPoisoning?: ProtoPoisoning
+  /**
+   * Maximum total UTF-8 bytes one SSE response emits. Default 4 MiB. Past it the response ends with an
+   * `error` frame (`output_limit`); with an `evidenceLog` the run keeps recording and a reconnect with
+   * `Last-Event-ID` reads on from where this response stopped, otherwise the run is cancelled.
+   */
+  readonly maxOutputBytes?: number
   /**
    * Build the ports for one request - the model, durable state store, approval transport, capabilities,
    * and budgets. Receives the route context so the caller can scope every port to the request subject.
@@ -107,13 +118,20 @@ export function mountAgent<
   const maxBytes = options.maxBodyBytes ?? DEFAULT_MAX_BODY_BYTES
   const proto = options.protoPoisoning ?? "reject"
   const runOptions = options.maxTurns === undefined ? {} : { maxTurns: options.maxTurns }
+  const maxOutputBytes = options.maxOutputBytes ?? DEFAULT_MAX_OUTPUT_BYTES
+  if (
+    !Number.isSafeInteger(maxOutputBytes) ||
+    maxOutputBytes < 1024 ||
+    maxOutputBytes > MAX_OUTPUT_BYTES
+  )
+    throw new RangeError(`mountAgent maxOutputBytes must be between 1024 and ${MAX_OUTPUT_BYTES}`)
 
   app.post(path, (c) =>
     readBodyFramed<Response>(
       c.req,
       maxBytes,
       proto,
-      (parsed) => execute(c, parsed, options, runOptions),
+      (parsed) => execute(c, parsed, options, runOptions, maxOutputBytes),
       (rejection) => render(rejection),
       () => render(plainError(400, "bad_request")),
     ).catch(() => render(plainError(400, "bad_request"))),
@@ -125,6 +143,7 @@ async function execute<InputSchema extends StandardSchemaV1, OutputSchema extend
   parsed: unknown,
   options: MountAgentOptions<InputSchema, OutputSchema>,
   runOptions: { readonly maxTurns?: number },
+  maxOutputBytes: number,
 ): Promise<Response> {
   const body = asRecord(parsed)
   if (body === undefined) return jsonResponse(400, { error: "body_must_be_object" })
@@ -148,7 +167,7 @@ async function execute<InputSchema extends StandardSchemaV1, OutputSchema extend
       return jsonResponse(400, { error: "invalid_last_event_id" })
     const replay = await log.replay(turnId, Number(lastEventId))
     if (replay === undefined) return jsonResponse(409, { error: "replay_unavailable" })
-    return sseReplayResponse(turnId, replay)
+    return sseReplayResponse(turnId, replay, maxOutputBytes)
   }
 
   const basePorts = await options.ports(c)
@@ -190,7 +209,7 @@ async function execute<InputSchema extends StandardSchemaV1, OutputSchema extend
       return jsonResponse(409, { error: "resume_unavailable" })
     }
   }
-  return sseResponse(turnId, start, log)
+  return sseResponse(turnId, start, log, maxOutputBytes)
 }
 
 function sseResponse<Output>(
@@ -200,6 +219,7 @@ function sseResponse<Output>(
     disconnect: AbortSignal | undefined,
   ) => Promise<AgentRunResult<Output>>,
   log: AgentEvidenceLog | undefined,
+  maxOutputBytes: number,
 ): Response {
   const stream = createAgentEvidenceStream()
   // `id:` frames are only meaningful when a log can serve the reconnect they invite.
@@ -210,9 +230,18 @@ function sseResponse<Output>(
   let detached = false
   const body = new ReadableStream<Uint8Array>({
     async start(controller) {
-      const deliver = sseSender(controller)
+      const deliver = sseSender(controller, maxOutputBytes)
       const send: typeof deliver = (event, data, id) => {
-        if (!detached) deliver(event, data, id)
+        if (detached) return
+        try {
+          deliver(event, data, id)
+        } catch (error) {
+          if (!(error instanceof SseOutputLimitError)) throw error
+          // The budget belongs to this response, not to the logged run's authoritative result.
+          detached = true
+          disconnect?.abort()
+          endOverBudget(deliver, controller, turnId)
+        }
       }
       const run = start(stream, disconnect?.signal)
       // The evidence stream must terminate whether the run resolves or throws, or the `for await`
@@ -254,10 +283,26 @@ function withDisconnect(
   return AbortSignal.any([signal, disconnect])
 }
 
-function sseReplayResponse(turnId: string, replay: AgentEvidenceReplay): Response {
+function sseReplayResponse(
+  turnId: string,
+  replay: AgentEvidenceReplay,
+  maxOutputBytes: number,
+): Response {
+  let detached = false
   const body = new ReadableStream<Uint8Array>({
     async start(controller) {
-      const send = sseSender(controller)
+      const deliver = sseSender(controller, maxOutputBytes)
+      const send: typeof deliver = (event, data, id) => {
+        if (detached) return
+        try {
+          deliver(event, data, id)
+        } catch (error) {
+          if (!(error instanceof SseOutputLimitError)) throw error
+          detached = true
+          void replay.live?.return?.()
+          endOverBudget(deliver, controller, turnId)
+        }
+      }
       try {
         for (const evidence of replay.evidence) send("step", evidence, evidence.seq)
         if (replay.live !== undefined)
@@ -268,21 +313,49 @@ function sseReplayResponse(turnId: string, replay: AgentEvidenceReplay): Respons
       } catch {
         send("error", { turnId, error: "run_failed" })
       } finally {
-        controller.close()
+        if (!detached) controller.close()
       }
     },
+    cancel() {
+      detached = true
+      // Only this view detaches; the logged run keeps recording for the next reconnect.
+      void replay.live?.return?.()
+    },
   })
   return sseHeaders(body)
 }
 
+type SseSend = (event: string, data: unknown, id?: number) => void
+
+class SseOutputLimitError extends Error {}
+
 function sseSender(
   controller: ReadableStreamDefaultController<Uint8Array>,
-): (event: string, data: unknown, id?: number) => void {
+  maxOutputBytes: number,
+): SseSend {
   const encoder = new TextEncoder()
+  let emittedBytes = 0
   return (event, data, id) => {
     const head = id === undefined ? "" : `id: ${id}\n`
-    controller.enqueue(encoder.encode(`${head}event: ${event}\ndata: ${JSON.stringify(data)}\n\n`))
+    const frame = encoder.encode(`${head}event: ${event}\ndata: ${JSON.stringify(data)}\n\n`)
+    if (frame.byteLength > maxOutputBytes || emittedBytes > maxOutputBytes - frame.byteLength)
+      throw new SseOutputLimitError("mountAgent SSE output limit exceeded")
+    emittedBytes += frame.byteLength
+    controller.enqueue(frame)
+  }
+}
+
+function endOverBudget(
+  deliver: SseSend,
+  controller: ReadableStreamDefaultController<Uint8Array>,
+  turnId: string,
+): void {
+  try {
+    deliver("error", { turnId, error: "output_limit" })
+  } catch {
+    // A fully spent budget cannot carry the error frame; closing still ends this response.
   }
+  controller.close()
 }
 
 function sseHeaders(body: ReadableStream<Uint8Array>): Response {