discovery: add configurable AWS SSM install timeout (#69102)

docs/pages/includes/auto-discovery/aws-ec2-advanced-config.mdx

@@ -2,6 +2,34 @@
 
 (!docs/pages/includes/auto-discovery/server-advanced-config.mdx matcher="aws"!)
 
+### Configure the SSM installation timeout
+
+By default, the Discovery Service waits up to five minutes for an SSM
+installation attempt. It allows an additional minute for AWS to publish the
+command's terminal state and reserves up to one more minute to collect and
+report the result.
+
+Set `TELEPORT_UNSTABLE_AWS_INSTALL_TIMEOUT` on the Discovery Service to change
+the five-minute timeout. The value must use Go duration syntax, such as
+`3m45s`. Values are limited to a minimum of 10 seconds and a maximum of 90
+minutes. Invalid values use the five-minute default.
+
+For example:
+
+```code
+TELEPORT_UNSTABLE_AWS_INSTALL_TIMEOUT=3m45s
+```
+
+<Admonition type="warning">
+
+This environment variable is unstable and may change in a future release.
+
+The timeout bounds how long the Discovery Service waits for the installation.
+It does not override the execution timeout in a custom SSM document or
+necessarily cancel a command that AWS has already accepted.
+
+</Admonition>
+
 ### Use a custom installation script
 
 (!docs/pages/includes/server-access/custom-installer-aws-config.mdx!)

lib/srv/discovery/status.go

@@ -382,42 +382,39 @@ func (ars *awsResourcesStatus) incrementEnrolled(g awsResourceGroup, count int)
 // ReportEC2SSMInstallationResult is called when discovery gets the result of running the installation script in a EC2 instance.
 // It will emit an audit event with the result and update the DiscoveryConfig status
 func (s *Server) ReportEC2SSMInstallationResult(ctx context.Context, result *server.SSMInstallationResult) error {
-	if err := s.Emitter.EmitAuditEvent(ctx, result.SSMRunEvent); err != nil {
-		return trace.Wrap(err)
-	}
-
 	// Only failed runs are counted.
 	// Successful ones only mean that the teleport was installed in the target host.
 	// If they succeed in joining the cluster, during the next iteration, they will be countd as "enrolled"
-	if result.SSMRunEvent.Metadata.Code == libevents.SSMRunSuccessCode {
-		return nil
+	if result.SSMRunEvent.Metadata.Code != libevents.SSMRunSuccessCode {
+		// Record failure status before emitting the audit event. Audit emission
+		// errors are reported to the caller, but must not hide the failed
+		// enrollment from DiscoveryConfig status or UserTasks.
+		s.awsEC2ResourcesStatus.incrementFailed(awsResourceGroup{
+			discoveryConfigName: result.DiscoveryConfigName,
+			integration:         result.IntegrationName,
+		}, 1)
+
+		s.awsEC2Tasks.addFailedEnrollment(
+			awsEC2TaskKey{
+				integration:     result.IntegrationName,
+				issueType:       result.IssueType,
+				accountID:       result.SSMRunEvent.AccountID,
+				region:          result.SSMRunEvent.Region,
+				ssmDocument:     result.SSMDocumentName,
+				installerScript: result.InstallerScript,
+			},
+			usertasksv1.DiscoverEC2Instance_builder{
+				InvocationUrl:   result.SSMRunEvent.InvocationURL,
+				DiscoveryConfig: result.DiscoveryConfigName,
+				DiscoveryGroup:  s.DiscoveryGroup,
+				SyncTime:        timestamppb.New(s.clock.Now()),
+				InstanceId:      result.SSMRunEvent.InstanceID,
+				Name:            result.InstanceName,
+			}.Build(),
+		)
 	}
 
-	s.awsEC2ResourcesStatus.incrementFailed(awsResourceGroup{
-		discoveryConfigName: result.DiscoveryConfigName,
-		integration:         result.IntegrationName,
-	}, 1)
-
-	s.awsEC2Tasks.addFailedEnrollment(
-		awsEC2TaskKey{
-			integration:     result.IntegrationName,
-			issueType:       result.IssueType,
-			accountID:       result.SSMRunEvent.AccountID,
-			region:          result.SSMRunEvent.Region,
-			ssmDocument:     result.SSMDocumentName,
-			installerScript: result.InstallerScript,
-		},
-		usertasksv1.DiscoverEC2Instance_builder{
-			InvocationUrl:   result.SSMRunEvent.InvocationURL,
-			DiscoveryConfig: result.DiscoveryConfigName,
-			DiscoveryGroup:  s.DiscoveryGroup,
-			SyncTime:        timestamppb.New(s.clock.Now()),
-			InstanceId:      result.SSMRunEvent.InstanceID,
-			Name:            result.InstanceName,
-		}.Build(),
-	)
-
-	return nil
+	return trace.Wrap(s.Emitter.EmitAuditEvent(ctx, result.SSMRunEvent))
 }
 
 // awsEC2Tasks contains the Discover EC2 User Tasks that must be reported to the user.

lib/srv/discovery/status_test.go

@@ -20,6 +20,7 @@ package discovery
 
 import (
 	"context"
+	"errors"
 	"maps"
 	"slices"
 	"strings"
@@ -37,10 +38,74 @@ import (
 	usertasksv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/usertasks/v1"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/types/discoveryconfig"
+	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/api/types/usertasks"
+	libevents "github.com/gravitational/teleport/lib/events"
+	"github.com/gravitational/teleport/lib/srv/server"
 	"github.com/gravitational/teleport/lib/utils/log/logtest"
 )
 
+type failingEmitter struct {
+	err error
+}
+
+func (e failingEmitter) EmitAuditEvent(context.Context, apievents.AuditEvent) error {
+	return e.err
+}
+
+func TestReportEC2SSMInstallationResultRecordsFailureOnAuditError(t *testing.T) {
+	emitErr := errors.New("audit emission failed")
+	clock := clockwork.NewFakeClock()
+	group := awsResourceGroup{
+		discoveryConfigName: "discovery-config",
+		integration:         "integration",
+	}
+	s := &Server{
+		Config: &Config{
+			Emitter:        failingEmitter{err: emitErr},
+			DiscoveryGroup: "discovery-group",
+			clock:          clock,
+		},
+		awsEC2ResourcesStatus: newAWSResourceStatusCollector("ec2"),
+	}
+	s.awsEC2ResourcesStatus.iterationStarted([]awsResourceGroup{group}, clock.Now())
+
+	result := &server.SSMInstallationResult{
+		SSMRunEvent: &apievents.SSMRun{
+			Metadata: apievents.Metadata{
+				Code: libevents.SSMRunFailCode,
+			},
+			AccountID:     "123456789012",
+			Region:        "eu-west-2",
+			InstanceID:    "i-1234567890abcdef0",
+			InvocationURL: "https://example.com/invocation",
+		},
+		IntegrationName:     group.integration,
+		DiscoveryConfigName: group.discoveryConfigName,
+		IssueType:           usertasks.AutoDiscoverEC2IssueSSMInvocationFailure,
+		SSMDocumentName:     "document",
+		InstallerScript:     "installer",
+		InstanceName:        "instance",
+	}
+
+	err := s.ReportEC2SSMInstallationResult(t.Context(), result)
+	require.ErrorIs(t, err, emitErr)
+	require.Equal(t, 1, s.awsEC2ResourcesStatus.awsResourcesResults[group].failed)
+
+	taskKey := awsEC2TaskKey{
+		integration:     result.IntegrationName,
+		issueType:       result.IssueType,
+		accountID:       result.SSMRunEvent.AccountID,
+		region:          result.SSMRunEvent.Region,
+		ssmDocument:     result.SSMDocumentName,
+		installerScript: result.InstallerScript,
+	}
+	task := s.awsEC2Tasks.instancesIssues[taskKey]
+	require.NotNil(t, task)
+	require.Contains(t, task.GetInstances(), result.SSMRunEvent.InstanceID)
+	require.Contains(t, s.awsEC2Tasks.issuesSyncQueue, taskKey)
+}
+
 func TestTruncateErrorMessage(t *testing.T) {
 	for _, tt := range []struct {
 		name     string

lib/srv/server/ssm_install.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"log/slog"
 	"maps"
+	"os"
 	"slices"
 	"strings"
 	"time"
@@ -45,6 +46,18 @@ import (
 )
 
 const (
+	// awsInstallTimeoutEnvVar configures how long an AWS SSM installation
+	// attempt may execute before timing out.
+	awsInstallTimeoutEnvVar = "TELEPORT_UNSTABLE_AWS_INSTALL_TIMEOUT"
+
+	defaultAWSInstallTimeout = 5 * time.Minute
+	minAWSInstallTimeout     = 10 * time.Second
+	maxAWSInstallTimeout     = 90 * time.Minute
+	// awsSSMResultGracePeriod gives SSM time to publish the terminal command
+	// status and Teleport time to collect the final result after installation
+	// timeout. It is applied once for each phase.
+	awsSSMResultGracePeriod = time.Minute
+
 	// waiterTimedOutErrorMessage is the error message returned by the AWS SDK command
 	// executed waiter when it times out.
 	waiterTimedOutErrorMessage = "exceeded max wait time for CommandExecuted waiter"
@@ -53,19 +66,6 @@ const (
 	//nolint:misspell // ignore Cancelled and Cancelling
 	// executed waiter when the command state transitions to one of Cancelled, TimedOut, Failed or Cancelling.
 	waiterTransitionedToFailureErrorMessage = "waiter state transitioned to Failure"
-	// waitTimeoutPad is extra waiter headroom beyond the installer's join-failure timeout.
-	// The installer decides success/failure based on whether join completes in time.
-	// This pad avoids reporting a temporary "still running" SSM state as the final outcome.
-	// Tradeoff: EC2 install handling is still synchronous in discovery, so a wedged
-	// SSM invocation can delay later groups until waitTimeout elapses.
-	//
-	// TODO(carlisia): Decouple SSM result waiting from synchronous discovery
-	// iterations so a wedged invocation does not stall later EC2 groups.
-	waitTimeoutPad = 10 * time.Minute
-
-	// waitTimeout is how long we wait for AWS to report a terminal command state
-	// for the installer result.
-	waitTimeout = installer.JoinFailureTimeout + waitTimeoutPad
 
 	// maxSSMRunOutputChars limits stdout/stderr size while preserving the most recent diagnostics.
 	// 24_000 matches the documented per-field cap for SSMRun stdout/stderr in the event schema,
@@ -73,6 +73,14 @@ const (
 	maxSSMRunOutputChars = 24_000
 )
 
+func getAWSInstallTimeout() time.Duration {
+	if timeout, err := time.ParseDuration(os.Getenv(awsInstallTimeoutEnvVar)); err == nil {
+		// Keep the timeout bounds consistent with Azure VM installation timeouts.
+		return min(maxAWSInstallTimeout, max(minAWSInstallTimeout, timeout))
+	}
+	return defaultAWSInstallTimeout
+}
+
 // SSMClient is the subset of the AWS SSM API required for EC2 discovery.
 type SSMClient interface {
 	ssm.DescribeInstanceInformationAPIClient
@@ -231,8 +239,21 @@ func (si *SSMInstaller) Run(ctx context.Context, req SSMRunRequest) error {
 		return nil
 	}
 
+	installTimeout := getAWSInstallTimeout()
+	waiterMaxWait := installTimeout + awsSSMResultGracePeriod
+	// EC2 install handling remains synchronous in discovery, so a wedged SSM
+	// invocation can delay later groups until the configured timeout and SSM
+	// status headroom elapse.
+	// TODO(carlisia): Decouple SSM result waiting from synchronous discovery
+	// iterations so a wedged invocation does not stall later EC2 groups.
+	// Give SSM time to publish a terminal state after the install timeout, then
+	// keep the context alive beyond the waiter so Teleport can collect and
+	// report that state.
+	installCtx, cancel := context.WithTimeout(ctx, waiterMaxWait+awsSSMResultGracePeriod)
+	defer cancel()
+
 	validInstanceIDs := instanceIDsFrom(validInstances)
-	output, err := req.SSM.SendCommand(ctx, &ssm.SendCommandInput{
+	output, err := req.SSM.SendCommand(installCtx, &ssm.SendCommandInput{
 		DocumentName: aws.String(req.DocumentName),
 		InstanceIds:  validInstanceIDs,
 		Parameters:   params,
@@ -260,7 +281,7 @@ func (si *SSMInstaller) Run(ctx context.Context, req SSMRunRequest) error {
 		// As a best effort, we try to call ssm.SendCommand again but this time without the "sshdConfigPath" param
 		// We must not remove the Param "sshdConfigPath" beforehand because customers might be using custom SSM Documents for ec2 auto discovery.
 		delete(params, ParamSSHDConfigPath)
-		output, err = req.SSM.SendCommand(ctx, &ssm.SendCommandInput{
+		output, err = req.SSM.SendCommand(installCtx, &ssm.SendCommandInput{
 			DocumentName: aws.String(req.DocumentName),
 			InstanceIds:  validInstanceIDs,
 			Parameters:   params,
@@ -270,8 +291,11 @@ func (si *SSMInstaller) Run(ctx context.Context, req SSMRunRequest) error {
 		}
 	}
 
-	g, ctx := errgroup.WithContext(ctx)
-	g.SetLimit(10)
+	installDeadline, _ := installCtx.Deadline()
+	g, groupCtx := errgroup.WithContext(ctx)
+	// SendCommand accepts at most awsEC2APIChunkSize instances. Start every
+	// waiter in the batch so the configured timeout bounds the whole batch.
+	g.SetLimit(awsEC2APIChunkSize)
 	for instanceID, instanceName := range validInstances {
 		metadata := instanceMetadata{
 			InstanceID:   instanceID,
@@ -283,7 +307,40 @@ func (si *SSMInstaller) Run(ctx context.Context, req SSMRunRequest) error {
 			}
 		}
 		g.Go(func() error {
-			return trace.Wrap(si.checkCommand(ctx, req, output.Command.CommandId, metadata))
+			log := si.Logger.With(
+				"command_id", aws.ToString(output.Command.CommandId),
+				"instance_id", metadata.InstanceID,
+			)
+			// Give every instance an independent context with the same absolute
+			// deadline. A local waiter or collection failure must not cancel the
+			// other instances in this SendCommand batch.
+			instanceCtx, cancel := context.WithDeadline(groupCtx, installDeadline)
+			defer cancel()
+
+			result, err := si.checkCommand(instanceCtx, req, output.Command.CommandId, metadata, waiterMaxWait)
+			if err != nil {
+				// Preserve caller cancellation as a Run error. Once SendCommand has
+				// succeeded, all other errors describe this instance only.
+				if ctxErr := groupCtx.Err(); ctxErr != nil {
+					return trace.Wrap(ctxErr)
+				}
+
+				log.WarnContext(ctx,
+					"Failed to collect SSM installation result",
+					"error", err,
+				)
+				result = failedSSMCommandInstallationResult(req, output.Command.CommandId, metadata)
+			}
+
+			if err := si.ReportSSMInstallationResultFunc(groupCtx, result); err != nil {
+				// Reporting is also instance-local. Returning this error would cause
+				// discovery to mark every instance in the batch as failed again.
+				log.ErrorContext(ctx,
+					"Failed to report SSM installation result",
+					"error", err,
+				)
+			}
+			return nil
 		})
 	}
 	return trace.Wrap(g.Wait())
@@ -323,6 +380,17 @@ func invalidSSMInstanceInstallationResult(req SSMRunRequest, instanceMetadata in
 	}
 }
 
+func failedSSMCommandInstallationResult(req SSMRunRequest, commandID *string, instanceMetadata instanceMetadata) *SSMInstallationResult {
+	result := invalidSSMInstanceInstallationResult(
+		req,
+		instanceMetadata,
+		"Failed to collect the SSM command result.",
+		usertasks.AutoDiscoverEC2IssueSSMInvocationFailure,
+	)
+	result.SSMRunEvent.CommandID = aws.ToString(commandID)
+	return result
+}
+
 func (si *SSMInstaller) emitInvalidInstanceEvents(ctx context.Context, req SSMRunRequest, instanceIDsState *instanceIDsSSMState) error {
 	var errs []error
 	for instanceID, instanceName := range instanceIDsState.missing {
@@ -452,11 +520,11 @@ func (si *SSMInstaller) describeSSMAgentState(ctx context.Context, req SSMRunReq
 	return ret, nil
 }
 
-func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, commandID *string, instanceMetadata instanceMetadata) error {
+func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, commandID *string, instanceMetadata instanceMetadata, waiterMaxWait time.Duration) (*SSMInstallationResult, error) {
 	err := si.getWaiter(req.SSM).Wait(ctx, &ssm.GetCommandInvocationInput{
 		CommandId:  commandID,
 		InstanceId: aws.String(instanceMetadata.InstanceID),
-	}, waitTimeout)
+	}, waiterMaxWait)
 	switch {
 	case err == nil:
 		// Command executed successfully.
@@ -472,9 +540,14 @@ func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, com
 		// The command might still be Pending or InProgress.
 		// Ignoring this error allows us to report that status back to the user.
 
+	case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil:
+		// The waiter can return a wrapped context deadline error when its
+		// internal timeout expires. The outer context remains available for
+		// collecting and reporting this instance's current invocation state.
+
 	default:
 		// For every other unknown error, return the error.
-		return trace.Wrap(err)
+		return nil, trace.Wrap(err)
 	}
 
 	invocationSteps, err := si.getInvocationSteps(ctx, req, commandID, aws.String(instanceMetadata.InstanceID))
@@ -491,7 +564,7 @@ func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, com
 		invocationSteps = awslib.EC2DiscoverySSMDocumentSteps
 
 	case err != nil:
-		return trace.Wrap(err)
+		return nil, trace.Wrap(err)
 	}
 
 	for i, step := range invocationSteps {
@@ -504,22 +577,22 @@ func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, com
 				// If that's the case, emit an event with the overall invocation result (ignoring specific steps' stdout and stderr).
 				outcome, err = si.getCommandStepOutcome(ctx, "" /*no step*/, req, commandID, instanceMetadata)
 				if err != nil {
-					return trace.Wrap(err)
+					return nil, trace.Wrap(err)
 				}
-				return trace.Wrap(si.reportCommandStepOutcome(ctx, req, instanceMetadata, outcome))
+				return commandStepInstallationResult(req, instanceMetadata, outcome), nil
 			}
 
-			return trace.Wrap(err)
+			return nil, trace.Wrap(err)
 		}
 
 		// Emit an event for the first failed step or for the latest step.
 		lastStep := i+1 == len(invocationSteps)
 		if outcome.SSMRunEvent.Metadata.Code != libevents.SSMRunSuccessCode || lastStep {
-			return trace.Wrap(si.reportCommandStepOutcome(ctx, req, instanceMetadata, outcome))
+			return commandStepInstallationResult(req, instanceMetadata, outcome), nil
 		}
 	}
 
-	return nil
+	return nil, trace.BadParameter("no command invocation steps were found")
 }
 
 func (si *SSMInstaller) getInvocationSteps(ctx context.Context, req SSMRunRequest, commandID, instanceID *string) ([]string, error) {
@@ -579,8 +652,8 @@ type commandStepOutcome struct {
 	IssueType   string
 }
 
-func (si *SSMInstaller) reportCommandStepOutcome(ctx context.Context, req SSMRunRequest, instanceMetadata instanceMetadata, outcome commandStepOutcome) error {
-	return si.ReportSSMInstallationResultFunc(ctx, &SSMInstallationResult{
+func commandStepInstallationResult(req SSMRunRequest, instanceMetadata instanceMetadata, outcome commandStepOutcome) *SSMInstallationResult {
+	return &SSMInstallationResult{
 		SSMRunEvent:         outcome.SSMRunEvent,
 		IntegrationName:     req.IntegrationName,
 		DiscoveryConfigName: req.DiscoveryConfigName,
@@ -588,7 +661,7 @@ func (si *SSMInstaller) reportCommandStepOutcome(ctx context.Context, req SSMRun
 		SSMDocumentName:     req.DocumentName,
 		InstallerScript:     req.InstallerScriptName(),
 		InstanceName:        instanceMetadata.InstanceName,
-	})
+	}
 }
 
 func (si *SSMInstaller) getCommandStepOutcome(ctx context.Context, step string, req SSMRunRequest, commandID *string, instanceMetadata instanceMetadata) (commandStepOutcome, error) {

lib/srv/server/ssm_install_test.go

@@ -23,7 +23,9 @@ import (
 	"fmt"
 	"slices"
 	"strings"
+	"sync"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -32,6 +34,7 @@ import (
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
 
+	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/api/types/usertasks"
 	libevent "github.com/gravitational/teleport/lib/events"
@@ -41,11 +44,18 @@ import (
 
 type mockSSMClient struct {
 	SSMClient
-	commandOutput          *ssm.SendCommandOutput
-	waiterTimeout          bool
-	commandInvokeOutput    map[string]*ssm.GetCommandInvocationOutput
-	describeOutput         *ssm.DescribeInstanceInformationOutput
-	listCommandInvocations *ssm.ListCommandInvocationsOutput
+	commandOutput                    *ssm.SendCommandOutput
+	waiterTimeout                    bool
+	waiterMaxWaitError               error
+	waitForContext                   bool
+	waitForContextInstance           map[string]bool
+	waiterCompletionDelay            time.Duration
+	waiterStarted                    chan struct{}
+	getCommandWaitForContextInstance map[string]bool
+	commandInvokeOutput              map[string]*ssm.GetCommandInvocationOutput
+	commandInvokeByInstance          map[string]*ssm.GetCommandInvocationOutput
+	describeOutput                   *ssm.DescribeInstanceInformationOutput
+	listCommandInvocations           *ssm.ListCommandInvocationsOutput
 }
 
 func TestTrimToRecentTail(t *testing.T) {
@@ -118,6 +128,46 @@ const docWithoutSSHDConfigPathParam = "ssmdocument-without-sshdConfigPath-param"
 
 const docWithoutEnvParam = "ssmdocument-without-env-param"
 
+func TestGetAWSInstallTimeout(t *testing.T) {
+	tests := []struct {
+		name  string
+		value string
+		want  time.Duration
+	}{
+		{
+			name: "unset uses default",
+			want: defaultAWSInstallTimeout,
+		},
+		{
+			name:  "valid duration",
+			value: "3m45s",
+			want:  3*time.Minute + 45*time.Second,
+		},
+		{
+			name:  "invalid duration uses default",
+			value: "invalid",
+			want:  defaultAWSInstallTimeout,
+		},
+		{
+			name:  "duration is clamped to minimum",
+			value: "1s",
+			want:  minAWSInstallTimeout,
+		},
+		{
+			name:  "duration is clamped to maximum",
+			value: "2h",
+			want:  maxAWSInstallTimeout,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv(awsInstallTimeoutEnvVar, tt.value)
+			require.Equal(t, tt.want, getAWSInstallTimeout())
+		})
+	}
+}
+
 func (sm *mockSSMClient) SendCommand(_ context.Context, input *ssm.SendCommandInput, _ ...func(*ssm.Options)) (*ssm.SendCommandOutput, error) {
 	if _, hasExtraParam := input.Parameters["sshdConfigPath"]; hasExtraParam && aws.ToString(input.DocumentName) == docWithoutSSHDConfigPathParam {
 		return nil, fmt.Errorf("InvalidParameters: document %s does not support parameters", docWithoutSSHDConfigPathParam)
@@ -130,7 +180,14 @@ func (sm *mockSSMClient) SendCommand(_ context.Context, input *ssm.SendCommandIn
 	return sm.commandOutput, nil
 }
 
-func (sm *mockSSMClient) GetCommandInvocation(_ context.Context, input *ssm.GetCommandInvocationInput, _ ...func(*ssm.Options)) (*ssm.GetCommandInvocationOutput, error) {
+func (sm *mockSSMClient) GetCommandInvocation(ctx context.Context, input *ssm.GetCommandInvocationInput, _ ...func(*ssm.Options)) (*ssm.GetCommandInvocationOutput, error) {
+	if sm.getCommandWaitForContextInstance[aws.ToString(input.InstanceId)] {
+		<-ctx.Done()
+		return nil, ctx.Err()
+	}
+	if result, found := sm.commandInvokeByInstance[aws.ToString(input.InstanceId)]; found {
+		return result, nil
+	}
 	if stepResult, found := sm.commandInvokeOutput[aws.ToString(input.PluginName)]; found {
 		return stepResult, nil
 	}
@@ -152,6 +209,31 @@ func (sm *mockSSMClient) ListCommandInvocations(_ context.Context, input *ssm.Li
 }
 
 func (sm *mockSSMClient) Wait(ctx context.Context, params *ssm.GetCommandInvocationInput, maxWaitDur time.Duration, optFns ...func(*ssm.CommandExecutedWaiterOptions)) error {
+	if sm.waiterStarted != nil {
+		sm.waiterStarted <- struct{}{}
+	}
+
+	if sm.waiterCompletionDelay > 0 {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(sm.waiterCompletionDelay):
+			return nil
+		}
+	}
+
+	if sm.waitForContext || sm.waitForContextInstance[aws.ToString(params.InstanceId)] {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(maxWaitDur):
+			if sm.waiterMaxWaitError != nil {
+				return sm.waiterMaxWaitError
+			}
+			return trace.Errorf(waiterTimedOutErrorMessage)
+		}
+	}
+
 	if sm.waiterTimeout {
 		return trace.Errorf(waiterTimedOutErrorMessage)
 	}
@@ -169,11 +251,375 @@ func (sm *mockSSMClient) Wait(ctx context.Context, params *ssm.GetCommandInvocat
 	return nil
 }
 
+func TestSSMInstallerHungCommandReportsPerInstanceResults(t *testing.T) {
+	const installTimeout = 3*time.Minute + 45*time.Second
+	t.Setenv(awsInstallTimeoutEnvVar, installTimeout.String())
+
+	synctest.Test(t, func(t *testing.T) {
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waitForContextInstance: map[string]bool{
+				"instance-hung": true,
+			},
+			commandInvokeByInstance: map[string]*ssm.GetCommandInvocationOutput{
+				"instance-success": {
+					Status:       ssmtypes.CommandInvocationStatusSuccess,
+					ResponseCode: 0,
+				},
+				"instance-hung": {
+					Status:       ssmtypes.CommandInvocationStatusInProgress,
+					ResponseCode: -1,
+				},
+			},
+		}
+
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		err = inst.Run(t.Context(), SSMRunRequest{
+			DocumentName: types.AWSSSMDocumentRunShellScript,
+			SSM:          client,
+			Instances: []EC2Instance{
+				{InstanceID: "instance-success"},
+				{InstanceID: "instance-hung"},
+			},
+		})
+		require.NoError(t, err)
+
+		require.Len(t, installationResults.installations, 2)
+		resultsByInstance := make(map[string]*SSMInstallationResult, 2)
+		for _, result := range installationResults.installations {
+			resultsByInstance[result.SSMRunEvent.InstanceID] = result
+		}
+		require.Equal(t, string(ssmtypes.CommandInvocationStatusSuccess), resultsByInstance["instance-success"].SSMRunEvent.Status)
+		require.Equal(t, libevent.SSMRunSuccessCode, resultsByInstance["instance-success"].SSMRunEvent.Metadata.Code)
+		require.Equal(t, string(ssmtypes.CommandInvocationStatusInProgress), resultsByInstance["instance-hung"].SSMRunEvent.Status)
+		require.Equal(t, libevent.SSMRunFailCode, resultsByInstance["instance-hung"].SSMRunEvent.Metadata.Code)
+	})
+}
+
+func TestSSMInstallerBatchCompletesWithinSingleWaitPeriod(t *testing.T) {
+	const installTimeout = 3*time.Minute + 45*time.Second
+	t.Setenv(awsInstallTimeoutEnvVar, installTimeout.String())
+
+	synctest.Test(t, func(t *testing.T) {
+		start := time.Now()
+		instances := make([]EC2Instance, 0, awsEC2APIChunkSize)
+		commandInvokeByInstance := make(map[string]*ssm.GetCommandInvocationOutput, awsEC2APIChunkSize)
+		for i := range awsEC2APIChunkSize {
+			instanceID := fmt.Sprintf("instance-%d", i)
+			instances = append(instances, EC2Instance{InstanceID: instanceID})
+			commandInvokeByInstance[instanceID] = &ssm.GetCommandInvocationOutput{
+				Status:       ssmtypes.CommandInvocationStatusInProgress,
+				ResponseCode: -1,
+			}
+		}
+
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waitForContext:          true,
+			commandInvokeByInstance: commandInvokeByInstance,
+		}
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		err = inst.Run(t.Context(), SSMRunRequest{
+			DocumentName: types.AWSSSMDocumentRunShellScript,
+			SSM:          client,
+			Instances:    instances,
+		})
+		require.NoError(t, err)
+		require.Equal(t, installTimeout+awsSSMResultGracePeriod, time.Since(start))
+		require.Len(t, installationResults.installations, awsEC2APIChunkSize)
+
+		for _, result := range installationResults.installations {
+			require.Equal(t, string(ssmtypes.CommandInvocationStatusInProgress), result.SSMRunEvent.Status)
+			require.Equal(t, libevent.SSMRunFailCode, result.SSMRunEvent.Metadata.Code)
+		}
+	})
+}
+
+func TestSSMInstallerCollectionTimeoutIsInstanceLocal(t *testing.T) {
+	const installTimeout = 3*time.Minute + 45*time.Second
+	const blockedInstanceID = "instance-blocked"
+	t.Setenv(awsInstallTimeoutEnvVar, installTimeout.String())
+
+	synctest.Test(t, func(t *testing.T) {
+		start := time.Now()
+		instances := make([]EC2Instance, 0, awsEC2APIChunkSize)
+		commandInvokeByInstance := make(map[string]*ssm.GetCommandInvocationOutput, awsEC2APIChunkSize)
+		for i := range awsEC2APIChunkSize - 1 {
+			instanceID := fmt.Sprintf("instance-%d", i)
+			instances = append(instances, EC2Instance{InstanceID: instanceID})
+			commandInvokeByInstance[instanceID] = &ssm.GetCommandInvocationOutput{
+				Status:       ssmtypes.CommandInvocationStatusSuccess,
+				ResponseCode: 0,
+			}
+		}
+		instances = append(instances, EC2Instance{InstanceID: blockedInstanceID})
+
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waitForContext: true,
+			getCommandWaitForContextInstance: map[string]bool{
+				blockedInstanceID: true,
+			},
+			commandInvokeByInstance: commandInvokeByInstance,
+		}
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		err = inst.Run(t.Context(), SSMRunRequest{
+			DocumentName: types.AWSSSMDocumentRunShellScript,
+			SSM:          client,
+			Instances:    instances,
+		})
+		require.NoError(t, err)
+		require.Equal(t, installTimeout+2*awsSSMResultGracePeriod, time.Since(start))
+
+		require.Len(t, installationResults.installations, awsEC2APIChunkSize)
+		resultsByInstance := make(map[string]*SSMInstallationResult, awsEC2APIChunkSize)
+		for _, result := range installationResults.installations {
+			instanceID := result.SSMRunEvent.InstanceID
+			require.NotContains(t, resultsByInstance, instanceID)
+			resultsByInstance[instanceID] = result
+		}
+
+		for _, instance := range instances {
+			result := resultsByInstance[instance.InstanceID]
+			require.NotNil(t, result)
+			if instance.InstanceID == blockedInstanceID {
+				require.Equal(t, libevent.SSMRunFailCode, result.SSMRunEvent.Metadata.Code)
+				require.Equal(t, int64(-1), result.SSMRunEvent.ExitCode)
+				require.Equal(t, usertasks.AutoDiscoverEC2IssueSSMInvocationFailure, result.IssueType)
+				continue
+			}
+
+			require.Equal(t, libevent.SSMRunSuccessCode, result.SSMRunEvent.Metadata.Code)
+			require.Equal(t, string(ssmtypes.CommandInvocationStatusSuccess), result.SSMRunEvent.Status)
+		}
+	})
+}
+
+func TestSSMInstallerReportErrorIsInstanceLocal(t *testing.T) {
+	client := &mockSSMClient{
+		commandOutput: &ssm.SendCommandOutput{
+			Command: &ssmtypes.Command{
+				CommandId: aws.String("command-id"),
+			},
+		},
+		commandInvokeByInstance: map[string]*ssm.GetCommandInvocationOutput{
+			"instance-report-error": {
+				Status:       ssmtypes.CommandInvocationStatusSuccess,
+				ResponseCode: 0,
+			},
+			"instance-success": {
+				Status:       ssmtypes.CommandInvocationStatusSuccess,
+				ResponseCode: 0,
+			},
+		},
+	}
+
+	var mu sync.Mutex
+	reportedInstances := make(map[string]int, 2)
+	inst, err := NewSSMInstaller(SSMInstallerConfig{
+		ReportSSMInstallationResultFunc: func(_ context.Context, result *SSMInstallationResult) error {
+			mu.Lock()
+			defer mu.Unlock()
+			reportedInstances[result.SSMRunEvent.InstanceID]++
+			if result.SSMRunEvent.InstanceID == "instance-report-error" {
+				return fmt.Errorf("report failed")
+			}
+			return nil
+		},
+		getWaiter: func(SSMClient) commandWaiter {
+			return client
+		},
+	})
+	require.NoError(t, err)
+
+	err = inst.Run(t.Context(), SSMRunRequest{
+		DocumentName: types.AWSSSMDocumentRunShellScript,
+		SSM:          client,
+		Instances: []EC2Instance{
+			{InstanceID: "instance-report-error"},
+			{InstanceID: "instance-success"},
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, map[string]int{
+		"instance-report-error": 1,
+		"instance-success":      1,
+	}, reportedInstances)
+}
+
+func TestSSMInstallerWaitsThroughStatusHeadroom(t *testing.T) {
+	const installTimeout = 3*time.Minute + 45*time.Second
+	const terminalStatusDelay = installTimeout + 30*time.Second
+	t.Setenv(awsInstallTimeoutEnvVar, installTimeout.String())
+
+	synctest.Test(t, func(t *testing.T) {
+		start := time.Now()
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waiterCompletionDelay: terminalStatusDelay,
+			commandInvokeByInstance: map[string]*ssm.GetCommandInvocationOutput{
+				"instance-id": {
+					Status:       ssmtypes.CommandInvocationStatusTimedOut,
+					ResponseCode: -1,
+				},
+			},
+		}
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		err = inst.Run(t.Context(), SSMRunRequest{
+			DocumentName: types.AWSSSMDocumentRunShellScript,
+			SSM:          client,
+			Instances: []EC2Instance{
+				{InstanceID: "instance-id"},
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, terminalStatusDelay, time.Since(start))
+
+		require.Len(t, installationResults.installations, 1)
+		require.Equal(t, string(ssmtypes.CommandInvocationStatusTimedOut), installationResults.installations[0].SSMRunEvent.Status)
+		require.Equal(t, libevent.SSMRunFailCode, installationResults.installations[0].SSMRunEvent.Metadata.Code)
+	})
+}
+
+func TestSSMInstallerCallerCancellation(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waitForContext: true,
+			waiterStarted:  make(chan struct{}, 1),
+		}
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		errCh := make(chan error, 1)
+		go func() {
+			errCh <- inst.Run(ctx, SSMRunRequest{
+				DocumentName: types.AWSSSMDocumentRunShellScript,
+				SSM:          client,
+				Instances: []EC2Instance{
+					{InstanceID: "instance-id"},
+				},
+			})
+		}()
+
+		<-client.waiterStarted
+		cancel()
+		require.ErrorIs(t, <-errCh, context.Canceled)
+		require.Empty(t, installationResults.installations)
+	})
+}
+
+func TestSSMInstallerHonorsTimeoutAboveLegacyWaiterLimit(t *testing.T) {
+	t.Setenv(awsInstallTimeoutEnvVar, maxAWSInstallTimeout.String())
+
+	synctest.Test(t, func(t *testing.T) {
+		start := time.Now()
+		client := &mockSSMClient{
+			commandOutput: &ssm.SendCommandOutput{
+				Command: &ssmtypes.Command{
+					CommandId: aws.String("command-id"),
+				},
+			},
+			waitForContext:     true,
+			waiterMaxWaitError: fmt.Errorf("request canceled while waiting: %w", context.DeadlineExceeded),
+			commandInvokeByInstance: map[string]*ssm.GetCommandInvocationOutput{
+				"instance-id": {
+					Status:       ssmtypes.CommandInvocationStatusInProgress,
+					ResponseCode: -1,
+				},
+			},
+		}
+		installationResults := &mockInstallationResults{}
+		inst, err := NewSSMInstaller(SSMInstallerConfig{
+			ReportSSMInstallationResultFunc: installationResults.ReportInstallationResult,
+			getWaiter: func(SSMClient) commandWaiter {
+				return client
+			},
+		})
+		require.NoError(t, err)
+
+		err = inst.Run(t.Context(), SSMRunRequest{
+			DocumentName: types.AWSSSMDocumentRunShellScript,
+			SSM:          client,
+			Instances: []EC2Instance{
+				{InstanceID: "instance-id"},
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, maxAWSInstallTimeout+awsSSMResultGracePeriod, time.Since(start))
+
+		require.Len(t, installationResults.installations, 1)
+		require.Equal(t, string(ssmtypes.CommandInvocationStatusInProgress), installationResults.installations[0].SSMRunEvent.Status)
+	})
+}
+
 type mockInstallationResults struct {
+	mu            sync.Mutex
 	installations []*SSMInstallationResult
 }
 
 func (me *mockInstallationResults) ReportInstallationResult(ctx context.Context, result *SSMInstallationResult) error {
+	me.mu.Lock()
+	defer me.mu.Unlock()
 	me.installations = append(me.installations, result)
 	return nil
 }