Session management and prompt handling improvements (#2577)

Co-authored-by: GitHub Action <action@github.com>
This commit is contained in:
Dax
2025-09-13 05:46:14 -04:00
committed by GitHub
parent 535230dce4
commit 9bb25a9260
9 changed files with 1755 additions and 1682 deletions

View File

@@ -0,0 +1,120 @@
import { generateText, type ModelMessage } from "ai"
import { Session } from "."
import { Identifier } from "../id/id"
import { Instance } from "../project/instance"
import { Provider } from "../provider/provider"
import { defer } from "../util/defer"
import { MessageV2 } from "./message-v2"
import { SystemPrompt } from "./system"
import { Bus } from "../bus"
import z from "zod"
import type { ModelsDev } from "../provider/models"
import { SessionPrompt } from "./prompt"
export namespace SessionCompaction {
export const Event = {
Compacted: Bus.event(
"session.compacted",
z.object({
sessionID: z.string(),
}),
),
}
export function isOverflow(input: { tokens: MessageV2.Assistant["tokens"]; model: ModelsDev.Model }) {
const count = input.tokens.input + input.tokens.cache.read + input.tokens.output
const output = Math.min(input.model.limit.output, SessionPrompt.OUTPUT_TOKEN_MAX) || SessionPrompt.OUTPUT_TOKEN_MAX
const usable = input.model.limit.context - output
return count > usable / 2
}
export async function run(input: { sessionID: string; providerID: string; modelID: string }) {
await Session.update(input.sessionID, (draft) => {
draft.time.compacting = Date.now()
})
await using _ = defer(async () => {
await Session.update(input.sessionID, (draft) => {
draft.time.compacting = undefined
})
})
const toSummarize = await Session.messages(input.sessionID).then(MessageV2.filterSummarized)
const model = await Provider.getModel(input.providerID, input.modelID)
const system = [
...SystemPrompt.summarize(model.providerID),
...(await SystemPrompt.environment()),
...(await SystemPrompt.custom()),
]
const msg = (await Session.updateMessage({
id: Identifier.ascending("message"),
role: "assistant",
sessionID: input.sessionID,
system,
mode: "build",
path: {
cwd: Instance.directory,
root: Instance.worktree,
},
cost: 0,
tokens: {
output: 0,
input: 0,
reasoning: 0,
cache: { read: 0, write: 0 },
},
modelID: input.modelID,
providerID: model.providerID,
time: {
created: Date.now(),
},
})) as MessageV2.Assistant
const generated = await generateText({
maxRetries: 10,
model: model.language,
messages: [
...system.map(
(x): ModelMessage => ({
role: "system",
content: x,
}),
),
...MessageV2.toModelMessage(toSummarize),
{
role: "user",
content: [
{
type: "text",
text: "Provide a detailed but concise summary of our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next.",
},
],
},
],
})
const usage = Session.getUsage(model.info, generated.usage, generated.providerMetadata)
msg.cost += usage.cost
msg.tokens = usage.tokens
msg.summary = true
msg.time.completed = Date.now()
await Session.updateMessage(msg)
const part = await Session.updatePart({
type: "text",
sessionID: input.sessionID,
messageID: msg.id,
id: Identifier.ascending("part"),
text: generated.text,
time: {
start: Date.now(),
end: Date.now(),
},
})
Bus.publish(Event.Compacted, {
sessionID: input.sessionID,
})
return {
info: msg,
parts: [part],
}
}
}