diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index e1ed0bb6de..9f3585f992 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -2,13 +2,19 @@ ## 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 -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. + +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: @@ -18,7 +24,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 | +| `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,15 +46,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] - 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 @@ -74,23 +81,47 @@ 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 context 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 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 `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. Foreign keys require the parent to share the +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, 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 +history boundary, so a checkpoint created later at an already inherited boundary +also appears. Listing uses checkpoint-ID pagination across local and inherited +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 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 d409e0a867..47a7bc60b0 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 @@ -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 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/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/internal/database/agent_instances.go b/go/core/internal/database/agent_instances.go index c581542a2b..2743dcf7d1 100644 --- a/go/core/internal/database/agent_instances.go +++ b/go/core/internal/database/agent_instances.go @@ -405,15 +405,25 @@ 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) VALUES ($1, $2, $3) - `, historyID, instance.Creator, instance.ContextId); err != nil { - return agentInstanceRow{}, fmt.Errorf("insert A2A context: %w", err) + 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, ` INSERT INTO agent_instance (id, user_id, request_id, context_id, history_id, prepared_revision, 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..c6adff91e4 --- /dev/null +++ b/go/core/internal/database/checkpoint_lineage_test.go @@ -0,0 +1,387 @@ +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 checkpointNextTurn(t *testing.T, client *Client, instance *apiv1alpha1.AgentInstance) *apiv1alpha1.Checkpoint { + t.Helper() + ctx := t.Context() + 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 +} + +func TestAgentHistoryConstraints(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) + 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 + }{ + {"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, boundary.HistorySequence) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "23503", pgErr.Code) + require.Equal(t, "agent_history_parent_binding_fkey", 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 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} { + 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()) + 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(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) + 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) + require.NoError(t, client.DeleteAgentInstanceCheckpoint(ctx, deleting.Id, "alice")) + c3 := checkpointNextTurn(t, client, source) + c4 := checkpointNextTurn(t, client, branch) + nested := fork(c4) + c5 := checkpointNextTurn(t, client, branch) + assertListed(source.Id, c1, c2, c3) + assertListed(branch.Id, c1, c2, c4, c5) + assertListed(nested.Id, c1, c2, c4) + + 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) + } + }) + } + } + t.Run("before_deletion", assertPages) + 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) + assertListed(source.Id, c1, c2, c3) + 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) + // 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)) + 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} { + _, _, 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..d8f9aa9274 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 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) @@ -80,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 } @@ -89,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") } @@ -392,20 +397,37 @@ 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. 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, ` - 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 - AND user_id = $2 - AND state = 'READY' - AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) - ORDER BY id - LIMIT $4 + WITH RECURSIVE lineage (history_id, boundary) AS ( + 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 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 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), ) @@ -425,8 +447,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 +458,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 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) + } checkpoint, err := toAgentInstanceCheckpoint(row) if err != nil { return err @@ -449,9 +478,11 @@ 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 agent_history h + 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 { return err diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index bb74830838..05115dc006 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, 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 a2a_context (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 a2a_context (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 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) } @@ -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()) @@ -701,7 +700,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, 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/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/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/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 0a86455ae3..dd884e5ed2 100644 --- a/go/core/pkg/migrations/core/000001_initial.sql +++ b/go/core/pkg/migrations/core/000001_initial.sql @@ -65,13 +65,28 @@ 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 ( - id UUID PRIMARY KEY, - user_id TEXT NOT NULL CHECK (user_id <> ''), - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - context_id UUID NOT NULL, - CONSTRAINT a2a_context_binding_key UNIQUE (id, context_id) +-- 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, + -- 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), + 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) ); +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 ( id UUID PRIMARY KEY, @@ -86,14 +101,14 @@ 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')), 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 ON agent_instance_checkpoint (source_instance_id) WHERE state = 'CREATING'; @@ -121,7 +136,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 a2a_context(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', @@ -145,7 +161,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, @@ -179,7 +195,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(), @@ -280,7 +296,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..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 a2a_context (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 3f426c17b1..fe1fb4df94 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,46 @@ 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) + } + // 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{ + CheckpointId: id, + }); status.Code(err) != codes.NotFound { + t.Fatalf("delete retained checkpoint %s = %v, want NotFound", id, err) + } } } 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 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) }; });