Fixed test cases for job timeout (#5332)

Co-authored-by: Sanju Yadav <sanjuyadav@microsoft.com>

Author: sanjuyadav24

src/Test/L1/Worker/L1TestBase.cs

@@ -250,10 +250,24 @@ protected async Task<TestResults> RunWorker(Pipelines.AgentJobRequestMessage mes
 
             await SetupMessage(_l1HostContext, message);
 
-            using (var cts = new CancellationTokenSource())
+            var cts = new CancellationTokenSource();
+            try
             {
                 cts.CancelAfter((int)JobTimeout.TotalMilliseconds);
-                return await RunWorker(_l1HostContext, message, cts.Token);
+                var result = await RunWorker(_l1HostContext, message, cts.Token);
+                
+                // If job timed out, give it a moment to complete finalization
+                if (result.TimedOut)
+                {
+                    await Task.Delay(100); // Allow 100ms for cleanup to complete
+                }
+                
+                return result;
+            }
+            finally
+            {
+                // Dispose after ensuring cleanup had time to complete
+                cts?.Dispose();
             }
         }
 
@@ -333,6 +347,24 @@ await processChannel.SendAsync(
                 }
                 else
                 {
+                    // Timeout occurred - give worker task a moment to complete cleanup gracefully
+                    try
+                    {
+                        // Wait up to 2 seconds for graceful shutdown after timeout
+                        using (var gracefulShutdownCts = new CancellationTokenSource(TimeSpan.FromSeconds(2)))
+                        {
+                            await workerTask.WaitAsync(gracefulShutdownCts.Token);
+                        }
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        // Worker didn't complete gracefully within 2 seconds, that's expected for timeout tests
+                    }
+                    catch (Exception)
+                    {
+                        // Other exceptions during shutdown are also expected in timeout scenarios
+                    }
+
                     return new TestResults
                     {
                         TimedOut = true