From 0bff5bf389d5d623d354ed6608235e6685bb781a Mon Sep 17 00:00:00 2001
From: Aiden Cline <aidenpcline@gmail.com>
Date: Mon, 29 Jun 2026 00:39:48 -0500
Subject: [PATCH] fix(opencode): avoid deferred MCP search at step limit

---
 packages/opencode/src/session/mcp-tools.ts   |  3 ++-
 packages/opencode/src/session/prompt.ts      |  1 +
 packages/opencode/src/session/tools.ts       |  1 +
 packages/opencode/test/session/tools.test.ts | 27 ++++++++++++++------
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/packages/opencode/src/session/mcp-tools.ts b/packages/opencode/src/session/mcp-tools.ts
index cd04f50575..c1a7b4f0c0 100644
--- a/packages/opencode/src/session/mcp-tools.ts
+++ b/packages/opencode/src/session/mcp-tools.ts
@@ -51,6 +51,7 @@ interface Input {
   bypassAgentCheck: boolean
   messages: SessionV1.WithParts[]
   promptOps: TaskPromptOps
+  canDeferMcpTools?: boolean
 }
 
 interface DeferredToolDescriptor {
@@ -243,7 +244,7 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In
     Object.entries(mcpTools).filter(([key]) => userTools?.[key] !== false && !mcpDisabled.has(key)),
   )
   const deferredDescriptors =
-    flags.experimentalToolSearch && Object.keys(allowedMcpTools).length > 0
+    (input.canDeferMcpTools ?? true) && flags.experimentalToolSearch && Object.keys(allowedMcpTools).length > 0
       ? yield* deferredToolDescriptors(allowedMcpTools)
       : []
   const deferMcpTools =
diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts
index 116839323e..dbc35babad 100644
--- a/packages/opencode/src/session/prompt.ts
+++ b/packages/opencode/src/session/prompt.ts
@@ -1231,6 +1231,7 @@ export const layer = Layer.effect(
               bypassAgentCheck,
               messages: msgs,
               promptOps,
+              canDeferMcpTools: maxSteps - step >= 2,
             }).pipe(
               Effect.provideService(Plugin.Service, plugin),
               Effect.provideService(Permission.Service, permission),
diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts
index 4c4c660e01..49374f24e6 100644
--- a/packages/opencode/src/session/tools.ts
+++ b/packages/opencode/src/session/tools.ts
@@ -28,6 +28,7 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
   bypassAgentCheck: boolean
   messages: SessionV1.WithParts[]
   promptOps: TaskPromptOps
+  canDeferMcpTools?: boolean
 }) {
   const tools: Record<string, Tool> = {}
   const run = yield* EffectBridge.make()
diff --git a/packages/opencode/test/session/tools.test.ts b/packages/opencode/test/session/tools.test.ts
index d0d03871d5..37646d3794 100644
--- a/packages/opencode/test/session/tools.test.ts
+++ b/packages/opencode/test/session/tools.test.ts
@@ -132,15 +132,16 @@ const belowThresholdIt = makeIt({
   queryDescription: "Natural language analytics query",
 })
 
-function resolveTools(messages: SessionV1.WithParts[] = []) {
+function resolveTools(input: { messages?: SessionV1.WithParts[]; canDeferMcpTools?: boolean } = {}) {
   return SessionTools.resolve({
     agent,
     model,
     session,
     processor,
     bypassAgentCheck: false,
-    messages,
+    messages: input.messages ?? [],
     promptOps,
+    canDeferMcpTools: input.canDeferMcpTools,
   })
 }
 
@@ -189,12 +190,14 @@ describe("session.tools", () => {
 
   deferredIt.instance("does not expose per-message disabled MCP tools through deferred search", () =>
     Effect.gen(function* () {
-      const tools = yield* resolveTools([
-        {
-          info: { ...user, tools: { posthog_feature_flags: false } },
-          parts: [],
-        },
-      ])
+      const tools = yield* resolveTools({
+        messages: [
+          {
+            info: { ...user, tools: { posthog_feature_flags: false } },
+            parts: [],
+          },
+        ],
+      })
 
       const search = tools.search_deferred_tools.execute
       if (!search) throw new Error("missing search_deferred_tools executor")
@@ -229,4 +232,12 @@ describe("session.tools", () => {
       expect(Object.keys(tools).sort()).toEqual(["posthog_feature_flags", "posthog_query_trends"])
     }),
   )
+
+  deferredIt.instance("keeps MCP tools direct when deferred search cannot finish before the last step", () =>
+    Effect.gen(function* () {
+      const tools = yield* resolveTools({ canDeferMcpTools: false })
+
+      expect(Object.keys(tools).sort()).toEqual(["posthog_feature_flags", "posthog_query_trends"])
+    }),
+  )
 })