wx-win-agent/node-agent/WxAgent.Windows/WechatDoctor.cs

using System.Diagnostics;
using FlaUI.Core;
using FlaUI.Core.Definitions;
using FlaUI.UIA3;
using WxAgent.Core;

namespace WxAgent.Windows;

public sealed record ProcessDiagnostic(
    int Id,
    long WorkingSetBytes,
    int SessionId,
    string? Version,
    bool CanReadMemory,
    string? AccessError,
    int HandleCount,
    int ThreadCount);

public sealed record DoctorReport(
    string OperatingSystem,
    bool UserInteractive,
    int CurrentSessionId,
    IReadOnlyList<ProcessDiagnostic> Processes,
    bool WindowFound,
    IReadOnlyDictionary<string, bool> RequiredControls,
    int DataRootCount,
    IReadOnlyList<WxAgentErrorCode> Errors,
    bool InputDesktopAvailable = false);

public static class WechatDoctor
{
    public static DoctorReport Run(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        var errors = new List<WxAgentErrorCode>();
        var inputDesktopAvailable = WechatDesktop.IsInputAvailable();
        if (!inputDesktopAvailable) errors.Add(WxAgentErrorCode.SessionLocked);
        using var currentProcess = Process.GetCurrentProcess();
        var rawProcesses = Process.GetProcessesByName("Weixin");
        ProcessDiagnostic[] processes;
        try
        {
            var diagnostics = new List<ProcessDiagnostic>();
            foreach (var process in rawProcesses)
            {
                cancellationToken.ThrowIfCancellationRequested();
                try { diagnostics.Add(CreateProcessDiagnostic(process, cancellationToken)); }
                catch (InvalidOperationException) { /* Process exited during sampling. */ }
                catch (System.ComponentModel.Win32Exception) { /* Process is no longer accessible. */ }
            }
            processes = diagnostics.OrderByDescending(process => process.WorkingSetBytes).ToArray();
        }
        finally
        {
            foreach (var process in rawProcesses) process.Dispose();
        }

        if (processes.Length == 0)
        {
            errors.Add(WxAgentErrorCode.WechatNotRunning);
        }
        else if (processes.All(process => !process.CanReadMemory))
        {
            errors.Add(WxAgentErrorCode.ProcessAccessDenied);
        }

        var controls = WechatLocators.RequiredAutomationIds.ToDictionary(id => id, _ => false, StringComparer.Ordinal);
        var windowFound = false;
        try
        {
            using var automation = new UIA3Automation();
            var window = FindWechatWindow(automation);
            windowFound = window is not null;
            if (window is not null)
            {
                foreach (var id in controls.Keys.ToArray())
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    controls[id] = window.FindFirstDescendant(condition => condition.ByAutomationId(id)) is not null;
                }
            }
        }
        catch (OperationCanceledException)
        {
            throw;
        }
        catch
        {
            windowFound = false;
        }

        if (!windowFound && processes.Length > 0)
        {
            errors.Add(processes.All(process => process.SessionId != currentProcess.SessionId)
                ? WxAgentErrorCode.SessionLocked
                : WxAgentErrorCode.WindowNotFound);
        }
        else if (windowFound && WechatLocators.ClassifyFoundWindow(controls) is { } uiError)
        {
            errors.Add(uiError);
        }

        var dataRootCount = WechatDatabaseDiscovery.FindAccountRoots(cancellationToken: cancellationToken).Count;
        if (dataRootCount == 0)
        {
            errors.Add(WxAgentErrorCode.DataRootNotFound);
        }

        return new DoctorReport(
            Environment.OSVersion.VersionString,
            Environment.UserInteractive,
            currentProcess.SessionId,
            processes,
            windowFound,
            controls,
            dataRootCount,
            errors.Distinct().ToArray(),
            inputDesktopAvailable);
    }

    private static readonly AsyncLocal<int?> TargetProcess = new();
    private static readonly AsyncLocal<long?> TargetWindow = new();

    internal static int? TargetProcessId => TargetProcess.Value;

    internal static IDisposable UseWindow(int processId, long windowHandle) => UseTarget(processId, windowHandle);

    private static IDisposable UseTarget(int? processId, long? windowHandle)
    {
        var previousProcess = TargetProcess.Value;
        var previousWindow = TargetWindow.Value;
        TargetProcess.Value = processId;
        TargetWindow.Value = windowHandle;
        return new TargetScope(previousProcess, previousWindow);
    }

    private sealed class TargetScope(int? previousProcess, long? previousWindow) : IDisposable
    {
        public void Dispose()
        {
            TargetProcess.Value = previousProcess;
            TargetWindow.Value = previousWindow;
        }
    }

    internal static FlaUI.Core.AutomationElements.AutomationElement? FindWechatWindow(AutomationBase automation)
    {
        FlaUI.Core.AutomationElements.AutomationElement? fallback = null;
        var processId = TargetProcess.Value;
        var windowHandle = TargetWindow.Value;
        var mainWindows = WechatNativeWindow.Enumerate().Where(window =>
            (processId is null || window.ProcessId == processId) &&
            (windowHandle is null || window.Handle == windowHandle) &&
            window.Visible && !window.Minimized && WechatLocators.IsNativeMainWindow(window.ClassName, window.Title));
        foreach (var native in mainWindows)
        {
            try
            {
                var candidate = automation.FromHandle((nint)native.Handle);
                if (candidate.FindFirstDescendant(condition => condition.ByAutomationId(WechatLocators.MainView)) is not null)
                    return candidate;
                fallback ??= candidate;
            }
            catch
            {
                // Ignore transient or closing subwindows.
            }
        }
        return fallback;
    }

    internal static FlaUI.Core.AutomationElements.AutomationElement? FindNamedWechatWindow(
        AutomationBase automation, string title, CancellationToken cancellationToken)
    {
        var matches = new List<FlaUI.Core.AutomationElements.AutomationElement>();
        foreach (var native in WechatNativeWindow.Enumerate().Where(window => window.Visible))
        {
            cancellationToken.ThrowIfCancellationRequested();
            try
            {
                var candidate = automation.FromHandle((nint)native.Handle);
                if (string.Equals(candidate.Name, title, StringComparison.Ordinal)) matches.Add(candidate);
            }
            catch { /* A window may disappear during enumeration. */ }
        }
        if (matches.Count > 1)
            throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "Multiple WeChat windows have the requested title; refusing an ambiguous target.");
        return matches.SingleOrDefault();
    }

    private static ProcessDiagnostic CreateProcessDiagnostic(Process process, CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        string? version = null;
        try
        {
            version = process.MainModule?.FileVersionInfo.FileVersion;
        }
        catch
        {
        }

        var access = ProcessMemoryScanner.Probe(process.Id);
        var handleCount = 0;
        var threadCount = 0;
        try { handleCount = process.HandleCount; } catch { }
        try { threadCount = process.Threads.Count; } catch { }
        return new ProcessDiagnostic(process.Id, process.WorkingSet64, process.SessionId, version, access.Success, access.Error, handleCount, threadCount);
    }
}