Add support for subagent hooks

src/extension/intents/node/toolCallingLoop.ts

@@ -7,7 +7,7 @@ import * as l10n from '@vscode/l10n';
 import { Raw } from '@vscode/prompt-tsx';
 import type { CancellationToken, ChatRequest, ChatResponseProgressPart, ChatResponseReferencePart, ChatResponseStream, ChatResult, LanguageModelToolInformation, Progress } from 'vscode';
 import { IAuthenticationChatUpgradeService } from '../../../platform/authentication/common/authenticationUpgrade';
-import { IChatHookService, StopHookInput, StopHookOutput } from '../../../platform/chat/common/chatHookService';
+import { IChatHookService, StopHookInput, StopHookOutput, SubagentStartHookInput, SubagentStartHookOutput, SubagentStopHookInput, SubagentStopHookOutput } from '../../../platform/chat/common/chatHookService';
 import { FetchStreamSource, IResponsePart } from '../../../platform/chat/common/chatMLFetcher';
 import { CanceledResult, ChatFetchResponseType, ChatResponse } from '../../../platform/chat/common/commonTypes';
 import { IConfigurationService } from '../../../platform/configuration/common/configurationService';
@@ -100,6 +100,24 @@ interface StopHookResult {
 	readonly reason?: string;
 }
 
+interface SubagentStartHookResult {
+	/**
+	 * Additional context to add to the subagent's context, if any.
+	 */
+	readonly additionalContext?: string;
+}
+
+interface SubagentStopHookResult {
+	/**
+	 * Whether the subagent should continue (not stop).
+	 */
+	readonly shouldContinue: boolean;
+	/**
+	 * The reason the subagent should continue, if shouldContinue is true.
+	 */
+	readonly reason?: string;
+}
+
 /**
  * This is a base class that can be used to implement a tool calling loop
  * against a model. It requires only that you build a prompt and is decoupled
@@ -112,6 +130,7 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 	private toolCallResults: Record<string, LanguageModelToolResult2> = Object.create(null);
 	private toolCallRounds: IToolCallRound[] = [];
 	private stopHookReason: string | undefined;
+	private additionalHookContext: string | undefined;
 
 	private readonly _onDidBuildPrompt = this._register(new Emitter<{ result: IBuildPromptResult; tools: LanguageModelToolInformation[]; promptTokenLength: number; toolTokenCount: number }>());
 	public readonly onDidBuildPrompt = this._onDidBuildPrompt.event;
@@ -186,6 +205,7 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 			isContinuation,
 			hasStopHookQuery,
 			modeInstructions: this.options.request.modeInstructions2,
+			additionalHookContext: this.additionalHookContext,
 		};
 	}
 
@@ -248,6 +268,98 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 		this._logService.trace(`[ToolCallingLoop] Stop hook blocked stopping: ${reason}`);
 	}
 
+	/**
+	 * Called when a subagent starts to allow hooks to provide additional context.
+	 * @param input The subagent start hook input containing agent_id and agent_type
+	 * @param token Cancellation token
+	 * @returns Result containing additional context from hooks
+	 */
+	protected async executeSubagentStartHook(input: SubagentStartHookInput, token: CancellationToken | PauseController): Promise<SubagentStartHookResult> {
+		try {
+			const results = await this._chatHookService.executeHook('SubagentStart', {
+				toolInvocationToken: this.options.request.toolInvocationToken,
+				input: input
+			}, token);
+
+			// Collect additionalContext from all successful hook results
+			const additionalContexts: string[] = [];
+			for (const result of results) {
+				if (result.success === true) {
+					const output = result.output;
+					if (typeof output === 'object' && output !== null) {
+						const hookOutput = output as SubagentStartHookOutput;
+						if (hookOutput.additionalContext) {
+							additionalContexts.push(hookOutput.additionalContext);
+							this._logService.trace(`[ToolCallingLoop] SubagentStart hook provided context: ${hookOutput.additionalContext.substring(0, 100)}...`);
+						}
+					}
+				} else if (result.success === false) {
+					const errorMessage = typeof result.output === 'string' ? result.output : 'Unknown error';
+					this._logService.error(`[ToolCallingLoop] SubagentStart hook error: ${errorMessage}`);
+				}
+			}
+
+			return {
+				additionalContext: additionalContexts.length > 0 ? additionalContexts.join('\n') : undefined
+			};
+		} catch (error) {
+			this._logService.error('[ToolCallingLoop] Error executing SubagentStart hook', error);
+			return {};
+		}
+	}
+
+	/**
+	 * Called before a subagent stops to give hooks a chance to block the stop.
+	 * @param input The subagent stop hook input containing agent_id, agent_type, and stop_hook_active flag
+	 * @param outputStream The output stream for displaying messages
+	 * @param token Cancellation token
+	 * @returns Result indicating whether to continue and the reason
+	 */
+	protected async executeSubagentStopHook(input: SubagentStopHookInput, outputStream: ChatResponseStream | undefined, token: CancellationToken | PauseController): Promise<SubagentStopHookResult> {
+		try {
+			const results = await this._chatHookService.executeHook('SubagentStop', {
+				toolInvocationToken: this.options.request.toolInvocationToken,
+				input: input
+			}, token);
+
+			// Check for blocking responses
+			for (const result of results) {
+				if (result.success === true) {
+					const output = result.output;
+					if (typeof output === 'object' && output !== null) {
+						const hookOutput = output as SubagentStopHookOutput;
+						this._logService.trace(`[ToolCallingLoop] Checking SubagentStop hook output: decision=${hookOutput.decision}, reason=${hookOutput.reason}`);
+						if (hookOutput.decision === 'block' && hookOutput.reason) {
+							this._logService.trace(`[ToolCallingLoop] SubagentStop hook blocked: ${hookOutput.reason}`);
+							return { shouldContinue: true, reason: hookOutput.reason };
+						}
+					}
+				} else if (result.success === false) {
+					const errorMessage = typeof result.output === 'string' ? result.output : 'Unknown error';
+					this._logService.error(`[ToolCallingLoop] SubagentStop hook error: ${errorMessage}`);
+				}
+			}
+
+			return { shouldContinue: false };
+		} catch (error) {
+			this._logService.error('[ToolCallingLoop] Error executing SubagentStop hook', error);
+			return { shouldContinue: false };
+		}
+	}
+
+	/**
+	 * Shows a message when the subagent stop hook blocks the subagent from stopping.
+	 * Override in subclasses to customize the display.
+	 * @param outputStream The output stream for displaying messages
+	 * @param reason The reason the subagent stop hook blocked stopping
+	 */
+	protected showSubagentStopHookBlockedMessage(outputStream: ChatResponseStream | undefined, reason: string): void {
+		if (outputStream) {
+			outputStream.markdown('\n\n' + l10n.t('**Subagent stop hook:** {0}', reason) + '\n\n');
+		}
+		this._logService.trace(`[ToolCallingLoop] SubagentStop hook blocked stopping: ${reason}`);
+	}
+
 	private async throwIfCancelled(token: CancellationToken | PauseController) {
 		if (await this.checkAsync(token)) {
 			throw new CancellationError();
@@ -260,6 +372,18 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 		let lastRequestMessagesStartingIndexForRun: number | undefined;
 		let stopHookActive = false;
 
+		// Execute SubagentStart hook for subagent requests to get additional context
+		if (this.options.request.subAgentInvocationId) {
+			const startHookResult = await this.executeSubagentStartHook({
+				agent_id: this.options.request.subAgentInvocationId,
+				agent_type: this.options.request.subAgentName ?? 'default'
+			}, token);
+			if (startHookResult.additionalContext) {
+				this.additionalHookContext = startHookResult.additionalContext;
+				this._logService.info(`[ToolCallingLoop] SubagentStart hook provided context for subagent ${this.options.request.subAgentInvocationId}`);
+			}
+		}
+
 		while (true) {
 			if (lastResult && i++ >= this.options.toolCallLimit) {
 				lastResult = this.hitToolCallLimit(outputStream, lastResult);
@@ -279,16 +403,34 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 				this.toolCallRounds.push(result.round);
 				if (!result.round.toolCalls.length || result.response.type !== ChatFetchResponseType.Success) {
 					// Before stopping, execute the stop hook
-					const stopHookResult = await this.executeStopHook({ stop_hook_active: stopHookActive }, outputStream, token);
-					this._logService.info(`[ToolCallingLoop] Stop hook result: shouldContinue=${stopHookResult.shouldContinue}, reason=${stopHookResult.reason}`);
-					if (stopHookResult.shouldContinue && stopHookResult.reason) {
-						// The stop hook blocked stopping - show reason and continue
-						this.showStopHookBlockedMessage(outputStream, stopHookResult.reason);
-						// Store the reason so it can be passed to the model in the next prompt
-						this.stopHookReason = stopHookResult.reason;
-						this._logService.info(`[ToolCallingLoop] Stop hook blocked, continuing with reason: ${stopHookResult.reason}`);
-						stopHookActive = true;
-						continue;
+					if (this.options.request.subAgentInvocationId) {
+						const stopHookResult = await this.executeSubagentStopHook({
+							agent_id: this.options.request.subAgentInvocationId,
+							agent_type: this.options.request.subAgentName ?? 'default',
+							stop_hook_active: stopHookActive
+						}, outputStream, token);
+						this._logService.info(`[ToolCallingLoop] Subagent stop hook result: shouldContinue=${stopHookResult.shouldContinue}, reason=${stopHookResult.reason}`);
+						if (stopHookResult.shouldContinue && stopHookResult.reason) {
+							// The stop hook blocked stopping - show reason and continue
+							this.showSubagentStopHookBlockedMessage(outputStream, stopHookResult.reason);
+							// Store the reason so it can be passed to the model in the next prompt
+							this.stopHookReason = stopHookResult.reason;
+							this._logService.info(`[ToolCallingLoop] Subagent stop hook blocked, continuing with reason: ${stopHookResult.reason}`);
+							stopHookActive = true;
+							continue;
+						}
+					} else {
+						const stopHookResult = await this.executeStopHook({ stop_hook_active: stopHookActive }, outputStream, token);
+						this._logService.info(`[ToolCallingLoop] Stop hook result: shouldContinue=${stopHookResult.shouldContinue}, reason=${stopHookResult.reason}`);
+						if (stopHookResult.shouldContinue && stopHookResult.reason) {
+							// The stop hook blocked stopping - show reason and continue
+							this.showStopHookBlockedMessage(outputStream, stopHookResult.reason);
+							// Store the reason so it can be passed to the model in the next prompt
+							this.stopHookReason = stopHookResult.reason;
+							this._logService.info(`[ToolCallingLoop] Stop hook blocked, continuing with reason: ${stopHookResult.reason}`);
+							stopHookActive = true;
+							continue;
+						}
 					}
 					lastResult = lastResult;
 					break;

src/extension/prompt/common/intents.ts

@@ -101,6 +101,10 @@ export interface IBuildPromptContext {
 	 * continuation that requires a specific user message.
 	 */
 	readonly hasStopHookQuery?: boolean;
+	/**
+	 * Additional context provided by a hook.
+	 */
+	readonly additionalHookContext?: string;
 }
 
 export const IBuildPromptContext = createServiceIdentifier<IBuildPromptContext>('IBuildPromptContext');

src/extension/prompts/node/agent/agentPrompt.tsx

@@ -279,6 +279,8 @@ export interface AgentUserMessageProps extends BasePromptElementProps, AgentUser
 	readonly sessionResource?: string;
 	/** When true, indicates this is a stop hook continuation where the stop hook query is rendered as a separate message. */
 	readonly hasStopHookQuery?: boolean;
+	/** Additional context provided by SubagentStart hooks. */
+	readonly additionalHookContext?: string;
 }
 
 export function getUserMessagePropsFromTurn(turn: Turn, endpoint: IChatEndpoint, customizations?: AgentUserMessageCustomizations): AgentUserMessageProps {
@@ -307,6 +309,7 @@ export function getUserMessagePropsFromAgentProps(agentProps: AgentPromptProps,
 		enableCacheBreakpoints: agentProps.enableCacheBreakpoints,
 		editedFileEvents: agentProps.promptContext.editedFileEvents,
 		hasStopHookQuery: agentProps.promptContext.hasStopHookQuery,
+		additionalHookContext: agentProps.promptContext.additionalHookContext,
 		// TODO:@roblourens
 		sessionId: (agentProps.promptContext.tools?.toolInvocationToken as any)?.sessionId,
 		sessionResource: (agentProps.promptContext.tools?.toolInvocationToken as any)?.sessionResource,
@@ -377,6 +380,7 @@ export class AgentUserMessage extends PromptElement<AgentUserMessageProps> {
 						<NotebookSummaryChange />
 						{hasTerminalTool && <TerminalStatePromptElement sessionId={this.props.sessionId} />}
 						{hasTodoTool && <TodoListContextPrompt sessionResource={this.props.sessionResource} />}
+						{this.props.additionalHookContext && <AdditionalHookContextPrompt context={this.props.additionalHookContext} />}
 					</Tag>
 					<CurrentEditorContext endpoint={this.props.endpoint} />
 					<Tag name='reminderInstructions'>
@@ -457,6 +461,19 @@ class CurrentDatePrompt extends PromptElement<BasePromptElementProps> {
 	}
 }
 
+interface AdditionalHookContextPromptProps extends BasePromptElementProps {
+	readonly context: string;
+}
+
+/**
+ * Renders additional context provided by hooks.
+ */
+class AdditionalHookContextPrompt extends PromptElement<AdditionalHookContextPromptProps> {
+	render(state: void, sizing: PromptSizing) {
+		return <>Additional instructions from hooks: {this.props.context}</>;
+	}
+}
+
 interface SkillAdherenceReminderProps extends BasePromptElementProps {
 	readonly chatVariables: ChatVariablesCollection;
 }

src/platform/chat/common/chatHookService.ts

@@ -61,4 +61,62 @@ export interface StopHookOutput {
 	readonly reason?: string;
 }
 
+/**
+ * Input passed to the SubagentStart hook.
+ */
+export interface SubagentStartHookInput {
+	/**
+	 * The unique identifier for the subagent.
+	 */
+	readonly agent_id: string;
+	/**
+	 * The agent name (built-in agents like "Plan" or custom agent names).
+	 */
+	readonly agent_type: string;
+}
+
+/**
+ * Output from the SubagentStart hook.
+ */
+export interface SubagentStartHookOutput {
+	/**
+	 * Additional context to add to the subagent's context.
+	 */
+	readonly additionalContext?: string;
+}
+
+/**
+ * Input passed to the SubagentStop hook.
+ */
+export interface SubagentStopHookInput {
+	/**
+	 * The unique identifier for the subagent.
+	 */
+	readonly agent_id: string;
+	/**
+	 * The agent name (built-in agents like "Plan" or custom agent names).
+	 */
+	readonly agent_type: string;
+	/**
+	 * True when the agent is already continuing as a result of a stop hook.
+	 * Check this value or process the transcript to prevent the agent from running indefinitely.
+	 */
+	readonly stop_hook_active: boolean;
+}
+
+/**
+ * Output from the SubagentStop hook.
+ */
+export interface SubagentStopHookOutput {
+	/**
+	 * Set to "block" to prevent the agent from stopping.
+	 * Omit or set to undefined to allow the agent to stop.
+	 */
+	readonly decision?: 'block';
+	/**
+	 * Required when decision is "block". Tells the agent why it should continue.
+	 */
+	readonly reason?: string;
+}
+
 //#endregion