From 15bc8ab92033a9ede2c73bb9b27f9d5d089f67dd Mon Sep 17 00:00:00 2001
From: Aiden Cline <aidenpcline@gmail.com>
Date: Sun, 28 Jun 2026 21:24:28 -0500
Subject: [PATCH] fix(opencode): avoid deferred MCP search near step limit

---
 packages/opencode/src/session/mcp-tools.ts   |  2 ++
 packages/opencode/src/session/prompt.ts      |  5 ++++-
 packages/opencode/src/session/tools.ts       |  1 +
 packages/opencode/test/session/tools.test.ts | 11 ++++++++++-
 4 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/packages/opencode/src/session/mcp-tools.ts b/packages/opencode/src/session/mcp-tools.ts
index 64a64dcd4b..4dcc206c11 100644
--- a/packages/opencode/src/session/mcp-tools.ts
+++ b/packages/opencode/src/session/mcp-tools.ts
@@ -53,6 +53,7 @@ interface Input {
   bypassAgentCheck: boolean
   messages: SessionV1.WithParts[]
   promptOps: TaskPromptOps
+  remainingSteps?: number
 }
 
 interface DeferredToolDescriptor {
@@ -220,6 +221,7 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In
       : []
   const deferMcpTools =
     deferredDescriptors.length > 0 &&
+    (input.remainingSteps ?? Infinity) >= 3 &&
     Token.estimate(JSON.stringify(deferredDescriptors.map(deferredToolEstimatePayload))) >=
       MIN_DEFERRED_MCP_SCHEMA_TOKENS
 
diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts
index 574ab6f937..1b89867f3c 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,
+              remainingSteps: Number.isFinite(maxSteps) ? maxSteps - step + 1 : undefined,
             }).pipe(
               Effect.provideService(Plugin.Service, plugin),
               Effect.provideService(Permission.Service, permission),
@@ -1265,7 +1266,9 @@ export const layer = Layer.effect(
               ...env,
               ...instructions,
               ...(mcpInstructions ? [mcpInstructions] : []),
-              ...(tools.search_deferred_tools && tools.call_deferred_tool ? [SessionTools.DEFERRED_TOOL_SYSTEM_PROMPT] : []),
+              ...(tools.search_deferred_tools && tools.call_deferred_tool
+                ? [SessionTools.DEFERRED_TOOL_SYSTEM_PROMPT]
+                : []),
               ...(skills ? [skills] : []),
             ]
             const format = lastUser.format ?? { type: "text" as const }
diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts
index 38368f988a..f79cee1345 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
+  remainingSteps?: number
 }) {
   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 7e5a3f2a55..95bd0f6fb5 100644
--- a/packages/opencode/test/session/tools.test.ts
+++ b/packages/opencode/test/session/tools.test.ts
@@ -119,7 +119,7 @@ const belowThresholdIt = makeIt({
   queryDescription: "Natural language analytics query",
 })
 
-function resolveTools() {
+function resolveTools(input: { remainingSteps?: number } = {}) {
   return SessionTools.resolve({
     agent,
     model,
@@ -128,6 +128,7 @@ function resolveTools() {
     bypassAgentCheck: false,
     messages: [],
     promptOps,
+    remainingSteps: input.remainingSteps,
   })
 }
 
@@ -172,4 +173,12 @@ describe("session.tools", () => {
       expect(Object.keys(tools).sort()).toEqual(["posthog_feature_flags", "posthog_query_trends"])
     }),
   )
+
+  deferredIt.instance("keeps MCP tools direct when the remaining step budget is too small", () =>
+    Effect.gen(function* () {
+      const tools = yield* resolveTools({ remainingSteps: 2 })
+
+      expect(Object.keys(tools).sort()).toEqual(["posthog_feature_flags", "posthog_query_trends"])
+    }),
+  )
 })