-
Notifications
You must be signed in to change notification settings - Fork 186
/
useStreamCallback.tsx
78 lines (67 loc) · 2.29 KB
/
useStreamCallback.tsx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import {
MutableRefObject,
createContext,
useContext,
useEffect,
useRef,
} from "react";
import { StreamCallback } from "./types";
export const AppCallbackContext = createContext<MutableRefObject<{
onStart: Exclude<StreamCallback["onStart"], undefined>[];
onChunk: Exclude<StreamCallback["onChunk"], undefined>[];
onSuccess: Exclude<StreamCallback["onSuccess"], undefined>[];
onError: Exclude<StreamCallback["onError"], undefined>[];
}> | null>(null);
export function useAppStreamCallbacks() {
// callbacks handling
const context = useRef<{
onStart: Exclude<StreamCallback["onStart"], undefined>[];
onChunk: Exclude<StreamCallback["onChunk"], undefined>[];
onSuccess: Exclude<StreamCallback["onSuccess"], undefined>[];
onError: Exclude<StreamCallback["onError"], undefined>[];
}>({ onStart: [], onChunk: [], onSuccess: [], onError: [] });
const callbacks: StreamCallback = {
onStart(...args) {
for (const callback of context.current.onStart) {
callback(...args);
}
},
onChunk(...args) {
for (const callback of context.current.onChunk) {
callback(...args);
}
},
onSuccess(...args) {
for (const callback of context.current.onSuccess) {
callback(...args);
}
},
onError(...args) {
for (const callback of context.current.onError) {
callback(...args);
}
},
};
return { context, callbacks };
}
export function useStreamCallback<
Type extends "onStart" | "onChunk" | "onSuccess" | "onError"
>(type: Type, callback: Exclude<StreamCallback[Type], undefined>) {
type CallbackType = Exclude<StreamCallback[Type], undefined>;
const appCbRef = useContext(AppCallbackContext);
const callbackRef = useRef<CallbackType>(callback);
callbackRef.current = callback;
useEffect(() => {
// @ts-expect-error Not sure why I can't expand the tuple
const current = (...args) => callbackRef.current?.(...args);
appCbRef?.current[type].push(current);
return () => {
if (!appCbRef) return;
// @ts-expect-error Assingability issues due to the tuple object
// eslint-disable-next-line react-hooks/exhaustive-deps
appCbRef.current[type] = appCbRef.current[type].filter(
(callbacks) => callbacks !== current
);
};
}, [type, appCbRef]);
}