Lumi/companion/plugins/Lumi.Companion.Transcription/ObsBridgePipe.cs

using System.Buffers.Binary;
using System.IO.Pipes;
using System.Collections.Concurrent;
using System.Text.Json;
using Lumi.Companion.Protocol;

namespace Lumi.Companion.Transcription;

public sealed class ObsBridgePipe : IAsyncDisposable
{
    private static readonly HashSet<string> AllowedTypes = ["hello", "source_list", "source_state", "selection_state", "obs_state", "health", "stream_test_service_state", "stream_test_metrics"];
    private readonly string _pipeName;
    private readonly Func<JsonElement, Task> _onMessage;
    private readonly Func<ReadOnlyMemory<byte>, Task> _onAudio;
    private readonly CancellationTokenSource _lifetime = new();
    private readonly SemaphoreSlim _writeLock = new(1, 1);
    private readonly ConcurrentDictionary<string, TaskCompletionSource<JsonElement>> _requests = new();
    private Task? _loop;
    private Stream? _connection;

    public ObsBridgePipe(string userSidHash, Func<JsonElement, Task> onMessage, Func<ReadOnlyMemory<byte>, Task> onAudio)
    {
        _pipeName = $"Lumi.Companion.ObsBridge.v1.{userSidHash}";
        _onMessage = onMessage; _onAudio = onAudio;
    }
    public event Action<bool>? ConnectionChanged;
    public void Start() => _loop ??= Task.Run(() => ListenAsync(_lifetime.Token));
    private async Task ListenAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            await using var pipe = new NamedPipeServerStream(_pipeName, PipeDirection.InOut, 1, PipeTransmissionMode.Byte,
                PipeOptions.Asynchronous | PipeOptions.CurrentUserOnly, 65536, 65536);
            await pipe.WaitForConnectionAsync(cancellationToken);
            _connection = pipe;
            ConnectionChanged?.Invoke(true);
            try { await ReadConnectionAsync(pipe, cancellationToken); }
            catch (Exception) when (!cancellationToken.IsCancellationRequested) { /* reconnect without taking down the shell */ }
            finally { _connection = null; ConnectionChanged?.Invoke(false); }
        }
    }
    private async Task ReadConnectionAsync(Stream stream, CancellationToken cancellationToken)
    {
        var lengthBytes = new byte[4];
        while (!cancellationToken.IsCancellationRequested)
        {
            await ReadExactlyAsync(stream, lengthBytes, cancellationToken);
            var length = BinaryPrimitives.ReadUInt32LittleEndian(lengthBytes);
            if (length is 0 or > ProtocolV1.MaxJsonBytes) throw new InvalidDataException("OBS bridge message exceeds its limit.");
            var message = new byte[(int)length];
            await ReadExactlyAsync(stream, message, cancellationToken);
            if (length >= 4 && message.AsSpan(0, 4).SequenceEqual("LACP"u8))
            {
                if (length > ProtocolV1.AudioHeaderBytes + ProtocolV1.MaxAudioPayloadBytes) throw new InvalidDataException("OBS audio frame exceeds its limit.");
                await _onAudio(message); continue;
            }
            using var json = JsonDocument.Parse(message);
            var type = json.RootElement.TryGetProperty("type", out var property) ? property.GetString() : null;
            if (type is null || !AllowedTypes.Contains(type)) throw new InvalidDataException("OBS bridge message type is not allowed.");
            var clone = json.RootElement.Clone();
            if (clone.TryGetProperty("request_id", out var requestId) && requestId.ValueKind == JsonValueKind.String &&
                _requests.TryRemove(requestId.GetString()!, out var pending)) pending.TrySetResult(clone);
            await _onMessage(clone);
        }
    }
    public static async Task WriteAsync(Stream stream, object message, CancellationToken cancellationToken)
    {
        var body = JsonSerializer.SerializeToUtf8Bytes(message, ProtocolV1.JsonOptions);
        if (body.Length > ProtocolV1.MaxJsonBytes) throw new InvalidDataException("Companion IPC message exceeds its limit.");
        var prefix = new byte[4]; BinaryPrimitives.WriteUInt32LittleEndian(prefix, (uint)body.Length);
        await stream.WriteAsync(prefix, cancellationToken); await stream.WriteAsync(body, cancellationToken); await stream.FlushAsync(cancellationToken);
    }
    public async Task<bool> SendAsync(object message, CancellationToken cancellationToken = default)
    {
        var connection = _connection;
        if (connection is null) return false;
        await _writeLock.WaitAsync(cancellationToken);
        try
        {
            if (!ReferenceEquals(connection, _connection)) return false;
            await WriteAsync(connection, message, cancellationToken);
            return true;
        }
        catch (IOException) { return false; }
        finally { _writeLock.Release(); }
    }
    public async Task<JsonElement> RequestAsync(string type, object payload, TimeSpan timeout, CancellationToken cancellationToken = default)
    {
        var requestId = Guid.NewGuid().ToString();
        var completion = new TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously);
        if (!_requests.TryAdd(requestId, completion)) throw new InvalidOperationException("OBS request identity could not be reserved.");
        using var timeoutSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutSource.CancelAfter(timeout);
        using var registration = timeoutSource.Token.Register(() => completion.TrySetCanceled(timeoutSource.Token));
        try
        {
            if (!await SendAsync(new { type, request_id = requestId, payload }, timeoutSource.Token))
                throw new InvalidOperationException("OBS is not connected to Lumi Companion.");
            return await completion.Task;
        }
        finally
        {
            _requests.TryRemove(requestId, out _);
        }
    }
    private static async Task ReadExactlyAsync(Stream stream, Memory<byte> buffer, CancellationToken cancellationToken)
    {
        var read = 0;
        while (read < buffer.Length)
        {
            var count = await stream.ReadAsync(buffer[read..], cancellationToken);
            if (count == 0) throw new EndOfStreamException();
            read += count;
        }
    }
    public async ValueTask DisposeAsync() { _lifetime.Cancel(); foreach (var request in _requests.Values) request.TrySetCanceled(); _requests.Clear(); if (_loop is not null) try { await _loop; } catch (OperationCanceledException) { } _writeLock.Dispose(); _lifetime.Dispose(); }
}