goodbuddy/src/main/agent/goodbuddy-harness-control-plane.test.ts

import { describe, expect, it, vi } from 'vitest'
import { Context } from '@deepseek-ai/cordis'
import type { Stream } from '@agentclientprotocol/sdk'
import { resolve } from 'node:path'
import {
  GOODBUDDY_HANDSHAKE,
  GOODBUDDY_PREPARE,
  GoodBuddyCredentialProvider,
  GoodBuddyHarnessControlPlane,
  GoodBuddySandboxRetryLedger,
  createBoundedAcpStream
} from './goodbuddy-harness-control-plane'

function execution(
  callId: string,
  name: string,
  args: Record<string, unknown>
) {
  return {
    callId,
    rootCallId: callId,
    name,
    arguments: args,
    signal: new AbortController().signal,
    token: Symbol('execution')
  } as never
}

const sandboxDenied = {
  isError: false,
  value: {
    sandbox: {
      denied: true
    }
  },
  content: []
} as const

function controlPlane() {
  return new GoodBuddyHarnessControlPlane({} as Context, {
    provider: 'goodbuddy',
    model: 'deepseek-test',
    workspace: resolve('workspace'),
    harnessVersion: '0.1.0-rc.6',
    sandbox: { provider: 'test', enforcement: 'full' },
    credentialRefs: ['GOODBUDDY_API_KEY'],
    skills: []
  })
}

function stubAgentContext() {
  const listeners = new Map<
    string,
    (...args: unknown[]) => unknown
  >()
  const extNotification = vi.fn(async () => undefined)
  const handle = {
    agent: {
      session: {
        id: 'session-output',
        header: { id: 'session-output' },
        events: []
      },
      cancel: vi.fn()
    }
  }
  const ctx = {
    on: vi.fn(
      (
        name: string,
        listener: (...args: unknown[]) => unknown
      ) => {
        listeners.set(name, listener)
        return vi.fn()
      }
    )
  } as unknown as Context
  const subject = new GoodBuddyHarnessControlPlane(ctx, {
    provider: 'goodbuddy',
    model: 'deepseek-test',
    workspace: resolve('workspace'),
    harnessVersion: '0.1.0-rc.6',
    sandbox: { provider: 'test', enforcement: 'full' },
    credentialRefs: ['GOODBUDDY_API_KEY'],
    skills: [],
    maxEventCharacters: 10_000,
    maxRequestCharacters: 180
  })
  const internals = subject as unknown as {
    connection: {
      extNotification: typeof extNotification
    }
    sessions: Map<
      string,
      {
        handle: typeof handle
        inflight: {
          requestId: string
          messageId: string
          resolve: (reason: string) => void
          reject: (error: unknown) => void
          emittedCharacters: number
          eventTail: Promise<void>
          eventError?: unknown
        }
      }
    >
    observeSessions(): void
  }
  internals.connection = { extNotification }
  internals.sessions.set('session-output', {
    handle,
    inflight: {
      requestId: 'request-output',
      messageId: 'message-output',
      resolve: vi.fn(),
      reject: vi.fn(),
      emittedCharacters: 0,
      eventTail: Promise.resolve()
    }
  })
  internals.observeSessions()
  return { listeners, extNotification, handle, internals }
}

describe('GoodBuddy Harness internal control plane', () => {
  it('requires a versioned handshake before privileged extensions', async () => {
    const subject = controlPlane()

    await expect(
      subject.extensionMethod(GOODBUDDY_PREPARE, {
        sessionId: 'session',
        requestId: 'request',
        mode: 'execute'
      })
    ).rejects.toThrow('GoodBuddy handshake is required')
    await expect(
      subject.extensionMethod(GOODBUDDY_HANDSHAKE, {
        controlProtocolVersion: 9
      })
    ).rejects.toThrow(
      'incompatible GoodBuddy Harness control protocol'
    )
    await expect(
      subject.extensionMethod(GOODBUDDY_HANDSHAKE, {
        controlProtocolVersion: 1
      })
    ).resolves.toMatchObject({
      controlProtocolVersion: 1,
      supports: {
        cancellation: true,
        sessionRelease: true,
        oneShotApproval: true,
        credentialResolution: true
      },
      sandbox: { enforcement: 'full' }
    })
  })

  it('keeps credentials memory-only, allowlisted, and read-only', async () => {
    const provider = new GoodBuddyCredentialProvider(
      new Context(),
      new Set(['GOODBUDDY_API_KEY'])
    )
    const resolver = vi
      .fn()
      .mockResolvedValue('secret-from-main')
    provider.bind(resolver)

    await expect(
      provider.resolve('GOODBUDDY_API_KEY' as never)
    ).resolves.toEqual({
      value: 'secret-from-main',
      source: 'goodbuddy-main'
    })
    await expect(
      provider.resolve('OTHER_KEY' as never)
    ).resolves.toBeUndefined()
    expect(resolver).toHaveBeenCalledTimes(1)
    await expect(
      provider.set('GOODBUDDY_API_KEY' as never, 'x')
    ).rejects.toThrow('read-only')
  })

  it('fails closed on oversized inbound and outbound ACP frames', async () => {
    const inbound = new TransformStream<
      Record<string, unknown>,
      Record<string, unknown>
    >()
    const outbound = new TransformStream<
      Record<string, unknown>,
      Record<string, unknown>
    >()
    const stream = createBoundedAcpStream(
      ({
        readable: inbound.readable,
        writable: outbound.writable
      } as unknown as Stream),
      16
    )
    const inputWriter = inbound.writable.getWriter()
    const reader = stream.readable.getReader()
    const read = reader.read()
    await inputWriter.write({ value: 'too-long-for-frame' })
    await expect(read).rejects.toThrow('input frame exceeds')

    const writer = stream.writable.getWriter()
    await expect(
      writer.write({ value: 'too-long-for-frame' } as never)
    ).rejects.toThrow('output frame exceeds')
  })

  it('counts the complete emitted envelope against the request limit', async () => {
    const { listeners, extNotification, handle, internals } =
      stubAgentContext()
    const sessionEvent = listeners.get('session/event')!
    sessionEvent(
      handle.agent.session,
      {
        type: 'assistant/chunk',
        data: {
          chunk: {
            type: 'text-delta',
            text: 'x'.repeat(80)
          }
        }
      }
    )
    sessionEvent(
      handle.agent.session,
      {
        type: 'assistant/chunk',
        data: {
          chunk: {
            type: 'usage',
            usage: {
              inputTokens: 1,
              outputTokens: 1,
              cacheReadTokens: 0,
              cacheWriteTokens: 0
            }
          }
        }
      }
    )
    await internals.sessions.get('session-output')!.inflight.eventTail

    expect(extNotification).toHaveBeenCalledTimes(1)
    expect(handle.agent.cancel).toHaveBeenCalledWith({
      kind: 'user'
    })
    expect(
      internals.sessions.get('session-output')!.inflight.eventError
    ).toEqual(
      new Error(
        'GoodBuddy Harness control request output exceeds safety limit'
      )
    )
    expect(
      internals.sessions.get('session-output')!.inflight.emittedCharacters
    ).toBeGreaterThan(180)
  })

  it('requires a matching real denial and consumes it once', () => {
    const ledger = new GoodBuddySandboxRetryLedger()
    const deniedArguments = {
      command: 'type C:\\outside\\file.txt',
      description: 'Read an outside file'
    }
    const retry = {
      ...deniedArguments,
      sandbox_permissions: 'danger-full-access',
      justification: 'The requested file is outside the workspace.'
    }

    expect(ledger.consumeRetry('pwsh', retry)).toBe(false)
    ledger.record(
      execution('denial-1', 'pwsh', deniedArguments),
      sandboxDenied as never
    )
    expect(
      ledger.consumeRetry('pwsh', {
        ...retry,
        command: 'type C:\\different\\file.txt'
      })
    ).toBe(false)
    expect(ledger.consumeRetry('bash', retry)).toBe(false)
    expect(ledger.consumeRetry('pwsh', retry)).toBe(true)
    expect(ledger.consumeRetry('pwsh', retry)).toBe(false)
  })

  it('rejects non-denials, narrow escalation, and reordered ambiguity', () => {
    const ledger = new GoodBuddySandboxRetryLedger()
    const deniedArguments = {
      description: 'Read an outside file',
      command: 'cat /outside/file'
    }
    ledger.record(execution('success', 'bash', deniedArguments), {
      isError: false,
      value: {},
      content: []
    } as never)
    expect(
      ledger.consumeRetry('bash', {
        command: 'cat /outside/file',
        description: 'Read an outside file',
        sandbox_permissions: 'danger-full-access',
        justification: 'The requested file is outside the workspace.'
      })
    ).toBe(false)

    ledger.record(
      execution('denial-2', 'bash', deniedArguments),
      sandboxDenied as never
    )
    expect(
      ledger.consumeRetry('bash', {
        command: 'cat /outside/file',
        description: 'Read an outside file',
        sandbox_permissions: 'workspace-write',
        justification: 'Retry in workspace-write.'
      })
    ).toBe(false)
    expect(
      ledger.consumeRetry('bash', {
        command: 'cat /outside/file',
        description: 'Read an outside file',
        sandbox_permissions: 'danger-full-access',
        justification: 'The requested file is outside the workspace.'
      })
    ).toBe(true)
  })
})