From 5cb633a48e5bdacb56cea3a75cceac8b656560e2 Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:16:49 -0500 Subject: [PATCH] feat(core): support pinned Code Mode tools (#39550) --- packages/core/src/codemode/catalog.ts | 24 +++++++++++++---- packages/core/src/codemode/tool.ts | 21 ++++++++++----- packages/core/test/codemode.test.ts | 2 ++ packages/core/test/codemode/catalog.test.ts | 27 ++++++++++++++++++- .../test/session-runner-tool-registry.test.ts | 2 +- packages/core/test/session-runner.test.ts | 2 +- packages/schema/src/tool.ts | 15 +++++++++-- 7 files changed, 77 insertions(+), 16 deletions(-) diff --git a/packages/core/src/codemode/catalog.ts b/packages/core/src/codemode/catalog.ts index eafc1568949..00d1b71e2d9 100644 --- a/packages/core/src/codemode/catalog.ts +++ b/packages/core/src/codemode/catalog.ts @@ -6,6 +6,7 @@ export const Entry = Schema.Struct({ path: Schema.String, description: Schema.String, signature: Schema.String, + pinned: Schema.optionalKey(Schema.Boolean), }) export type Entry = typeof Entry.Type @@ -56,26 +57,39 @@ export function summarize(entries: ReadonlyArray, budget = INLINE_BUDGET) if (left.path > right.path) return 1 return 0 }) + const ranked = rankListings(listings) + const pinned = new Set( + namespaceEntries + .filter((entry) => entry.pinned) + .map((entry) => listings.find((listing) => listing.path === entry.path)) + .filter((listing) => listing !== undefined), + ) return { name, listings, - selectionOrder: rankListings(listings), - selectedListings: new Set(), + selectionOrder: ranked.filter((candidate) => !pinned.has(candidate.listing)), + selectedListings: pinned, + selectionIndex: 0, } }) const active = new Set(namespaces) - let remaining = budget + let remaining = + budget - + namespaces + .flatMap((namespace) => namespace.listings.filter((listing) => namespace.selectedListings.has(listing))) + .reduce((total, listing) => total + Math.round(listing.line.length / CHARACTERS_PER_TOKEN), 0) while (active.size > 0) { for (const namespace of active) { - const candidate = namespace.selectionOrder[namespace.selectedListings.size] + const candidate = namespace.selectionOrder[namespace.selectionIndex] if (!candidate || candidate.cost > remaining) { active.delete(namespace) continue } namespace.selectedListings.add(candidate.listing) + namespace.selectionIndex += 1 remaining -= candidate.cost - if (namespace.selectedListings.size === namespace.selectionOrder.length) active.delete(namespace) + if (namespace.selectionIndex === namespace.selectionOrder.length) active.delete(namespace) } } diff --git a/packages/core/src/codemode/tool.ts b/packages/core/src/codemode/tool.ts index b0e6d7b3192..b0a13ac58ff 100644 --- a/packages/core/src/codemode/tool.ts +++ b/packages/core/src/codemode/tool.ts @@ -138,7 +138,14 @@ export const create = ( } export const catalog = (registrations: ReadonlyMap) => { - return runtime(registrations, () => Effect.fail(toolError("Execute context is unavailable"))).catalog() + const pinned = new Set( + Array.from(registrations.values()) + .filter((registration) => registration.options?.pinned === true) + .map(qualifiedName), + ) + return runtime(registrations, () => Effect.fail(toolError("Execute context is unavailable"))) + .catalog() + .map((entry) => ({ ...entry, pinned: pinned.has(entry.path) })) } function runtime( @@ -149,11 +156,7 @@ function runtime( const tools: Record> = {} for (const [name, registration] of registrations) { const child = definition(registration) - const normalized = registration.name.replace(/[^a-zA-Z0-9_-]/g, "_") - const path = - registration.options?.namespace === undefined - ? normalized - : `${registration.options.namespace}.${normalized}` + const path = qualifiedName(registration) tools[path] = Tool.make({ description: child.description, input: child.inputSchema, @@ -164,6 +167,12 @@ function runtime( return CodeMode.make({ tools, ...hooks }) } +function qualifiedName(registration: Info) { + const normalized = registration.name.replace(/[^a-zA-Z0-9_-]/g, "_") + if (registration.options?.namespace === undefined) return normalized + return `${registration.options.namespace}.${normalized}` +} + // Tool inputs arrive as parsed JSON, so the JSON value cast is a boundary fact. function displayInput(input: unknown): Record | undefined { if (input === null || input === undefined) return diff --git a/packages/core/test/codemode.test.ts b/packages/core/test/codemode.test.ts index 2c3776b7412..25a5f87941d 100644 --- a/packages/core/test/codemode.test.ts +++ b/packages/core/test/codemode.test.ts @@ -16,6 +16,7 @@ describe("CodeMode", () => { description: "Echo text", input: Schema.Struct({ text: Schema.String }), output: Schema.String, + options: { pinned: true }, execute: ({ text }) => Effect.succeed({ output: text }), }), ) @@ -27,6 +28,7 @@ describe("CodeMode", () => { path: "echo", description: "Echo text", signature: "tools.echo(input: {\n text: string,\n}): Promise", + pinned: true, }, ]) }).pipe( diff --git a/packages/core/test/codemode/catalog.test.ts b/packages/core/test/codemode/catalog.test.ts index 1168c7ee67b..afa1f063895 100644 --- a/packages/core/test/codemode/catalog.test.ts +++ b/packages/core/test/codemode/catalog.test.ts @@ -2,10 +2,11 @@ import { describe, expect, test } from "bun:test" import { CodeModeCatalog } from "@opencode-ai/core/codemode/catalog" import { CodeModeInstructions } from "@opencode-ai/core/codemode/instructions" -const entry = (path: string, description: string, signature?: string): CodeModeCatalog.Entry => ({ +const entry = (path: string, description: string, signature?: string, pinned = false): CodeModeCatalog.Entry => ({ path, description, signature: signature ?? `tools.${path}(input: {\n q: string,\n}): Promise`, + pinned, }) const lookup = entry( @@ -46,6 +47,30 @@ describe("CodeModeCatalog.summarize", () => { expect(catalog.namespaces.every((namespace) => namespace.entries.length === 0)).toBe(true) }) + test("always retains pinned tools beyond the inline budget", () => { + const pinned = [ + entry("alpha.first", "First", undefined, true), + entry("beta.second", "Second", undefined, true), + ] + const catalog = CodeModeCatalog.summarize([...pinned, entry("alpha.unpinned", "Unpinned")], 0) + + expect(catalog.shown).toBe(2) + expect(catalog.namespaces.flatMap((namespace) => namespace.entries.map((item) => item.path))).toEqual([ + "alpha.first", + "beta.second", + ]) + }) + + test("spends the budget remaining after pinned tools on unpinned tools", () => { + const pinned = entry("alpha.pinned", "Pinned", undefined, true) + const unpinned = entry("beta.unpinned", "Unpinned") + const pinCost = Math.round(` - ${pinned.signature} // Pinned`.length / 4) + const unpinnedCost = Math.round(` - ${unpinned.signature} // Unpinned`.length / 4) + + expect(CodeModeCatalog.summarize([pinned, unpinned], pinCost + unpinnedCost).shown).toBe(2) + expect(CodeModeCatalog.summarize([pinned, unpinned], pinCost + unpinnedCost - 1).shown).toBe(1) + }) + test("retains only the rendered portion of inline descriptions", () => { const catalog = CodeModeCatalog.summarize([entry("alpha.one", `Summary\n${"detail".repeat(10_000)}`)]) expect(catalog.namespaces[0]?.entries[0]?.line).toEndWith("// Summary") diff --git a/packages/core/test/session-runner-tool-registry.test.ts b/packages/core/test/session-runner-tool-registry.test.ts index 1398a749791..aacbb700692 100644 --- a/packages/core/test/session-runner-tool-registry.test.ts +++ b/packages/core/test/session-runner-tool-registry.test.ts @@ -67,7 +67,7 @@ const transform = ( ) => service.transform((draft) => Object.entries(tools).forEach(([name, tool]) => - draft.add({ ...tool, name, options: { ...tool.options, ...options } }), + draft.add({ ...tool, name, options: options ?? tool.options }), ), ) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index c6b7d7b9c7d..6bbf8d90a75 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -231,7 +231,7 @@ const permission = Layer.succeed( const transformTools = (registry: Tool.Interface, tools: Readonly>, options?: Tool.Options) => registry.transform((draft) => Object.entries(tools).forEach(([name, tool]) => - draft.add({ ...tool, name, options: { ...tool.options, ...options } }), + draft.add({ ...tool, name, options: options ?? tool.options }), ), ) const echo = Layer.effectDiscard( diff --git a/packages/schema/src/tool.ts b/packages/schema/src/tool.ts index 1787ebc1839..f0f122ba2ec 100644 --- a/packages/schema/src/tool.ts +++ b/packages/schema/src/tool.ts @@ -19,12 +19,23 @@ export interface Context { readonly progress: (update: Metadata) => Effect.Effect } -export interface Options { +interface BaseOptions { readonly namespace?: string - readonly codemode?: boolean readonly permission?: string } +export type Options = BaseOptions & + ( + | { + readonly codemode?: true + readonly pinned?: boolean + } + | { + readonly codemode: boolean + readonly pinned?: never + } + ) + export type ValueSchema = | Schema.Codec | (StandardSchemaV1 & StandardJSONSchemaV1)