diff --git a/containerm/mgr/mgr_internal_init.go b/containerm/mgr/mgr_internal_init.go
index db9338e..3c545b5 100644
--- a/containerm/mgr/mgr_internal_init.go
+++ b/containerm/mgr/mgr_internal_init.go
@@ -47,7 +47,6 @@ func newContainerMgr(metaPath string, execPath string, defaultCtrsStopTimeout ti
 		restartCtrsMgrCache:    newRestartMgrCache(),
 		containerRepository:    &ctrRepository,
 	}
-	fmt.Println("mgr:", defaultCtrsStopTimeout.String())
 	ctrClient.SetContainerExitHooks(manager.exitedAndRelease)
 
 	return manager, nil
diff --git a/containerm/updateagent/update_operation.go b/containerm/updateagent/update_operation.go
index 43fcc95..bb7c4c3 100644
--- a/containerm/updateagent/update_operation.go
+++ b/containerm/updateagent/update_operation.go
@@ -215,9 +215,10 @@ var baselineCommandHandlers = map[types.CommandType]struct {
 		commandHandler:         activate,
 	},
 	types.CommandRollback: {
-		expectedBaselineStatus: []types.StatusType{types.BaselineStatusActivationFailure, types.BaselineStatusActivationSuccess},
-		baselineFailureStatus:  types.BaselineStatusRollbackFailure,
-		commandHandler:         rollback,
+		expectedBaselineStatus: []types.StatusType{types.BaselineStatusDownloadSuccess, types.BaselineStatusUpdateSuccess,
+			types.BaselineStatusActivationFailure, types.BaselineStatusActivationSuccess},
+		baselineFailureStatus: types.BaselineStatusRollbackFailure,
+		commandHandler:        rollback,
 	},
 	types.CommandCleanup: {
 		baselineFailureStatus: types.BaselineStatusCleanup,