Skip to content

Commit 7cc3533

Browse files
authored
fix(trace): replace images pasted while the agent is busy (#3575)
Signed-off-by: Miguel Martinez Trivino <miguel@chainloop.dev>
1 parent 7474d3b commit 7cc3533

4 files changed

Lines changed: 156 additions & 33 deletions

File tree

‎app/cli/internal/trace/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ What the session was asked to build, and which skills it used.
6767
| Local spec sources read at push | A capture whose source is a local file holds only a header and a placeholder. Each push reads the file and records its current content, so the evidence follows the edits of the file. Only a regular text file of up to 1 MiB is read. Otherwise the push keeps the body that the agent wrote, or drops the capture with a warning when there is no body (spec issue-3561). | Yes | Yes | Yes | Yes |
6868
| Skills tracking | The evidence lists the skills that the session used, with a copy of each skill as it ran. | Yes | No | Partial | Partial |
6969
| Spec source pointers | Exact copies of spec sources in the transcript (full file reads and file writes) are replaced with a pointer to the spec material, so the evidence does not hold the same content two or three times (spec issue-3556). | Yes | No | No | No |
70-
| Pasted image pointers | Each distinct image that the user pasted into the session is stored once as an image spec entry, from the data in the transcript, with the title "Pasted image N" and no role. A repeated paste of the same image adds no entry. The inline image in the transcript is replaced with a pointer to that material. An image whose data cannot be decoded, or that the push could not store, stays inline, and the push continues. When the agent captured the same image, its entry is kept and no second entry is added (spec issue-3569). | Yes | No | No | No |
70+
| Pasted image pointers | Each distinct image that the user pasted into the session, including images in prompts sent while the agent was busy, is stored once as an image spec entry, from the data in the transcript, with the title "Pasted image N" and no role. A repeated paste of the same image adds no entry. The inline image in the transcript is replaced with a pointer to that material. An image whose data cannot be decoded, or that the push could not store, stays inline, and the push continues. When the agent captured the same image, its entry is kept and no second entry is added (spec issue-3569). | Yes | No | No | No |
7171

7272
Why not Yes:
7373

‎app/cli/internal/trace/claude/pastedimages.go‎

Lines changed: 90 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ var _ trace.PastedImageFinder = (*Provider)(nil)
3131
const (
3232
blockImage = "image"
3333
sourceBase64 = "base64"
34+
// recordTypeAttachment is a record that Claude Code adds to the
35+
// conversation, such as a prompt that the user queued.
36+
recordTypeAttachment = "attachment"
37+
attachmentQueuedCommand = "queued_command"
38+
originHuman = "human"
3439
)
3540

3641
// PastedImages returns the images that the user pasted into the session
@@ -43,8 +48,11 @@ func (p *Provider) PastedImages(raw map[string][]json.RawMessage) []trace.Pasted
4348
position := 0
4449
for _, key := range streamOrder(raw) {
4550
for _, l := range raw[key] {
46-
e, blocks := pastedImages(l)
47-
for _, b := range blocks {
51+
pl := pastedImages(l)
52+
if pl == nil {
53+
continue
54+
}
55+
for _, b := range pl.images {
4856
position++
4957
data, err := b.decode()
5058
if err != nil {
@@ -65,7 +73,7 @@ func (p *Provider) PastedImages(raw map[string][]json.RawMessage) []trace.Pasted
6573
Digest: digest,
6674
MediaType: b.block.object("source").str("media_type"),
6775
Number: number,
68-
Timestamp: normalizeTimestamp(e.fields.str("timestamp")),
76+
Timestamp: normalizeTimestamp(pl.fields.str("timestamp")),
6977
})
7078
}
7179
}
@@ -81,13 +89,14 @@ func (p *Provider) PastedImages(raw map[string][]json.RawMessage) []trace.Pasted
8189
func replacePastedImages(raw map[string][]json.RawMessage, sources pointer.Sources, report pointer.Report) {
8290
for _, lines := range raw {
8391
for i, l := range lines {
84-
e, blocks := pastedImages(l)
85-
if len(blocks) == 0 {
92+
pl := pastedImages(l)
93+
if pl == nil {
8694
continue
8795
}
8896

8997
res := report.Finder(pointer.FinderPastedImage)
90-
for _, b := range blocks {
98+
changed := false
99+
for _, b := range pl.images {
91100
data, err := b.decode()
92101
if err != nil {
93102
res.Skip(pointer.ReasonInvalidImage, 1)
@@ -104,23 +113,23 @@ func replacePastedImages(raw map[string][]json.RawMessage, sources pointer.Sourc
104113
res.BytesRemoved += len(b.block["source"])
105114
b.block["source"] = pointer.New(material)
106115
res.Replaced++
107-
e.changed = true
116+
changed = true
108117
}
109118

110-
if !e.changed {
119+
if !changed {
111120
continue
112121
}
113122
// A line that cannot be written again keeps its content inline.
114-
if out, err := e.encode(); err == nil {
123+
if out, err := pl.encode(); err == nil {
115124
lines[i] = out
116125
}
117126
}
118127
}
119128
}
120129

121-
// pastedImage is one pasted image block of a transcript entry.
130+
// pastedImage is one pasted image block of a transcript line.
122131
type pastedImage struct {
123-
// block is the image block, which belongs to the content of the entry.
132+
// block is the image block, which belongs to the prompt of the line.
124133
block object
125134
// number is the number of the paste that the agent showed to the user,
126135
// or 0 when it showed none.
@@ -146,40 +155,90 @@ func (b pastedImage) decode() ([]byte, error) {
146155
return data, nil
147156
}
148157

149-
// pastedImages decodes a line that is a user prompt with pasted images, and
150-
// returns the entry and its image blocks with inline data, in order. It
151-
// returns no blocks for any other line.
152-
func pastedImages(line json.RawMessage) (*entry, []pastedImage) {
158+
// pastedLine is a transcript line that holds pasted images.
159+
type pastedLine struct {
160+
fields object
161+
images []pastedImage
162+
// encode writes the line again, with its changed image blocks.
163+
encode func() (json.RawMessage, error)
164+
}
165+
166+
// pastedImages decodes a line that holds pasted images, or returns nil. Claude
167+
// Code writes a prompt in two shapes: a user message, or, when the user sent
168+
// it while the agent was busy, a queued command attachment. A queued prompt
169+
// is not written again as a user message, so both shapes hold pastes.
170+
func pastedImages(line json.RawMessage) *pastedLine {
153171
if !bytes.Contains(line, []byte(`"`+blockImage+`"`)) || !bytes.Contains(line, []byte(`"`+sourceBase64+`"`)) {
154-
return nil, nil
172+
return nil
155173
}
156174

157-
e := decodeEntry(line)
158-
if e == nil || !isUserPrompt(e) {
159-
return nil, nil
160-
}
175+
fields := decodeObject(line)
176+
switch fields.str("type") {
177+
case recordTypeUser:
178+
e := newEntry(fields)
179+
if e == nil || !isUserPrompt(e) {
180+
return nil
181+
}
161182

162-
var blocks []pastedImage
163-
for _, b := range e.content {
164-
if b.str("type") != blockImage {
165-
continue
183+
return newPastedLine(fields, e.content, fields["imagePasteIds"], e.encode)
184+
case recordTypeAttachment:
185+
att := fields.object("attachment")
186+
if !isQueuedPrompt(att) {
187+
return nil
166188
}
167-
if b.object("source").str("type") != sourceBase64 {
168-
continue
189+
prompt := att.objects("prompt")
190+
191+
return newPastedLine(fields, prompt, att["imagePasteIds"], func() (json.RawMessage, error) {
192+
raw, err := marshal(prompt)
193+
if err != nil {
194+
return nil, err
195+
}
196+
att["prompt"] = raw
197+
if fields["attachment"], err = marshal(att); err != nil {
198+
return nil, err
199+
}
200+
201+
return marshal(fields)
202+
})
203+
}
204+
205+
return nil
206+
}
207+
208+
// newPastedLine selects the image blocks with inline data of a prompt, in
209+
// order, or returns nil when there is none.
210+
func newPastedLine(fields object, prompt []object, pasteIDs json.RawMessage, encode func() (json.RawMessage, error)) *pastedLine {
211+
var images []pastedImage
212+
for _, b := range prompt {
213+
if b.str("type") == blockImage && b.object("source").str("type") == sourceBase64 {
214+
images = append(images, pastedImage{block: b})
169215
}
170-
blocks = append(blocks, pastedImage{block: b})
216+
}
217+
if len(images) == 0 {
218+
return nil
171219
}
172220

173221
// Claude Code records the paste numbers that the user saw, one for each
174222
// image of the prompt. When they do not match the images, none is used.
175223
var ids []int
176-
if v, ok := e.fields["imagePasteIds"]; ok && json.Unmarshal(v, &ids) == nil && len(ids) == len(blocks) {
177-
for i := range blocks {
178-
blocks[i].number = ids[i]
224+
if json.Unmarshal(pasteIDs, &ids) == nil && len(ids) == len(images) {
225+
for i := range images {
226+
images[i].number = ids[i]
179227
}
180228
}
181229

182-
return e, blocks
230+
return &pastedLine{fields: fields, images: images, encode: encode}
231+
}
232+
233+
// isQueuedPrompt reports whether an attachment is a prompt that the user
234+
// sent while the agent was busy, and not one that Claude Code queued itself.
235+
func isQueuedPrompt(att object) bool {
236+
if att.str("type") != attachmentQueuedCommand {
237+
return false
238+
}
239+
origin := att.object("origin")
240+
241+
return origin == nil || origin.str("kind") == originHuman
183242
}
184243

185244
// isUserPrompt reports whether the entry is a message that the user wrote,

‎app/cli/internal/trace/claude/pastedimages_test.go‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,22 @@ func paste(text string, pasteIDs []int, images ...obj) obj {
6161
return e
6262
}
6363

64+
// queuedPaste is a prompt with pasted images that the user sent while the
65+
// agent was busy, as Claude Code writes it: an attachment that holds the
66+
// prompt blocks and the paste numbers. origin is the kind of its origin.
67+
func queuedPaste(origin string, pasteIDs []int, images ...obj) obj {
68+
prompt := make([]any, 0, 1+len(images))
69+
prompt = append(prompt, obj{keyType: resultText, resultText: "look at this"})
70+
for _, img := range images {
71+
prompt = append(prompt, img)
72+
}
73+
74+
return obj{keyType: "attachment", "timestamp": pastedTime, "attachment": obj{
75+
keyType: "queued_command", "prompt": prompt, "imagePasteIds": pasteIDs,
76+
"commandMode": "prompt", "origin": obj{"kind": origin}, "humanTurn": true,
77+
}}
78+
}
79+
6480
func TestPastedImages(t *testing.T) {
6581
testCases := []struct {
6682
name string
@@ -116,6 +132,33 @@ func TestPastedImages(t *testing.T) {
116132
{Data: pastedBytes, MediaType: pngType, Number: 1, Timestamp: pastedUTC},
117133
},
118134
},
135+
{
136+
name: "an image pasted while the agent was busy keeps its paste number",
137+
raw: map[string][]obj{mainStream: {
138+
paste("[Image #1] first", []int{1}, imageBlock(pngType, pastedBytes)),
139+
queuedPaste("human", []int{2}, imageBlock("image/jpeg", otherPasted)),
140+
}},
141+
want: []trace.PastedImage{
142+
{Data: pastedBytes, MediaType: pngType, Number: 1, Timestamp: pastedUTC},
143+
{Data: otherPasted, MediaType: "image/jpeg", Number: 2, Timestamp: pastedUTC},
144+
},
145+
},
146+
{
147+
name: "a queued prompt that a human did not send is not a paste",
148+
raw: map[string][]obj{mainStream: {
149+
queuedPaste("task-notification", []int{1}, imageBlock(pngType, pastedBytes)),
150+
}},
151+
},
152+
{
153+
name: "another attachment with an image is not a paste",
154+
raw: map[string][]obj{mainStream: {
155+
func() obj {
156+
e := queuedPaste("human", []int{1}, imageBlock(pngType, pastedBytes))
157+
e["attachment"].(obj)[keyType] = "skill_listing"
158+
return e
159+
}(),
160+
}},
161+
},
119162
{
120163
name: "an image in the result of a file read is not a pasted image",
121164
raw: map[string][]obj{mainStream: {
@@ -224,6 +267,22 @@ func TestReplacePastedImages(t *testing.T) {
224267
},
225268
wantReport: map[string]pointer.Result{pointer.FinderPastedImage: {Replaced: 2}},
226269
},
270+
{
271+
name: "an image pasted while the agent was busy becomes a pointer in the attachment",
272+
session: []obj{
273+
queuedPaste("human", []int{3}, imageBlock(pngType, pastedBytes)),
274+
},
275+
check: func(t *testing.T, out []json.RawMessage) {
276+
e := decode(t, out[0])
277+
assert.Equal(t, "look at this", at(t, e, "attachment", "prompt", 0, resultText))
278+
assert.Equal(t, blockImage, at(t, e, "attachment", "prompt", 1, keyType))
279+
assert.Equal(t, ptr(pastedMaterial), at(t, e, "attachment", "prompt", 1, keySource))
280+
// The other fields of the attachment are kept.
281+
assert.Equal(t, []any{float64(3)}, at(t, e, "attachment", "imagePasteIds"))
282+
assert.Equal(t, "queued_command", at(t, e, "attachment", keyType))
283+
},
284+
wantReport: map[string]pointer.Result{pointer.FinderPastedImage: {Replaced: 1}},
285+
},
227286
{
228287
// The file read finder looks at the image of the result, not
229288
// the pasted image finder.

‎app/cli/internal/trace/claude/speccopies.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,12 @@ type pendingWrite struct {
442442

443443
// decodeEntry decodes a line whose message holds blocks, or returns nil.
444444
func decodeEntry(line json.RawMessage) *entry {
445-
fields := decodeObject(line)
445+
return newEntry(decodeObject(line))
446+
}
447+
448+
// newEntry returns the entry of the decoded fields of a line whose message
449+
// holds blocks, or nil.
450+
func newEntry(fields object) *entry {
446451
message := fields.object("message")
447452
content := message.objects("content")
448453
if content == nil {

0 commit comments

Comments
 (0)