oh-my-opencode/src/features/skill-mcp-manager/disconnect-cleanup.test.ts

import { Client } from "@modelcontextprotocol/sdk/client/index.js"
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"
import { afterEach, describe, expect, it } from "bun:test"
import { disconnectSession, registerProcessCleanup, unregisterProcessCleanup } from "./cleanup"
import type { ManagedClient, SkillMcpManagerState } from "./types"

const trackedStates: SkillMcpManagerState[] = []

afterEach(() => {
  for (const state of trackedStates) {
    unregisterProcessCleanup(state)
  }

  trackedStates.length = 0
})

function createState(): SkillMcpManagerState {
  const state: SkillMcpManagerState = {
    clients: new Map(),
    pendingConnections: new Map(),
    disconnectedSessions: new Map(),
    authProviders: new Map(),
    cleanupRegistered: false,
    cleanupInterval: null,
    cleanupHandlers: [],
    idleTimeoutMs: 5 * 60 * 1000,
    shutdownGeneration: 0,
    inFlightConnections: new Map(),
    disposed: false,
  }

  trackedStates.push(state)
  return state
}

function createManagedClient(skillName: string): ManagedClient {
  return {
    client: new Client(
      { name: `test-${skillName}`, version: "1.0.0" },
      { capabilities: {} }
    ),
    transport: new StreamableHTTPClientTransport(new URL("https://example.com/mcp")),
    skillName,
    lastUsedAt: Date.now(),
    connectionType: "http",
  }
}

describe("disconnectSession cleanup registration", () => {
  it("#given state with 1 client and cleanup registered #when disconnectSession removes last client #then process cleanup handlers are unregistered", async () => {
    // given
    const state = createState()
    const signalIntCountBeforeRegister = process.listenerCount("SIGINT")
    const signalTermCountBeforeRegister = process.listenerCount("SIGTERM")

    state.clients.set("session-1:skill-1:server-1", createManagedClient("skill-1"))
    registerProcessCleanup(state)

    // when
    await disconnectSession(state, "session-1")

    // then
    expect(state.cleanupRegistered).toBe(false)
    expect(state.cleanupHandlers).toEqual([])
    expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister)
    expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister)
  })

  it("#given state with 2 clients in different sessions #when disconnectSession removes one session #then process cleanup handlers remain registered", async () => {
    // given
    const state = createState()
    const signalIntCountBeforeRegister = process.listenerCount("SIGINT")
    const signalTermCountBeforeRegister = process.listenerCount("SIGTERM")

    state.clients.set("session-1:skill-1:server-1", createManagedClient("skill-1"))
    state.clients.set("session-2:skill-2:server-2", createManagedClient("skill-2"))
    registerProcessCleanup(state)

    // when
    await disconnectSession(state, "session-1")

    // then
    expect(state.clients.has("session-2:skill-2:server-2")).toBe(true)
    expect(state.cleanupRegistered).toBe(true)
    expect(state.cleanupHandlers).toHaveLength(2)
    expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister + 1)
    expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister + 1)
  })

  it("#given state with 2 clients in different sessions #when both sessions disconnected #then process cleanup handlers are unregistered", async () => {
    // given
    const state = createState()
    const signalIntCountBeforeRegister = process.listenerCount("SIGINT")
    const signalTermCountBeforeRegister = process.listenerCount("SIGTERM")

    state.clients.set("session-1:skill-1:server-1", createManagedClient("skill-1"))
    state.clients.set("session-2:skill-2:server-2", createManagedClient("skill-2"))
    registerProcessCleanup(state)

    // when
    await disconnectSession(state, "session-1")
    await disconnectSession(state, "session-2")

    // then
    expect(state.clients.size).toBe(0)
    expect(state.cleanupRegistered).toBe(false)
    expect(state.cleanupHandlers).toEqual([])
    expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister)
    expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister)
  })
})