examples/web_ui/frontend/src/context/AudioContext.tsx

import { createContext, useContext, useEffect, useMemo, useRef, useSyncExternalStore } from 'react';
import type { ReactNode } from 'react';

import { StreamingAudioManager } from '@/utils/streamingAudio';
import type { StreamingAudioState } from '@/utils/streamingAudio';

const AudioContext = createContext<StreamingAudioManager | null>(null);

interface ReplayController {
	play: (el: HTMLAudioElement) => void;
	stop: () => void;
}

const ReplayContext = createContext<ReplayController | null>(null);

/**

 * Provides a {@link StreamingAudioManager} to the component tree. The manager

 * collects DATA_BLOCK_* events for audio blocks and exposes per-block state

 * that ``MessageBubble`` consumes via {@link useAudioBlock}.

 *

 * Mount this once around any subtree that renders assistant messages.

 */
export function AudioProvider({ children }: { children: ReactNode }) {
	const manager = useMemo(() => new StreamingAudioManager(), []);
	useEffect(() => () => manager.disposeAll(), [manager]);

	const currentRef = useRef<HTMLAudioElement | null>(null);
	const replay: ReplayController = useMemo(
		() => ({
			play(el: HTMLAudioElement) {
				manager.stopLivePlayback();
				if (currentRef.current && currentRef.current !== el) {
					currentRef.current.pause();
					currentRef.current.currentTime = 0;
				}
				currentRef.current = el;
			},
			stop() {
				currentRef.current = null;
			},
		}),
		[manager],
	);

	return (
		<AudioContext.Provider value={manager}>

			<ReplayContext.Provider value={replay}>{children}</ReplayContext.Provider>

		</AudioContext.Provider>
	);
}

/**

 * Access the streaming audio manager. Returns ``null`` outside of an

 * ``AudioProvider`` — callers must handle that, since audio handling is

 * an optional feature.

 */
export function useAudioManager(): StreamingAudioManager | null {
	return useContext(AudioContext);
}

/**

 * Access the replay controller that ensures only one audio element plays

 * at a time across all message bubbles.

 */
export function useReplayController(): ReplayController | null {
	return useContext(ReplayContext);
}

/**

 * Subscribe to the streaming state for a single audio DataBlock.

 * Re-renders when the block transitions from ``streaming`` to ``ready``.

 *

 * Returns ``null`` if the block isn't being tracked (e.g. a historical

 * message loaded from the server, where the bytes are already complete).

 */
export function useAudioBlock(blockId: string | undefined): StreamingAudioState | null {
	const manager = useAudioManager();
	return useSyncExternalStore(
		(fn) => {
			if (!manager || !blockId) return () => undefined;
			return manager.subscribe(blockId, fn);
		},
		() => (manager && blockId ? manager.getState(blockId) : null),
		() => null,
	);
}