bidi streaming improvements

src/app/components/call-controls/call-controls.component.html

@@ -5,7 +5,7 @@
     class="video-rec-btn"
     [class.recording]="isVideoRecording"
     [matTooltip]="isVideoRecording ? i18n.turnOffCamTooltip : i18n.useCamTooltip"
-    [disabled]="!isBidiStreamingEnabled"
+    [disabled]="!isBidiStreamingEnabled || !useLive"
   >
     <mat-icon>videocam</mat-icon>
   </button>
@@ -22,12 +22,12 @@
     (click)="onCallClick()"
     class="audio-rec-btn"
     [class.recording]="isAudioRecording"
-    [disabled]="!isBidiStreamingEnabled"
+    [disabled]="!isBidiStreamingEnabled || !useLive"
   >
     <mat-icon>{{ isAudioRecording ? 'call_end' : 'call' }}</mat-icon>
   </button>
 
-  @if (showFlags && !isAudioRecording) {
+  @if (useLive && showFlags && !isAudioRecording) {
     <div class="flags-panel">
       <div class="flags-title">Live Flags</div>
       <div class="flag-item">

src/app/components/call-controls/call-controls.component.scss

@@ -19,6 +19,10 @@ button {
   }
 }
 
+button.audio-rec-btn:not(.recording):disabled {
+  color: gray !important;
+}
+
 button.audio-rec-btn:not(.recording) {
   color: #34a853 !important;
 }

src/app/components/call-controls/call-controls.component.ts

@@ -40,6 +40,7 @@ export class CallControlsComponent {
   @Input() isVideoRecording = false;
   @Input() micVolume = 0;
   @Input() isBidiStreamingEnabled: boolean | null = false;
+  @Input() useLive: boolean | null = false;
 
   @Output() readonly toggleAudioRecording = new EventEmitter<LiveFlags>();
   @Output() readonly toggleVideoRecording = new EventEmitter<void>();

src/app/components/chat-panel/chat-panel.component.html

@@ -207,6 +207,7 @@ <h3 style="margin: 0; color: var(--mat-sys-primary);">Evaluation Result</h3>
     <div class="chat-input-actions-right">
       <app-call-controls [isAudioRecording]="isAudioRecording" [isVideoRecording]="isVideoRecording"
         [micVolume]="micVolume" [isBidiStreamingEnabled]="(isBidiStreamingEnabledObs | async) ?? false"
+        [useLive]="useLive"
         (toggleAudioRecording)="toggleAudioRecording.emit($event)"
         (toggleVideoRecording)="toggleVideoRecording.emit()"></app-call-controls>
     </div>

src/app/components/chat-panel/chat-panel.component.ts

@@ -134,7 +134,8 @@ export class ChatPanelComponent implements OnChanges, AfterViewInit {
   @Input() agentGraphData: any = null;
   @Input() isEditFunctionArgsEnabled: boolean = false;
   @Input() isTokenStreamingEnabled: boolean = false;
-  @Input() useSse: boolean = false;
+  @Input() useStreaming: boolean = false;
+  @Input() useLive: boolean = false;
   @Input() userInput: string = '';
   @Input() userEditEvalCaseMessage: string = '';
   @Input() selectedFiles: { file: File; url: string }[] = [];
@@ -178,7 +179,7 @@ export class ChatPanelComponent implements OnChanges, AfterViewInit {
   @Output() readonly toggleVideoRecording = new EventEmitter<void>();
   @Output() readonly longRunningResponseComplete = new EventEmitter<AgentRunRequest>();
   @Output() readonly toggleHideIntermediateEvents = new EventEmitter<void>();
-  @Output() readonly toggleSse = new EventEmitter<void>();
+  @Output() readonly toggleStreaming = new EventEmitter<void>();
 
   @ViewChild('videoContainer', { read: ElementRef }) videoContainer!: ElementRef;
   @ViewChild('autoScroll') scrollContainer!: ElementRef;

src/app/components/chat/chat.component.html

@@ -556,15 +556,26 @@
   }
 
   @if ((isTokenStreamingEnabledObs | async) && canEditSession()) {
-  <button mat-button (click)="toggleSse()" matTooltip="Enable real-time token streaming from the server"
-    [style.color]="useSse() ? 'var(--mat-sys-primary)' : 'var(--mat-sys-on-surface-variant)'"
+  <button mat-button (click)="toggleStreaming()" matTooltip="Enable real-time token streaming from the server"
+    [style.color]="useStreaming() ? 'var(--mat-sys-primary)' : 'var(--mat-sys-on-surface-variant)'"
     style="height: 32px; line-height: 32px; padding: 0 12px; border-radius: 16px;">
     <mat-icon
       style="font-size: 20px; width: 20px; height: 20px; line-height: 20px; margin-right: 4px; vertical-align: middle;">{{
-      useSse() ? 'check_circle' : 'radio_button_unchecked' }}</mat-icon>
+      useStreaming() ? 'check_circle' : 'radio_button_unchecked' }}</mat-icon>
     <span style="font-size: 13px; font-weight: 500; vertical-align: middle;">Streaming</span>
   </button>
   }
+
+  @if ((isBidiStreamingEnabledObs | async) && canEditSession()) {
+  <button mat-button (click)="toggleLive()" matTooltip="Enable live connection to the server"
+    [style.color]="useLive() ? 'var(--mat-sys-primary)' : 'var(--mat-sys-on-surface-variant)'"
+    style="height: 32px; line-height: 32px; padding: 0 12px; border-radius: 16px;">
+    <mat-icon
+      style="font-size: 20px; width: 20px; height: 20px; line-height: 20px; margin-right: 4px; vertical-align: middle;">{{
+      useLive() ? 'check_circle' : 'radio_button_unchecked' }}</mat-icon>
+    <span style="font-size: 13px; font-weight: 500; vertical-align: middle;">Live</span>
+  </button>
+  }
 </div>
 
         @switch (chatType()) {
@@ -579,8 +590,8 @@
               (toggleHideIntermediateEvents)="toggleHideIntermediateEvents()"
               [traceData]="traceData"
               [isTokenStreamingEnabled]="(isTokenStreamingEnabledObs | async) ?? false"
-              [useSse]="useSse()"
-              (toggleSse)="toggleSse()"
+              [useStreaming]="useStreaming()"
+              (toggleStreaming)="toggleStreaming()"
               [isChatMode]="true"
               [selectedFiles]="selectedFiles"
               [updatedSessionState]="updatedSessionState()"
@@ -609,6 +620,7 @@
               [invocationDisplayMap]="invocationDisplayMap()"
               [viewMode]="viewMode()"
               [shouldShowEvent]="shouldShowEventFn"
+              [useLive]="useLive()"
             ></app-chat-panel>
           }
           @case ('eval-case') {
@@ -642,6 +654,7 @@
               [invocationDisplayMap]="invocationDisplayMap()"
               [viewMode]="viewMode()"
               [shouldShowEvent]="shouldShowEventFn"
+              [useLive]="useLive()"
             ></app-chat-panel>
           }
           @case ('eval-result') {
@@ -723,6 +736,7 @@
                     [invocationDisplayMap]="invocationDisplayMap()"
                     [viewMode]="viewMode()"
                     [shouldShowEvent]="shouldShowEventFn"
+                    [useLive]="useLive()"
                   ></app-chat-panel>
                 </div>
                 <div class="side-panel-half">
@@ -736,8 +750,8 @@
                     (toggleHideIntermediateEvents)="toggleHideIntermediateEvents()"
                     [traceData]="traceData"
                     [isTokenStreamingEnabled]="(isTokenStreamingEnabledObs | async) ?? false"
-                    [useSse]="useSse()"
-                    (toggleSse)="toggleSse()"
+                    [useStreaming]="useStreaming()"
+                    (toggleStreaming)="toggleStreaming()"
                     [isChatMode]="false"
                     [evalCase]="evalCase"
                     [isEvalEditMode]="isEvalEditMode()"
@@ -775,6 +789,7 @@
                     [invocationDisplayMap]="invocationDisplayMap()"
                     [viewMode]="viewMode()"
                     [shouldShowEvent]="shouldShowEventFn"
+                    [useLive]="useLive()"
                   ></app-chat-panel>
                 </div>
               </div>
@@ -796,6 +811,7 @@
                 [invocationDisplayMap]="invocationDisplayMap()"
                 [viewMode]="viewMode()"
                 [shouldShowEvent]="shouldShowEventFn"
+                [useLive]="useLive()"
               ></app-chat-panel>
             }
           }

src/app/components/chat/chat.component.ts

@@ -289,7 +289,8 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
   showBuilderAssistant = true;
   showAppSelectorDrawer = false;
   showSessionSelectorDrawer = false;
-  useSse = signal(window.localStorage.getItem('adk-use-sse') === 'true');
+  useStreaming = signal(window.localStorage.getItem('adk-use-streaming') === 'true');
+  useLive = signal(window.localStorage.getItem('adk-use-live') === 'true');
   currentSessionState: SessionState | undefined = {};
   root_agent = ROOT_AGENT;
   updatedSessionState: WritableSignal<any> = signal(null);
@@ -528,9 +529,6 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     window.localStorage.setItem('adk-hide-intermediate-events', String(newVal));
   }
 
-  // TODO: Remove this once backend supports restarting bidi streaming.
-  sessionHasUsedBidi = new Set<string>();
-
   eventData = new Map<string, any>();
   traceData: any[] = [];
   renderedEventGraph: SafeHtml | undefined;
@@ -630,6 +628,8 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     this.featureFlagService.isApplicationSelectorEnabled();
   readonly isTokenStreamingEnabledObs: Observable<boolean> =
     this.featureFlagService.isTokenStreamingEnabled();
+  readonly isBidiStreamingEnabledObs: Observable<boolean> =
+    this.featureFlagService.isBidiStreamingEnabled();
   readonly isExportSessionEnabledObs: Observable<boolean> =
     this.featureFlagService.isExportSessionEnabled();
   readonly isNewSessionButtonEnabledObs: Observable<boolean> =
@@ -1097,7 +1097,7 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
       userId: this.userId,
       sessionId: this.sessionId,
       newMessage: content,
-      streaming: this.useSse(),
+      streaming: this.useStreaming(),
       stateDelta: this.updatedSessionState(),
     };
     if (functionCallEventId) {
@@ -1110,6 +1110,10 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
 
   submitAgentRunRequest(req: AgentRunRequest) {
     this.autoSelectLatestEvent = true;
+    if (this.useLive()) {
+      this.streamChatService.sendMessage(req, {});
+      return;
+    }
     this.agentService.runSse(req).subscribe({
       next: async (chunkJson: any) => {
         if (chunkJson.error) {
@@ -1150,7 +1154,7 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     });
   }
 
-  private appendEventRow(apiEvent: any, reverseOrder: boolean = false) {
+  private appendEventRow(apiEvent: AdkEvent, reverseOrder: boolean = false) {
     if (apiEvent.inputTranscription !== undefined) {
       apiEvent.author = 'user';
     } else if (apiEvent.outputTranscription !== undefined) {
@@ -1173,7 +1177,7 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     if (apiEvent?.longRunningToolIds && apiEvent.longRunningToolIds.length > 0) {
       const startIndex = this.longRunningEvents.length;
       this.getAsyncFunctionsFromParts(
-        apiEvent.longRunningToolIds, apiEvent.content.parts, apiEvent.invocationId);
+        apiEvent.longRunningToolIds, apiEvent.content.parts, apiEvent.invocationId!!);
 
       // Store event ID for later reference
       this.functionCallEventId = apiEvent.id;
@@ -1815,11 +1819,6 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
   }
 
   async startAudioRecording(flags?: LiveFlags) {
-    if (this.sessionId && this.sessionHasUsedBidi.has(this.sessionId)) {
-      this.openSnackBar(BIDI_STREAMING_RESTART_WARNING, 'OK');
-      return;
-    }
-
     // Lazily create a real session if it does not exist
     const isSessionActive = await this.ensureSessionActive();
     if (!isSessionActive) {
@@ -1833,7 +1832,6 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
       sessionId: this.sessionId,
       flags: flags,
     });
-    this.sessionHasUsedBidi.add(this.sessionId);
   }
 
   stopAudioRecording() {
@@ -2723,9 +2721,14 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     this.selectedFiles.splice(index, 1);
   }
 
-  toggleSse() {
-    this.useSse.set(!this.useSse());
-    window.localStorage.setItem('adk-use-sse', String(this.useSse()));
+  toggleStreaming() {
+    this.useStreaming.set(!this.useStreaming());
+    window.localStorage.setItem('adk-use-streaming', String(this.useStreaming()));
+  }
+
+  toggleLive() {
+    this.useLive.set(!this.useLive());
+    window.localStorage.setItem('adk-use-live', String(this.useLive()));
   }
 
   enterBuilderMode() {

src/app/components/long-running-response/long-running-response.ts

@@ -39,6 +39,7 @@ import {MatIcon} from '@angular/material/icon';
     MatButton,
     MatIcon,
     NgxJsonViewerModule,
+    MarkdownComponent,
   ],
 })
 export class LongRunningResponseComponent implements OnChanges {

src/app/core/models/types.ts

@@ -128,11 +128,12 @@ export declare interface Event extends LlmResponse {
   nodeInfo?: { path?: string;[key: string]: any; };
   data?: any;
   output?: { result?: any; };
-  inputTranscription?: { text: string; };
-  outputTranscription?: { text: string; };
+  inputTranscription?: { text: string; finished?: boolean; };
+  outputTranscription?: { text: string; finished?: boolean; };
   usageMetadata?: any;
   interrupted?: boolean;
   turnComplete?: boolean;
+  partial?: boolean;
 }
 
 export interface ComputerUsePayload {

src/app/core/services/interfaces/stream-chat.ts

@@ -17,6 +17,8 @@
 
 import {ElementRef, InjectionToken} from '@angular/core';
 import {Observable} from 'rxjs';
+import {Event} from '../../models/types';
+import {AgentRunRequest} from '../../models/AgentRunRequest';
 
 export const STREAM_CHAT_SERVICE =
   new InjectionToken<StreamChatService>('StreamChatService');
@@ -51,4 +53,5 @@ export declare abstract class StreamChatService {
   abstract stopVideoStreaming(videoContainer: ElementRef): void;
   abstract onStreamClose(): Observable<string>;
   abstract closeStream(): void;
+  abstract sendMessage(req: AgentRunRequest, flags?: LiveFlags): void;
 }

src/app/core/services/stream-chat.service.ts

@@ -26,6 +26,8 @@ import {VIDEO_SERVICE} from './interfaces/video';
 import {WEBSOCKET_SERVICE} from './interfaces/websocket';
 import {VideoService} from './video.service';
 import {WebSocketService} from './websocket.service';
+import {map,filter} from 'rxjs/operators';
+import {AgentRunRequest} from '../models/AgentRunRequest';
 
 /**
  * Service for supporting live streaming with audio/video.
@@ -39,9 +41,25 @@ export class StreamChatService implements StreamChatServiceInterface {
   private readonly webSocketService = inject(WEBSOCKET_SERVICE);
   private audioIntervalId: number|undefined = undefined;
   private videoIntervalId: number|undefined = undefined;
+  private currentUrl: string|undefined = undefined;
 
   constructor() {}
 
+  connect(serverUrl: string) {
+    if (this.currentUrl !== serverUrl) {
+      this.webSocketService.connect(serverUrl);
+      this.currentUrl = serverUrl;
+    }
+  }
+
+  sendMessage({ appName, userId, sessionId, newMessage }: AgentRunRequest, flags?: LiveFlags): void {
+    this.connect(this.getWsUrl(appName, userId, sessionId!!, flags));
+    const request: LiveRequest = {
+      content: newMessage,
+    };
+    this.webSocketService.sendMessage(request);
+  }
+
   private getWsUrl(appName: string, userId: string, sessionId: string, flags?: LiveFlags): string {
     const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws';
     let url = `${protocol}://${URLUtil.getWSServerUrl()}/run_live?app_name=${appName}&user_id=${userId}&session_id=${sessionId}`;
@@ -68,14 +86,13 @@ export class StreamChatService implements StreamChatServiceInterface {
     sessionId,
     flags,
   }: {appName: string; userId: string; sessionId: string; flags?: LiveFlags;}) {
-    this.webSocketService.connect(this.getWsUrl(appName, userId, sessionId, flags));
+    this.connect(this.getWsUrl(appName, userId, sessionId, flags));
 
     await this.startAudioStreaming();
   }
 
   stopAudioChat() {
     this.stopAudioStreaming();
-    this.webSocketService.closeConnection();
   }
 
   private async startAudioStreaming() {
@@ -118,7 +135,7 @@ export class StreamChatService implements StreamChatServiceInterface {
     videoContainer: ElementRef;
     flags?: LiveFlags;
   }) {
-    this.webSocketService.connect(this.getWsUrl(appName, userId, sessionId, flags));
+    this.connect(this.getWsUrl(appName, userId, sessionId, flags));
 
     await this.startAudioStreaming();
     await this.startVideoStreaming(videoContainer);
@@ -127,7 +144,6 @@ export class StreamChatService implements StreamChatServiceInterface {
   stopVideoChat(videoContainer: ElementRef) {
     this.stopAudioStreaming();
     this.stopVideoStreaming(videoContainer);
-    this.webSocketService.closeConnection();
   }
 
   async startVideoStreaming(videoContainer: ElementRef) {

src/app/core/services/websocket.service.ts

@@ -66,7 +66,9 @@ export class WebSocketService implements WebSocketServiceInterface {
   }
 
   sendMessage(data: LiveRequest) {
-    data.blob.data = this.arrayBufferToBase64(data.blob.data.buffer);
+    if (data.blob?.data) {
+      data.blob.data = this.arrayBufferToBase64(data.blob.data.buffer);
+    }
     if (!this.socket$ || this.socket$.closed) {
       console.error('WebSocket is not open.');
       return;
@@ -101,7 +103,8 @@ export class WebSocketService implements WebSocketServiceInterface {
     if (
       msg['content'] &&
       msg['content']['parts'] &&
-      msg['content']['parts'][0]['inlineData']
+      msg['content']['parts'][0]['inlineData'] &&
+      msg['content']['parts'][0]['inlineData']['mimeType']?.startsWith('audio/pcm')
     ) {
       const pcmBytes = this.base64ToUint8Array(
           msg['content']['parts'][0]['inlineData']['data'],