browser_mcp/pool.go

package main

import (
	"strings"
	"sync"

	"fn-registry/functions/browser"
)

// connPool reusa conexiones CDP entre invocaciones de tools. Clave = puerto CDP.
// Una conexión = una sesión viva a una tab "page". Mantenerla evita pagar el
// handshake WebSocket en cada tool y preserva estado (event handlers, contexto).
//
// El pool también registra el PID del Chrome que el MCP LANZÓ por puerto
// (mapa `pids`). Sin ese PID, cerrar la conexión solo suelta el WebSocket y deja
// el proceso chromium huérfano (~789 MiB RSS cada uno) — ese era el leak de RAM.
// Con el PID registrado, `drop`/`closeAll` matan el grupo de proceso completo.
// Un puerto SIN pid registrado (p.ej. el navegador diario del usuario en 9222,
// que el MCP no lanzó) nunca se mata: solo se suelta el WebSocket.
type connPool struct {
	mu         sync.Mutex
	conns      map[int]*browser.CDPConn
	pids       map[int]int                // puerto -> PID del Chrome lanzado por el MCP (solo los SUYOS)
	cancels    map[int]func()             // cancels de handlers persistentes (handle_dialog)
	dialogLogs map[int]*browser.DialogLog // log de diálogos auto-respondidos por puerto
}

func newConnPool() *connPool {
	return &connPool{
		conns:      map[int]*browser.CDPConn{},
		pids:       map[int]int{},
		cancels:    map[int]func(){},
		dialogLogs: map[int]*browser.DialogLog{},
	}
}

func (p *connPool) get(port int) (*browser.CDPConn, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if c, ok := p.conns[port]; ok && c != nil {
		return c, nil
	}
	c, err := browser.CdpConnect(port)
	if err != nil {
		return nil, err
	}
	p.conns[port] = c
	return c, nil
}

// setPID registra el PID del Chrome que el MCP lanzó en este puerto. A partir de
// aquí drop/closeAll podrán matar ese proceso (es nuestro).
func (p *connPool) setPID(port, pid int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.pids[port] = pid
}

// getPID devuelve el PID registrado para el puerto (y si existe). pid<=0 o
// ausente significa que el MCP no lanzó ningún Chrome propio en ese puerto.
func (p *connPool) getPID(port int) (int, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pid, ok := p.pids[port]
	return pid, ok
}

// clearPID olvida el PID de un puerto sin matar nada (p.ej. el proceso ya murió).
func (p *connPool) clearPID(port int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	delete(p.pids, port)
}

// launchedCount devuelve cuántos Chrome propios tiene vivos el MCP (uno por
// puerto registrado). Alimenta el tope de instancias en handleLaunch.
func (p *connPool) launchedCount() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.pids)
}

// releaseConn cierra SOLO el WebSocket pooled del puerto (pid=0, no mata Chrome)
// y lo borra del mapa, PRESERVANDO el PID registrado. Cancela el handler de
// diálogo de esa sesión (está atado a la conexión que se suelta). Lo usan el
// retry de withConn y connectTarget: necesitan reconectar al MISMO Chrome, no
// matarlo.
func (p *connPool) releaseConn(port int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if cancel, ok := p.cancels[port]; ok && cancel != nil {
		cancel()
		delete(p.cancels, port)
	}
	delete(p.dialogLogs, port)
	if c, ok := p.conns[port]; ok && c != nil {
		// pid=0: solo soltar el WebSocket. El Chrome sigue vivo para reconectar.
		_ = browser.CdpClose(c, 0)
		delete(p.conns, port)
	}
}

// drop cierra la sesión del puerto Y mata el Chrome SI lo lanzó el MCP (pid
// registrado). Para un Chrome externo (sin pid registrado, p.ej. el navegador
// diario en 9222) pasa pid=0 a CdpClose: solo cierra el WebSocket, NUNCA mata el
// navegador del usuario. Limpia todas las entradas del puerto.
func (p *connPool) drop(port int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if cancel, ok := p.cancels[port]; ok && cancel != nil {
		cancel()
		delete(p.cancels, port)
	}
	delete(p.dialogLogs, port)

	pid := p.pids[port] // 0 si el MCP no lanzó este Chrome
	c := p.conns[port]
	// CdpClose mata el grupo de proceso completo SOLO si pid>0 (Setpgid=true en
	// ChromeLaunch). Con c!=nil cierra además el WebSocket; con pid<=0 no toca el
	// proceso.
	_ = browser.CdpClose(c, pid)
	delete(p.conns, port)
	delete(p.pids, port)
}

// connectTarget descarta la conexión actual del puerto y reconecta a un target
// determinista (por id o substring de URL). Asegura que el agente opera sobre una
// pestaña conocida y no sobre "la primera al azar". Usa releaseConn (NO drop):
// cambiar de pestaña no debe matar el Chrome, es el mismo navegador.
func (p *connPool) connectTarget(port int, match string) (*browser.CDPConn, error) {
	p.releaseConn(port)
	c, err := browser.CdpConnectTarget("localhost", port, match)
	if err != nil {
		return nil, err
	}
	p.mu.Lock()
	p.conns[port] = c
	p.mu.Unlock()
	return c, nil
}

// setDialog guarda el cancel y el DialogLog del auto-handler de diálogos del
// puerto. Si ya había uno armado, lo cancela primero.
func (p *connPool) setDialog(port int, cancel func(), dlog *browser.DialogLog) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if old := p.cancels[port]; old != nil {
		old()
	}
	p.cancels[port] = cancel
	p.dialogLogs[port] = dlog
}

// dialogSnapshot devuelve el estado del log de diálogos del puerto (0,"","" si
// no hay handler armado).
func (p *connPool) dialogSnapshot(port int) (int, string, string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if dl := p.dialogLogs[port]; dl != nil {
		return dl.Snapshot()
	}
	return 0, "", ""
}

// closeAll cierra todas las conexiones y mata TODOS los Chrome que el MCP lanzó
// (pid registrado). Se llama con defer en main() (cierre por EOF de stdio) y
// desde el handler de señales (SIGTERM/SIGINT). Idempotente: vacía los mapas, así
// que una segunda llamada no hace nada. Un Chrome externo (sin pid) no se mata.
func (p *connPool) closeAll() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for port, c := range p.conns {
		if cancel := p.cancels[port]; cancel != nil {
			cancel()
		}
		_ = browser.CdpClose(c, p.pids[port]) // mata nuestro Chrome; pid=0 para externos
		delete(p.pids, port)                  // marcado como ya cerrado
	}
	// Matar también los Chrome propios cuya conexión ya fue soltada (releaseConn
	// preserva el pid pero borra la conn): pid registrado sin conn viva.
	for port, pid := range p.pids {
		if pid > 0 {
			_ = browser.CdpClose(nil, pid)
		}
		_ = port
	}
	p.conns = map[int]*browser.CDPConn{}
	p.pids = map[int]int{}
	p.cancels = map[int]func(){}
	p.dialogLogs = map[int]*browser.DialogLog{}
}

// isConnErr reconoce errores de conexión CDP muerta para reintentar UNA vez.
func isConnErr(err error) bool {
	s := err.Error()
	return strings.Contains(s, "connection close") || strings.Contains(s, "broken pipe") ||
		strings.Contains(s, "use of closed") || strings.Contains(s, "ws read") || strings.Contains(s, "EOF")
}