feat: AI SDK v6 support (#18433)

This commit is contained in:
Aiden Cline
2026-03-27 15:24:30 -05:00
committed by GitHub
parent 7a7643c86a
commit c33d9996f0
36 changed files with 1290 additions and 1155 deletions

View File

@@ -215,7 +215,7 @@ When constructing the summary, try to stick to this template:
tools: {},
system: [],
messages: [
...MessageV2.toModelMessages(msgs, model, { stripMedia: true }),
...(await MessageV2.toModelMessages(msgs, model, { stripMedia: true })),
{
role: "user",
content: [

View File

@@ -1,16 +1,6 @@
import { Installation } from "@/installation"
import { Provider } from "@/provider/provider"
import { Log } from "@/util/log"
import {
streamText,
wrapLanguageModel,
type ModelMessage,
type StreamTextResult,
type Tool,
type ToolSet,
tool,
jsonSchema,
} from "ai"
import { streamText, wrapLanguageModel, type ModelMessage, type Tool, tool, jsonSchema } from "ai"
import { mergeDeep, pipe } from "remeda"
import { GitLabWorkflowLanguageModel } from "gitlab-ai-provider"
import { ProviderTransform } from "@/provider/transform"
@@ -23,6 +13,7 @@ import { SystemPrompt } from "./system"
import { Flag } from "@/flag/flag"
import { Permission } from "@/permission"
import { Auth } from "@/auth"
import { Installation } from "@/installation"
export namespace LLM {
const log = Log.create({ service: "llm" })
@@ -43,8 +34,6 @@ export namespace LLM {
toolChoice?: "auto" | "required" | "none"
}
export type StreamOutput = StreamTextResult<ToolSet, unknown>
export async function stream(input: StreamInput) {
const l = log
.clone()
@@ -273,8 +262,10 @@ export namespace LLM {
model: language,
middleware: [
{
specificationVersion: "v3" as const,
async transformParams(args) {
if (args.type === "stream") {
// TODO: verify that LanguageModelV3Prompt is still compat here!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
// @ts-expect-error
args.params.prompt = ProviderTransform.message(args.params.prompt, input.model, options)
}

View File

@@ -573,11 +573,11 @@ export namespace MessageV2 {
}))
}
export function toModelMessages(
export async function toModelMessages(
input: WithParts[],
model: Provider.Model,
options?: { stripMedia?: boolean },
): ModelMessage[] {
): Promise<ModelMessage[]> {
const result: UIMessage[] = []
const toolNames = new Set<string>()
// Track media from tool results that need to be injected as user messages
@@ -601,7 +601,8 @@ export namespace MessageV2 {
return false
})()
const toModelOutput = (output: unknown) => {
const toModelOutput = (options: { toolCallId: string; input: unknown; output: unknown }) => {
const output = options.output
if (typeof output === "string") {
return { type: "text", value: output }
}
@@ -799,7 +800,7 @@ export namespace MessageV2 {
const tools = Object.fromEntries(Array.from(toolNames).map((toolName) => [toolName, { toModelOutput }]))
return convertToModelMessages(
return await convertToModelMessages(
result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
{
//@ts-expect-error (convertToModelMessages expects a ToolSet but only actually needs tools[name]?.toModelOutput)
@@ -871,7 +872,13 @@ export namespace MessageV2 {
db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).orderBy(PartTable.id).all(),
)
return rows.map(
(row) => ({ ...row.data, id: row.id, sessionID: row.session_id, messageID: row.message_id }) as MessageV2.Part,
(row) =>
({
...row.data,
id: row.id,
sessionID: row.session_id,
messageID: row.message_id,
}) as MessageV2.Part,
)
})

View File

@@ -11,7 +11,7 @@ import { Session } from "."
import { Agent } from "../agent/agent"
import { Provider } from "../provider/provider"
import { ModelID, ProviderID } from "../provider/schema"
import { type Tool as AITool, tool, jsonSchema, type ToolCallOptions, asSchema } from "ai"
import { type Tool as AITool, tool, jsonSchema, type ToolExecutionOptions, asSchema } from "ai"
import { SessionCompaction } from "./compaction"
import { Instance } from "../project/instance"
import { Bus } from "../bus"
@@ -321,7 +321,13 @@ export namespace SessionPrompt {
if (!lastUser) throw new Error("No user message found in stream. This should never happen.")
if (
lastAssistant?.finish &&
!["tool-calls", "unknown"].includes(lastAssistant.finish) &&
![
"tool-calls",
// in v6 unknown became other but other existed in v5 too and was distinctly different
// I think there are certain providers that used to have bad stop reasons, not rlly sure which
// ones if any still have this?
// "unknown",
].includes(lastAssistant.finish) &&
lastUser.id < lastAssistant.id
) {
log.info("exiting loop", { sessionID })
@@ -692,7 +698,7 @@ export namespace SessionPrompt {
sessionID,
system,
messages: [
...MessageV2.toModelMessages(msgs, model),
...(await MessageV2.toModelMessages(msgs, model)),
...(isLastStep
? [
{
@@ -775,7 +781,7 @@ export namespace SessionPrompt {
using _ = log.time("resolveTools")
const tools: Record<string, AITool> = {}
const context = (args: any, options: ToolCallOptions): Tool.Context => ({
const context = (args: any, options: ToolExecutionOptions): Tool.Context => ({
sessionID: input.session.id,
abort: options.abortSignal!,
messageID: input.processor.message.id,
@@ -861,7 +867,8 @@ export namespace SessionPrompt {
const execute = item.execute
if (!execute) continue
const transformed = ProviderTransform.schema(input.model, asSchema(item.inputSchema).jsonSchema)
const schema = await asSchema(item.inputSchema).jsonSchema
const transformed = ProviderTransform.schema(input.model, schema)
item.inputSchema = jsonSchema(transformed)
// Wrap execute to add plugin hooks and format output
item.execute = async (args, opts) => {
@@ -974,10 +981,10 @@ export namespace SessionPrompt {
metadata: { valid: true },
}
},
toModelOutput(result) {
toModelOutput({ output }) {
return {
type: "text",
value: result.output,
value: output.output,
}
},
})
@@ -2010,28 +2017,28 @@ NOTE: At any point in time through this workflow you should feel free to ask the
(await Provider.getSmallModel(input.providerID)) ?? (await Provider.getModel(input.providerID, input.modelID))
)
})
const result = await LLM.stream({
agent,
user: firstRealUser.info as MessageV2.User,
system: [],
small: true,
tools: {},
model,
abort: new AbortController().signal,
sessionID: input.session.id,
retries: 2,
messages: [
{
role: "user",
content: "Generate a title for this conversation:\n",
},
...(hasOnlySubtaskParts
? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }]
: MessageV2.toModelMessages(contextMessages, model)),
],
})
const text = await result.text.catch((err) => log.error("failed to generate title", { error: err }))
if (text) {
try {
const result = await LLM.stream({
agent,
user: firstRealUser.info as MessageV2.User,
system: [],
small: true,
tools: {},
model,
abort: new AbortController().signal,
sessionID: input.session.id,
retries: 2,
messages: [
{
role: "user",
content: "Generate a title for this conversation:\n",
},
...(hasOnlySubtaskParts
? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }]
: await MessageV2.toModelMessages(contextMessages, model)),
],
})
const text = await result.text
const cleaned = text
.replace(/<think>[\s\S]*?<\/think>\s*/g, "")
.split("\n")
@@ -2044,6 +2051,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the
if (NotFoundError.isInstance(err)) return
throw err
})
} catch (error) {
log.error("failed to generate title", { error })
}
}
}