voice-cat/clients/windows/VoiceCat.App/Audio/AudioSessionEnumerator.cs

using System.Diagnostics;
using System.Runtime.InteropServices;

namespace VoiceCat.App.Audio;

// ── Scope types (mirror macOS ScreenAudioScope / ScreenAudioSelection) ────────

public abstract record AppAudioScope;
// ExcludeSelf=true captures the whole system render mix minus VoiceCat's own process tree
// (kills the self-echo loop). Routed through the external-feed mixer as a single EXCLUDE
// capture; ExcludeSelf=false keeps the core's whole-device loopback path.
public sealed record EntireDesktop(bool ExcludeSelf = false) : AppAudioScope;
public sealed record OnlyApps(IReadOnlyList<int> Pids, IReadOnlyList<string> Names) : AppAudioScope;
public sealed record AllExceptApps(IReadOnlyList<int> Pids, IReadOnlyList<string> Names) : AppAudioScope;

// IsPlaying = a process in this app's executable group currently has an *active* audio
// session on some render endpoint. Purely informational for the picker — capture works
// on any PID regardless (process-loopback yields silence until the app plays).
public sealed record AudioAppInfo(int Pid, string DisplayName, bool IsPlaying);

// ── Enumerator ─────────────────────────────────────────────────────────────────

public static class AudioSessionEnumerator
{
    // Returns every app in the user's interactive session — windowed or not — so any of
    // them can be picked for capture even before it starts producing audio. Process
    // loopback (see ProcessLoopbackCapture) targets a PID and its child tree, so a silent
    // selection simply starts working the moment that app plays.
    //
    // Apps are deduped by executable (multiple PIDs of the same program collapse to one
    // row whose PID is the process-tree root). Session-0 services and VoiceCat itself are
    // excluded. Apps currently producing audio are flagged IsPlaying and sorted first.
    public static IReadOnlyList<AudioAppInfo> GetAudioApps()
    {
        int selfPid    = Environment.ProcessId;
        int sessionId  = SafeCurrentSessionId();
        var playingPids = GetActiveAudioPids();

        // Group by executable name; keep the best representative PID per group.
        var groups = new Dictionary<string, AppGroup>(StringComparer.OrdinalIgnoreCase);

        foreach (var proc in Process.GetProcesses())
        {
            try
            {
                if (proc.Id == selfPid) continue;
                if (proc.SessionId != sessionId) continue;   // drop session-0 services
                string name = proc.ProcessName;
                if (string.IsNullOrEmpty(name)) continue;

                bool hasWindow = proc.MainWindowHandle != IntPtr.Zero;
                bool playing   = playingPids.Contains(proc.Id);
                string title   = hasWindow ? SafeWindowTitle(proc) : "";
                string display = title.Length > 0 ? $"{name} — {title}" : name;

                if (groups.TryGetValue(name, out var g))
                {
                    g.AnyPlaying |= playing;
                    // Prefer a windowed PID (the process-tree root) as the capture target;
                    // among non-windowed, prefer a currently-playing PID.
                    bool better = (hasWindow && !g.HasWindow)
                               || (!g.HasWindow && !hasWindow && playing && !g.RepPlaying);
                    if (better)
                    {
                        g.Pid        = proc.Id;
                        g.Display    = display;
                        g.HasWindow  = hasWindow;
                        g.RepPlaying = playing;
                    }
                }
                else
                {
                    groups[name] = new AppGroup
                    {
                        Pid        = proc.Id,
                        Display    = display,
                        HasWindow  = hasWindow,
                        RepPlaying = playing,
                        AnyPlaying = playing,
                    };
                }
            }
            catch { /* protected/exited process — skip */ }
            finally { proc.Dispose(); }
        }

        return groups.Values
            .Select(g => new AudioAppInfo(g.Pid, g.Display, g.AnyPlaying))
            .OrderByDescending(a => a.IsPlaying)
            .ThenBy(a => a.DisplayName, StringComparer.OrdinalIgnoreCase)
            .ToList();
    }

    private sealed class AppGroup
    {
        public int    Pid;
        public string Display    = "";
        public bool   HasWindow;
        public bool   RepPlaying; // representative PID is playing
        public bool   AnyPlaying; // any PID in the group is playing
    }

    private static int SafeCurrentSessionId()
    {
        try { using var me = Process.GetCurrentProcess(); return me.SessionId; }
        catch { return 1; } // typical interactive session fallback
    }

    private static string SafeWindowTitle(Process proc)
    {
        try { return proc.MainWindowTitle; }
        catch { return ""; }
    }

    // ── WASAPI: PIDs with an active render session (any endpoint) ─────────────
    // Scans ALL active render endpoints, not just the default — an app routed to a
    // secondary device still counts as "playing now".

    private static HashSet<int> GetActiveAudioPids()
    {
        var pids = new HashSet<int>();
        IMMDeviceEnumerator? enumerator = null;
        IMMDeviceCollection? devices    = null;

        try
        {
            enumerator = (IMMDeviceEnumerator)Activator.CreateInstance(
                Type.GetTypeFromCLSID(new Guid("BCDE0395-E52F-467C-8E3D-C4579291692E"))!)!;

            if (enumerator.EnumAudioEndpoints(0 /*eRender*/, 0x1 /*DEVICE_STATE_ACTIVE*/,
                    out devices) < 0 || devices == null)
                return pids;

            devices.GetCount(out int devCount);
            for (int d = 0; d < devCount; d++)
                CollectActivePids(devices, d, pids);
        }
        catch { /* no audio device or WASAPI unavailable — ignore */ }
        finally
        {
            if (devices    != null) Marshal.ReleaseComObject(devices);
            if (enumerator != null) Marshal.ReleaseComObject(enumerator);
        }

        return pids;
    }

    private static void CollectActivePids(IMMDeviceCollection devices, int index, HashSet<int> pids)
    {
        IMMDevice? device                = null;
        IAudioSessionManager2? manager    = null;
        IAudioSessionEnumerator? sessions = null;

        try
        {
            if (devices.Item(index, out device) < 0 || device == null) return;

            var mgr2Iid = new Guid("77AA99A0-1BD6-484F-8BC7-2C654C9A9B6F");
            if (device.Activate(ref mgr2Iid, 0x17 /*CLSCTX_ALL*/, IntPtr.Zero, out object mgr) < 0)
                return;
            manager = (IAudioSessionManager2)mgr;

            if (manager.GetSessionEnumerator(out sessions) < 0 || sessions == null) return;
            sessions.GetCount(out int count);

            for (int i = 0; i < count; i++)
            {
                IAudioSessionControl? ctrl = null;
                try
                {
                    if (sessions.GetSession(i, out ctrl) < 0 || ctrl == null) continue;
                    ctrl.GetState(out int state);
                    if (state != 1 /*AudioSessionStateActive*/) continue;

                    var ctrl2 = (IAudioSessionControl2)ctrl;
                    ctrl2.GetProcessId(out uint pid);
                    if (pid != 0) pids.Add((int)pid);
                }
                catch { /* stale session */ }
                finally { if (ctrl != null) Marshal.ReleaseComObject(ctrl); }
            }
        }
        catch { /* device went away */ }
        finally
        {
            if (sessions != null) Marshal.ReleaseComObject(sessions);
            if (manager  != null) Marshal.ReleaseComObject(manager);
            if (device   != null) Marshal.ReleaseComObject(device);
        }
    }

    // ── COM interface declarations ─────────────────────────────────────────────

    [ComImport, Guid("A95664D2-9614-4F35-A746-DE8DB63617E6"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IMMDeviceEnumerator
    {
        [PreserveSig] int EnumAudioEndpoints(int dataFlow, int stateMask, out IMMDeviceCollection devices);
        [PreserveSig] int GetDefaultAudioEndpoint(int dataFlow, int role, out IMMDevice endpoint);
        [PreserveSig] int GetDevice([MarshalAs(UnmanagedType.LPWStr)] string id, out IMMDevice device);
        [PreserveSig] int RegisterEndpointNotificationCallback(IntPtr client);
        [PreserveSig] int UnregisterEndpointNotificationCallback(IntPtr client);
    }

    [ComImport, Guid("0BD7A1BE-7A1A-44DB-8397-CC5392387B5E"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IMMDeviceCollection
    {
        [PreserveSig] int GetCount(out int count);
        [PreserveSig] int Item(int index, out IMMDevice device);
    }

    [ComImport, Guid("D666063F-1587-4E43-81F1-B948E807363F"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IMMDevice
    {
        [PreserveSig] int Activate(ref Guid iid, int clsCtx, IntPtr activationParams,
            [MarshalAs(UnmanagedType.IUnknown)] out object ppInterface);
        [PreserveSig] int OpenPropertyStore(int stgmAccess, out IntPtr propStore);
        [PreserveSig] int GetId([MarshalAs(UnmanagedType.LPWStr)] out string id);
        [PreserveSig] int GetState(out int state);
    }

    [ComImport, Guid("77AA99A0-1BD6-484F-8BC7-2C654C9A9B6F"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IAudioSessionManager2
    {
        [PreserveSig] int GetAudioSessionControl(ref Guid audioSessionGuid, int streamFlags,
            out IAudioSessionControl session);
        [PreserveSig] int GetSimpleAudioVolume(ref Guid audioSessionGuid, int streamFlags,
            out IntPtr audioVolume);
        [PreserveSig] int GetSessionEnumerator(out IAudioSessionEnumerator sessionEnum);
        [PreserveSig] int RegisterSessionNotification(IntPtr notification);
        [PreserveSig] int UnregisterSessionNotification(IntPtr notification);
        [PreserveSig] int RegisterDuckNotification([MarshalAs(UnmanagedType.LPWStr)] string sessionID,
            IntPtr notification);
        [PreserveSig] int UnregisterDuckNotification(IntPtr notification);
    }

    [ComImport, Guid("E2F5BB11-0570-40CA-ACDD-3AA01277DEE8"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IAudioSessionEnumerator
    {
        [PreserveSig] int GetCount(out int sessionCount);
        [PreserveSig] int GetSession(int sessionIndex, out IAudioSessionControl session);
    }

    [ComImport, Guid("F4B1A599-7266-4319-A8CA-E70ACB11E8CD"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IAudioSessionControl
    {
        [PreserveSig] int GetState(out int state);
        [PreserveSig] int GetDisplayName([MarshalAs(UnmanagedType.LPWStr)] out string name);
        [PreserveSig] int SetDisplayName([MarshalAs(UnmanagedType.LPWStr)] string name,
            ref Guid eventContext);
        [PreserveSig] int GetIconPath([MarshalAs(UnmanagedType.LPWStr)] out string iconPath);
        [PreserveSig] int SetIconPath([MarshalAs(UnmanagedType.LPWStr)] string iconPath,
            ref Guid eventContext);
        [PreserveSig] int GetGroupingParam(out Guid groupingParam);
        [PreserveSig] int SetGroupingParam(ref Guid groupingParam, ref Guid eventContext);
        [PreserveSig] int RegisterAudioSessionNotification(IntPtr notification);
        [PreserveSig] int UnregisterAudioSessionNotification(IntPtr notification);
    }

    [ComImport, Guid("BFB7FF88-7239-4FC9-8FA2-07C950BE9C6D"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    internal interface IAudioSessionControl2 : IAudioSessionControl
    {
        [PreserveSig] int GetSessionIdentifier([MarshalAs(UnmanagedType.LPWStr)] out string id);
        [PreserveSig] int GetSessionInstanceIdentifier([MarshalAs(UnmanagedType.LPWStr)] out string id);
        [PreserveSig] int GetProcessId(out uint pid);
        [PreserveSig] int IsSystemSoundsSession();
        [PreserveSig] int SetDuckingPreference(bool optOut);
    }
}