125 lines
3.4 KiB
TypeScript

import type { PluginInput } from "@opencode-ai/plugin"
import { getMessageIdsFromSDK } from "./message-storage-directory"
import { TRUNCATION_MESSAGE } from "./storage-paths"
import type { ToolResultInfo } from "./tool-part-types"
import { patchPart } from "../../shared/opencode-http-api"
import { log } from "../../shared/logger"
import { normalizeSDKResponse } from "../../shared"
type OpencodeClient = PluginInput["client"]
interface SDKToolPart {
id: string
type: string
callID?: string
tool?: string
state?: {
status?: string
input?: Record<string, unknown>
output?: string
error?: string
time?: { start?: number; end?: number; compacted?: number }
}
}
interface SDKMessage {
info?: { id?: string }
parts?: SDKToolPart[]
}
export async function findToolResultsBySizeFromSDK(
client: OpencodeClient,
sessionID: string
): Promise<ToolResultInfo[]> {
try {
const response = await client.session.messages({ path: { id: sessionID } })
const messages = normalizeSDKResponse(response, [] as SDKMessage[], { preferResponseOnMissingData: true })
const results: ToolResultInfo[] = []
for (const msg of messages) {
const messageID = msg.info?.id
if (!messageID || !msg.parts) continue
for (const part of msg.parts) {
if (part.type === "tool" && part.state?.output && !part.state?.time?.compacted && part.tool) {
results.push({
partPath: "",
partId: part.id,
messageID,
toolName: part.tool,
outputSize: part.state.output.length,
})
}
}
}
return results.sort((a, b) => b.outputSize - a.outputSize)
} catch {
return []
}
}
export async function truncateToolResultAsync(
client: OpencodeClient,
sessionID: string,
messageID: string,
partId: string,
part: SDKToolPart
): Promise<{ success: boolean; toolName?: string; originalSize?: number }> {
if (!part.state?.output) return { success: false }
const originalSize = part.state.output.length
const toolName = part.tool
const updatedPart: Record<string, unknown> = {
...part,
state: {
...part.state,
output: TRUNCATION_MESSAGE,
time: {
...(part.state.time ?? { start: Date.now() }),
compacted: Date.now(),
},
},
}
try {
const patched = await patchPart(client, sessionID, messageID, partId, updatedPart)
if (!patched) return { success: false }
return { success: true, toolName, originalSize }
} catch (error) {
log("[context-window-recovery] truncateToolResultAsync failed", { error: String(error) })
return { success: false }
}
}
export async function countTruncatedResultsFromSDK(
client: OpencodeClient,
sessionID: string
): Promise<number> {
try {
const response = await client.session.messages({ path: { id: sessionID } })
const messages = normalizeSDKResponse(response, [] as SDKMessage[], { preferResponseOnMissingData: true })
let count = 0
for (const msg of messages) {
if (!msg.parts) continue
for (const part of msg.parts) {
if (part.type === "tool" && part.state?.time?.compacted) count++
}
}
return count
} catch {
return 0
}
}
export async function getTotalToolOutputSizeFromSDK(
client: OpencodeClient,
sessionID: string
): Promise<number> {
const results = await findToolResultsBySizeFromSDK(client, sessionID)
return results.reduce((sum, result) => sum + result.outputSize, 0)
}