@nifrajs/coding-agent 3.3.0 → 3.5.0

package/src/rpc.ts

@@ -11,7 +11,7 @@ import { CodingAgentHost, type CodingAgentHostOptions } from "./host.ts"
 import { readBoundedText } from "./process.ts"
 import type { ContextWindowOptions, SessionStore } from "./sessions.ts"
 import type { UiExtensionHost, UiExtensionManifest } from "./ui.ts"
-import { runNifraVerification } from "./verification.ts"
+import { isSafeReviewDiff, runNifraReview, runNifraVerification } from "./verification.ts"
 
 export interface CodingAgentRpcServerOptions {
   readonly backend: AgentBackend
@@ -25,6 +25,8 @@ export interface CodingAgentRpcServerOptions {
   /** Include bounded exception stacks in RPC failures only for trusted loopback debugging. */
   readonly exposeErrorStacks?: boolean
   readonly maxBodyBytes?: number
+  /** Maximum UTF-8 bytes in one non-streaming JSON response. Defaults to 4 MiB. */
+  readonly maxResponseBytes?: number
   readonly sessionStore?: SessionStore
   readonly contextWindow?: ContextWindowOptions
   readonly extensions?: ExtensionHost
@@ -50,6 +52,15 @@ interface RpcServerLike {
   stop(closeActiveConnections?: boolean): void
 }
 
+const MAX_RPC_BODY_BYTES = 64 * 1024 * 1024
+const DEFAULT_MAX_RPC_RESPONSE_BYTES = 4 * 1024 * 1024
+const MAX_RPC_RESPONSE_BYTES = 64 * 1024 * 1024
+const SESSION_ID_PATTERN = /^[A-Za-z0-9._:-]{1,128}$/
+const MIN_REMOTE_TOKEN_BYTES = 32
+const AUTH_FAILURE_WINDOW_MS = 60_000
+const AUTH_FAILURE_THRESHOLD = 3
+const AUTH_MAX_BACKOFF_MS = 2_000
+
 /**
  * Minimal loopback RPC surface for the CLI, Workbench, CI clients, and a future mobile companion.
  * Turn output is SSE so clients do not need a WebSocket dependency; every event is still a versioned
@@ -60,7 +71,11 @@ export class CodingAgentRpcServer {
   private readonly options: CodingAgentRpcServerOptions
   private readonly token: string
   private readonly maxBodyBytes: number
+  private readonly maxResponseBytes: number
   private listener: RpcServerLike | undefined
+  private authFailures = 0
+  private authLastFailureAt = 0
+  private authBackoffUntil = 0
 
   constructor(options: CodingAgentRpcServerOptions) {
     this.options = Object.freeze({ ...options })
@@ -78,12 +93,25 @@ export class CodingAgentRpcServer {
       exposeErrorStacks: options.exposeErrorStacks === true,
       ...(options.verification === undefined ? {} : { verification: options.verification }),
     })
-    this.token = options.authToken ?? crypto.randomUUID().replaceAll("-", "")
+    this.token = options.authToken ?? randomAuthToken()
     if (!/^[A-Za-z0-9._~-]{16,256}$/.test(this.token))
       throw new TypeError("agent rpc: authToken must be a bounded token")
     this.maxBodyBytes = options.maxBodyBytes ?? 1_048_576
-    if (!Number.isSafeInteger(this.maxBodyBytes) || this.maxBodyBytes < 1024)
-      throw new RangeError("agent rpc: maxBodyBytes must be at least 1024")
+    if (
+      !Number.isSafeInteger(this.maxBodyBytes) ||
+      this.maxBodyBytes < 1024 ||
+      this.maxBodyBytes > MAX_RPC_BODY_BYTES
+    )
+      throw new RangeError(`agent rpc: maxBodyBytes must be between 1024 and ${MAX_RPC_BODY_BYTES}`)
+    this.maxResponseBytes = options.maxResponseBytes ?? DEFAULT_MAX_RPC_RESPONSE_BYTES
+    if (
+      !Number.isSafeInteger(this.maxResponseBytes) ||
+      this.maxResponseBytes < 1024 ||
+      this.maxResponseBytes > MAX_RPC_RESPONSE_BYTES
+    )
+      throw new RangeError(
+        `agent rpc: maxResponseBytes must be between 1024 and ${MAX_RPC_RESPONSE_BYTES}`,
+      )
   }
 
   async start(): Promise<CodingAgentRpcServerHandle> {
@@ -91,6 +119,14 @@ export class CodingAgentRpcServer {
     const hostname = this.options.hostname ?? "127.0.0.1"
     if (!this.options.allowRemote && !isLoopbackHost(hostname))
       throw new Error("agent rpc: remote binding requires allowRemote: true")
+    if (
+      this.options.authToken !== undefined &&
+      !isLoopbackHost(hostname) &&
+      new TextEncoder().encode(this.token).byteLength < MIN_REMOTE_TOKEN_BYTES
+    )
+      throw new Error(
+        `agent rpc: explicit remote authToken must be at least ${MIN_REMOTE_TOKEN_BYTES} bytes`,
+      )
     if (this.options.exposeErrorStacks === true && this.options.allowRemote === true)
       throw new Error("agent rpc: exposeErrorStacks is only allowed for local-only binding")
     this.listener = Bun.serve({
@@ -111,8 +147,13 @@ export class CodingAgentRpcServer {
   async stop(): Promise<void> {
     const listener = this.listener
     this.listener = undefined
-    await this.host.stop("rpc server stopped")
-    listener?.stop(true)
+    try {
+      await this.host.stop("rpc server stopped")
+    } finally {
+      // The audit write is best effort from a transport-lifecycle perspective. Even when the host
+      // reports a persistence failure, the listening socket must not remain reachable or leaked.
+      listener?.stop(true)
+    }
   }
 
   private async fetch(request: Request): Promise<Response> {
@@ -120,27 +161,53 @@ export class CodingAgentRpcServer {
     if (request.method === "OPTIONS") return new Response(null, { status: 204, headers: cors })
     const url = new URL(request.url)
     if (url.pathname === "/health" && request.method === "GET")
-      return json({ ok: true, protocol: 1 }, 200, cors)
-    if (!authorized(request, this.token))
-      return json(
+      return this.response({ ok: true, protocol: 1 }, 200, cors)
+    const backoffMs = this.authBackoffMs()
+    if (backoffMs > 0) {
+      const throttledCors = new Headers(cors)
+      throttledCors.set("retry-after", String(Math.ceil(backoffMs / 1000)))
+      return this.response(
+        { error: { code: "auth_throttled", message: "too many failed authorization attempts" } },
+        429,
+        throttledCors,
+      )
+    }
+    if (!authorized(request, this.token)) {
+      this.noteAuthFailure()
+      return this.response(
         { error: { code: "unauthorized", message: "agent RPC authorization required" } },
         401,
         cors,
       )
+    }
+    this.noteAuthSuccess()
     if (url.pathname !== "/rpc" || request.method !== "POST")
-      return json(
+      return this.response(
         { error: { code: "not_found", message: "unknown agent RPC endpoint" } },
         404,
         cors,
       )
-    const body = await this.readRequest(request)
+    let body: Awaited<ReturnType<CodingAgentRpcServer["readRequest"]>>
+    try {
+      body = await this.readRequest(request)
+    } catch {
+      return this.response(
+        { error: { code: "invalid_request", message: "request body could not be read" } },
+        400,
+        cors,
+      )
+    }
     if (!body.ok)
-      return json({ error: { code: "invalid_request", message: body.error } }, body.status, cors)
+      return this.response(
+        { error: { code: "invalid_request", message: body.error } },
+        body.status,
+        cors,
+      )
     let rpc: RpcRequest
     try {
       rpc = parseRpcRequest(body.text)
     } catch (error) {
-      return json(
+      return this.response(
         { error: this.protocolError("invalid_request", error, "request body is invalid") },
         400,
         cors,
@@ -150,7 +217,7 @@ export class CodingAgentRpcServer {
     try {
       return await this.dispatch(rpc, request, cors)
     } catch (error) {
-      return json(
+      return this.response(
         { error: this.protocolError("rpc_failed", error, "request could not be completed") },
         500,
         cors,
@@ -164,15 +231,27 @@ export class CodingAgentRpcServer {
     httpRequest: Request,
     cors: Headers,
   ): Promise<Response> {
+    const json = (
+      value: unknown,
+      status: number,
+      inherited: Headers,
+      includeErrorStacks = false,
+    ): Response => this.response(value, status, inherited, includeErrorStacks)
     switch (request.method) {
       case "session.create": {
         const params = record(request.params)
+        if (params.sessionId !== undefined && !isValidSessionId(params.sessionId))
+          return json(
+            { error: { code: "invalid_session", message: "sessionId is invalid" } },
+            422,
+            cors,
+          )
         const snapshot =
           this.host.snapshot ??
           (await this.host.start({
             cwd: this.options.cwd,
             backend: this.options.backend.info.name,
-            ...(typeof params.sessionId === "string" ? { sessionId: params.sessionId } : {}),
+            ...(params.sessionId === undefined ? {} : { sessionId: params.sessionId }),
           }))
         return json(snapshot, 200, cors)
       }
@@ -182,10 +261,7 @@ export class CodingAgentRpcServer {
       }
       case "session.resume": {
         const params = record(request.params)
-        if (
-          typeof params.sessionId !== "string" ||
-          !/^[A-Za-z0-9._:-]{1,128}$/.test(params.sessionId)
-        )
+        if (!isValidSessionId(params.sessionId))
           return json(
             { error: { code: "invalid_session", message: "sessionId is required" } },
             422,
@@ -233,7 +309,7 @@ export class CodingAgentRpcServer {
       case "session.fork": {
         this.requireSnapshot()
         const params = record(request.params)
-        if (params.targetSessionId !== undefined && typeof params.targetSessionId !== "string")
+        if (params.targetSessionId !== undefined && !isValidSessionId(params.targetSessionId))
           return json(
             { error: { code: "invalid_session", message: "targetSessionId must be a string" } },
             422,
@@ -259,8 +335,11 @@ export class CodingAgentRpcServer {
             cors,
           )
         const stream = this.host.prompt(params.message, httpRequest.signal)
+        const sessionId = this.host.snapshot?.id
         const encoder = new TextEncoder()
         const includeErrorStacks = this.options.exposeErrorStacks === true
+        const maxResponseBytes = this.maxResponseBytes
+        let finished = false
         const turnError = (error: unknown) => ({
           code: "turn_failed",
           ...publicErrorDetails(error, "turn could not be completed", includeErrorStacks),
@@ -268,18 +347,43 @@ export class CodingAgentRpcServer {
         const body = new ReadableStream<Uint8Array>({
           async start(controller) {
             try {
-              for await (const event of stream)
+              for await (const event of stream) {
+                if (finished) return
+                const serialized = serializeRpcValue(event, includeErrorStacks)
+                if (serializedByteLength(`data: ${serialized}\n\n`) > maxResponseBytes)
+                  throw new Error("turn event exceeded the configured response limit")
+                controller.enqueue(encoder.encode(`data: ${serialized}\n\n`))
+              }
+              if (!finished) {
+                finished = true
+                controller.close()
+              }
+            } catch (error) {
+              if (finished) return
+              finished = true
+              try {
+                const detail = JSON.stringify(turnError(error))
+                const frame = `event: error\ndata: ${detail}\n\n`
                 controller.enqueue(
-                  encoder.encode(`data: ${serializeRpcValue(event, includeErrorStacks)}\n\n`),
+                  encoder.encode(
+                    serializedByteLength(frame) <= maxResponseBytes
+                      ? frame
+                      : 'event: error\ndata: {"code":"turn_failed","message":"turn could not be completed"}\n\n',
+                  ),
                 )
-              controller.close()
-            } catch (error) {
-              controller.enqueue(
-                encoder.encode(`event: error\ndata: ${JSON.stringify(turnError(error))}\n\n`),
-              )
-              controller.close()
+                controller.close()
+              } catch {
+                // A disconnected client has already cancelled the stream; there is no response
+                // channel left on which to report the producer error.
+              }
             }
           },
+          cancel: async () => {
+            if (finished) return
+            finished = true
+            if (sessionId !== undefined)
+              await this.host.backend.cancel(sessionId, "client disconnected").catch(() => {})
+          },
         })
         return new Response(body, {
           status: 200,
@@ -465,6 +569,38 @@ export class CodingAgentRpcServer {
           this.options.exposeErrorStacks === true,
         )
       }
+      case "review.run": {
+        const params = record(request.params)
+        const allowed = new Set(["strict", "diff"])
+        if (Object.keys(params).some((key) => !allowed.has(key)))
+          return json(
+            { error: { code: "invalid_review", message: "review accepts only strict and diff" } },
+            422,
+            cors,
+          )
+        if (params.strict !== undefined && typeof params.strict !== "boolean")
+          return json(
+            { error: { code: "invalid_review", message: "strict must be a boolean" } },
+            422,
+            cors,
+          )
+        if (params.diff !== undefined && !isSafeReviewDiff(params.diff))
+          return json(
+            { error: { code: "invalid_review", message: "diff must be a safe Git ref" } },
+            422,
+            cors,
+          )
+        return json(
+          await runNifraReview({
+            cwd: this.options.cwd,
+            ...(this.options.verification === undefined ? {} : this.options.verification),
+            ...(params.strict === undefined ? {} : { strict: params.strict }),
+            ...(params.diff === undefined ? {} : { diff: params.diff }),
+          }),
+          200,
+          cors,
+        )
+      }
       case "project.diff": {
         this.requireSnapshot()
         return json(
@@ -495,6 +631,43 @@ export class CodingAgentRpcServer {
     return this.host.snapshot
   }
 
+  private authBackoffMs(): number {
+    const now = Date.now()
+    if (this.authFailures === 0) return 0
+    if (now - this.authLastFailureAt > AUTH_FAILURE_WINDOW_MS) {
+      this.noteAuthSuccess()
+      return 0
+    }
+    return Math.max(0, this.authBackoffUntil - now)
+  }
+
+  private noteAuthFailure(): void {
+    const now = Date.now()
+    if (now - this.authLastFailureAt > AUTH_FAILURE_WINDOW_MS) this.authFailures = 0
+    this.authLastFailureAt = now
+    this.authFailures = Math.min(this.authFailures + 1, 31)
+    if (this.authFailures >= AUTH_FAILURE_THRESHOLD) {
+      const exponent = this.authFailures - AUTH_FAILURE_THRESHOLD
+      const delay = Math.min(AUTH_MAX_BACKOFF_MS, 250 * 2 ** exponent)
+      this.authBackoffUntil = now + delay
+    }
+  }
+
+  private noteAuthSuccess(): void {
+    this.authFailures = 0
+    this.authLastFailureAt = 0
+    this.authBackoffUntil = 0
+  }
+
+  private response(
+    value: unknown,
+    status: number,
+    inherited: Headers,
+    includeErrorStacks = false,
+  ): Response {
+    return json(value, status, inherited, includeErrorStacks, this.maxResponseBytes)
+  }
+
   private protocolError(
     code: string,
     error: unknown,
@@ -509,15 +682,34 @@ export class CodingAgentRpcServer {
     | { readonly ok: true; readonly text: string }
     | { readonly ok: false; readonly status: 400 | 413; readonly error: string }
   > {
-    const length = Number(request.headers.get("content-length") ?? "0")
-    if (Number.isFinite(length) && length > this.maxBodyBytes)
-      return { ok: false, status: 413, error: "request body is too large" }
+    const declared = request.headers.get("content-length")
+    if (declared !== null) {
+      const length = decimalContentLength(declared)
+      if (length === undefined)
+        return { ok: false, status: 400, error: "content-length is invalid" }
+      if (length > this.maxBodyBytes)
+        return { ok: false, status: 413, error: "request body is too large" }
+    }
     const result = await readBoundedText(request.body, this.maxBodyBytes)
     if (result.truncated) return { ok: false, status: 413, error: "request body is too large" }
     return { ok: true, text: result.text }
   }
 }
 
+function decimalContentLength(value: string): number | undefined {
+  if (!/^\d+$/.test(value)) return undefined
+  let length = 0
+  for (const character of value) {
+    length = length * 10 + (character.charCodeAt(0) - 48)
+    if (!Number.isSafeInteger(length)) return Number.POSITIVE_INFINITY
+  }
+  return length
+}
+
+function isValidSessionId(value: unknown): value is string {
+  return typeof value === "string" && SESSION_ID_PATTERN.test(value)
+}
+
 function parseRpcRequest(text: string): RpcRequest {
   let value: unknown
   try {
@@ -545,6 +737,13 @@ function authorized(request: Request, token: string): boolean {
   return constantTimeEqual(header.slice(7), token)
 }
 
+function randomAuthToken(): string {
+  const bytes = crypto.getRandomValues(new Uint8Array(MIN_REMOTE_TOKEN_BYTES))
+  let binary = ""
+  for (const byte of bytes) binary += String.fromCharCode(byte)
+  return btoa(binary).replaceAll("+", "-").replaceAll("/", "_").replaceAll("=", "")
+}
+
 function constantTimeEqual(left: string, right: string): boolean {
   const a = new TextEncoder().encode(left)
   const b = new TextEncoder().encode(right)
@@ -580,19 +779,48 @@ function json(
   status: number,
   inherited: Headers,
   includeErrorStacks = false,
+  maxBytes = DEFAULT_MAX_RPC_RESPONSE_BYTES,
 ): Response {
   const headers = new Headers(inherited)
   headers.set("content-type", "application/json; charset=utf-8")
   // lgtm [js/stack-trace-exposure] the default replacer removes stack fields from protocol data;
   // stacks are serialized only in explicitly opt-in, loopback-only diagnostics mode.
-  return new Response(serializeRpcValue(value, includeErrorStacks), { status, headers })
+  let body: string
+  let responseStatus = status
+  try {
+    body = serializeRpcValue(value, includeErrorStacks)
+  } catch {
+    body = JSON.stringify({
+      error: {
+        code: "response_serialization_failed",
+        message: "RPC response could not be serialized",
+      },
+    })
+    responseStatus = 500
+  }
+  if (serializedByteLength(body) > maxBytes) {
+    body = JSON.stringify({
+      error: {
+        code: "response_too_large",
+        message: "RPC response exceeded the configured size limit",
+      },
+    })
+    responseStatus = 500
+  }
+  return new Response(body, { status: responseStatus, headers })
 }
 
 function serializeRpcValue(value: unknown, includeErrorStacks: boolean): string {
-  return JSON.stringify(
+  const serialized = JSON.stringify(
     value,
     includeErrorStacks
       ? undefined
       : (key: string, nested: unknown) => (key === "stack" ? undefined : nested),
   )
+  if (serialized === undefined) throw new TypeError("RPC value is not JSON serializable")
+  return serialized
+}
+
+function serializedByteLength(value: string): number {
+  return new TextEncoder().encode(value).byteLength
 }