From e934d5e23ed5751efa03784641147f8359216cd4 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 24 Jul 2025 17:08:35 -0700
Subject: [PATCH] chore: retain the source code from the underlying tools
 (#756)

---
 src/loop/loop.ts         | 10 ++++------
 src/loop/loopClaude.ts   |  5 +----
 src/loop/loopOpenAI.ts   |  5 +----
 src/loop/main.ts         |  9 +++++----
 src/loopTools/context.ts | 16 ++++++++++++++--
 5 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/src/loop/loop.ts b/src/loop/loop.ts
index 5546ff3..43a8e2e 100644
--- a/src/loop/loop.ts
+++ b/src/loop/loop.ts
@@ -47,7 +47,7 @@ export interface LLMDelegate {
   checkDoneToolCall(toolCall: LLMToolCall): string | null;
 }
 
-export async function runTask(delegate: LLMDelegate, client: Client, task: string, oneShot: boolean = false): Promise<string> {
+export async function runTask(delegate: LLMDelegate, client: Client, task: string, oneShot: boolean = false): Promise<LLMMessage[]> {
   const { tools } = await client.listTools();
   const taskContent = oneShot ? `Perform following task: ${task}.` : `Perform following task: ${task}. Once the task is complete, call the "done" tool.`;
   const conversation = delegate.createConversation(taskContent, tools, oneShot);
@@ -60,10 +60,9 @@ export async function runTask(delegate: LLMDelegate, client: Client, task: strin
 
     const toolResults: Array<{ toolCallId: string; content: string; isError?: boolean }> = [];
     for (const toolCall of toolCalls) {
-      // Check if this is the "done" tool
       const doneResult = delegate.checkDoneToolCall(toolCall);
       if (doneResult !== null)
-        return doneResult;
+        return conversation.messages;
 
       const { name, arguments: args, id } = toolCall;
       try {
@@ -100,10 +99,9 @@ export async function runTask(delegate: LLMDelegate, client: Client, task: strin
       }
     }
 
+    delegate.addToolResults(conversation, toolResults);
     if (oneShot)
-      return toolResults.map(result => result.content).join('\n');
-    else
-      delegate.addToolResults(conversation, toolResults);
+      return conversation.messages;
   }
 
   throw new Error('Failed to perform step, max attempts reached');
diff --git a/src/loop/loopClaude.ts b/src/loop/loopClaude.ts
index 2fd3ed2..92dc1d5 100644
--- a/src/loop/loopClaude.ts
+++ b/src/loop/loopClaude.ts
@@ -44,10 +44,7 @@ export class ClaudeDelegate implements LLMDelegate {
         description: 'Call this tool when the task is complete.',
         inputSchema: {
           type: 'object',
-          properties: {
-            result: { type: 'string', description: 'The result of the task.' },
-          },
-          required: ['result'],
+          properties: {},
         },
       });
     }
diff --git a/src/loop/loopOpenAI.ts b/src/loop/loopOpenAI.ts
index 224b19d..81e2ef2 100644
--- a/src/loop/loopOpenAI.ts
+++ b/src/loop/loopOpenAI.ts
@@ -44,10 +44,7 @@ export class OpenAIDelegate implements LLMDelegate {
         description: 'Call this tool when the task is complete.',
         inputSchema: {
           type: 'object',
-          properties: {
-            result: { type: 'string', description: 'The result of the task.' },
-          },
-          required: ['result'],
+          properties: {},
         },
       });
     }
diff --git a/src/loop/main.ts b/src/loop/main.ts
index 8ad22be..b909d4b 100644
--- a/src/loop/main.ts
+++ b/src/loop/main.ts
@@ -49,10 +49,11 @@ async function run(delegate: LLMDelegate) {
   await client.connect(transport);
   await client.ping();
 
-  let lastResult: string | undefined;
-  for (const task of tasks)
-    lastResult = await runTask(delegate, client, task);
-  console.log(lastResult);
+  for (const task of tasks) {
+    const messages = await runTask(delegate, client, task);
+    for (const message of messages)
+      console.log(`${message.role}: ${message.content}`);
+  }
   await client.close();
 }
 
diff --git a/src/loopTools/context.ts b/src/loopTools/context.ts
index 6f54e73..59d2b64 100644
--- a/src/loopTools/context.ts
+++ b/src/loopTools/context.ts
@@ -53,9 +53,21 @@ export class Context {
   }
 
   async runTask(task: string, oneShot: boolean = false): Promise<mcpServer.ToolResponse> {
-    const result = await runTask(this._delegate, this._client!, task, oneShot);
+    const messages = await runTask(this._delegate, this._client!, task, oneShot);
+    const lines: string[] = [];
+
+    // Skip the first message, which is the user's task.
+    for (const message of messages.slice(1)) {
+      // Trim out all page snapshots.
+      if (!message.content.trim())
+        continue;
+      const index = oneShot ? -1 : message.content.indexOf('### Page state');
+      const trimmedContent = index === -1 ? message.content : message.content.substring(0, index);
+      lines.push(`[${message.role}]:`, trimmedContent);
+    }
+
     return {
-      content: [{ type: 'text', text: result }],
+      content: [{ type: 'text', text: lines.join('\n') }],
     };
   }