forked from anomalyco/opencode
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcompaction.ts
More file actions
328 lines (300 loc) · 10.9 KB
/
compaction.ts
File metadata and controls
328 lines (300 loc) · 10.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
import { BusEvent } from "@/bus/bus-event"
import { Bus } from "@/bus"
import { Session } from "."
import { Identifier } from "../id/id"
import { Instance } from "../project/instance"
import { Provider } from "../provider/provider"
import { MessageV2 } from "./message-v2"
import z from "zod"
import { Token } from "../util/token"
import { Log } from "../util/log"
import { SessionProcessor } from "./processor"
import { fn } from "@/util/fn"
import { Agent } from "@/agent/agent"
import { Plugin } from "@/plugin"
import { Config } from "@/config/config"
import { ProviderTransform } from "@/provider/transform"
export namespace SessionCompaction {
const log = Log.create({ service: "session.compaction" })
export const Event = {
Compacted: BusEvent.define(
"session.compacted",
z.object({
sessionID: z.string(),
}),
),
}
const COMPACTION_BUFFER = 20_000
export async function isOverflow(input: { tokens: MessageV2.Assistant["tokens"]; model: Provider.Model }) {
const config = await Config.get()
if (config.compaction?.auto === false) return false
const context = input.model.limit.context
if (context === 0) return false
const count =
input.tokens.total ||
input.tokens.input + input.tokens.output + input.tokens.cache.read + input.tokens.cache.write
const reserved =
config.compaction?.reserved ?? Math.min(COMPACTION_BUFFER, ProviderTransform.maxOutputTokens(input.model))
const usable = input.model.limit.input
? input.model.limit.input - reserved
: context - ProviderTransform.maxOutputTokens(input.model)
return count >= usable
}
export const PRUNE_MINIMUM = 20_000
export const PRUNE_PROTECT = 40_000
const PRUNE_PROTECTED_TOOLS = ["skill"]
// goes backwards through parts until there are 40_000 tokens worth of tool
// calls. then erases output of previous tool calls. idea is to throw away old
// tool calls that are no longer relevant.
export async function prune(input: { sessionID: string }) {
const config = await Config.get()
if (config.compaction?.prune === false) return
log.info("pruning")
const msgs = await Session.messages({ sessionID: input.sessionID })
let total = 0
let pruned = 0
const toPrune = []
let turns = 0
loop: for (let msgIndex = msgs.length - 1; msgIndex >= 0; msgIndex--) {
const msg = msgs[msgIndex]
if (msg.info.role === "user") turns++
if (turns < 2) continue
if (msg.info.role === "assistant" && msg.info.summary) break loop
for (let partIndex = msg.parts.length - 1; partIndex >= 0; partIndex--) {
const part = msg.parts[partIndex]
if (part.type === "tool")
if (part.state.status === "completed") {
if (PRUNE_PROTECTED_TOOLS.includes(part.tool)) continue
if (part.state.time.compacted) break loop
const estimate = Token.estimate(part.state.output)
total += estimate
if (total > PRUNE_PROTECT) {
pruned += estimate
toPrune.push(part)
}
}
}
}
log.info("found", { pruned, total })
if (pruned > PRUNE_MINIMUM) {
for (const part of toPrune) {
if (part.state.status === "completed") {
part.state.time.compacted = Date.now()
await Session.updatePart(part)
}
}
log.info("pruned", { count: toPrune.length })
}
}
export async function process(input: {
parentID: string
messages: MessageV2.WithParts[]
sessionID: string
abort: AbortSignal
auto: boolean
overflow?: boolean
}) {
const userMessage = input.messages.findLast((m) => m.info.id === input.parentID)!.info as MessageV2.User
let messages = input.messages
let replay: MessageV2.WithParts | undefined
if (input.overflow) {
const idx = input.messages.findIndex((m) => m.info.id === input.parentID)
for (let i = idx - 1; i >= 0; i--) {
const msg = input.messages[i]
if (msg.info.role === "user" && !msg.parts.some((p) => p.type === "compaction")) {
replay = msg
messages = input.messages.slice(0, i)
break
}
}
const hasContent =
replay && messages.some((m) => m.info.role === "user" && !m.parts.some((p) => p.type === "compaction"))
if (!hasContent) {
replay = undefined
messages = input.messages
}
}
const agent = await Agent.get("compaction")
const model = agent.model
? await Provider.getModel(agent.model.providerID, agent.model.modelID)
: await Provider.getModel(userMessage.model.providerID, userMessage.model.modelID)
const msg = (await Session.updateMessage({
id: Identifier.ascending("message"),
role: "assistant",
parentID: input.parentID,
sessionID: input.sessionID,
mode: "compaction",
agent: "compaction",
variant: userMessage.variant,
summary: true,
path: {
cwd: Instance.directory,
root: Instance.worktree,
},
cost: 0,
tokens: {
output: 0,
input: 0,
reasoning: 0,
cache: { read: 0, write: 0 },
},
modelID: model.id,
providerID: model.providerID,
time: {
created: Date.now(),
},
})) as MessageV2.Assistant
const processor = SessionProcessor.create({
assistantMessage: msg,
sessionID: input.sessionID,
model,
abort: input.abort,
})
// Allow plugins to inject context or replace compaction prompt
const compacting = await Plugin.trigger(
"experimental.session.compacting",
{ sessionID: input.sessionID },
{ context: [], prompt: undefined },
)
const defaultPrompt = `Provide a detailed prompt for continuing 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.
The summary that you construct will be used so that another agent can read it and continue the work.
When constructing the summary, try to stick to this template:
---
## Goal
[What goal(s) is the user trying to accomplish?]
## Instructions
- [What important instructions did the user give you that are relevant]
- [If there is a plan or spec, include information about it so next agent can continue using it]
## Discoveries
[What notable things were learned during this conversation that would be useful for the next agent to know when continuing the work]
## Accomplished
[What work has been completed, what work is still in progress, and what work is left?]
## Relevant files / directories
[Construct a structured list of relevant files that have been read, edited, or created that pertain to the task at hand. If all the files in a directory are relevant, include the path to the directory.]
---`
const promptText = compacting.prompt ?? [defaultPrompt, ...compacting.context].join("\n\n")
const result = await processor.process({
user: userMessage,
agent,
abort: input.abort,
sessionID: input.sessionID,
tools: {},
system: [],
messages: [
...MessageV2.toModelMessages(messages, model, { stripMedia: true }),
{
role: "user",
content: [
{
type: "text",
text: promptText,
},
],
},
],
model,
})
if (result === "compact") {
processor.message.error = new MessageV2.ContextOverflowError({
message: replay
? "Conversation history too large to compact - exceeds model context limit"
: "Session too large to compact - context exceeds model limit even after stripping media",
}).toObject()
processor.message.finish = "error"
await Session.updateMessage(processor.message)
return "stop"
}
if (result === "continue" && input.auto) {
if (replay) {
const original = replay.info as MessageV2.User
const replayMsg = await Session.updateMessage({
id: Identifier.ascending("message"),
role: "user",
sessionID: input.sessionID,
time: { created: Date.now() },
agent: original.agent,
model: original.model,
format: original.format,
tools: original.tools,
system: original.system,
variant: original.variant,
})
for (const part of replay.parts) {
if (part.type === "compaction") continue
const replayPart =
part.type === "file" && MessageV2.isMedia(part.mime)
? { type: "text" as const, text: `[Attached ${part.mime}: ${part.filename ?? "file"}]` }
: part
await Session.updatePart({
...replayPart,
id: Identifier.ascending("part"),
messageID: replayMsg.id,
sessionID: input.sessionID,
})
}
} else {
const continueMsg = await Session.updateMessage({
id: Identifier.ascending("message"),
role: "user",
sessionID: input.sessionID,
time: { created: Date.now() },
agent: userMessage.agent,
model: userMessage.model,
})
const text =
(input.overflow
? "The previous request exceeded the provider's size limit due to large media attachments. The conversation was compacted and media files were removed from context. If the user was asking about attached images or files, explain that the attachments were too large to process and suggest they try again with smaller or fewer files.\n\n"
: "") +
"Continue if you have next steps, or stop and ask for clarification if you are unsure how to proceed."
await Session.updatePart({
id: Identifier.ascending("part"),
messageID: continueMsg.id,
sessionID: input.sessionID,
type: "text",
synthetic: true,
text,
time: {
start: Date.now(),
end: Date.now(),
},
})
}
}
if (processor.message.error) return "stop"
Bus.publish(Event.Compacted, { sessionID: input.sessionID })
return "continue"
}
export const create = fn(
z.object({
sessionID: Identifier.schema("session"),
agent: z.string(),
model: z.object({
providerID: z.string(),
modelID: z.string(),
}),
auto: z.boolean(),
overflow: z.boolean().optional(),
}),
async (input) => {
const msg = await Session.updateMessage({
id: Identifier.ascending("message"),
role: "user",
model: input.model,
sessionID: input.sessionID,
agent: input.agent,
time: {
created: Date.now(),
},
})
await Session.updatePart({
id: Identifier.ascending("part"),
messageID: msg.id,
sessionID: msg.sessionID,
type: "compaction",
auto: input.auto,
overflow: input.overflow,
})
},
)
}