handle an AbortSignal

Author: julien-c

packages/mcp-client/cli.ts

@@ -35,10 +35,6 @@ if (process.env.EXPERIMENTAL_HF_MCP_SERVER) {
 	});
 }
 
-let SAMPLE_INPUT = process.env.USE_SAMPLE_INPUT
-	? `generate a haiku about Hugging Face and save it to a file name hf.txt on my Desktop`
-	: undefined;
-
 async function main() {
 	if (!process.env.HF_TOKEN) {
 		console.error(`a valid HF_TOKEN must be present in the env`);
@@ -53,14 +49,33 @@ async function main() {
 	});
 
 	const rl = readline.createInterface({ input: stdin, output: stdout });
+	let abortController = new AbortController();
+	let waitingForInput = false;
+	async function waitForInput() {
+		waitingForInput = true;
+		const input = await rl.question("> ");
+		waitingForInput = false;
+		return input;
+	}
 	rl.on("SIGINT", async () => {
-		await agent.cleanup();
-		stdout.write("\n");
-		rl.close();
+		if (waitingForInput) {
+			// close the whole process
+			await agent.cleanup();
+			stdout.write("\n");
+			rl.close();
+		} else {
+			// otherwise, it means a request is underway
+			abortController.abort();
+			abortController = new AbortController();
+			stdout.write(ANSI.GRAY);
+			stdout.write("Ctrl+C a second time to exit");
+			stdout.write(ANSI.RESET);
+		}
 	});
-	process.on("uncaughtException", () => {
+	process.on("uncaughtException", (err) => {
 		stdout.write("\n");
 		rl.close();
+		throw err;
 	});
 
 	await agent.loadTools();
@@ -72,8 +87,8 @@ async function main() {
 	stdout.write("\n");
 
 	while (true) {
-		const input = await rl.question("> ");
-		for await (const chunk of agent.run(input)) {
+		const input = await waitForInput();
+		for await (const chunk of agent.run(input, { abortSignal: abortController.signal })) {
 			if ("choices" in chunk) {
 				const delta = chunk.choices[0]?.delta;
 				if (delta.content) {

packages/mcp-client/src/Agent.ts

@@ -72,7 +72,10 @@ export class Agent extends McpClient {
 		return this.addMcpServers(this.servers);
 	}
 
-	async *run(input: string): AsyncGenerator<ChatCompletionStreamOutput | ChatCompletionInputMessageTool> {
+	async *run(
+		input: string,
+		opts: { abortSignal?: AbortSignal } = {}
+	): AsyncGenerator<ChatCompletionStreamOutput | ChatCompletionInputMessageTool> {
 		this.messages.push({
 			role: "user",
 			content: input,
@@ -81,10 +84,18 @@ export class Agent extends McpClient {
 		let numOfTurns = 0;
 		let nextTurnShouldCallTools = true;
 		while (true) {
-			yield* this.processSingleTurnWithTools(this.messages, {
-				exitLoopTools,
-				exitIfFirstChunkNoTool: numOfTurns > 0 && nextTurnShouldCallTools,
-			});
+			try {
+				yield* this.processSingleTurnWithTools(this.messages, {
+					exitLoopTools,
+					exitIfFirstChunkNoTool: numOfTurns > 0 && nextTurnShouldCallTools,
+					abortSignal: opts.abortSignal,
+				});
+			} catch (err) {
+				if (err instanceof Error && err.message === "AbortError") {
+					return;
+				}
+				throw err;
+			}
 			numOfTurns++;
 			// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 			const currentLast = this.messages.at(-1)!;

packages/mcp-client/src/McpClient.ts

@@ -72,7 +72,11 @@ export class McpClient {
 
 	async *processSingleTurnWithTools(
 		messages: ChatCompletionInputMessage[],
-		opts: { exitLoopTools?: ChatCompletionInputTool[]; exitIfFirstChunkNoTool?: boolean } = {}
+		opts: {
+			exitLoopTools?: ChatCompletionInputTool[];
+			exitIfFirstChunkNoTool?: boolean;
+			abortSignal?: AbortSignal;
+		} = {}
 	): AsyncGenerator<ChatCompletionStreamOutput | ChatCompletionInputMessageTool> {
 		debug("start of single turn");
 
@@ -82,6 +86,7 @@ export class McpClient {
 			messages,
 			tools: opts.exitLoopTools ? [...opts.exitLoopTools, ...this.availableTools] : this.availableTools,
 			tool_choice: "auto",
+			signal: opts.abortSignal,
 		});
 
 		const message = {
@@ -92,6 +97,9 @@ export class McpClient {
 		let numOfChunks = 0;
 
 		for await (const chunk of stream) {
+			if (opts.abortSignal?.aborted) {
+				throw new Error("AbortError");
+			}
 			yield chunk;
 			debug(chunk.choices[0]);
 			numOfChunks++;
@@ -143,7 +151,7 @@ export class McpClient {
 			/// Get the appropriate session for this tool
 			const client = this.clients.get(toolName);
 			if (client) {
-				const result = await client.callTool({ name: toolName, arguments: toolArgs });
+				const result = await client.callTool({ name: toolName, arguments: toolArgs, signal: opts.abortSignal });
 				toolMessage.content = (result.content as Array<{ text: string }>)[0].text;
 			} else {
 				toolMessage.content = `Error: No session found for tool: ${toolName}`;