diff --git a/docs/architecture/a2a-gateway.md b/docs/architecture/a2a-gateway.md index ef605561d2..ff8367f631 100644 --- a/docs/architecture/a2a-gateway.md +++ b/docs/architecture/a2a-gateway.md @@ -80,6 +80,20 @@ and history in one transaction. Clients can read completion even when runtime pause/suspend is slow, unavailable, or has not started. TaskStore performs no runtime lifecycle calls and runs no background workers. +Explicit Session deletion can proceed even if a runtime died before acknowledging +its saved boundary ([#3010](https://github.com/kagent-dev/kagent/issues/3010)). +It closes task admission and retains the deletion intent until Substrate confirms +compute deletion. It does not publish the pending boundary or recover the same +conversation. Unexpired dispatches, claimed idle work, and checkpoint creation +still block deletion; a failed backend call leaves deletion pending for retry. +Deletion admission and settlement take the same session lock. If deletion wins, +a late acknowledgement cannot publish the pending state/history or enqueue a +push notification, including after deletion preparation fails. If settlement +committed first, acknowledging that published version remains idempotent until +final deletion. +Idle expiration retains the native-cleanup requirement and never applies this +explicit-deletion exception. + An independent Session lifecycle worker pauses waiting actors or suspends terminal ones and records their matching snapshot. A new turn can supersede idle work that has not been claimed. Once claimed, pause/suspend blocks new execution, diff --git a/go/core/internal/database/active_task_test.go b/go/core/internal/database/active_task_test.go index add942cd9e..b0897330c2 100644 --- a/go/core/internal/database/active_task_test.go +++ b/go/core/internal/database/active_task_test.go @@ -32,6 +32,8 @@ func TestCheckpointCreationBlocksSessionTaskWrites(t *testing.T) { next.ContextID = session.GetContextId() _, err = client.CreateRuntimeTask(ctx, session.GetId(), taskMutationHash("next-request"), next, "") require.ErrorIs(t, err, ErrConflict) + _, err = client.BeginSessionOperation(ctx, session.GetId(), apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, ErrConflict, "explicit deletion must not supersede checkpoint creation") require.ErrorIs(t, saveRuntimeTask(t, client, session.GetId(), task, task, nil), ErrFailedPrecondition) _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") require.ErrorIs(t, err, ErrNotFound) diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 2821feee36..bb1efdc16d 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -252,7 +252,7 @@ func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1 if pending { return fmt.Errorf("checkpoint requires all tasks to be terminal: %w", ErrFailedPrecondition) } - if err := requireSettledRuntime(ctx, tx, session.HistoryID, ""); err != nil { + if err := requireRuntimeAdmission(ctx, tx, session.HistoryID, "", false); err != nil { if errors.Is(err, ErrFailedPrecondition) { return ErrSnapshotPending } diff --git a/go/core/internal/database/session_expiration.go b/go/core/internal/database/session_expiration.go index 87c2d2cd43..e6c838d97c 100644 --- a/go/core/internal/database/session_expiration.go +++ b/go/core/internal/database/session_expiration.go @@ -81,8 +81,9 @@ func (c *Client) BeginIdleSessionDeletion(ctx context.Context, id string, before return deletion, nil } -// requireIdleSession runs under the session lock, before lifecycle admission -// checks dispatch, native cleanup, quiescence, and checkpoint creation. +// requireIdleSession runs under the session lock. Unlike explicit deletion, +// expiration must wait for native cleanup. Lifecycle admission additionally +// checks dispatch, claimed idle work, and checkpoint creation. func requireIdleSession(ctx context.Context, tx pgx.Tx, row sessionRow, idleSince, before time.Time) error { if idleSince.After(before) || row.State == "RUNTIME_STATE_DELETED" || row.State == "RUNTIME_STATE_DELETING" || row.Operation != "RUNTIME_OPERATION_NONE" { return ErrConflict @@ -90,6 +91,7 @@ func requireIdleSession(ctx context.Context, tx pgx.Tx, row sessionRow, idleSinc busy, err := queryOne(ctx, tx, ` SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND state NOT IN ('TASK_STATE_COMPLETED', 'TASK_STATE_CANCELED', 'TASK_STATE_FAILED', 'TASK_STATE_REJECTED')) + OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND NOT published) `, pgx.RowTo[bool], row.HistoryID) if err != nil { return err diff --git a/go/core/internal/database/session_expiration_test.go b/go/core/internal/database/session_expiration_test.go index 799f929897..40248639ce 100644 --- a/go/core/internal/database/session_expiration_test.go +++ b/go/core/internal/database/session_expiration_test.go @@ -6,6 +6,7 @@ import ( "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/stretchr/testify/require" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" @@ -55,7 +56,15 @@ func TestSessionExpirationTaskAdmission(t *testing.T) { func TestSessionExpirationRechecksActivity(t *testing.T) { client, session := expirationFixture(t) ctx := t.Context() - before := time.Now() + // Event timestamps come from PostgreSQL. Use its clock for the precise + // before/after cutoffs too: a container's clock can lag the test host. + databaseTime := func() time.Time { + t.Helper() + now, err := queryOne(ctx, client.db, `SELECT clock_timestamp()`, pgx.RowTo[time.Time]) + require.NoError(t, err) + return now + } + before := databaseTime() ids, err := client.ListIdleSessions(ctx, before, "", 100) require.NoError(t, err) require.Equal(t, []string{session.Id}, ids) @@ -64,13 +73,13 @@ func TestSessionExpirationRechecksActivity(t *testing.T) { task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted version, err := client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("completed"), task, "") require.NoError(t, err) - savedBefore := time.Now() - _, err = client.BeginIdleSessionDeletion(ctx, session.Id, time.Now()) + savedBefore := databaseTime() + _, err = client.BeginIdleSessionDeletion(ctx, session.Id, databaseTime()) require.ErrorIs(t, err, ErrConflict, "unpublished runtime cleanup blocks expiration") require.NoError(t, client.SettleSessionTask(ctx, session.Id, string(task.ID), version)) work, err := client.ClaimSessionQuiescence(ctx) require.NoError(t, err) - _, err = client.BeginIdleSessionDeletion(ctx, session.Id, time.Now()) + _, err = client.BeginIdleSessionDeletion(ctx, session.Id, databaseTime()) require.ErrorIs(t, err, ErrFailedPrecondition, "a claimed suspension blocks expiration") require.NoError(t, client.FinishSessionQuiescence(ctx, work, &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshot", ContentScope: "FULL"})) _, err = client.BeginIdleSessionDeletion(ctx, session.Id, before) @@ -93,6 +102,8 @@ func TestSessionExpirationFencesDispatchAndRetries(t *testing.T) { require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, dispatch, "message")) _, err := client.BeginIdleSessionDeletion(ctx, session.Id, time.Now()) require.ErrorIs(t, err, ErrFailedPrecondition) + _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, ErrFailedPrecondition, "explicit deletion must not supersede a live dispatch") _, err = client.RevokeSessionDispatch(ctx, session.Id, dispatch, "message") require.NoError(t, err) _, err = client.BeginIdleSessionDeletion(ctx, session.Id, time.Time{}) diff --git a/go/core/internal/database/session_operations.go b/go/core/internal/database/session_operations.go index 918083bd16..c47bfc0bef 100644 --- a/go/core/internal/database/session_operations.go +++ b/go/core/internal/database/session_operations.go @@ -87,7 +87,11 @@ func beginSessionOperation(ctx context.Context, tx pgx.Tx, row sessionRow, kind // admitAgentLifecycle runs under the same row lock as task and checkpoint // admission. Common transitions have no knowledge of these agent-only records. func admitAgentLifecycle(ctx context.Context, tx pgx.Tx, row sessionRow, kind apiv1.RuntimeOperation) error { - if err := requireSettledRuntime(ctx, tx, row.HistoryID, ""); err != nil { + // Explicit deletion can supersede an unpublished boundary (#3010, case B). + // It closes admission and retains its intent until compute deletion succeeds; + // it does not publish the boundary or permit reuse of the conversation. + // Claimed lifecycle work remains protected until recovery can fence its owner. + if err := requireRuntimeAdmission(ctx, tx, row.HistoryID, "", kind == apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE); err != nil { return err } busy, err := queryOne(ctx, tx, ` diff --git a/go/core/internal/database/session_tasks.go b/go/core/internal/database/session_tasks.go index 017e24cdaf..adfeac1cc0 100644 --- a/go/core/internal/database/session_tasks.go +++ b/go/core/internal/database/session_tasks.go @@ -61,7 +61,7 @@ func (c *Client) ReserveSessionDispatch(ctx context.Context, sessionID string, d if blocked { return ErrConflict } - if err := requireSettledRuntime(ctx, tx, session.HistoryID, ""); err != nil { + if err := requireRuntimeAdmission(ctx, tx, session.HistoryID, "", false); err != nil { if errors.Is(err, ErrFailedPrecondition) { return ErrDispatchBusy } @@ -201,7 +201,7 @@ func (c *Client) writeRuntimeTask(ctx context.Context, sessionID string, expecte return fmt.Errorf("dispatch expired or was revoked before acceptance: %w", ErrFailedPrecondition) } } - if err := requireSettledRuntime(ctx, tx, session.HistoryID, dispatchID); err != nil { + if err := requireRuntimeAdmission(ctx, tx, session.HistoryID, dispatchID, false); err != nil { return err } // New execution supersedes idle work that no worker has claimed yet. @@ -809,16 +809,17 @@ func readTaskMessages(ctx context.Context, db dbExecutor, historyID uuid.UUID, t `, pgx.RowToStructByName[taskHistoryRow], historyID, taskIDs, historyLength, includeUnpublished) } -// requireSettledRuntime rejects unfinished native cleanup, claimed idle work, or -// another unexpired dispatch. A writer may pass its own dispatch ID. +// requireRuntimeAdmission rejects claimed idle work and another unexpired +// dispatch. A writer may pass its own dispatch ID. Unpublished events also block +// admission unless allowUnpublished is set for explicit deletion. // Callers hold the session lock so new execution and lifecycle claims cannot race. -func requireSettledRuntime(ctx context.Context, db dbExecutor, historyID uuid.UUID, dispatchID string) error { +func requireRuntimeAdmission(ctx context.Context, db dbExecutor, historyID uuid.UUID, dispatchID string, allowUnpublished bool) error { pending, err := queryOne(ctx, db, ` SELECT EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 - AND (NOT published OR (quiescence_pending AND quiescence_executor_id IS NOT NULL))) + AND ((NOT published AND NOT $3::boolean) OR (quiescence_pending AND quiescence_executor_id IS NOT NULL))) OR EXISTS (SELECT 1 FROM session WHERE history_id = $1 AND dispatch_expires_at > clock_timestamp() AND ($2::text = '' OR dispatch_id::text <> $2)) - `, pgx.RowTo[bool], historyID, dispatchID) + `, pgx.RowTo[bool], historyID, dispatchID, allowUnpublished) if err != nil { return err } diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index bb80540a14..75784771dc 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -13,7 +13,8 @@ import ( // SettleSessionTask publishes a saved terminal/waiting task and its history // atomically after native cleanup. It does not wait for runtime pause or snapshots. // Retries cannot republish an old state over a later turn. Callers authenticate -// runtime authority; a missing or deleted session returns ErrNotFound. +// runtime authority; a missing or deleted session returns ErrNotFound. Deletion +// supersedes unpublished boundaries; their late acknowledgements return ErrConflict. func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string, version int64) error { return c.withTx(ctx, func(tx pgx.Tx) error { session, err := lockSession(ctx, tx, sessionID) @@ -42,6 +43,12 @@ func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string if row.Published { return nil } + // Deletion keeps admission closed even if preparation clears the + // operation. Reject unpublished acknowledgements in that state so + // discarded results cannot become public or enqueue callbacks. + if session.State == "RUNTIME_STATE_DELETING" { + return ErrConflict + } stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) if err != nil { return err diff --git a/go/core/internal/database/task_store_test.go b/go/core/internal/database/task_store_test.go index f9319b7ead..80a920dcd1 100644 --- a/go/core/internal/database/task_store_test.go +++ b/go/core/internal/database/task_store_test.go @@ -182,7 +182,8 @@ func TestRuntimeCompletionDoesNotWaitForSnapshot(t *testing.T) { require.Equal(t, a2a.TaskStateWorking, public.Status.State) _, err = client.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) require.ErrorIs(t, err, ErrConflict) - require.ErrorIs(t, deleteSession(t.Context(), client, session.Id), ErrFailedPrecondition) + _, err = client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND) + require.ErrorIs(t, err, ErrFailedPrecondition) _, err = client.ClaimSessionQuiescence(t.Context()) require.ErrorIs(t, err, ErrNotFound) // The native cleanup callback has not finished. require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(waiting.ID), version)) @@ -218,6 +219,69 @@ func TestRuntimeCompletionDoesNotWaitForSnapshot(t *testing.T) { require.ErrorIs(t, err, ErrNotFound) } +func TestSettlementSerializesWithSessionDeletion(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + ctx := t.Context() + sessionFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + for _, ordering := range []string{"delete_first", "settle_first", "concurrent"} { + t.Run(ordering, func(t *testing.T) { + session, _, err := client.CreateSession(ctx, newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + session, err = markSessionReady(ctx, client, session.Id, "agent.example") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), uuid.NewString()) + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted + require.NoError(t, client.RegisterSessionPushNotification(ctx, session.Id, task.History[0].ID, "", &a2a.PushConfig{ID: "callback", URL: "http://receiver"})) + version, err := client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("completed"), task, "") + require.NoError(t, err) + settle := func() error { return client.SettleSessionTask(ctx, session.Id, string(task.ID), version) } + remove := func() error { + _, err := client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + return err + } + var settleErr error + switch ordering { + case "delete_first": + require.NoError(t, remove()) + settleErr = settle() + require.ErrorIs(t, settleErr, ErrConflict) + case "settle_first": + settleErr = settle() + require.NoError(t, settleErr) + require.NoError(t, remove()) + case "concurrent": + start := make(chan struct{}) + settled, deleted := make(chan error, 1), make(chan error, 1) + go func() { <-start; settled <- settle() }() + go func() { <-start; deleted <- remove() }() + close(start) + require.NoError(t, <-deleted) + settleErr = <-settled + if settleErr != nil { + require.ErrorIs(t, settleErr, ErrConflict) + } + } + public, err := client.GetSessionTask(ctx, session.Id, string(task.ID), nil) + require.NoError(t, err) + delivery, err := client.ClaimDuePushDelivery(ctx) + require.NoError(t, err) + if settleErr == nil { + require.Equal(t, a2a.TaskStateCompleted, public.Status.State) + require.NotNil(t, delivery) + require.Equal(t, string(task.ID), delivery.TaskID) + require.NoError(t, client.FinishPushDelivery(ctx, *delivery, true)) + // Acknowledging a publication that won the race remains idempotent. + require.NoError(t, settle()) + } else { + require.Equal(t, a2a.TaskStateSubmitted, public.Status.State) + require.Nil(t, delivery) + } + _, err = client.ClaimSessionQuiescence(ctx) + require.ErrorIs(t, err, ErrNotFound, "deletion excludes background runtime work") + }) + } +} + func TestRuntimeForkRetainsOnlyTheCheckpointBoundary(t *testing.T) { client := NewClient(setupTestDB(t), "public") sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") diff --git a/go/core/internal/grpcserver/adk_settlement_recovery_test.go b/go/core/internal/grpcserver/adk_settlement_recovery_test.go new file mode 100644 index 0000000000..3e1e584f54 --- /dev/null +++ b/go/core/internal/grpcserver/adk_settlement_recovery_test.go @@ -0,0 +1,205 @@ +package grpcserver + +import ( + "context" + "errors" + "iter" + "log/slog" + "net" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v1" + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/runner" + adksession "google.golang.org/adk/v2/session" + "google.golang.org/genai" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + + adka2a "github.com/kagent-dev/kagent/go/adk/pkg/a2a" + "github.com/kagent-dev/kagent/go/adk/pkg/controllerclient" + runtimetaskstore "github.com/kagent-dev/kagent/go/adk/pkg/taskstore" + apia2a "github.com/kagent-dev/kagent/go/api/a2a" + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + runtimev1 "github.com/kagent-dev/kagent/go/api/gen/kagent/runtime/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/dbtest" + authimpl "github.com/kagent-dev/kagent/go/core/internal/httpserver/auth" + "github.com/kagent-dev/kagent/go/core/internal/service/taskstore" +) + +type unavailableSettlementStore struct { + *database.Client + fail atomic.Bool + attempts atomic.Int32 +} + +func (s *unavailableSettlementStore) SettleSessionTask(ctx context.Context, sessionID, taskID string, version int64) error { + s.attempts.Add(1) + if s.fail.Load() { + return status.Error(codes.Unavailable, "injected settlement outage") + } + return s.Client.SettleSessionTask(ctx, sessionID, taskID, version) +} + +type settlementProbeModel struct{ calls atomic.Int32 } + +var _ model.LLM = (*settlementProbeModel)(nil) +var _ taskstore.Store = (*unavailableSettlementStore)(nil) + +type cleanupObservedExecutor struct { + a2asrv.AgentExecutor + a2asrv.AgentExecutionCleaner + cleaned chan struct{} +} + +var _ a2asrv.AgentExecutor = (*cleanupObservedExecutor)(nil) +var _ a2asrv.AgentExecutionCleaner = (*cleanupObservedExecutor)(nil) + +func (e *cleanupObservedExecutor) Cleanup(ctx context.Context, input *a2asrv.ExecutorContext, result a2a.SendMessageResult, err error) { + e.AgentExecutionCleaner.Cleanup(ctx, input, result, err) + close(e.cleaned) +} + +func (*settlementProbeModel) Name() string { return "settlement-probe" } + +func (m *settlementProbeModel) GenerateContent(context.Context, *model.LLMRequest, bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + m.calls.Add(1) + yield(&model.LLMResponse{Content: genai.NewContentFromText("finished", genai.RoleModel)}, nil) + } +} + +// Exercise the production Go ADK executor, TaskStore client, gRPC service and +// PostgreSQL together, with an in-memory ADK session and a fixed model response. +// https://github.com/kagent-dev/kagent/issues/3010 case B tracks this failure. +// This verifies DELETE admission, not Actor deletion or same-session recovery. +func TestGoADKDeletionAfterSettlementOutage(t *testing.T) { + for _, fail := range []bool{false, true} { + name := "healthy" + if fail { + name = "settlement_outage" + } + t.Run(name, func(t *testing.T) { + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + dbtest.MigrateT(t, dsn, false) + db, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn)}) + require.NoError(t, err) + t.Cleanup(db.Close) + store := &unavailableSettlementStore{Client: database.NewClient(db, "public")} + store.fail.Store(fail) + session := createTaskStoreSession(t, store.Client) + listener := bufconn.Listen(DefaultMaxMessageSize) + server, err := New(Config{ + Listener: listener, SystemService: testSystemService(), + Authenticator: &authimpl.InsecureAuthenticator{}, RuntimeAuthenticator: &taskstore.Authenticator{}, + TaskStoreService: taskstore.NewService(store, nil), + }) + require.NoError(t, err) + serverCtx, stop := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- server.Start(serverCtx) }() + t.Cleanup(func() { stop(); require.NoError(t, <-done) }) + controller, err := controllerclient.New(controllerclient.Config{ + APIURL: "http://api.test", DialOptions: []grpc.DialOption{ + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), + }, + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, controller.Close()) }) + identityDir := t.TempDir() + for file, value := range map[string]string{"name": "session-" + session.Id, "atespace": "team-a", "uid": "actor-uid"} { + require.NoError(t, os.WriteFile(filepath.Join(identityDir, file), []byte(value), 0o600)) + } + runtimeStore := runtimetaskstore.New(controller, filepath.Join(identityDir, "name")) + llm := &settlementProbeModel{} + agent, err := llmagent.New(llmagent.Config{Name: "probe", Model: llm}) + require.NoError(t, err) + native, err := adka2a.NewKAgentExecutor(adka2a.KAgentExecutorConfig{ + RunnerConfig: runner.Config{AppName: "probe", Agent: agent, SessionService: adksession.InMemoryService()}, + AppName: "probe", Logger: slog.New(slog.DiscardHandler), + }) + require.NoError(t, err) + wrapped := runtimeStore.WrapExecutor(native, "", nil) + observed := &cleanupObservedExecutor{ + AgentExecutor: wrapped, AgentExecutionCleaner: wrapped, cleaned: make(chan struct{}), + } + handler := a2asrv.NewHandler(observed, a2asrv.WithTaskStore(runtimeStore), + a2asrv.WithCallInterceptors(adka2a.HITLActivationInterceptor(), adka2a.UserIDCallInterceptor(), wrapped)) + runtimeListener := bufconn.Listen(DefaultMaxMessageSize) + runtimeServer := grpc.NewServer() + a2agrpc.NewHandler(handler).RegisterWith(runtimeServer) + runtimeDone := make(chan error, 1) + go func() { runtimeDone <- runtimeServer.Serve(runtimeListener) }() + t.Cleanup(func() { + runtimeServer.Stop() + if err := <-runtimeDone; !errors.Is(err, grpc.ErrServerStopped) { + require.NoError(t, err, "runtime A2A server failed") + } + }) + conn, err := grpc.NewClient("passthrough:///runtime", grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return runtimeListener.Dial() })) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second) + defer cancel() + response, err := a2apb.NewA2AServiceClient(conn).SendMessage(ctx, &a2apb.SendMessageRequest{ + Message: &a2apb.Message{MessageId: uuid.NewString(), ContextId: session.ContextId, Role: a2apb.Role_ROLE_USER, + Parts: []*a2apb.Part{{Content: &a2apb.Part_Text{Text: "finish one turn"}}}}, + }) + require.NoError(t, err) + select { + case <-observed.cleaned: + case <-ctx.Done(): + t.Fatal("ADK cleanup did not finish", ctx.Err()) + } + require.EqualValues(t, 1, llm.calls.Load()) + taskID := response.GetTask().GetId() + require.NotEmpty(t, taskID) + stored, version, err := store.GetVersionedSessionTask(t.Context(), session.Id, taskID) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) + public, err := store.GetSessionTask(t.Context(), session.Id, taskID, nil) + require.NoError(t, err) + if fail { + require.EqualValues(t, 4, store.attempts.Load()) + require.Equal(t, a2a.TaskStateWorking, public.Status.State) + // Restoring availability after Cleanup cannot retry the boundary. + store.fail.Store(false) + _, err = store.GetSettledSessionTask(t.Context(), session.Id, taskID, nil) + require.ErrorIs(t, err, database.ErrConflict) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), ""), database.ErrConflict) + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound, "idle worker cannot recover an unpublished boundary") + } else { + require.EqualValues(t, 1, store.attempts.Load()) + require.Equal(t, a2a.TaskStateCompleted, public.Status.State) + } + t.Logf("Go ADK calls=%d; SettleTask attempts=%d; stored=%s; public=%s; cleanup finished", llm.calls.Load(), store.attempts.Load(), stored.Status.State, public.Status.State) + _, err = store.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + t.Logf("DELETE admission: %v", err) + require.NoError(t, err, "explicit deletion must remain available after settlement retries are exhausted") + privateCtx := metadata.NewOutgoingContext(t.Context(), metadata.Pairs(apia2a.InsecureRuntimeIdentityHeader, "team-a/session-"+session.Id+"/actor-uid")) + _, err = controller.TaskStoreService().SettleTask(privateCtx, &runtimev1.TaskStoreServiceSettleTaskRequest{SessionId: session.Id, TaskId: taskID, Version: version}) + if fail { + require.Equal(t, codes.Aborted, status.Code(err), "late cleanup cannot publish after deletion admission") + } else { + require.NoError(t, err, "an already published acknowledgement remains idempotent") + } + }) + } +} diff --git a/go/core/internal/service/session/deletion_recovery_test.go b/go/core/internal/service/session/deletion_recovery_test.go new file mode 100644 index 0000000000..30606503cd --- /dev/null +++ b/go/core/internal/service/session/deletion_recovery_test.go @@ -0,0 +1,175 @@ +package session + +import ( + "context" + "crypto/sha256" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +// An adapter can die after saving a boundary and before SettleTask. Explicit +// deletion must remain available without claiming that the task completed. +func TestDeleteSessionWithUnpublishedRuntimeBoundary(t *testing.T) { + for _, state := range []a2a.TaskState{ + a2a.TaskStateCompleted, a2a.TaskStateCanceled, a2a.TaskStateFailed, + a2a.TaskStateRejected, a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired, + } { + for _, firstEvent := range []bool{false, true} { + name := string(state) + "/update" + if firstEvent { + name = string(state) + "/create" + } + t.Run(name, func(t *testing.T) { + store, session := lifecycleFixture(t) + actors := &deletionFaultActors{ + lifecycleTestActors: &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}}, + effectBeforeError: firstEvent, + readFailure: state == a2a.TaskStateFailed, + suspendFailure: state == a2a.TaskStateCanceled, + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + session, err := workflow.Create(t.Context(), session) + require.NoError(t, err) + name := substrate.ActorName(session.Id) + actor, err := actors.ResumeActor(t.Context(), "team-a", name) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("work")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + require.NoError(t, store.RegisterSessionPushNotification(t.Context(), session.Id, message.ID, "", &a2a.PushConfig{ID: "callback", URL: "http://receiver"})) + var version int64 + if !firstEvent { + digest := sha256.Sum256([]byte("start")) + version, err = store.CreateRuntimeTask(t.Context(), session.Id, digest[:], task, "") + require.NoError(t, err) + } + task.Status = a2a.TaskStatus{State: state, Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("saved boundary"))} + digest := sha256.Sum256([]byte("boundary")) + var boundaryVersion int64 + if firstEvent { + boundaryVersion, err = store.CreateRuntimeTask(t.Context(), session.Id, digest[:], task, "") + } else { + boundaryVersion, err = store.UpdateSessionTask(t.Context(), session.Id, version, digest[:], task, task, "") + } + require.NoError(t, err) + + _, err = store.BeginIdleSessionDeletion(t.Context(), session.Id, time.Now().Add(time.Hour)) + require.ErrorIs(t, err, database.ErrConflict, "TTL must not apply the explicit deletion exception") + // Timeout is not cleanup proof. Admission closes, but the task's + // unpublished status remains hidden and the delete intent survives. + actors.fail = true + _, err = workflow.Delete(t.Context(), session) + require.Error(t, err) + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETING, current.State) + if actors.readFailure { + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation, "preparation failure clears the operation but must not reopen settlement") + } + _, err = store.BeginIdleSessionDeletion(t.Context(), session.Id, time.Now().Add(time.Hour)) + require.ErrorIs(t, err, database.ErrConflict, "TTL must not take over a user's deletion intent") + public, err := store.GetSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateSubmitted, public.Status.State) + _, err = store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.ErrorIs(t, err, database.ErrConflict) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), ""), database.ErrConflict) + // Deletion also fences a cleanup acknowledgement that was delayed + // in flight. It must not publish the boundary or enqueue a late push. + require.ErrorIs(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), boundaryVersion), database.ErrConflict) + public, err = store.GetSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateSubmitted, public.Status.State) + require.Nil(t, public.Status.Message) + delivery, err := store.ClaimDuePushDelivery(t.Context()) + require.NoError(t, err) + require.Nil(t, delivery, "discarded boundary must not enqueue a callback") + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound) + // A runtime with stale credentials may still be alive after an + // ambiguous delete. Neither a fresh task nor a new update is admitted. + late := a2a.NewSubmittedTask(message, message) + lateDigest := sha256.Sum256([]byte("late runtime write")) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, lateDigest[:], late, "") + require.ErrorIs(t, err, database.ErrConflict) + update := *task + update.Status = a2a.TaskStatus{State: a2a.TaskStateWorking} + _, err = store.UpdateSessionTask(t.Context(), session.Id, boundaryVersion, lateDigest[:], &update, &update, "") + require.ErrorIs(t, err, database.ErrConflict) + + // A fresh controller reconciles both a failed deletion and a lost + // response after the Actor was already deleted. + actors.fail = false + workflow = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + deleted, err := workflow.Delete(t.Context(), current) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.State) + require.Empty(t, deleted.A2AAuthority) + _, err = actors.GetActor(t.Context(), "team-a", name) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = store.GetSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.ErrorIs(t, err, database.ErrNotFound) + _, err = store.GetSessionForRuntime(t.Context(), session.Id, actor.GetMetadata().GetUid()) + require.ErrorIs(t, err, database.ErrNotFound) + require.ErrorIs(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), boundaryVersion), database.ErrNotFound) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, lateDigest[:], late, "") + require.ErrorIs(t, err, database.ErrNotFound) + again, err := workflow.Delete(t.Context(), current) + require.NoError(t, err) + require.Equal(t, deleted.State, again.State) + }) + } + } +} + +type deletionFaultActors struct { + *lifecycleTestActors + fail bool + effectBeforeError bool + readFailure bool + suspendFailure bool +} + +var _ substrate.LifecycleClient = (*deletionFaultActors)(nil) + +func (a *deletionFaultActors) GetActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + if a.fail && a.readFailure { + return nil, status.Error(codes.Unavailable, "injected preparation failure") + } + return a.lifecycleTestActors.GetActor(ctx, atespace, name) +} + +func (a *deletionFaultActors) SuspendActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + if a.fail && a.suspendFailure { + if a.effectBeforeError { + if _, err := a.lifecycleTestActors.SuspendActor(ctx, atespace, name); err != nil { + return nil, err + } + } + return nil, status.Error(codes.Unavailable, "injected suspension failure") + } + return a.lifecycleTestActors.SuspendActor(ctx, atespace, name) +} + +func (a *deletionFaultActors) DeleteActor(ctx context.Context, atespace, name string) error { + if a.fail { + if a.effectBeforeError { + if err := a.lifecycleTestActors.DeleteActor(ctx, atespace, name); err != nil { + return err + } + } + return status.Error(codes.Unavailable, "injected deletion failure") + } + return a.lifecycleTestActors.DeleteActor(ctx, atespace, name) +} diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index d5cb2b9717..3d4fa914d9 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -54,6 +54,14 @@ The suite exercises the public API against a clean Kind installation. It does not reconcile Kubernetes resources itself: installation creates the Harness fixtures, and each test owns the templates and API resources it creates. +`TestSessionUnpublishedBoundaryDeletion` keeps a live runtime blocked at a model +request and stages a terminal or waiting boundary through its private TaskStore +API without acknowledging cleanup. It checks that public deletion removes the +Actor, revokes stale runtime writes, and prevents sends from recreating compute. +The waiting case also discards a successful deletion reply and retries through +a fresh public call. This uses the shared harness matrix and can be selected for +one harness using the same subtest filters as other interaction cases. + Shared AgentTemplate and Session tests use `forEachHarness`, which runs every case on kagent, Codex, Claude, and configured BYO (the Go ADK image through the BYO compiler). Each harness gets a named subtest and independent resources, diff --git a/go/core/test/e2e/unpublished_boundary_test.go b/go/core/test/e2e/unpublished_boundary_test.go new file mode 100644 index 0000000000..11a42917af --- /dev/null +++ b/go/core/test/e2e/unpublished_boundary_test.go @@ -0,0 +1,121 @@ +package e2e_test + +import ( + "context" + "fmt" + "testing" + "time" + + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + apia2a "github.com/kagent-dev/kagent/go/api/a2a" + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + runtimev1 "github.com/kagent-dev/kagent/go/api/gen/kagent/runtime/v1" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +// Stage the cleanup gap through TaskStore while a real runtime is blocked at its +// model request. No settlement acknowledgement is sent. Explicit deletion must +// stop and remove that Actor even though the public task is still WORKING. +func TestSessionUnpublishedBoundaryDeletion(t *testing.T) { + t.Parallel() + forEachHarness(t, func(t *testing.T, harness testHarness) { + t.Parallel() + for _, boundary := range []a2apb.TaskState{a2apb.TaskState_TASK_STATE_COMPLETED, a2apb.TaskState_TASK_STATE_INPUT_REQUIRED} { + t.Run(boundary.String(), func(t *testing.T) { + target := interactionTarget(t) + modelURL, started := startBlockingInteractionMock(t) + fixture := newInteractionFixture(t, harness, target, modelURL) + _, request := newMessageRequest(t, "Wait for cancellation") + request.Tenant, request.Message.ContextId = fixture.tenant, fixture.contextID + stream, err := fixture.client.SendStreamingMessage(fixture.ctx, request) + require.NoError(t, err) + select { + case <-started: + case <-time.After(time.Minute): + t.Fatal("runtime did not start its model request") + } + var taskID string + for taskID == "" { + event, err := stream.Recv() + require.NoError(t, err) + taskID = event.GetTask().GetId() + } + actor, err := findSubstrateActor(fixture.ctx, fixture.system, "kagent", substrate.ActorName(fixture.sessionID)) + require.NoError(t, err) + require.NotNil(t, actor) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_RUNNING, actor.GetStatus().GetState()) + privateCtx := metadata.NewOutgoingContext(fixture.ctx, metadata.Pairs(apia2a.InsecureRuntimeIdentityHeader, + fmt.Sprintf("kagent/%s/%s", actor.GetMetadata().GetName(), actor.GetMetadata().GetUid()))) + store := runtimev1.NewTaskStoreServiceClient(newControllerConn(t, target)) + stored, err := store.GetTask(privateCtx, &runtimev1.TaskStoreServiceGetTaskRequest{SessionId: fixture.sessionID, TaskId: taskID}) + require.NoError(t, err) + task := proto.CloneOf(stored.Stored.Task) + task.Status = &a2apb.TaskStatus{State: boundary} + saved, err := store.UpdateTask(privateCtx, &runtimev1.TaskStoreServiceUpdateTaskRequest{ + SessionId: fixture.sessionID, Task: task, ExpectedVersion: stored.Stored.Version, + }) + require.NoError(t, err) + public, err := fixture.client.GetTask(fixture.ctx, &a2apb.GetTaskRequest{Tenant: fixture.tenant, Id: taskID}) + require.NoError(t, err) + require.Equal(t, a2apb.TaskState_TASK_STATE_WORKING, public.Status.State) + remove := &apiv1.DeleteSessionRequest{SessionId: fixture.sessionID} + if boundary == a2apb.TaskState_TASK_STATE_INPUT_REQUIRED { + // Lose the public reply after the real backend deletion commits. + // A fresh retry observes NotFound and cannot recreate compute. + loss := &lostDeleteReplyConn{ClientConnInterface: newControllerConn(t, target)} + _, err = apiv1.NewSessionServiceClient(loss).DeleteSession(fixture.ctx, remove) + require.True(t, loss.lost) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = fixture.sessions.DeleteSession(fixture.ctx, remove) + require.Equal(t, codes.NotFound, status.Code(err)) + } else { + deleted, err := fixture.sessions.DeleteSession(fixture.ctx, remove) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.Session.State) + } + actor, err = findSubstrateActor(fixture.ctx, fixture.system, "kagent", substrate.ActorName(fixture.sessionID)) + require.NoError(t, err) + require.Nil(t, actor) + _, err = fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = store.SettleTask(privateCtx, &runtimev1.TaskStoreServiceSettleTaskRequest{SessionId: fixture.sessionID, TaskId: taskID, Version: saved.Version}) + require.Equal(t, codes.NotFound, status.Code(err), "deleted runtime authority cannot publish the staged boundary") + task.Status.State = a2apb.TaskState_TASK_STATE_WORKING + _, err = store.UpdateTask(privateCtx, &runtimev1.TaskStoreServiceUpdateTaskRequest{SessionId: fixture.sessionID, Task: task, ExpectedVersion: saved.Version}) + require.Equal(t, codes.NotFound, status.Code(err)) + _, send := newMessageRequest(t, "Do not recreate deleted compute") + send.Tenant, send.Message.ContextId = fixture.tenant, fixture.contextID + _, err = fixture.client.SendMessage(fixture.ctx, send) + require.Equal(t, codes.PermissionDenied, status.Code(err), "deleted contexts no longer authorize execution") + actor, err = findSubstrateActor(fixture.ctx, fixture.system, "kagent", substrate.ActorName(fixture.sessionID)) + require.NoError(t, err) + require.Nil(t, actor, "rejected sends must not recreate the Actor") + t.Logf("unpublished %s deleted with its live %s Actor; stale writes rejected", boundary, harness.runtimeLabel) + }) + } + }) +} + +type lostDeleteReplyConn struct { + grpc.ClientConnInterface + lost bool +} + +var _ grpc.ClientConnInterface = (*lostDeleteReplyConn)(nil) + +func (c *lostDeleteReplyConn) Invoke(ctx context.Context, method string, args, reply any, opts ...grpc.CallOption) error { + err := c.ClientConnInterface.Invoke(ctx, method, args, reply, opts...) + if err == nil && method == apiv1.SessionService_DeleteSession_FullMethodName && !c.lost { + c.lost = true + return status.Error(codes.Unavailable, "injected loss of a committed deletion reply") + } + return err +}