fix: preserve plugin tool state.attachments in UI and materializer (#2367)

* preserve attachments in UI and materializer

* add support for non-image attachments

* preserve attachments in UI and materializer

* improve attachment rendering and state synchronization

* simplify attachment filtering and key handling

* prevent rendering tool attachments without URLs

* require `f.url` in `imageAttachments` to match gallery indices
This commit is contained in:
FrostiDrinks
2026-07-22 10:46:40 +03:00
committed by GitHub
parent 733bd37c37
commit c29b2d0849
13 changed files with 283 additions and 5 deletions
@@ -165,6 +165,154 @@ describe("materializeSessionSnapshots", () => {
expect(mergedPart.state?.time?.start).toBe(1000)
expect(mergedPart.state?.time?.end).toBe(2000)
})
test("preserves state.attachments from existing part when completed snapshot lacks them", () => {
const livePart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: {
status: "completed",
output: "done",
time: { start: 100, end: 200 },
attachments: [{ id: "att-1", type: "file", mime: "image/png", url: "data:image/png,..." }],
},
} as unknown as Part
const snapshotPart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: { status: "completed", output: "done", time: { start: 100, end: 200 } },
} as unknown as Part
const state = {
message: { ses_1: [message("msg_1")] },
part: { msg_1: [livePart] },
}
const result = materializeSessionSnapshots(
state,
"ses_1",
[{ info: message("msg_1"), parts: [snapshotPart] }],
)
const mergedPart = result.part.msg_1[0] as { state?: { attachments?: Array<unknown> } }
expect(mergedPart.state?.attachments).toHaveLength(1)
expect((mergedPart.state?.attachments?.[0] as { id?: string })?.id).toBe("att-1")
})
test("preserves state.attachments during streaming merge when snapshot has no end time", () => {
const livePart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: {
status: "running",
time: { start: 100 },
attachments: [{ id: "att-1", type: "file", mime: "image/png", url: "data:image/png,..." }],
},
} as unknown as Part
const snapshotPart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: { status: "running", time: { start: 100 } },
} as unknown as Part
const state = {
message: { ses_1: [message("msg_1")] },
part: { msg_1: [livePart] },
}
const result = materializeSessionSnapshots(
state,
"ses_1",
[{ info: message("msg_1"), parts: [snapshotPart] }],
)
const mergedPart = result.part.msg_1[0] as { state?: { attachments?: Array<unknown> } }
expect(mergedPart.state?.attachments).toHaveLength(1)
expect((mergedPart.state?.attachments?.[0] as { id?: string })?.id).toBe("att-1")
})
test("preserves both state.attachments and state.time.start during streaming merge when snapshot lacks both", () => {
const livePart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: {
status: "running",
time: { start: 100 },
attachments: [{ id: "att-1", type: "file", mime: "image/png", url: "data:image/png,..." }],
},
} as unknown as Part
const snapshotPart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: { status: "running" },
} as unknown as Part
const state = {
message: { ses_1: [message("msg_1")] },
part: { msg_1: [livePart] },
}
const result = materializeSessionSnapshots(
state,
"ses_1",
[{ info: message("msg_1"), parts: [snapshotPart] }],
)
const mergedPart = result.part.msg_1[0] as { state?: { attachments?: Array<unknown>; time?: { start?: number; end?: number } } }
expect(mergedPart.state?.attachments).toHaveLength(1)
expect((mergedPart.state?.attachments?.[0] as { id?: string })?.id).toBe("att-1")
expect(mergedPart.state?.time?.start).toBe(100)
})
test("does not merge existing state.attachments when snapshot has its own", () => {
const livePart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: {
status: "completed",
output: "done",
time: { start: 100, end: 200 },
attachments: [{ id: "att-old", type: "file", mime: "image/png", url: "data:image/png,..." }],
},
} as unknown as Part
const snapshotPart = {
id: "prt_1",
messageID: "msg_1",
sessionID: "ses_1",
type: "tool",
state: {
status: "completed",
output: "done",
time: { start: 100, end: 200 },
attachments: [{ id: "att-new", type: "file", mime: "image/jpeg", url: "data:image/jpeg,..." }],
},
} as unknown as Part
const state = {
message: { ses_1: [message("msg_1")] },
part: { msg_1: [livePart] },
}
const result = materializeSessionSnapshots(
state,
"ses_1",
[{ info: message("msg_1"), parts: [snapshotPart] }],
)
const mergedPart = result.part.msg_1[0] as { state?: { attachments?: Array<unknown> } }
expect(mergedPart.state?.attachments).toHaveLength(1)
expect((mergedPart.state?.attachments?.[0] as { id?: string })?.id).toBe("att-new")
})
})
describe("getSessionMaterializationStatus", () => {
+30 -3
View File
@@ -108,6 +108,13 @@ function getStringField(part: Part, field: "text" | "output"): string | undefine
return typeof value === "string" ? value : undefined
}
function getPartStateAttachments(part: Part): Array<unknown> | undefined {
const state = (part as Record<string, unknown>).state as Record<string, unknown> | undefined
if (!state) return undefined
const attachments = state.attachments
return Array.isArray(attachments) && attachments.length > 0 ? attachments : undefined
}
function hasLiveStreamingField(part: Part): boolean {
if (getPartEndTime(part) !== undefined) return false
return STREAMING_PART_FIELDS.some((field) => {
@@ -126,7 +133,18 @@ function getPartStateTime(part: Part): { start?: number; end?: number } | undefi
}
function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
if (!existing || getPartEndTime(next) !== undefined) return next
if (!existing) return next
if (getPartEndTime(next) !== undefined) {
const existingAttachments = getPartStateAttachments(existing)
if (existingAttachments && !getPartStateAttachments(next)) {
const nextRecord = { ...next }
const nextState = { ...((next as Record<string, unknown>).state as Record<string, unknown> ?? {}), attachments: existingAttachments }
;(nextRecord as Record<string, unknown>).state = nextState
return nextRecord
}
return next
}
let merged: Part = next
for (const field of STREAMING_PART_FIELDS) {
@@ -142,6 +160,15 @@ function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
mergedRecord[field] = existingValue
}
const existingAttachments = getPartStateAttachments(existing)
if (existingAttachments && !getPartStateAttachments(next)) {
if (merged === next) merged = { ...next }
const mergedRecord = merged as Record<string, unknown>
const nextState = (next as Record<string, unknown>).state as Record<string, unknown> | undefined
const newState = { ...(nextState ?? {}), attachments: existingAttachments }
mergedRecord.state = newState
}
const existingTime = getPartStateTime(existing)
if (existingTime) {
const nextTime = getPartStateTime(next)
@@ -150,8 +177,8 @@ function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
if (preservedStart !== nextTime?.start || preservedEnd !== nextTime?.end) {
if (merged === next) merged = { ...next }
const mergedRecord = merged as Record<string, unknown>
const nextState = (next as Record<string, unknown>).state as Record<string, unknown> | undefined
const newState = { ...(nextState ?? {}), time: { start: preservedStart, end: preservedEnd } }
const currentState = (mergedRecord.state as Record<string, unknown> | undefined) ?? (next as Record<string, unknown>).state as Record<string, unknown> | undefined
const newState = { ...(currentState ?? {}), time: { start: preservedStart, end: preservedEnd } }
mergedRecord.state = newState
}
}