fix(context-agent): add status callbacks back

vscode/src/chat/agentic/CodyChatAgent.ts

@@ -3,6 +3,7 @@ import {
     type ChatClient,
     type ContextItem,
     type Message,
+    type ProcessingStep,
     type PromptMixin,
     type PromptString,
     newPromptMixin,
@@ -19,18 +20,34 @@ export abstract class CodyChatAgent {
     protected readonly promptMixins: PromptMixin[] = []
     protected readonly toolHandlers: Map<string, CodyTool>
     protected statusCallback?: ToolStatusCallback
-    protected postMessageCallback?: (model: string) => void
+    private stepsManager?: ProcessManager
 
     constructor(
         protected readonly chatBuilder: ChatBuilder,
         protected readonly chatClient: Pick<ChatClient, 'chat'>,
         protected readonly tools: CodyTool[],
-        protected context: ContextItem[] = []
+        protected context: ContextItem[],
+        statusUpdateCallback?: (steps: ProcessingStep[]) => void
     ) {
         // Initialize handlers and mixins in constructor
         this.toolHandlers = new Map(tools.map(tool => [tool.config.tags.tag.toString(), tool]))
         this.initializeMultiplexer()
         this.promptMixins.push(newPromptMixin(this.buildPrompt()))
+
+        this.stepsManager = new ProcessManager(steps => {
+            statusUpdateCallback?.(steps)
+        })
+        this.statusCallback = {
+            onStart: () => {
+                this.stepsManager?.initializeStep()
+            },
+            onStream: (toolName, content) => {
+                this.stepsManager?.addStep(toolName, content)
+            },
+            onComplete: (toolName, error) => {
+                this.stepsManager?.completeStep(toolName, error)
+            },
+        }
     }
 
     protected initializeMultiplexer(): void {
@@ -86,29 +103,6 @@ export abstract class CodyChatAgent {
         return new DefaultPrompter(explicitMentions, implicitMentions.slice(-MAX_SEARCH_ITEMS))
     }
 
-    public setStatusCallback(postMessage: (model: string) => void): void {
-        this.postMessageCallback = postMessage
-        const model = this.chatBuilder.selectedModel ?? ''
-
-        // Create a steps manager to handle state updates efficiently
-        const stepsManager = new ProcessManager(steps => {
-            this.chatBuilder.setLastMessageProcesses(steps)
-            this.postMessageCallback?.(model)
-        })
-
-        this.statusCallback = {
-            onStart: () => {
-                stepsManager.initializeStep()
-            },
-            onStream: (toolName, content) => {
-                stepsManager.addStep(toolName, content)
-            },
-            onComplete: (toolName, error) => {
-                stepsManager.completeStep(toolName, error)
-            },
-        }
-    }
-
     // Abstract methods that must be implemented by derived classes
     protected abstract buildPrompt(): PromptString
 }

vscode/src/chat/chat-view/ChatController.ts

@@ -18,6 +18,7 @@ import {
     type Guardrails,
     ModelUsage,
     type NLSSearchDynamicFilter,
+    type ProcessingStep,
     PromptString,
     type SerializedChatInteraction,
     type SerializedChatTranscript,
@@ -842,6 +843,9 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                         messageInProgress = message
                         this.postViewTranscript(message)
                     },
+                    postStatuses: (steps: ProcessingStep[]): void => {
+                        this.chatBuilder.setLastMessageProcesses(steps)
+                    },
                     postDone: (op?: { abort: boolean }): void => {
                         if (op?.abort) {
                             this.handleAbort()

vscode/src/chat/chat-view/handlers/ChatHandler.ts

@@ -57,6 +57,7 @@ export class ChatHandler implements AgentHandler {
             { text: inputText, mentions },
             editorState,
             chatBuilder,
+            delegate,
             signal
         )
         if (contextResult.error) {
@@ -223,6 +224,7 @@ export class ChatHandler implements AgentHandler {
         { text, mentions }: HumanInput,
         editorState: SerializedPromptEditorState | null,
         _chatBuilder: ChatBuilder,
+        _delegate: AgentHandlerDelegate,
         signal?: AbortSignal
     ): Promise<{
         contextItems?: ContextItem[]

vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts

@@ -1,6 +1,7 @@
 import {
     type ContextItem,
     FeatureFlag,
+    type ProcessingStep,
     type SerializedPromptEditorState,
     featureFlagProvider,
     storeLastValue,
@@ -13,7 +14,7 @@ import type { ChatControllerOptions } from '../ChatController'
 import type { ContextRetriever } from '../ContextRetriever'
 import type { HumanInput } from '../context'
 import { ChatHandler } from './ChatHandler'
-import type { AgentHandler } from './interfaces'
+import type { AgentHandler, AgentHandlerDelegate } from './interfaces'
 
 export class DeepCodyHandler extends ChatHandler implements AgentHandler {
     constructor(
@@ -38,6 +39,7 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler {
         { text, mentions }: HumanInput,
         editorState: SerializedPromptEditorState | null,
         chatBuilder: ChatBuilder,
+        delegate: AgentHandlerDelegate,
         signal: AbortSignal
     ): Promise<{
         contextItems?: ContextItem[]
@@ -49,6 +51,7 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler {
             { text, mentions },
             editorState,
             chatBuilder,
+            delegate,
             signal
         )
         const isEnabled = chatBuilder.getMessages().length < 4
@@ -66,13 +69,14 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler {
         }
 
         const baseContext = baseContextResult.contextItems ?? []
-        const codyAgent = new DeepCodyAgent(
+        const agent = new DeepCodyAgent(
             chatBuilder,
             this.chatClient,
             this.toolProvider.getTools(),
-            baseContext
+            baseContext,
+            (steps: ProcessingStep[]) => delegate.postStatuses(steps)
         )
-        const agenticContext = await codyAgent.getContext(requestID, signal)
+        const agenticContext = await agent.getContext(requestID, signal)
         return { contextItems: [...baseContext, ...agenticContext] }
     }
 }

vscode/src/chat/chat-view/handlers/interfaces.ts

@@ -2,6 +2,7 @@ import type { Span } from '@opentelemetry/api'
 import type {
     ChatMessage,
     ContextItem,
+    ProcessingStep,
     PromptString,
     SerializedPromptEditorState,
 } from '@sourcegraph/cody-shared'
@@ -24,6 +25,7 @@ export interface AgentTools {
  */
 export interface AgentHandlerDelegate {
     postError(error: Error, type?: MessageErrorType): void
+    postStatuses(steps: ProcessingStep[]): void
     postMessageInProgress(message: ChatMessage): void
     postDone(ops?: { abort: boolean }): void
 }