@@ -87,6 +87,19 @@ export type AIChatState = {
8787 error : boolean ;
8888} ;
8989
90+ export type AIChatEvent =
91+ | { type : 'open' }
92+ | { type : 'postMessage' ; message : string }
93+ | { type : 'clear' }
94+ | { type : 'close' } ;
95+
96+ type AIChatEventData < T extends AIChatEvent [ 'type' ] > = Omit <
97+ Extract < AIChatEvent , { type : T } > ,
98+ 'type'
99+ > ;
100+
101+ type AIChatEventListener = ( input ?: Omit < AIChatEvent , 'type' > ) => void ;
102+
90103export type AIChatController = {
91104 /** Open the dialog */
92105 open : ( ) => void ;
@@ -97,7 +110,10 @@ export type AIChatController = {
97110 /** Clear the conversation */
98111 clear : ( ) => void ;
99112 /** Register an event listener */
100- on : ( event : 'postMessage' , listener : ( input : { message : string } ) => void ) => ( ) => void ;
113+ on : < T extends AIChatEvent [ 'type' ] > (
114+ event : T ,
115+ listener : ( input ?: AIChatEventData < T > ) => void
116+ ) => ( ) => void ;
101117} ;
102118
103119const AIChatControllerContext = React . createContext < AIChatController | null > ( null ) ;
@@ -125,6 +141,17 @@ export function useAIChatState(): AIChatState {
125141 return state ;
126142}
127143
144+ function notify (
145+ listeners : AIChatEventListener [ ] | undefined ,
146+ input : Omit < AIChatEvent , 'type' >
147+ ) : void {
148+ if ( ! listeners ) return ;
149+ // Defer event listeners to next tick so React can process state updates first
150+ setTimeout ( ( ) => {
151+ listeners . forEach ( ( listener ) => listener ( input ) ) ;
152+ } , 0 ) ;
153+ }
154+
128155/**
129156 * Provide the controller to interact with the AI chat.
130157 */
@@ -140,9 +167,7 @@ export function AIChatProvider(props: {
140167 const language = useLanguage ( ) ;
141168
142169 // Event listeners storage
143- const eventsRef = React . useRef < Map < 'postMessage' , Array < ( input : { message : string } ) => void > > > (
144- new Map ( )
145- ) ;
170+ const eventsRef = React . useRef < Map < AIChatEvent [ 'type' ] , AIChatEventListener [ ] > > ( new Map ( ) ) ;
146171
147172 // Open AI chat and sync with search state
148173 const onOpen = React . useCallback ( ( ) => {
@@ -156,6 +181,8 @@ export function AIChatProvider(props: {
156181 scope : prev ?. scope ?? 'default' ,
157182 open : false , // Close search popover when opening chat
158183 } ) ) ;
184+
185+ notify ( eventsRef . current . get ( 'open' ) , { } ) ;
159186 } , [ setSearchState ] ) ;
160187
161188 // Close AI chat and clear ask parameter
@@ -169,6 +196,8 @@ export function AIChatProvider(props: {
169196 scope : prev ?. scope ?? 'default' ,
170197 open : false ,
171198 } ) ) ;
199+
200+ notify ( eventsRef . current . get ( 'close' ) , { } ) ;
172201 } , [ setSearchState ] ) ;
173202
174203 // Stream a message with the AI backend
@@ -386,11 +415,7 @@ export function AIChatProvider(props: {
386415 } ) ) ;
387416 }
388417
389- // Defer event listeners to next tick so React can process state updates first
390- setTimeout ( ( ) => {
391- const listeners = eventsRef . current . get ( 'postMessage' ) || [ ] ;
392- listeners . forEach ( ( listener ) => listener ( input ) ) ;
393- } , 0 ) ;
418+ notify ( eventsRef . current . get ( 'postMessage' ) , { message : input . message } ) ;
394419
395420 if ( query === input . message ) {
396421 // Return early if the message is the same as the previous message
@@ -458,9 +483,12 @@ export function AIChatProvider(props: {
458483 } , [ setSearchState ] ) ;
459484
460485 const onEvent = React . useCallback (
461- ( event : 'postMessage' , listener : ( input : { message : string } ) => void ) => {
486+ < T extends AIChatEvent [ 'type' ] > (
487+ event : T ,
488+ listener : ( input ?: AIChatEventData < T > ) => void
489+ ) => {
462490 const listeners = eventsRef . current . get ( event ) || [ ] ;
463- listeners . push ( listener ) ;
491+ listeners . push ( listener as AIChatEventListener ) ;
464492 eventsRef . current . set ( event , listeners ) ;
465493 return ( ) => {
466494 const currentListeners = eventsRef . current . get ( event ) || [ ] ;
0 commit comments