1
0
Эх сурвалжийг харах

feat(core): support pinned Code Mode tools (#39550)

Aiden Cline 2 долоо хоног өмнө
parent
commit
5cb633a48e

+ 19 - 5
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<Entry>, 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<typeof Listing.Type>(),
+        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)
     }
   }
 

+ 15 - 6
packages/core/src/codemode/tool.ts

@@ -138,7 +138,14 @@ export const create = (
 }
 
 export const catalog = (registrations: ReadonlyMap<string, Info>) => {
-  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<string, Tool.Tool<never>> = {}
   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<typeof tools>({ 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<string, typeof Schema.Json.Type> | undefined {
   if (input === null || input === undefined) return

+ 2 - 0
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<string>",
+          pinned: true,
         },
       ])
     }).pipe(

+ 26 - 1
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<string>`,
+  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")

+ 1 - 1
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 }),
     ),
   )
 

+ 1 - 1
packages/core/test/session-runner.test.ts

@@ -231,7 +231,7 @@ const permission = Layer.succeed(
 const transformTools = (registry: Tool.Interface, tools: Readonly<Record<string, Info>>, 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(

+ 13 - 2
packages/schema/src/tool.ts

@@ -19,12 +19,23 @@ export interface Context {
   readonly progress: (update: Metadata) => Effect.Effect<void>
 }
 
-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<A = unknown> =
   | Schema.Codec<A, any>
   | (StandardSchemaV1<any, A> & StandardJSONSchemaV1<any, A>)