diff --git a/packages/sdk/src/realtime/client.ts b/packages/sdk/src/realtime/client.ts index 4baac21..017239c 100644 --- a/packages/sdk/src/realtime/client.ts +++ b/packages/sdk/src/realtime/client.ts @@ -12,7 +12,7 @@ import { type SubscribeEvents, type SubscribeOptions, } from "./subscribe-client"; -import type { ConnectionState, SessionIdMessage } from "./types"; +import type { ConnectionState, GenerationTickMessage, SessionIdMessage } from "./types"; import { WebRTCManager } from "./webrtc-manager"; async function blobToBase64(blob: Blob): Promise { @@ -99,6 +99,7 @@ export type RealTimeClientConnectOptions = z.infer { }; manager.getWebsocketMessageEmitter().on("sessionId", sessionIdListener); + const tickListener = (msg: GenerationTickMessage) => { + emitOrBuffer("generationTick", { seconds: msg.seconds }); + }; + manager.getWebsocketMessageEmitter().on("generationTick", tickListener); + await manager.connect(inputStream); const methods = realtimeMethods(manager, imageToBase64); diff --git a/packages/sdk/src/realtime/types.ts b/packages/sdk/src/realtime/types.ts index 89d17b6..acd6b39 100644 --- a/packages/sdk/src/realtime/types.ts +++ b/packages/sdk/src/realtime/types.ts @@ -64,6 +64,17 @@ export type GenerationStartedMessage = { type: "generation_started"; }; +export type GenerationTickMessage = { + type: "generation_tick"; + seconds: number; +}; + +export type GenerationEndedMessage = { + type: "generation_ended"; + seconds: number; + reason: string; +}; + export type SessionIdMessage = { type: "session_id"; session_id: string; @@ -84,6 +95,8 @@ export type IncomingWebRTCMessage = | ErrorMessage | SetImageAckMessage | GenerationStartedMessage + | GenerationTickMessage + | GenerationEndedMessage | SessionIdMessage; // Outgoing message types (to server) diff --git a/packages/sdk/src/realtime/webrtc-connection.ts b/packages/sdk/src/realtime/webrtc-connection.ts index 15a74cc..a412615 100644 --- a/packages/sdk/src/realtime/webrtc-connection.ts +++ b/packages/sdk/src/realtime/webrtc-connection.ts @@ -2,6 +2,7 @@ import mitt from "mitt"; import { buildUserAgent } from "../utils/user-agent"; import type { ConnectionState, + GenerationTickMessage, IncomingWebRTCMessage, OutgoingWebRTCMessage, PromptAckMessage, @@ -29,6 +30,7 @@ type WsMessageEvents = { promptAck: PromptAckMessage; setImageAck: SetImageAckMessage; sessionId: SessionIdMessage; + generationTick: GenerationTickMessage; }; export class WebRTCConnection { @@ -157,6 +159,17 @@ export class WebRTCConnection { return; } + if (msg.type === "generation_tick") { + this.websocketMessagesEmitter.emit("generationTick", msg); + return; + } + + if (msg.type === "generation_ended") { + // Handled internally — not surfaced as a public event. + // Devs use connectionChange for disconnect and error for insufficient credits. + return; + } + if (msg.type === "session_id") { this.websocketMessagesEmitter.emit("sessionId", msg); return;