From 2cbaeb1e961e62e9cca67e754f9c397de9635168 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 15:46:00 +0530
Subject: [PATCH 01/11] Initial changes for stream command

---
 .../internal/cmd/job.go                       |   1 +
 .../internal/cmd/job_stream.go                |  90 +++++++
 .../internal/service/stream_service.go        | 255 ++++++++++++++++++
 .../pkg/client/history.go                     |  83 ++++++
 .../pkg/client/logs.go                        |  48 ++++
 .../pkg/models/history.go                     |  45 ++--
 6 files changed, 503 insertions(+), 19 deletions(-)
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go
index 45baf7101b7..ccdac9a8f16 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go
@@ -35,6 +35,7 @@ func newJobCommand() *cobra.Command {
 	cmd.AddCommand(newJobDeleteCommand())
 	cmd.AddCommand(newJobCancelCommand())
 	cmd.AddCommand(newJobValidateCommand())
+	cmd.AddCommand(newJobStreamCommand())
 
 	return cmd
 }
diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go
new file mode 100644
index 00000000000..34e7ff4e790
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go
@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"fmt"
+
+	"azure.ai.customtraining/internal/service"
+	"azure.ai.customtraining/internal/utils"
+	"azure.ai.customtraining/pkg/client"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
+	"github.com/spf13/cobra"
+)
+
+func newJobStreamCommand() *cobra.Command {
+	var name string
+
+	cmd := &cobra.Command{
+		Use:   "stream",
+		Short: "Stream logs from a running training job",
+		Long:  "Stream live log output from a training job. Polls for new log content\nuntil the job reaches a terminal state.\n\nExample:\n  azd ai training job stream --name my-job-123",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := azdext.WithAccessToken(cmd.Context())
+
+			if name == "" {
+				return fmt.Errorf("--name is required")
+			}
+
+			azdClient, err := azdext.NewAzdClient()
+			if err != nil {
+				return fmt.Errorf("failed to create azd client: %w", err)
+			}
+			defer azdClient.Close()
+
+			envValues, err := utils.GetEnvironmentValues(ctx, azdClient)
+			if err != nil {
+				return fmt.Errorf("failed to get environment values: %w", err)
+			}
+
+			accountName := envValues[utils.EnvAzureAccountName]
+			projectName := envValues[utils.EnvAzureProjectName]
+			tenantID := envValues[utils.EnvAzureTenantID]
+
+			if accountName == "" || projectName == "" {
+				return fmt.Errorf("environment not configured. Run 'azd ai training init' first")
+			}
+
+			credential, err := azidentity.NewAzureDeveloperCLICredential(
+				&azidentity.AzureDeveloperCLICredentialOptions{
+					TenantID:                   tenantID,
+					AdditionallyAllowedTenants: []string{"*"},
+				},
+			)
+			if err != nil {
+				return fmt.Errorf("failed to create azure credential: %w", err)
+			}
+
+			endpoint := buildProjectEndpoint(accountName, projectName)
+			apiClient, err := client.NewClient(endpoint, credential)
+			if err != nil {
+				return fmt.Errorf("failed to create API client: %w", err)
+			}
+			apiClient.SetDebugBody(rootFlags.Debug)
+
+			streamSvc := service.NewStreamService(apiClient)
+			result, err := streamSvc.StreamJobLogs(ctx, name)
+			if err != nil {
+				return err
+			}
+
+			fmt.Println()
+			fmt.Println("Execution Summary")
+			fmt.Printf("RunId: %s\n", result.JobName)
+			fmt.Printf("Status: %s\n", result.Status)
+			if result.StudioURL != "" {
+				fmt.Printf("Web View: %s\n", result.StudioURL)
+			}
+
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&name, "name", "", "Job name/ID (required)")
+
+	return cmd
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
new file mode 100644
index 00000000000..5b3d5a46d2c
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -0,0 +1,255 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package service
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sort"
+	"strings"
+	"time"
+
+	"azure.ai.customtraining/pkg/client"
+)
+
+// StreamService handles polling and streaming log output from a running job.
+type StreamService struct {
+	client *client.Client
+}
+
+// NewStreamService creates a new stream service.
+func NewStreamService(apiClient *client.Client) *StreamService {
+	return &StreamService{client: apiClient}
+}
+
+// terminalStates are job statuses that indicate the job has finished.
+var terminalStates = map[string]bool{
+	"Completed":     true,
+	"Failed":        true,
+	"Canceled":      true,
+	"NotResponding": true,
+	"Paused":        true,
+}
+
+// activeStates are job statuses where streaming is applicable.
+var activeStates = map[string]bool{
+	"NotStarted":   true,
+	"Queued":       true,
+	"Preparing":    true,
+	"Provisioning": true,
+	"Starting":     true,
+	"Running":      true,
+	"Finalizing":   true,
+}
+
+// StreamResult contains the final state of a streamed job.
+type StreamResult struct {
+	JobName   string
+	Status    string
+	StudioURL string
+}
+
+// StreamJobLogs polls the job and streams log output until the job reaches a terminal state.
+func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*StreamResult, error) {
+	fmt.Fprintf(os.Stderr, "Stream logs for job %s...\n\n", jobName)
+
+	offsets := make(map[string]int64)
+
+	const (
+		initialInterval    = 2 * time.Second
+		maxInterval        = 5 * time.Second
+		jobCheckFrequency  = 10
+		maxConsecutiveErrs = 3
+		initialTailBytes   = int64(8192)
+	)
+
+	pollInterval := initialInterval
+	pollCount := 0
+	consecutiveErrs := 0
+	trackingEndpoint := ""
+
+	for {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+
+		var jobStatus string
+		var studioURL string
+		if pollCount%jobCheckFrequency == 0 {
+			job, err := s.client.GetJob(ctx, jobName)
+			if err != nil {
+				consecutiveErrs++
+				if consecutiveErrs >= maxConsecutiveErrs {
+					return nil, fmt.Errorf("failed to get job status after %d retries: %w", maxConsecutiveErrs, err)
+				}
+				time.Sleep(pollInterval)
+				pollCount++
+				continue
+			}
+			consecutiveErrs = 0
+
+			jobStatus = job.Properties.Status
+			studioURL = extractServiceEndpoint(job.Properties.Services, "Studio")
+
+			if trackingEndpoint == "" {
+				trackingEndpoint = extractServiceEndpoint(job.Properties.Services, "Tracking")
+			}
+
+			if terminalStates[jobStatus] {
+				if trackingEndpoint != "" {
+					s.flushLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+				}
+				return &StreamResult{
+					JobName:   jobName,
+					Status:    jobStatus,
+					StudioURL: studioURL,
+				}, nil
+			}
+
+			if !activeStates[jobStatus] {
+				fmt.Fprintf(os.Stderr, "Job status: %s, waiting...\n", jobStatus)
+				time.Sleep(pollInterval)
+				pollCount++
+				continue
+			}
+		}
+
+		if trackingEndpoint != "" {
+			hasNewContent, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+			if err != nil {
+				consecutiveErrs++
+				if consecutiveErrs >= maxConsecutiveErrs {
+					return nil, fmt.Errorf("failed to stream logs after %d retries: %w", maxConsecutiveErrs, err)
+				}
+			} else {
+				consecutiveErrs = 0
+				if hasNewContent {
+					pollInterval = initialInterval
+				} else {
+					pollInterval = min(pollInterval+time.Second, maxInterval)
+				}
+			}
+		} else {
+			fmt.Fprintf(os.Stderr, "Waiting for job to initialize...\n")
+			pollInterval = maxInterval
+		}
+
+		pollCount++
+		time.Sleep(pollInterval)
+	}
+}
+
+// pollAndPrintLogs fetches run history details and streams new log content.
+func (s *StreamService) pollAndPrintLogs(
+	ctx context.Context,
+	trackingEndpoint string,
+	jobName string,
+	offsets map[string]int64,
+	initialTailBytes int64,
+) (bool, error) {
+	details, err := s.client.GetRunHistoryDetails(ctx, trackingEndpoint, jobName)
+	if err != nil {
+		return false, err
+	}
+	if details == nil || len(details.LogFiles) == 0 {
+		return false, nil
+	}
+
+	fileNames := make([]string, 0, len(details.LogFiles))
+	for name := range details.LogFiles {
+		fileNames = append(fileNames, name)
+	}
+	sort.Strings(fileNames)
+
+	hasNewContent := false
+	for _, fileName := range fileNames {
+		sasURI := details.LogFiles[fileName]
+
+		offset, exists := offsets[fileName]
+		if !exists {
+			offset = -initialTailBytes
+		}
+
+		content, bytesRead, err := s.fetchLogChunk(ctx, sasURI, offset)
+		if err != nil {
+			continue
+		}
+
+		if bytesRead > 0 && content != "" {
+			fmt.Printf("\n--- %s ---\n", fileName)
+			fmt.Print(content)
+			if !strings.HasSuffix(content, "\n") {
+				fmt.Println()
+			}
+			hasNewContent = true
+
+			if offset < 0 {
+				offsets[fileName] = bytesRead
+			} else {
+				offsets[fileName] = offset + bytesRead
+			}
+		} else if !exists {
+			offsets[fileName] = 0
+		}
+	}
+
+	return hasNewContent, nil
+}
+
+// fetchLogChunk retrieves log content, handling initial tail (negative offset).
+func (s *StreamService) fetchLogChunk(ctx context.Context, sasURI string, offset int64) (string, int64, error) {
+	if offset < 0 {
+		content, bytesRead, err := s.client.GetLogContent(ctx, sasURI, 0)
+		if err != nil {
+			return "", 0, err
+		}
+		tailSize := -offset
+		if int64(len(content)) > tailSize {
+			trimmed := content[int64(len(content))-tailSize:]
+			if idx := strings.Index(trimmed, "\n"); idx >= 0 {
+				trimmed = trimmed[idx+1:]
+			}
+			return trimmed, bytesRead, nil
+		}
+		return content, bytesRead, nil
+	}
+
+	return s.client.GetLogContent(ctx, sasURI, offset)
+}
+
+// flushLogs does a final poll to capture any remaining log output.
+func (s *StreamService) flushLogs(
+	ctx context.Context,
+	trackingEndpoint string,
+	jobName string,
+	offsets map[string]int64,
+	initialTailBytes int64,
+) {
+	_, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+}
+
+// extractServiceEndpoint extracts the endpoint URL from the job services map.
+func extractServiceEndpoint(services map[string]interface{}, serviceName string) string {
+	if services == nil {
+		return ""
+	}
+	svc, ok := services[serviceName]
+	if !ok {
+		return ""
+	}
+	svcMap, ok := svc.(map[string]interface{})
+	if !ok {
+		return ""
+	}
+	endpoint, ok := svcMap["endpoint"]
+	if !ok {
+		return ""
+	}
+	str, ok := endpoint.(string)
+	if !ok {
+		return ""
+	}
+	return str
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go
index 7836f62cf0d..acfe9eb04d4 100644
--- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go
@@ -8,6 +8,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/url"
+	"strings"
 
 	"azure.ai.customtraining/pkg/models"
 )
@@ -36,3 +38,84 @@ func (c *Client) GetRunHistory(ctx context.Context, runID string) (*models.RunHi
 
 	return &result, nil
 }
+
+// GetRunHistoryDetails retrieves detailed run information including log file SAS URIs.
+// This calls the AML history service directly using the tracking endpoint from the job response.
+// GET https://{region}.api.azureml.ms/history/v1.0/{workspace}/runs/{runId}/details
+func (c *Client) GetRunHistoryDetails(ctx context.Context, trackingEndpoint string, runID string) (*models.RunHistoryDetails, error) {
+	baseURL, workspacePath, err := parseTrackingEndpoint(trackingEndpoint)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse tracking endpoint: %w", err)
+	}
+
+	reqURL := fmt.Sprintf("%s/history/v1.0%s/runs/%s/details", baseURL, workspacePath, url.PathEscape(runID))
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	if err := c.addAuth(ctx, req, DataPlaneScope); err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	if c.debugBody {
+		fmt.Printf("[DEBUG] GET %s\n", reqURL)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusNotFound {
+		return nil, nil
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, c.HandleError(resp)
+	}
+
+	var result models.RunHistoryDetails
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode run history details response: %w", err)
+	}
+
+	return &result, nil
+}
+
+// parseTrackingEndpoint extracts the base URL and workspace path from a tracking endpoint.
+// Input format:  azureml://{region}.api.azureml.ms/mlflow/v1.0/{workspace-path}?...
+// Returns:       https://{region}.api.azureml.ms, /{workspace-path}
+func parseTrackingEndpoint(trackingEndpoint string) (string, string, error) {
+	// Remove the azureml:// prefix
+	endpoint := strings.TrimPrefix(trackingEndpoint, "azureml://")
+	// Remove any trailing query params
+	if idx := strings.Index(endpoint, "?"); idx != -1 {
+		endpoint = endpoint[:idx]
+	}
+
+	// Parse: {host}/mlflow/v1.0/{subscription-path}
+	parsed, err := url.Parse("https://" + endpoint)
+	if err != nil {
+		return "", "", fmt.Errorf("invalid tracking endpoint: %w", err)
+	}
+
+	host := parsed.Host
+	if host == "" {
+		return "", "", fmt.Errorf("invalid tracking endpoint: missing host")
+	}
+
+	// Extract workspace path: everything after /mlflow/v1.0/ or /mlflow/v2.0/
+	path := parsed.Path
+	for _, prefix := range []string{"/mlflow/v1.0", "/mlflow/v2.0"} {
+		if strings.HasPrefix(path, prefix) {
+			workspacePath := strings.TrimPrefix(path, prefix)
+			return "https://" + host, workspacePath, nil
+		}
+	}
+
+	return "", "", fmt.Errorf("invalid tracking endpoint: expected mlflow path prefix in %q", path)
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
new file mode 100644
index 00000000000..935706d9bd7
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package client
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+)
+
+// GetLogContent fetches log file content from a SAS URI starting at the given byte offset.
+// Returns the content and the total number of bytes read (to update the offset).
+// No authentication is needed since the URL contains a SAS token.
+func (c *Client) GetLogContent(ctx context.Context, sasURI string, offset int64) (string, int64, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sasURI, nil)
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to create log request: %w", err)
+	}
+
+	if offset > 0 {
+		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", offset))
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return "", 0, fmt.Errorf("log request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// 416 Range Not Satisfiable means no new content beyond the offset
+	if resp.StatusCode == http.StatusRequestedRangeNotSatisfiable {
+		return "", 0, nil
+	}
+
+	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
+		return "", 0, fmt.Errorf("log request returned status %d", resp.StatusCode)
+	}
+
+	// Cap read to 1MB per poll to avoid memory issues with very large logs
+	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to read log content: %w", err)
+	}
+
+	return string(body), int64(len(body)), nil
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go b/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go
index 5a709bf224e..23a17f31b8a 100644
--- a/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go
@@ -5,25 +5,25 @@ package models
 
 // RunHistory represents the run history details for a job.
 type RunHistory struct {
-	RunID           string            `json:"runId,omitempty"`
-	RunUUID         string            `json:"runUuid,omitempty"`
-	RootRunID       string            `json:"rootRunId,omitempty"`
-	Status          string            `json:"status,omitempty"`
-	StartTimeUTC    string            `json:"startTimeUtc,omitempty"`
-	EndTimeUTC      string            `json:"endTimeUtc,omitempty"`
-	Duration        string            `json:"duration,omitempty"`
-	ComputeDuration string            `json:"computeDuration,omitempty"`
-	CreatedUTC      string            `json:"createdUtc,omitempty"`
-	LastModifiedUTC string            `json:"lastModifiedUtc,omitempty"`
-	DisplayName     string            `json:"displayName,omitempty"`
-	Description     string            `json:"description,omitempty"`
-	Target          string            `json:"target,omitempty"`
-	RunType         string            `json:"runType,omitempty"`
-	Error           *RunHistoryError  `json:"error,omitempty"`
-	CreatedBy       *RunHistoryUser   `json:"createdBy,omitempty"`
-	Compute         *RunHistoryCompute `json:"compute,omitempty"`
-	Properties      map[string]string `json:"properties,omitempty"`
-	Tags            map[string]string `json:"tags,omitempty"`
+	RunID           string                     `json:"runId,omitempty"`
+	RunUUID         string                     `json:"runUuid,omitempty"`
+	RootRunID       string                     `json:"rootRunId,omitempty"`
+	Status          string                     `json:"status,omitempty"`
+	StartTimeUTC    string                     `json:"startTimeUtc,omitempty"`
+	EndTimeUTC      string                     `json:"endTimeUtc,omitempty"`
+	Duration        string                     `json:"duration,omitempty"`
+	ComputeDuration string                     `json:"computeDuration,omitempty"`
+	CreatedUTC      string                     `json:"createdUtc,omitempty"`
+	LastModifiedUTC string                     `json:"lastModifiedUtc,omitempty"`
+	DisplayName     string                     `json:"displayName,omitempty"`
+	Description     string                     `json:"description,omitempty"`
+	Target          string                     `json:"target,omitempty"`
+	RunType         string                     `json:"runType,omitempty"`
+	Error           *RunHistoryError           `json:"error,omitempty"`
+	CreatedBy       *RunHistoryUser            `json:"createdBy,omitempty"`
+	Compute         *RunHistoryCompute         `json:"compute,omitempty"`
+	Properties      map[string]string          `json:"properties,omitempty"`
+	Tags            map[string]string          `json:"tags,omitempty"`
 	Inputs          map[string]RunHistoryAsset `json:"inputs,omitempty"`
 	Outputs         map[string]RunHistoryAsset `json:"outputs,omitempty"`
 }
@@ -68,3 +68,10 @@ type RunHistoryList struct {
 	Value    []RunHistory `json:"value"`
 	NextLink string       `json:"nextLink,omitempty"`
 }
+
+// RunHistoryDetails represents the details response from the AML history service.
+// GET .../runs/{runId}/details
+type RunHistoryDetails struct {
+	RunID    string            `json:"runId,omitempty"`
+	LogFiles map[string]string `json:"logFiles,omitempty"` // filename → SAS URI
+}

From 254cfa0ca65e7b88ec163395e15846f647e19599 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 19:42:00 +0530
Subject: [PATCH 02/11] Wrap debug logs

---
 .../extensions/azure.ai.customtraining/pkg/client/client.go   | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go
index 8784444cdae..e22b513e906 100644
--- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go
@@ -88,7 +88,9 @@ func (c *Client) doDataPlaneWithVersion(ctx context.Context, method, path, apiVe
 		reqURL += fmt.Sprintf("&%s=%s", queryParams[i], url.QueryEscape(queryParams[i+1]))
 	}
 
-	fmt.Printf("[DEBUG] %s %s\n", method, reqURL)
+	if c.debugBody {
+		fmt.Printf("[DEBUG] %s %s\n", method, reqURL)
+	}
 
 	var bodyReader io.Reader
 	if body != nil {

From 29aa9f903aaff063c35a6e2a91d606ecd0e917a6 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 20:17:38 +0530
Subject: [PATCH 03/11] Add regex to only log specific files as in AML | Use
 line count method instead of offset to avoid abrupt line breaks and match
 with AML experience

---
 .../internal/service/stream_service.go        | 115 +++++++++---------
 .../pkg/client/logs.go                        |  16 +--
 2 files changed, 63 insertions(+), 68 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index 5b3d5a46d2c..c31d98b12c7 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"regexp"
 	"sort"
 	"strings"
 	"time"
@@ -24,6 +25,14 @@ func NewStreamService(apiClient *client.Client) *StreamService {
 	return &StreamService{client: apiClient}
 }
 
+// Log file patterns matching the Azure ML SDK behavior.
+// Primary: Common Runtime user logs.
+// Fallback: Legacy azureml-logs for older compute targets.
+var (
+	commonRuntimeLogPattern = regexp.MustCompile(`user_logs/std_log[\D]*[0]*(?:_ps)?\.txt`)
+	legacyLogPattern        = regexp.MustCompile(`azureml-logs/[\d]{2}.+\.txt`)
+)
+
 // terminalStates are job statuses that indicate the job has finished.
 var terminalStates = map[string]bool{
 	"Completed":     true,
@@ -53,16 +62,17 @@ type StreamResult struct {
 
 // StreamJobLogs polls the job and streams log output until the job reaches a terminal state.
 func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*StreamResult, error) {
-	fmt.Fprintf(os.Stderr, "Stream logs for job %s...\n\n", jobName)
+	fmt.Fprintf(os.Stderr, "Streaming logs for job: %s\n\n", jobName)
 
-	offsets := make(map[string]int64)
+	// Line-count tracking per file, matching the Azure ML SDK approach.
+	// Each poll downloads full content, skips already-printed lines, prints the rest.
+	processedLines := make(map[string]int)
 
 	const (
 		initialInterval    = 2 * time.Second
 		maxInterval        = 5 * time.Second
 		jobCheckFrequency  = 10
 		maxConsecutiveErrs = 3
-		initialTailBytes   = int64(8192)
 	)
 
 	pollInterval := initialInterval
@@ -99,7 +109,7 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 
 			if terminalStates[jobStatus] {
 				if trackingEndpoint != "" {
-					s.flushLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+					s.flushLogs(ctx, trackingEndpoint, jobName, processedLines)
 				}
 				return &StreamResult{
 					JobName:   jobName,
@@ -117,7 +127,7 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 		}
 
 		if trackingEndpoint != "" {
-			hasNewContent, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+			hasNewContent, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
 			if err != nil {
 				consecutiveErrs++
 				if consecutiveErrs >= maxConsecutiveErrs {
@@ -141,13 +151,34 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 	}
 }
 
-// pollAndPrintLogs fetches run history details and streams new log content.
+// filterLogFiles selects streamable log files from the history details.
+// Matches Common Runtime user logs first; falls back to legacy azureml-logs.
+// Returns matched file names sorted alphabetically.
+func filterLogFiles(logFiles map[string]string) []string {
+	var matched []string
+	for name := range logFiles {
+		if commonRuntimeLogPattern.MatchString(name) {
+			matched = append(matched, name)
+		}
+	}
+	if len(matched) == 0 {
+		// Fallback to legacy log pattern for older compute targets
+		for name := range logFiles {
+			if legacyLogPattern.MatchString(name) {
+				matched = append(matched, name)
+			}
+		}
+	}
+	sort.Strings(matched)
+	return matched
+}
+
+// pollAndPrintLogs fetches run history details and prints only new log lines.
 func (s *StreamService) pollAndPrintLogs(
 	ctx context.Context,
 	trackingEndpoint string,
 	jobName string,
-	offsets map[string]int64,
-	initialTailBytes int64,
+	processedLines map[string]int,
 ) (bool, error) {
 	details, err := s.client.GetRunHistoryDetails(ctx, trackingEndpoint, jobName)
 	if err != nil {
@@ -157,66 +188,39 @@ func (s *StreamService) pollAndPrintLogs(
 		return false, nil
 	}
 
-	fileNames := make([]string, 0, len(details.LogFiles))
-	for name := range details.LogFiles {
-		fileNames = append(fileNames, name)
+	fileNames := filterLogFiles(details.LogFiles)
+	if len(fileNames) == 0 {
+		return false, nil
 	}
-	sort.Strings(fileNames)
 
 	hasNewContent := false
 	for _, fileName := range fileNames {
 		sasURI := details.LogFiles[fileName]
 
-		offset, exists := offsets[fileName]
-		if !exists {
-			offset = -initialTailBytes
-		}
-
-		content, bytesRead, err := s.fetchLogChunk(ctx, sasURI, offset)
-		if err != nil {
+		content, _, err := s.client.GetLogContent(ctx, sasURI, 0)
+		if err != nil || content == "" {
 			continue
 		}
 
-		if bytesRead > 0 && content != "" {
-			fmt.Printf("\n--- %s ---\n", fileName)
-			fmt.Print(content)
-			if !strings.HasSuffix(content, "\n") {
-				fmt.Println()
-			}
-			hasNewContent = true
-
-			if offset < 0 {
-				offsets[fileName] = bytesRead
-			} else {
-				offsets[fileName] = offset + bytesRead
-			}
-		} else if !exists {
-			offsets[fileName] = 0
+		lines := strings.Split(content, "\n")
+		// Remove trailing empty element from final newline
+		if len(lines) > 0 && lines[len(lines)-1] == "" {
+			lines = lines[:len(lines)-1]
 		}
-	}
-
-	return hasNewContent, nil
-}
 
-// fetchLogChunk retrieves log content, handling initial tail (negative offset).
-func (s *StreamService) fetchLogChunk(ctx context.Context, sasURI string, offset int64) (string, int64, error) {
-	if offset < 0 {
-		content, bytesRead, err := s.client.GetLogContent(ctx, sasURI, 0)
-		if err != nil {
-			return "", 0, err
+		previousLines := processedLines[fileName]
+		if len(lines) <= previousLines {
+			continue
 		}
-		tailSize := -offset
-		if int64(len(content)) > tailSize {
-			trimmed := content[int64(len(content))-tailSize:]
-			if idx := strings.Index(trimmed, "\n"); idx >= 0 {
-				trimmed = trimmed[idx+1:]
-			}
-			return trimmed, bytesRead, nil
+
+		for _, line := range lines[previousLines:] {
+			fmt.Println(line)
 		}
-		return content, bytesRead, nil
+		hasNewContent = true
+		processedLines[fileName] = len(lines)
 	}
 
-	return s.client.GetLogContent(ctx, sasURI, offset)
+	return hasNewContent, nil
 }
 
 // flushLogs does a final poll to capture any remaining log output.
@@ -224,10 +228,9 @@ func (s *StreamService) flushLogs(
 	ctx context.Context,
 	trackingEndpoint string,
 	jobName string,
-	offsets map[string]int64,
-	initialTailBytes int64,
+	processedLines map[string]int,
 ) {
-	_, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, offsets, initialTailBytes)
+	_, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
 }
 
 // extractServiceEndpoint extracts the endpoint URL from the job services map.
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
index 935706d9bd7..c6aa2b4ded0 100644
--- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
@@ -10,8 +10,9 @@ import (
 	"net/http"
 )
 
-// GetLogContent fetches log file content from a SAS URI starting at the given byte offset.
-// Returns the content and the total number of bytes read (to update the offset).
+// GetLogContent fetches log file content from a SAS URI.
+// Downloads the full content on each call (no Range header). The caller tracks
+// which lines have already been printed, matching the Azure ML SDK approach.
 // No authentication is needed since the URL contains a SAS token.
 func (c *Client) GetLogContent(ctx context.Context, sasURI string, offset int64) (string, int64, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sasURI, nil)
@@ -19,22 +20,13 @@ func (c *Client) GetLogContent(ctx context.Context, sasURI string, offset int64)
 		return "", 0, fmt.Errorf("failed to create log request: %w", err)
 	}
 
-	if offset > 0 {
-		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", offset))
-	}
-
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return "", 0, fmt.Errorf("log request failed: %w", err)
 	}
 	defer resp.Body.Close()
 
-	// 416 Range Not Satisfiable means no new content beyond the offset
-	if resp.StatusCode == http.StatusRequestedRangeNotSatisfiable {
-		return "", 0, nil
-	}
-
-	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
+	if resp.StatusCode != http.StatusOK {
 		return "", 0, fmt.Errorf("log request returned status %d", resp.StatusCode)
 	}
 

From 2592919315efeee0754b1f8628cb78b345a9cfd8 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 20:25:09 +0530
Subject: [PATCH 04/11] Skip download for jobs already in terminal state

---
 .../azure.ai.customtraining/internal/service/stream_service.go  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index c31d98b12c7..193a815268e 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -108,7 +108,7 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 			}
 
 			if terminalStates[jobStatus] {
-				if trackingEndpoint != "" {
+				if trackingEndpoint != "" && pollCount > 0 {
 					s.flushLogs(ctx, trackingEndpoint, jobName, processedLines)
 				}
 				return &StreamResult{

From 24012767f2490f864e8e9c2297a908bc48064577 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 20:36:20 +0530
Subject: [PATCH 05/11] Use Sigmoid for Poll interval - to reduce API load for
 multi-hour jobs

---
 .../internal/service/stream_service.go        | 106 +++++++++---------
 1 file changed, 56 insertions(+), 50 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index 193a815268e..907dc436bae 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -6,6 +6,7 @@ package service
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"regexp"
 	"sort"
@@ -68,66 +69,58 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 	// Each poll downloads full content, skips already-printed lines, prints the rest.
 	processedLines := make(map[string]int)
 
-	const (
-		initialInterval    = 2 * time.Second
-		maxInterval        = 5 * time.Second
-		jobCheckFrequency  = 10
-		maxConsecutiveErrs = 3
-	)
+	const maxConsecutiveErrs = 3
 
-	pollInterval := initialInterval
-	pollCount := 0
+	startTime := time.Now()
 	consecutiveErrs := 0
 	trackingEndpoint := ""
+	firstPoll := true
 
 	for {
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
 
-		var jobStatus string
-		var studioURL string
-		if pollCount%jobCheckFrequency == 0 {
-			job, err := s.client.GetJob(ctx, jobName)
-			if err != nil {
-				consecutiveErrs++
-				if consecutiveErrs >= maxConsecutiveErrs {
-					return nil, fmt.Errorf("failed to get job status after %d retries: %w", maxConsecutiveErrs, err)
-				}
-				time.Sleep(pollInterval)
-				pollCount++
-				continue
+		// Check job status every poll
+		job, err := s.client.GetJob(ctx, jobName)
+		if err != nil {
+			consecutiveErrs++
+			if consecutiveErrs >= maxConsecutiveErrs {
+				return nil, fmt.Errorf("failed to get job status after %d retries: %w", maxConsecutiveErrs, err)
 			}
-			consecutiveErrs = 0
+			time.Sleep(pollInterval(startTime))
+			continue
+		}
+		consecutiveErrs = 0
 
-			jobStatus = job.Properties.Status
-			studioURL = extractServiceEndpoint(job.Properties.Services, "Studio")
+		jobStatus := job.Properties.Status
+		studioURL := extractServiceEndpoint(job.Properties.Services, "Studio")
 
-			if trackingEndpoint == "" {
-				trackingEndpoint = extractServiceEndpoint(job.Properties.Services, "Tracking")
-			}
+		if trackingEndpoint == "" {
+			trackingEndpoint = extractServiceEndpoint(job.Properties.Services, "Tracking")
+		}
 
-			if terminalStates[jobStatus] {
-				if trackingEndpoint != "" && pollCount > 0 {
-					s.flushLogs(ctx, trackingEndpoint, jobName, processedLines)
-				}
-				return &StreamResult{
-					JobName:   jobName,
-					Status:    jobStatus,
-					StudioURL: studioURL,
-				}, nil
+		if terminalStates[jobStatus] {
+			if trackingEndpoint != "" && !firstPoll {
+				s.flushLogs(ctx, trackingEndpoint, jobName, processedLines)
 			}
+			return &StreamResult{
+				JobName:   jobName,
+				Status:    jobStatus,
+				StudioURL: studioURL,
+			}, nil
+		}
 
-			if !activeStates[jobStatus] {
-				fmt.Fprintf(os.Stderr, "Job status: %s, waiting...\n", jobStatus)
-				time.Sleep(pollInterval)
-				pollCount++
-				continue
-			}
+		if !activeStates[jobStatus] {
+			fmt.Fprintf(os.Stderr, "Job status: %s, waiting...\n", jobStatus)
+			time.Sleep(pollInterval(startTime))
+			firstPoll = false
+			continue
 		}
 
+		// Stream logs if tracking endpoint is available
 		if trackingEndpoint != "" {
-			hasNewContent, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
+			_, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
 			if err != nil {
 				consecutiveErrs++
 				if consecutiveErrs >= maxConsecutiveErrs {
@@ -135,20 +128,33 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 				}
 			} else {
 				consecutiveErrs = 0
-				if hasNewContent {
-					pollInterval = initialInterval
-				} else {
-					pollInterval = min(pollInterval+time.Second, maxInterval)
-				}
 			}
 		} else {
 			fmt.Fprintf(os.Stderr, "Waiting for job to initialize...\n")
-			pollInterval = maxInterval
 		}
 
-		pollCount++
-		time.Sleep(pollInterval)
+		firstPoll = false
+		time.Sleep(pollInterval(startTime))
+	}
+}
+
+// pollInterval returns the polling interval using a sigmoid curve from 2s → 60s
+// based on elapsed time since streaming started. Matches the Azure ML SDK approach.
+//
+//	duration = MAX / (1.0 + 100 * exp(-elapsed_seconds / 20.0))
+//	return max(MIN, duration)
+func pollInterval(startTime time.Time) time.Duration {
+	const (
+		minInterval = 2 * time.Second
+		maxInterval = 60 * time.Second
+	)
+	elapsed := time.Since(startTime).Seconds()
+	durationSec := 60.0 / (1.0 + 100.0*math.Exp(-elapsed/20.0))
+	duration := time.Duration(durationSec * float64(time.Second))
+	if duration < minInterval {
+		return minInterval
 	}
+	return duration
 }
 
 // filterLogFiles selects streamable log files from the history details.

From 0d7672231d99e3e598b5a02fb3d897abdecc60af Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 20:55:54 +0530
Subject: [PATCH 06/11] Nit: Refactor - rename logs.go to blob.go

---
 .../internal/service/stream_service.go        |  2 +-
 .../pkg/client/blob.go                        | 40 +++++++++++++++++++
 .../pkg/client/logs.go                        | 40 -------------------
 3 files changed, 41 insertions(+), 41 deletions(-)
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go
 delete mode 100644 cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index 907dc436bae..23fdf0dc9bf 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -203,7 +203,7 @@ func (s *StreamService) pollAndPrintLogs(
 	for _, fileName := range fileNames {
 		sasURI := details.LogFiles[fileName]
 
-		content, _, err := s.client.GetLogContent(ctx, sasURI, 0)
+		content, _, err := s.client.GetBlobContent(ctx, sasURI)
 		if err != nil || content == "" {
 			continue
 		}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go
new file mode 100644
index 00000000000..568d9d94b43
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go
@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package client
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+)
+
+// GetBlobContent downloads content from a SAS URI.
+// Fetches the full content on each call. No authentication is needed since the
+// URL contains a SAS token. The caller is responsible for tracking incremental
+// progress (e.g., line counts for log streaming).
+func (c *Client) GetBlobContent(ctx context.Context, sasURI string) (string, int64, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sasURI, nil)
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to create blob request: %w", err)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return "", 0, fmt.Errorf("blob request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", 0, fmt.Errorf("blob request returned status %d", resp.StatusCode)
+	}
+
+	// Cap read to 1MB per call to avoid memory issues with very large blobs
+	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to read blob content: %w", err)
+	}
+
+	return string(body), int64(len(body)), nil
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
deleted file mode 100644
index c6aa2b4ded0..00000000000
--- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/logs.go
+++ /dev/null
@@ -1,40 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-
-package client
-
-import (
-	"context"
-	"fmt"
-	"io"
-	"net/http"
-)
-
-// GetLogContent fetches log file content from a SAS URI.
-// Downloads the full content on each call (no Range header). The caller tracks
-// which lines have already been printed, matching the Azure ML SDK approach.
-// No authentication is needed since the URL contains a SAS token.
-func (c *Client) GetLogContent(ctx context.Context, sasURI string, offset int64) (string, int64, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sasURI, nil)
-	if err != nil {
-		return "", 0, fmt.Errorf("failed to create log request: %w", err)
-	}
-
-	resp, err := c.httpClient.Do(req)
-	if err != nil {
-		return "", 0, fmt.Errorf("log request failed: %w", err)
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return "", 0, fmt.Errorf("log request returned status %d", resp.StatusCode)
-	}
-
-	// Cap read to 1MB per poll to avoid memory issues with very large logs
-	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
-	if err != nil {
-		return "", 0, fmt.Errorf("failed to read log content: %w", err)
-	}
-
-	return string(body), int64(len(body)), nil
-}

From da62d5c76d2e2f8a881219cf74e9637222b6caec Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Tue, 28 Apr 2026 21:13:34 +0530
Subject: [PATCH 07/11] Add header for file - enhancement from AML - for multi
 file streaming, add header every time for better readability

---
 .../internal/service/stream_service.go        | 42 +++++++++++++++----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index 23fdf0dc9bf..cada40fdc90 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -75,6 +75,7 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 	consecutiveErrs := 0
 	trackingEndpoint := ""
 	firstPoll := true
+	multiFile := false // latches to true once >1 file is seen
 
 	for {
 		if ctx.Err() != nil {
@@ -102,7 +103,7 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 
 		if terminalStates[jobStatus] {
 			if trackingEndpoint != "" && !firstPoll {
-				s.flushLogs(ctx, trackingEndpoint, jobName, processedLines)
+				s.flushLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile)
 			}
 			return &StreamResult{
 				JobName:   jobName,
@@ -120,7 +121,11 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 
 		// Stream logs if tracking endpoint is available
 		if trackingEndpoint != "" {
-			_, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
+			var newMultiFile bool
+			_, newMultiFile, err = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile)
+			if newMultiFile {
+				multiFile = true
+			}
 			if err != nil {
 				consecutiveErrs++
 				if consecutiveErrs >= maxConsecutiveErrs {
@@ -179,24 +184,39 @@ func filterLogFiles(logFiles map[string]string) []string {
 	return matched
 }
 
+// printFileHeader prints a visual separator for a log file, matching the Azure ML SDK style.
+func printFileHeader(fileName string) {
+	fmt.Println()
+	fmt.Printf("Streaming %s\n", fileName)
+	fmt.Println(strings.Repeat("=", len("Streaming ")+len(fileName)))
+	fmt.Println()
+}
+
 // pollAndPrintLogs fetches run history details and prints only new log lines.
+// Returns whether new content was printed and whether multiple files were seen.
 func (s *StreamService) pollAndPrintLogs(
 	ctx context.Context,
 	trackingEndpoint string,
 	jobName string,
 	processedLines map[string]int,
-) (bool, error) {
+	multiFile bool,
+) (bool, bool, error) {
 	details, err := s.client.GetRunHistoryDetails(ctx, trackingEndpoint, jobName)
 	if err != nil {
-		return false, err
+		return false, multiFile, err
 	}
 	if details == nil || len(details.LogFiles) == 0 {
-		return false, nil
+		return false, multiFile, nil
 	}
 
 	fileNames := filterLogFiles(details.LogFiles)
 	if len(fileNames) == 0 {
-		return false, nil
+		return false, multiFile, nil
+	}
+
+	// Latch multiFile once we see >1 file — formatting stays consistent thereafter
+	if len(fileNames) > 1 {
+		multiFile = true
 	}
 
 	hasNewContent := false
@@ -219,6 +239,11 @@ func (s *StreamService) pollAndPrintLogs(
 			continue
 		}
 
+		// Print file header: always on first encounter, or on every switch for multi-file jobs
+		if previousLines == 0 || multiFile {
+			printFileHeader(fileName)
+		}
+
 		for _, line := range lines[previousLines:] {
 			fmt.Println(line)
 		}
@@ -226,7 +251,7 @@ func (s *StreamService) pollAndPrintLogs(
 		processedLines[fileName] = len(lines)
 	}
 
-	return hasNewContent, nil
+	return hasNewContent, multiFile, nil
 }
 
 // flushLogs does a final poll to capture any remaining log output.
@@ -235,8 +260,9 @@ func (s *StreamService) flushLogs(
 	trackingEndpoint string,
 	jobName string,
 	processedLines map[string]int,
+	multiFile bool,
 ) {
-	_, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines)
+	_, _, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile)
 }
 
 // extractServiceEndpoint extracts the endpoint URL from the job services map.

From 2e3633492b9c9dac8d93ab70ddf13d589cdbec49 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Wed, 29 Apr 2026 10:00:25 +0530
Subject: [PATCH 08/11] Log warnings on error so users know logs may be
 incomplete

---
 .../internal/service/stream_service.go                | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index cada40fdc90..ac87008287e 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -224,7 +224,11 @@ func (s *StreamService) pollAndPrintLogs(
 		sasURI := details.LogFiles[fileName]
 
 		content, _, err := s.client.GetBlobContent(ctx, sasURI)
-		if err != nil || content == "" {
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Warning: failed to read log file %s: %v\n", fileName, err)
+			continue
+		}
+		if content == "" {
 			continue
 		}
 
@@ -262,7 +266,10 @@ func (s *StreamService) flushLogs(
 	processedLines map[string]int,
 	multiFile bool,
 ) {
-	_, _, _ = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile)
+	_, _, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Warning: failed to flush final logs: %v\n", err)
+	}
 }
 
 // extractServiceEndpoint extracts the endpoint URL from the job services map.

From 6d355db16faebac7e2c8fdf32d42dffa866ccae1 Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Wed, 29 Apr 2026 10:00:42 +0530
Subject: [PATCH 09/11] Add UTs

---
 .../internal/service/stream_service_test.go   | 287 ++++++++++++++++++
 .../pkg/client/history_test.go                |  68 +++++
 2 files changed, 355 insertions(+)
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
 create mode 100644 cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
new file mode 100644
index 00000000000..abbe156f9b8
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
@@ -0,0 +1,287 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package service
+
+import (
+	"testing"
+	"time"
+)
+
+func TestPollInterval_SigmoidBoundaries(t *testing.T) {
+	tests := []struct {
+		name        string
+		elapsedSec  float64
+		wantMinSec  float64
+		wantMaxSec  float64
+		description string
+	}{
+		{
+			name:        "start returns minimum",
+			elapsedSec:  0,
+			wantMinSec:  2.0,
+			wantMaxSec:  2.0,
+			description: "at t=0 the sigmoid value is ~0.59s, clamped to 2s minimum",
+		},
+		{
+			name:        "10 seconds still near minimum",
+			elapsedSec:  10,
+			wantMinSec:  2.0,
+			wantMaxSec:  3.0,
+			description: "early in the curve, interval stays low",
+		},
+		{
+			name:        "60 seconds mid range",
+			elapsedSec:  60,
+			wantMinSec:  5.0,
+			wantMaxSec:  15.0,
+			description: "at 60s the sigmoid is in the transition zone",
+		},
+		{
+			name:        "120 seconds approaching max",
+			elapsedSec:  120,
+			wantMinSec:  45.0,
+			wantMaxSec:  60.0,
+			description: "at 120s the sigmoid approaches the 60s asymptote",
+		},
+		{
+			name:        "300 seconds at max",
+			elapsedSec:  300,
+			wantMinSec:  59.0,
+			wantMaxSec:  60.0,
+			description: "well past the inflection, saturated at max",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			startTime := time.Now().Add(-time.Duration(tt.elapsedSec * float64(time.Second)))
+			got := pollInterval(startTime)
+			gotSec := got.Seconds()
+			if gotSec < tt.wantMinSec || gotSec > tt.wantMaxSec {
+				t.Errorf("pollInterval(elapsed=%vs) = %vs, want [%v, %v]s",
+					tt.elapsedSec, gotSec, tt.wantMinSec, tt.wantMaxSec)
+			}
+		})
+	}
+}
+
+func TestPollInterval_NeverBelowMinimum(t *testing.T) {
+	for elapsed := 0; elapsed <= 300; elapsed += 5 {
+		startTime := time.Now().Add(-time.Duration(elapsed) * time.Second)
+		got := pollInterval(startTime)
+		if got < 2*time.Second {
+			t.Errorf("pollInterval(elapsed=%ds) = %v, below 2s minimum", elapsed, got)
+		}
+		if got > 60*time.Second {
+			t.Errorf("pollInterval(elapsed=%ds) = %v, above 60s maximum", elapsed, got)
+		}
+	}
+}
+
+func TestFilterLogFiles_PrimaryPattern(t *testing.T) {
+	tests := []struct {
+		name     string
+		logFiles map[string]string
+		want     []string
+	}{
+		{
+			name:     "empty map returns nil",
+			logFiles: map[string]string{},
+			want:     nil,
+		},
+		{
+			name: "matches single user_logs std_log",
+			logFiles: map[string]string{
+				"user_logs/std_log.txt": "https://blob.core/std_log.txt?sas",
+			},
+			want: []string{"user_logs/std_log.txt"},
+		},
+		{
+			name: "matches std_log_ps variant",
+			logFiles: map[string]string{
+				"user_logs/std_log_ps.txt": "https://blob.core/std_log_ps.txt?sas",
+			},
+			want: []string{"user_logs/std_log_ps.txt"},
+		},
+		{
+			name: "excludes activity logs and system files",
+			logFiles: map[string]string{
+				"user_logs/std_log.txt":                         "https://blob.core/std_log.txt?sas",
+				"system_logs/virtualcluster_activity_log.txt":   "https://blob.core/activity.txt?sas",
+				"system_logs/cs_capability/cs-capability-0.txt": "https://blob.core/cs.txt?sas",
+			},
+			want: []string{"user_logs/std_log.txt"},
+		},
+		{
+			name: "multiple user_logs sorted alphabetically",
+			logFiles: map[string]string{
+				"user_logs/std_log.txt":    "https://blob.core/std_log.txt?sas",
+				"user_logs/std_log_ps.txt": "https://blob.core/std_log_ps.txt?sas",
+			},
+			want: []string{"user_logs/std_log.txt", "user_logs/std_log_ps.txt"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := filterLogFiles(tt.logFiles)
+			if !slicesEqual(got, tt.want) {
+				t.Errorf("filterLogFiles() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestFilterLogFiles_FallbackToLegacy(t *testing.T) {
+	tests := []struct {
+		name     string
+		logFiles map[string]string
+		want     []string
+	}{
+		{
+			name: "falls back to azureml-logs when no user_logs match",
+			logFiles: map[string]string{
+				"azureml-logs/70_driver_log.txt":  "https://blob.core/70.txt?sas",
+				"azureml-logs/55_azureml-log.txt": "https://blob.core/55.txt?sas",
+			},
+			want: []string{"azureml-logs/55_azureml-log.txt", "azureml-logs/70_driver_log.txt"},
+		},
+		{
+			name: "prefers user_logs over azureml-logs when both exist",
+			logFiles: map[string]string{
+				"user_logs/std_log.txt":          "https://blob.core/std_log.txt?sas",
+				"azureml-logs/70_driver_log.txt": "https://blob.core/70.txt?sas",
+			},
+			want: []string{"user_logs/std_log.txt"},
+		},
+		{
+			name: "no matching files returns nil",
+			logFiles: map[string]string{
+				"system_logs/something.txt": "https://blob.core/sys.txt?sas",
+				"random_file.log":           "https://blob.core/random.txt?sas",
+			},
+			want: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := filterLogFiles(tt.logFiles)
+			if !slicesEqual(got, tt.want) {
+				t.Errorf("filterLogFiles() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestExtractServiceEndpoint(t *testing.T) {
+	tests := []struct {
+		name        string
+		services    map[string]any
+		serviceName string
+		want        string
+	}{
+		{
+			name:        "nil services",
+			services:    nil,
+			serviceName: "Studio",
+			want:        "",
+		},
+		{
+			name:        "missing service",
+			services:    map[string]any{"Tracking": map[string]any{"endpoint": "https://tracking"}},
+			serviceName: "Studio",
+			want:        "",
+		},
+		{
+			name:        "service not a map",
+			services:    map[string]any{"Studio": "not-a-map"},
+			serviceName: "Studio",
+			want:        "",
+		},
+		{
+			name:        "missing endpoint key",
+			services:    map[string]any{"Studio": map[string]any{"other": "value"}},
+			serviceName: "Studio",
+			want:        "",
+		},
+		{
+			name:        "endpoint not a string",
+			services:    map[string]any{"Studio": map[string]any{"endpoint": 42}},
+			serviceName: "Studio",
+			want:        "",
+		},
+		{
+			name: "valid Studio endpoint",
+			services: map[string]any{
+				"Studio": map[string]any{
+					"endpoint": "https://ml.azure.com/runs/my-job?wsid=/sub/rg/ws",
+				},
+			},
+			serviceName: "Studio",
+			want:        "https://ml.azure.com/runs/my-job?wsid=/sub/rg/ws",
+		},
+		{
+			name: "valid Tracking endpoint",
+			services: map[string]any{
+				"Tracking": map[string]any{
+					"endpoint": "azureml://eastus.api.azureml.ms/mlflow/v1.0/sub/rg/ws",
+				},
+			},
+			serviceName: "Tracking",
+			want:        "azureml://eastus.api.azureml.ms/mlflow/v1.0/sub/rg/ws",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := extractServiceEndpoint(tt.services, tt.serviceName)
+			if got != tt.want {
+				t.Errorf("extractServiceEndpoint(%q) = %q, want %q", tt.serviceName, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestTerminalAndActiveStates(t *testing.T) {
+	// Verify terminal states don't overlap with active states
+	for state := range terminalStates {
+		if activeStates[state] {
+			t.Errorf("state %q is both terminal and active", state)
+		}
+	}
+
+	// Verify expected terminal states
+	expectedTerminal := []string{"Completed", "Failed", "Canceled", "NotResponding", "Paused"}
+	for _, s := range expectedTerminal {
+		if !terminalStates[s] {
+			t.Errorf("expected %q to be a terminal state", s)
+		}
+	}
+
+	// Verify expected active states
+	expectedActive := []string{"NotStarted", "Queued", "Preparing", "Provisioning", "Starting", "Running", "Finalizing"}
+	for _, s := range expectedActive {
+		if !activeStates[s] {
+			t.Errorf("expected %q to be an active state", s)
+		}
+	}
+}
+
+// slicesEqual compares two string slices including nil vs empty.
+func slicesEqual(a, b []string) bool {
+	if len(a) == 0 && len(b) == 0 {
+		// Treat nil and empty as equal
+		return (a == nil) == (b == nil) || (len(a) == 0 && len(b) == 0)
+	}
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go
new file mode 100644
index 00000000000..52bfdaa540f
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go
@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package client
+
+import (
+	"testing"
+)
+
+func TestParseTrackingEndpoint(t *testing.T) {
+	tests := []struct {
+		name         string
+		endpoint     string
+		wantBaseURL  string
+		wantWorkPath string
+		wantErr      bool
+	}{
+		{
+			name:         "standard v1.0 endpoint",
+			endpoint:     "azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws",
+			wantBaseURL:  "https://eastus.api.azureml.ms",
+			wantWorkPath: "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws",
+		},
+		{
+			name:         "v2.0 endpoint",
+			endpoint:     "azureml://westus2.api.azureml.ms/mlflow/v2.0/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws",
+			wantBaseURL:  "https://westus2.api.azureml.ms",
+			wantWorkPath: "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws",
+		},
+		{
+			name:         "endpoint with query params stripped",
+			endpoint:     "azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/sub/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws?extra=param",
+			wantBaseURL:  "https://eastus.api.azureml.ms",
+			wantWorkPath: "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws",
+		},
+		{
+			name:     "missing mlflow prefix",
+			endpoint: "azureml://eastus.api.azureml.ms/other/v1.0/subscriptions/sub/resourceGroups/rg",
+			wantErr:  true,
+		},
+		{
+			name:     "empty string",
+			endpoint: "",
+			wantErr:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			baseURL, workPath, err := parseTrackingEndpoint(tt.endpoint)
+			if tt.wantErr {
+				if err == nil {
+					t.Errorf("parseTrackingEndpoint(%q) expected error, got baseURL=%q workPath=%q", tt.endpoint, baseURL, workPath)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("parseTrackingEndpoint(%q) unexpected error: %v", tt.endpoint, err)
+			}
+			if baseURL != tt.wantBaseURL {
+				t.Errorf("baseURL = %q, want %q", baseURL, tt.wantBaseURL)
+			}
+			if workPath != tt.wantWorkPath {
+				t.Errorf("workPath = %q, want %q", workPath, tt.wantWorkPath)
+			}
+		})
+	}
+}

From 636c97cb8aa46b6d068ad512e115dbaeddd1a25f Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Wed, 29 Apr 2026 10:14:12 +0530
Subject: [PATCH 10/11] Nit

---
 cli/azd/extensions/azure.ai.customtraining/extension.yaml       | 2 +-
 cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/extension.yaml b/cli/azd/extensions/azure.ai.customtraining/extension.yaml
index 8874c1a9b8a..aaf39c21014 100644
--- a/cli/azd/extensions/azure.ai.customtraining/extension.yaml
+++ b/cli/azd/extensions/azure.ai.customtraining/extension.yaml
@@ -1,7 +1,7 @@
 id: azure.ai.customtraining
 namespace: ai.training
 displayName: Azure AI Custom Training (Preview)
-description: Extension for Azure AI Foundry custom training jobs. (Preview)
+description: Extension for Microsoft Foundry custom training jobs. (Preview)
 usage: azd ai training <command> [options]
 version: 0.0.1-preview
 language: go
diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go
index 25bc52d4e4f..bfc10f8f485 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go
@@ -17,7 +17,7 @@ var rootFlags rootFlagsDefinition
 func NewRootCommand() *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use:           "training <command> [options]",
-		Short:         "Extension for Azure AI Foundry custom training jobs. (Preview)",
+		Short:         "Extension for Microsoft Foundry custom training jobs. (Preview)",
 		SilenceUsage:  true,
 		SilenceErrors: true,
 		CompletionOptions: cobra.CompletionOptions{

From 5b9865b4c5a9cd33198797f73b58807253f3cd6e Mon Sep 17 00:00:00 2001
From: Saanika Gupta <gsaanika@gmail.com>
Date: Wed, 29 Apr 2026 10:29:09 +0530
Subject: [PATCH 11/11] Honor context cancellation during sleep

---
 .../internal/service/stream_service.go        | 22 +++++++++++--
 .../internal/service/stream_service_test.go   | 31 +++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
index ac87008287e..69cc232bf23 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go
@@ -89,7 +89,9 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 			if consecutiveErrs >= maxConsecutiveErrs {
 				return nil, fmt.Errorf("failed to get job status after %d retries: %w", maxConsecutiveErrs, err)
 			}
-			time.Sleep(pollInterval(startTime))
+			if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil {
+				return nil, err
+			}
 			continue
 		}
 		consecutiveErrs = 0
@@ -114,7 +116,9 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 
 		if !activeStates[jobStatus] {
 			fmt.Fprintf(os.Stderr, "Job status: %s, waiting...\n", jobStatus)
-			time.Sleep(pollInterval(startTime))
+			if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil {
+				return nil, err
+			}
 			firstPoll = false
 			continue
 		}
@@ -139,7 +143,19 @@ func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*Str
 		}
 
 		firstPoll = false
-		time.Sleep(pollInterval(startTime))
+		if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil {
+			return nil, err
+		}
+	}
+}
+
+// sleepWithContext sleeps for the given duration, returning early if ctx is cancelled.
+func sleepWithContext(ctx context.Context, d time.Duration) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-time.After(d):
+		return nil
 	}
 }
 
diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
index abbe156f9b8..cfdea80e757 100644
--- a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
+++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go
@@ -4,6 +4,7 @@
 package service
 
 import (
+	"context"
 	"testing"
 	"time"
 )
@@ -79,6 +80,36 @@ func TestPollInterval_NeverBelowMinimum(t *testing.T) {
 	}
 }
 
+func TestSleepWithContext_CancelledImmediately(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel before sleeping
+
+	start := time.Now()
+	err := sleepWithContext(ctx, 10*time.Second)
+	elapsed := time.Since(start)
+
+	if err == nil {
+		t.Error("expected error from cancelled context, got nil")
+	}
+	if elapsed > 500*time.Millisecond {
+		t.Errorf("sleepWithContext took %v, should have returned immediately", elapsed)
+	}
+}
+
+func TestSleepWithContext_CompletesNormally(t *testing.T) {
+	ctx := context.Background()
+	start := time.Now()
+	err := sleepWithContext(ctx, 50*time.Millisecond)
+	elapsed := time.Since(start)
+
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+	if elapsed < 50*time.Millisecond {
+		t.Errorf("sleepWithContext returned too early: %v", elapsed)
+	}
+}
+
 func TestFilterLogFiles_PrimaryPattern(t *testing.T) {
 	tests := []struct {
 		name     string