feat(companion): add typed agent stream subscription helpers

This commit is contained in:
William Valentin
2026-02-16 19:08:13 -08:00
parent fc0dd8b73a
commit 0a43abd531
5 changed files with 62 additions and 1 deletions
+1
View File
@@ -1,6 +1,7 @@
export {
CompanionRuntimeClient,
GatewayRpcError,
COMPANION_EVENT_NAMES,
} from './runtimeClient.js';
export {
MacOSCompanionClient,
+28
View File
@@ -196,6 +196,34 @@ describe('CompanionRuntimeClient', () => {
expect(streamHandler).toHaveBeenCalledTimes(1);
});
it('supports subscribeAgentStream and subscribeAgentTyping helpers', () => {
const client = new CompanionRuntimeClient({
url: 'ws://127.0.0.1:1',
});
const streamHandler = vi.fn();
const typingHandler = vi.fn();
client.subscribeAgentStream(streamHandler);
client.subscribeAgentTyping(typingHandler);
(client as unknown as { handleMessage: (raw: string) => void }).handleMessage(
JSON.stringify({
id: 52,
event: 'agent.stream',
data: { token: 'x' },
}),
);
(client as unknown as { handleMessage: (raw: string) => void }).handleMessage(
JSON.stringify({
id: 53,
event: 'agent.typing',
data: { active: true },
}),
);
expect(streamHandler).toHaveBeenCalledWith({ token: 'x' });
expect(typingHandler).toHaveBeenCalledWith({ active: true });
});
it('clears all event subscriptions', () => {
const client = new CompanionRuntimeClient({
url: 'ws://127.0.0.1:1',
+18
View File
@@ -45,6 +45,12 @@ export type CompanionEventHandler = (event: string, data: unknown) => void;
export type CompanionTypedEventHandler<TData = unknown> = (data: TData) => void;
export type CompanionEventPredicate<TData = unknown> = (data: TData) => boolean;
export const COMPANION_EVENT_NAMES = {
agentStream: 'agent.stream',
agentTyping: 'agent.typing',
contextWarning: 'context_warning',
} as const;
export interface RegisterNodeInput {
nodeId: string;
role: string;
@@ -385,6 +391,18 @@ export class CompanionRuntimeClient {
});
}
subscribeAgentStream<TData = unknown>(
handler: CompanionTypedEventHandler<TData>,
): () => void {
return this.subscribeEvent<TData>(COMPANION_EVENT_NAMES.agentStream, handler);
}
subscribeAgentTyping<TData = unknown>(
handler: CompanionTypedEventHandler<TData>,
): () => void {
return this.subscribeEvent<TData>(COMPANION_EVENT_NAMES.agentTyping, handler);
}
waitForEvent<TData = unknown>(
eventName: string,
options?: {