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:
@@ -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", () => {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user