DSH-Arch-Setup/plugins/dsh-touched-git/lib/git-write-guard.js

import { randomUUID } from 'node:crypto'
import { canonicalWriteTarget, checkGitWriteScope } from './git-write-scope.js'

function isFileMutation(execution) {
  return execution?.name === 'write' || execution?.name === 'edit'
}

function sessionIdentity(execution) {
  return {
    sessionId: execution.agent?.session.header.id,
    cwd: execution.agent?.session.header.cwd,
  }
}

function failureMessage(error) {
  return error instanceof Error ? error.message : String(error)
}

function throwIfCancelled(execution) {
  execution.signal?.throwIfAborted()
}

function warningContext(warnings) {
  return {
    id: randomUUID(),
    role: 'user',
    content: warnings.map(text => ({ type: 'text', text })),
    source: { kind: 'plugin', plugin: 'dsh-touched-git' },
  }
}

function augmentDecision(execution, result, decision, warnings) {
  if (warnings.length === 0) return decision

  const needsContext = execution.parent !== undefined || decision.kind === 'block' || Object.hasOwn(decision, 'value')
  const additionalContexts = needsContext
    ? [...decision.additionalContexts ?? [], warningContext(warnings)]
    : decision.additionalContexts

  if (decision.kind === 'block' || Object.hasOwn(decision, 'value')) {
    return { ...decision, additionalContexts }
  }

  return {
    ...decision,
    content: [...decision.content ?? result.content, ...warnings.map(text => ({ type: 'text', text }))],
    ...(additionalContexts === undefined ? {} : { additionalContexts }),
  }
}

export function registerGitWriteGuard(context, { store, onSuccessfulMutation }) {
  const certificates = new Map()
  let disposed = false

  function denialReason(execution) {
    if (!isFileMutation(execution)) return undefined

    const certificate = certificates.get(execution.token)
    if (!certificate || disposed) return 'Write blocked: the Git scope pre-execution check did not complete.'
    if (certificate.denial) return certificate.denial

    const { sessionId, cwd } = sessionIdentity(execution)
    if (sessionId !== certificate.sessionId || cwd !== certificate.cwd) {
      return 'Write blocked: the session or working directory changed after the Git scope check.'
    }

    try {
      if (certificate.repository && !store.isRegistered(sessionId, certificate.repository)) {
        return `Write blocked: ${certificate.repository} is not registered in this session’s Git Working Dirs menu. Call add_touched_git_directory for this repository or worktree before writing.`
      }
    } catch (error) {
      return `Write blocked: Git registration could not be checked: ${failureMessage(error)}`
    }
    return undefined
  }

  async function classify(execution, target) {
    throwIfCancelled(execution)
    const identity = sessionIdentity(execution)
    if (typeof identity.sessionId !== 'string' || identity.sessionId === '' || typeof identity.cwd !== 'string' || identity.cwd === '') {
      throw new Error('Write blocked: an active session with a working directory is required.')
    }

    const filePath = await context.fs.processPath(target)
    throwIfCancelled(execution)
    const scope = await checkGitWriteScope({ filePath, sessionId: identity.sessionId, store, signal: execution.signal })
    throwIfCancelled(execution)
    if (disposed) throw new Error('Write blocked: the Git scope guard was disposed during the check.')

    const current = sessionIdentity(execution)
    if (identity.sessionId !== current.sessionId || identity.cwd !== current.cwd) {
      throw new Error('Write blocked: the session or working directory changed during the Git scope check.')
    }

    const certificate = { ...identity, ...scope }
    certificates.set(execution.token, certificate)
    const denial = denialReason(execution)
    if (denial) throw new Error(denial)
    return certificate
  }

  async function beforeExecution(execution, next) {
    if (!isFileMutation(execution)) return next()

    certificates.delete(execution.token)
    const decision = await next()
    if (decision.kind === 'deny') return decision

    try {
      throwIfCancelled(execution)
      const identity = sessionIdentity(execution)
      if (typeof identity.sessionId !== 'string' || identity.sessionId === '' || typeof identity.cwd !== 'string' || identity.cwd === '') {
        throw new Error('Write blocked: an active session with a working directory is required.')
      }
      const requestedPath = execution.arguments.file_path
      const parentSegment = /(?:^|[\\/])\.\.(?:[\\/]|$)/
      const cwd = parentSegment.test(identity.cwd) || parentSegment.test(requestedPath)
        ? await canonicalWriteTarget(identity.cwd, execution.signal)
        : identity.cwd
      const target = await context.fs.resolve(requestedPath, { cwd, signal: execution.signal })
      if (sessionIdentity(execution).sessionId !== identity.sessionId || sessionIdentity(execution).cwd !== identity.cwd) {
        throw new Error('Write blocked: the session or working directory changed during path resolution.')
      }
      await classify(execution, target)
    } catch (error) {
      if (!disposed) certificates.set(execution.token, { denial: failureMessage(error) })
      throwIfCancelled(execution)
    }
    return decision
  }

  async function beforeMutation(target, execution, next) {
    if (!isFileMutation(execution)) return next()

    throwIfCancelled(execution)
    const denial = denialReason(execution)
    if (denial) throw new Error(denial)

    try {
      await classify(execution, target)
    } catch (error) {
      if (!disposed) certificates.set(execution.token, { denial: failureMessage(error) })
      throw error
    }

    const intent = await next()
    throwIfCancelled(execution)
    const finalDenial = denialReason(execution)
    if (finalDenial) throw new Error(finalDenial)
    return intent
  }

  async function afterExecution(execution, result, next) {
    const decision = await next()
    if (!isFileMutation(execution)) return decision

    const certificate = certificates.get(execution.token)
    if (!certificate || certificate.denial) return decision

    const warnings = certificate.warning ? [certificate.warning] : []
    if (!result.isError && !certificate.reported && onSuccessfulMutation) {
      certificate.reported = true
      try {
        await onSuccessfulMutation({
          sessionId: certificate.sessionId,
          callId: execution.callId,
          toolName: execution.name,
          filePath: certificate.filePath,
          repository: certificate.repository,
          warning: certificate.warning,
        })
      } catch (error) {
        warnings.push(`WARNING: The ${execution.name} operation completed, but its Git mutation record could not be saved: ${failureMessage(error)}`)
      }
    }
    return augmentDecision(execution, result, decision, warnings)
  }

  const disposers = []
  try {
    disposers.push(context.on('tools/pre-execute', beforeExecution))
    disposers.push(context.tools.guard(denialReason))
    disposers.push(context.on('fs/write-intent', beforeMutation, true))
    disposers.push(context.on('fs/edit-intent', beforeMutation, true))
    disposers.push(context.on('tools/post-execute', afterExecution, true))
    disposers.push(context.on('tools/result', execution => certificates.delete(execution.token)))
  } catch (error) {
    disposed = true
    certificates.clear()
    for (const dispose of disposers.reverse()) dispose()
    throw error
  }

  return () => {
    if (disposed) return
    disposed = true
    certificates.clear()
    for (const dispose of disposers.reverse()) dispose()
  }
}