fix: audio track undo buffer rehydration

src/core/audio-interface/KGAudioInterface.ts

@@ -339,6 +339,20 @@ export class KGAudioInterface {
     }
   }
 
+  /**
+   * Check whether an audio track currently has a player bus.
+   */
+  public hasTrackAudioPlayerBus(trackId: string): boolean {
+    return this.trackAudioPlayerBuses.has(trackId);
+  }
+
+  /**
+   * Check whether an audio buffer is loaded on a specific track's player bus.
+   */
+  public hasAudioBufferForTrack(trackId: string, audioFileId: string): boolean {
+    return this.trackAudioPlayerBuses.get(trackId)?.hasBuffer(audioFileId) ?? false;
+  }
+
   /**
    * Load an audio buffer into a track's player bus
    */

src/stores/projectStore.test.ts

@@ -2,6 +2,8 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { act } from '@testing-library/react';
 import { KGTrack } from '../core/track/KGTrack';
 import { KGMidiTrack } from '../core/track/KGMidiTrack';
+import { KGAudioTrack } from '../core/track/KGAudioTrack';
+import { KGAudioRegion } from '../core/region/KGAudioRegion';
 import { createDefaultGlobalTracks } from '../core/global-track';
 
 const pianoRollStateMocks = vi.hoisted(() => ({
@@ -9,6 +11,25 @@ const pianoRollStateMocks = vi.hoisted(() => ({
   setPianoRollZoom: vi.fn(),
 }));
 
+const audioStorageMocks = vi.hoisted(() => ({
+  loadAudioFile: vi.fn(),
+}));
+
+const toneMocks = vi.hoisted(() => {
+  const decodeAudioData = vi.fn();
+  const toneBufferSet = vi.fn();
+
+  class MockToneAudioBuffer {
+    public set = toneBufferSet;
+  }
+
+  return {
+    decodeAudioData,
+    toneBufferSet,
+    ToneAudioBuffer: MockToneAudioBuffer,
+  };
+});
+
 let mockTracks: KGTrack[] = [new KGMidiTrack('Track 1', 0, 'acoustic_grand_piano')];
 let mockIsMetronomeEnabled = false;
 let mockShowGlobalTracks = false;
@@ -44,6 +65,9 @@ const mockAudioInterface = {
   removeTrackSynth: vi.fn(),
   removeTrackAudioPlayerBus: vi.fn(),
   createTrackAudioPlayerBus: vi.fn().mockResolvedValue(undefined),
+  hasTrackAudioPlayerBus: vi.fn().mockReturnValue(true),
+  hasAudioBufferForTrack: vi.fn().mockReturnValue(false),
+  getAudioBuffer: vi.fn(),
   loadAudioBufferForTrack: vi.fn(),
   createTrackSynth: vi.fn(),
   setTrackVolume: vi.fn(),
@@ -96,6 +120,21 @@ vi.mock('../core/audio-interface/KGAudioInterface', () => ({
   },
 }));
 
+vi.mock('../core/io/KGAudioFileStorage', () => ({
+  KGAudioFileStorage: {
+    loadAudioFile: audioStorageMocks.loadAudioFile,
+  },
+}));
+
+vi.mock('tone', () => ({
+  getContext: () => ({
+    rawContext: {
+      decodeAudioData: toneMocks.decodeAudioData,
+    },
+  }),
+  ToneAudioBuffer: toneMocks.ToneAudioBuffer,
+}));
+
 vi.mock('../core/config/ConfigManager', () => ({
   ConfigManager: {
     instance: () => ({
@@ -139,7 +178,18 @@ describe('projectStore piano roll state', () => {
     mockAudioInterface.cancelAudioRecording.mockResolvedValue(undefined);
     mockAudioInterface.getTransportPosition.mockReset();
     mockAudioInterface.getTransportPosition.mockReturnValue(8);
+    mockAudioInterface.hasTrackAudioPlayerBus.mockReset();
+    mockAudioInterface.hasTrackAudioPlayerBus.mockReturnValue(true);
+    mockAudioInterface.hasAudioBufferForTrack.mockReset();
+    mockAudioInterface.hasAudioBufferForTrack.mockReturnValue(false);
+    mockAudioInterface.getAudioBuffer.mockReset();
+    mockAudioInterface.createTrackAudioPlayerBus.mockReset();
+    mockAudioInterface.createTrackAudioPlayerBus.mockResolvedValue(undefined);
+    mockAudioInterface.loadAudioBufferForTrack.mockReset();
     mockAudioInterface.setMetronomeEnabled.mockReset();
+    audioStorageMocks.loadAudioFile.mockReset();
+    toneMocks.decodeAudioData.mockReset();
+    toneMocks.toneBufferSet.mockReset();
     mockIsMetronomeEnabled = false;
     mockShowGlobalTracks = false;
     mockProject.setIsMetronomeEnabled.mockClear();
@@ -376,6 +426,96 @@ describe('projectStore piano roll state', () => {
     expect(useProjectStore.getState().trackAutomationRedrawVersion).toBe(initialVersion + 2);
   });
 
+  it('rehydrates missing audio buffers during refreshProjectState', async () => {
+    const audioTrack = new KGAudioTrack('Audio 1', 1);
+    audioTrack.setTrackIndex(0);
+    audioTrack.setRegions([
+      new KGAudioRegion('audio-region-1', '1', 0, 'clip.wav', 0, 4, 'audio-file-1.wav', 'clip.wav', 2.5),
+    ]);
+    mockTracks = [audioTrack];
+
+    const decodedBuffer = { duration: 2.5 } as AudioBuffer;
+    audioStorageMocks.loadAudioFile.mockResolvedValue(new ArrayBuffer(16));
+    toneMocks.decodeAudioData.mockResolvedValue(decodedBuffer);
+
+    const { useProjectStore } = await import('./projectStore');
+
+    await act(async () => {
+      useProjectStore.getState().refreshProjectState();
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    expect(audioStorageMocks.loadAudioFile).toHaveBeenCalledWith('Test Project', 'audio-file-1.wav');
+    expect(toneMocks.decodeAudioData).toHaveBeenCalled();
+    expect(mockAudioInterface.loadAudioBufferForTrack).toHaveBeenCalledWith(
+      '1',
+      'audio-file-1.wav',
+      expect.any(toneMocks.ToneAudioBuffer),
+    );
+  });
+
+  it('reloads a restored audio track buffer on undo', async () => {
+    const restoredTrack = new KGAudioTrack('Audio 1', 1);
+    restoredTrack.setTrackIndex(0);
+    const restoredRegion = new KGAudioRegion(
+      'audio-region-1',
+      '1',
+      0,
+      'clip.wav',
+      0,
+      4,
+      'audio-file-1.wav',
+      'clip.wav',
+      2.5,
+    );
+    restoredTrack.setRegions([restoredRegion]);
+
+    mockTracks = [];
+    mockCore.undo.mockImplementationOnce(() => {
+      mockTracks = [restoredTrack];
+      return true;
+    });
+
+    const decodedBuffer = { duration: 2.5 } as AudioBuffer;
+    audioStorageMocks.loadAudioFile.mockResolvedValue(new ArrayBuffer(16));
+    toneMocks.decodeAudioData.mockResolvedValue(decodedBuffer);
+
+    const { useProjectStore } = await import('./projectStore');
+    const initialWaveformVersion = useProjectStore.getState().audioWaveformRedrawVersion;
+
+    await act(async () => {
+      useProjectStore.getState().undo();
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    const state = useProjectStore.getState();
+    expect(state.tracks).toHaveLength(1);
+    const restoredTrackState = state.tracks[0] as KGAudioTrack;
+    expect((restoredTrackState.getRegions()[0] as KGAudioRegion).getAudioFileId()).toBe('audio-file-1.wav');
+    expect(mockAudioInterface.loadAudioBufferForTrack).toHaveBeenCalledWith(
+      '1',
+      'audio-file-1.wav',
+      expect.any(toneMocks.ToneAudioBuffer),
+    );
+    expect(state.audioWaveformRedrawVersion).toBeGreaterThan(initialWaveformVersion);
+  });
+
+  it('does not attempt audio buffer hydration for MIDI-only undo', async () => {
+    mockTracks = [new KGMidiTrack('Track 1', 1, 'acoustic_grand_piano')];
+
+    const { useProjectStore } = await import('./projectStore');
+
+    await act(async () => {
+      useProjectStore.getState().undo();
+      await Promise.resolve();
+    });
+
+    expect(audioStorageMocks.loadAudioFile).not.toHaveBeenCalled();
+    expect(mockAudioInterface.loadAudioBufferForTrack).not.toHaveBeenCalled();
+  });
+
   it('restores Chat after closing Settings when Chat was active on entry', async () => {
     const { useProjectStore } = await import('./projectStore');
 

src/stores/projectStore.ts

@@ -317,6 +317,66 @@ function getAudioRecordingExtension(mimeType: string): string {
   return 'webm';
 }
 
+const pendingAudioBufferHydrations = new Set<string>();
+
+async function decodeStoredAudioFile(arrayBuffer: ArrayBuffer): Promise<Tone.ToneAudioBuffer> {
+  const audioContext = Tone.getContext().rawContext as AudioContext;
+  const decoded = await audioContext.decodeAudioData(arrayBuffer);
+  const toneBuffer = new Tone.ToneAudioBuffer();
+  toneBuffer.set(decoded);
+  return toneBuffer;
+}
+
+async function hydrateAudioTrackBuffers(project: KGProject): Promise<boolean> {
+  const audioInterface = KGAudioInterface.instance();
+  const projectName = project.getName();
+  let hydratedAnyBuffer = false;
+
+  for (const track of project.getTracks()) {
+    if (track.getCurrentType() !== 'KGAudioTrack') {
+      continue;
+    }
+
+    const audioTrack = track as KGAudioTrack;
+    const trackId = audioTrack.getId().toString();
+
+    if (!audioInterface.hasTrackAudioPlayerBus(trackId)) {
+      await audioInterface.createTrackAudioPlayerBus(trackId, audioTrack.getVolume());
+    }
+
+    for (const region of audioTrack.getRegions()) {
+      if (region.getCurrentType() !== 'KGAudioRegion') {
+        continue;
+      }
+
+      const audioRegion = region as KGAudioRegion;
+      const audioFileId = audioRegion.getAudioFileId();
+      if (!audioFileId || audioInterface.hasAudioBufferForTrack(trackId, audioFileId)) {
+        continue;
+      }
+
+      const hydrationKey = `${projectName}:${trackId}:${audioFileId}`;
+      if (pendingAudioBufferHydrations.has(hydrationKey)) {
+        continue;
+      }
+
+      pendingAudioBufferHydrations.add(hydrationKey);
+      try {
+        const arrayBuffer = await KGAudioFileStorage.loadAudioFile(projectName, audioFileId);
+        const toneBuffer = await decodeStoredAudioFile(arrayBuffer);
+        audioInterface.loadAudioBufferForTrack(trackId, audioFileId, toneBuffer);
+        hydratedAnyBuffer = true;
+      } catch (err) {
+        console.error(`Failed to load audio file ${audioFileId}:`, err);
+      } finally {
+        pendingAudioBufferHydrations.delete(hydrationKey);
+      }
+    }
+  }
+
+  return hydratedAnyBuffer;
+}
+
 // Create the store
 export const useProjectStore = create<ProjectState>((set, get) => {
   const currentProject = KGCore.instance().getCurrentProject();
@@ -887,34 +947,12 @@ export const useProjectStore = create<ProjectState>((set, get) => {
         });
 
         // Create synths/buses for all tracks (with their stored volumes)
-        const projectName = projectToLoad.getName();
         for (const track of tracks) {
           const trackId = track.getId().toString();
 
           if (track.getCurrentType() === 'KGAudioTrack') {
-            // Audio track: create player bus and load audio buffers
+            // Audio track: create player bus; buffers are hydrated in a shared pass below
             await audioInterface.createTrackAudioPlayerBus(trackId, track.getVolume());
-
-            // Load audio buffers for all regions in this audio track
-            const audioTrack = track as KGAudioTrack;
-            for (const region of audioTrack.getRegions()) {
-              if (region.getCurrentType() === 'KGAudioRegion') {
-                const audioRegion = region as KGAudioRegion;
-                const audioFileId = audioRegion.getAudioFileId();
-                if (audioFileId) {
-                  try {
-                    const arrayBuffer = await KGAudioFileStorage.loadAudioFile(projectName, audioFileId);
-                    const audioContext = Tone.getContext().rawContext as AudioContext;
-                    const decoded = await audioContext.decodeAudioData(arrayBuffer);
-                    const toneBuffer = new Tone.ToneAudioBuffer();
-                    toneBuffer.set(decoded);
-                    audioInterface.loadAudioBufferForTrack(trackId, audioFileId, toneBuffer);
-                  } catch (err) {
-                    console.error(`Failed to load audio file ${audioFileId}:`, err);
-                  }
-                }
-              }
-            }
           } else {
             // MIDI track: create sampler-based audio bus
             let instrument: InstrumentType = 'acoustic_grand_piano';
@@ -926,6 +964,8 @@ export const useProjectStore = create<ProjectState>((set, get) => {
           }
         }
 
+        await hydrateAudioTrackBuffers(projectToLoad);
+
         // Reapply restored mute/solo state after all buses exist so solo logic can be
         // computed against the full track set.
         for (const track of tracks) {
@@ -1926,6 +1966,17 @@ export const useProjectStore = create<ProjectState>((set, get) => {
       const actions = get();
       actions.syncUndoRedoState();
       actions.syncSelectionFromCore();
+
+      void hydrateAudioTrackBuffers(project).then((hydratedAnyBuffer) => {
+        if (!hydratedAnyBuffer) {
+          return;
+        }
+
+        set(state => ({
+          tracks: [...project.getTracks()] as KGTrack[],
+          audioWaveformRedrawVersion: state.audioWaveformRedrawVersion + 1,
+        }));
+      });
     },
 
     // Initialize store with configuration values