mizan/frontends/mizan-webview-channels/src/index.ts

/**
 * @mizan/webview-channels — RPC + pub/sub over a VSCode webview's
 * postMessage channel. Mirror of mizan-react/src/channels (WebSocket-based)
 * adapted for the in-process bidirectional channel a webview shares with
 * its extension host.
 *
 * The webview process can only acquire the VSCode API once, so this module
 * and `@mizan/webview-transport` coexist by sharing the singleton:
 * whichever loads first acquires the API. Both packages post their own
 * envelope kinds and filter incoming messages by the `kind` field.
 *
 * Envelope shapes (this side ↔ extension-host):
 *
 *   webview → host:
 *     { kind: 'rpc',         id, fn, args }
 *     { kind: 'subscribe',   id, channel, params? }
 *     { kind: 'unsubscribe', id }
 *
 *   host → webview:
 *     { kind: 'rpc-response', id, ok: true,  data }
 *     { kind: 'rpc-response', id, ok: false, error: { code, message, details? } }
 *     { kind: 'event',        subscription, payload }
 *     { kind: 'close',        subscription, reason? }
 *
 * Long-running ops (e.g. kata's claude-p shell-out) live as channel
 * subscriptions — the host streams events as work progresses, then closes
 * the subscription on completion. RPC is reserved for short request/reply.
 */

declare global {
    function acquireVsCodeApi(): VsCodeApi
}

interface VsCodeApi {
    postMessage(msg: unknown): void
    setState(state: unknown): void
    getState(): unknown
}

let _api: VsCodeApi | null = null

function api(): VsCodeApi {
    if (_api === null) {
        if (typeof acquireVsCodeApi !== 'function') {
            throw new Error(
                '@mizan/webview-channels: acquireVsCodeApi is not defined. ' +
                'This module runs inside a VSCode webview only.',
            )
        }
        _api = acquireVsCodeApi()
    }
    return _api
}

// ─── Errors ────────────────────────────────────────────────────────────────

export class WebviewRPCError extends Error {
    code: string
    details?: Record<string, unknown>

    constructor(code: string, message: string, details?: Record<string, unknown>) {
        super(message)
        this.name = 'WebviewRPCError'
        this.code = code
        this.details = details
    }
}

// ─── Envelope types ────────────────────────────────────────────────────────

type RPCEnvelope = { kind: 'rpc'; id: string; fn: string; args: unknown }
type SubscribeEnvelope = { kind: 'subscribe'; id: string; channel: string; params?: Record<string, unknown> }
type UnsubscribeEnvelope = { kind: 'unsubscribe'; id: string }

type RPCResponseOk = { kind: 'rpc-response'; id: string; ok: true; data: unknown }
type RPCResponseErr = {
    kind: 'rpc-response'
    id: string
    ok: false
    error: { code: string; message: string; details?: Record<string, unknown> }
}
type EventEnvelope = { kind: 'event'; subscription: string; payload: unknown }
type CloseEnvelope = { kind: 'close'; subscription: string; reason?: string }

type Incoming = RPCResponseOk | RPCResponseErr | EventEnvelope | CloseEnvelope

// ─── Subscription handle ───────────────────────────────────────────────────

export interface Subscription {
    readonly id: string
    readonly channel: string
    readonly params: Record<string, unknown> | undefined
    /** Called for every event the host pushes on this subscription. */
    onEvent(handler: (payload: unknown) => void): () => void
    /** Called once when the host closes the subscription. */
    onClose(handler: (reason?: string) => void): () => void
    /** Send an unsubscribe envelope to the host; local handlers stop firing. */
    unsubscribe(): void
    /** True until the host has closed the subscription or the consumer has unsubscribed. */
    readonly active: boolean
}

// ─── Connection ────────────────────────────────────────────────────────────

export class WebviewChannelConnection {
    private pendingRPCs = new Map<string, {
        resolve: (data: unknown) => void
        reject: (err: WebviewRPCError) => void
    }>()
    private subscriptions = new Map<string, SubscriptionImpl>()
    private installed = false
    private counter = 0

    private install(): void {
        if (this.installed) return
        this.installed = true
        window.addEventListener('message', (e: MessageEvent) => this.handle(e.data))
    }

    private nextId(prefix: string): string {
        return `${prefix}_${++this.counter}_${Date.now()}`
    }

    private handle(msg: unknown): void {
        if (!msg || typeof msg !== 'object') return
        const m = msg as Incoming
        switch (m.kind) {
            case 'rpc-response': {
                const pen = this.pendingRPCs.get(m.id)
                if (!pen) return
                this.pendingRPCs.delete(m.id)
                if (m.ok) {
                    pen.resolve(m.data)
                } else {
                    pen.reject(new WebviewRPCError(
                        m.error.code,
                        m.error.message,
                        m.error.details,
                    ))
                }
                return
            }
            case 'event': {
                const sub = this.subscriptions.get(m.subscription)
                if (sub) sub.dispatchEvent(m.payload)
                return
            }
            case 'close': {
                const sub = this.subscriptions.get(m.subscription)
                if (sub) {
                    sub.dispatchClose(m.reason)
                    this.subscriptions.delete(m.subscription)
                }
                return
            }
        }
    }

    /**
     * Call a host-registered RPC function. Resolves with the function's
     * return value or throws `WebviewRPCError` on failure.
     */
    rpc<TInput, TOutput>(fn: string, args: TInput): Promise<TOutput> {
        this.install()
        return new Promise<TOutput>((resolve, reject) => {
            const id = this.nextId('rpc')
            this.pendingRPCs.set(id, {
                resolve: resolve as (data: unknown) => void,
                reject,
            })
            const env: RPCEnvelope = { kind: 'rpc', id, fn, args }
            api().postMessage(env)
        })
    }

    /**
     * Open a long-running subscription on a host-registered channel.
     * The returned Subscription handle receives events until the host
     * closes it or the consumer calls `unsubscribe()`.
     */
    subscribe(channel: string, params?: Record<string, unknown>): Subscription {
        this.install()
        const id = this.nextId('sub')
        const sub = new SubscriptionImpl(id, channel, params, () => {
            this.subscriptions.delete(id)
            const env: UnsubscribeEnvelope = { kind: 'unsubscribe', id }
            api().postMessage(env)
        })
        this.subscriptions.set(id, sub)
        const env: SubscribeEnvelope = { kind: 'subscribe', id, channel, params }
        api().postMessage(env)
        return sub
    }
}

class SubscriptionImpl implements Subscription {
    private eventHandlers = new Set<(payload: unknown) => void>()
    private closeHandlers = new Set<(reason?: string) => void>()
    private _active = true

    constructor(
        public readonly id: string,
        public readonly channel: string,
        public readonly params: Record<string, unknown> | undefined,
        private readonly onLocalUnsubscribe: () => void,
    ) {}

    get active(): boolean {
        return this._active
    }

    onEvent(handler: (payload: unknown) => void): () => void {
        this.eventHandlers.add(handler)
        return () => this.eventHandlers.delete(handler)
    }

    onClose(handler: (reason?: string) => void): () => void {
        this.closeHandlers.add(handler)
        return () => this.closeHandlers.delete(handler)
    }

    unsubscribe(): void {
        if (!this._active) return
        this._active = false
        this.onLocalUnsubscribe()
        this.closeHandlers.forEach(h => h('unsubscribed'))
        this.eventHandlers.clear()
        this.closeHandlers.clear()
    }

    dispatchEvent(payload: unknown): void {
        if (!this._active) return
        this.eventHandlers.forEach(h => h(payload))
    }

    dispatchClose(reason?: string): void {
        if (!this._active) return
        this._active = false
        this.closeHandlers.forEach(h => h(reason))
        this.eventHandlers.clear()
        this.closeHandlers.clear()
    }
}

// ─── Singleton ─────────────────────────────────────────────────────────────

let _default: WebviewChannelConnection | null = null

export function getDefaultChannelConnection(): WebviewChannelConnection {
    if (_default === null) _default = new WebviewChannelConnection()
    return _default
}