Public/Stop-WorkloadSession.ps1

function Stop-WorkloadSession {
    <#
    .SYNOPSIS
        Terminates WorkloadsSessionHost processes based on configurable strategy.

    .DESCRIPTION
        Intelligently stops WorkloadsSessionHost.exe instances using one of three
        strategies:

        - MemoryThreshold (default): Kills processes exceeding the per-process
          memory threshold, OR when aggregate memory across all instances exceeds
          the aggregate threshold.
        - Timer: Kills processes that have been running longer than TimerSeconds.
          This mirrors the original WorkloadManager/SafeAutentic behavior.
        - All: Immediately terminates all instances.

        Supports -WhatIf and -Confirm for safe previewing. Logs all actions to
        the Windows Event Log via Write-WinslopFixLog.

    .PARAMETER Strategy
        The termination strategy to use. Defaults to MemoryThreshold.

    .PARAMETER PerProcessThresholdMB
        Maximum memory (in MB) allowed per individual process before termination.
        Only applies to the MemoryThreshold strategy. Default: 512.

    .PARAMETER AggregateThresholdMB
        Maximum total memory (in MB) allowed across all instances before the
        highest-memory processes are terminated. Only applies to the
        MemoryThreshold strategy. Default: 2048.

    .PARAMETER TimerSeconds
        Maximum age (in seconds) a process is allowed to run before termination.
        Only applies to the Timer strategy. Default: 60.

    .EXAMPLE
        Stop-WorkloadSession -WhatIf

        Preview which processes would be terminated using default memory thresholds.

    .EXAMPLE
        Stop-WorkloadSession -Strategy Timer -TimerSeconds 120

        Kill any WorkloadsSessionHost process older than 2 minutes.

    .EXAMPLE
        Stop-WorkloadSession -Strategy All -Confirm:$false

        Immediately terminate all instances without prompting.

    .OUTPUTS
        PSCustomObject[]
        Each object contains: PID, MemoryFreedMB, AgeSeconds, Reason, Result.

    .LINK
        https://github.com/DailenG/WinslopFix
    #>

    [CmdletBinding(SupportsShouldProcess, ConfirmImpact = 'High')]
    [OutputType([PSCustomObject[]])]
    param(
        [Parameter()]
        [ValidateSet('MemoryThreshold', 'Timer', 'All')]
        [string]$Strategy = 'MemoryThreshold',

        [Parameter()]
        [ValidateRange(64, 8192)]
        [int]$PerProcessThresholdMB = 512,

        [Parameter()]
        [ValidateRange(128, 16384)]
        [int]$AggregateThresholdMB = 2048,

        [Parameter()]
        [ValidateRange(10, 3600)]
        [int]$TimerSeconds = 60,

        [Parameter()]
        [string]$ProcessName = 'WorkloadsSessionHost'
    )

    process {
        $now       = Get-Date
        $processes = @(Get-WinslopFixProcess -ProcessName $ProcessName -ErrorAction SilentlyContinue)
        $results   = [System.Collections.Generic.List[PSCustomObject]]::new()

        if ($processes.Count -eq 0) {
            Write-Verbose "No '$ProcessName' processes found. Nothing to do."
            return
        }

        Write-Verbose "Found $($processes.Count) '$ProcessName' process(es). Strategy: $Strategy"

        # Build enriched process list
        $enriched = foreach ($proc in $processes) {
            [PSCustomObject]@{
                Process    = $proc
                PID        = $proc.Id
                MemoryMB   = [math]::Round($proc.PrivateMemorySize64 / 1MB, 1)
                AgeSeconds = [math]::Round(($now - $proc.StartTime).TotalSeconds, 0)
            }
        }

        $totalMemoryMB = ($enriched | Measure-Object -Property MemoryMB -Sum).Sum

        # Determine which processes to kill based on strategy
        $targets = switch ($Strategy) {
            'All' {
                $enriched | ForEach-Object {
                    $_ | Add-Member -NotePropertyName 'Reason' -NotePropertyValue 'Strategy: All' -PassThru
                }
            }

            'Timer' {
                $enriched | Where-Object { $_.AgeSeconds -ge $TimerSeconds } | ForEach-Object {
                    $reason = "Age $($_.AgeSeconds)s exceeds timer threshold of ${TimerSeconds}s"
                    $_ | Add-Member -NotePropertyName 'Reason' -NotePropertyValue $reason -PassThru
                }
            }

            'MemoryThreshold' {
                $candidates = [System.Collections.Generic.List[PSCustomObject]]::new()

                # Individual threshold violations
                foreach ($item in $enriched) {
                    if ($item.MemoryMB -ge $PerProcessThresholdMB) {
                        $reason = "Memory $($item.MemoryMB)MB exceeds per-process threshold of ${PerProcessThresholdMB}MB"
                        $item | Add-Member -NotePropertyName 'Reason' -NotePropertyValue $reason
                        $candidates.Add($item)
                    }
                }

                # Aggregate threshold — kill highest consumers first until under budget
                if ($totalMemoryMB -ge $AggregateThresholdMB) {
                    $sorted = $enriched | Sort-Object MemoryMB -Descending
                    $runningTotal = $totalMemoryMB

                    foreach ($item in $sorted) {
                        if ($runningTotal -lt $AggregateThresholdMB) { break }
                        if ($item.PID -notin $candidates.PID) {
                            $reason = "Aggregate memory ${totalMemoryMB}MB exceeds threshold of ${AggregateThresholdMB}MB (this process: $($item.MemoryMB)MB)"
                            $item | Add-Member -NotePropertyName 'Reason' -NotePropertyValue $reason
                            $candidates.Add($item)
                        }
                        $runningTotal -= $item.MemoryMB
                    }
                }

                $candidates
            }
        }

        if (-not $targets -or @($targets).Count -eq 0) {
            Write-Verbose 'No processes met termination criteria.'
            return
        }

        # Execute termination
        foreach ($target in $targets) {
            $description = "PID $($target.PID) ($($target.MemoryMB)MB, $($target.AgeSeconds)s old)"

            if ($PSCmdlet.ShouldProcess($description, 'Terminate WorkloadsSessionHost')) {
                try {
                    $target.Process.Kill()

                    $eventMsg = "Terminated $ProcessName PID $($target.PID). " +
                                "Memory: $($target.MemoryMB)MB. Age: $($target.AgeSeconds)s. " +
                                "Reason: $($target.Reason)"

                    $eventId = switch ($Strategy) {
                        'MemoryThreshold' { 3000 }
                        'Timer'           { 3001 }
                        'All'             { 3002 }
                    }

                    Write-WinslopFixLog -Message $eventMsg -EventId $eventId

                    $results.Add([PSCustomObject]@{
                        PSTypeName    = 'WinslopFix.KillResult'
                        PID           = $target.PID
                        MemoryFreedMB = $target.MemoryMB
                        AgeSeconds    = $target.AgeSeconds
                        Reason        = $target.Reason
                        Result        = 'Terminated'
                        Timestamp     = (Get-Date).ToString('o')
                    })
                }
                catch {
                    Write-Warning "Failed to terminate PID $($target.PID): $($_.Exception.Message)"

                    Write-WinslopFixLog -Message "Failed to terminate PID $($target.PID): $($_.Exception.Message)" `
                        -EventId 9001 -EntryType Error

                    $results.Add([PSCustomObject]@{
                        PSTypeName    = 'WinslopFix.KillResult'
                        PID           = $target.PID
                        MemoryFreedMB = 0
                        AgeSeconds    = $target.AgeSeconds
                        Reason        = $target.Reason
                        Result        = "Failed: $($_.Exception.Message)"
                        Timestamp     = (Get-Date).ToString('o')
                    })
                }
            }
        }

        return $results
    }
}