Prechádzať zdrojové kódy

fix(core): preserve cache for session generation

Dax Raad 3 týždňov pred
rodič
commit
89e3141079

+ 28 - 5
packages/core/src/session/generate-node.ts

@@ -12,6 +12,7 @@ import { SessionGenerate } from "./generate"
 import { SessionHistory } from "./history"
 import { SessionModelHeaders } from "./model-headers"
 import { SessionRunnerModel } from "./runner/model"
+import { ToolRegistry } from "../tool/registry"
 import PROMPT_DEFAULT from "./runner/prompt/base.txt"
 import { toLLMMessages } from "./runner/to-llm-message"
 
@@ -23,6 +24,7 @@ export const layer = Layer.effect(
     const hooks = yield* PluginHooks.Service
     const llm = yield* LLMClient.Service
     const models = yield* SessionRunnerModel.Service
+    const registry = yield* ToolRegistry.Service
     const app = yield* App.Metadata
 
     return SessionGenerate.Service.of({
@@ -34,6 +36,9 @@ export const layer = Layer.effect(
         const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(selection.session.id)
           ? selection.session.id.slice(4)
           : selection.session.id
+        const executableTools = yield* registry.materialize(selection.agent.info.permissions)
+        const toolDefinitions = executableTools.definitions
+        const toolsByName = new Map(toolDefinitions.map((tool) => [tool.name, tool]))
         const contextEvent = yield* hooks.trigger("session", "context", {
           sessionID: selection.session.id,
           agent: selection.agent.id,
@@ -46,24 +51,34 @@ export const layer = Layer.effect(
             ...(history.instructionUpdate ? [Message.system(history.instructionUpdate)] : []),
             Message.user(input.prompt),
           ],
-          tools: {},
+          tools: Object.fromEntries(
+            toolDefinitions.map((tool) => [tool.name, { description: tool.description, input: { ...tool.inputSchema } }]),
+          ),
+        })
+        const hookedTools = Object.entries(contextEvent.tools).flatMap(([name, tool]) => {
+          const registered = toolsByName.get(name)
+          return registered
+            ? [Object.assign({}, registered, { description: tool.description, inputSchema: tool.input })]
+            : []
         })
         yield* Effect.logInfo("sending session generation request", {
           sessionID: selection.session.id,
           providerID: model.ref.providerID,
           modelID: model.ref.id,
         })
-        return (yield* llm.generate(
+        const response = yield* llm.generate(
           LLM.request({
             model: model.model,
             http: { headers: SessionModelHeaders.make(selection.session, app) },
             providerOptions: { openai: { promptCacheKey } },
             system: contextEvent.system,
             messages: contextEvent.messages,
-            tools: [],
+            tools: hookedTools,
             toolChoice: "none",
           }),
-        )).text
+        )
+        yield* Effect.logInfo("session generation usage diagnostic", { usage: response.usage })
+        return response.text
       }),
     })
   }),
@@ -72,5 +87,13 @@ export const layer = Layer.effect(
 export const node = makeLocationNode({
   service: SessionGenerate.Service,
   layer,
-  deps: [SessionContext.node, Database.node, PluginHooks.node, SessionRunnerModel.node, App.node, llmClient],
+  deps: [
+    SessionContext.node,
+    Database.node,
+    PluginHooks.node,
+    SessionRunnerModel.node,
+    ToolRegistry.node,
+    App.node,
+    llmClient,
+  ],
 })

+ 14 - 2
packages/core/test/session-generate.test.ts

@@ -1,5 +1,5 @@
 import { expect } from "bun:test"
-import { LLMClient, LLMEvent, LLMResponse, Model, SystemPart, type LLMRequest } from "@opencode-ai/ai"
+import { LLMClient, LLMEvent, LLMResponse, Model, SystemPart, ToolDefinition, type LLMRequest } from "@opencode-ai/ai"
 import { OpenAIChat } from "@opencode-ai/ai/protocols"
 import { AgentV2 } from "@opencode-ai/core/agent"
 import { Database } from "@opencode-ai/core/database/database"
@@ -38,6 +38,7 @@ import { SessionStore } from "@opencode-ai/core/session/store"
 import { SkillInstructions } from "@opencode-ai/core/skill/instructions"
 import { PluginHooks } from "@opencode-ai/core/plugin/hooks"
 import { PluginSupervisor } from "@opencode-ai/core/plugin/supervisor"
+import { ToolRegistry } from "@opencode-ai/core/tool/registry"
 import { asc, eq } from "drizzle-orm"
 import { Effect, Layer, Schema, Stream } from "effect"
 import { testEffect } from "./lib/effect"
@@ -92,6 +93,15 @@ const skills = Layer.mock(SkillInstructions.Service, { load: () => Effect.succee
 const references = Layer.mock(ReferenceInstructions.Service, { load: () => Effect.succeed(Instructions.empty) })
 const mcp = Layer.mock(McpInstructions.Service, { load: () => Effect.succeed(Instructions.empty) })
 const plugins = Layer.mock(PluginSupervisor.Service, { flush: Effect.void })
+const tools = Layer.mock(ToolRegistry.Service, {
+  materialize: () =>
+    Effect.succeed({
+      definitions: [ToolDefinition.make({ name: "lookup", description: "Lookup", inputSchema: { type: "object" } })],
+      settle: () => Effect.die(new Error("unused")),
+    }),
+  register: () => Effect.die(new Error("unused")),
+  registerBatch: () => Effect.die(new Error("unused")),
+})
 
 const it = testEffect(
   AppNodeBuilder.build(
@@ -114,6 +124,7 @@ const it = testEffect(
       [ReferenceInstructions.node, references],
       [McpInstructions.node, mcp],
       [PluginSupervisor.node, plugins],
+      [ToolRegistry.node, tools],
       [Location.node, Location.boundNode({ directory: AbsolutePath.make("/project") })],
     ],
   ),
@@ -259,6 +270,7 @@ it.effect("generates from fresh settled Session context without durable mutation
     yield* hooks.register("session", "context", (event) =>
       Effect.sync(() => {
         event.system = [SystemPart.make("Hooked system"), ...event.system]
+        if (event.tools.lookup) event.tools.lookup.description = "Hooked lookup"
       }),
     )
 
@@ -287,7 +299,7 @@ it.effect("generates from fresh settled Session context without durable mutation
           : [],
       ),
     ).toEqual(["Settled partial answer"])
-    expect(requests[0]?.tools).toEqual([])
+    expect(requests[0]?.tools).toMatchObject([{ name: "lookup", description: "Hooked lookup" }])
     expect(requests[0]?.toolChoice).toMatchObject({ type: "none" })
     expect(yield* durableState(db, sessionID)).toEqual(before)
   }),