Просмотр исходного кода

refactor(opencode): extract session LLM request prep (#28560)

Kit Langton 3 месяцев назад
Родитель
Сommit
fb9d69ef62
2 измененных файлов с 248 добавлено и 191 удалено
  1. 40 191
      packages/opencode/src/session/llm.ts
  2. 208 0
      packages/opencode/src/session/llm/request.ts

+ 40 - 191
packages/opencode/src/session/llm.ts

@@ -1,41 +1,34 @@
 import { Provider } from "@/provider/provider"
 import * as Log from "@opencode-ai/core/util/log"
-import { Context, Effect, Layer, Record } from "effect"
+import { Context, Effect, Layer } from "effect"
 import * as Stream from "effect/Stream"
-import { streamText, wrapLanguageModel, type ModelMessage, type Tool, tool as aiTool, jsonSchema } from "ai"
+import { streamText, wrapLanguageModel, type ModelMessage, type Tool } from "ai"
 import type { LLMEvent } from "@opencode-ai/llm"
 import { LLMClient, RequestExecutor, WebSocketExecutor } from "@opencode-ai/llm/route"
 import type { LLMClientService } from "@opencode-ai/llm/route"
-import { mergeDeep } from "remeda"
 import { GitLabWorkflowLanguageModel } from "gitlab-ai-provider"
 import { ProviderTransform } from "@/provider/transform"
 import { Config } from "@/config/config"
-import { InstanceState } from "@/effect/instance-state"
 import type { Agent } from "@/agent/agent"
 import type { MessageV2 } from "./message-v2"
 import { Plugin } from "@/plugin"
-import { SystemPrompt } from "./system"
 import { Permission } from "@/permission"
 import { PermissionID } from "@/permission/schema"
 import { Bus } from "@/bus"
 import { Wildcard } from "@/util/wildcard"
 import { SessionID } from "@/session/schema"
 import { Auth } from "@/auth"
-import { InstallationVersion } from "@opencode-ai/core/installation/version"
 import { EffectBridge } from "@/effect/bridge"
 import { RuntimeFlags } from "@/effect/runtime-flags"
 import * as Option from "effect/Option"
 import * as OtelTracer from "@effect/opentelemetry/Tracer"
 import { LLMAISDK } from "./llm/ai-sdk"
 import { LLMNativeRuntime } from "./llm/native-runtime"
+import { LLMRequestPrep } from "./llm/request"
 
 const log = Log.create({ service: "llm" })
 export const OUTPUT_TOKEN_MAX = ProviderTransform.OUTPUT_TOKEN_MAX
 
-// Avoid re-instantiating remeda's deep merge types in this hot LLM path; the runtime behavior is still mergeDeep.
-const mergeOptions = (target: Record<string, any>, source: Record<string, any> | undefined): Record<string, any> =>
-  mergeDeep(target, source ?? {}) as Record<string, any>
-
 export type StreamInput = {
   user: MessageV2.User
   sessionID: string
@@ -106,123 +99,15 @@ const live: Layer.Layer<
         { concurrency: "unbounded" },
       )
 
-      // TODO: move this to a proper hook
-      const isOpenaiOauth = item.id === "openai" && info?.type === "oauth"
-
-      const system: string[] = []
-      system.push(
-        [
-          // use agent prompt otherwise provider prompt
-          ...(input.agent.prompt ? [input.agent.prompt] : SystemPrompt.provider(input.model)),
-          // any custom prompt passed into this call
-          ...input.system,
-          // any custom prompt from last user message
-          ...(input.user.system ? [input.user.system] : []),
-        ]
-          .filter((x) => x)
-          .join("\n"),
-      )
-
-      const header = system[0]
-      yield* plugin.trigger(
-        "experimental.chat.system.transform",
-        { sessionID: input.sessionID, model: input.model },
-        { system },
-      )
-      // rejoin to maintain 2-part structure for caching if header unchanged
-      if (system.length > 2 && system[0] === header) {
-        const rest = system.slice(1)
-        system.length = 0
-        system.push(header, rest.join("\n"))
-      }
-
-      const variant =
-        !input.small && input.model.variants && input.user.model.variant
-          ? input.model.variants[input.user.model.variant]
-          : {}
-      const base = input.small
-        ? ProviderTransform.smallOptions(input.model)
-        : ProviderTransform.options({
-            model: input.model,
-            sessionID: input.sessionID,
-            providerOptions: item.options,
-          })
-      const options = mergeOptions(mergeOptions(mergeOptions(base, input.model.options), input.agent.options), variant)
-      if (isOpenaiOauth) {
-        options.instructions = system.join("\n")
-      }
-
       const isWorkflow = language instanceof GitLabWorkflowLanguageModel
-      const messages = isOpenaiOauth
-        ? input.messages
-        : isWorkflow
-          ? input.messages
-          : [
-              ...system.map(
-                (x): ModelMessage => ({
-                  role: "system",
-                  content: x,
-                }),
-              ),
-              ...input.messages,
-            ]
-
-      const params = yield* plugin.trigger(
-        "chat.params",
-        {
-          sessionID: input.sessionID,
-          agent: input.agent.name,
-          model: input.model,
-          provider: item,
-          message: input.user,
-        },
-        {
-          temperature: input.model.capabilities.temperature
-            ? (input.agent.temperature ?? ProviderTransform.temperature(input.model))
-            : undefined,
-          topP: input.agent.topP ?? ProviderTransform.topP(input.model),
-          topK: ProviderTransform.topK(input.model),
-          maxOutputTokens: ProviderTransform.maxOutputTokens(input.model, flags.outputTokenMax),
-          options,
-        },
-      )
-
-      const { headers } = yield* plugin.trigger(
-        "chat.headers",
-        {
-          sessionID: input.sessionID,
-          agent: input.agent.name,
-          model: input.model,
-          provider: item,
-          message: input.user,
-        },
-        {
-          headers: {},
-        },
-      )
-
-      const tools = resolveTools(input)
-
-      // GitHub Copilot may require the tools parameter when message history contains
-      // tool calls but no tools are active (e.g. compaction). Inject a stub tool that
-      // is never meant to be invoked. LiteLLM-backed providers are excluded.
-      if (
-        input.model.providerID.includes("github-copilot") &&
-        Object.keys(tools).length === 0 &&
-        hasToolCalls(input.messages)
-      ) {
-        tools["_noop"] = aiTool({
-          description: "Do not call this tool. It exists only for API compatibility and must never be invoked.",
-          inputSchema: jsonSchema({
-            type: "object",
-            properties: {
-              reason: { type: "string", description: "Unused" },
-            },
-          }),
-          execute: async () => ({ output: "", title: "", metadata: {} }),
-        })
-      }
-      const sortedTools = Object.fromEntries(Object.entries(tools).toSorted(([a], [b]) => a.localeCompare(b)))
+      const prepared = yield* LLMRequestPrep.prepare({
+        ...input,
+        provider: item,
+        auth: info,
+        plugin,
+        flags,
+        isWorkflow,
+      })
 
       // Wire up toolExecutor for DWS workflow models so that tool calls
       // from the workflow service are executed via opencode's tool system
@@ -234,9 +119,9 @@ const live: Layer.Layer<
           approvalHandler?: (approvalTools: { name: string; args: string }[]) => Promise<{ approved: boolean }>
         }
         workflowModel.sessionID = input.sessionID
-        workflowModel.systemPrompt = system.join("\n")
+        workflowModel.systemPrompt = prepared.system.join("\n")
         workflowModel.toolExecutor = async (toolName, argsJson, _requestID) => {
-          const t = sortedTools[toolName]
+          const t = prepared.tools[toolName]
           if (!t || !t.execute) {
             return { result: "", error: `Unknown tool: ${toolName}` }
           }
@@ -258,7 +143,7 @@ const live: Layer.Layer<
         }
 
         const ruleset = Permission.merge(input.agent.permission ?? [], input.permission ?? [])
-        workflowModel.sessionPreapprovedTools = Object.keys(sortedTools).filter((name) => {
+        workflowModel.sessionPreapprovedTools = Object.keys(prepared.tools).filter((name) => {
           const match = ruleset.findLast((rule) => Wildcard.match(name, rule.permission))
           return !match || match.action !== "ask"
         })
@@ -327,28 +212,6 @@ const live: Layer.Layer<
           })
         : undefined
 
-      const opencodeProjectID = input.model.providerID.startsWith("opencode")
-        ? (yield* InstanceState.context).project.id
-        : undefined
-
-      const requestHeaders = {
-        ...(input.model.providerID.startsWith("opencode")
-          ? {
-              ...(opencodeProjectID ? { "x-opencode-project": opencodeProjectID } : {}),
-              "x-opencode-session": input.sessionID,
-              "x-opencode-request": input.user.id,
-              "x-opencode-client": flags.client,
-              "User-Agent": `opencode/${InstallationVersion}`,
-            }
-          : {
-              "x-session-affinity": input.sessionID,
-              ...(input.parentSessionID ? { "x-parent-session-id": input.parentSessionID } : {}),
-              "User-Agent": `opencode/${InstallationVersion}`,
-            }),
-        ...input.model.headers,
-        ...headers,
-      }
-
       // Runtime seam: native is an opt-in adapter over @opencode-ai/llm. It
       // either returns a ready LLMEvent stream or a concrete fallback reason.
       if (flags.experimentalNativeLlm) {
@@ -357,17 +220,17 @@ const live: Layer.Layer<
           provider: item,
           auth: info,
           llmClient,
-          isOpenaiOauth,
-          system,
-          messages,
-          tools: sortedTools,
+          isOpenaiOauth: prepared.isOpenaiOauth,
+          system: prepared.system,
+          messages: prepared.messages,
+          tools: prepared.tools,
           toolChoice: input.toolChoice,
-          temperature: params.temperature,
-          topP: params.topP,
-          topK: params.topK,
-          maxOutputTokens: params.maxOutputTokens,
-          providerOptions: params.options,
-          headers: requestHeaders,
+          temperature: prepared.params.temperature,
+          topP: prepared.params.topP,
+          topK: prepared.params.topK,
+          maxOutputTokens: prepared.params.maxOutputTokens,
+          providerOptions: prepared.params.options,
+          headers: prepared.headers,
           abort: input.abort,
         })
         if (native.type === "supported") {
@@ -413,7 +276,7 @@ const live: Layer.Layer<
           },
           async experimental_repairToolCall(failed) {
             const lower = failed.toolCall.toolName.toLowerCase()
-            if (lower !== failed.toolCall.toolName && sortedTools[lower]) {
+            if (lower !== failed.toolCall.toolName && prepared.tools[lower]) {
               l.info("repairing tool call", {
                 tool: failed.toolCall.toolName,
                 repaired: lower,
@@ -432,18 +295,18 @@ const live: Layer.Layer<
               toolName: "invalid",
             }
           },
-          temperature: params.temperature,
-          topP: params.topP,
-          topK: params.topK,
-          providerOptions: ProviderTransform.providerOptions(input.model, params.options),
-          activeTools: Object.keys(sortedTools).filter((x) => x !== "invalid"),
-          tools: sortedTools,
+          temperature: prepared.params.temperature,
+          topP: prepared.params.topP,
+          topK: prepared.params.topK,
+          providerOptions: ProviderTransform.providerOptions(input.model, prepared.params.options),
+          activeTools: Object.keys(prepared.tools).filter((x) => x !== "invalid"),
+          tools: prepared.tools,
           toolChoice: input.toolChoice,
-          maxOutputTokens: params.maxOutputTokens,
+          maxOutputTokens: prepared.params.maxOutputTokens,
           abortSignal: input.abort,
-          headers: requestHeaders,
+          headers: prepared.headers,
           maxRetries: input.retries ?? 0,
-          messages,
+          messages: prepared.messages,
           model: wrapLanguageModel({
             model: language,
             middleware: [
@@ -452,7 +315,11 @@ const live: Layer.Layer<
                 async transformParams(args) {
                   if (args.type === "stream") {
                     // @ts-expect-error
-                    args.params.prompt = ProviderTransform.message(args.params.prompt, input.model, options)
+                    args.params.prompt = ProviderTransform.message(
+                      args.params.prompt,
+                      input.model,
+                      prepared.messageTransformOptions,
+                    )
                   }
                   return args.params
                 },
@@ -517,24 +384,6 @@ export const defaultLayer = Layer.suspend(() =>
   ),
 )
 
-function resolveTools(input: Pick<StreamInput, "tools" | "agent" | "permission" | "user">) {
-  const disabled = Permission.disabled(
-    Object.keys(input.tools),
-    Permission.merge(input.agent.permission, input.permission ?? []),
-  )
-  return Record.filter(input.tools, (_, k) => input.user.tools?.[k] !== false && !disabled.has(k))
-}
-
-// Check if messages contain any tool-call content
-// Used to determine if a dummy tool should be added (GitHub Copilot only; see stream()).
-export function hasToolCalls(messages: ModelMessage[]): boolean {
-  for (const msg of messages) {
-    if (!Array.isArray(msg.content)) continue
-    for (const part of msg.content) {
-      if (part.type === "tool-call" || part.type === "tool-result") return true
-    }
-  }
-  return false
-}
+export const hasToolCalls = LLMRequestPrep.hasToolCalls
 
 export * as LLM from "./llm"

+ 208 - 0
packages/opencode/src/session/llm/request.ts

@@ -0,0 +1,208 @@
+import type { Auth } from "@/auth"
+import type { RuntimeFlags } from "@/effect/runtime-flags"
+import { InstanceState } from "@/effect/instance-state"
+import { Permission } from "@/permission"
+import type { Agent } from "@/agent/agent"
+import type { MessageV2 } from "../message-v2"
+import type { Provider } from "@/provider/provider"
+import { ProviderTransform } from "@/provider/transform"
+import { SystemPrompt } from "../system"
+import { InstallationVersion } from "@opencode-ai/core/installation/version"
+import { Effect, Record } from "effect"
+import { jsonSchema, tool as aiTool, type ModelMessage, type Tool } from "ai"
+import type { Plugin } from "@/plugin"
+import { mergeDeep } from "remeda"
+
+const USER_AGENT = `opencode/${InstallationVersion}`
+
+type PrepareInput = {
+  readonly user: MessageV2.User
+  readonly sessionID: string
+  readonly parentSessionID?: string
+  readonly model: Provider.Model
+  readonly agent: Agent.Info
+  readonly permission?: Permission.Ruleset
+  readonly system: string[]
+  readonly messages: ModelMessage[]
+  readonly small?: boolean
+  readonly tools: Record<string, Tool>
+  readonly provider: Provider.Info
+  readonly auth: Auth.Info | undefined
+  readonly plugin: Plugin.Interface
+  readonly flags: RuntimeFlags.Info
+  readonly isWorkflow: boolean
+}
+
+export type Prepared = {
+  readonly isOpenaiOauth: boolean
+  readonly system: string[]
+  readonly messages: ModelMessage[]
+  readonly tools: Record<string, Tool>
+  readonly params: {
+    readonly temperature?: number
+    readonly topP?: number
+    readonly topK?: number
+    readonly maxOutputTokens?: number
+    readonly options: Record<string, any>
+  }
+  readonly messageTransformOptions: Record<string, any>
+  readonly headers: Record<string, string>
+}
+
+const mergeOptions = (target: Record<string, any>, source: Record<string, any> | undefined): Record<string, any> =>
+  mergeDeep(target, source ?? {}) as Record<string, any>
+
+export const prepare = Effect.fn("LLMRequestPrep.prepare")(function* (input: PrepareInput) {
+  const isOpenaiOauth = input.provider.id === "openai" && input.auth?.type === "oauth"
+  const system = [
+    [
+      ...(input.agent.prompt ? [input.agent.prompt] : SystemPrompt.provider(input.model)),
+      ...input.system,
+      ...(input.user.system ? [input.user.system] : []),
+    ]
+      .filter((x) => x)
+      .join("\n"),
+  ]
+
+  const header = system[0]
+  yield* input.plugin.trigger(
+    "experimental.chat.system.transform",
+    { sessionID: input.sessionID, model: input.model },
+    { system },
+  )
+  if (system.length > 2 && system[0] === header) {
+    const rest = system.slice(1)
+    system.length = 0
+    system.push(header, rest.join("\n"))
+  }
+
+  const variant =
+    !input.small && input.model.variants && input.user.model.variant
+      ? input.model.variants[input.user.model.variant]
+      : {}
+  const base = input.small
+    ? ProviderTransform.smallOptions(input.model)
+    : ProviderTransform.options({
+        model: input.model,
+        sessionID: input.sessionID,
+        providerOptions: input.provider.options,
+      })
+  const options = mergeOptions(mergeOptions(mergeOptions(base, input.model.options), input.agent.options), variant)
+  if (isOpenaiOauth) options.instructions = system.join("\n")
+
+  const messages =
+    isOpenaiOauth || input.isWorkflow
+      ? input.messages
+      : [
+          ...system.map(
+            (x): ModelMessage => ({
+              role: "system",
+              content: x,
+            }),
+          ),
+          ...input.messages,
+        ]
+
+  const params = yield* input.plugin.trigger(
+    "chat.params",
+    {
+      sessionID: input.sessionID,
+      agent: input.agent.name,
+      model: input.model,
+      provider: input.provider,
+      message: input.user,
+    },
+    {
+      temperature: input.model.capabilities.temperature
+        ? (input.agent.temperature ?? ProviderTransform.temperature(input.model))
+        : undefined,
+      topP: input.agent.topP ?? ProviderTransform.topP(input.model),
+      topK: ProviderTransform.topK(input.model),
+      maxOutputTokens: ProviderTransform.maxOutputTokens(input.model, input.flags.outputTokenMax),
+      options,
+    },
+  )
+
+  const { headers } = yield* input.plugin.trigger(
+    "chat.headers",
+    {
+      sessionID: input.sessionID,
+      agent: input.agent.name,
+      model: input.model,
+      provider: input.provider,
+      message: input.user,
+    },
+    {
+      headers: {},
+    },
+  )
+
+  const tools = resolveTools(input)
+  if (
+    input.model.providerID.includes("github-copilot") &&
+    Object.keys(tools).length === 0 &&
+    hasToolCalls(input.messages)
+  ) {
+    // Copilot needs a tools field when replaying prior tool calls, even if no tools are currently enabled.
+    tools["_noop"] = aiTool({
+      description: "Do not call this tool. It exists only for API compatibility and must never be invoked.",
+      inputSchema: jsonSchema({
+        type: "object",
+        properties: {
+          reason: { type: "string", description: "Unused" },
+        },
+      }),
+      execute: async () => ({ output: "", title: "", metadata: {} }),
+    })
+  }
+
+  const opencodeProjectID = input.model.providerID.startsWith("opencode")
+    ? (yield* InstanceState.context).project.id
+    : undefined
+
+  return {
+    isOpenaiOauth,
+    system,
+    messages,
+    tools: Object.fromEntries(Object.entries(tools).toSorted(([a], [b]) => a.localeCompare(b))),
+    params,
+    messageTransformOptions: options,
+    headers: {
+      ...(input.model.providerID.startsWith("opencode")
+        ? {
+            ...(opencodeProjectID ? { "x-opencode-project": opencodeProjectID } : {}),
+            "x-opencode-session": input.sessionID,
+            "x-opencode-request": input.user.id,
+            "x-opencode-client": input.flags.client,
+            "User-Agent": USER_AGENT,
+          }
+        : {
+            "x-session-affinity": input.sessionID,
+            ...(input.parentSessionID ? { "x-parent-session-id": input.parentSessionID } : {}),
+            "User-Agent": USER_AGENT,
+          }),
+      ...input.model.headers,
+      ...headers,
+    },
+  }
+})
+
+function resolveTools(input: Pick<PrepareInput, "tools" | "agent" | "permission" | "user">) {
+  const disabled = Permission.disabled(
+    Object.keys(input.tools),
+    Permission.merge(input.agent.permission, input.permission ?? []),
+  )
+  return Record.filter(input.tools, (_, k) => input.user.tools?.[k] !== false && !disabled.has(k))
+}
+
+export function hasToolCalls(messages: ModelMessage[]): boolean {
+  for (const msg of messages) {
+    if (!Array.isArray(msg.content)) continue
+    for (const part of msg.content) {
+      if (part.type === "tool-call" || part.type === "tool-result") return true
+    }
+  }
+  return false
+}
+
+export * as LLMRequestPrep from "./request"