Dev

Services/OrchestratorService.cs

@@ -607,15 +607,15 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task,
                             CheckRunCompletion(run);
                         }
                     }
-                    await _store.UpsertTaskAsync(run.Name, task);
-                    await _store.UpsertRunAsync(run);
+                    PersistTaskAndRunAsync(run, task);
                     return;
                 }
 
                 lock (_lock)
                 {
                     task.Status = "Running";
                 }
+                // Pre-script "Running" write is awaited — it's the durability marker for crash recovery.
                 await _store.UpsertTaskAsync(run.Name, task);
 
                 try
@@ -646,8 +646,11 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task,
                         task.Parameters = null!;
                         CheckRunCompletion(run);
                     }
-                    await _store.UpsertTaskAsync(run.Name, task);
-                    await _store.UpsertRunAsync(run);
+                    // Post-script writes are fire-and-forget so the JobManager slot releases
+                    // immediately and the dispatch loop can hand the worker to the next task.
+                    // Crash recovery still works: the next startup re-reads task state from the
+                    // table and re-runs anything not marked Completed (idempotent).
+                    PersistTaskAndRunAsync(run, task);
 
                     _logger.LogDebug("[Scheduler] Task completed: {TaskId}", task.Id);
                 }
@@ -668,15 +671,32 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task,
                         task.Parameters = null!;
                         CheckRunCompletion(run);
                     }
-                    await _store.UpsertTaskAsync(run.Name, task);
-                    await _store.UpsertRunAsync(run);
+                    PersistTaskAndRunAsync(run, task);
                     _logger.LogError(ex, "[Scheduler] Task failed: {TaskId}", task.Id);
                     throw; // Let JobManager also track the failure
                 }
             }
         );
     }
 
+    /// <summary>
+    /// Fire-and-forget persistence of task + run state. Callers do not await this — it lets the
+    /// JobManager slot release immediately so the dispatch loop can hand the worker to the next
+    /// task. Errors are logged; on host crash, ResumeInterruptedRunsAsync re-derives state from
+    /// whatever made it to the table (writes are idempotent).
+    /// </summary>
+    private void PersistTaskAndRunAsync(OrchestratorRun run, OrchestratorTaskItem task)
+    {
+        _ = Task.Run(async () =>
+        {
+            try { await _store.UpsertTaskAsync(run.Name, task); }
+            catch (Exception ex) { _logger.LogWarning(ex, "[Scheduler] Background UpsertTask failed for {Run}/{Task}", run.Name, task.Id); }
+
+            try { await _store.UpsertRunAsync(run); }
+            catch (Exception ex) { _logger.LogWarning(ex, "[Scheduler] Background UpsertRun failed for {Run}", run.Name); }
+        });
+    }
+
     private void LogRunStatus(OrchestratorRun run)
     {
         var elapsed = DateTime.UtcNow - run.StartedUtc;

Services/PowerShellRunnerService.cs

@@ -293,6 +293,26 @@ private async Task<ScriptResult> ExecuteHttpScriptInternal(string route, Hashtab
         }
     }
 
+    /// <summary>
+    /// Drain pending orchestrator/queue triggers off the calling thread. Bridges are thread-safe
+    /// concurrent queues, so multiple in-flight drain calls are fine — each TryDequeue serialises.
+    /// </summary>
+    private void DrainBridgesInBackground()
+    {
+        _ = Task.Run(async () =>
+        {
+            try
+            {
+                await OrchestratorBridge.DrainPendingAsync();
+                QueueBridge.DrainPending();
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "[Scheduler] Background bridge drain failed");
+            }
+        });
+    }
+
     /// <summary>
     /// Execute a script by function name (for scheduler / orchestrator). No HTTP context needed.
     /// Runs on the background pool.
@@ -380,10 +400,6 @@ public async Task ExecuteScript(string functionName, Dictionary<string, object>?
             sw.Stop();
             _logger.LogInformation("[Scheduler] {InvocationId} {Function} completed {Ms}ms",
                 invocation.Id, functionName, sw.ElapsedMilliseconds);
-
-            // Process any orchestrator/queue triggers queued during execution
-            await OrchestratorBridge.DrainPendingAsync();
-            QueueBridge.DrainPending();
         }
         catch (OperationCanceledException) when (sw.ElapsedMilliseconds > 0)
         {
@@ -410,6 +426,11 @@ public async Task ExecuteScript(string functionName, Dictionary<string, object>?
             if (onVerbose != null) worker.Streams.Verbose.DataAdded -= onVerbose;
             _pool.Reclaim(worker, isHttp: false, faulted: exceptionOccurred);
         }
+
+        // Worker has been returned to the pool — drain any orchestrator/queue triggers
+        // the script enqueued in the background so the next job can grab the worker now
+        // instead of waiting for child-run table writes.
+        DrainBridgesInBackground();
     }
 
     /// <summary>
@@ -483,14 +504,14 @@ public async Task<string> ExecuteScriptWithOutput(string functionName, Dictionar
                 : null;
             var results = await worker.InvokeAsync(resolvedName, psParams, cts?.Token ?? default);
 
-            // Process any orchestrator/queue triggers queued during execution
-            await OrchestratorBridge.DrainPendingAsync();
-            QueueBridge.DrainPending();
-
             sw.Stop();
             _logger.LogInformation("[Planner] {InvocationId} {Function} completed {Ms}ms",
                 invocation.Id, functionName, sw.ElapsedMilliseconds);
-            return string.Join("\n", (results ?? new Collection<PSObject>()).Select(r => r?.ToString() ?? ""));
+            var output = string.Join("\n", (results ?? new Collection<PSObject>()).Select(r => r?.ToString() ?? ""));
+            // Drain triggered child orchestrators/queue commands after returning — they should
+            // not block the planner's caller. (Note: Reclaim still happens in finally below.)
+            DrainBridgesInBackground();
+            return output;
         }
         catch (OperationCanceledException) when (sw.ElapsedMilliseconds > 0)
         {

Services/PowerShellWorker.cs

@@ -105,16 +105,39 @@ class HttpResponseContext {
 ");
         }
 
-        // Load shared assemblies from config
+        // Load shared assemblies from config.
+        // ISS-level registration happens in PowerShellWorkerPool.RegisterSharedAssemblies(); this
+        // runtime LoadFile is a defence-in-depth fallback that surfaces any load failure to the log
+        // (RunScript silently swallows streams, so we run a labelled invocation here instead).
         foreach (var asmRelPath in settings.Worker.SharedAssemblies)
         {
+            if (string.IsNullOrWhiteSpace(asmRelPath)) continue;
             var asmPath = Path.Combine(apiBasePath, asmRelPath).Replace("\\", "/");
-            RunScript($@"
-if (Test-Path '{asmPath}') {{
-    if (-not ([System.AppDomain]::CurrentDomain.GetAssemblies().Location -contains '{asmPath}')) {{
+            var asmLabel = Path.GetFileNameWithoutExtension(asmPath);
+            try
+            {
+                _pwsh.AddScript($@"
+if (Test-Path -LiteralPath '{asmPath}') {{
+    if (-not ([System.AppDomain]::CurrentDomain.GetAssemblies() | Where-Object {{ $_.Location -ieq '{asmPath}' }})) {{
         [void][Reflection.Assembly]::LoadFile('{asmPath}')
     }}
-}}");
+    [System.AppDomain]::CurrentDomain.GetAssemblies() | Where-Object {{ $_.Location -ieq '{asmPath}' }} | Select-Object -First 1 -ExpandProperty FullName
+}} else {{
+    Write-Error ""SharedAssembly not found: {asmPath}""
+}}").Invoke();
+
+                foreach (var err in _pwsh.Streams.Error)
+                    _logger.LogError("Worker{Id}: SharedAssembly '{Label}' load error: {Error}", Id, asmLabel, err.ToString());
+                foreach (var warn in _pwsh.Streams.Warning)
+                    _logger.LogWarning("Worker{Id}: SharedAssembly '{Label}' warning: {Message}", Id, asmLabel, warn.Message);
+
+                _logger.LogDebug("Worker{Id}: SharedAssembly '{Label}' available at {Path}", Id, asmLabel, asmPath);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "Worker{Id}: SharedAssembly '{Label}' load threw", Id, asmLabel);
+            }
+            finally { _pwsh.Commands.Clear(); _pwsh.Streams.ClearStreams(); }
         }
 
         // Deploy background scripts as Function:\ items.

Services/PowerShellWorkerPool.cs

@@ -440,13 +440,32 @@ public void Reclaim(PowerShellWorker worker, bool isHttp, bool faulted = false)
             var oldId = worker.Id;
             WorkerMetricsBridge.DeregisterWorker(oldId);
             worker.Dispose();
-            var cloned = isHttp ? _httpClonedState : _bgClonedState;
-            var iss = cloned != null ? BuildClonedISS(cloned) : BuildISS(isHttp: isHttp);
-            worker = new PowerShellWorker(Interlocked.Increment(ref _nextId), iss, _logger);
-            worker.Initialize(_repo, _apiBasePath, _settings);
-            WorkerMetricsBridge.RegisterWorker(worker.Id, isHttp);
-            _logger.LogInformation("[Pool] Replaced W{OldId} → W{NewId} ({Type})",
-                oldId, worker.Id, isHttp ? "HTTP" : "BG");
+
+            // Build the replacement off the calling thread so the dispatch loop is not held
+            // for the 6-13s ISS rebuild + Initialize() cost. Capacity briefly drops by 1 (the
+            // pool is short one worker until this Task completes), but the thread that just
+            // finished a task returns to the dispatch loop immediately.
+            var ish = isHttp;
+            _ = Task.Run(() =>
+            {
+                try
+                {
+                    var cloned = ish ? _httpClonedState : _bgClonedState;
+                    var iss = cloned != null ? BuildClonedISS(cloned) : BuildISS(isHttp: ish);
+                    var fresh = new PowerShellWorker(Interlocked.Increment(ref _nextId), iss, _logger);
+                    fresh.Initialize(_repo, _apiBasePath, _settings);
+                    WorkerMetricsBridge.RegisterWorker(fresh.Id, ish);
+                    if (ish) _httpPool.Add(fresh); else _bgPool.Add(fresh);
+                    _logger.LogInformation("[Pool] Replaced W{OldId} → W{NewId} ({Type}) (background recycle)",
+                        oldId, fresh.Id, ish ? "HTTP" : "BG");
+                }
+                catch (Exception ex)
+                {
+                    _logger.LogError(ex, "[Pool] Background recycle failed for W{OldId} ({Type}); pool capacity reduced by 1",
+                        oldId, ish ? "HTTP" : "BG");
+                }
+            });
+            return;
         }
 
         if (isHttp) _httpPool.Add(worker); else _bgPool.Add(worker);
@@ -458,6 +477,43 @@ private InitialSessionState BuildISS(bool isHttp)
         return BuildISSForModules(allowList.Count > 0 ? allowList : null);
     }
 
+    /// <summary>
+    /// Register configured SharedAssemblies on the InitialSessionState so the runspace's
+    /// type resolver knows about them from creation. This complements the runtime
+    /// [Reflection.Assembly]::LoadFile script in PowerShellWorker.Initialize() and ensures
+    /// type literals (e.g. [CIPP.TestDataCache]) resolve in cloned/recycled runspaces.
+    /// Logs each path so silent load failures during cold-start become observable.
+    /// </summary>
+    private void RegisterSharedAssemblies(InitialSessionState iss, string buildContext)
+    {
+        if (_settings.Worker.SharedAssemblies.Count == 0) return;
+
+        foreach (var asmRelPath in _settings.Worker.SharedAssemblies)
+        {
+            if (string.IsNullOrWhiteSpace(asmRelPath)) continue;
+
+            var asmPath = Path.GetFullPath(Path.Combine(_apiBasePath, asmRelPath));
+            if (!File.Exists(asmPath))
+            {
+                _logger.LogError("[Pool] SharedAssembly missing for {Context}: {Path}", buildContext, asmPath);
+                continue;
+            }
+
+            try
+            {
+                var asmName = Path.GetFileNameWithoutExtension(asmPath);
+                iss.Assemblies.Add(new SessionStateAssemblyEntry(asmName, asmPath));
+                _logger.LogDebug("[Pool] SharedAssembly registered for {Context}: {Name} ({Path})",
+                    buildContext, asmName, asmPath);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "[Pool] SharedAssembly registration failed for {Context}: {Path}",
+                    buildContext, asmPath);
+            }
+        }
+    }
+
     /// <summary>
     /// Build an ISS that imports only the specified modules.
     /// If moduleList is null, imports all modules (minus SkipModules).
@@ -468,6 +524,8 @@ private InitialSessionState BuildISSForModules(List<string>? moduleList)
         if (OperatingSystem.IsWindows())
             iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass;
 
+        RegisterSharedAssemblies(iss, "BuildISSForModules");
+
         // Copy environment variables into the runspace
         foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables())
             iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null));
@@ -515,6 +573,8 @@ private InitialSessionState BuildClonedISSWithModules(ExportedModuleState baseSt
         if (OperatingSystem.IsWindows())
             iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass;
 
+        RegisterSharedAssemblies(iss, "BuildClonedISSWithModules");
+
         foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables())
             iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null));
 
@@ -581,6 +641,8 @@ private InitialSessionState BuildClonedISS(ExportedModuleState state)
         if (OperatingSystem.IsWindows())
             iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass;
 
+        RegisterSharedAssemblies(iss, "BuildClonedISS");
+
         // Copy environment variables
         foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables())
             iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null));

Services/SetupPages.cs

@@ -216,8 +216,7 @@ function setTenantMode(section, isMulti) {
             }
             document.getElementById('manual-redirect-uri').textContent = window.location.origin + '/.auth/login/aad/callback';
             if (setupState.isEasyAuthConfigured) {
-                showBanner('Authentication is already configured. Redirecting...');
-                setTimeout(() => window.location.href = '/', 2000);
+                showBanner('Authentication is already configured. This page will redirect shortly.');
                 return;
             }
             if (!setupState.isRunningInAppService || !setupState.hasManagedIdentity) {
@@ -459,57 +458,52 @@ async function submitManual() {
         }
     }
     function showRestartScreen() {
-        // Stop background status polling — restart screen has its own polling
+        // Stop background status polling
         clearInterval(statusPollTimer);
-        // Hide setup sections, show restart polling UI
+        // Hide setup sections, show restart status
         document.getElementById('user-section').classList.add('hidden');
         document.getElementById('auto-section').classList.add('hidden');
         document.getElementById('manual-section').classList.add('hidden');
         document.getElementById('divider-user').classList.add('hidden');
         document.getElementById('divider-auth').classList.add('hidden');
         document.querySelector('.subtitle').textContent = '';
-        document.getElementById('page-title').textContent = 'Restarting...';
+        document.getElementById('page-title').textContent = 'Setup Complete';
 
         const banner = document.getElementById('status-banner');
         banner.innerHTML = '<div style="text-align:center;">' +
             '<div style="margin-bottom:0.75rem;">Authentication has been configured. The application is restarting to apply changes.</div>' +
             '<div id="restart-spinner" style="font-size:1.5rem;margin-bottom:0.5rem;">&#8987;</div>' +
-            '<div id="restart-status" style="font-size:0.85rem;color:#94a3b8;">Waiting for the application to come back online...</div>' +
+            '<div id="restart-status" style="font-size:0.85rem;color:#94a3b8;">Waiting for the new container to come online...</div>' +
             '</div>';
         banner.classList.remove('hidden');
 
-        let attempts = 0;
-        const maxAttempts = 120; // 10 min at 5s intervals
+        // Poll /api/setup/status — once the new container comes online with EasyAuth
+        // configured, isEasyAuthConfigured will be true and setup mode will no longer
+        // be active. At that point redirect to / and let the app take over.
         const pollInterval = 5000;
-
         const pollRestart = async () => {
-            attempts++;
-            const statusEl = document.getElementById('restart-status');
             try {
-                const res = await fetch('/api/setup/health', { cache: 'no-store' });
+                const res = await fetch('/api/setup/status', { cache: 'no-store' });
                 if (res.ok) {
                     const data = await res.json();
-                    if (data.ready) {
-                        statusEl.textContent = 'Application is ready! Redirecting...';
+                    if (data.isEasyAuthConfigured && !data.isSetupCompleted) {
+                        // New container is online with EasyAuth active — redirect to app
+                        document.getElementById('restart-status').textContent = 'Application is ready!';
                         document.getElementById('restart-spinner').textContent = '\u2705';
-                        setTimeout(() => window.location.href = '/', 1500);
+                        setTimeout(() => window.location.href = '/', 1000);
                         return;
                     }
+                    // Still on the old container (isSetupCompleted=true) — keep waiting
+                    document.getElementById('restart-status').textContent = 'Auth configured, waiting for container restart...';
                 }
             } catch (e) {
-                // App still restarting — expected
-            }
-            if (attempts >= maxAttempts) {
-                statusEl.textContent = 'The application is taking longer than expected. Try refreshing the page manually.';
-                document.getElementById('restart-spinner').textContent = '\u26A0\uFE0F';
-                return;
+                // Container is cycling — expected during restart
+                document.getElementById('restart-status').textContent = 'Container restarting...';
             }
-            statusEl.textContent = 'Waiting for the application to come back online... (' + attempts + ')';
             setTimeout(pollRestart, pollInterval);
         };
 
-        // Start polling after a short delay to let the restart begin
-        setTimeout(pollRestart, 8000);
+        setTimeout(pollRestart, 5000);
     }
 </script>
 </body>