Merge pull request #4205 from code-yeongyu/fix/comment-checker-apply-patch-payloads

fix(comment-checker): handle apply_patch payloads

src/hooks/comment-checker/apply-patch-edits.ts

@@ -0,0 +1,180 @@
+import type { ApplyPatchEdit } from "./cli-runner"
+
+type ApplyPatchFileMetadata = {
+  readonly filePath: string
+  readonly movePath?: string
+  readonly before: string
+  readonly after: string
+  readonly type?: string
+}
+
+type ApplyPatchAccumulator = {
+  operation: "add" | "update" | "delete"
+  filePath: string
+  movePath?: string
+  oldLines: string[]
+  newLines: string[]
+}
+
+export function extractApplyPatchEdits(
+  details: unknown,
+  args?: Record<string, unknown>,
+): ApplyPatchEdit[] {
+  const metadataEdits = getApplyPatchMetadataFiles(details)
+    .filter((file) => file.type?.toLowerCase() !== "delete")
+    .map((file) => ({
+      filePath: file.movePath ?? file.filePath,
+      before: file.before,
+      after: file.after,
+    }))
+
+  if (metadataEdits.length > 0) return metadataEdits
+
+  const patch = args === undefined ? undefined : getString(args, ["patchText", "input", "patch", "command"])
+  if (patch === undefined) return []
+
+  return parseApplyPatchEdits(patch)
+}
+
+function getApplyPatchMetadataFiles(details: unknown): ApplyPatchFileMetadata[] {
+  if (!isRecord(details)) return []
+
+  const direct = readApplyPatchMetadataFiles(details["files"])
+  if (direct.length > 0) return direct
+
+  const resultDetails = details["result"]
+  const result = isRecord(resultDetails) ? readApplyPatchMetadataFiles(resultDetails["files"]) : []
+  if (result.length > 0) return result
+
+  const metadataDetails = details["metadata"]
+  return isRecord(metadataDetails) ? readApplyPatchMetadataFiles(metadataDetails["files"]) : []
+}
+
+function readApplyPatchMetadataFiles(value: unknown): ApplyPatchFileMetadata[] {
+  if (!Array.isArray(value)) return []
+
+  const files: ApplyPatchFileMetadata[] = []
+  for (const item of value) {
+    if (!isRecord(item)) continue
+
+    const filePath = getString(item, ["filePath", "file_path", "path"])
+    const movePath = getString(item, ["movePath", "move_path"])
+    const before = getString(item, ["before", "old", "oldString", "old_string"])
+    const after = getString(item, ["after", "new", "newString", "new_string"])
+    const type = getString(item, ["type", "operation"])
+
+    if (filePath === undefined || before === undefined || after === undefined) continue
+
+    files.push({
+      filePath,
+      before,
+      after,
+      ...(movePath === undefined ? {} : { movePath }),
+      ...(type === undefined ? {} : { type }),
+    })
+  }
+
+  return files
+}
+
+function parseApplyPatchEdits(patch: string): ApplyPatchEdit[] {
+  const edits: ApplyPatchEdit[] = []
+  let current: ApplyPatchAccumulator | undefined
+
+  const flush = (): void => {
+    if (current === undefined) return
+
+    if (current.operation === "add") {
+      const after = joinPatchLines(current.newLines)
+      if (after.length > 0) {
+        edits.push({ filePath: current.filePath, before: "", after })
+      }
+    }
+
+    if (current.operation === "update") {
+      const after = joinPatchLines(current.newLines)
+      if (after.length > 0) {
+        edits.push({
+          filePath: current.movePath ?? current.filePath,
+          before: joinPatchLines(current.oldLines),
+          after,
+        })
+      }
+    }
+
+    current = undefined
+  }
+
+  for (const line of patch.split(/\r?\n/)) {
+    if (line === "*** Begin Patch" || line === "*** End Patch") continue
+
+    if (line.startsWith("*** Add File: ")) {
+      flush()
+      current = makeAccumulator("add", line.slice("*** Add File: ".length).trim())
+      continue
+    }
+
+    if (line.startsWith("*** Update File: ")) {
+      flush()
+      current = makeAccumulator("update", line.slice("*** Update File: ".length).trim())
+      continue
+    }
+
+    if (line.startsWith("*** Delete File: ")) {
+      flush()
+      current = makeAccumulator("delete", line.slice("*** Delete File: ".length).trim())
+      continue
+    }
+
+    if (line.startsWith("*** Move to: ")) {
+      if (current?.operation === "update") {
+        current.movePath = line.slice("*** Move to: ".length).trim()
+      }
+      continue
+    }
+
+    if (current === undefined || line.startsWith("@@")) continue
+
+    if (current.operation === "add") {
+      if (line.startsWith("+")) current.newLines.push(line.slice(1))
+      continue
+    }
+
+    if (current.operation === "update") {
+      if (line.startsWith("-")) current.oldLines.push(line.slice(1))
+      if (line.startsWith("+")) current.newLines.push(line.slice(1))
+    }
+  }
+
+  flush()
+  return edits
+}
+
+function makeAccumulator(
+  operation: ApplyPatchAccumulator["operation"],
+  filePath: string,
+): ApplyPatchAccumulator {
+  return {
+    operation,
+    filePath,
+    oldLines: [],
+    newLines: [],
+  }
+}
+
+function getString(input: Record<string, unknown>, keys: readonly string[]): string | undefined {
+  for (const key of keys) {
+    const value = input[key]
+    if (typeof value === "string") return value
+  }
+
+  return undefined
+}
+
+function joinPatchLines(lines: readonly string[]): string {
+  return lines.length === 0 ? "" : `${lines.join("\n")}\n`
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null
+}

src/hooks/comment-checker/hook.apply-patch.test.ts

@@ -82,4 +82,130 @@ describe("comment-checker apply_patch integration", () => {
     // then
     expect(processApplyPatchEditsWithCli).toHaveBeenCalledTimes(0)
   })
+
+  it("#given apply_patch metadata nested under result #when hook runs #then checks edited files", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "apply_patch", sessionID: "ses_test", callID: "call_test" }
+    const output = {
+      title: "ok",
+      output: "Success",
+      metadata: {
+        result: {
+          files: [
+            {
+              path: "/repo/src/result.ts",
+              old: "const a = 1\n",
+              new: "// result comment\nconst a = 1\n",
+              operation: "update",
+            },
+          ],
+        },
+      },
+    }
+
+    // when
+    await hooks["tool.execute.after"](input, output)
+
+    // then
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledTimes(1)
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledWith(
+      "ses_test",
+      [
+        {
+          filePath: "/repo/src/result.ts",
+          before: "const a = 1\n",
+          after: "// result comment\nconst a = 1\n",
+        },
+      ],
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
+
+  it("#given apply_patch metadata nested under metadata #when hook runs #then checks edited files", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "apply_patch", sessionID: "ses_test", callID: "call_test" }
+    const output = {
+      title: "ok",
+      output: "Success",
+      metadata: {
+        metadata: {
+          files: [
+            {
+              file_path: "/repo/src/metadata.ts",
+              old_string: "const b = 1\n",
+              new_string: "// metadata comment\nconst b = 1\n",
+              type: "update",
+            },
+          ],
+        },
+      },
+    }
+
+    // when
+    await hooks["tool.execute.after"](input, output)
+
+    // then
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledTimes(1)
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledWith(
+      "ses_test",
+      [
+        {
+          filePath: "/repo/src/metadata.ts",
+          before: "const b = 1\n",
+          after: "// metadata comment\nconst b = 1\n",
+        },
+      ],
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
+
+  it("#given apply_patch patchText args without metadata #when hook runs #then parses patch edits", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = {
+      tool: "apply_patch",
+      sessionID: "ses_test",
+      callID: "call_test",
+      args: {
+        patchText: [
+          "*** Begin Patch",
+          "*** Update File: /repo/src/raw.ts",
+          "@@",
+          "-const c = 1",
+          "+// raw comment",
+          "+const c = 1",
+          "*** End Patch",
+        ].join("\n"),
+      },
+    }
+    const output = { title: "ok", output: "Success", metadata: {} }
+
+    // when
+    await hooks["tool.execute.after"](input, output)
+
+    // then
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledTimes(1)
+    expect(processApplyPatchEditsWithCli).toHaveBeenCalledWith(
+      "ses_test",
+      [
+        {
+          filePath: "/repo/src/raw.ts",
+          before: "const c = 1\n",
+          after: "// raw comment\nconst c = 1\n",
+        },
+      ],
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
 })

src/hooks/comment-checker/hook.before-after.test.ts

@@ -0,0 +1,133 @@
+import { afterAll, afterEach, beforeEach, describe, expect, it, mock } from "bun:test"
+
+const processWithCli = mock(async () => {})
+
+mock.module("./cli-runner", () => ({
+  initializeCommentCheckerCli: () => {},
+  getCommentCheckerCliPathPromise: () => Promise.resolve("/tmp/fake-comment-checker"),
+  isCliPathUsable: () => true,
+  processWithCli,
+  processApplyPatchEditsWithCli: async () => {},
+}))
+
+afterAll(() => {
+  mock.restore()
+})
+
+const { createCommentCheckerHooks } = await import("./hook")
+const { stopPendingCallCleanup } = await import("./pending-calls")
+const { _resetCommentCheckerInitializationForTesting } = await import("./initialization-gate")
+
+describe("comment-checker mutation tool routing", () => {
+  beforeEach(() => {
+    processWithCli.mockClear()
+    stopPendingCallCleanup()
+    _resetCommentCheckerInitializationForTesting()
+  })
+
+  afterEach(() => {
+    stopPendingCallCleanup()
+    _resetCommentCheckerInitializationForTesting()
+  })
+
+  it("#given write tool with filePath #when before and after hooks run #then it checks the pending write", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "write", sessionID: "ses_test", callID: "call_write" }
+
+    // when
+    await hooks["tool.execute.before"](input, {
+      args: { filePath: "/repo/src/write.ts", content: "// write comment\nconst a = 1\n" },
+    })
+    await hooks["tool.execute.after"](input, { title: "ok", output: "Success", metadata: {} })
+
+    // then
+    expect(processWithCli).toHaveBeenCalledTimes(1)
+    expect(processWithCli).toHaveBeenCalledWith(
+      input,
+      expect.objectContaining({
+        filePath: "/repo/src/write.ts",
+        content: "// write comment\nconst a = 1\n",
+        tool: "write",
+        sessionID: "ses_test",
+      }),
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
+
+  it("#given edit tool with file_path #when before and after hooks run #then it checks the pending edit", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "edit", sessionID: "ses_test", callID: "call_edit" }
+
+    // when
+    await hooks["tool.execute.before"](input, {
+      args: {
+        file_path: "/repo/src/edit.ts",
+        old_string: "const b = 1\n",
+        new_string: "// edit comment\nconst b = 1\n",
+      },
+    })
+    await hooks["tool.execute.after"](input, { title: "ok", output: "Success", metadata: {} })
+
+    // then
+    expect(processWithCli).toHaveBeenCalledTimes(1)
+    expect(processWithCli).toHaveBeenCalledWith(
+      input,
+      expect.objectContaining({
+        filePath: "/repo/src/edit.ts",
+        oldString: "const b = 1\n",
+        newString: "// edit comment\nconst b = 1\n",
+        tool: "edit",
+      }),
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
+
+  it("#given multiedit tool with path #when before and after hooks run #then it checks the pending multiedit", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "multiedit", sessionID: "ses_test", callID: "call_multiedit" }
+    const edits = [{ old_string: "const c = 1\n", new_string: "// multiedit comment\nconst c = 1\n" }]
+
+    // when
+    await hooks["tool.execute.before"](input, {
+      args: { path: "/repo/src/multiedit.ts", edits },
+    })
+    await hooks["tool.execute.after"](input, { title: "ok", output: "Success", metadata: {} })
+
+    // then
+    expect(processWithCli).toHaveBeenCalledTimes(1)
+    expect(processWithCli).toHaveBeenCalledWith(
+      input,
+      expect.objectContaining({
+        filePath: "/repo/src/multiedit.ts",
+        edits,
+        tool: "multiedit",
+      }),
+      expect.any(Object),
+      "/tmp/fake-comment-checker",
+      undefined,
+      expect.any(Function),
+    )
+  })
+
+  it("#given non-mutation tool #when before and after hooks run #then it does not run the checker", async () => {
+    // given
+    const hooks = createCommentCheckerHooks()
+    const input = { tool: "read", sessionID: "ses_test", callID: "call_read" }
+
+    // when
+    await hooks["tool.execute.before"](input, { args: { filePath: "/repo/src/read.ts" } })
+    await hooks["tool.execute.after"](input, { title: "ok", output: "Success", metadata: {} })
+
+    // then
+    expect(processWithCli).toHaveBeenCalledTimes(0)
+  })
+})

src/hooks/comment-checker/hook.ts

@@ -1,20 +1,6 @@
 import type { PendingCall } from "./types"
 import type { CommentCheckerConfig } from "../../config/schema"
 
-import z from "zod"
-
-const ApplyPatchMetadataSchema = z.object({
-  files: z.array(
-    z.object({
-      filePath: z.string(),
-      movePath: z.string().optional(),
-      before: z.string(),
-      after: z.string(),
-      type: z.string().optional(),
-    }),
-  ),
-})
-
 import {
   initializeCommentCheckerCli,
   getCommentCheckerCliPathPromise,
@@ -22,6 +8,7 @@ import {
   processWithCli,
   processApplyPatchEditsWithCli,
 } from "./cli-runner"
+import { extractApplyPatchEdits } from "./apply-patch-edits"
 import {
   registerPendingCall,
   startPendingCallCleanup,
@@ -104,7 +91,7 @@ export function createCommentCheckerHooks(config?: CommentCheckerConfig) {
     },
 
     "tool.execute.after": async (
-      input: { tool: string; sessionID: string; callID: string },
+      input: { tool: string; sessionID: string; callID: string; args?: Record<string, unknown> },
       output: { title: string; output: string; metadata: unknown },
     ): Promise<void> => {
       debugLog("tool.execute.after:", { tool: input.tool, callID: input.callID })
@@ -126,20 +113,7 @@ export function createCommentCheckerHooks(config?: CommentCheckerConfig) {
 
 
       if (toolLower === "apply_patch") {
-        const parsed = ApplyPatchMetadataSchema.safeParse(output.metadata)
+        const edits = extractApplyPatchEdits(output.metadata, input.args)
-        if (!parsed.success) {
-          debugLog("apply_patch metadata schema mismatch, skipping")
-          return
-        }
-
-        const edits = parsed.data.files
-          .filter((f) => f.type !== "delete")
-          .map((f) => ({
-            filePath: f.movePath ?? f.filePath,
-            before: f.before,
-            after: f.after,
-          }))
-
         if (edits.length === 0) {
           debugLog("apply_patch had no editable files, skipping")
           return

src/hooks/comment-checker/initialization-gate.ts

@@ -5,3 +5,7 @@ export function ensureCommentCheckerInitialization(initializer: () => void): voi
   initialized = true
   initializer()
 }
+
+export function _resetCommentCheckerInitializationForTesting(): void {
+  initialized = false
+}

src/plugin/tool-execute-after.test.ts

@@ -146,4 +146,30 @@ describe("createToolExecuteAfterHandler", () => {
     // then
     expect(output).toEqual({ title: "result", output: "read output", metadata: {} })
   })
+
+  it("#given after input includes tool args #when comment checker runs #then it receives the args", async () => {
+    // given
+    let seenArgs: Record<string, unknown> | undefined
+    const handler = createToolExecuteAfterHandler({
+      ctx: { directory: "/repo" } as never,
+      hooks: {
+        commentChecker: {
+          "tool.execute.after": async (input) => {
+            seenArgs = input.args
+          },
+        },
+      } as never,
+    })
+
+    const args = { patchText: "*** Begin Patch\n*** End Patch" }
+
+    // when
+    await handler(
+      { tool: "apply_patch", sessionID: "ses_parent", callID: "call_patch", args },
+      { title: "result", output: "Success", metadata: {} },
+    )
+
+    // then
+    expect(seenArgs).toBe(args)
+  })
 })

src/plugin/tool-execute-after.ts

@@ -6,6 +6,21 @@ import type { PluginContext } from "./types"
 
 const VERIFICATION_ATTEMPT_PATTERN = /<ulw_verification_attempt_id>(.*?)<\/ulw_verification_attempt_id>/i
 
+type ToolExecuteAfterInput = {
+  readonly tool: string
+  readonly sessionID: string
+  readonly callID?: string
+  readonly callId?: string
+  readonly call_id?: string
+  readonly args?: Record<string, unknown>
+}
+
+type ToolExecuteAfterOutput = {
+  title: string
+  output: string
+  metadata: Record<string, unknown>
+}
+
 function getMetadataString(metadata: Record<string, unknown> | undefined, keys: string[]): string | undefined {
   for (const key of keys) {
     const value = metadata?.[key]
@@ -29,10 +44,8 @@ export function createToolExecuteAfterHandler(args: {
   ctx: PluginContext
   hooks: CreatedHooks
 }): (
-  input: { tool: string; sessionID: string; callID: string },
+  input: ToolExecuteAfterInput,
-  output:
+  output: ToolExecuteAfterOutput | undefined,
-    | { title: string; output: string; metadata: Record<string, unknown> }
-    | undefined,
 ) => Promise<void> {
   const { ctx, hooks } = args
 
@@ -40,8 +53,8 @@ export function createToolExecuteAfterHandler(args: {
   // We must treat their identity as a best-effort correlation key, not a guaranteed public contract.
 
   return async (
-    input: { tool: string; sessionID: string; callID?: string; callId?: string; call_id?: string },
+    input: ToolExecuteAfterInput,
-    output: { title: string; output: string; metadata: Record<string, unknown> } | undefined,
+    output: ToolExecuteAfterOutput | undefined,
   ): Promise<void> => {
     if (!output) return
 
@@ -49,6 +62,7 @@ export function createToolExecuteAfterHandler(args: {
       tool: input.tool,
       sessionID: input.sessionID,
       callID: input.callID ?? input.callId ?? input.call_id ?? "",
+      ...(input.args === undefined ? {} : { args: input.args }),
     }
 
     const nativeSessionId = getMetadataString(output.metadata, ["sessionId", "sessionID", "session_id"])