feat(server): refactor copilot (#14892)

#### PR Dependency Tree


* **PR #14892** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

packages/backend/server/src/plugins/copilot/controller.ts

@@ -13,22 +13,17 @@ import type { Request, Response } from 'express';
 import {
   BehaviorSubject,
   catchError,
-  connect,
   filter,
   finalize,
   from,
-  ignoreElements,
   interval,
   lastValueFrom,
   map,
   merge,
-  mergeMap,
   Observable,
-  reduce,
   Subject,
   take,
   takeUntil,
-  tap,
 } from 'rxjs';
 
 import {
@@ -36,28 +31,18 @@ import {
   BlobNotFound,
   CallMetric,
   Config,
-  CopilotSessionNotFound,
   mapSseError,
   metrics,
-  NoCopilotProviderAvailable,
   UnsplashIsNotConfigured,
 } from '../../base';
-import { ServerFeature, ServerService } from '../../core';
 import { CurrentUser, Public } from '../../core/auth';
-import { CopilotContextService } from './context/service';
-import { CopilotProviderFactory } from './providers/factory';
-import type { CopilotProvider } from './providers/provider';
 import {
-  ModelInputType,
-  ModelOutputType,
-  type StreamObject,
-} from './providers/types';
-import { StreamObjectParser } from './providers/utils';
-import { ChatSession, ChatSessionService } from './session';
+  ActionStreamHost,
+  projectActionEventToChatEvent,
+} from './runtime/hosts/action-stream-host';
+import { TurnOrchestrator } from './runtime/turn-orchestrator';
 import { CopilotStorage } from './storage';
-import { ChatMessage, ChatQuerySchema } from './types';
-import { getSignal, getTools } from './utils';
-import { CopilotWorkflowService, GraphExecutorState } from './workflow';
+import { getSignal } from './utils';
 
 export interface ChatEvent {
   type: 'event' | 'attachment' | 'message' | 'error' | 'ping';
@@ -74,11 +59,8 @@ export class CopilotController implements BeforeApplicationShutdown {
 
   constructor(
     private readonly config: Config,
-    private readonly server: ServerService,
-    private readonly chatSession: ChatSessionService,
-    private readonly context: CopilotContextService,
-    private readonly provider: CopilotProviderFactory,
-    private readonly workflow: CopilotWorkflowService,
+    private readonly orchestrator: TurnOrchestrator,
+    private readonly actionStreams: ActionStreamHost,
     private readonly storage: CopilotStorage
   ) {}
 
@@ -92,85 +74,6 @@ export class CopilotController implements BeforeApplicationShutdown {
     this.ongoingStreamCount$.complete();
   }
 
-  private async chooseProvider(
-    outputType: ModelOutputType,
-    userId: string,
-    sessionId: string,
-    messageId?: string,
-    modelId?: string
-  ): Promise<{
-    provider: CopilotProvider;
-    model: string;
-    hasAttachment: boolean;
-  }> {
-    const [, session] = await Promise.all([
-      this.chatSession.checkQuota(userId),
-      this.chatSession.get(sessionId),
-    ]);
-
-    if (!session || session.config.userId !== userId) {
-      throw new CopilotSessionNotFound();
-    }
-
-    const model = await session.resolveModel(
-      this.server.features.includes(ServerFeature.Payment),
-      modelId
-    );
-
-    const hasAttachment = messageId
-      ? !!(await session.getMessageById(messageId)).attachments?.length
-      : false;
-
-    const provider = await this.provider.getProvider({
-      outputType,
-      modelId: model,
-    });
-    if (!provider) {
-      throw new NoCopilotProviderAvailable({ modelId: model });
-    }
-
-    return { provider, model, hasAttachment };
-  }
-
-  private async appendSessionMessage(
-    sessionId: string,
-    messageId?: string,
-    retry = false
-  ): Promise<[ChatMessage | undefined, ChatSession]> {
-    const session = await this.chatSession.get(sessionId);
-    if (!session) {
-      throw new CopilotSessionNotFound();
-    }
-
-    let latestMessage = undefined;
-    if (!messageId || retry) {
-      // revert the latest message generated by the assistant
-      // if messageId is provided, we will also revert latest user message
-      await this.chatSession.revertLatestMessage(sessionId, !!messageId);
-      session.revertLatestMessage(!!messageId);
-      if (!messageId) {
-        latestMessage = session.latestUserMessage;
-      }
-    }
-
-    if (messageId) {
-      await session.pushByMessageId(messageId);
-    }
-
-    return [latestMessage, session];
-  }
-
-  private parseNumber(value: string | string[] | undefined) {
-    if (!value) {
-      return undefined;
-    }
-    const num = Number.parseInt(Array.isArray(value) ? value[0] : value, 10);
-    if (Number.isNaN(num)) {
-      return undefined;
-    }
-    return num;
-  }
-
   private mergePingStream(
     messageId: string,
     source$: Observable<ChatEvent>
@@ -184,59 +87,12 @@ export class CopilotController implements BeforeApplicationShutdown {
     return merge(source$.pipe(finalize(() => subject$.next(null))), ping$);
   }
 
-  private async prepareChatSession(
-    user: CurrentUser,
-    sessionId: string,
-    query: Record<string, string | string[]>,
-    outputType: ModelOutputType
-  ) {
-    let { messageId, retry, modelId, params } = ChatQuerySchema.parse(query);
+  private toMessageEvent(messageId: string | undefined, data: string | object) {
+    return { type: 'message' as const, id: messageId, data };
+  }
 
-    const { provider, model } = await this.chooseProvider(
-      outputType,
-      user.id,
-      sessionId,
-      messageId,
-      modelId
-    );
-
-    const [latestMessage, session] = await this.appendSessionMessage(
-      sessionId,
-      messageId,
-      retry
-    );
-
-    const context = await this.context.getBySessionId(sessionId);
-    const contextParams =
-      (Array.isArray(context?.files) && context.files.length > 0) ||
-      (Array.isArray(context?.blobs) && context.blobs.length > 0)
-        ? {
-            contextFiles: [
-              ...context.files,
-              ...(await context.getBlobMetadata()),
-            ],
-          }
-        : {};
-    const lastParams = latestMessage
-      ? {
-          ...latestMessage.params,
-          content: latestMessage.content,
-          attachments: latestMessage.attachments,
-        }
-      : {};
-
-    const finalMessage = session.finish({
-      ...params,
-      ...lastParams,
-      ...contextParams,
-    });
-
-    return {
-      provider,
-      model,
-      session,
-      finalMessage,
-    };
+  private toAttachmentEvent(messageId: string | undefined, data: string) {
+    return { type: 'attachment' as const, id: messageId, data };
   }
 
   @Sse('/chat/:sessionId/stream')
@@ -250,19 +106,6 @@ export class CopilotController implements BeforeApplicationShutdown {
     const info: any = { sessionId, params: query, throwInStream: false };
 
     try {
-      const { provider, model, session, finalMessage } =
-        await this.prepareChatSession(
-          user,
-          sessionId,
-          query,
-          ModelOutputType.Text
-        );
-
-      info.model = model;
-      info.finalMessage = finalMessage.filter(m => m.role !== 'system');
-      metrics.ai.counter('chat_stream_calls').add(1, { model });
-      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
-
       const { signal, onConnectionClosed } = getSignal(req);
       let endBeforePromiseResolve = false;
       onConnectionClosed(isAborted => {
@@ -271,51 +114,23 @@ export class CopilotController implements BeforeApplicationShutdown {
         }
       });
 
-      const { messageId, reasoning, webSearch, toolsConfig } =
-        ChatQuerySchema.parse(query);
+      const prepared = await this.orchestrator.streamText(
+        user.id,
+        sessionId,
+        query,
+        signal,
+        () => endBeforePromiseResolve
+      );
 
-      const source$ = from(
-        provider.streamText({ modelId: model }, finalMessage, {
-          ...session.config.promptConfig,
-          signal,
-          user: user.id,
-          session: session.config.sessionId,
-          workspace: session.config.workspaceId,
-          reasoning,
-          webSearch,
-          tools: getTools(session.config.promptConfig?.tools, toolsConfig),
-        })
-      ).pipe(
-        connect(shared$ =>
-          merge(
-            // actual chat event stream
-            shared$.pipe(
-              map(data => ({ type: 'message' as const, id: messageId, data }))
-            ),
-            // save the generated text to the session
-            shared$.pipe(
-              reduce((acc, chunk) => acc + chunk, ''),
-              tap(buffer => {
-                session.push({
-                  role: 'assistant',
-                  content: endBeforePromiseResolve
-                    ? '> Request aborted'
-                    : buffer,
-                  createdAt: new Date(),
-                });
-                void session
-                  .save()
-                  .catch(err =>
-                    this.logger.error(
-                      'Failed to save session in sse stream',
-                      err
-                    )
-                  );
-              }),
-              ignoreElements()
-            )
-          )
-        ),
+      info.model = prepared.model;
+      info.finalMessage = prepared.finalMessage.filter(
+        m => m.role !== 'system'
+      );
+      metrics.ai.counter('chat_stream_calls').add(1, { model: prepared.model });
+      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
+
+      const source$ = from(prepared.stream).pipe(
+        map(data => this.toMessageEvent(prepared.messageId, data)),
         catchError(e => {
           metrics.ai.counter('chat_stream_errors').add(1);
           info.throwInStream = true;
@@ -326,7 +141,7 @@ export class CopilotController implements BeforeApplicationShutdown {
         })
       );
 
-      return this.mergePingStream(messageId || '', source$);
+      return this.mergePingStream(prepared.messageId || '', source$);
     } catch (err) {
       metrics.ai.counter('chat_stream_errors').add(1, info);
       return mapSseError(err, info);
@@ -344,19 +159,6 @@ export class CopilotController implements BeforeApplicationShutdown {
     const info: any = { sessionId, params: query, throwInStream: false };
 
     try {
-      const { provider, model, session, finalMessage } =
-        await this.prepareChatSession(
-          user,
-          sessionId,
-          query,
-          ModelOutputType.Object
-        );
-
-      info.model = model;
-      info.finalMessage = finalMessage.filter(m => m.role !== 'system');
-      metrics.ai.counter('chat_object_stream_calls').add(1, { model });
-      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
-
       const { signal, onConnectionClosed } = getSignal(req);
       let endBeforePromiseResolve = false;
       onConnectionClosed(isAborted => {
@@ -365,55 +167,25 @@ export class CopilotController implements BeforeApplicationShutdown {
         }
       });
 
-      const { messageId, reasoning, webSearch, toolsConfig } =
-        ChatQuerySchema.parse(query);
+      const prepared = await this.orchestrator.streamObject(
+        user.id,
+        sessionId,
+        query,
+        signal,
+        () => endBeforePromiseResolve
+      );
 
-      const source$ = from(
-        provider.streamObject({ modelId: model }, finalMessage, {
-          ...session.config.promptConfig,
-          signal,
-          user: user.id,
-          session: session.config.sessionId,
-          workspace: session.config.workspaceId,
-          reasoning,
-          webSearch,
-          tools: getTools(session.config.promptConfig?.tools, toolsConfig),
-        })
-      ).pipe(
-        connect(shared$ =>
-          merge(
-            // actual chat event stream
-            shared$.pipe(
-              map(data => ({ type: 'message' as const, id: messageId, data }))
-            ),
-            // save the generated text to the session
-            shared$.pipe(
-              reduce((acc, chunk) => acc.concat([chunk]), [] as StreamObject[]),
-              tap(result => {
-                const parser = new StreamObjectParser();
-                const streamObjects = parser.mergeTextDelta(result);
-                const content = parser.mergeContent(streamObjects);
-                session.push({
-                  role: 'assistant',
-                  content: endBeforePromiseResolve
-                    ? '> Request aborted'
-                    : content,
-                  streamObjects: endBeforePromiseResolve ? null : streamObjects,
-                  createdAt: new Date(),
-                });
-                void session
-                  .save()
-                  .catch(err =>
-                    this.logger.error(
-                      'Failed to save session in sse stream',
-                      err
-                    )
-                  );
-              }),
-              ignoreElements()
-            )
-          )
-        ),
+      info.model = prepared.model;
+      info.finalMessage = prepared.finalMessage.filter(
+        m => m.role !== 'system'
+      );
+      metrics.ai.counter('chat_object_stream_calls').add(1, {
+        model: prepared.model,
+      });
+      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
+
+      const source$ = from(prepared.stream).pipe(
+        map(data => this.toMessageEvent(prepared.messageId, data)),
         catchError(e => {
           metrics.ai.counter('chat_object_stream_errors').add(1);
           info.throwInStream = true;
@@ -424,16 +196,16 @@ export class CopilotController implements BeforeApplicationShutdown {
         })
       );
 
-      return this.mergePingStream(messageId || '', source$);
+      return this.mergePingStream(prepared.messageId || '', source$);
     } catch (err) {
       metrics.ai.counter('chat_object_stream_errors').add(1, info);
       return mapSseError(err, info);
     }
   }
 
-  @Sse('/chat/:sessionId/workflow')
-  @CallMetric('ai', 'chat_workflow', { timer: true })
-  async chatWorkflow(
+  @Sse('/actions/:sessionId/stream')
+  @CallMetric('ai', 'action_stream', { timer: true })
+  async actionStream(
     @CurrentUser() user: CurrentUser,
     @Req() req: Request,
     @Param('sessionId') sessionId: string,
@@ -441,103 +213,26 @@ export class CopilotController implements BeforeApplicationShutdown {
   ): Promise<Observable<ChatEvent>> {
     const info: any = { sessionId, params: query, throwInStream: false };
     try {
-      let { messageId, params } = ChatQuerySchema.parse(query);
+      const { signal } = getSignal(req);
 
-      const [, session] = await this.appendSessionMessage(sessionId, messageId);
-      info.model = session.model;
-
-      metrics.ai.counter('workflow_calls').add(1, { model: session.model });
-
-      const latestMessage = session.stashMessages.findLast(
-        m => m.role === 'user'
+      const prepared = await this.actionStreams.stream(
+        user.id,
+        sessionId,
+        query,
+        signal
       );
-      if (latestMessage) {
-        params = Object.assign({}, params, latestMessage.params, {
-          content: latestMessage.content,
-          attachments: latestMessage.attachments,
-        });
-      }
+      info.actionId = prepared.actionId;
+      info.actionVersion = prepared.actionVersion;
+      metrics.ai.counter('action_stream_calls').add(1, {
+        actionId: prepared.actionId,
+        actionVersion: prepared.actionVersion,
+      });
       this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
 
-      const { signal, onConnectionClosed } = getSignal(req);
-      let endBeforePromiseResolve = false;
-      onConnectionClosed(isAborted => {
-        if (isAborted) {
-          endBeforePromiseResolve = true;
-        }
-      });
-
-      const source$ = from(
-        this.workflow.runGraph(params, session.model, {
-          ...session.config.promptConfig,
-          signal,
-          user: user.id,
-          session: session.config.sessionId,
-          workspace: session.config.workspaceId,
-        })
-      ).pipe(
-        connect(shared$ =>
-          merge(
-            // actual chat event stream
-            shared$.pipe(
-              map(data => {
-                switch (data.status) {
-                  case GraphExecutorState.EmitContent:
-                    return {
-                      type: 'message' as const,
-                      id: messageId,
-                      data: data.content,
-                    };
-                  case GraphExecutorState.EmitAttachment:
-                    return {
-                      type: 'attachment' as const,
-                      id: messageId,
-                      data: data.attachment,
-                    };
-                  default:
-                    return {
-                      type: 'event' as const,
-                      id: messageId,
-                      data: {
-                        status: data.status,
-                        id: data.node.id,
-                        type: data.node.config.nodeType,
-                      },
-                    };
-                }
-              })
-            ),
-            // save the generated text to the session
-            shared$.pipe(
-              reduce((acc, chunk) => {
-                if (chunk.status === GraphExecutorState.EmitContent) {
-                  acc += chunk.content;
-                }
-                return acc;
-              }, ''),
-              tap(content => {
-                session.push({
-                  role: 'assistant',
-                  content: endBeforePromiseResolve
-                    ? '> Request aborted'
-                    : content,
-                  createdAt: new Date(),
-                });
-                void session
-                  .save()
-                  .catch(err =>
-                    this.logger.error(
-                      'Failed to save session in sse stream',
-                      err
-                    )
-                  );
-              }),
-              ignoreElements()
-            )
-          )
-        ),
+      const source$ = from(prepared.stream).pipe(
+        map(data => projectActionEventToChatEvent(prepared.messageId, data)),
         catchError(e => {
-          metrics.ai.counter('workflow_errors').add(1, info);
+          metrics.ai.counter('action_stream_errors').add(1, info);
           info.throwInStream = true;
           return mapSseError(e, info);
         }),
@@ -546,9 +241,9 @@ export class CopilotController implements BeforeApplicationShutdown {
         )
       );
 
-      return this.mergePingStream(messageId || '', source$);
+      return this.mergePingStream(prepared.messageId || '', source$);
     } catch (err) {
-      metrics.ai.counter('workflow_errors').add(1, info);
+      metrics.ai.counter('action_stream_errors').add(1, info);
       return mapSseError(err, info);
     }
   }
@@ -563,36 +258,6 @@ export class CopilotController implements BeforeApplicationShutdown {
   ): Promise<Observable<ChatEvent>> {
     const info: any = { sessionId, params: query, throwInStream: false };
     try {
-      let { messageId, params } = ChatQuerySchema.parse(query);
-
-      const { provider, model, hasAttachment } = await this.chooseProvider(
-        ModelOutputType.Image,
-        user.id,
-        sessionId,
-        messageId
-      );
-
-      const [latestMessage, session] = await this.appendSessionMessage(
-        sessionId,
-        messageId
-      );
-      info.model = model;
-      metrics.ai.counter('images_stream_calls').add(1, { model });
-
-      if (latestMessage) {
-        params = Object.assign({}, params, latestMessage.params, {
-          content: latestMessage.content,
-          attachments: latestMessage.attachments,
-        });
-      }
-
-      const handleRemoteLink = this.storage.handleRemoteLink.bind(
-        this.storage,
-        user.id,
-        sessionId
-      );
-      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
-
       const { signal, onConnectionClosed } = getSignal(req);
       let endBeforePromiseResolve = false;
       onConnectionClosed(isAborted => {
@@ -601,59 +266,22 @@ export class CopilotController implements BeforeApplicationShutdown {
         }
       });
 
-      const source$ = from(
-        provider.streamImages(
-          {
-            modelId: model,
-            inputTypes: hasAttachment
-              ? [ModelInputType.Image]
-              : [ModelInputType.Text],
-          },
-          session.finish(params),
-          {
-            ...session.config.promptConfig,
-            quality: params.quality || undefined,
-            seed: this.parseNumber(params.seed),
-            signal,
-            user: user.id,
-            session: session.config.sessionId,
-            workspace: session.config.workspaceId,
-          }
-        )
-      ).pipe(
-        mergeMap(handleRemoteLink),
-        connect(shared$ =>
-          merge(
-            // actual chat event stream
-            shared$.pipe(
-              map(attachment => ({
-                type: 'attachment' as const,
-                id: messageId,
-                data: attachment,
-              }))
-            ),
-            // save the generated text to the session
-            shared$.pipe(
-              reduce((acc, chunk) => acc.concat([chunk]), [] as string[]),
-              tap(attachments => {
-                session.push({
-                  role: 'assistant',
-                  content: endBeforePromiseResolve ? '> Request aborted' : '',
-                  attachments: endBeforePromiseResolve ? [] : attachments,
-                  createdAt: new Date(),
-                });
-                void session
-                  .save()
-                  .catch(err =>
-                    this.logger.error(
-                      'Failed to save session in sse stream',
-                      err
-                    )
-                  );
-              }),
-              ignoreElements()
-            )
-          )
+      const prepared = await this.orchestrator.streamImages(
+        user.id,
+        sessionId,
+        query,
+        signal,
+        () => endBeforePromiseResolve
+      );
+      info.model = prepared.model;
+      metrics.ai.counter('images_stream_calls').add(1, {
+        model: prepared.model,
+      });
+      this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
+
+      const source$ = from(prepared.stream).pipe(
+        map(attachment =>
+          this.toAttachmentEvent(prepared.messageId, attachment)
         ),
         catchError(e => {
           metrics.ai.counter('images_stream_errors').add(1, info);
@@ -665,7 +293,7 @@ export class CopilotController implements BeforeApplicationShutdown {
         )
       );
 
-      return this.mergePingStream(messageId || '', source$);
+      return this.mergePingStream(prepared.messageId || '', source$);
     } catch (err) {
       metrics.ai.counter('images_stream_errors').add(1, info);
       return mapSseError(err, info);