Skip to content

Commit 4f7afeb

Browse files
CharlesthebirdCharlesthebird
committed
refactor(core): store a checkpoint's name in its payload
The name was split out into the `source_name` column and stripped from the protobuf, so every read had to put it back and the column's name no longer described what it held. One representation instead: `Checkpoint.Name` is persisted like every other field, rename follows the transactional pattern `UpdateAgentInstanceName` already uses, and forks take their name from the decoded checkpoint. Reservation decodes the source instance again. That call went when the source name lookup did, and it was also the guard rejecting a corrupt stored payload — without it a checkpoint could be reserved against an instance that cannot be read. The default-name assertions now compare against literal strings, so a wrong constant in `defaultCheckpointName` fails them, and rename is covered against creating, failed and deleting checkpoints. Signed-off-by: Nicholas Bucher <[email protected]>
1 parent 90dadcd commit 4f7afeb

4 files changed

Lines changed: 69 additions & 57 deletions

File tree

‎go/core/internal/database/checkpoints.go‎

Lines changed: 37 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,8 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re
4646
if err != nil {
4747
return fmt.Errorf("lock checkpoint: %w", err)
4848
}
49-
if _, err := toAgentInstanceCheckpoint(checkpoint); err != nil {
49+
source, err := toAgentInstanceCheckpoint(checkpoint)
50+
if err != nil {
5051
return err
5152
}
5253
if checkpoint.PreparedRevision == nil {
@@ -67,7 +68,7 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re
6768
now := timestamppb.Now()
6869
instance := &apiv1alpha1.AgentInstance{
6970
Id: instanceID, Creator: userID, ContextId: sourceContextID.String(),
70-
Name: checkpoint.SourceName,
71+
Name: source.GetName(),
7172
Harness: &apiv1alpha1.ResourceReference{Namespace: revision.Namespace, Name: revision.HarnessName},
7273
AgentTemplate: &apiv1alpha1.ResourceReference{Namespace: revision.Namespace, Name: revision.AgentTemplateName},
7374
PreparedRevision: *checkpoint.PreparedRevision,
@@ -190,6 +191,10 @@ func (c *Client) ReserveAgentInstanceCheckpoint(ctx context.Context, checkpoint
190191
if err != nil {
191192
return fmt.Errorf("lock AgentInstance %s: %w", checkpoint.GetAgentInstanceId(), err)
192193
}
194+
// Decoding rejects a corrupt stored instance before it becomes a checkpoint source.
195+
if _, err := toAgentInstance(instance); err != nil {
196+
return err
197+
}
193198
if instance.State != "AGENT_INSTANCE_STATE_READY" || instance.Operation != "AGENT_INSTANCE_OPERATION_UNSPECIFIED" {
194199
return fmt.Errorf("AgentInstance %s cannot checkpoint in state %s with operation %s: %w", checkpoint.GetAgentInstanceId(), instance.State, instance.Operation, ErrConflict)
195200
}
@@ -243,24 +248,24 @@ func (c *Client) ReserveAgentInstanceCheckpoint(ctx context.Context, checkpoint
243248
value.State = apiv1alpha1.CheckpointState_CHECKPOINT_STATE_CREATING
244249
value.CreatedAt = timestamppb.Now()
245250
value.Failure = nil
246-
data, err := marshalCheckpoint(value)
251+
value.Name = defaultCheckpointName(sourceID.String(), boundary.ID)
252+
data, err := proto.Marshal(value)
247253
if err != nil {
248254
return fmt.Errorf("encode checkpoint: %w", err)
249255
}
250256
row, err := queryOne(ctx, tx, `
251257
INSERT INTO agent_instance_checkpoint (id, source_instance_id, user_id, request_id, head_task_id,
252258
history_sequence, snapshot_atespace, snapshot_uri, snapshot_content_scope, source_history_id,
253-
prepared_revision, data, source_name, state) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9,
254-
$10, $11, $12, $13, 'CREATING')
259+
prepared_revision, data, state) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9,
260+
$10, $11, $12, 'CREATING')
255261
ON CONFLICT DO NOTHING
256262
RETURNING id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace,
257-
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision,
258-
source_name
263+
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision
259264
`,
260265
pgx.RowToStructByName[agentInstanceCheckpointRow], checkpoint.GetId(),
261266
checkpoint.GetAgentInstanceId(), userID, requestID, boundary.ID, *boundary.HistorySequence,
262267
*boundary.SnapshotAtespace, *boundary.SnapshotURI, *boundary.SnapshotContentScope, instance.HistoryID,
263-
instance.PreparedRevision, data, defaultCheckpointName(sourceID.String(), boundary.ID),
268+
instance.PreparedRevision, data,
264269
)
265270
if errors.Is(err, pgx.ErrNoRows) {
266271
existing, existingErr := readCheckpointRequest(ctx, tx, userID, requestID)
@@ -320,7 +325,7 @@ func (c *Client) FinalizeAgentInstanceCheckpoint(ctx context.Context, id, tagUID
320325
result.State = apiv1alpha1.CheckpointState_CHECKPOINT_STATE_FAILED
321326
result.Failure = &apiv1alpha1.Failure{Reason: "SnapshotTagFailed", Message: failure}
322327
}
323-
data, err := marshalCheckpoint(result)
328+
data, err := proto.Marshal(result)
324329
if err != nil {
325330
return fmt.Errorf("encode checkpoint: %w", err)
326331
}
@@ -386,8 +391,8 @@ func checkpointSnapshot(row agentInstanceCheckpointRow) *AgentInstanceTaskSnapsh
386391
func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, userID, afterID string, limit int) ([]*apiv1alpha1.Checkpoint, error) {
387392
rows, err := queryMany(ctx, c.db, `
388393
SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace,
389-
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision,
390-
source_name FROM agent_instance_checkpoint
394+
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision
395+
FROM agent_instance_checkpoint
391396
WHERE source_instance_id = $1
392397
AND user_id = $2
393398
AND state = 'READY'
@@ -428,7 +433,7 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use
428433
return err
429434
}
430435
checkpoint.State = apiv1alpha1.CheckpointState_CHECKPOINT_STATE_DELETING
431-
data, err := marshalCheckpoint(checkpoint)
436+
data, err := proto.Marshal(checkpoint)
432437
if err != nil {
433438
return fmt.Errorf("encode checkpoint: %w", err)
434439
}
@@ -469,15 +474,6 @@ func (c *Client) DeleteAgentInstanceCheckpoint(ctx context.Context, id, userID s
469474
return nil
470475
}
471476

472-
// marshalCheckpoint encodes a checkpoint for the data column with its name left out.
473-
// The name lives in its own column so a rename is one UPDATE; a copy in the payload as
474-
// well would disagree with that column the moment anybody renamed the checkpoint.
475-
func marshalCheckpoint(checkpoint *apiv1alpha1.Checkpoint) ([]byte, error) {
476-
stored := proto.Clone(checkpoint).(*apiv1alpha1.Checkpoint)
477-
stored.Name = ""
478-
return proto.Marshal(stored)
479-
}
480-
481477
// defaultCheckpointName names a boundary by the conversation it was taken from and the
482478
// turn it sits at, the two things that tell two boundaries apart. Callers pass the
483479
// parsed id rather than what arrived on the wire, so that clearing a name restores the
@@ -499,23 +495,31 @@ func (c *Client) UpdateCheckpointName(ctx context.Context, id, userID, name stri
499495
if err != nil {
500496
return notFoundOr(err)
501497
}
498+
checkpoint, err := toAgentInstanceCheckpoint(row)
499+
if err != nil {
500+
return err
501+
}
502502
if name == "" {
503503
name = defaultCheckpointName(row.SourceInstanceID.String(), row.HeadTaskID)
504504
}
505+
checkpoint.Name = name
506+
data, err := proto.Marshal(checkpoint)
507+
if err != nil {
508+
return fmt.Errorf("encode checkpoint: %w", err)
509+
}
505510
tag, err := tx.Exec(ctx, `
506511
UPDATE agent_instance_checkpoint
507-
SET source_name = $3
508-
WHERE id = $1 AND user_id = $2
509-
`, row.ID, userID, name)
512+
SET data = $1
513+
WHERE id = $2 AND user_id = $3
514+
`, data, row.ID, userID)
510515
if err != nil {
511516
return err
512517
}
513518
if tag.RowsAffected() != 1 {
514519
return ErrNotFound
515520
}
516-
row.SourceName = name
517-
result, err = toAgentInstanceCheckpoint(row)
518-
return err
521+
result = checkpoint
522+
return nil
519523
})
520524
if err != nil {
521525
return nil, fmt.Errorf("rename checkpoint %s: %w", id, err)
@@ -535,9 +539,6 @@ func toAgentInstanceCheckpoint(row agentInstanceCheckpointRow) (*apiv1alpha1.Che
535539
strings.TrimPrefix(checkpoint.GetState().String(), "CHECKPOINT_STATE_") != row.State {
536540
return nil, fmt.Errorf("checkpoint %s payload disagrees with indexed columns", row.ID)
537541
}
538-
// The name lives in its own column so a rename stays one UPDATE, which also keeps
539-
// it out of the payload cross-check above.
540-
checkpoint.Name = row.SourceName
541542
return checkpoint, nil
542543
}
543544

@@ -556,15 +557,6 @@ type agentInstanceCheckpointRow struct {
556557
Data []byte
557558
SourceHistoryID uuid.UUID
558559
PreparedRevision *string
559-
/*
560-
* The checkpoint's own display name, and the name its forks are given.
561-
*
562-
* The column predates the name being something a reader chooses — it held a copy of
563-
* the source conversation's title so a fork could still be named after a conversation
564-
* that had since been deleted. The job is the same one, so it was kept rather than a
565-
* second column added beside it; only what fills it has changed.
566-
*/
567-
SourceName string
568560
}
569561

570562
// lockCheckpoint locks a checkpoint until the caller's transaction ends, optionally
@@ -573,8 +565,8 @@ type agentInstanceCheckpointRow struct {
573565
func lockCheckpoint(ctx context.Context, db pgx.Tx, id string, allUsers bool, userID string, state *string) (agentInstanceCheckpointRow, error) {
574566
return queryOne(ctx, db, `
575567
SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace,
576-
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision,
577-
source_name FROM agent_instance_checkpoint
568+
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision
569+
FROM agent_instance_checkpoint
578570
WHERE id = $1
579571
-- Only internal finalization explicitly opts out of owner filtering.
580572
AND ($2::boolean OR user_id = $3)
@@ -621,8 +613,8 @@ func insertReplayedTask(ctx context.Context, db dbExecutor, task agentInstanceTa
621613
func readCheckpoint(ctx context.Context, db dbExecutor, id, userID string, state *string) (agentInstanceCheckpointRow, error) {
622614
return queryOne(ctx, db, `
623615
SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace,
624-
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision,
625-
source_name FROM agent_instance_checkpoint
616+
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision
617+
FROM agent_instance_checkpoint
626618
WHERE id = $1 AND user_id = $2
627619
-- Lifecycle work also reads creating and deleting checkpoints.
628620
AND ($3::text IS NULL OR state = $3)
@@ -634,8 +626,8 @@ func readCheckpoint(ctx context.Context, db dbExecutor, id, userID string, state
634626
func readCheckpointRequest(ctx context.Context, db dbExecutor, userID, requestID string) (agentInstanceCheckpointRow, error) {
635627
return queryOne(ctx, db, `
636628
SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace,
637-
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision,
638-
source_name FROM agent_instance_checkpoint
629+
snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision
630+
FROM agent_instance_checkpoint
639631
WHERE user_id = $1 AND request_id = $2
640632
`, pgx.RowToStructByName[agentInstanceCheckpointRow], userID, requestID)
641633
}

‎go/core/internal/database/client_agent_instance_test.go‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,10 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
325325
if _, _, err := client.GetAgentInstanceCheckpointSnapshot(ctx, checkpoint.GetId(), "mallory"); !errors.Is(err, ErrNotFound) {
326326
t.Fatalf("snapshot lookup by another user = %v, want not found", err)
327327
}
328-
require.Equal(t, defaultCheckpointName(instanceID, "task-1"), checkpoint.GetName())
328+
require.Equal(t, "11111111-1111-4111-8111-111111111111-task-1", checkpoint.GetName())
329+
if _, err := client.UpdateCheckpointName(ctx, checkpoint.GetId(), "alice", "too early"); !errors.Is(err, ErrNotFound) {
330+
t.Fatalf("rename of a creating checkpoint = %v, want not found", err)
331+
}
329332
if checkpoint.HeadTaskId != "task-1" || snapshot == nil ||
330333
*snapshot != (AgentInstanceTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/snapshot-1", ContentScope: "DATA"}) || checkpoint.HistorySequence == 0 {
331334
t.Fatalf("checkpoint boundary = %+v", checkpoint)
@@ -363,7 +366,7 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
363366
require.Equal(t, "Before the detour", stored.GetName())
364367
restored, err := client.UpdateCheckpointName(ctx, checkpoint.GetId(), "alice", "")
365368
require.NoError(t, err)
366-
require.Equal(t, defaultCheckpointName(instanceID, "task-1"), restored.GetName())
369+
require.Equal(t, "11111111-1111-4111-8111-111111111111-task-1", restored.GetName())
367370
retained, tagUID, err := client.GetAgentInstanceCheckpointSnapshot(ctx, checkpoint.GetId(), "alice")
368371
if err != nil || tagUID != "tag-uid" || retained.URI != "s3://tags/checkpoint" || retained.ContentScope != snapshot.ContentScope {
369372
t.Fatalf("checkpoint tag = %q, error %v", tagUID, err)
@@ -379,6 +382,9 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
379382
if err != nil || failed.State != apiv1alpha1.CheckpointState_CHECKPOINT_STATE_FAILED || failed.GetFailure().GetMessage() != "tag creation failed" {
380383
t.Fatalf("failed checkpoint = %+v, error %v", failed, err)
381384
}
385+
if _, err := client.UpdateCheckpointName(ctx, failed.GetId(), "alice", "doomed"); !errors.Is(err, ErrNotFound) {
386+
t.Fatalf("rename of a failed checkpoint = %v, want not found", err)
387+
}
382388
if err := client.DeleteAgentInstance(ctx, instanceID); err != nil {
383389
t.Fatal(err)
384390
}
@@ -400,6 +406,9 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
400406
if _, err := client.GetAgentInstanceCheckpoint(ctx, checkpoint.GetId(), "alice"); !errors.Is(err, ErrNotFound) {
401407
t.Fatalf("deleting checkpoint is publicly visible: %v", err)
402408
}
409+
if _, err := client.UpdateCheckpointName(ctx, checkpoint.GetId(), "alice", "too late"); !errors.Is(err, ErrNotFound) {
410+
t.Fatalf("rename of a deleting checkpoint = %v, want not found", err)
411+
}
403412
deletingSnapshot, deletingTagUID, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.GetId(), "alice")
404413
if err != nil || deletingSnapshot == nil || *deletingSnapshot != *retained || deletingTagUID != "tag-uid" {
405414
t.Fatalf("deleting snapshot = %+v, tag = %q, error %v", deletingSnapshot, deletingTagUID, err)
@@ -409,6 +418,25 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
409418
}
410419
}
411420

421+
func TestReserveAgentInstanceCheckpointRejectsCorruptSource(t *testing.T) {
422+
db := setupTestDB(t)
423+
ctx := context.Background()
424+
client := NewClient(db)
425+
instanceID := "99999999-9999-4999-8999-999999999999"
426+
if _, err := db.Exec(ctx, `
427+
INSERT INTO a2a_context (id, user_id, context_id) VALUES ($1, 'alice', $1)
428+
`, instanceID); err != nil {
429+
t.Fatal(err)
430+
}
431+
if _, err := db.Exec(ctx, `
432+
INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ($1, 'alice', 'instance-request', $1, $1, 'AGENT_INSTANCE_STATE_READY', $2)
433+
`, instanceID, []byte{0xff}); err != nil {
434+
t.Fatal(err)
435+
}
436+
_, _, err := client.ReserveAgentInstanceCheckpoint(ctx, &apiv1alpha1.Checkpoint{Id: uuid.NewString(), AgentInstanceId: instanceID}, "alice", "checkpoint-request")
437+
require.ErrorContains(t, err, "decode AgentInstance")
438+
}
439+
412440
func TestForkAgentInstanceCopiesBoundedHistory(t *testing.T) {
413441
db := setupTestDB(t)
414442
client := NewClient(db)
@@ -564,7 +592,7 @@ func TestForkAgentInstanceCopiesBoundedHistory(t *testing.T) {
564592
t.Fatal(err)
565593
}
566594
fork2, created, err := client.ForkAgentInstance(ctx, checkpoint2.GetId(), "alice", "fork-request-2", fork2ID)
567-
require.Equal(t, defaultCheckpointName(fork.GetId(), "task-1"), fork2.GetName())
595+
require.Equal(t, "77777777-7777-4777-8777-777777777777-task-1", fork2.GetName())
568596
if err != nil || !created || fork2.GetId() != fork2ID {
569597
t.Fatalf("fork of fork = %+v, created %v, error %v", fork2, created, err)
570598
}

‎go/core/internal/database/protobuf_test.go‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -209,14 +209,7 @@ func TestProtobufPersistenceLifecycle(t *testing.T) {
209209
require.NoError(t, err)
210210
storedCheckpoint := &apiv1alpha1.Checkpoint{}
211211
require.NoError(t, proto.Unmarshal(checkpointRow.Data, storedCheckpoint))
212-
// The name is the one field the payload does not carry: it lives in its own column
213-
// so a rename is a single UPDATE, and a copy here would disagree with that column
214-
// the moment anybody renamed the checkpoint.
215-
require.Empty(t, storedCheckpoint.GetName())
216-
require.Equal(t, checkpointRow.SourceName, checkpoint.GetName())
217-
withoutName := proto.Clone(checkpoint).(*apiv1alpha1.Checkpoint)
218-
withoutName.Name = ""
219-
require.True(t, proto.Equal(withoutName, storedCheckpoint))
212+
require.True(t, proto.Equal(checkpoint, storedCheckpoint))
220213
require.Equal(t, checkpointRequest.ProtoReflect().GetUnknown(), checkpoint.ProtoReflect().GetUnknown())
221214
_, err = client.GetAgentInstanceCheckpoint(ctx, checkpoint.Id, "mallory")
222215
require.ErrorIs(t, err, ErrNotFound)

‎go/core/pkg/migrations/core/000001_initial.sql‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,6 @@ CREATE TABLE agent_instance_checkpoint (
8787
data BYTEA NOT NULL,
8888
source_history_id UUID NOT NULL REFERENCES a2a_context(id) ON DELETE RESTRICT,
8989
prepared_revision TEXT REFERENCES runtime_revision(revision) ON DELETE RESTRICT,
90-
source_name TEXT NOT NULL DEFAULT '',
9190
CHECK (snapshot_content_scope IN ('FULL', 'DATA')),
9291
CHECK (state IN ('CREATING', 'READY', 'FAILED', 'DELETING')),
9392
UNIQUE (user_id, request_id)

0 commit comments

Comments
 (0)