From 2441aced5eb01402e0c83185b3dfd4f37af58929 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Wed, 9 Sep 2026 19:18:47 +0000 Subject: [PATCH 1/9] fix: retain checkpoint visibility across forks Signed-off-by: Eitan Yarmush --- go/core/internal/database/agent_instances.go | 4 +- .../database/checkpoint_lineage_test.go | 122 ++++++++++++++++++ go/core/internal/database/checkpoints.go | 47 +++++-- go/core/internal/database/protobuf_test.go | 10 ++ .../pkg/migrations/core/000001_initial.sql | 8 ++ go/core/test/e2e/interaction_test.go | 39 ++++-- 6 files changed, 205 insertions(+), 25 deletions(-) create mode 100644 go/core/internal/database/checkpoint_lineage_test.go diff --git a/go/core/internal/database/agent_instances.go b/go/core/internal/database/agent_instances.go index c581542a2b..93a0ac52a2 100644 --- a/go/core/internal/database/agent_instances.go +++ b/go/core/internal/database/agent_instances.go @@ -411,8 +411,8 @@ func insertAgentInstanceRecords(ctx context.Context, db dbExecutor, instance *ap return agentInstanceRow{}, err } if err := execSQL(ctx, db, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ($1, $2, $3) - `, historyID, instance.Creator, instance.ContextId); err != nil { + INSERT INTO a2a_context (id, user_id, context_id, source_checkpoint_id) VALUES ($1, $2, $3, $4) + `, historyID, instance.Creator, instance.ContextId, sourceCheckpointID); err != nil { return agentInstanceRow{}, fmt.Errorf("insert A2A context: %w", err) } return queryOne(ctx, db, ` diff --git a/go/core/internal/database/checkpoint_lineage_test.go b/go/core/internal/database/checkpoint_lineage_test.go new file mode 100644 index 0000000000..c342b93c2c --- /dev/null +++ b/go/core/internal/database/checkpoint_lineage_test.go @@ -0,0 +1,122 @@ +package database + +import ( + "slices" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" +) + +func TestCheckpointLineage(t *testing.T) { + client := NewClient(setupTestDB(t)) + ctx := t.Context() + agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + source, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + require.NoError(t, err) + source, err = markAgentInstanceReady(ctx, client, source.Id, "source.example") + require.NoError(t, err) + + checkpointNextTurn := func(instance *apiv1alpha1.AgentInstance) *apiv1alpha1.Checkpoint { + t.Helper() + task := newAgentInstanceTask(uuid.NewString(), uuid.NewString()) + task.ContextID = instance.ContextId + _, _, err := client.CreateAgentInstanceTask(ctx, instance.Id, []byte(task.ID), task) + require.NoError(t, err) + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, client.StoreAgentInstanceTaskEvent(ctx, instance.Id, task, task, + &AgentInstanceTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/" + string(task.ID), ContentScope: "DATA"})) + checkpoint, _, err := client.ReserveAgentInstanceCheckpoint(ctx, + &apiv1alpha1.Checkpoint{Id: uuid.NewString(), AgentInstanceId: instance.Id}, "alice", uuid.NewString()) + require.NoError(t, err) + checkpoint, err = client.FinalizeAgentInstanceCheckpoint(ctx, checkpoint.Id, "tag-"+checkpoint.Id, "s3://tags/"+checkpoint.Id, "") + require.NoError(t, err) + return checkpoint + } + fork := func(checkpoint *apiv1alpha1.Checkpoint) *apiv1alpha1.AgentInstance { + t.Helper() + instance, created, err := client.ForkAgentInstance(ctx, checkpoint.Id, "alice", uuid.NewString(), uuid.NewString()) + require.NoError(t, err) + require.True(t, created) + instance, err = markAgentInstanceReady(ctx, client, instance.Id, "fork.example") + require.NoError(t, err) + return instance + } + assertListed := func(instanceID string, want ...*apiv1alpha1.Checkpoint) { + t.Helper() + got, err := client.ListAgentInstanceCheckpoints(ctx, instanceID, "alice", "", 100) + require.NoError(t, err) + require.Len(t, got, len(want)) + for _, checkpoint := range want { + require.True(t, slices.ContainsFunc(got, func(candidate *apiv1alpha1.Checkpoint) bool { + return proto.Equal(checkpoint, candidate) + }), "missing checkpoint %s with its original provenance", checkpoint.Id) + } + } + + deleting := checkpointNextTurn(source) + c1 := checkpointNextTurn(source) + c2 := checkpointNextTurn(source) + _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") + require.NoError(t, err) + branch := fork(c2) + // A fork created after deletion begins must not prevent its cleanup from retrying. + _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") + require.NoError(t, err) + require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice")) + c3 := checkpointNextTurn(source) + c4 := checkpointNextTurn(branch) + nested := fork(c4) + c5 := checkpointNextTurn(branch) + assertListed(source.Id, c1, c2, c3) + assertListed(branch.Id, c1, c2, c4, c5) + assertListed(nested.Id, c1, c2, c4) + + // Paging uses checkpoint IDs across every ancestor, without duplicates. + wantIDs := []string{c1.Id, c2.Id, c4.Id} + slices.Sort(wantIDs) + var afterID string + for _, id := range wantIDs { + page, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "alice", afterID, 1) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, id, page[0].Id) + afterID = id + } + page, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "alice", afterID, 1) + require.NoError(t, err) + require.Empty(t, page) + unauthorized, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "mallory", "", 100) + require.NoError(t, err) + require.Empty(t, unauthorized) + assertListed(uuid.NewString()) + + require.NoError(t, client.DeleteAgentInstance(ctx, source.Id)) + require.NoError(t, client.DeleteAgentInstance(ctx, branch.Id)) + assertListed(nested.Id, c1, c2, c4) + // No fork starts at c1 yet; it is retained because it precedes c2 in the inherited history. + _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, c1.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + // A checkpoint inherited by a fork remains independently usable as a fork source. + earlier := fork(c1) + assertListed(earlier.Id, c1) + require.NoError(t, client.DeleteAgentInstance(ctx, earlier.Id)) + require.NoError(t, client.DeleteAgentInstance(ctx, nested.Id)) + + // Retained histories protect the entire inherited prefix, even after instance deletion. + for _, checkpoint := range []*apiv1alpha1.Checkpoint{c1, c2, c4} { + _, _, err := client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + got, err := client.GetAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.True(t, proto.Equal(checkpoint, got)) + } + for _, checkpoint := range []*apiv1alpha1.Checkpoint{c3, c5} { + _, _, err := client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice")) + } +} diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 0cec934fd8..27a68a4c26 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -62,8 +62,10 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re if err != nil { return fmt.Errorf("get checkpoint runtime revision: %w", err) } + // Serialize new lineage with deletion of any checkpoint in its inherited prefix. + // NO KEY UPDATE leaves foreign-key checks free to read the history. sourceContextID, err := queryOne(ctx, tx, ` - SELECT context_id FROM a2a_context WHERE id = $1 + SELECT context_id FROM a2a_context WHERE id = $1 FOR NO KEY UPDATE `, pgx.RowTo[uuid.UUID], checkpoint.SourceHistoryID) if err != nil { return fmt.Errorf("get checkpoint context: %w", err) @@ -392,15 +394,28 @@ func checkpointSnapshot(row agentInstanceCheckpointRow) *AgentInstanceTaskSnapsh } } -// ListAgentInstanceCheckpoints returns an owner's READY checkpoints for the source -// instance in ascending ID order after afterID, up to limit. Retained checkpoints remain -// listable after the source instance is deleted. +// ListAgentInstanceCheckpoints returns an owner's READY checkpoints in the instance's +// history and inherited prefixes, in ascending ID order after afterID, up to limit. +// Checkpoint fields retain their source provenance. Locally created checkpoints remain +// listable after the instance is deleted. func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, userID, afterID string, limit int) ([]*apiv1alpha1.Checkpoint, error) { rows, err := queryMany(ctx, c.db, ` + WITH RECURSIVE lineage (history_id, boundary) AS ( + SELECT history_id, NULL::bigint FROM agent_instance WHERE id = $1 AND user_id = $2 + UNION + SELECT c.source_history_id, c.history_sequence + FROM lineage l + JOIN a2a_context h ON h.id = l.history_id + JOIN agent_instance_checkpoint c ON c.id = h.source_checkpoint_id + WHERE h.user_id = $2 AND c.user_id = $2 + ) SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace, snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision - FROM agent_instance_checkpoint - WHERE source_instance_id = $1 + FROM agent_instance_checkpoint c + WHERE (source_instance_id = $1 OR EXISTS ( + SELECT 1 FROM lineage l WHERE l.history_id = c.source_history_id + AND (l.boundary IS NULL OR c.history_sequence <= l.boundary) + )) AND user_id = $2 AND state = 'READY' AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) @@ -425,8 +440,9 @@ func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, u // BeginDeleteAgentInstanceCheckpoint marks an owned READY checkpoint DELETING and returns // the snapshot and tag identity for external cleanup. Retrying a DELETING checkpoint is -// allowed. A fork reference, another lifecycle state, or a missing/unowned checkpoint -// returns ErrNotFound. +// allowed even if new forks have since inherited the boundary. An inherited history +// reference blocks starting deletion; another state or a missing/unowned checkpoint +// also returns ErrNotFound. func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, userID string) (*AgentInstanceTaskSnapshot, string, error) { var snapshot *AgentInstanceTaskSnapshot var tagUID string @@ -435,6 +451,12 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use if err != nil { return notFoundOr(err) } + // Take the same history lock as ForkAgentInstance before checking its descendants. + if _, err := queryOne(ctx, tx, ` + SELECT id FROM a2a_context WHERE id = $1 FOR NO KEY UPDATE + `, pgx.RowTo[uuid.UUID], row.SourceHistoryID); err != nil { + return fmt.Errorf("lock checkpoint history: %w", err) + } checkpoint, err := toAgentInstanceCheckpoint(row) if err != nil { return err @@ -449,9 +471,12 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use SET state = 'DELETING', data = $3 WHERE agent_instance_checkpoint.id = $1 AND agent_instance_checkpoint.user_id = $2 AND agent_instance_checkpoint.state IN ('READY', 'DELETING') - AND NOT EXISTS ( - SELECT 1 FROM agent_instance i WHERE i.pinned_checkpoint_id = agent_instance_checkpoint.id - ) + AND (agent_instance_checkpoint.state = 'DELETING' OR NOT EXISTS ( + SELECT 1 FROM a2a_context h + JOIN agent_instance_checkpoint parent ON parent.id = h.source_checkpoint_id + WHERE parent.source_history_id = agent_instance_checkpoint.source_history_id + AND parent.history_sequence >= agent_instance_checkpoint.history_sequence + )) `, row.ID, userID, data) if err != nil { return err diff --git a/go/core/internal/database/protobuf_test.go b/go/core/internal/database/protobuf_test.go index a841e8fddd..b7e8fd0fd6 100644 --- a/go/core/internal/database/protobuf_test.go +++ b/go/core/internal/database/protobuf_test.go @@ -238,6 +238,16 @@ func TestProtobufPersistenceLifecycle(t *testing.T) { require.Equal(t, a2apb.TaskState_TASK_STATE_COMPLETED, forkTask.Status.State) require.NoError(t, client.DeleteAgentInstance(ctx, fork.Id)) _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + // Exercise deletion with unknown fields on a boundary outside the retained fork. + require.NoError(t, client.StoreAgentInstanceTaskEvent(ctx, instance.Id, task, task, + &AgentInstanceTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/later", ContentScope: "DATA"})) + checkpointRequest.Id = uuid.NewString() + checkpoint, _, err = client.ReserveAgentInstanceCheckpoint(ctx, checkpointRequest, "alice", "later-checkpoint") + require.NoError(t, err) + checkpoint, err = client.FinalizeAgentInstanceCheckpoint(ctx, checkpoint.Id, "later-tag", "s3://tags/later", "") + require.NoError(t, err) + _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.Id, "alice") require.NoError(t, err) checkpointRow, err = readCheckpoint(ctx, q, checkpoint.Id, "alice", nil) require.NoError(t, err) diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index 0a86455ae3..c522ee8f60 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -94,10 +94,17 @@ CREATE TABLE agent_instance_checkpoint ( ); CREATE INDEX agent_instance_checkpoint_list_idx ON agent_instance_checkpoint (source_instance_id, id); +CREATE INDEX agent_instance_checkpoint_history_idx + ON agent_instance_checkpoint (source_history_id, history_sequence); CREATE UNIQUE INDEX agent_instance_checkpoint_one_creating_idx ON agent_instance_checkpoint (source_instance_id) WHERE state = 'CREATING'; +ALTER TABLE a2a_context ADD COLUMN source_checkpoint_id UUID + REFERENCES agent_instance_checkpoint(id) ON DELETE RESTRICT; +CREATE INDEX a2a_context_source_checkpoint_idx ON a2a_context (source_checkpoint_id) + WHERE source_checkpoint_id IS NOT NULL; + CREATE TABLE agent_instance ( id UUID PRIMARY KEY, user_id TEXT NOT NULL CHECK (user_id <> ''), @@ -279,6 +286,7 @@ DROP TABLE agent_instance_share; DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; +ALTER TABLE a2a_context DROP COLUMN source_checkpoint_id; DROP TABLE agent_instance_checkpoint; DROP TABLE a2a_context; DROP TABLE agent_template_harness_pair; diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index 3f426c17b1..7e7b408ee1 100644 --- a/go/core/test/e2e/interaction_test.go +++ b/go/core/test/e2e/interaction_test.go @@ -207,6 +207,12 @@ func TestAgentInstanceCheckpoint(t *testing.T) { if fork.GetId() == fixture.instanceID || fork.GetState() != apiv1alpha1.AgentInstanceState_AGENT_INSTANCE_STATE_READY { t.Fatalf("fork = %+v", fork) } + listed, err = fixture.checkpoints.ListCheckpoints(fixture.ctx, &apiv1alpha1.ListCheckpointsRequest{ + AgentInstanceId: fork.GetId(), + }) + if err != nil || len(listed.GetCheckpoints()) != 1 || listed.GetCheckpoints()[0].GetId() != checkpoint.GetId() { + t.Fatalf("inherited checkpoints = %+v, error %v", listed.GetCheckpoints(), err) + } t.Cleanup(func() { cleanupCtx, cleanupCancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute) defer cleanupCancel() @@ -275,14 +281,6 @@ func TestAgentInstanceCheckpoint(t *testing.T) { if nestedTask.Status.State != a2atype.TaskStateCompleted || !strings.Contains(taskText(nestedTask), "The answer is 4.") { t.Fatalf("nested fork lost checkpoint memory: %+v", nestedTask) } - if _, err := fixture.instances.DeleteAgentInstance(nestedCtx, &apiv1alpha1.DeleteAgentInstanceRequest{AgentInstanceId: nestedID}); err != nil { - t.Fatalf("delete nested fork: %v", err) - } - if _, err := fixture.checkpoints.DeleteCheckpoint(forkCtx, &apiv1alpha1.DeleteCheckpointRequest{ - CheckpointId: forkCheckpoint.GetCheckpoint().GetId(), - }); err != nil { - t.Fatalf("delete fresh fork checkpoint: %v", err) - } forkFixture := &interactionFixture{ctx: forkCtx, client: fixture.client} _, _, forkTask := forkFixture.send(t, "What was the answer before the checkpoint?") if forkTask.Status.State != a2atype.TaskStateCompleted || !strings.Contains(taskText(forkTask), "The answer is 4.") { @@ -293,10 +291,27 @@ func TestAgentInstanceCheckpoint(t *testing.T) { }); err != nil { t.Fatalf("delete fork AgentInstance: %v", err) } - if _, err := fixture.checkpoints.DeleteCheckpoint(fixture.ctx, &apiv1alpha1.DeleteCheckpointRequest{ - CheckpointId: checkpoint.GetId(), - }); err != nil { - t.Fatalf("delete checkpoint: %v", err) + listed, err = fixture.checkpoints.ListCheckpoints(nestedCtx, &apiv1alpha1.ListCheckpointsRequest{AgentInstanceId: nestedID}) + if err != nil || len(listed.GetCheckpoints()) != 2 { + t.Fatalf("nested checkpoints after ancestor deletion = %+v, error %v", listed.GetCheckpoints(), err) + } + wantIDs := map[string]bool{checkpoint.GetId(): true, forkCheckpoint.GetCheckpoint().GetId(): true} + for _, inherited := range listed.GetCheckpoints() { + if !wantIDs[inherited.GetId()] { + t.Fatalf("unexpected inherited checkpoint %s", inherited.GetId()) + } + delete(wantIDs, inherited.GetId()) + } + if _, err := fixture.instances.DeleteAgentInstance(nestedCtx, &apiv1alpha1.DeleteAgentInstanceRequest{AgentInstanceId: nestedID}); err != nil { + t.Fatalf("delete nested fork: %v", err) + } + // Deleting instances retains their histories and the checkpoints they inherit. + for _, id := range []string{checkpoint.GetId(), forkCheckpoint.GetCheckpoint().GetId()} { + if _, err := fixture.checkpoints.DeleteCheckpoint(fixture.ctx, &apiv1alpha1.DeleteCheckpointRequest{ + CheckpointId: id, + }); status.Code(err) != codes.NotFound { + t.Fatalf("delete retained checkpoint %s = %v, want NotFound", id, err) + } } } From 8f6b2208c86dd2d34b19cb9f8e7609aec865310f Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Wed, 9 Sep 2026 20:12:08 +0000 Subject: [PATCH 2/9] refactor: inline checkpoint lineage column in baseline schema Signed-off-by: Eitan Yarmush --- .../pkg/migrations/core/000001_initial.sql | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index c522ee8f60..efd0826314 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -66,12 +66,15 @@ CREATE INDEX agent_template_harness_pair_name_idx ON agent_template_harness_pair (namespace, agent_template_name, harness_name); CREATE TABLE a2a_context ( - id UUID PRIMARY KEY, - user_id TEXT NOT NULL CHECK (user_id <> ''), - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - context_id UUID NOT NULL, + id UUID PRIMARY KEY, + user_id TEXT NOT NULL CHECK (user_id <> ''), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + context_id UUID NOT NULL, + source_checkpoint_id UUID, CONSTRAINT a2a_context_binding_key UNIQUE (id, context_id) ); +CREATE INDEX a2a_context_source_checkpoint_idx ON a2a_context (source_checkpoint_id) + WHERE source_checkpoint_id IS NOT NULL; CREATE TABLE agent_instance_checkpoint ( id UUID PRIMARY KEY, @@ -100,10 +103,9 @@ CREATE UNIQUE INDEX agent_instance_checkpoint_one_creating_idx ON agent_instance_checkpoint (source_instance_id) WHERE state = 'CREATING'; -ALTER TABLE a2a_context ADD COLUMN source_checkpoint_id UUID - REFERENCES agent_instance_checkpoint(id) ON DELETE RESTRICT; -CREATE INDEX a2a_context_source_checkpoint_idx ON a2a_context (source_checkpoint_id) - WHERE source_checkpoint_id IS NOT NULL; +-- History and checkpoints reference each other, so this foreign key follows both tables. +ALTER TABLE a2a_context ADD CONSTRAINT a2a_context_source_checkpoint_id_fkey + FOREIGN KEY (source_checkpoint_id) REFERENCES agent_instance_checkpoint(id) ON DELETE RESTRICT; CREATE TABLE agent_instance ( id UUID PRIMARY KEY, @@ -286,7 +288,7 @@ DROP TABLE agent_instance_share; DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; -ALTER TABLE a2a_context DROP COLUMN source_checkpoint_id; +ALTER TABLE a2a_context DROP CONSTRAINT a2a_context_source_checkpoint_id_fkey; DROP TABLE agent_instance_checkpoint; DROP TABLE a2a_context; DROP TABLE agent_template_harness_pair; From eade96e25a5a64ffadd597f82b7c846a8aede886 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Wed, 9 Sep 2026 21:39:47 +0000 Subject: [PATCH 3/9] refactor: store fork ancestry directly on histories Signed-off-by: Eitan Yarmush --- .../persistence-checkpoints-and-forks.md | 24 +++++++++- docs/plans/api-v2-execution-plan.md | 2 +- go/core/internal/database/agent_instances.go | 16 +++++-- .../database/checkpoint_lineage_test.go | 47 ++++++++++++++----- go/core/internal/database/checkpoints.go | 46 ++++++++++-------- .../pkg/migrations/core/000001_initial.sql | 25 +++++----- 6 files changed, 109 insertions(+), 51 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index e1ed0bb6de..d93b6d4bc1 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -18,7 +18,7 @@ The core PostgreSQL records are: | `agent_template_harness_pair` | Pair status and latest successful revision | | `agent_instance` | Compute identity, pinned revision, lifecycle phase, and Actor identity | | `agent_instance_share` | Instance authorization grants | -| `a2a_context` | Durable history scope and its wire A2A context binding | +| `a2a_context` | Durable history scope, wire A2A context binding, and parent history/cutoff | | `agent_instance_task` | Rebuildable current A2A task state and query indexes | | `agent_instance_task_event` | Authoritative append-only task and message events, with creation and runtime-boundary metadata | | `agent_instance_checkpoint` | Named immutable snapshot/history boundary | @@ -48,6 +48,7 @@ flowchart TD CHECKPOINT --> TAG[Substrate snapshot tag] CHECKPOINT --> FORK[forked AgentInstance] FORK --> NEWCTX[new history scope, same wire context] + NEWCTX -->|parent history + cutoff| CONTEXT EVENT -->|copy through checkpoint cutoff| NEWCTX ``` @@ -79,10 +80,29 @@ covers the latest history event, including when an older paused task resumes. Th deleted while its context and checkpoint remain. Deletion first hides the checkpoint, then deletes its snapshot tag, then removes -the row. A checkpoint referenced by a fork cannot be deleted. Substrate deletes the Tag's copied snapshot with the Tag. +the row. A checkpoint inside any retained fork history's inherited prefix cannot +start deletion. Deleting instances keeps those histories, so inherited checkpoints +and their Tags remain protected even after all related instances are removed. +There is currently no history garbage collection. A deletion already in progress +can finish or retry if a later fork inherits its boundary. Substrate deletes the +Tag's copied snapshot with the Tag. ## Forking +History ancestry is recorded directly on `a2a_context`: `parent_history_id` and +`parent_history_sequence` identify the source history and the cutoff copied from +it. Roots have neither field. Fork creation sets both fields atomically with the +new instance and never changes them. Each new history points to an existing +parent; ancestry traversal needs no checkpoint rows. `agent_instance.source_checkpoint_id` +separately retains fork-request identity and runtime provenance while the instance exists. + +Checkpoint listing follows parent histories and their cutoffs, returning owned +ready checkpoints with their original source provenance. Membership depends on the +history boundary, so a checkpoint created later at an already inherited boundary +also appears. Listing uses checkpoint-ID pagination across local and inherited +results. After an instance is deleted, listing by its former ID returns only its +locally created checkpoints; surviving descendants still follow retained ancestry. + Forking creates a new AgentInstance authority and durable history scope. It preserves wire context, task, message, artifact IDs, and request deduplication metadata while copying events through the saved cutoff and reconstructing task diff --git a/docs/plans/api-v2-execution-plan.md b/docs/plans/api-v2-execution-plan.md index d409e0a867..41b6c85276 100644 --- a/docs/plans/api-v2-execution-plan.md +++ b/docs/plans/api-v2-execution-plan.md @@ -444,7 +444,7 @@ Implement `ForkAgentInstance`: - Create a new Actor identity from the checkpoint's retained snapshot tag. - Create a new AgentInstance, A2A authority, and creator ownership. - Keep source instance, history, and snapshots immutable. -- Represent inherited history through copy-on-write projections: deterministic fork-local Task IDs and the new context ID reference immutable source payloads and lineage without duplicating content. +- Copy events through the checkpoint cutoff and rebuild private task projections, preserving wire context and task IDs for paused-runtime continuity. Record ancestry with a direct parent history and cutoff; sharing event storage is a separate future optimization. - New Tasks append only to the fork. - Return ready only after the cloned Actor resumes and passes health. diff --git a/go/core/internal/database/agent_instances.go b/go/core/internal/database/agent_instances.go index 93a0ac52a2..30acf1a7f4 100644 --- a/go/core/internal/database/agent_instances.go +++ b/go/core/internal/database/agent_instances.go @@ -405,14 +405,24 @@ func readAgentInstance(ctx context.Context, db dbExecutor, id string) (agentInst // insertAgentInstanceRecords stores a new instance and its independent history together. // Callers supply a transaction and a prepared CREATING/CREATE instance; a duplicate // creator/requestID returns pgx.ErrNoRows so the caller can roll back and resolve it. -func insertAgentInstanceRecords(ctx context.Context, db dbExecutor, instance *apiv1alpha1.AgentInstance, requestID string, historyID uuid.UUID, sourceCheckpointID *uuid.UUID) (agentInstanceRow, error) { +// A fork records immutable parent history/cutoff metadata separately from the instance's +// source checkpoint, so ancestry does not depend on checkpoint rows or instance lifetime. +func insertAgentInstanceRecords(ctx context.Context, db dbExecutor, instance *apiv1alpha1.AgentInstance, requestID string, historyID uuid.UUID, source *agentInstanceCheckpointRow) (agentInstanceRow, error) { data, err := marshalAgentInstance(instance) if err != nil { return agentInstanceRow{}, err } + var sourceCheckpointID, parentHistoryID *uuid.UUID + var parentHistorySequence *int64 + if source != nil { + sourceCheckpointID = &source.ID + parentHistoryID = &source.SourceHistoryID + parentHistorySequence = &source.HistorySequence + } if err := execSQL(ctx, db, ` - INSERT INTO a2a_context (id, user_id, context_id, source_checkpoint_id) VALUES ($1, $2, $3, $4) - `, historyID, instance.Creator, instance.ContextId, sourceCheckpointID); err != nil { + INSERT INTO a2a_context (id, user_id, context_id, parent_history_id, parent_history_sequence) + VALUES ($1, $2, $3, $4, $5) + `, historyID, instance.Creator, instance.ContextId, parentHistoryID, parentHistorySequence); err != nil { return agentInstanceRow{}, fmt.Errorf("insert A2A context: %w", err) } return queryOne(ctx, db, ` diff --git a/go/core/internal/database/checkpoint_lineage_test.go b/go/core/internal/database/checkpoint_lineage_test.go index c342b93c2c..7ce311860c 100644 --- a/go/core/internal/database/checkpoint_lineage_test.go +++ b/go/core/internal/database/checkpoint_lineage_test.go @@ -63,6 +63,7 @@ func TestCheckpointLineage(t *testing.T) { _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") require.NoError(t, err) branch := fork(c2) + assertListed(branch.Id, c1, c2) // A fork created after deletion begins must not prevent its cleanup from retrying. _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") require.NoError(t, err) @@ -75,20 +76,38 @@ func TestCheckpointLineage(t *testing.T) { assertListed(branch.Id, c1, c2, c4, c5) assertListed(nested.Id, c1, c2, c4) - // Paging uses checkpoint IDs across every ancestor, without duplicates. - wantIDs := []string{c1.Id, c2.Id, c4.Id} - slices.Sort(wantIDs) - var afterID string - for _, id := range wantIDs { - page, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "alice", afterID, 1) - require.NoError(t, err) - require.Len(t, page, 1) - require.Equal(t, id, page[0].Id) - afterID = id + // Paging merges local and inherited checkpoints without repeating local candidates. + for _, tt := range []struct { + name string + instanceID string + wantIDs []string + }{ + {"root", source.Id, []string{c1.Id, c2.Id, c3.Id}}, + {"mixed", branch.Id, []string{c1.Id, c2.Id, c4.Id, c5.Id}}, + {"inherited", nested.Id, []string{c1.Id, c2.Id, c4.Id}}, + } { + t.Run(tt.name, func(t *testing.T) { + slices.Sort(tt.wantIDs) + for _, pageSize := range []int{1, 2, 10} { + var afterID string + var gotIDs []string + for { + page, err := client.ListAgentInstanceCheckpoints(ctx, tt.instanceID, "alice", afterID, pageSize) + require.NoError(t, err) + require.LessOrEqual(t, len(page), pageSize) + if len(page) == 0 { + break + } + for _, checkpoint := range page { + require.Greater(t, checkpoint.Id, afterID) + gotIDs = append(gotIDs, checkpoint.Id) + afterID = checkpoint.Id + } + } + require.Equal(t, tt.wantIDs, gotIDs) + } + }) } - page, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "alice", afterID, 1) - require.NoError(t, err) - require.Empty(t, page) unauthorized, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "mallory", "", 100) require.NoError(t, err) require.Empty(t, unauthorized) @@ -97,6 +116,8 @@ func TestCheckpointLineage(t *testing.T) { require.NoError(t, client.DeleteAgentInstance(ctx, source.Id)) require.NoError(t, client.DeleteAgentInstance(ctx, branch.Id)) assertListed(nested.Id, c1, c2, c4) + assertListed(source.Id, c1, c2, c3) + assertListed(branch.Id, c4, c5) // No fork starts at c1 yet; it is retained because it precedes c2 in the inherited history. _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, c1.Id, "alice") require.ErrorIs(t, err, ErrNotFound) diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 27a68a4c26..dcfe555a6a 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -82,7 +82,7 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re Operation: apiv1alpha1.AgentInstanceOperation_AGENT_INSTANCE_OPERATION_CREATE, CreatedAt: now, UpdatedAt: now, } - row, err = insertAgentInstanceRecords(ctx, tx, instance, requestID, historyID, &checkpoint.ID) + row, err = insertAgentInstanceRecords(ctx, tx, instance, requestID, historyID, &checkpoint) if err != nil { return err } @@ -399,28 +399,37 @@ func checkpointSnapshot(row agentInstanceCheckpointRow) *AgentInstanceTaskSnapsh // Checkpoint fields retain their source provenance. Locally created checkpoints remain // listable after the instance is deleted. func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, userID, afterID string, limit int) ([]*apiv1alpha1.Checkpoint, error) { + // Limit candidates per history before fetching payloads. Filtering the entire + // checkpoint table by lineage can otherwise turn a small page into a full scan. rows, err := queryMany(ctx, c.db, ` WITH RECURSIVE lineage (history_id, boundary) AS ( SELECT history_id, NULL::bigint FROM agent_instance WHERE id = $1 AND user_id = $2 UNION - SELECT c.source_history_id, c.history_sequence + SELECT h.parent_history_id, h.parent_history_sequence FROM lineage l JOIN a2a_context h ON h.id = l.history_id - JOIN agent_instance_checkpoint c ON c.id = h.source_checkpoint_id - WHERE h.user_id = $2 AND c.user_id = $2 + WHERE h.user_id = $2 AND h.parent_history_id IS NOT NULL + ), page AS MATERIALIZED ( + (SELECT id FROM agent_instance_checkpoint + WHERE source_instance_id = $1 AND user_id = $2 AND state = 'READY' + AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) + ORDER BY id LIMIT $4) + UNION + SELECT c.id FROM lineage l + CROSS JOIN LATERAL ( + SELECT id FROM agent_instance_checkpoint + WHERE source_history_id = l.history_id AND user_id = $2 AND state = 'READY' + AND (l.boundary IS NULL OR history_sequence <= l.boundary) + AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) + ORDER BY id LIMIT $4 + ) c + ORDER BY id LIMIT $4 ) - SELECT id, source_instance_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace, - snapshot_uri, snapshot_content_scope, tag_uid, state, data, source_history_id, prepared_revision - FROM agent_instance_checkpoint c - WHERE (source_instance_id = $1 OR EXISTS ( - SELECT 1 FROM lineage l WHERE l.history_id = c.source_history_id - AND (l.boundary IS NULL OR c.history_sequence <= l.boundary) - )) - AND user_id = $2 - AND state = 'READY' - AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) - ORDER BY id - LIMIT $4 + SELECT c.id, c.source_instance_id, c.user_id, c.request_id, c.head_task_id, c.history_sequence, + c.snapshot_atespace, c.snapshot_uri, c.snapshot_content_scope, c.tag_uid, c.state, c.data, + c.source_history_id, c.prepared_revision + FROM page p JOIN agent_instance_checkpoint c ON c.id = p.id + ORDER BY c.id `, pgx.RowToStructByName[agentInstanceCheckpointRow], instanceID, userID, afterID, int32(limit), ) @@ -473,9 +482,8 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use AND agent_instance_checkpoint.state IN ('READY', 'DELETING') AND (agent_instance_checkpoint.state = 'DELETING' OR NOT EXISTS ( SELECT 1 FROM a2a_context h - JOIN agent_instance_checkpoint parent ON parent.id = h.source_checkpoint_id - WHERE parent.source_history_id = agent_instance_checkpoint.source_history_id - AND parent.history_sequence >= agent_instance_checkpoint.history_sequence + WHERE h.parent_history_id = agent_instance_checkpoint.source_history_id + AND h.parent_history_sequence >= agent_instance_checkpoint.history_sequence )) `, row.ID, userID, data) if err != nil { diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index efd0826314..ea3649098b 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -66,15 +66,19 @@ CREATE INDEX agent_template_harness_pair_name_idx ON agent_template_harness_pair (namespace, agent_template_name, harness_name); CREATE TABLE a2a_context ( - id UUID PRIMARY KEY, - user_id TEXT NOT NULL CHECK (user_id <> ''), - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - context_id UUID NOT NULL, - source_checkpoint_id UUID, - CONSTRAINT a2a_context_binding_key UNIQUE (id, context_id) + id UUID PRIMARY KEY, + user_id TEXT NOT NULL CHECK (user_id <> ''), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + context_id UUID NOT NULL, + parent_history_id UUID REFERENCES a2a_context(id) ON DELETE RESTRICT, + parent_history_sequence BIGINT, + CONSTRAINT a2a_context_binding_key UNIQUE (id, context_id), + CHECK ((parent_history_id IS NULL) = (parent_history_sequence IS NULL)), + CHECK (parent_history_id <> id), + CHECK (parent_history_sequence > 0) ); -CREATE INDEX a2a_context_source_checkpoint_idx ON a2a_context (source_checkpoint_id) - WHERE source_checkpoint_id IS NOT NULL; +CREATE INDEX a2a_context_parent_history_idx ON a2a_context (parent_history_id, parent_history_sequence) + WHERE parent_history_id IS NOT NULL; CREATE TABLE agent_instance_checkpoint ( id UUID PRIMARY KEY, @@ -103,10 +107,6 @@ CREATE UNIQUE INDEX agent_instance_checkpoint_one_creating_idx ON agent_instance_checkpoint (source_instance_id) WHERE state = 'CREATING'; --- History and checkpoints reference each other, so this foreign key follows both tables. -ALTER TABLE a2a_context ADD CONSTRAINT a2a_context_source_checkpoint_id_fkey - FOREIGN KEY (source_checkpoint_id) REFERENCES agent_instance_checkpoint(id) ON DELETE RESTRICT; - CREATE TABLE agent_instance ( id UUID PRIMARY KEY, user_id TEXT NOT NULL CHECK (user_id <> ''), @@ -288,7 +288,6 @@ DROP TABLE agent_instance_share; DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; -ALTER TABLE a2a_context DROP CONSTRAINT a2a_context_source_checkpoint_id_fkey; DROP TABLE agent_instance_checkpoint; DROP TABLE a2a_context; DROP TABLE agent_template_harness_pair; From 03fc7d7223c9d2c2fed8a0d4c10ed3dee95a0cdb Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Thu, 10 Sep 2026 14:08:13 +0000 Subject: [PATCH 4/9] refactor: name durable history table agent_history Signed-off-by: Eitan Yarmush --- .../persistence-checkpoints-and-forks.md | 22 +++++++++---------- docs/plans/api-v2-execution-plan.md | 4 ++-- go/core/internal/database/agent_instances.go | 4 ++-- go/core/internal/database/checkpoints.go | 8 +++---- .../database/client_agent_instance_test.go | 10 ++++----- go/core/internal/database/client_test.go | 2 +- .../pkg/migrations/core/000001_initial.sql | 19 ++++++++-------- go/core/pkg/migrations/runner_test.go | 2 +- 8 files changed, 36 insertions(+), 35 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index d93b6d4bc1..913b37d917 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -2,12 +2,12 @@ ## Durable interaction model -`AgentInstance` represents ephemeral compute. An A2A context durably owns its +`AgentInstance` represents ephemeral compute. An `agent_history` row durably owns its tasks and ordered events, allowing interaction history to remain as an audit trail after compute is removed. New instances allocate independent instance, wire A2A context, and durable history IDs. `agent_instance.history_id` selects the history; `agent_instance.context_id` binds its public context. A composite foreign key -ensures that binding agrees with `a2a_context`. A history belongs to at most one +ensures that binding agrees with `agent_history`. A history belongs to at most one live instance, while multiple fork authorities may use the same wire context. The core PostgreSQL records are: @@ -18,7 +18,7 @@ The core PostgreSQL records are: | `agent_template_harness_pair` | Pair status and latest successful revision | | `agent_instance` | Compute identity, pinned revision, lifecycle phase, and Actor identity | | `agent_instance_share` | Instance authorization grants | -| `a2a_context` | Durable history scope, wire A2A context binding, and parent history/cutoff | +| `agent_history` | One durable history branch, its wire A2A context binding, and parent history/cutoff | | `agent_instance_task` | Rebuildable current A2A task state and query indexes | | `agent_instance_task_event` | Authoritative append-only task and message events, with creation and runtime-boundary metadata | | `agent_instance_checkpoint` | Named immutable snapshot/history boundary | @@ -40,16 +40,16 @@ without duplicating SQL. Transaction boundaries remain with the owning operation flowchart TD PAIR[Harness + AgentTemplate pair] --> REV[runtime revision] REV --> INSTANCE[AgentInstance] - INSTANCE --> CONTEXT[history scope + wire context] - CONTEXT --> EVENT[immutable ordered task events] + INSTANCE -->|history_id| HISTORY[agent_history] + HISTORY --> EVENT[immutable ordered task events] EVENT -->|replay| TASK[materialized task views] - CONTEXT --> CHECKPOINT[checkpoint boundary] + HISTORY --> CHECKPOINT[checkpoint boundary] REV --> CHECKPOINT CHECKPOINT --> TAG[Substrate snapshot tag] CHECKPOINT --> FORK[forked AgentInstance] - FORK --> NEWCTX[new history scope, same wire context] - NEWCTX -->|parent history + cutoff| CONTEXT - EVENT -->|copy through checkpoint cutoff| NEWCTX + FORK -->|history_id| NEWHISTORY[new agent_history, same wire context] + NEWHISTORY -->|parent history + cutoff| HISTORY + EVENT -->|copy through checkpoint cutoff| NEWHISTORY ``` ## Checkpoint creation @@ -77,7 +77,7 @@ cutoff in the same transaction as the runtime boundary reference. Later replies append events beyond that cutoff and cannot change the saved task state. The head identifies the task whose snapshot covers the latest history event, including when an older paused task resumes. The source AgentInstance may be -deleted while its context and checkpoint remain. +deleted while its history and checkpoint remain. Deletion first hides the checkpoint, then deletes its snapshot tag, then removes the row. A checkpoint inside any retained fork history's inherited prefix cannot @@ -89,7 +89,7 @@ Tag's copied snapshot with the Tag. ## Forking -History ancestry is recorded directly on `a2a_context`: `parent_history_id` and +History ancestry is recorded directly on `agent_history`: `parent_history_id` and `parent_history_sequence` identify the source history and the cutoff copied from it. Roots have neither field. Fork creation sets both fields atomically with the new instance and never changes them. Each new history points to an existing diff --git a/docs/plans/api-v2-execution-plan.md b/docs/plans/api-v2-execution-plan.md index 41b6c85276..2d2b4ec976 100644 --- a/docs/plans/api-v2-execution-plan.md +++ b/docs/plans/api-v2-execution-plan.md @@ -15,8 +15,8 @@ Public invariants: - `Harness` and `AgentTemplate` are `kagent.dev/v1alpha3` CRDs. - `AgentInstance` is a PostgreSQL-backed gRPC resource. -- One AgentInstance owns one rooted template tree and one A2A context. -- `A2A context_id == AgentInstance.id`. +- One AgentInstance owns one rooted template tree and one private history branch in `agent_history`. +- `AgentInstance.context_id` is the public A2A context ID; forks preserve it while using distinct instance and history IDs. - A2A owns interaction and history; AgentInstanceService owns catalog, lifecycle, metadata, and sharing. - Every quiescent A2A turn is durably suspended to an exact snapshot before its final state is published; physical Actor suspension does not change a ready diff --git a/go/core/internal/database/agent_instances.go b/go/core/internal/database/agent_instances.go index 30acf1a7f4..af3fab9667 100644 --- a/go/core/internal/database/agent_instances.go +++ b/go/core/internal/database/agent_instances.go @@ -420,10 +420,10 @@ func insertAgentInstanceRecords(ctx context.Context, db dbExecutor, instance *ap parentHistorySequence = &source.HistorySequence } if err := execSQL(ctx, db, ` - INSERT INTO a2a_context (id, user_id, context_id, parent_history_id, parent_history_sequence) + INSERT INTO agent_history (id, user_id, context_id, parent_history_id, parent_history_sequence) VALUES ($1, $2, $3, $4, $5) `, historyID, instance.Creator, instance.ContextId, parentHistoryID, parentHistorySequence); err != nil { - return agentInstanceRow{}, fmt.Errorf("insert A2A context: %w", err) + return agentInstanceRow{}, fmt.Errorf("insert agent history: %w", err) } return queryOne(ctx, db, ` INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, prepared_revision, diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index dcfe555a6a..b4c2fd1825 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -65,7 +65,7 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re // Serialize new lineage with deletion of any checkpoint in its inherited prefix. // NO KEY UPDATE leaves foreign-key checks free to read the history. sourceContextID, err := queryOne(ctx, tx, ` - SELECT context_id FROM a2a_context WHERE id = $1 FOR NO KEY UPDATE + SELECT context_id FROM agent_history WHERE id = $1 FOR NO KEY UPDATE `, pgx.RowTo[uuid.UUID], checkpoint.SourceHistoryID) if err != nil { return fmt.Errorf("get checkpoint context: %w", err) @@ -407,7 +407,7 @@ func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, u UNION SELECT h.parent_history_id, h.parent_history_sequence FROM lineage l - JOIN a2a_context h ON h.id = l.history_id + JOIN agent_history h ON h.id = l.history_id WHERE h.user_id = $2 AND h.parent_history_id IS NOT NULL ), page AS MATERIALIZED ( (SELECT id FROM agent_instance_checkpoint @@ -462,7 +462,7 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use } // Take the same history lock as ForkAgentInstance before checking its descendants. if _, err := queryOne(ctx, tx, ` - SELECT id FROM a2a_context WHERE id = $1 FOR NO KEY UPDATE + SELECT id FROM agent_history WHERE id = $1 FOR NO KEY UPDATE `, pgx.RowTo[uuid.UUID], row.SourceHistoryID); err != nil { return fmt.Errorf("lock checkpoint history: %w", err) } @@ -481,7 +481,7 @@ func (c *Client) BeginDeleteAgentInstanceCheckpoint(ctx context.Context, id, use WHERE agent_instance_checkpoint.id = $1 AND agent_instance_checkpoint.user_id = $2 AND agent_instance_checkpoint.state IN ('READY', 'DELETING') AND (agent_instance_checkpoint.state = 'DELETING' OR NOT EXISTS ( - SELECT 1 FROM a2a_context h + SELECT 1 FROM agent_history h WHERE h.parent_history_id = agent_instance_checkpoint.source_history_id AND h.parent_history_sequence >= agent_instance_checkpoint.history_sequence )) diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index bb74830838..3e24233e38 100644 --- a/go/core/internal/database/client_agent_instance_test.go +++ b/go/core/internal/database/client_agent_instance_test.go @@ -97,7 +97,7 @@ func TestAgentInstanceTasksAreDurableAndExclusive(t *testing.T) { db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) @@ -189,7 +189,7 @@ func TestConcurrentAgentInstanceMessageReplay(t *testing.T) { db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) @@ -239,7 +239,7 @@ func TestAgentInstanceReplyArchivesStatusMessageAtomically(t *testing.T) { ctx := context.Background() instanceID := "11111111-1111-4111-8111-111111111111" if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\\x00') `); err != nil { t.Fatal(err) @@ -289,7 +289,7 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) { t.Fatal(err) } if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ($1, 'alice', $1) + INSERT INTO agent_history (id, user_id, context_id) VALUES ($1, 'alice', $1) `, instanceID); err != nil { t.Fatal(err) } @@ -701,7 +701,7 @@ func TestInterruptActiveAgentInstanceTaskRequiresMatchingTaskAndReusesSlot(t *te db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) diff --git a/go/core/internal/database/client_test.go b/go/core/internal/database/client_test.go index 1cfbe4ec73..19020a6759 100644 --- a/go/core/internal/database/client_test.go +++ b/go/core/internal/database/client_test.go @@ -112,7 +112,7 @@ func setupTestDB(t *testing.T) *pgxpool.Pool { scheduled_run, tool, toolserver, memory, agent_instance_share, - agent_instance, a2a_context, agent_template_harness_pair, runtime_revision + agent_instance, agent_history, agent_template_harness_pair, runtime_revision RESTART IDENTITY CASCADE `) require.NoError(t, err, "Failed to truncate test tables") diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index ea3649098b..4a6d324023 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -65,19 +65,20 @@ CREATE TABLE agent_template_harness_pair ( CREATE INDEX agent_template_harness_pair_name_idx ON agent_template_harness_pair (namespace, agent_template_name, harness_name); -CREATE TABLE a2a_context ( +-- Each row owns one history branch; forks preserve the public A2A context_id. +CREATE TABLE agent_history ( id UUID PRIMARY KEY, user_id TEXT NOT NULL CHECK (user_id <> ''), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), context_id UUID NOT NULL, - parent_history_id UUID REFERENCES a2a_context(id) ON DELETE RESTRICT, + parent_history_id UUID REFERENCES agent_history(id) ON DELETE RESTRICT, parent_history_sequence BIGINT, - CONSTRAINT a2a_context_binding_key UNIQUE (id, context_id), + CONSTRAINT agent_history_binding_key UNIQUE (id, context_id), CHECK ((parent_history_id IS NULL) = (parent_history_sequence IS NULL)), CHECK (parent_history_id <> id), CHECK (parent_history_sequence > 0) ); -CREATE INDEX a2a_context_parent_history_idx ON a2a_context (parent_history_id, parent_history_sequence) +CREATE INDEX agent_history_parent_history_idx ON agent_history (parent_history_id, parent_history_sequence) WHERE parent_history_id IS NOT NULL; CREATE TABLE agent_instance_checkpoint ( @@ -93,7 +94,7 @@ CREATE TABLE agent_instance_checkpoint ( tag_uid TEXT NOT NULL DEFAULT '', state TEXT NOT NULL, data BYTEA NOT NULL, - source_history_id UUID NOT NULL REFERENCES a2a_context(id) ON DELETE RESTRICT, + source_history_id UUID NOT NULL REFERENCES agent_history(id) ON DELETE RESTRICT, prepared_revision TEXT REFERENCES runtime_revision(revision) ON DELETE RESTRICT, CHECK (snapshot_content_scope IN ('FULL', 'DATA')), CHECK (state IN ('CREATING', 'READY', 'FAILED', 'DELETING')), @@ -130,7 +131,7 @@ CREATE TABLE agent_instance ( (prepared_revision IS NULL AND operation = 'AGENT_INSTANCE_OPERATION_UNSPECIFIED')), history_id UUID NOT NULL, CONSTRAINT agent_instance_context_binding_fkey - FOREIGN KEY (history_id, context_id) REFERENCES a2a_context(id, context_id) ON DELETE RESTRICT, + FOREIGN KEY (history_id, context_id) REFERENCES agent_history(id, context_id) ON DELETE RESTRICT, CONSTRAINT agent_instance_history_key UNIQUE (history_id), CONSTRAINT agent_instance_operation_check CHECK (operation IN ('AGENT_INSTANCE_OPERATION_UNSPECIFIED', 'AGENT_INSTANCE_OPERATION_CREATE', @@ -154,7 +155,7 @@ CREATE INDEX agent_instance_share_instance_idx ON agent_instance_share (instance_id, id); CREATE TABLE agent_instance_task ( - history_id UUID CONSTRAINT agent_instance_task_instance_id_not_null NOT NULL REFERENCES a2a_context(id) ON DELETE CASCADE, + history_id UUID CONSTRAINT agent_instance_task_instance_id_not_null NOT NULL REFERENCES agent_history(id) ON DELETE CASCADE, id TEXT NOT NULL, state TEXT NOT NULL, status_timestamp TIMESTAMPTZ, @@ -188,7 +189,7 @@ CREATE UNIQUE INDEX agent_instance_task_message_idx CREATE TABLE agent_instance_task_event ( sequence BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, - history_id UUID CONSTRAINT agent_instance_task_event_instance_id_not_null NOT NULL REFERENCES a2a_context(id) ON DELETE CASCADE, + history_id UUID CONSTRAINT agent_instance_task_event_instance_id_not_null NOT NULL REFERENCES agent_history(id) ON DELETE CASCADE, task_id TEXT, data BYTEA NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), @@ -289,7 +290,7 @@ DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; DROP TABLE agent_instance_checkpoint; -DROP TABLE a2a_context; +DROP TABLE agent_history; DROP TABLE agent_template_harness_pair; DROP TABLE runtime_revision; DROP TABLE toolserver; diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 5a04c0f202..d436f66e85 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -211,7 +211,7 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { contextID := "00000000-0000-0000-0000-000000000001" instanceID := "00000000-0000-0000-0000-000000000002" historyID := "00000000-0000-0000-0000-000000000003" - execSQL(t, dsn, "INSERT INTO a2a_context (id, user_id, context_id) VALUES ($1, 'user', $2)", historyID, contextID) + execSQL(t, dsn, "INSERT INTO agent_history (id, user_id, context_id) VALUES ($1, 'user', $2)", historyID, contextID) execSQL(t, dsn, "INSERT INTO agent_instance (id, user_id, request_id, state, data, context_id, history_id) VALUES ($1, 'user', 'request', 'AGENT_INSTANCE_STATE_READY', $2, $3, $4)", instanceID, []byte{}, contextID, historyID) execSQL(t, dsn, "INSERT INTO agent_instance_task (history_id, id, state, data) VALUES ($1, 'task', 'TASK_STATE_INPUT_REQUIRED', $2)", historyID, []byte{}) scheduleID := "00000000-0000-0000-0000-000000000003" From ca4ace173ce0de70e55d3a462a6593dbba711eb2 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Thu, 10 Sep 2026 16:21:44 +0000 Subject: [PATCH 5/9] test: preserve checkpoint lifecycle coverage after rebase Retain the runtime-revision lifecycle checks previously carried by a merge commit. Align fork documentation with main requiring terminal task boundaries for checkpoints. Signed-off-by: Eitan Yarmush --- .../persistence-checkpoints-and-forks.md | 13 ++-- docs/plans/api-v2-execution-plan.md | 2 +- go/core/test/e2e/runtime_revision_test.go | 72 ++++++++++++------- 3 files changed, 55 insertions(+), 32 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 913b37d917..2452840a31 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -75,9 +75,9 @@ The checkpoint retains source-instance provenance, source history, prepared revision, name, head task, and history sequence. Reservation saves an event cutoff in the same transaction as the runtime boundary reference. Later replies append events beyond that cutoff and cannot change the saved task state. -The head identifies the task whose snapshot -covers the latest history event, including when an older paused task resumes. The source AgentInstance may be -deleted while its history and checkpoint remain. +The head identifies the task whose snapshot covers the latest history event, +including when an older paused task resumes and reaches a terminal state. The +source AgentInstance may be deleted while its history and checkpoint remain. Deletion first hides the checkpoint, then deletes its snapshot tag, then removes the row. A checkpoint inside any retained fork history's inherited prefix cannot @@ -107,10 +107,9 @@ Forking creates a new AgentInstance authority and durable history scope. It preserves wire context, task, message, artifact IDs, and request deduplication metadata while copying events through the saved cutoff and reconstructing task views from those events. It never reads the source's current task views. It creates a -separate Actor from the checkpoint's snapshot tag. Private runtime session IDs and -opaque paused-tool references therefore remain valid without runtime-specific -rewriting. New work appends only to the fork's history; source history and the -checkpoint remain immutable. The copied head boundary +separate Actor from the checkpoint's snapshot tag. Private runtime session IDs +remain valid without runtime-specific rewriting. New work appends only to the +fork's history; it does not mutate the source history or checkpoint. The copied head boundary uses the retained Tag URI, allowing a fresh fork to be checkpointed before its first turn. Fork creation verifies the Tag UID and URI as well as the Actor's source Tag, suspended state, template, and external snapshot. diff --git a/docs/plans/api-v2-execution-plan.md b/docs/plans/api-v2-execution-plan.md index 2d2b4ec976..47a7bc60b0 100644 --- a/docs/plans/api-v2-execution-plan.md +++ b/docs/plans/api-v2-execution-plan.md @@ -444,7 +444,7 @@ Implement `ForkAgentInstance`: - Create a new Actor identity from the checkpoint's retained snapshot tag. - Create a new AgentInstance, A2A authority, and creator ownership. - Keep source instance, history, and snapshots immutable. -- Copy events through the checkpoint cutoff and rebuild private task projections, preserving wire context and task IDs for paused-runtime continuity. Record ancestry with a direct parent history and cutoff; sharing event storage is a separate future optimization. +- Copy events through the checkpoint cutoff and rebuild private task projections, preserving wire context and task IDs for runtime continuity. Record ancestry with a direct parent history and cutoff; sharing event storage is a separate future optimization. - New Tasks append only to the fork. - Return ready only after the cloned Actor resumes and passes health. diff --git a/go/core/test/e2e/runtime_revision_test.go b/go/core/test/e2e/runtime_revision_test.go index a8b326166d..6091dca76a 100644 --- a/go/core/test/e2e/runtime_revision_test.go +++ b/go/core/test/e2e/runtime_revision_test.go @@ -29,6 +29,21 @@ import ( // invalid edits, template retirement, checkpoint retention, and later preparation. func TestRuntimeRevisionLifecycle(t *testing.T) { t.Parallel() + for _, tt := range []struct { + name string + fork bool + }{ + {"unforked_checkpoint_allows_collection", false}, + {"fork_history_retains_checkpoint", true}, + } { + t.Run(tt.name, func(t *testing.T) { + testRuntimeRevisionLifecycle(t, tt.fork) + }) + } +} + +func testRuntimeRevisionLifecycle(t *testing.T, fork bool) { + t.Helper() target := interactionTarget(t) modelURL := startInteractionMock(t) templateName := createInteractionTemplate(t, modelURL) @@ -147,33 +162,42 @@ func TestRuntimeRevisionLifecycle(t *testing.T) { deleteInstance(response.GetAgentInstance().GetId()) return false, nil })) - forked, err := checkpoints.ForkAgentInstance(ctx, &apiv1alpha1.ForkAgentInstanceRequest{CheckpointId: checkpointID, RequestId: uuid.NewString()}) - require.NoError(t, err, "a checkpoint must retain runnable inputs after its source and template are deleted") - forkID := forked.GetAgentInstance().GetId() - t.Cleanup(func() { deleteInstance(forkID) }) - send(forkID) - deleteInstance(forkID) + if fork { + forked, err := checkpoints.ForkAgentInstance(ctx, &apiv1alpha1.ForkAgentInstanceRequest{CheckpointId: checkpointID, RequestId: uuid.NewString()}) + require.NoError(t, err, "a checkpoint must retain runnable inputs after its source and template are deleted") + forkID := forked.GetAgentInstance().GetId() + t.Cleanup(func() { deleteInstance(forkID) }) + send(forkID) + deleteInstance(forkID) + } _, err = checkpoints.DeleteCheckpoint(ctx, &apiv1alpha1.DeleteCheckpointRequest{CheckpointId: checkpointID}) - require.NoError(t, err) - - // GC must remove the template after the final - // checkpoint disappears, without another template event to drive cleanup. - require.NoError(t, wait.PollUntilContextTimeout(ctx, time.Second, 3*time.Minute, true, func(ctx context.Context) (bool, error) { - backend, err := system.GetSubstrateSummary(ctx, &apiv1alpha1.GetSubstrateSummaryRequest{Namespace: "kagent", Atespace: runtimeNamespace}) - if err != nil { - return false, err - } - require.Empty(t, backend.GetAteApiError()) - for _, actorTemplate := range backend.GetActorTemplates() { - if actorTemplate.GetMetadata().GetAtespace() == runtimeNamespace && actorTemplate.GetMetadata().GetName() == runtimeName { - return false, nil + if fork { + require.Equal(t, codes.NotFound, status.Code(err), "retained fork history must protect the checkpoint after instance deletion") + retained, err := checkpoints.GetCheckpoint(ctx, &apiv1alpha1.GetCheckpointRequest{CheckpointId: checkpointID}) + require.NoError(t, err) + require.Equal(t, apiv1alpha1.CheckpointState_CHECKPOINT_STATE_READY, retained.GetCheckpoint().GetState()) + } else { + require.NoError(t, err) + // GC must remove both the template and golden actor after the final + // checkpoint disappears, without another template event to drive cleanup. + require.NoError(t, wait.PollUntilContextTimeout(ctx, time.Second, 3*time.Minute, true, func(ctx context.Context) (bool, error) { + backend, err := system.GetSubstrateSummary(ctx, &apiv1alpha1.GetSubstrateSummaryRequest{Namespace: "kagent", Atespace: runtimeNamespace}) + if err != nil { + return false, err } - } - actor, err := findSubstrateActor(ctx, system, "ate-golden", goldenActorID) - return actor == nil, err - }), "final checkpoint deletion must eventually collect its runtime without template changes") + require.Empty(t, backend.GetAteApiError()) + for _, actorTemplate := range backend.GetActorTemplates() { + if actorTemplate.GetMetadata().GetAtespace() == runtimeNamespace && actorTemplate.GetMetadata().GetName() == runtimeName { + return false, nil + } + } + actor, err := findSubstrateActor(ctx, system, "ate-golden", goldenActorID) + return actor == nil, err + }), "final checkpoint deletion must eventually collect its runtime without template changes") + } - // Recreating the name must prepare a new identity after collection. + // Recreating the name must prepare a new identity, including when the old + // revision remains pinned by retained fork history. replacement := template.DeepCopy() replacement.ObjectMeta = metav1.ObjectMeta{Namespace: template.Namespace, Name: template.Name, Labels: template.Labels} replacement.Spec.ModelConfig.Name = originalModelName From a28b825ad9cae4cd1070866801c433ca19ff6d2e Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Sun, 20 Sep 2026 13:24:55 +0000 Subject: [PATCH 6/9] fix: preserve checkpoint lineage after instance deletion Signed-off-by: Eitan Yarmush --- .../persistence-checkpoints-and-forks.md | 16 +- go/core/internal/database/agent_instances.go | 6 +- .../database/checkpoint_lineage_test.go | 294 ++++++++++++++---- go/core/internal/database/checkpoints.go | 9 +- .../database/client_agent_instance_test.go | 10 +- .../database/runtime_revisions_test.go | 14 +- .../pkg/migrations/core/000001_initial.sql | 24 +- go/core/pkg/migrations/runner_test.go | 2 +- go/core/test/e2e/interaction_test.go | 19 ++ 9 files changed, 307 insertions(+), 87 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 2452840a31..a852affddd 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -7,8 +7,10 @@ tasks and ordered events, allowing interaction history to remain as an audit tra after compute is removed. New instances allocate independent instance, wire A2A context, and durable history IDs. `agent_instance.history_id` selects the history; `agent_instance.context_id` binds its public context. A composite foreign key -ensures that binding agrees with `agent_history`. A history belongs to at most one -live instance, while multiple fork authorities may use the same wire context. +ensures the instance ID, owner, and context agree with `agent_history`. The history +retains its unique originating `instance_id` after compute is deleted, so that ID +still resolves the complete checkpoint lineage. Multiple fork authorities may use +the same wire context. The core PostgreSQL records are: @@ -92,16 +94,18 @@ Tag's copied snapshot with the Tag. History ancestry is recorded directly on `agent_history`: `parent_history_id` and `parent_history_sequence` identify the source history and the cutoff copied from it. Roots have neither field. Fork creation sets both fields atomically with the -new instance and never changes them. Each new history points to an existing -parent; ancestry traversal needs no checkpoint rows. `agent_instance.source_checkpoint_id` +new instance and never changes them. Foreign keys require the parent to share the +child's owner and context, and the cutoff to identify an event in that parent. +Each new history points to an existing parent; ancestry traversal needs no +checkpoint rows. `agent_instance.source_checkpoint_id` separately retains fork-request identity and runtime provenance while the instance exists. Checkpoint listing follows parent histories and their cutoffs, returning owned ready checkpoints with their original source provenance. Membership depends on the history boundary, so a checkpoint created later at an already inherited boundary also appears. Listing uses checkpoint-ID pagination across local and inherited -results. After an instance is deleted, listing by its former ID returns only its -locally created checkpoints; surviving descendants still follow retained ancestry. +results. Listing starts from the history's retained instance ID, so deleting an +instance does not change its local or inherited checkpoint results. Forking creates a new AgentInstance authority and durable history scope. It preserves wire context, task, message, artifact IDs, and request deduplication diff --git a/go/core/internal/database/agent_instances.go b/go/core/internal/database/agent_instances.go index af3fab9667..2743dcf7d1 100644 --- a/go/core/internal/database/agent_instances.go +++ b/go/core/internal/database/agent_instances.go @@ -420,9 +420,9 @@ func insertAgentInstanceRecords(ctx context.Context, db dbExecutor, instance *ap parentHistorySequence = &source.HistorySequence } if err := execSQL(ctx, db, ` - INSERT INTO agent_history (id, user_id, context_id, parent_history_id, parent_history_sequence) - VALUES ($1, $2, $3, $4, $5) - `, historyID, instance.Creator, instance.ContextId, parentHistoryID, parentHistorySequence); err != nil { + INSERT INTO agent_history (id, instance_id, user_id, context_id, parent_history_id, parent_history_sequence) + VALUES ($1, $2, $3, $4, $5, $6) + `, historyID, instance.Id, instance.Creator, instance.ContextId, parentHistoryID, parentHistorySequence); err != nil { return agentInstanceRow{}, fmt.Errorf("insert agent history: %w", err) } return queryOne(ctx, db, ` diff --git a/go/core/internal/database/checkpoint_lineage_test.go b/go/core/internal/database/checkpoint_lineage_test.go index 7ce311860c..c5f15d70d0 100644 --- a/go/core/internal/database/checkpoint_lineage_test.go +++ b/go/core/internal/database/checkpoint_lineage_test.go @@ -1,41 +1,222 @@ package database import ( + "context" + "fmt" + "math" "slices" + "sync" "testing" + "time" "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1" "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" ) -func TestCheckpointLineage(t *testing.T) { - client := NewClient(setupTestDB(t)) +func checkpointNextTurn(t *testing.T, client *Client, instance *apiv1alpha1.AgentInstance) *apiv1alpha1.Checkpoint { + t.Helper() ctx := t.Context() - agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") - source, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + task := newAgentInstanceTask(uuid.NewString(), uuid.NewString()) + task.ContextID = instance.ContextId + _, _, err := client.CreateAgentInstanceTask(ctx, instance.Id, []byte(task.ID), task) require.NoError(t, err) - source, err = markAgentInstanceReady(ctx, client, source.Id, "source.example") + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, client.StoreAgentInstanceTaskEvent(ctx, instance.Id, task, task, + &AgentInstanceTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/" + string(task.ID), ContentScope: "DATA"})) + checkpoint, _, err := client.ReserveAgentInstanceCheckpoint(ctx, + &apiv1alpha1.Checkpoint{Id: uuid.NewString(), AgentInstanceId: instance.Id}, "alice", uuid.NewString()) require.NoError(t, err) + checkpoint, err = client.FinalizeAgentInstanceCheckpoint(ctx, checkpoint.Id, "tag-"+checkpoint.Id, "s3://tags/"+checkpoint.Id, "") + require.NoError(t, err) + return checkpoint +} - checkpointNextTurn := func(instance *apiv1alpha1.AgentInstance) *apiv1alpha1.Checkpoint { +func TestAgentHistoryConstraints(t *testing.T) { + client := NewClient(setupTestDB(t)) + ctx := t.Context() + agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + newSource := func() (*apiv1alpha1.AgentInstance, agentInstanceCheckpointRow) { t.Helper() - task := newAgentInstanceTask(uuid.NewString(), uuid.NewString()) - task.ContextID = instance.ContextId - _, _, err := client.CreateAgentInstanceTask(ctx, instance.Id, []byte(task.ID), task) + instance, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) require.NoError(t, err) - task.Status.State = a2a.TaskStateCompleted - require.NoError(t, client.StoreAgentInstanceTaskEvent(ctx, instance.Id, task, task, - &AgentInstanceTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/" + string(task.ID), ContentScope: "DATA"})) - checkpoint, _, err := client.ReserveAgentInstanceCheckpoint(ctx, - &apiv1alpha1.Checkpoint{Id: uuid.NewString(), AgentInstanceId: instance.Id}, "alice", uuid.NewString()) + instance, err = markAgentInstanceReady(ctx, client, instance.Id, "source.example") require.NoError(t, err) - checkpoint, err = client.FinalizeAgentInstanceCheckpoint(ctx, checkpoint.Id, "tag-"+checkpoint.Id, "s3://tags/"+checkpoint.Id, "") + checkpoint := checkpointNextTurn(t, client, instance) + row, err := readCheckpoint(ctx, client.db, checkpoint.Id, "alice", nil) require.NoError(t, err) - return checkpoint + return instance, row + } + source, boundary := newSource() + _, otherBoundary := newSource() + for _, tt := range []struct { + name string + owner string + contextID string + sequence int64 + constraint string + }{ + {"different owner", "mallory", source.ContextId, boundary.HistorySequence, "agent_history_parent_binding_fkey"}, + {"different context", "alice", uuid.NewString(), boundary.HistorySequence, "agent_history_parent_binding_fkey"}, + {"missing event", "alice", source.ContextId, math.MaxInt64, "agent_history_parent_event_fkey"}, + {"another history event", "alice", source.ContextId, otherBoundary.HistorySequence, "agent_history_parent_event_fkey"}, + } { + t.Run(tt.name, func(t *testing.T) { + // Deliberately bypass the store to prove PostgreSQL rejects invalid ancestry. + _, err := client.db.Exec(ctx, ` + INSERT INTO agent_history (id, instance_id, user_id, context_id, parent_history_id, parent_history_sequence) + VALUES ($1, $2, $3, $4, $5, $6) + `, uuid.New(), uuid.New(), tt.owner, tt.contextID, boundary.SourceHistoryID, tt.sequence) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "23503", pgErr.Code) + require.Equal(t, tt.constraint, pgErr.ConstraintName) + }) + } + for _, column := range []string{"id", "user_id", "context_id"} { + t.Run("instance binding/"+column, func(t *testing.T) { + _, err := client.db.Exec(ctx, "UPDATE agent_instance SET "+column+" = $1 WHERE id = $2", uuid.NewString(), source.Id) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "23503", pgErr.Code) + require.Equal(t, "agent_instance_context_binding_fkey", pgErr.ConstraintName) + }) + } + require.NoError(t, client.DeleteAgentInstance(ctx, source.Id)) + _, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(source.Id, "assistant", "kagent", "replacement"), uuid.NewString()) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "23505", pgErr.Code, "a deleted instance ID must not be rebound to another history") +} + +func TestCheckpointDeletionSerializesWithFork(t *testing.T) { + for _, sameCheckpoint := range []bool{true, false} { + for _, forkFirst := range []bool{true, false} { + t.Run(fmt.Sprintf("same_checkpoint=%t/fork_first=%t", sameCheckpoint, forkFirst), func(t *testing.T) { + pool := setupTestDB(t) + client := NewClient(pool) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + source, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + require.NoError(t, err) + source, err = markAgentInstanceReady(ctx, client, source.Id, "source.example") + require.NoError(t, err) + deletedCheckpoint := checkpointNextTurn(t, client, source) + forkCheckpoint := deletedCheckpoint + if !sameCheckpoint { + forkCheckpoint = checkpointNextTurn(t, client, source) + } + // Pause the first operation while it holds the history lock. With different + // checkpoints, no checkpoint-row lock can serialize these operations for us. + barrier := &queryBarrier{ + query: "FROM agent_history WHERE id = $1 FOR NO KEY UPDATE", afterQuery: true, + reached: make(chan struct{}), resume: make(chan struct{}), + } + var resume sync.Once + defer resume.Do(func() { close(barrier.resume) }) + config := pool.Config() + config.ConnConfig.Tracer = barrier + blockedPool, err := pgxpool.NewWithConfig(ctx, config) + require.NoError(t, err) + t.Cleanup(blockedPool.Close) + forking, deleting := client, NewClient(blockedPool) + if forkFirst { + forking, deleting = deleting, forking + } + forkID := uuid.NewString() + forked, deleted := make(chan error, 1), make(chan error, 1) + fork := func() { + _, _, err := forking.ForkAgentInstance(ctx, forkCheckpoint.Id, "alice", uuid.NewString(), forkID) + forked <- err + } + deleteCheckpoint := func() { + _, _, err := deleting.BeginDeleteAgentInstanceCheckpoint(ctx, deletedCheckpoint.Id, "alice") + deleted <- err + } + if forkFirst { + go fork() + } else { + go deleteCheckpoint() + } + select { + case <-barrier.reached: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + if forkFirst { + go deleteCheckpoint() + } else { + go fork() + } + waitingQuery := "%FROM agent_history WHERE id = $1 FOR NO KEY UPDATE%" + if sameCheckpoint { + waitingQuery = "%FROM agent_instance_checkpoint%FOR UPDATE%" + } + require.Eventually(t, func() bool { + var waiting bool + err := pool.QueryRow(ctx, `SELECT EXISTS ( + SELECT 1 FROM pg_stat_activity WHERE datname = current_database() + AND query LIKE $1 AND cardinality(pg_blocking_pids(pid)) > 0 + )`, waitingQuery).Scan(&waiting) + return err == nil && waiting + }, 5*time.Second, 10*time.Millisecond) + resume.Do(func() { close(barrier.resume) }) + forkErr, deleteErr := <-forked, <-deleted + if forkFirst { + require.NoError(t, forkErr) + require.ErrorIs(t, deleteErr, ErrNotFound) + checkpoint, err := client.GetAgentInstanceCheckpoint(ctx, deletedCheckpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1alpha1.CheckpointState_CHECKPOINT_STATE_READY, checkpoint.State) + listed, err := client.ListAgentInstanceCheckpoints(ctx, forkID, "alice", "", 100) + require.NoError(t, err) + require.Contains(t, listedCheckpointIDs(listed), deletedCheckpoint.Id) + } else { + require.NoError(t, deleteErr) + if sameCheckpoint { + require.ErrorIs(t, forkErr, ErrNotFound) + _, err := client.GetAgentInstance(ctx, forkID, "alice") + require.ErrorIs(t, err, ErrNotFound) + } else { + require.NoError(t, forkErr) + listed, err := client.ListAgentInstanceCheckpoints(ctx, forkID, "alice", "", 100) + require.NoError(t, err) + require.Equal(t, []string{forkCheckpoint.Id}, listedCheckpointIDs(listed)) + } + // Starting deletion first permits cleanup retries even after a later fork. + _, _, err := client.BeginDeleteAgentInstanceCheckpoint(ctx, deletedCheckpoint.Id, "alice") + require.NoError(t, err) + require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, deletedCheckpoint.Id, "alice")) + _, err = client.GetAgentInstanceCheckpoint(ctx, deletedCheckpoint.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + } + }) + } } +} + +func listedCheckpointIDs(checkpoints []*apiv1alpha1.Checkpoint) []string { + ids := make([]string, len(checkpoints)) + for i, checkpoint := range checkpoints { + ids[i] = checkpoint.Id + } + return ids +} + +func TestCheckpointLineage(t *testing.T) { + client := NewClient(setupTestDB(t)) + ctx := t.Context() + agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + source, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + require.NoError(t, err) + source, err = markAgentInstanceReady(ctx, client, source.Id, "source.example") + require.NoError(t, err) + fork := func(checkpoint *apiv1alpha1.Checkpoint) *apiv1alpha1.AgentInstance { t.Helper() instance, created, err := client.ForkAgentInstance(ctx, checkpoint.Id, "alice", uuid.NewString(), uuid.NewString()) @@ -57,9 +238,9 @@ func TestCheckpointLineage(t *testing.T) { } } - deleting := checkpointNextTurn(source) - c1 := checkpointNextTurn(source) - c2 := checkpointNextTurn(source) + deleting := checkpointNextTurn(t, client, source) + c1 := checkpointNextTurn(t, client, source) + c2 := checkpointNextTurn(t, client, source) _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") require.NoError(t, err) branch := fork(c2) @@ -68,46 +249,50 @@ func TestCheckpointLineage(t *testing.T) { _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice") require.NoError(t, err) require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice")) - c3 := checkpointNextTurn(source) - c4 := checkpointNextTurn(branch) + c3 := checkpointNextTurn(t, client, source) + c4 := checkpointNextTurn(t, client, branch) nested := fork(c4) - c5 := checkpointNextTurn(branch) + c5 := checkpointNextTurn(t, client, branch) assertListed(source.Id, c1, c2, c3) assertListed(branch.Id, c1, c2, c4, c5) assertListed(nested.Id, c1, c2, c4) - // Paging merges local and inherited checkpoints without repeating local candidates. - for _, tt := range []struct { - name string - instanceID string - wantIDs []string - }{ - {"root", source.Id, []string{c1.Id, c2.Id, c3.Id}}, - {"mixed", branch.Id, []string{c1.Id, c2.Id, c4.Id, c5.Id}}, - {"inherited", nested.Id, []string{c1.Id, c2.Id, c4.Id}}, - } { - t.Run(tt.name, func(t *testing.T) { - slices.Sort(tt.wantIDs) - for _, pageSize := range []int{1, 2, 10} { - var afterID string - var gotIDs []string - for { - page, err := client.ListAgentInstanceCheckpoints(ctx, tt.instanceID, "alice", afterID, pageSize) - require.NoError(t, err) - require.LessOrEqual(t, len(page), pageSize) - if len(page) == 0 { - break - } - for _, checkpoint := range page { - require.Greater(t, checkpoint.Id, afterID) - gotIDs = append(gotIDs, checkpoint.Id) - afterID = checkpoint.Id + assertPages := func(t *testing.T) { + t.Helper() + // Paging merges local and inherited checkpoints without repeating local candidates. + for _, tt := range []struct { + name string + instanceID string + wantIDs []string + }{ + {"root", source.Id, []string{c1.Id, c2.Id, c3.Id}}, + {"mixed", branch.Id, []string{c1.Id, c2.Id, c4.Id, c5.Id}}, + {"inherited", nested.Id, []string{c1.Id, c2.Id, c4.Id}}, + } { + t.Run(tt.name, func(t *testing.T) { + slices.Sort(tt.wantIDs) + for _, pageSize := range []int{1, 2, 10} { + var afterID string + var gotIDs []string + for { + page, err := client.ListAgentInstanceCheckpoints(ctx, tt.instanceID, "alice", afterID, pageSize) + require.NoError(t, err) + require.LessOrEqual(t, len(page), pageSize) + if len(page) == 0 { + break + } + for _, checkpoint := range page { + require.Greater(t, checkpoint.Id, afterID) + gotIDs = append(gotIDs, checkpoint.Id) + afterID = checkpoint.Id + } } + require.Equal(t, tt.wantIDs, gotIDs) } - require.Equal(t, tt.wantIDs, gotIDs) - } - }) + }) + } } + t.Run("before_deletion", assertPages) unauthorized, err := client.ListAgentInstanceCheckpoints(ctx, nested.Id, "mallory", "", 100) require.NoError(t, err) require.Empty(t, unauthorized) @@ -117,7 +302,11 @@ func TestCheckpointLineage(t *testing.T) { require.NoError(t, client.DeleteAgentInstance(ctx, branch.Id)) assertListed(nested.Id, c1, c2, c4) assertListed(source.Id, c1, c2, c3) - assertListed(branch.Id, c4, c5) + assertListed(branch.Id, c1, c2, c4, c5) + t.Run("after_deletion", assertPages) + unauthorized, err = client.ListAgentInstanceCheckpoints(ctx, branch.Id, "mallory", "", 100) + require.NoError(t, err) + require.Empty(t, unauthorized) // No fork starts at c1 yet; it is retained because it precedes c2 in the inherited history. _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, c1.Id, "alice") require.ErrorIs(t, err, ErrNotFound) @@ -126,6 +315,7 @@ func TestCheckpointLineage(t *testing.T) { assertListed(earlier.Id, c1) require.NoError(t, client.DeleteAgentInstance(ctx, earlier.Id)) require.NoError(t, client.DeleteAgentInstance(ctx, nested.Id)) + assertListed(nested.Id, c1, c2, c4) // Retained histories protect the entire inherited prefix, even after instance deletion. for _, checkpoint := range []*apiv1alpha1.Checkpoint{c1, c2, c4} { diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index b4c2fd1825..b50336e2ea 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -396,25 +396,20 @@ func checkpointSnapshot(row agentInstanceCheckpointRow) *AgentInstanceTaskSnapsh // ListAgentInstanceCheckpoints returns an owner's READY checkpoints in the instance's // history and inherited prefixes, in ascending ID order after afterID, up to limit. -// Checkpoint fields retain their source provenance. Locally created checkpoints remain +// Checkpoint fields retain their source provenance. The complete lineage remains // listable after the instance is deleted. func (c *Client) ListAgentInstanceCheckpoints(ctx context.Context, instanceID, userID, afterID string, limit int) ([]*apiv1alpha1.Checkpoint, error) { // Limit candidates per history before fetching payloads. Filtering the entire // checkpoint table by lineage can otherwise turn a small page into a full scan. rows, err := queryMany(ctx, c.db, ` WITH RECURSIVE lineage (history_id, boundary) AS ( - SELECT history_id, NULL::bigint FROM agent_instance WHERE id = $1 AND user_id = $2 + SELECT id, NULL::bigint FROM agent_history WHERE instance_id = $1 AND user_id = $2 UNION SELECT h.parent_history_id, h.parent_history_sequence FROM lineage l JOIN agent_history h ON h.id = l.history_id WHERE h.user_id = $2 AND h.parent_history_id IS NOT NULL ), page AS MATERIALIZED ( - (SELECT id FROM agent_instance_checkpoint - WHERE source_instance_id = $1 AND user_id = $2 AND state = 'READY' - AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) - ORDER BY id LIMIT $4) - UNION SELECT c.id FROM lineage l CROSS JOIN LATERAL ( SELECT id FROM agent_instance_checkpoint diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index 3e24233e38..97a04d2d5c 100644 --- a/go/core/internal/database/client_agent_instance_test.go +++ b/go/core/internal/database/client_agent_instance_test.go @@ -97,7 +97,7 @@ func TestAgentInstanceTasksAreDurableAndExclusive(t *testing.T) { db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) @@ -189,7 +189,7 @@ func TestConcurrentAgentInstanceMessageReplay(t *testing.T) { db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) @@ -239,7 +239,7 @@ func TestAgentInstanceReplyArchivesStatusMessageAtomically(t *testing.T) { ctx := context.Background() instanceID := "11111111-1111-4111-8111-111111111111" if _, err := db.Exec(ctx, ` - INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\\x00') `); err != nil { t.Fatal(err) @@ -289,7 +289,7 @@ func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) { t.Fatal(err) } if _, err := db.Exec(ctx, ` - INSERT INTO agent_history (id, user_id, context_id) VALUES ($1, 'alice', $1) + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ($1, $1, 'alice', $1) `, instanceID); err != nil { t.Fatal(err) } @@ -701,7 +701,7 @@ func TestInterruptActiveAgentInstanceTaskRequiresMatchingTaskAndReusesSlot(t *te db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` - INSERT INTO agent_history (id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ('11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'alice', '11111111-1111-4111-8111-111111111111'); INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, state, data) VALUES ('11111111-1111-4111-8111-111111111111', 'alice', 'request-1', '11111111-1111-4111-8111-111111111111', '11111111-1111-4111-8111-111111111111', 'AGENT_INSTANCE_STATE_READY', '\x') `); err != nil { t.Fatal(err) diff --git a/go/core/internal/database/runtime_revisions_test.go b/go/core/internal/database/runtime_revisions_test.go index 9b128721c4..bb3d6ddefb 100644 --- a/go/core/internal/database/runtime_revisions_test.go +++ b/go/core/internal/database/runtime_revisions_test.go @@ -176,16 +176,16 @@ func TestRuntimeRevisionPairReplacement(t *testing.T) { require.NoError(t, client.DeleteRuntimeRevision(ctx, "old", "old-actor-uid")) } -// Pause the real store before or after it locks a revision. Exercise both -// commit orderings without replacing persistence queries with test-only writes. -type runtimeReferenceBarrier struct { +// Pause a real store query before or after execution to control commit ordering. +// Tests still exercise production persistence queries and PostgreSQL locks. +type queryBarrier struct { query string afterQuery bool reached chan struct{} resume chan struct{} } -func (b *runtimeReferenceBarrier) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { +func (b *queryBarrier) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { if strings.Contains(data.SQL, b.query) { if !b.afterQuery { close(b.reached) @@ -196,7 +196,7 @@ func (b *runtimeReferenceBarrier) TraceQueryStart(ctx context.Context, _ *pgx.Co return ctx } -func (b *runtimeReferenceBarrier) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryEndData) { +func (b *queryBarrier) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryEndData) { if b.afterQuery && ctx.Value(b) != nil { close(b.reached) <-b.resume @@ -217,7 +217,7 @@ func TestRuntimeRevisionDeletionSerializesWithReferenceAcquisition(t *testing.T) query = "FOR UPDATE OF r" require.NoError(t, client.RetirePairIdentities(ctx, "team-a", "assistant", "kagent", nil)) } - barrier := &runtimeReferenceBarrier{query: query, afterQuery: referenceFirst, reached: make(chan struct{}), resume: make(chan struct{})} + barrier := &queryBarrier{query: query, afterQuery: referenceFirst, reached: make(chan struct{}), resume: make(chan struct{})} var resume sync.Once defer resume.Do(func() { close(barrier.resume) }) config := pool.Config() @@ -355,7 +355,7 @@ func TestRuntimeRevisionFinalizationSerializesWithPairWrites(t *testing.T) { require.NoError(t, err) require.NotNil(t, claimed) - barrier := &runtimeReferenceBarrier{ + barrier := &queryBarrier{ query: "FOR UPDATE OF r", reached: make(chan struct{}), resume: make(chan struct{}), } if operation == "record" { diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index 4a6d324023..99542a1a17 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -68,12 +68,18 @@ CREATE INDEX agent_template_harness_pair_name_idx -- Each row owns one history branch; forks preserve the public A2A context_id. CREATE TABLE agent_history ( id UUID PRIMARY KEY, + -- Historical provenance, retained after the instance is deleted. + instance_id UUID NOT NULL UNIQUE, user_id TEXT NOT NULL CHECK (user_id <> ''), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), context_id UUID NOT NULL, - parent_history_id UUID REFERENCES agent_history(id) ON DELETE RESTRICT, + parent_history_id UUID, parent_history_sequence BIGINT, - CONSTRAINT agent_history_binding_key UNIQUE (id, context_id), + CONSTRAINT agent_history_binding_key UNIQUE (id, user_id, context_id), + CONSTRAINT agent_history_instance_binding_key UNIQUE (id, instance_id, user_id, context_id), + CONSTRAINT agent_history_parent_binding_fkey + FOREIGN KEY (parent_history_id, user_id, context_id) + REFERENCES agent_history(id, user_id, context_id) ON DELETE RESTRICT, CHECK ((parent_history_id IS NULL) = (parent_history_sequence IS NULL)), CHECK (parent_history_id <> id), CHECK (parent_history_sequence > 0) @@ -131,7 +137,8 @@ CREATE TABLE agent_instance ( (prepared_revision IS NULL AND operation = 'AGENT_INSTANCE_OPERATION_UNSPECIFIED')), history_id UUID NOT NULL, CONSTRAINT agent_instance_context_binding_fkey - FOREIGN KEY (history_id, context_id) REFERENCES agent_history(id, context_id) ON DELETE RESTRICT, + FOREIGN KEY (history_id, id, user_id, context_id) + REFERENCES agent_history(id, instance_id, user_id, context_id) ON DELETE RESTRICT, CONSTRAINT agent_instance_history_key UNIQUE (history_id), CONSTRAINT agent_instance_operation_check CHECK (operation IN ('AGENT_INSTANCE_OPERATION_UNSPECIFIED', 'AGENT_INSTANCE_OPERATION_CREATE', @@ -211,12 +218,17 @@ CREATE UNIQUE INDEX agent_instance_task_event_creation_idx ON agent_instance_task_event (history_id, task_id) WHERE task_position IS NOT NULL; CREATE UNIQUE INDEX agent_instance_task_event_position_idx ON agent_instance_task_event (history_id, task_position) WHERE task_position IS NOT NULL; -CREATE INDEX agent_instance_task_event_instance_sequence_idx +CREATE UNIQUE INDEX agent_instance_task_event_instance_sequence_idx ON agent_instance_task_event (history_id, sequence); CREATE UNIQUE INDEX agent_instance_task_event_message_idx ON agent_instance_task_event (history_id, task_id, message_id) WHERE message_id IS NOT NULL; +-- Histories own events, so the reverse boundary reference must follow both tables. +ALTER TABLE agent_history ADD CONSTRAINT agent_history_parent_event_fkey + FOREIGN KEY (parent_history_id, parent_history_sequence) + REFERENCES agent_instance_task_event(history_id, sequence) ON DELETE RESTRICT; + CREATE TABLE scheduled_run ( id UUID PRIMARY KEY, creator TEXT NOT NULL CHECK (creator <> ''), @@ -286,11 +298,11 @@ DROP TABLE scheduled_run; DROP VIEW unreferenced_runtime_revision; DROP TABLE agent_instance_share; -DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; DROP TABLE agent_instance_checkpoint; -DROP TABLE agent_history; +-- Drop both ends of the history/event references together. +DROP TABLE agent_instance_task_event, agent_history; DROP TABLE agent_template_harness_pair; DROP TABLE runtime_revision; DROP TABLE toolserver; diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index d436f66e85..bc622d7daa 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -211,7 +211,7 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { contextID := "00000000-0000-0000-0000-000000000001" instanceID := "00000000-0000-0000-0000-000000000002" historyID := "00000000-0000-0000-0000-000000000003" - execSQL(t, dsn, "INSERT INTO agent_history (id, user_id, context_id) VALUES ($1, 'user', $2)", historyID, contextID) + execSQL(t, dsn, "INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ($1, $2, 'user', $3)", historyID, instanceID, contextID) execSQL(t, dsn, "INSERT INTO agent_instance (id, user_id, request_id, state, data, context_id, history_id) VALUES ($1, 'user', 'request', 'AGENT_INSTANCE_STATE_READY', $2, $3, $4)", instanceID, []byte{}, contextID, historyID) execSQL(t, dsn, "INSERT INTO agent_instance_task (history_id, id, state, data) VALUES ($1, 'task', 'TASK_STATE_INPUT_REQUIRED', $2)", historyID, []byte{}) scheduleID := "00000000-0000-0000-0000-000000000003" diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index 7e7b408ee1..fe1fb4df94 100644 --- a/go/core/test/e2e/interaction_test.go +++ b/go/core/test/e2e/interaction_test.go @@ -305,6 +305,25 @@ func TestAgentInstanceCheckpoint(t *testing.T) { if _, err := fixture.instances.DeleteAgentInstance(nestedCtx, &apiv1alpha1.DeleteAgentInstanceRequest{AgentInstanceId: nestedID}); err != nil { t.Fatalf("delete nested fork: %v", err) } + // Both deleted instance IDs still resolve their complete checkpoint lineage, + // including a nested fork that never created a checkpoint of its own. + for _, instanceID := range []string{fork.GetId(), nestedID} { + remaining := map[string]bool{checkpoint.GetId(): true, forkCheckpoint.GetCheckpoint().GetId(): true} + var pageToken string + for page := range 2 { + listed, err := fixture.checkpoints.ListCheckpoints(fixture.ctx, &apiv1alpha1.ListCheckpointsRequest{ + AgentInstanceId: instanceID, Page: &apiv1alpha1.PageRequest{Limit: 1, PageToken: pageToken}, + }) + if err != nil || len(listed.GetCheckpoints()) != 1 || !remaining[listed.GetCheckpoints()[0].GetId()] { + t.Fatalf("deleted instance %s checkpoint page = %+v, error %v", instanceID, listed, err) + } + delete(remaining, listed.GetCheckpoints()[0].GetId()) + pageToken = listed.GetPage().GetNextPageToken() + if (pageToken == "") != (page == 1) { + t.Fatalf("deleted instance %s page %d token = %q", instanceID, page, pageToken) + } + } + } // Deleting instances retains their histories and the checkpoints they inherit. for _, id := range []string{checkpoint.GetId(), forkCheckpoint.GetCheckpoint().GetId()} { if _, err := fixture.checkpoints.DeleteCheckpoint(fixture.ctx, &apiv1alpha1.DeleteCheckpointRequest{ From c99bddd57afd7bd2fa2dd485601d81e8f6613ca3 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Mon, 21 Sep 2026 12:29:29 +0000 Subject: [PATCH 7/9] refactor: validate fork cutoffs in the store transaction Signed-off-by: Eitan Yarmush --- .../persistence-checkpoints-and-forks.md | 6 +- .../database/checkpoint_lineage_test.go | 104 +++++++++++++----- go/core/internal/database/checkpoints.go | 3 + .../pkg/migrations/core/000001_initial.sql | 12 +- 4 files changed, 91 insertions(+), 34 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index a852affddd..096d1fae2e 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -95,7 +95,11 @@ History ancestry is recorded directly on `agent_history`: `parent_history_id` an `parent_history_sequence` identify the source history and the cutoff copied from it. Roots have neither field. Fork creation sets both fields atomically with the new instance and never changes them. Foreign keys require the parent to share the -child's owner and context, and the cutoff to identify an event in that parent. +child's owner and context. The fork transaction reads only parent-history events +and verifies the exact cutoff and its task/snapshot boundary before committing; +failure rolls back the new history and instance. There is no reverse foreign key +from history to an event. Events remain retained with their history; future event +garbage collection must preserve inherited boundaries and coordinate with forks. Each new history points to an existing parent; ancestry traversal needs no checkpoint rows. `agent_instance.source_checkpoint_id` separately retains fork-request identity and runtime provenance while the instance exists. diff --git a/go/core/internal/database/checkpoint_lineage_test.go b/go/core/internal/database/checkpoint_lineage_test.go index c5f15d70d0..c6adff91e4 100644 --- a/go/core/internal/database/checkpoint_lineage_test.go +++ b/go/core/internal/database/checkpoint_lineage_test.go @@ -40,41 +40,31 @@ func TestAgentHistoryConstraints(t *testing.T) { client := NewClient(setupTestDB(t)) ctx := t.Context() agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") - newSource := func() (*apiv1alpha1.AgentInstance, agentInstanceCheckpointRow) { - t.Helper() - instance, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) - require.NoError(t, err) - instance, err = markAgentInstanceReady(ctx, client, instance.Id, "source.example") - require.NoError(t, err) - checkpoint := checkpointNextTurn(t, client, instance) - row, err := readCheckpoint(ctx, client.db, checkpoint.Id, "alice", nil) - require.NoError(t, err) - return instance, row - } - source, boundary := newSource() - _, otherBoundary := newSource() + source, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + require.NoError(t, err) + source, err = markAgentInstanceReady(ctx, client, source.Id, "source.example") + require.NoError(t, err) + checkpoint := checkpointNextTurn(t, client, source) + boundary, err := readCheckpoint(ctx, client.db, checkpoint.Id, "alice", nil) + require.NoError(t, err) for _, tt := range []struct { - name string - owner string - contextID string - sequence int64 - constraint string + name string + owner string + contextID string }{ - {"different owner", "mallory", source.ContextId, boundary.HistorySequence, "agent_history_parent_binding_fkey"}, - {"different context", "alice", uuid.NewString(), boundary.HistorySequence, "agent_history_parent_binding_fkey"}, - {"missing event", "alice", source.ContextId, math.MaxInt64, "agent_history_parent_event_fkey"}, - {"another history event", "alice", source.ContextId, otherBoundary.HistorySequence, "agent_history_parent_event_fkey"}, + {"different owner", "mallory", source.ContextId}, + {"different context", "alice", uuid.NewString()}, } { t.Run(tt.name, func(t *testing.T) { // Deliberately bypass the store to prove PostgreSQL rejects invalid ancestry. _, err := client.db.Exec(ctx, ` INSERT INTO agent_history (id, instance_id, user_id, context_id, parent_history_id, parent_history_sequence) VALUES ($1, $2, $3, $4, $5, $6) - `, uuid.New(), uuid.New(), tt.owner, tt.contextID, boundary.SourceHistoryID, tt.sequence) + `, uuid.New(), uuid.New(), tt.owner, tt.contextID, boundary.SourceHistoryID, boundary.HistorySequence) var pgErr *pgconn.PgError require.ErrorAs(t, err, &pgErr) require.Equal(t, "23503", pgErr.Code) - require.Equal(t, tt.constraint, pgErr.ConstraintName) + require.Equal(t, "agent_history_parent_binding_fkey", pgErr.ConstraintName) }) } for _, column := range []string{"id", "user_id", "context_id"} { @@ -87,12 +77,76 @@ func TestAgentHistoryConstraints(t *testing.T) { }) } require.NoError(t, client.DeleteAgentInstance(ctx, source.Id)) - _, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(source.Id, "assistant", "kagent", "replacement"), uuid.NewString()) + _, _, err = client.CreateAgentInstance(ctx, newAgentInstanceRequest(source.Id, "assistant", "kagent", "replacement"), uuid.NewString()) var pgErr *pgconn.PgError require.ErrorAs(t, err, &pgErr) require.Equal(t, "23505", pgErr.Code, "a deleted instance ID must not be rebound to another history") } +func TestForkAgentInstanceRejectsInvalidBoundary(t *testing.T) { + client := NewClient(setupTestDB(t)) + ctx := t.Context() + agentInstanceFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + newCheckpoint := func() *apiv1alpha1.Checkpoint { + t.Helper() + instance, _, err := client.CreateAgentInstance(ctx, newAgentInstanceRequest(uuid.NewString(), "assistant", "kagent", "source"), uuid.NewString()) + require.NoError(t, err) + instance, err = markAgentInstanceReady(ctx, client, instance.Id, "source.example") + require.NoError(t, err) + return checkpointNextTurn(t, client, instance) + } + source, other := newCheckpoint(), newCheckpoint() + for _, tt := range []struct { + name string + sequence uint64 + headTask string + wantErr string + }{ + {"missing event", math.MaxInt64, source.HeadTaskId, "checkpoint history boundary is missing"}, + {"another history event", other.HistorySequence, other.HeadTaskId, "checkpoint history boundary is missing"}, + {"wrong head task", source.HistorySequence, uuid.NewString(), "checkpoint runtime boundary is inconsistent"}, + } { + t.Run(tt.name, func(t *testing.T) { + // Corrupt both the payload and indexes so the test reaches event validation, + // rather than failing the independent protobuf/index consistency check. + invalid := proto.Clone(source).(*apiv1alpha1.Checkpoint) + invalid.HistorySequence, invalid.HeadTaskId = tt.sequence, tt.headTask + data, err := proto.Marshal(invalid) + require.NoError(t, err) + _, err = client.db.Exec(ctx, ` + UPDATE agent_instance_checkpoint SET history_sequence = $2, head_task_id = $3, data = $4 WHERE id = $1 + `, source.Id, int64(invalid.HistorySequence), invalid.HeadTaskId, data) + require.NoError(t, err) + forkID, requestID := uuid.NewString(), uuid.NewString() + fork, created, err := client.ForkAgentInstance(ctx, source.Id, "alice", requestID, forkID) + require.ErrorContains(t, err, tt.wantErr) + require.Nil(t, fork) + require.False(t, created) + _, err = client.GetAgentInstance(ctx, forkID, "alice") + require.ErrorIs(t, err, ErrNotFound) + var historyExists bool + require.NoError(t, client.db.QueryRow(ctx, "SELECT EXISTS (SELECT 1 FROM agent_history WHERE instance_id = $1)", forkID).Scan(&historyExists)) + require.False(t, historyExists, "failed validation must roll back the durable lineage") + + // Repair the boundary and retry the same request/instance IDs: no failed + // reservation or history may prevent the valid fork from being created. + data, err = proto.Marshal(source) + require.NoError(t, err) + _, err = client.db.Exec(ctx, ` + UPDATE agent_instance_checkpoint SET history_sequence = $2, head_task_id = $3, data = $4 WHERE id = $1 + `, source.Id, int64(source.HistorySequence), source.HeadTaskId, data) + require.NoError(t, err) + fork, created, err = client.ForkAgentInstance(ctx, source.Id, "alice", requestID, forkID) + require.NoError(t, err) + require.True(t, created) + require.Equal(t, forkID, fork.Id) + listed, err := client.ListAgentInstanceCheckpoints(ctx, forkID, "alice", "", 100) + require.NoError(t, err) + require.Equal(t, []string{source.Id}, listedCheckpointIDs(listed)) + }) + } +} + func TestCheckpointDeletionSerializesWithFork(t *testing.T) { for _, sameCheckpoint := range []bool{true, false} { for _, forkFirst := range []bool{true, false} { diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index b50336e2ea..d8f9aa9274 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -91,6 +91,9 @@ func (c *Client) ForkAgentInstance(ctx context.Context, checkpointID, userID, re if err != nil { return fmt.Errorf("list checkpoint events: %w", err) } + // The query is scoped to the parent history. Validate the exact cutoff here + // rather than creating a reverse history-to-event foreign key. A failure + // rolls back the new instance and history along with the rest of the fork. if len(events) == 0 || events[len(events)-1].Sequence != checkpoint.HistorySequence { return fmt.Errorf("checkpoint history boundary is missing") } diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index 99542a1a17..55713cb87e 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -74,6 +74,7 @@ CREATE TABLE agent_history ( created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), context_id UUID NOT NULL, parent_history_id UUID, + -- ForkAgentInstance validates the cutoff against the parent's events in its transaction. parent_history_sequence BIGINT, CONSTRAINT agent_history_binding_key UNIQUE (id, user_id, context_id), CONSTRAINT agent_history_instance_binding_key UNIQUE (id, instance_id, user_id, context_id), @@ -218,17 +219,12 @@ CREATE UNIQUE INDEX agent_instance_task_event_creation_idx ON agent_instance_task_event (history_id, task_id) WHERE task_position IS NOT NULL; CREATE UNIQUE INDEX agent_instance_task_event_position_idx ON agent_instance_task_event (history_id, task_position) WHERE task_position IS NOT NULL; -CREATE UNIQUE INDEX agent_instance_task_event_instance_sequence_idx +CREATE INDEX agent_instance_task_event_instance_sequence_idx ON agent_instance_task_event (history_id, sequence); CREATE UNIQUE INDEX agent_instance_task_event_message_idx ON agent_instance_task_event (history_id, task_id, message_id) WHERE message_id IS NOT NULL; --- Histories own events, so the reverse boundary reference must follow both tables. -ALTER TABLE agent_history ADD CONSTRAINT agent_history_parent_event_fkey - FOREIGN KEY (parent_history_id, parent_history_sequence) - REFERENCES agent_instance_task_event(history_id, sequence) ON DELETE RESTRICT; - CREATE TABLE scheduled_run ( id UUID PRIMARY KEY, creator TEXT NOT NULL CHECK (creator <> ''), @@ -298,11 +294,11 @@ DROP TABLE scheduled_run; DROP VIEW unreferenced_runtime_revision; DROP TABLE agent_instance_share; +DROP TABLE agent_instance_task_event; DROP TABLE agent_instance_task; DROP TABLE agent_instance; DROP TABLE agent_instance_checkpoint; --- Drop both ends of the history/event references together. -DROP TABLE agent_instance_task_event, agent_history; +DROP TABLE agent_history; DROP TABLE agent_template_harness_pair; DROP TABLE runtime_revision; DROP TABLE toolserver; From b39513ea6464bc9a85cd4c350d21fe79d434b643 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Mon, 21 Sep 2026 12:34:48 +0000 Subject: [PATCH 8/9] test: align checkpoint lineage with instance tombstones Signed-off-by: Eitan Yarmush --- docs/architecture/persistence-checkpoints-and-forks.md | 8 ++++++-- go/core/internal/database/client_agent_instance_test.go | 9 ++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 096d1fae2e..9f3585f992 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -12,6 +12,10 @@ retains its unique originating `instance_id` after compute is deleted, so that I still resolves the complete checkpoint lineage. Multiple fork authorities may use the same wire context. +Deleting an instance leaves a hidden tombstone for its owner and creation-request +identity, while releasing its runtime references and shares. A retry of that +creation request cannot recreate the deleted instance. + The core PostgreSQL records are: | Record | Purpose | @@ -101,8 +105,8 @@ failure rolls back the new history and instance. There is no reverse foreign key from history to an event. Events remain retained with their history; future event garbage collection must preserve inherited boundaries and coordinate with forks. Each new history points to an existing parent; ancestry traversal needs no -checkpoint rows. `agent_instance.source_checkpoint_id` -separately retains fork-request identity and runtime provenance while the instance exists. +checkpoint rows. `agent_instance.source_checkpoint_id` separately retains +fork-request identity and runtime provenance, including on deletion tombstones. Checkpoint listing follows parent histories and their cutoffs, returning owned ready checkpoints with their original source provenance. Membership depends on the diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index 97a04d2d5c..05115dc006 100644 --- a/go/core/internal/database/client_agent_instance_test.go +++ b/go/core/internal/database/client_agent_instance_test.go @@ -424,7 +424,7 @@ func TestReserveAgentInstanceCheckpointRejectsCorruptSource(t *testing.T) { client := NewClient(db) instanceID := "99999999-9999-4999-8999-999999999999" if _, err := db.Exec(ctx, ` - INSERT INTO a2a_context (id, user_id, context_id) VALUES ($1, 'alice', $1) + INSERT INTO agent_history (id, instance_id, user_id, context_id) VALUES ($1, $1, 'alice', $1) `, instanceID); err != nil { t.Fatal(err) } @@ -596,11 +596,10 @@ func TestForkAgentInstanceCopiesBoundedHistory(t *testing.T) { if err != nil || !created || fork2.GetId() != fork2ID { t.Fatalf("fork of fork = %+v, created %v, error %v", fork2, created, err) } - // Deleting a fork releases its checkpoint pin but preserves its original - // request identity, even after that checkpoint itself has been removed. + // Deleting a fork preserves its original request identity. Its retained + // history still protects the inherited checkpoint after compute is removed. _, _, err = client.BeginDeleteAgentInstanceCheckpoint(ctx, checkpoint.GetId(), "alice") - require.NoError(t, err) - require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, checkpoint.GetId(), "alice")) + require.ErrorIs(t, err, ErrNotFound) _, _, err = client.ForkAgentInstance(ctx, checkpoint.GetId(), "alice", "fork-request-1", uuid.NewString()) require.ErrorIs(t, err, ErrFailedPrecondition) _, _, err = client.ForkAgentInstance(ctx, checkpoint2.GetId(), "alice", "fork-request-1", uuid.NewString()) From 5c703e05d009a990b30a42fff2bc3104fdfbe96b Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Tue, 22 Sep 2026 10:30:31 +0000 Subject: [PATCH 9/9] fix: keep inherited checkpoint actions scoped to their origin Signed-off-by: Eitan Yarmush --- .../gen/kagent/api/v1alpha1/checkpoints.pb.go | 14 ++-- .../api/v1alpha1/checkpoints_grpc.pb.go | 6 ++ .../pkg/migrations/core/000001_initial.sql | 2 - proto/kagent/api/v1alpha1/checkpoints.proto | 7 ++ ui/playwright/tests/chat/checkpoints.spec.ts | 47 +++++++++-- ui/src/api/domain/checkpoints.ts | 2 +- ui/src/components/chat/ChatTranscript.tsx | 16 ++-- .../components/chat/SnapshotDetailsModal.tsx | 79 ++++++++++--------- .../kagent/api/v1alpha1/checkpoints_pb.ts | 10 +++ ui/src/mocks/checkpointLineage.test.ts | 43 ++++++++++ ui/src/mocks/state.ts | 33 +++++++- ui/src/mocks/transport.ts | 6 +- 12 files changed, 204 insertions(+), 61 deletions(-) create mode 100644 ui/src/mocks/checkpointLineage.test.ts diff --git a/go/api/gen/kagent/api/v1alpha1/checkpoints.pb.go b/go/api/gen/kagent/api/v1alpha1/checkpoints.pb.go index 0125e4e954..3cce694f46 100644 --- a/go/api/gen/kagent/api/v1alpha1/checkpoints.pb.go +++ b/go/api/gen/kagent/api/v1alpha1/checkpoints.pb.go @@ -79,8 +79,10 @@ func (CheckpointState) EnumDescriptor() ([]byte, []int) { } type Checkpoint struct { - state protoimpl.MessageState `protogen:"open.v1"` - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + // The instance that created this checkpoint. For an inherited checkpoint this + // differs from the instance requested in ListCheckpoints. AgentInstanceId string `protobuf:"bytes,2,opt,name=agent_instance_id,json=agentInstanceId,proto3" json:"agent_instance_id,omitempty"` HeadTaskId string `protobuf:"bytes,3,opt,name=head_task_id,json=headTaskId,proto3" json:"head_task_id,omitempty"` HistorySequence uint64 `protobuf:"varint,4,opt,name=history_sequence,json=historySequence,proto3" json:"history_sequence,omitempty"` @@ -366,9 +368,11 @@ func (x *GetCheckpointResponse) GetCheckpoint() *Checkpoint { } type ListCheckpointsRequest struct { - state protoimpl.MessageState `protogen:"open.v1"` - AgentInstanceId string `protobuf:"bytes,1,opt,name=agent_instance_id,json=agentInstanceId,proto3" json:"agent_instance_id,omitempty"` - Page *PageRequest `protobuf:"bytes,2,opt,name=page,proto3" json:"page,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + // The instance whose local and inherited checkpoints to list. Its retained + // history remains listable after the instance is deleted. + AgentInstanceId string `protobuf:"bytes,1,opt,name=agent_instance_id,json=agentInstanceId,proto3" json:"agent_instance_id,omitempty"` + Page *PageRequest `protobuf:"bytes,2,opt,name=page,proto3" json:"page,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } diff --git a/go/api/gen/kagent/api/v1alpha1/checkpoints_grpc.pb.go b/go/api/gen/kagent/api/v1alpha1/checkpoints_grpc.pb.go index 0592c1d240..32f9723bed 100644 --- a/go/api/gen/kagent/api/v1alpha1/checkpoints_grpc.pb.go +++ b/go/api/gen/kagent/api/v1alpha1/checkpoints_grpc.pb.go @@ -33,6 +33,9 @@ const ( type CheckpointServiceClient interface { CreateCheckpoint(ctx context.Context, in *CreateCheckpointRequest, opts ...grpc.CallOption) (*CreateCheckpointResponse, error) GetCheckpoint(ctx context.Context, in *GetCheckpointRequest, opts ...grpc.CallOption) (*GetCheckpointResponse, error) + // Lists ready checkpoints from the instance's history and inherited prefixes. + // Checkpoints retain their originating instance IDs, including after deletion + // of the originating or listed instance. ListCheckpoints(ctx context.Context, in *ListCheckpointsRequest, opts ...grpc.CallOption) (*ListCheckpointsResponse, error) DeleteCheckpoint(ctx context.Context, in *DeleteCheckpointRequest, opts ...grpc.CallOption) (*DeleteCheckpointResponse, error) ForkAgentInstance(ctx context.Context, in *ForkAgentInstanceRequest, opts ...grpc.CallOption) (*ForkAgentInstanceResponse, error) @@ -113,6 +116,9 @@ func (c *checkpointServiceClient) UpdateCheckpointName(ctx context.Context, in * type CheckpointServiceServer interface { CreateCheckpoint(context.Context, *CreateCheckpointRequest) (*CreateCheckpointResponse, error) GetCheckpoint(context.Context, *GetCheckpointRequest) (*GetCheckpointResponse, error) + // Lists ready checkpoints from the instance's history and inherited prefixes. + // Checkpoints retain their originating instance IDs, including after deletion + // of the originating or listed instance. ListCheckpoints(context.Context, *ListCheckpointsRequest) (*ListCheckpointsResponse, error) DeleteCheckpoint(context.Context, *DeleteCheckpointRequest) (*DeleteCheckpointResponse, error) ForkAgentInstance(context.Context, *ForkAgentInstanceRequest) (*ForkAgentInstanceResponse, error) diff --git a/go/core/pkg/migrations/core/000001_initial.sql b/go/core/pkg/migrations/core/000001_initial.sql index 55713cb87e..dd884e5ed2 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -107,8 +107,6 @@ CREATE TABLE agent_instance_checkpoint ( CHECK (state IN ('CREATING', 'READY', 'FAILED', 'DELETING')), UNIQUE (user_id, request_id) ); -CREATE INDEX agent_instance_checkpoint_list_idx - ON agent_instance_checkpoint (source_instance_id, id); CREATE INDEX agent_instance_checkpoint_history_idx ON agent_instance_checkpoint (source_history_id, history_sequence); CREATE UNIQUE INDEX agent_instance_checkpoint_one_creating_idx diff --git a/proto/kagent/api/v1alpha1/checkpoints.proto b/proto/kagent/api/v1alpha1/checkpoints.proto index cd5e575b52..f9da16aaac 100644 --- a/proto/kagent/api/v1alpha1/checkpoints.proto +++ b/proto/kagent/api/v1alpha1/checkpoints.proto @@ -12,6 +12,9 @@ option go_package = "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1 service CheckpointService { rpc CreateCheckpoint(CreateCheckpointRequest) returns (CreateCheckpointResponse); rpc GetCheckpoint(GetCheckpointRequest) returns (GetCheckpointResponse); + // Lists ready checkpoints from the instance's history and inherited prefixes. + // Checkpoints retain their originating instance IDs, including after deletion + // of the originating or listed instance. rpc ListCheckpoints(ListCheckpointsRequest) returns (ListCheckpointsResponse); rpc DeleteCheckpoint(DeleteCheckpointRequest) returns (DeleteCheckpointResponse); rpc ForkAgentInstance(ForkAgentInstanceRequest) returns (ForkAgentInstanceResponse); @@ -28,6 +31,8 @@ enum CheckpointState { message Checkpoint { string id = 1; + // The instance that created this checkpoint. For an inherited checkpoint this + // differs from the instance requested in ListCheckpoints. string agent_instance_id = 2; string head_task_id = 3; uint64 history_sequence = 4; @@ -61,6 +66,8 @@ message GetCheckpointResponse { } message ListCheckpointsRequest { + // The instance whose local and inherited checkpoints to list. Its retained + // history remains listable after the instance is deleted. string agent_instance_id = 1 [(buf.validate.field).string.uuid = true]; PageRequest page = 2; } diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index abe633e838..8eb2dbefa4 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -1,6 +1,7 @@ import { test, expect } from "../../fixtures/test"; import { agentChat, instances } from "../../helpers/app"; import { pressOnce } from "../../helpers/resource"; +import { sendAndAwaitTurn } from "../../helpers/chat"; /** The menu that is actually on screen: antd leaves a closed dropdown mounted. */ const openMenu = (page: import("@playwright/test").Page) => @@ -128,8 +129,8 @@ test("chat: the mark names itself, carries its controls, and opens its record", * * Three claims about one action, so one action proves all three. Forking the *seeded* * boundary once a second one exists leaves the later turn behind — which is the whole - * point of forking a boundary rather than a conversation — the copy arrives with no - * lines on it at all, and it is titled after the snapshot rather than after its + * point of forking a boundary rather than a conversation — the copy retains only + * inherited checkpoint lines, and it is titled after the snapshot rather than after its * source. That last one is the controller's job now: the page passes no name, so a * fork called anything else means the name never reached `ForkAgentInstance`. * @@ -137,9 +138,9 @@ test("chat: the mark names itself, carries its controls, and opens its record", * remembered against the message it was taken at, because the reader's newest message * has no turn id yet. A fork is handed copies of its source's messages under the same * ids — so without dropping that memory when the conversation changes, a fork opened - * from here drew a line it does not have. + * from here could draw a line beyond the inherited cutoff. */ -test("chat: a fork holds only what was above its line, takes the snapshot's name, and inherits none of its marks", async ({ +test("chat: a fork inherits its checkpoint prefix with view and fork actions only", async ({ page, }) => { const SNAPSHOT_NAME = "Before the second question"; @@ -183,8 +184,42 @@ test("chat: a fork holds only what was above its line, takes the snapshot's name await expect(page.getByTestId("chat-sessions")).toContainText(SNAPSHOT_NAME); // Only what was above the line: the second turn is below it. await expect(mine).toHaveCount(1, { timeout: 30_000 }); - // And none of the source page's marks came with it. - await expect(dividers(page)).toHaveCount(0); + await expect(dividers(page)).toHaveCount(1); + const inherited = dividers(page).first(); + await expect(inherited.getByTestId("chat-checkpoint-label")).toHaveText(`Snapshot “${SNAPSHOT_NAME}”`); + await expect(inherited.locator('[data-testid^="chat-checkpoint-fork-"]')).toBeEnabled(); + await expect(inherited.locator('[data-testid^="chat-checkpoint-rename-"]')).toHaveCount(0); + await expect(inherited.locator('[data-testid^="chat-checkpoint-delete-"]')).toHaveCount(0); + + await test.step("a checkpoint created on the fork still offers rename and delete", async () => { + await sendAndAwaitTurn(page, "A new turn on this fork."); + await expect(mine).toHaveCount(2); + await expect(page.getByTestId("chat-checkpoint")).toBeEnabled(); + await page.getByTestId("chat-checkpoint").click(); + await expect(dividers(page)).toHaveCount(2); + const local = dividers(page).last(); + await expect(local.locator('[data-testid^="chat-checkpoint-rename-"]')).toBeEnabled(); + await expect(local.locator('[data-testid^="chat-checkpoint-delete-"]')).toBeEnabled(); + await openSnapshot(page, local); + await expect(page.getByTestId("snapshot-details-rename")).toBeEnabled(); + await expect(page.getByTestId("snapshot-details-delete")).toBeEnabled(); + await page.keyboard.press("Escape"); + await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); + }); + + await test.step("inherited details allow viewing and another fork, but no mutations", async () => { + await openSnapshot(page, inherited); + await expect(page.getByTestId("snapshot-details-name")).toHaveText(SNAPSHOT_NAME); + await expect(page.getByTestId("snapshot-details-rename")).toHaveCount(0); + await expect(page.getByTestId("snapshot-details-delete")).toHaveCount(0); + const firstForkURL = page.url(); + await page.getByTestId("snapshot-details-fork").click(); + await expect(page).not.toHaveURL(firstForkURL); + await expect(rows).toHaveCount(before + 2); + // The new fork must load the inherited prefix and leave our local turn behind. + await expect(mine).toHaveCount(1); + await expect(dividers(page)).toHaveCount(1); + }); }); /* diff --git a/ui/src/api/domain/checkpoints.ts b/ui/src/api/domain/checkpoints.ts index ec3c8f37de..538f8ab706 100644 --- a/ui/src/api/domain/checkpoints.ts +++ b/ui/src/api/domain/checkpoints.ts @@ -29,6 +29,7 @@ export type CheckpointState = export interface Checkpoint { id: string; + /** Originating instance; inherited checkpoints keep this ID when listed in a fork. */ agentInstanceId: string; /** * What the reader calls it, and what a fork taken from it is named. @@ -50,4 +51,3 @@ export interface Checkpoint { export function canForkFrom(checkpoint: Checkpoint): boolean { return checkpoint.state === "ready"; } - diff --git a/ui/src/components/chat/ChatTranscript.tsx b/ui/src/components/chat/ChatTranscript.tsx index 52b37b59ec..56ca69dfbd 100644 --- a/ui/src/components/chat/ChatTranscript.tsx +++ b/ui/src/components/chat/ChatTranscript.tsx @@ -88,9 +88,7 @@ export function ChatTranscript({ /** * The conversation being shown, handed to each message. * - * Only used by contributions at the per-message point: a message id means something - * to this client, while anything asking a backend about the work behind a message is - * keyed by the conversation and the turn. + * Also determines which checkpoints were created here and may expose rename/delete. */ sessionId?: string; }) { @@ -354,6 +352,8 @@ export function ChatTranscript({ )); const checkpointId = group.checkpointId; if (checkpointId) { + const checkpoint = checkpointsById?.get(checkpointId); + const isLocal = Boolean(checkpoint && checkpoint.agentInstanceId === sessionId); drawn.push(
onOpenCheckpoint(checkpointId))} onFork={onFork && (() => onFork(checkpointId))} - onRename={onRenameCheckpoint && (() => onRenameCheckpoint(checkpointId))} - onDelete={onDeleteCheckpoint && (() => onDeleteCheckpoint(checkpointId))} + onRename={ + isLocal && onRenameCheckpoint ? () => onRenameCheckpoint(checkpointId) : undefined + } + onDelete={ + isLocal && onDeleteCheckpoint ? () => onDeleteCheckpoint(checkpointId) : undefined + } />
, ); diff --git a/ui/src/components/chat/SnapshotDetailsModal.tsx b/ui/src/components/chat/SnapshotDetailsModal.tsx index 249fbe39f5..1a7c3a6739 100644 --- a/ui/src/components/chat/SnapshotDetailsModal.tsx +++ b/ui/src/components/chat/SnapshotDetailsModal.tsx @@ -49,7 +49,7 @@ export function SnapshotDetailsModal({ onRenamed, }: { checkpoint: Checkpoint; - /** The conversation it was taken in, which the extension point is keyed by. */ + /** The conversation being viewed; inherited checkpoints originate elsewhere. */ instanceId: string; onClose: () => void; /** Opens a new conversation holding the transcript up to here. */ @@ -69,6 +69,7 @@ export function SnapshotDetailsModal({ const [isBusy, setBusy] = useState(false); const label = snapshotLabel(checkpoint); const forkable = canForkFrom(checkpoint); + const isLocal = checkpoint.agentInstanceId === instanceId; const state = STATE_APPEARANCE[checkpoint.state]; async function act(run: () => void | Promise) { @@ -106,16 +107,18 @@ export function SnapshotDetailsModal({ children: ( {label} - - - + + + ) : null} {/* Part of the record rather than something floating under it, which is what @@ -269,7 +274,7 @@ export function SnapshotDetailsModal({ {/* Mounted only while open, which is what seeds the box with the stored name without an effect to put it there — see `RenameConversationDialog`. */} - {isRenaming ? ( + {isLocal && isRenaming ? ( setRenaming(false)} diff --git a/ui/src/generated/kagent/api/v1alpha1/checkpoints_pb.ts b/ui/src/generated/kagent/api/v1alpha1/checkpoints_pb.ts index 2ee222ce2a..dc804c93d1 100644 --- a/ui/src/generated/kagent/api/v1alpha1/checkpoints_pb.ts +++ b/ui/src/generated/kagent/api/v1alpha1/checkpoints_pb.ts @@ -29,6 +29,9 @@ export type Checkpoint = Message<"kagent.api.v1alpha1.Checkpoint"> & { id: string; /** + * The instance that created this checkpoint. For an inherited checkpoint this + * differs from the instance requested in ListCheckpoints. + * * @generated from field: string agent_instance_id = 2; */ agentInstanceId: string; @@ -153,6 +156,9 @@ export const GetCheckpointResponseSchema: GenMessage = /* */ export type ListCheckpointsRequest = Message<"kagent.api.v1alpha1.ListCheckpointsRequest"> & { /** + * The instance whose local and inherited checkpoints to list. Its retained + * history remains listable after the instance is deleted. + * * @generated from field: string agent_instance_id = 1; */ agentInstanceId: string; @@ -360,6 +366,10 @@ export const CheckpointService: GenService<{ output: typeof GetCheckpointResponseSchema; }, /** + * Lists ready checkpoints from the instance's history and inherited prefixes. + * Checkpoints retain their originating instance IDs, including after deletion + * of the originating or listed instance. + * * @generated from rpc kagent.api.v1alpha1.CheckpointService.ListCheckpoints */ listCheckpoints: { diff --git a/ui/src/mocks/checkpointLineage.test.ts b/ui/src/mocks/checkpointLineage.test.ts new file mode 100644 index 0000000000..577ebb0087 --- /dev/null +++ b/ui/src/mocks/checkpointLineage.test.ts @@ -0,0 +1,43 @@ +import { beforeEach, expect, test } from "vitest"; +import { + deleteCheckpoint, + readCheckpoints, + recordCheckpointFork, + renameCheckpoint, + saveCheckpoint, + SEEDED_CHECKPOINT, +} from "./state"; + +beforeEach(() => window.sessionStorage.clear()); + +test("forks list their inherited prefix with original provenance, including later saves at that boundary", () => { + const source = SEEDED_CHECKPOINT; + saveCheckpoint({ ...source, id: "after-cutoff", headTaskId: "later-task" }); + recordCheckpointFork(source.agentInstanceId, "fork", [source.headTaskId]); + expect(readCheckpoints("fork")).toEqual([source]); + + const laterSave = saveCheckpoint({ ...source, id: "saved-after-fork" }); + const local = saveCheckpoint({ ...source, id: "fork-local", agentInstanceId: "fork", headTaskId: "fork-task" }); + recordCheckpointFork("fork", "nested", [source.headTaskId, local.headTaskId]); + expect(readCheckpoints("nested")).toEqual( + expect.arrayContaining([source, laterSave, local]), + ); + expect(readCheckpoints("nested")).toHaveLength(3); + expect(readCheckpoints("unrelated")).toEqual([]); + + renameCheckpoint(source.id, "Renamed at the source"); + expect(readCheckpoints("nested").find((row) => row.id === source.id)).toMatchObject({ + name: "Renamed at the source", agentInstanceId: source.agentInstanceId, + }); +}); + +test("retained prefixes block deletion but checkpoints after the cutoff remain deletable", () => { + const source = SEEDED_CHECKPOINT; + recordCheckpointFork(source.agentInstanceId, "fork", [source.headTaskId]); + expect(deleteCheckpoint(source.id)).toBe(false); + expect(readCheckpoints("fork")).toEqual([source]); + + const later = saveCheckpoint({ ...source, id: "after-cutoff", headTaskId: "later-task" }); + expect(deleteCheckpoint(later.id)).toBe(true); + expect(readCheckpoints(source.agentInstanceId)).toEqual([source]); +}); diff --git a/ui/src/mocks/state.ts b/ui/src/mocks/state.ts index b187b5f584..0466e6f6c8 100644 --- a/ui/src/mocks/state.ts +++ b/ui/src/mocks/state.ts @@ -520,9 +520,30 @@ function writeAll(rows: MockCheckpoint[]): void { } } -/** Every boundary saved against one conversation. */ +// Mock turns are completed, immutable boundaries. Retain the parent's task prefix +// alongside each fork, just as its transcript is copied through the selected turn. +const CHECKPOINT_FORKS_KEY = "kagent.mock.checkpointForks"; + +type MockCheckpointFork = { sourceId: string; forkId: string; taskIds: string[] }; + +function readCheckpointForks(): MockCheckpointFork[] { + return JSON.parse(window.sessionStorage.getItem(CHECKPOINT_FORKS_KEY) ?? "[]") as MockCheckpointFork[]; +} + +export function recordCheckpointFork(sourceId: string, forkId: string, taskIds: string[]): void { + window.sessionStorage.setItem(CHECKPOINT_FORKS_KEY, JSON.stringify([ + ...readCheckpointForks(), { sourceId, forkId, taskIds }, + ])); +} + +/** Local and inherited boundaries, preserving the originating instance and name. */ export function readCheckpoints(agentInstanceId: string): MockCheckpoint[] { - return readAll().filter((row) => row.agentInstanceId === agentInstanceId); + const local = readAll().filter((row) => row.agentInstanceId === agentInstanceId); + const parent = readCheckpointForks().find((fork) => fork.forkId === agentInstanceId); + const inherited = parent + ? readCheckpoints(parent.sourceId).filter((row) => parent.taskIds.includes(row.headTaskId)) + : []; + return [...local, ...inherited].sort((a, b) => a.id.localeCompare(b.id)); } export function checkpointById(id: string): MockCheckpoint | undefined { @@ -552,8 +573,14 @@ export function renameCheckpoint(id: string, name: string): MockCheckpoint | und /** Removes one, the way `DeleteCheckpoint` releases the snapshot it was holding. */ export function deleteCheckpoint(id: string): boolean { const rows = readAll(); + const checkpoint = rows.find((row) => row.id === id); + if (!checkpoint) return false; + for (const parent of readCheckpointForks()) { + if (parent.sourceId === checkpoint.agentInstanceId && parent.taskIds.includes(checkpoint.headTaskId)) { + return false; + } + } const kept = rows.filter((row) => row.id !== id); - if (kept.length === rows.length) return false; writeAll(kept); return true; } diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index 81723337d2..8f1112fe3c 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -153,11 +153,12 @@ import { deleteCheckpoint, generatedCheckpointName, readCheckpoints, + recordCheckpointFork, renameCheckpoint, saveCheckpoint, } from "./state"; import type { MockCheckpoint } from "./state"; -import { mockForkTranscript, mockLatestTaskId } from "@/api/chat/mockChatClient"; +import { mockForkTranscript, mockLatestTaskId, mockTranscriptOf } from "@/api/chat/mockChatClient"; /** What a fake is told about the call it is answering. */ interface MockCall { @@ -1025,6 +1026,9 @@ on(CheckpointService.method.forkAgentInstance, (input, call) => { updatedAt: now, }); mockForkTranscript(source.id, forked.id, checkpoint.headTaskId); + recordCheckpointFork(source.id, forked.id, + mockTranscriptOf(forked.id).flatMap((message) => message.taskId ? [message.taskId] : []), + ); return { agentInstance: agentInstanceMessage(forked) }; });