diff --git a/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenModels.swift b/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenModels.swift
index da00717a1..83a38763d 100644
--- a/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenModels.swift
+++ b/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenModels.swift
@@ -28,7 +28,7 @@ struct MediaEventsTimelineScreenViewState: BindableState {
     
     var groups = [MediaEventsTimelineGroup]()
     
-    var activeTimelineContextProvider: (() -> TimelineViewModel.Context)!
+    var activeTimelineContext: TimelineViewModel.Context
     
     var bindings: MediaEventsTimelineScreenViewStateBindings
 }
diff --git a/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenViewModel.swift b/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenViewModel.swift
index 24323bc2b..bca3a031f 100644
--- a/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenViewModel.swift
+++ b/ElementX/Sources/Screens/MediaEventsTimelineScreen/MediaEventsTimelineScreenViewModel.swift
@@ -37,7 +37,7 @@ class MediaEventsTimelineScreenViewModel: MediaEventsTimelineScreenViewModelType
     
     init(mediaTimelineViewModel: TimelineViewModelProtocol,
          filesTimelineViewModel: TimelineViewModelProtocol,
-         initialViewState: MediaEventsTimelineScreenViewState = .init(bindings: .init(screenMode: .media)),
+         initialScreenMode: MediaEventsTimelineScreenMode = .media,
          mediaProvider: MediaProviderProtocol,
          userIndicatorController: UserIndicatorControllerProtocol,
          appMediator: AppMediatorProtocol) {
@@ -47,14 +47,24 @@ class MediaEventsTimelineScreenViewModel: MediaEventsTimelineScreenViewModelType
         self.userIndicatorController = userIndicatorController
         self.appMediator = appMediator
         
-        super.init(initialViewState: initialViewState, mediaProvider: mediaProvider)
-        
-        state.activeTimelineContextProvider = { [weak self] in
-            guard let self else { fatalError() }
-            
-            return activeTimelineViewModel.context
+        let activeTimelineContext = switch initialScreenMode {
+        case .media: mediaTimelineViewModel.context
+        case .files: filesTimelineViewModel.context
         }
         
+        super.init(initialViewState: .init(activeTimelineContext: activeTimelineContext, bindings: .init(screenMode: initialScreenMode)), mediaProvider: mediaProvider)
+        
+        context.$viewState.map(\.bindings.screenMode)
+            .removeDuplicates()
+            .map {
+                switch $0 {
+                case .media: mediaTimelineViewModel.context
+                case .files: filesTimelineViewModel.context
+                }
+            }
+            .weakAssign(to: \.state.activeTimelineContext, on: self)
+            .store(in: &cancellables)
+        
         mediaTimelineViewModel.context.$viewState.sink { [weak self] timelineViewState in
             guard let self, state.bindings.screenMode == .media else {
                 return
diff --git a/ElementX/Sources/Screens/MediaEventsTimelineScreen/View/MediaEventsTimelineScreen.swift b/ElementX/Sources/Screens/MediaEventsTimelineScreen/View/MediaEventsTimelineScreen.swift
index 40d4e294e..1b6acc7ce 100644
--- a/ElementX/Sources/Screens/MediaEventsTimelineScreen/View/MediaEventsTimelineScreen.swift
+++ b/ElementX/Sources/Screens/MediaEventsTimelineScreen/View/MediaEventsTimelineScreen.swift
@@ -18,8 +18,8 @@ struct MediaEventsTimelineScreen: View {
             // Doesn't play well with the transformed scrollView
             .toolbarBackground(.visible, for: .navigationBar)
             .toolbar { toolbar }
-            .environmentObject(context.viewState.activeTimelineContextProvider())
-            .environment(\.timelineContext, context.viewState.activeTimelineContextProvider())
+            .environmentObject(context.viewState.activeTimelineContext)
+            .environment(\.timelineContext, context.viewState.activeTimelineContext)
             .onChange(of: context.screenMode) { _, _ in
                 context.send(viewAction: .changedScreenMode)
             }
@@ -137,7 +137,7 @@ struct MediaEventsTimelineScreen: View {
             AudioMediaEventsTimelineView(timelineItem: timelineItem)
         case .voice(let timelineItem):
             let defaultPlayerState = AudioPlayerState(id: .timelineItemIdentifier(timelineItem.id), title: L10n.commonVoiceMessage, duration: 0)
-            let playerState = context.viewState.activeTimelineContextProvider().viewState.audioPlayerStateProvider?(timelineItem.id) ?? defaultPlayerState
+            let playerState = context.viewState.activeTimelineContext.viewState.audioPlayerStateProvider?(timelineItem.id) ?? defaultPlayerState
             VoiceMessageMediaEventsTimelineView(timelineItem: timelineItem, playerState: playerState)
         default:
             EmptyView()
@@ -251,7 +251,7 @@ struct MediaEventsTimelineScreen_Previews: PreviewProvider, TestablePreview {
                                       screenMode: MediaEventsTimelineScreenMode) -> MediaEventsTimelineScreenViewModel {
         MediaEventsTimelineScreenViewModel(mediaTimelineViewModel: makeTimelineViewModel(empty: empty),
                                            filesTimelineViewModel: makeTimelineViewModel(empty: empty),
-                                           initialViewState: .init(bindings: .init(screenMode: screenMode)),
+                                           initialScreenMode: screenMode,
                                            mediaProvider: MediaProviderMock(configuration: .init()),
                                            userIndicatorController: UserIndicatorControllerMock(),
                                            appMediator: AppMediatorMock())