diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index ef0c48efbf..21bacfa738 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -154,6 +154,8 @@ jobs: port: http periodSeconds: 1 config: + sessions: + automaticRecovery: true push: allowHttp: true allowPrivateNetworks: true diff --git a/docs/architecture/process-configuration.md b/docs/architecture/process-configuration.md index d298b7100f..fafb49ac1c 100644 --- a/docs/architecture/process-configuration.md +++ b/docs/architecture/process-configuration.md @@ -114,6 +114,12 @@ runtimeDefaults: level: info ``` +`sessions.automaticRecovery` enables recovery of crashed runtimes during sends. +It defaults to false. The chart renders it only when set, so a rollback to a +release that predates the setting can still read the default config. Explicit +`RevertSession` calls remain available with either value. The setting does not +change migrations or stored data. + The controller connects with two logins. `database.owner` applies migrations at startup; `database.readWrite` serves every runtime query. Each assumes its `role` after connecting. Set `database.skipMigrations: true` when migrations are applied diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index a76cc94697..235a4df1c4 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -265,25 +265,82 @@ retires the claim in its transaction, so a late finish records nothing. Execution reverts the same Actor to its last durable snapshot without replaying input. A runtime without a durable snapshot returns to its initial state. An interrupted revert continues from REVERTING; a lost response can be reconciled -from SUSPENDED. Nothing retries a recovery in the background. An admitted -recovery survives API restarts as a durable REVERT operation, and the client -continues it by retrying RevertSession, as with Create and Delete. While it is -pending, `GetSession` reports FAILED/REVERT with failure reason `RuntimeCrashed`, -and interactions are refused because a lifecycle operation is in progress. +from SUSPENDED. An admitted recovery survives API restarts as a durable REVERT +operation, and the client continues it by retrying RevertSession, as with Create +and Delete. Nothing retries it in the background. While it is +pending, `GetSession` reports FAILED/REVERT with failure reason `RuntimeCrashed`. + +With `sessions.automaticRecovery` enabled in controller config, an A2A send also +starts recovery, so chat and API clients recover with the retry they already +make. Substrate cannot resume a crashed Actor, so a send to one fails before the +runtime accepts it. When the gateway proves the input was not accepted, it runs +the same Revert on that Session before answering `KAGENT_SEND_NOT_ACCEPTED`, and +the client's retry reaches the restored runtime. A crash during a turn leaves +that turn open, and an open turn refuses every send, so a send refused for an +open turn also runs Revert once and then tries its reservation again: a crashed +runtime is restored and the send proceeds, and a busy healthy runtime refuses it +as before. Other conflicts, such as a checkpoint in progress, do not start +recovery. A send while recovery is pending joins it. + +A worker can also be lost while the Actor suspends after a turn, for example +during a node drain or scale-down. The Actor is then CRASHED, the Session still +reads READY, and the idle worker keeps the suspension's claim because the outcome +is unknown; nothing takes the claim over on a timeout. Every later send waits out +the 10-second busy period and is told to retry, so the conversation would never +recover on its own. Reservation reports a held suspension claim as +`ErrDispatchSuspending`, a kind of `ErrDispatchBusy`. When that claim outlasts +the whole wait, which a healthy suspension does not, the send attempts the same +bounded recovery once. Admitting the recovery retires the claim. Revert leaves a +suspension that is still running alone, and one more reservation, without +another wait, proceeds or answers retry as before. + +A send waits at most 10 seconds on the recovery it starts. A recovery still +running then stays admitted: the attempt releases its claim, the send is +answered `KAGENT_SEND_NOT_ACCEPTED` with `state: RECOVERING` and a one-second +`retryAfterMs`, and the client's retry continues it. The reason is unchanged so +existing clients keep retrying. Reverting needs Session update access, as +`RevertSession` does. When the revert cannot run, whether for missing access, +contention, or an outage, the send keeps the answer it would give without +recovery, unless a recovery is already admitted, which answers RECOVERING. No +Session read calls Substrate. + +The setting is off by default. Sends then keep the refusals they had before +recovery existed, and only an explicit `RevertSession` recovers. Chat retries a +send refused as `KAGENT_SEND_NOT_ACCEPTED` with the same message ID, doubling +the wait from the gateway's `retryAfterMs` up to two seconds for 30 seconds, +and then offers Recover, which calls `RevertSession`. A crash during a turn +leaves that turn open, so with the setting off the next send is refused +outright while the Session still reads READY. The gateway names that refusal +`KAGENT_TURN_ACTIVE` in its ErrorInfo metadata, and chat offers Recover for it +too, because nothing else observes the crash. Other refusals keep a plain error. +Chat also offers Recover while a recovery is pending, because a recovery whose +attempt failed stays FAILED/REVERT until a client retries it. Recover returns a +healthy runtime unchanged and records a lost one, so it is safe to offer on a +guess. + +A missing, replaced, or retemplated Actor is not this Session's runtime, and +reverting it would restore the wrong one. Whichever Revert observes that, +directly or while executing an admitted recovery, records it in one transaction +guarded by the generation it observed: the Session becomes FAILED/NONE with +failure reason `RuntimeLost`, every open task fails, and a pending dispatch is +cleared. A newer generation makes the observation stale and leaves the record +alone. Later sends and reverts answer from PostgreSQL with a non-retryable error +and do not ask Substrate again. Delete still works. Delete supersedes an admitted recovery even after an attempt issued the revert. Other issued lifecycle work blocks Delete because it may still create something Delete would miss; a revert restores the recorded Actor in place and creates nothing. Substrate serializes the two calls, and a late revert fails against the -deleted Actor. A recovery whose Actor is missing or replaced stays pending, and -each retry returns -`FAILED_PRECONDITION` so the client deletes the Session. +deleted Actor. Recording a lost runtime supersedes an admitted recovery the +same way. A successful revert returns the Session to READY with the same Actor suspended. The next authorized interaction resumes it through Substrate ingress. Runtime changes after the retained snapshot are lost; the discarded turn receives no checkpoint at that older snapshot. Another revert returns the Session. +Chat shows a pending recovery and a lost runtime from the stored Session state. + ```mermaid sequenceDiagram participant Client diff --git a/docs/lifecycle-retries.md b/docs/lifecycle-retries.md index 452b09ad1b..d1227358ca 100644 --- a/docs/lifecycle-retries.md +++ b/docs/lifecycle-retries.md @@ -41,7 +41,10 @@ should likewise retain the `request_id` passed to `create_sandbox`. ## Recovering a crashed Session -Call `SessionService.RevertSession` with the Session ID to restore its recorded runtime to the last durable snapshot. +A send recovers a crashed runtime. It returns `KAGENT_SEND_NOT_ACCEPTED` after reverting, so retry the send as usual. +A send refused because a crashed turn is still active recovers the runtime first and then proceeds. +A send while recovery is pending joins it. Recovering through a send needs Session update access. +Call `SessionService.RevertSession` with the Session ID to restore the recorded runtime to its last durable snapshot without sending. The caller needs Session update access. The operation verifies Substrate's CRASHED state and the recorded Actor identity. Recovery fails interrupted turns and waiting tasks absent from the retained snapshot. It preserves durable terminal results and returns the Session to READY with its Actor suspended. @@ -51,15 +54,17 @@ A runtime without a snapshot restarts from its initial state. Retry an interrupted revert against the same Session ID; nothing retries it in the background. Pending recovery retains its operation across API restarts, and a retry joins it. While recovery is pending, `GetSession` reports state `FAILED`, operation `REVERT`, and failure reason `RuntimeCrashed`. -Interactions fail because a lifecycle operation is in progress. Call `RevertSession` again to finish the recovery. +Retry the send, or call `RevertSession` again, to finish the recovery. A retry after a revert whose response was lost returns the recovered Session. Checkpoint creation returns `ABORTED` until it finishes. A revert that observed the crash before another lifecycle operation began also returns `ABORTED`. So does a revert while another lifecycle operation is active. A revert of a runtime that is not crashed returns the Session unchanged. -A revert that can never succeed returns `FAILED_PRECONDITION`: the recorded Actor is missing, or a same-name Actor replaced it. -Stop retrying, delete the Session, then start another or fork one from a checkpoint. Delete succeeds even while the recovery is pending. +When the recorded Actor is missing, or a same-name Actor replaced it, the runtime cannot be recovered. +The Session is recorded as state `FAILED`, operation `NONE`, failure reason `RuntimeLost`, and its open tasks fail. +From then on, `RevertSession` returns `FAILED_PRECONDITION` and sends fail with an error that is not retryable. +Stop retrying, delete the Session, then start another or fork one from a checkpoint. Read the reason rather than the message text, which is not stable. ## Errors diff --git a/go/api/gen/kagent/api/v1/sessions.pb.go b/go/api/gen/kagent/api/v1/sessions.pb.go index ab556d6f9b..3a85696443 100644 --- a/go/api/gen/kagent/api/v1/sessions.pb.go +++ b/go/api/gen/kagent/api/v1/sessions.pb.go @@ -81,14 +81,17 @@ type Session struct { Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` Creator string `protobuf:"bytes,2,opt,name=creator,proto3" json:"creator,omitempty"` // Output only. Agent identity captured from this Session's pinned revision. - Agent *AgentProvenance `protobuf:"bytes,15,opt,name=agent,proto3" json:"agent,omitempty"` - PreparedRevision string `protobuf:"bytes,5,opt,name=prepared_revision,json=preparedRevision,proto3" json:"prepared_revision,omitempty"` - A2AAuthority string `protobuf:"bytes,6,opt,name=a2a_authority,json=a2aAuthority,proto3" json:"a2a_authority,omitempty"` - State RuntimeState `protobuf:"varint,7,opt,name=state,proto3,enum=kagent.api.v1.RuntimeState" json:"state,omitempty"` - Operation RuntimeOperation `protobuf:"varint,8,opt,name=operation,proto3,enum=kagent.api.v1.RuntimeOperation" json:"operation,omitempty"` - Failure *Failure `protobuf:"bytes,9,opt,name=failure,proto3" json:"failure,omitempty"` - CreatedAt *timestamppb.Timestamp `protobuf:"bytes,10,opt,name=created_at,json=createdAt,proto3" json:"created_at,omitempty"` - UpdatedAt *timestamppb.Timestamp `protobuf:"bytes,11,opt,name=updated_at,json=updatedAt,proto3" json:"updated_at,omitempty"` + Agent *AgentProvenance `protobuf:"bytes,15,opt,name=agent,proto3" json:"agent,omitempty"` + PreparedRevision string `protobuf:"bytes,5,opt,name=prepared_revision,json=preparedRevision,proto3" json:"prepared_revision,omitempty"` + A2AAuthority string `protobuf:"bytes,6,opt,name=a2a_authority,json=a2aAuthority,proto3" json:"a2a_authority,omitempty"` + State RuntimeState `protobuf:"varint,7,opt,name=state,proto3,enum=kagent.api.v1.RuntimeState" json:"state,omitempty"` + Operation RuntimeOperation `protobuf:"varint,8,opt,name=operation,proto3,enum=kagent.api.v1.RuntimeOperation" json:"operation,omitempty"` + // Output only. With state FAILED, reason RuntimeCrashed and operation REVERT + // mean a crash recovery is pending; reason RuntimeLost and operation NONE mean + // the recorded runtime is gone and the Session can only be deleted. + Failure *Failure `protobuf:"bytes,9,opt,name=failure,proto3" json:"failure,omitempty"` + CreatedAt *timestamppb.Timestamp `protobuf:"bytes,10,opt,name=created_at,json=createdAt,proto3" json:"created_at,omitempty"` + UpdatedAt *timestamppb.Timestamp `protobuf:"bytes,11,opt,name=updated_at,json=updatedAt,proto3" json:"updated_at,omitempty"` // Reader-supplied display name for the conversation. Empty means unnamed. Name string `protobuf:"bytes,13,opt,name=name,proto3" json:"name,omitempty"` // Output only. A2A conversation ID, equal to this session's ID. diff --git a/go/api/gen/kagent/config/v1/config.pb.go b/go/api/gen/kagent/config/v1/config.pb.go index 58b7630f7e..62635f3485 100644 --- a/go/api/gen/kagent/config/v1/config.pb.go +++ b/go/api/gen/kagent/config/v1/config.pb.go @@ -1250,6 +1250,8 @@ func (x *SubstrateConfig) GetActorJwtIssuer() string { type SessionsConfig struct { state protoimpl.MessageState `protogen:"open.v1"` + // Recover crashed runtimes during sends. Disabled by default; explicit revert remains available. + AutomaticRecovery bool `protobuf:"varint,5,opt,name=automatic_recovery,json=automaticRecovery,proto3" json:"automatic_recovery,omitempty"` // Zero disables idle expiration. IdleTtl *durationpb.Duration `protobuf:"bytes,1,opt,name=idle_ttl,json=idleTtl,proto3" json:"idle_ttl,omitempty"` // Zero leaves session shares unbounded. @@ -1290,6 +1292,13 @@ func (*SessionsConfig) Descriptor() ([]byte, []int) { return file_kagent_config_v1_config_proto_rawDescGZIP(), []int{17} } +func (x *SessionsConfig) GetAutomaticRecovery() bool { + if x != nil { + return x.AutomaticRecovery + } + return false +} + func (x *SessionsConfig) GetIdleTtl() *durationpb.Duration { if x != nil { return x.IdleTtl @@ -2243,8 +2252,9 @@ const file_kagent_config_v1_config_proto_rawDesc = "" + "\x10actor_jwt_issuer\x18\x05 \x01(\tB\x1a\xbaH\x17r\x152\x10^https://[^/?#]+\x88\x01\x01H\x02R\x0eactorJwtIssuer\x88\x01\x01B\x0f\n" + "\r_api_endpointB\r\n" + "\v_router_urlB\x13\n" + - "\x11_actor_jwt_issuer\"\xf6\x02\n" + - "\x0eSessionsConfig\x12F\n" + + "\x11_actor_jwt_issuer\"\xa5\x03\n" + + "\x0eSessionsConfig\x12-\n" + + "\x12automatic_recovery\x18\x05 \x01(\bR\x11automaticRecovery\x12F\n" + "\bidle_ttl\x18\x01 \x01(\v2\x19.google.protobuf.DurationB\x10\xbaH\r\xaa\x01\n" + "\"\x06\b\x84\xfa\x85\xae\"2\x00R\aidleTtl\x12O\n" + "\rshare_max_ttl\x18\x02 \x01(\v2\x19.google.protobuf.DurationB\x10\xbaH\r\xaa\x01\n" + diff --git a/go/core/internal/a2agateway/gateway.go b/go/core/internal/a2agateway/gateway.go index 33779d0547..6d7ea9e09f 100644 --- a/go/core/internal/a2agateway/gateway.go +++ b/go/core/internal/a2agateway/gateway.go @@ -24,6 +24,7 @@ type interactionService interface { PrepareCancelTask(context.Context, types.NamespacedName, *a2a.CancelTaskRequest) (*apiv1.Session, *a2a.Task, error) PrepareTaskSubscription(context.Context, types.NamespacedName, *a2a.SubscribeToTaskRequest) (*apiv1.Session, *a2a.Task, error) RevokeSend(context.Context, types.NamespacedName, *a2a.Message, uuid.UUID) (bool, error) + RecoverRejectedSend(context.Context, types.NamespacedName, *a2a.Message) error GetTaskByMessage(context.Context, types.NamespacedName, *a2a.Message) (*a2a.Task, error) GetSendResult(context.Context, types.NamespacedName, *a2a.Message, a2a.TaskID, *int) (*a2a.Task, error) GetCancelResult(context.Context, types.NamespacedName, a2a.TaskID) (*a2a.Task, error) diff --git a/go/core/internal/a2agateway/gateway_test.go b/go/core/internal/a2agateway/gateway_test.go index 7fb90bdde4..f8e1a737f4 100644 --- a/go/core/internal/a2agateway/gateway_test.go +++ b/go/core/internal/a2agateway/gateway_test.go @@ -19,8 +19,10 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/require" "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" "google.golang.org/protobuf/proto" @@ -55,6 +57,9 @@ type gatewayTestStore struct { initialID string listedIDs []string revoked bool + reverts int + onRevert func() + revertErr error session *apiv1.Session revision *database.RuntimeRevision err error @@ -808,6 +813,81 @@ func TestGatewayReportsOnlyRevokedAttemptsAsNotAccepted(t *testing.T) { if runtime.sendCalls != 1 { t.Fatalf("retried a transport failure: %d sends", runtime.sendCalls) } + if want := map[bool]int{false: 0, true: 1}[revoked]; store.reverts != want { + t.Fatalf("reverts = %d, want %d: only a refused send checks for a crashed runtime", store.reverts, want) + } + } +} + +func TestGatewaySendJoinsAnAdmittedCrashRecovery(t *testing.T) { + session := gatewayTestSession() + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT + store := &gatewayTestStore{session: session} + runtime := &gatewayTestRuntime{} + gateway := newTestGateway(store, &gatewayTestAuthorizer{}, &gatewayTestDialer{client: gatewayTestClient(t, runtime)}, gatewayTestURL) + _, err := gateway.SendMessage(gatewayTestContext(), gatewayTestRequest()) + var protocolError *a2atype.Error + if !errors.As(err, &protocolError) || protocolError.ErrorInfo().Value["metadata"].(map[string]string)["reason"] != "KAGENT_SEND_NOT_ACCEPTED" { + t.Fatalf("error = %v, want a retryable send", err) + } + if store.reverts != 1 || runtime.sendCalls != 0 { + t.Fatalf("reverts = %d, sends = %d: the send must join the recovery, not reach the runtime", store.reverts, runtime.sendCalls) + } +} + +// A crash during a turn leaves it open, so a send refused for an open turn probes +// the runtime once. Other conflicts, and a probe that cannot run, keep the refusal. +func TestGatewayBusySendRecoversACrashedTurn(t *testing.T) { + for _, test := range []struct { + name string + reserveErr error + crashed bool + revertErr error + reverts, sent int + reason string + }{ + {name: "a busy healthy runtime", reserveErr: database.ErrActiveTurn, reverts: 1, reason: "KAGENT_TURN_ACTIVE"}, + {name: "a crashed turn", reserveErr: database.ErrActiveTurn, crashed: true, reverts: 1, sent: 1}, + {name: "a revert that Substrate cannot run", reserveErr: database.ErrActiveTurn, + revertErr: status.Error(codes.Unavailable, "control plane unavailable"), reverts: 1, reason: "KAGENT_TURN_ACTIVE"}, + {name: "a checkpoint in progress", reserveErr: database.ErrConflict}, + } { + t.Run("Given "+test.name+", when a send is refused, assert that recovery runs only for an open turn", func(t *testing.T) { + store := &gatewayTestStore{session: gatewayTestSession(), reserveErr: test.reserveErr, taskErr: database.ErrNotFound, revertErr: test.revertErr} + if test.crashed { + store.onRevert = func() { store.reserveErr = nil } + } + runtime := &gatewayTestRuntime{} + gateway := newTestGateway(store, &gatewayTestAuthorizer{}, &gatewayTestDialer{client: gatewayTestClient(t, runtime)}, gatewayTestURL) + _, err := gateway.SendMessage(gatewayTestContext(), gatewayTestRequest()) + require.Equal(t, test.reverts, store.reverts) + require.Equal(t, test.sent, runtime.sendCalls) + if test.sent == 1 { + return + } + require.ErrorIs(t, err, a2atype.ErrUnsupportedOperation) + require.Equal(t, test.reason, errorInfoReason(err), "a refusal that is not a crash must not invite a retry loop") + }) + } +} + +func TestGatewaySendToLostRuntimeFailsFromStoredState(t *testing.T) { + session := gatewayTestSession() + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE + session.Failure = &apiv1.Failure{Reason: database.SessionRuntimeLost, Message: "Runtime cannot be recovered"} + store := &gatewayTestStore{session: session} + runtime := &gatewayTestRuntime{} + gateway := newTestGateway(store, &gatewayTestAuthorizer{}, &gatewayTestDialer{client: gatewayTestClient(t, runtime)}, gatewayTestURL) + _, err := gateway.SendMessage(gatewayTestContext(), gatewayTestRequest()) + var protocolError *a2atype.Error + if !errors.As(err, &protocolError) || !errors.Is(err, a2atype.ErrUnsupportedOperation) { + t.Fatalf("error = %v, want a non-retryable refusal", err) + } + if meta, _ := protocolError.ErrorInfo().Value["metadata"].(map[string]string); meta["reason"] == "KAGENT_SEND_NOT_ACCEPTED" { + t.Fatalf("a lost runtime must not invite a retry: %v", err) + } + if store.reverts != 0 || runtime.sendCalls != 0 { + t.Fatalf("reverts = %d, sends = %d: a recorded loss is answered from PostgreSQL", store.reverts, runtime.sendCalls) } } diff --git a/go/core/internal/a2agateway/interactions.go b/go/core/internal/a2agateway/interactions.go index cda95b7e22..47b8c7792f 100644 --- a/go/core/internal/a2agateway/interactions.go +++ b/go/core/internal/a2agateway/interactions.go @@ -14,7 +14,6 @@ import ( apia2a "github.com/kagent-dev/kagent/go/api/a2a" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" - sessionsvc "github.com/kagent-dev/kagent/go/core/internal/service/session" "github.com/kagent-dev/kagent/go/pkg/logging" ) @@ -187,7 +186,7 @@ func (g *Gateway) finishSend(ctx context.Context, agent types.NamespacedName, me if err != nil { logging.FromContext(ctx).ErrorContext(ctx, "release runtime dispatch", "session_id", message.ContextID, "error", err) } else if revoked && sendErr != nil && a2atype.ErrorReason(sendErr) == a2atype.ErrorReason(a2atype.ErrInternalError) { - return sessionsvc.ErrSendNotAccepted + return g.interactions.RecoverRejectedSend(ctx, agent, message) } return sendErr } diff --git a/go/core/internal/a2agateway/routing_test.go b/go/core/internal/a2agateway/routing_test.go index 7a062f307d..e1712fbd7e 100644 --- a/go/core/internal/a2agateway/routing_test.go +++ b/go/core/internal/a2agateway/routing_test.go @@ -2,6 +2,7 @@ package a2agateway import ( "context" + "errors" "testing" a2a "github.com/a2aproject/a2a-go/v2/a2a" @@ -43,18 +44,32 @@ func (s gatewayTestAgents) Get(ctx context.Context, ref types.NamespacedName) (* return &v1alpha3.Agent{ObjectMeta: metav1.ObjectMeta{Namespace: ref.Namespace, Name: ref.Name}, Status: v1alpha3.AgentStatus{LatestSuccessfulRevision: "revision-1"}}, nil } -type gatewayTestWorkflow struct{ *sessionsvc.ActorWorkflow } +type gatewayTestWorkflow struct { + *sessionsvc.ActorWorkflow + store *gatewayTestStore +} + +func (w gatewayTestWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + w.store.reverts++ + if w.store.revertErr != nil { + return nil, w.store.revertErr + } + if w.store.onRevert != nil { + w.store.onRevert() + } + return session, nil +} func (gatewayTestWorkflow) Create(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { return session, nil } func newTestSessions(store *gatewayTestStore, authorizer auth.Authorizer) *sessionsvc.Service { - return sessionsvc.NewService(store, authorizer, gatewayTestWorkflow{}) + return sessionsvc.NewService(store, authorizer, gatewayTestWorkflow{store: store}) } func newTestGateway(store *gatewayTestStore, authorizer auth.Authorizer, dialer *gatewayTestDialer, url string) a2asrv.RequestHandler { - interactions := sessionsvc.NewInteractionService(store, gatewayTestAgents{store, authorizer}, newTestSessions(store, authorizer), false) + interactions := sessionsvc.NewInteractionService(store, gatewayTestAgents{store, authorizer}, newTestSessions(store, authorizer), false, true) return New(interactions, dialer, url) } @@ -131,3 +146,50 @@ func TestListWithoutContextRestrictsShareToOneConversation(t *testing.T) { require.NoError(t, err) require.Equal(t, []string{gatewayTestID}, store.listedIDs) } + +func TestSendRecoveryDisabled(t *testing.T) { + for _, scenario := range []struct { + name string + pending, refused bool + reserveErr error + want error + reason string + }{ + {name: "an admitted recovery", pending: true, want: a2a.ErrUnsupportedOperation}, + {name: "an open turn", reserveErr: database.ErrActiveTurn, want: a2a.ErrUnsupportedOperation, reason: "KAGENT_TURN_ACTIVE"}, + {name: "a checkpoint in progress", reserveErr: database.ErrConflict, want: a2a.ErrUnsupportedOperation}, + {name: "a runtime refusal", refused: true, want: sessionsvc.ErrSendNotAccepted, reason: "KAGENT_SEND_NOT_ACCEPTED"}, + } { + t.Run("Given automatic recovery is disabled, when a send encounters "+scenario.name+", assert that it refuses without reverting", func(t *testing.T) { + session := gatewayTestSession() + if scenario.pending { + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT + } + store := &gatewayTestStore{session: session, revoked: scenario.refused, taskErr: database.ErrNotFound, reserveErr: scenario.reserveErr} + runtime := &gatewayTestRuntime{} + if scenario.refused { + runtime.onSend = func() error { return a2a.ErrInternalError } + } + authorizer := &gatewayTestAuthorizer{} + interactions := sessionsvc.NewInteractionService(store, gatewayTestAgents{store, authorizer}, newTestSessions(store, authorizer), false, false) + gateway := New(interactions, &gatewayTestDialer{client: gatewayTestClient(t, runtime)}, gatewayTestURL) + _, err := gateway.SendMessage(gatewayTestContext(), gatewayTestRequest()) + require.ErrorIs(t, err, scenario.want) + require.Equal(t, scenario.reason, errorInfoReason(err)) + require.Zero(t, store.reverts) + if scenario.reserveErr != nil { + require.Equal(t, 1, store.reserveCalls) + } + }) + } +} + +// errorInfoReason is the kagent reason a refusal carries in its ErrorInfo metadata. +func errorInfoReason(err error) string { + refusal, ok := errors.AsType[*a2a.Error](err) + if !ok { + return "" + } + meta, _ := refusal.ErrorInfo().Value["metadata"].(map[string]string) + return meta["reason"] +} diff --git a/go/core/internal/controller/scheduledrun/controller_test.go b/go/core/internal/controller/scheduledrun/controller_test.go index 77ef1b23e7..7a28bc178b 100644 --- a/go/core/internal/controller/scheduledrun/controller_test.go +++ b/go/core/internal/controller/scheduledrun/controller_test.go @@ -338,7 +338,7 @@ func TestControllerDispatchUsesBoundedInProcessAuthority(t *testing.T) { store := &scopedControllerStore{session: &apiv1.Session{Id: id, ContextId: id, Creator: "owner", Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team", Name: "assistant"}}, State: apiv1.RuntimeState_RUNTIME_STATE_READY, Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE}} policy := &controllerSessionPolicy{} sessions := sessionsvc.NewService(store, policy, nil) - interactions := sessionsvc.NewInteractionService(store, nil, sessions, false) + interactions := sessionsvc.NewInteractionService(store, nil, sessions, false, false) ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second) defer cancel() gateway := scopedControllerGateway{interactions: interactions, cancel: cancel, agent: types.NamespacedName{Namespace: "team", Name: "assistant"}} diff --git a/go/core/internal/database/errors.go b/go/core/internal/database/errors.go index 3184e03884..8a49ee4c23 100644 --- a/go/core/internal/database/errors.go +++ b/go/core/internal/database/errors.go @@ -21,3 +21,7 @@ var ( // ErrObjectDeleting reports that deletion has started. The object is unavailable // to new callers, so it also matches ErrNotFound. var ErrObjectDeleting = fmt.Errorf("object is being deleted: %w", ErrNotFound) + +// ErrActiveTurn reports a dispatch refused because an earlier turn is still open. +// It also matches ErrConflict. +var ErrActiveTurn = fmt.Errorf("an earlier turn is still open: %w", ErrConflict) diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go index d591f33b5b..22240ba80d 100644 --- a/go/core/internal/database/session_recovery.go +++ b/go/core/internal/database/session_recovery.go @@ -18,6 +18,22 @@ import ( // recorded Actor crashed. Transport failures never set it. const SessionRuntimeCrashed = "RuntimeCrashed" +// SessionRuntimeLost is the failure reason of a FAILED Session whose recorded +// Actor is missing or was replaced. Nothing can restore it; Delete still works. +const SessionRuntimeLost = "RuntimeLost" + +// RuntimeLost reports a Session recorded as having lost its runtime. +func RuntimeLost(session *apiv1.Session) bool { + return session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && + session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE && + session.GetFailure().GetReason() == SessionRuntimeLost +} + +// LifecycleIdle reports a READY Session that no lifecycle operation owns. +func LifecycleIdle(session *apiv1.Session) bool { + return session.State == apiv1.RuntimeState_RUNTIME_STATE_READY && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE +} + // CrashRecoveryAdmitted reports an admitted, unfinished crash recovery. Only // BeginCrashedSessionRecovery starts a REVERT, so the operation alone identifies it. func CrashRecoveryAdmitted(session *apiv1.Session) bool { @@ -119,7 +135,7 @@ func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *Se if err := execSQL(ctx, tx, `UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID, row.Data); err != nil { return err } - if err := settleCrashedSessionTasks(ctx, tx, row, snapshot); err != nil { + if err := settleCrashedSessionTasks(ctx, tx, row, snapshot, "The runtime crashed before this turn completed. Send a new message to continue."); err != nil { return err } operation, err = toSessionOperation(row) @@ -128,9 +144,63 @@ func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *Se return operation, err } +// RecordLostSessionRuntime fails a Session whose recorded Actor is missing or was +// replaced, so later sends and reverts answer from PostgreSQL instead of asking +// Substrate again. operationID is the generation the loss was observed under: the +// last completed operation of a READY Session, or its admitted REVERT. A newer +// generation returns ErrConflict, because the observation may predate it. +// +// Every open task is failed, including waits a snapshot captured: no runtime +// remains to answer them. Saved results are published first. +func (c *Client) RecordLostSessionRuntime(ctx context.Context, sessionID string, operationID uuid.UUID, message string) (*apiv1.Session, error) { + var session *apiv1.Session + err := c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockSession(ctx, tx, sessionID) + if err != nil { + return notFoundOr(err) + } + session, err = toSession(row) + if err != nil { + return err + } + if RuntimeLost(session) { + return nil + } + current := uuid.Nil + if row.OperationID != nil { + current = *row.OperationID + } + if current != operationID || (!LifecycleIdle(session) && !CrashRecoveryAdmitted(session)) { + return ErrConflict + } + row.State, row.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED.String(), apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE.String() + // A new generation tells an attempt still executing the REVERT that it was + // superseded, as Delete does. + row.OperationID, row.ExecutorID = new(uuid.New()), nil + if err := saveRuntimeLifecycle(ctx, tx, row.runtimeInstanceRow, runtimeKindAgent); err != nil { + return err + } + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE + session.Failure = &apiv1.Failure{Reason: SessionRuntimeLost, Message: message} + session.UpdatedAt = timestamppb.Now() + row.Data, err = marshalSession(session) + if err != nil { + return err + } + if err := execSQL(ctx, tx, `UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID, row.Data); err != nil { + return err + } + return settleCrashedSessionTasks(ctx, tx, row, SessionTaskSnapshot{}, "The runtime was lost before this turn completed. Start another conversation to continue.") + }) + if err != nil { + return nil, fmt.Errorf("record lost Session %s runtime: %w", sessionID, err) + } + return session, nil +} + // settleCrashedSessionTasks publishes saved results, keeps waits the snapshot // captured, and fails the rest of the open work the revert discards. -func settleCrashedSessionTasks(ctx context.Context, tx pgx.Tx, session sessionRow, snapshot SessionTaskSnapshot) error { +func settleCrashedSessionTasks(ctx context.Context, tx pgx.Tx, session sessionRow, snapshot SessionTaskSnapshot, reason string) error { if err := publishSavedSessionResults(ctx, tx, session); err != nil { return err } @@ -167,7 +237,7 @@ func settleCrashedSessionTasks(ctx context.Context, tx pgx.Tx, session sessionRo return err } task.Status = a2a.TaskStatus{State: a2a.TaskStateFailed, Timestamp: new(time.Now().UTC()), - Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("The runtime crashed before this turn completed. Send a new message to continue."))} + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart(reason))} apia2a.SetTimelinePosition(task.Status.Message, *task.Status.Timestamp) event := &a2a.TaskStatusUpdateEvent{TaskID: task.ID, ContextID: task.ContextID, Status: task.Status} if err := storeSessionTaskEvent(ctx, tx, session, task, event, false); err != nil { diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go index 0412f02c07..ba0be1b17e 100644 --- a/go/core/internal/database/session_recovery_test.go +++ b/go/core/internal/database/session_recovery_test.go @@ -142,6 +142,42 @@ func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { // A crash during post-turn suspension leaves the worker's claim issued, because // the worker cannot tell whether its suspension took effect. The confirmed crash // settles that: admission retires the claim instead of waiting for it. +// A suspension whose worker was lost keeps its claim. Reservation names that +// cause so a send can recover the Session; a crash recovery retires the claim. +func TestDispatchReportsUnsettledSuspension(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + dispatch := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(t.Context(), session.Id, dispatch, "")) + busy := client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "") + require.ErrorIs(t, busy, ErrDispatchBusy) + require.NotErrorIs(t, busy, ErrDispatchSuspending, "an ordinary dispatch in flight is not a suspension") + revoked, err := client.RevokeSessionDispatch(t.Context(), session.Id, dispatch, "unused") + require.NoError(t, err) + require.True(t, revoked) + + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + _, err = client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + suspending := client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "") + require.ErrorIs(t, suspending, ErrDispatchSuspending) + require.ErrorIs(t, suspending, ErrDispatchBusy, "callers that only wait keep treating it as busy") + + uid := "actor-" + session.Id + operation, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, uid), SessionTaskSnapshot{}) + require.NoError(t, err) + finishCrashedRecovery(t, client, session.Id, operation, uid) + require.NoError(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), ""), "recovery retires the claim") +} + func TestCrashedRecoveryRetiresIssuedIdleSuspension(t *testing.T) { client := NewClient(setupTestDB(t), "public") sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") @@ -327,3 +363,58 @@ func TestRevertCanOnlyJoinAnAdmittedRecovery(t *testing.T) { require.NoError(t, err) require.Equal(t, admitted.ID, joined.ID) } + +func TestRecordLostSessionRuntime(t *testing.T) { + for _, test := range []struct { + name string + admitted bool + }{ + {name: "observed by an idle revert"}, + {name: "observed by an admitted recovery's claimed attempt", admitted: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + uid := "actor-" + session.Id + generation := sessionGeneration(t, client, session.Id).OperationID + var executor uuid.UUID + if test.admitted { + operation, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, uid), SessionTaskSnapshot{}) + require.NoError(t, err) + executor = uuid.New() + claimed, err := client.ClaimSessionOperation(t.Context(), session.Id, operation.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + _, err = client.RecordLostSessionRuntime(t.Context(), session.Id, generation, "lost") + require.ErrorIs(t, err, ErrConflict, "an observation older than the admitted recovery is stale") + generation = operation.ID + } + _, err := client.RecordLostSessionRuntime(t.Context(), session.Id, uuid.New(), "lost") + require.ErrorIs(t, err, ErrConflict, "a stale generation records nothing") + + lost, err := client.RecordLostSessionRuntime(t.Context(), session.Id, generation, "Runtime cannot be recovered") + require.NoError(t, err) + require.True(t, RuntimeLost(lost)) + require.Equal(t, "Runtime cannot be recovered", lost.Failure.Message) + stored, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.True(t, proto.Equal(lost, stored)) + settled, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, settled.Status.State, "no runtime remains to answer the wait") + + replayed, err := client.RecordLostSessionRuntime(t.Context(), session.Id, uuid.New(), "again") + require.NoError(t, err) + require.True(t, proto.Equal(stored, replayed), "recording a loss twice changes nothing") + if test.admitted { + _, err = client.FinishSessionOperation(t.Context(), session.Id, generation, executor, "", uid, "") + require.ErrorIs(t, err, ErrConflict, "the abandoned revert cannot finish") + } + require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) + deletion, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err, "a lost runtime can still be deleted") + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, deletion.Instance.Operation) + }) + } +} diff --git a/go/core/internal/database/session_tasks.go b/go/core/internal/database/session_tasks.go index a51317f2c4..29b0cb2525 100644 --- a/go/core/internal/database/session_tasks.go +++ b/go/core/internal/database/session_tasks.go @@ -23,6 +23,18 @@ import ( // or idle lifecycle operation currently owns the session. var ErrDispatchBusy = errors.New("session temporarily unavailable for dispatch") +// ErrDispatchSuspending is ErrDispatchBusy because a post-turn suspension holds its +// claim. The suspension is running, or its worker was lost and its outcome is +// unknown: the idle worker keeps the claim rather than guess, and nothing takes it +// over on a timeout. Only a crash recovery or Delete retires it, so the send path +// needs to tell this cause apart from an ordinary busy dispatch. +var ErrDispatchSuspending = fmt.Errorf("runtime suspension has not settled: %w", ErrDispatchBusy) + +// errSuspensionUnsettled is requireSettledRuntime's failure when a claimed +// suspension is what keeps the runtime unsettled. It still matches +// ErrFailedPrecondition for every other caller. +var errSuspensionUnsettled = fmt.Errorf("runtime suspension has not settled: %w", ErrFailedPrecondition) + // ErrMessageAccepted means an earlier dispatch already persisted this initial input. var ErrMessageAccepted = errors.New("message already accepted") @@ -50,18 +62,29 @@ func (c *Client) ReserveSessionDispatch(ctx context.Context, sessionID string, d if session.State != "RUNTIME_STATE_READY" || session.Operation != "RUNTIME_OPERATION_NONE" { return ErrConflict } - blocked, err := queryOne(ctx, tx, ` + activeTurn, err := queryOne(ctx, tx, ` SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND state IN ('TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING')) - OR EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $2 AND state = 'CREATING') - `, pgx.RowTo[bool], session.HistoryID, session.ID) + `, pgx.RowTo[bool], session.HistoryID) + if err != nil { + return err + } + if activeTurn { + return ErrActiveTurn + } + checkpointing, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + `, pgx.RowTo[bool], session.ID) if err != nil { return err } - if blocked { + if checkpointing { return ErrConflict } if err := requireSettledRuntime(ctx, tx, session.HistoryID, ""); err != nil { + if errors.Is(err, errSuspensionUnsettled) { + return ErrDispatchSuspending + } if errors.Is(err, ErrFailedPrecondition) { return ErrDispatchBusy } @@ -821,16 +844,23 @@ func readTaskMessages(ctx context.Context, db dbExecutor, historyID uuid.UUID, t // another unexpired dispatch. A writer may pass its own dispatch ID. // 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 { - pending, err := queryOne(ctx, db, ` + type unsettled struct { + Suspending, Pending bool + } + state, 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 quiescence_pending AND quiescence_executor_id IS NOT NULL) AS suspending, + EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND NOT published) 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) + AND ($2::text = '' OR dispatch_id::text <> $2)) AS pending + `, pgx.RowToStructByName[unsettled], historyID, dispatchID) if err != nil { return err } - if pending { + if state.Suspending { + return errSuspensionUnsettled + } + if state.Pending { return fmt.Errorf("runtime cleanup or lifecycle work is not yet settled: %w", ErrFailedPrecondition) } return nil diff --git a/go/core/internal/grpcserver/scheduledrun_controller_test.go b/go/core/internal/grpcserver/scheduledrun_controller_test.go index 244ad99763..7f6aecb60a 100644 --- a/go/core/internal/grpcserver/scheduledrun_controller_test.go +++ b/go/core/internal/grpcserver/scheduledrun_controller_test.go @@ -356,7 +356,7 @@ func TestScheduledRunControllerThroughGRPC(t *testing.T) { if tc.want == apiv1.ScheduledRunExecutionState_SCHEDULED_RUN_EXECUTION_STATE_SUCCEEDED && !tc.stream { controllerStore.loseTaskLink = true } - interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, scheduledControllerAuthorizer{}, nil), false) + interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, scheduledControllerAuthorizer{}, healthyRuntimeWorkflow{}), false, false) controller := scheduledrun.NewController(controllerStore, workflow, a2agateway.New(interactions, dialer, "http://gateway.test"), time.Second) go func() { done <- controller.Start(ctx) }() t.Cleanup(func() { cancel(); require.NoError(t, <-done) }) diff --git a/go/core/internal/grpcserver/session_recovery_test.go b/go/core/internal/grpcserver/session_recovery_test.go index 5971703f9b..926e5e171c 100644 --- a/go/core/internal/grpcserver/session_recovery_test.go +++ b/go/core/internal/grpcserver/session_recovery_test.go @@ -147,14 +147,6 @@ func TestRevertSessionThroughGRPC(t *testing.T) { require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) control.actors[name].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED before := proto.CloneOf(control.actors[name]) - control.actors[name].Metadata.Uid = "replacement-uid" - _, err = client.RevertSession(owner, request) - require.Equal(t, codes.FailedPrecondition, status.Code(err), "a same-name replacement is not this runtime") - control.actors[name].Metadata.Uid = before.Metadata.Uid - delete(control.actors, name) - _, err = client.RevertSession(owner, request) - require.Equal(t, codes.FailedPrecondition, status.Code(err), "a missing runtime cannot authorize recovery") - control.actors[name] = proto.CloneOf(before) calls := 0 failOnce := true control.revert = func(actor *ateapipb.Actor) error { @@ -202,4 +194,64 @@ func TestRevertSessionThroughGRPC(t *testing.T) { require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) }) } + + for _, test := range []struct { + name string + admitted bool + lose func(*ateapipb.Actor) *ateapipb.Actor + }{ + {name: "replaced Actor", lose: func(actor *ateapipb.Actor) *ateapipb.Actor { + actor.Metadata.Uid = "replacement-uid" + return actor + }}, + {name: "missing Actor", lose: func(*ateapipb.Actor) *ateapipb.Actor { return nil }}, + {name: "Actor lost during an admitted recovery", admitted: true, lose: func(*ateapipb.Actor) *ateapipb.Actor { return nil }}, + } { + t.Run(test.name, func(t *testing.T) { + session := createTaskStoreSession(t, store) + name := substrate.ActorName(session.Id) + original := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: name, Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team-a", Name: "assistant-kagent-revision"}, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_CRASHED}, + } + control.actors[name] = proto.CloneOf(original) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + request := &apiv1.RevertSessionRequest{SessionId: session.Id} + control.revert = nil + if test.admitted { + control.revert = func(*ateapipb.Actor) error { return status.Error(codes.Unavailable, "temporary revert outage") } + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.Unavailable, status.Code(err)) + control.revert = nil + } + if lost := test.lose(control.actors[name]); lost == nil { + delete(control.actors, name) + } + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.FailedPrecondition, status.Code(err), "a replaced or missing runtime is not this Session's") + got, err := client.GetSession(owner, &apiv1.GetSessionRequest{SessionId: session.Id}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, got.Session.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, got.Session.Operation) + require.Equal(t, database.SessionRuntimeLost, got.Session.Failure.Reason) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, stored.Status.State, "no runtime remains to answer a wait") + + // The loss is recorded: even the original Actor reappearing does not + // reopen the Session. + control.actors[name] = proto.CloneOf(original) + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString()), database.ErrConflict) + }) + } } diff --git a/go/core/internal/grpcserver/task_store_test.go b/go/core/internal/grpcserver/task_store_test.go index 24082a062d..bed8502d89 100644 --- a/go/core/internal/grpcserver/task_store_test.go +++ b/go/core/internal/grpcserver/task_store_test.go @@ -287,7 +287,7 @@ func TestRuntimeTaskStoreThroughGRPC(t *testing.T) { var httpTransport a2aclient.Transport for i := range gateways { publicListener := bufconn.Listen(DefaultMaxMessageSize) - interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true) + interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, healthyRuntimeWorkflow{}), true, false) gateway := a2agateway.New(interactions, taskStoreRuntimeDialer{runtimeListener}, "http://gateway.test") public, err := New(Config{ Listener: publicListener, SystemService: testSystemService(), @@ -443,7 +443,7 @@ func TestRuntimeTaskStoreThroughGRPC(t *testing.T) { storedConfigs, err := store.ListTaskPushConfigs(t.Context(), id, parkedID, "", 10) require.NoError(t, err) require.Len(t, storedConfigs, 2, "stored configs: %+v", storedConfigs) - managedService := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true) + managedService := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, healthyRuntimeWorkflow{}), true, false) serviceCtx := auth.AuthSessionTo(publicCtx, &authimpl.SimpleSession{P: auth.Principal{ID: "alice"}}) servicePage, err := managedService.ListTaskPushConfigs(serviceCtx, types.NamespacedName{Namespace: "team-a", Name: "assistant"}, &a2a.ListTaskPushConfigRequest{TaskID: a2a.TaskID(parkedID), PageSize: 1}) @@ -719,3 +719,20 @@ func createTaskStoreSession(t *testing.T, store *database.Client) *apiv1.Session require.NoError(t, err) return session } + +// healthyRuntimeWorkflow stands in for Substrate when a send refused for active +// work checks whether the runtime crashed: the runtime is healthy, so Revert +// returns the Session unchanged and the send is refused as before. +type healthyRuntimeWorkflow struct{} + +func (healthyRuntimeWorkflow) Create(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, nil +} + +func (healthyRuntimeWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, nil +} + +func (healthyRuntimeWorkflow) Delete(_ context.Context, session *apiv1.Session, _ bool) (*apiv1.Session, error) { + return session, nil +} diff --git a/go/core/internal/service/session/expiration_test.go b/go/core/internal/service/session/expiration_test.go index 4acd1e36dc..49483a7c67 100644 --- a/go/core/internal/service/session/expiration_test.go +++ b/go/core/internal/service/session/expiration_test.go @@ -107,7 +107,7 @@ func TestSessionExpirationRetainsCheckpointHistory(t *testing.T) { store, source := lifecycleFixture(t) actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} _, checkpointID := lifecycleForkFixture(t, store, actors, source) - interactions := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false) + interactions := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false, false) agent := types.NamespacedName{Namespace: source.GetAgent().Ref.Namespace, Name: source.GetAgent().Ref.Name} // Real ownership filtering must hide an existing conversation in exactly // the same way as a missing or expired conversation, without reading history. diff --git a/go/core/internal/service/session/interactions.go b/go/core/internal/service/session/interactions.go index 464db60622..6bb985ae7e 100644 --- a/go/core/internal/service/session/interactions.go +++ b/go/core/internal/service/session/interactions.go @@ -48,14 +48,25 @@ type agentService interface { // state. It never connects to actors: the gateway owns dispatch and observation. // These operations form the persistence boundary for a separate API service. type InteractionService struct { - store interactionStore - agents agentService - sessions *Service - allowPushHTTP bool + store interactionStore + agents agentService + sessions *Service + allowPushHTTP bool + automaticRecovery bool + // recoveryWait bounds how long a send waits on the recovery it starts. + recoveryWait time.Duration + // dispatchWait bounds how long a send waits for a busy Session. + dispatchWait time.Duration } -func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP bool) *InteractionService { - return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP} +// sendRecoveryWait bounds a send's wait for crash recovery. A recovery still +// running when it ends stays admitted: the attempt releases its claim, and the +// client's retry continues it. +const sendRecoveryWait = 10 * time.Second + +func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP, automaticRecovery bool) *InteractionService { + return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP, automaticRecovery: automaticRecovery, + recoveryWait: sendRecoveryWait, dispatchWait: 10 * time.Second} } // PreparedSend either grants a new dispatch or returns an already accepted task. @@ -162,7 +173,46 @@ func (s *InteractionService) PrepareSend(ctx context.Context, agent types.Namesp config.PushConfig = nil req.Config = &config } - id, err := s.reserveDispatch(ctx, session.Id, initialID) + id, err := s.reserveDispatch(ctx, session.Id, initialID, s.dispatchWait) + if s.automaticRecovery && errors.Is(err, errSuspensionUnsettled) { + // Why: a worker lost while the Actor suspends after a turn leaves the + // Actor CRASHED, the Session READY, and the suspension's claim held, since + // its outcome is unknown and nothing takes the claim over on a timeout. + // Every later send then waits out the busy period and is told to retry, + // so the conversation never recovers on its own. A crash recovery is the + // only thing that retires that claim. Reaching here means the claim + // outlasted the whole wait, which a healthy suspension does not, so the + // send tries the bounded recovery once. Revert leaves a still-running + // suspension alone, and one more reservation, without another wait, + // proceeds or answers retry as before. + switch recoverErr := s.recoverForSend(ctx, session.Id); { + case recoverErr == nil: + id, err = s.reserveDispatch(ctx, session.Id, initialID, 0) + case !errors.Is(recoverErr, errRecoveryUnavailable): + return nil, recoverErr + } + } + if s.automaticRecovery && errors.Is(err, database.ErrActiveTurn) { + // A crash during a turn leaves that turn open, and an open turn refuses + // every send, so no send would reach the crashed Actor to recover it. + // Revert restores a crashed Actor and leaves a busy one alone, so one more + // reservation either proceeds or is refused as before. + switch recoverErr := s.recoverForSend(ctx, session.Id); { + case recoverErr == nil: + id, err = s.reserveDispatch(ctx, session.Id, initialID, s.dispatchWait) + case !errors.Is(recoverErr, errRecoveryUnavailable): + return nil, recoverErr + } + } + if errors.Is(err, errSuspensionUnsettled) { + err = ErrSendNotAccepted + } + switch { + case errors.Is(err, database.ErrActiveTurn): + err = errTurnActive + case errors.Is(err, database.ErrConflict): + err = interactionStoreError(ctx, err) + } if errors.Is(err, database.ErrMessageAccepted) { var historyLength *int if req.Config != nil { @@ -213,34 +263,56 @@ func (s *InteractionService) prepareTask(ctx context.Context, agent types.Namesp // reserveDispatch waits only before forwarding input. No runtime send is retried // on an ambiguous transport error, and no SQL lock is held during dispatch. -func (s *InteractionService) reserveDispatch(ctx context.Context, sessionID, initialID string) (uuid.UUID, error) { +// It waits at most wait for a busy Session, answering ErrSendNotAccepted after +// it, or errSuspensionUnsettled when a claimed suspension was still the cause. +func (s *InteractionService) reserveDispatch(ctx context.Context, sessionID, initialID string, wait time.Duration) (uuid.UUID, error) { id := uuid.New() - deadline := time.NewTimer(10 * time.Second) - defer deadline.Stop() + deadline := time.Now().Add(wait) for { err := s.store.ReserveSessionDispatch(ctx, sessionID, id, initialID) if err == nil { return id, nil } - if errors.Is(err, database.ErrMessageAccepted) { + if errors.Is(err, database.ErrMessageAccepted) || errors.Is(err, database.ErrConflict) { return uuid.Nil, err } if !errors.Is(err, database.ErrDispatchBusy) { return uuid.Nil, interactionStoreError(ctx, err) } - timer := time.NewTimer(100 * time.Millisecond) + if !time.Now().Before(deadline) { + if errors.Is(err, database.ErrDispatchSuspending) { + return uuid.Nil, errSuspensionUnsettled + } + return uuid.Nil, ErrSendNotAccepted + } + timer := time.NewTimer(min(100*time.Millisecond, time.Until(deadline))) select { case <-ctx.Done(): timer.Stop() return uuid.Nil, ctx.Err() - case <-deadline.C: - timer.Stop() - return uuid.Nil, ErrSendNotAccepted case <-timer.C: } } } +// errSuspensionUnsettled is reserveDispatch's answer when a claimed post-turn +// suspension still held the Session at the end of the wait. PrepareSend recovers +// it or answers ErrSendNotAccepted, so it never reaches a client. +var errSuspensionUnsettled = errors.New("runtime suspension has not settled") + +// errTurnActive refuses a send while an earlier turn is still open. A runtime +// that crashed mid-turn leaves its turn open, so unless the send recovers it, +// this is how the crash reaches the client. The reason lets a client offer +// RevertSession for this refusal and no other; the message is unchanged. +var errTurnActive = a2atype.NewError(a2atype.ErrUnsupportedOperation, database.ErrConflict.Error()). + WithErrorInfoMeta(map[string]string{"reason": "KAGENT_TURN_ACTIVE"}) + +// ErrRuntimeRecovering answers a send while a crash recovery is admitted but not +// finished. It keeps ErrSendNotAccepted's reason, so existing clients retry it, +// and adds state RECOVERING with a longer retry delay. +var ErrRuntimeRecovering = a2atype.NewError(a2atype.ErrUnsupportedOperation, "input was not accepted; the runtime is recovering from a crash"). + WithErrorInfoMeta(map[string]string{"reason": "KAGENT_SEND_NOT_ACCEPTED", "state": "RECOVERING", "retryAfterMs": "1000"}) + // ErrSendNotAccepted tells callers they may retry this input because no dispatch // occurred, or an unused dispatch was revoked before the runtime persisted it. var ErrSendNotAccepted = a2atype.NewError(a2atype.ErrUnsupportedOperation, "input was not accepted; retry after the session becomes available"). @@ -260,6 +332,61 @@ func (s *InteractionService) RevokeSend(ctx context.Context, agent types.Namespa return revoked, nil } +// RecoverRejectedSend runs after the runtime refused a send it never accepted. A +// crashed Actor cannot be resumed, so the refusal would repeat on every retry; +// reverting it first lets the caller's retry reach the restored runtime. +func (s *InteractionService) RecoverRejectedSend(ctx context.Context, agent types.NamespacedName, message *a2atype.Message) error { + if !s.automaticRecovery { + return ErrSendNotAccepted + } + session, err := s.messageSession(ctx, agent, message) + if err != nil { + return err + } + if err := s.recoverForSend(ctx, session.Id); err != nil && !errors.Is(err, errRecoveryUnavailable) { + return err + } + return ErrSendNotAccepted +} + +// ErrRuntimeUnrecoverable answers a send to a Session whose recorded Actor is +// missing or was replaced. Retrying cannot succeed. +var ErrRuntimeUnrecoverable = a2atype.NewError(a2atype.ErrUnsupportedOperation, "Session runtime cannot be recovered; start another conversation") + +// errRecoveryUnavailable reports that a send's recovery did not run and none is +// admitted, so the send keeps the answer it would give without recovery. +var errRecoveryUnavailable = errors.New("crash recovery unavailable for this send") + +// recoverForSend reverts a crashed runtime for a send that found it, waiting at +// most recoveryWait. It returns nil once the runtime is healthy or restored, and +// ErrRuntimeUnrecoverable for a lost one. Otherwise the stored Session decides: an +// admitted recovery answers ErrRuntimeRecovering, and anything else, including no +// update access, contention, or an outage before admission, answers +// errRecoveryUnavailable. Reverting needs RevertSession's update access, which +// sending alone does not grant. The attempt ignores the caller's cancellation, so +// a disconnect does not cut it short. +func (s *InteractionService) recoverForSend(ctx context.Context, sessionID string) error { + attempt, cancel := context.WithTimeout(context.WithoutCancel(ctx), s.recoveryWait) + defer cancel() + _, err := s.sessions.Revert(attempt, sessionID) + if err == nil { + return nil + } + if errors.Is(err, errRuntimeUnrecoverable) { + return ErrRuntimeUnrecoverable + } + logging.FromContext(ctx).WarnContext(ctx, "recover crashed runtime for a send", "session_id", sessionID, "error", err) + // The caller was already authorized for this Session; this read only chooses + // which refusal to return. + readCtx, cancelRead := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancelRead() + current, readErr := s.sessions.store.GetSessionByID(readCtx, sessionID) + if readErr == nil && database.CrashRecoveryAdmitted(current) { + return ErrRuntimeRecovering + } + return errRecoveryUnavailable +} + // GetTaskByMessage reads only the task that accepted this input. It uses the // send's permissions and Session scope, so recovery cannot expose another task. func (s *InteractionService) GetTaskByMessage(ctx context.Context, agent types.NamespacedName, message *a2atype.Message) (*a2atype.Task, error) { @@ -437,6 +564,15 @@ func (s *InteractionService) resolveSend(ctx context.Context, agent types.Namesp return nil, a2atype.NewError(a2atype.ErrInvalidRequest, "message context does not match task") } req.Message.ContextID = session.ContextId + if database.RuntimeLost(session) { + return nil, ErrRuntimeUnrecoverable + } + if s.automaticRecovery && database.CrashRecoveryAdmitted(session) { + if err := s.recoverForSend(ctx, session.Id); err != nil && !errors.Is(err, errRecoveryUnavailable) { + return nil, err + } + return nil, ErrSendNotAccepted + } if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return nil, a2atype.NewError(a2atype.ErrUnsupportedOperation, "Session cannot accept work during a lifecycle operation") } diff --git a/go/core/internal/service/session/interactions_test.go b/go/core/internal/service/session/interactions_test.go index 21421bc391..2c954f687d 100644 --- a/go/core/internal/service/session/interactions_test.go +++ b/go/core/internal/service/session/interactions_test.go @@ -59,7 +59,7 @@ func TestPrepareSendStampsInboundTimeline(t *testing.T) { if test.accepted { tasks.dispatchErr = database.ErrMessageAccepted } - service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false) + service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false, false) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("Approved")) message.ID, message.ContextID, message.TaskID = original.ID, id, test.taskID apia2a.SetTimelinePosition(message, time.Unix(1, 0)) @@ -99,7 +99,7 @@ func TestPrepareSendReservesTaskOpeningMessages(t *testing.T) { Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: agent.Namespace, Name: agent.Name}}} tasks := &interactionTestStore{task: &a2atype.Task{ID: "task", ContextID: id, Status: a2atype.TaskStatus{State: a2atype.TaskStateInputRequired}}} - service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false) + service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false, false) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) message.ID, message.ContextID, message.TaskID = "opening", id, test.taskID _, err := service.PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) @@ -244,7 +244,7 @@ func TestInteractionsEnforcePermissionsWithoutGateway(t *testing.T) { sessionStore.getResult = nil } tasks := &interactionTestStore{task: &a2atype.Task{ID: "task", ContextID: id, Status: a2atype.TaskStatus{State: a2atype.TaskStateCompleted}}} - service := NewInteractionService(tasks, nil, NewService(sessionStore, authorizer, nil), false) + service := NewInteractionService(tasks, nil, NewService(sessionStore, authorizer, nil), false, false) want := test.want _, shared := auth.ShareContextFrom(test.ctx) if shared && operation.verb != auth.VerbGet { @@ -273,7 +273,7 @@ func TestInteractionsEnforcePermissionsWithoutGateway(t *testing.T) { } func TestInteractionsRequireAgentForUnfilteredTaskList(t *testing.T) { - service := NewInteractionService(nil, nil, nil, false) + service := NewInteractionService(nil, nil, nil, false, false) for _, agent := range []types.NamespacedName{ {}, {Namespace: "team-a"}, @@ -298,7 +298,7 @@ func TestExistingConversationDoesNotRequireCurrentAgentGrant(t *testing.T) { ctx := serviceTestContext("alice") message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")) message.ContextID = id - _, err := NewInteractionService(nil, nil, sessions, false).messageSession(ctx, agent, message) + _, err := NewInteractionService(nil, nil, sessions, false, false).messageSession(ctx, agent, message) if denySession { if !errors.Is(err, a2atype.ErrUnauthorized) { t.Fatalf("error = %v, want Session denial", err) @@ -334,7 +334,7 @@ func TestScheduledMessageIDPrefixIsReserved(t *testing.T) { sessions := NewService(&serviceTestStore{getResult: stored}, &recordingAuthorizer{}, serviceTestWorkflow{}) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")) message.ID, message.ContextID = test.id, id - _, err := NewInteractionService(nil, nil, sessions, false).messageSession(test.ctx, agent, message) + _, err := NewInteractionService(nil, nil, sessions, false, false).messageSession(test.ctx, agent, message) if test.wantOK { require.NoError(t, err) } else { @@ -350,7 +350,7 @@ func TestScheduledMessageIDPrefixCannotStartAConversation(t *testing.T) { message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) message.ID = scheduling.MessageIDPrefix + uuid.NewString() agent := types.NamespacedName{Namespace: "team-a", Name: "assistant"} - _, err := NewInteractionService(nil, nil, sessions, false).PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) + _, err := NewInteractionService(nil, nil, sessions, false, false).PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) require.ErrorIs(t, err, a2atype.ErrInvalidParams) require.Nil(t, store.createInput) } diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go index 7d348775bf..a88569f39b 100644 --- a/go/core/internal/service/session/recovery.go +++ b/go/core/internal/service/session/recovery.go @@ -4,7 +4,9 @@ import ( "context" "errors" "fmt" + "time" + "github.com/google/uuid" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -18,21 +20,26 @@ import ( // Substrate and database failures must not wrap it, so they stay retryable. var errRuntimeUnrecoverable = fmt.Errorf("runtime cannot be recovered: %w", database.ErrFailedPrecondition) -func lifecycleIdle(session *apiv1.Session) bool { - return session.State == apiv1.RuntimeState_RUNTIME_STATE_READY && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE -} - type crashObservation struct { generation *database.SessionGeneration actor *ateapipb.Actor } +// actorLostError reports why the recorded Actor cannot be recovered, and the +// generation it was observed under. +type actorLostError struct { + operationID uuid.UUID + reason string +} + +func (e *actorLostError) Error() string { return "runtime is lost because " + e.reason } + func (o crashObservation) crashed() bool { return o.actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED } -// observeCrash treats a missing, replaced, or retemplated Actor as unrecoverable, -// never as a crash: reverting it would restore the wrong runtime. +// observeCrash reports a missing, replaced, or retemplated Actor as an +// actorLostError, never as a crash: reverting it would restore the wrong runtime. // // The generation is read before the Actor. A concurrent revert keeps the Actor UID // and returns to READY, so only the generation shows that an observation predates @@ -50,25 +57,46 @@ func (w *ActorWorkflow) observeCrash(ctx context.Context, session *apiv1.Session TemplateAtespace: generation.TemplateAtespace, TemplateName: generation.TemplateName} actor, err := w.actors.GetActor(ctx, binding.Atespace, binding.Name) if status.Code(err) == codes.NotFound { - return crashObservation{}, fmt.Errorf("runtime Actor is missing: %w", errRuntimeUnrecoverable) + return crashObservation{}, &actorLostError{operationID: generation.OperationID, reason: "its Actor is missing"} } if err != nil { return crashObservation{}, err } if !binding.Matches(actor) { - return crashObservation{}, fmt.Errorf("runtime Actor identity or template changed: %w", errRuntimeUnrecoverable) + return crashObservation{}, &actorLostError{operationID: generation.OperationID, reason: "its Actor identity or template changed"} } if actor.GetMetadata().GetUid() != generation.ActorUID { - return crashObservation{}, fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable) + return crashObservation{}, &actorLostError{operationID: generation.OperationID, reason: "its Actor was replaced"} } return crashObservation{generation: generation, actor: actor}, nil } +// recordLost stores a lost runtime so later sends and reverts fail from +// PostgreSQL. A newer generation makes the observation stale, so that conflict +// stays retryable rather than reporting the runtime unrecoverable. +func (w *ActorWorkflow) recordLost(ctx context.Context, sessionID string, operationID uuid.UUID, reason string) error { + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + message := "Runtime cannot be recovered because " + reason + "; start another conversation" + _, err := w.store.RecordLostSessionRuntime(ctx, sessionID, operationID, message) + if errors.Is(err, database.ErrNotFound) { + return fmt.Errorf("runtime Session was deleted: %w", errRuntimeUnrecoverable) + } + if err != nil { + return fmt.Errorf("runtime is lost because %s, but recording it failed: %w", reason, err) + } + return fmt.Errorf("%s: %w", message, errRuntimeUnrecoverable) +} + // Revert restores a confirmed crashed runtime without replaying its interrupted turn. // Nothing retries an admitted recovery in the background; the caller retries until // it succeeds. A runtime that is not crashed, including after a revert whose -// response was lost, is returned unchanged. +// response was lost, is returned unchanged. A lost runtime is recorded once and +// then answered from the stored Session. func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*apiv1.Session, error) { + if database.RuntimeLost(session) { + return nil, fmt.Errorf("%s: %w", session.GetFailure().GetMessage(), errRuntimeUnrecoverable) + } if database.CrashRecoveryAdmitted(session) { return w.run(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, false) } @@ -76,13 +104,16 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return nil, database.ErrConflict } - if !lifecycleIdle(session) { + if !database.LifecycleIdle(session) { return session, nil } // run bounds a joined attempt. Admission bounds its own reads and execution. ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) defer cancel() observed, err := w.observeCrash(ctx, session) + if lost, ok := errors.AsType[*actorLostError](err); ok { + return nil, w.recordLost(ctx, session.Id, lost.operationID, lost.reason) + } if err != nil { return nil, err } diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 2f0d05dcc5..90525bbe79 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -15,10 +15,13 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" + "k8s.io/apimachinery/pkg/types" 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/service/serviceerrors" "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" ) func (a *lifecycleTestActors) RevertActor(_ context.Context, atespace, name string) (*ateapipb.Actor, error) { @@ -238,10 +241,18 @@ func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { } current, err := store.GetSessionByID(t.Context(), session.Id) require.NoError(t, err) - require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) - require.Nil(t, current.Failure) stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) require.NoError(t, err) + if test.want != nil { + // The loss is recorded so later sends and reverts need not ask Substrate. + require.True(t, database.RuntimeLost(current), "state = %s/%s", current.State, current.Operation) + require.Equal(t, a2a.TaskStateFailed, stored.Status.State, "no runtime remains to finish the turn") + _, err = NewActorWorkflow(store, &recoveryTestActors{lifecycleTestActors: base, readErr: status.Error(codes.Internal, "Substrate must not be read")}, make(chan struct{}, 1), time.Second).Revert(t.Context(), current) + require.ErrorIs(t, err, errRuntimeUnrecoverable) + return + } + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Nil(t, current.Failure) require.Equal(t, a2a.TaskStateWorking, stored.Status.State) }) } @@ -499,7 +510,8 @@ func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { } // A revert restores the recorded Actor in place and creates nothing Delete could -// miss, so Delete supersedes an admitted recovery even after the revert was issued. +// miss, so Delete supersedes an admitted recovery even after the revert was issued, +// and deletes a Session whose recovery ended with a lost Actor. func TestDeleteSupersedesCrashRecovery(t *testing.T) { for _, test := range []struct { name string @@ -553,7 +565,11 @@ func TestDeleteSupersedesCrashRecovery(t *testing.T) { pending, err := store.GetSessionByID(t.Context(), session.Id) require.NoError(t, err) require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, pending.State) - require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, pending.Operation, "a failed attempt keeps the recovery for a retry") + if test.gone || test.replaced { + require.True(t, database.RuntimeLost(pending), "a lost Actor ends the recovery") + } else { + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, pending.Operation, "a failed attempt keeps the recovery for a retry") + } deleted, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Delete(t.Context(), pending, false) if test.replaced { require.NotErrorIs(t, err, database.ErrConflict) @@ -572,3 +588,186 @@ func TestDeleteSupersedesCrashRecovery(t *testing.T) { }) } } + +func TestSendRecoveryAnswersTheSend(t *testing.T) { + unavailable := status.Error(codes.Unavailable, "control plane unavailable") + denied := serviceerrors.NewPermissionDenied("denied", nil) + for _, test := range []struct { + name string + err error + authErr error + admitted bool + want error + }{ + {name: "the runtime is healthy or restored"}, + {name: "the runtime is lost", err: fmt.Errorf("runtime Actor is missing: %w", errRuntimeUnrecoverable), want: ErrRuntimeUnrecoverable}, + {name: "another operation holds the Session", err: database.ErrConflict, want: errRecoveryUnavailable}, + {name: "Substrate is unavailable before admission", err: unavailable, want: errRecoveryUnavailable}, + {name: "Substrate is unavailable after admission", err: unavailable, admitted: true, want: ErrRuntimeRecovering}, + // Sending needs only create access; reverting needs RevertSession's update access. + {name: "the caller cannot revert", authErr: denied, want: errRecoveryUnavailable}, + {name: "the caller cannot revert an admitted recovery", authErr: denied, admitted: true, want: ErrRuntimeRecovering}, + } { + t.Run("Given "+test.name+", when a send recovers, assert that it answers accordingly", func(t *testing.T) { + session := &apiv1.Session{Id: uuid.NewString(), State: apiv1.RuntimeState_RUNTIME_STATE_READY} + if test.admitted { + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT + } + authorizer := &verbAuthorizer{denied: auth.VerbUpdate, err: test.authErr} + sessions := NewService(&serviceTestStore{getResult: session}, authorizer, serviceTestWorkflow{err: test.err}) + interactions := NewInteractionService(nil, nil, sessions, false, true) + err := interactions.recoverForSend(serviceTestContext("alice"), session.Id) + if test.want == nil { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, test.want) + } + require.Equal(t, []auth.Verb{auth.VerbUpdate}, authorizer.checked, "recovery is authorized as RevertSession") + }) + } +} + +// A revert that outlasts the send's wait is cut off and answered as recovering; +// the admitted operation stays for the client's retry to continue. +func TestSendRecoveryIsBounded(t *testing.T) { + session := &apiv1.Session{Id: uuid.NewString(), State: apiv1.RuntimeState_RUNTIME_STATE_FAILED, + Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT} + workflow := &blockingRevertWorkflow{stopped: make(chan error, 1)} + interactions := NewInteractionService(nil, nil, NewService(&serviceTestStore{getResult: session}, serviceTestAuthorizer{}, workflow), false, true) + interactions.recoveryWait = 50 * time.Millisecond + ctx, cancel := context.WithCancel(serviceTestContext("alice")) + defer cancel() + started := time.Now() + err := interactions.recoverForSend(ctx, session.Id) + require.ErrorIs(t, err, ErrRuntimeRecovering) + require.Less(t, time.Since(started), 5*time.Second, "the send answers within its wait") + require.ErrorIs(t, <-workflow.stopped, context.DeadlineExceeded, "the attempt is cut off by the wait, not the caller") +} + +type blockingRevertWorkflow struct { + serviceTestWorkflow + stopped chan error +} + +func (w *blockingRevertWorkflow) Revert(ctx context.Context, _ *apiv1.Session) (*apiv1.Session, error) { + <-ctx.Done() + w.stopped <- ctx.Err() + return nil, fmt.Errorf("revert attempt stopped: %w", ctx.Err()) +} + +// verbAuthorizer denies one verb with err, so a test can grant send but not revert. +type verbAuthorizer struct { + denied auth.Verb + err error + checked []auth.Verb +} + +func (a *verbAuthorizer) Check(_ context.Context, _ auth.Principal, verb auth.Verb, _ auth.Resource) error { + a.checked = append(a.checked, verb) + if verb == a.denied { + return a.err + } + return nil +} + +// A refused send whose Session can no longer be read answers with that error, so +// a deleted Session or a revoked grant does not become an endless retry. +func TestRecoverRejectedSendReturnsSessionErrors(t *testing.T) { + for _, test := range []struct { + name string + store *serviceTestStore + authorizer serviceTestAuthorizer + want error + }{ + {name: "not found", store: &serviceTestStore{getErr: database.ErrNotFound}, want: a2a.ErrTaskNotFound}, + {name: "permission denied", store: &serviceTestStore{}, authorizer: serviceTestAuthorizer{err: serviceerrors.NewPermissionDenied("denied", nil)}, want: a2a.ErrUnauthorized}, + } { + t.Run(test.name, func(t *testing.T) { + interactions := NewInteractionService(nil, nil, NewService(test.store, test.authorizer, serviceTestWorkflow{}), false, true) + message := &a2a.Message{ID: "input", ContextID: uuid.NewString()} + err := interactions.RecoverRejectedSend(serviceTestContext("alice"), types.NamespacedName{Namespace: "team-a", Name: "assistant"}, message) + require.ErrorIs(t, err, test.want) + require.NotErrorIs(t, err, ErrSendNotAccepted) + }) + } +} + +// suspensionTestStore answers reservations with reserveErr until a revert clears it. +type suspensionTestStore struct { + *database.Client + reserveErr error + reserves int +} + +func (s *suspensionTestStore) ReserveSessionDispatch(context.Context, string, uuid.UUID, string) error { + s.reserves++ + return s.reserveErr +} + +// suspensionTestWorkflow reverts a crashed runtime by retiring the held claim. +type suspensionTestWorkflow struct { + serviceTestWorkflow + store *suspensionTestStore + crashed bool + reverts int +} + +func (w *suspensionTestWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + w.reverts++ + if w.err != nil { + return nil, w.err + } + if w.crashed { + w.store.reserveErr = nil + } + return session, nil +} + +// A worker lost during the post-turn suspension leaves the Actor CRASHED and the +// suspension's claim held, so every send waits and is told to retry. Only a send +// that outlasted the wait on that claim, with automatic recovery on, recovers it. +func TestSendRecoversACrashDuringSuspension(t *testing.T) { + for _, test := range []struct { + name string + disabled bool + reserveErr error + crashed bool + revertErr error + reverts int + dispatched bool + }{ + {name: "a crash during suspension", reserveErr: database.ErrDispatchSuspending, crashed: true, reverts: 1, dispatched: true}, + {name: "a suspension still running", reserveErr: database.ErrDispatchSuspending, reverts: 1}, + {name: "a revert that cannot run", reserveErr: database.ErrDispatchSuspending, crashed: true, + revertErr: status.Error(codes.Unavailable, "control plane unavailable"), reverts: 1}, + {name: "automatic recovery off", disabled: true, reserveErr: database.ErrDispatchSuspending, crashed: true}, + {name: "an ordinary dispatch in flight", reserveErr: database.ErrDispatchBusy, crashed: true}, + } { + t.Run("Given "+test.name+", when a send outlasts the busy wait, assert that it recovers only a crashed suspension", func(t *testing.T) { + id := uuid.NewString() + agent := types.NamespacedName{Namespace: "team-a", Name: "assistant"} + stored := &apiv1.Session{Id: id, ContextId: id, State: apiv1.RuntimeState_RUNTIME_STATE_READY, + Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, + Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: agent.Namespace, Name: agent.Name}}} + store := &suspensionTestStore{reserveErr: test.reserveErr} + workflow := &suspensionTestWorkflow{serviceTestWorkflow: serviceTestWorkflow{err: test.revertErr}, store: store, crashed: test.crashed} + interactions := NewInteractionService(store, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, workflow), false, !test.disabled) + const wait = 200 * time.Millisecond + interactions.dispatchWait = wait + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("What is 2+2?")) + message.ContextID = id + started := time.Now() + prepared, err := interactions.PrepareSend(serviceTestContext("alice"), agent, &a2a.SendMessageRequest{Message: message}) + elapsed := time.Since(started) + require.Equal(t, test.reverts, workflow.reverts) + if test.dispatched { + require.NoError(t, err) + require.NotEqual(t, uuid.Nil, prepared.DispatchID) + return + } + require.ErrorIs(t, err, ErrSendNotAccepted, "the client still sees an ordinary retryable refusal") + require.GreaterOrEqual(t, elapsed, wait, "a send first waits out the busy period") + require.Less(t, elapsed, 2*wait, "the reservation after a recovery attempt does not wait again") + }) + } +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index 14e3c61783..e47a10089e 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -17,6 +17,7 @@ import ( type workflowStore interface { GetSessionGeneration(context.Context, string) (*database.SessionGeneration, error) BeginCrashedSessionRecovery(context.Context, *database.SessionGeneration, database.SessionTaskSnapshot) (*database.SessionOperation, error) + RecordLostSessionRuntime(context.Context, string, uuid.UUID, string) (*apiv1.Session, error) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) HasPendingSessionQuiescence(context.Context) (bool, error) FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error @@ -224,9 +225,10 @@ func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database. if database.CrashRecoveryAdmitted(admitted.Instance) { // Admission already failed the interrupted turn, so clearing the operation // would leave a FAILED Session that a retried RevertSession could not - // continue. Delete can supersede the kept intent, even after a claim. + // continue. Delete can supersede the kept intent, even after a claim. A lost + // Actor ends the recovery, so it is recorded instead. if errors.Is(cause, substrate.ErrActorUnrecoverable) { - return nil, fmt.Errorf("%w: %w", cause, errRuntimeUnrecoverable) + return nil, fmt.Errorf("%w: %w", cause, w.recordLost(ctx, admitted.Instance.Id, admitted.ID, "its Actor is missing or was replaced")) } return nil, cause } diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index d992c9a2a9..3ee158d04c 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -384,7 +384,7 @@ func Run(ctx context.Context, input *configv1.ControllerConfig, opts Options) er return err } agents := kubecrud.NewService(manager.GetClient(), authorizer, &kagentv1alpha3.Agent{}, &kagentv1alpha3.AgentList{}, "Agent") - interactions := sessionsvc.NewInteractionService(store, agents, sessions, cfg.GetPush().GetAllowHttp()) + interactions := sessionsvc.NewInteractionService(store, agents, sessions, cfg.GetPush().GetAllowHttp(), cfg.GetSessions().GetAutomaticRecovery()) pushSender := sessionsvc.NewHTTPPushSender(cfg.GetPush().GetSendTimeout().AsDuration(), cfg.GetPush().GetAllowHttp(), cfg.GetPush().GetAllowPrivateNetworks()) pushIssuer := cfg.GetPush().GetIssuer() diff --git a/go/core/pkg/config/helm_test.go b/go/core/pkg/config/helm_test.go index 8c49153290..ccf972ee93 100644 --- a/go/core/pkg/config/helm_test.go +++ b/go/core/pkg/config/helm_test.go @@ -71,6 +71,11 @@ func TestHelmControllerConfig(t *testing.T) { require.Equal(t, "team", base.Kubernetes.GetNamespace()) require.Equal(t, "http://demo-controller.team.svc.cluster.local:8083", base.RuntimeDefaults.GetControlPlaneUrl()) require.NotEmpty(t, checksum) + require.False(t, base.GetSessions().GetAutomaticRecovery()) + recovery, recoveryChecksum, err := render("config:\n sessions: {automaticRecovery: true}\n") + require.NoError(t, err) + require.True(t, recovery.GetSessions().GetAutomaticRecovery()) + require.NotEqual(t, checksum, recoveryChecksum) require.Nil(t, base.GetServer().GetAuth().GetServiceAccount()) machine, _, err := render("config:\n server:\n auth:\n serviceAccount: {authority: cluster-a}\n") require.NoError(t, err) diff --git a/go/core/pkg/config/load_test.go b/go/core/pkg/config/load_test.go index 89d2dc2566..c0c6665068 100644 --- a/go/core/pkg/config/load_test.go +++ b/go/core/pkg/config/load_test.go @@ -211,3 +211,23 @@ func TestYAMLConfigReadsOnlyFirstDocument(t *testing.T) { require.NoError(t, config.Decode([]byte("logging: {level: debug}\n---\n["), cfg)) require.Equal(t, "debug", cfg.GetLogging().GetLevel()) } + +func TestSessionAutomaticRecovery(t *testing.T) { + for _, scenario := range []struct { + name, input string + want bool + }{ + {"omitted", "{}", false}, + {"disabled", "sessions: {automaticRecovery: false}", false}, + {"enabled", "sessions: {automaticRecovery: true}", true}, + } { + t.Run("Given controller config with automatic recovery "+scenario.name+", when it loads, assert that the setting is preserved", func(t *testing.T) { + cfg := &configv1.ControllerConfig{} + require.NoError(t, config.Decode([]byte(scenario.input), cfg)) + config.DefaultController(cfg) + setLogins(cfg) + require.NoError(t, config.Validate(cfg)) + require.Equal(t, scenario.want, cfg.GetSessions().GetAutomaticRecovery()) + }) + } +} diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index 34e3c4d8d2..24f811985f 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -35,6 +35,9 @@ Once a cluster API URL is set, missing `KAGENT_E2E_RUNTIME_IMAGE` fails the test before provisioning. The same digest-pinned Go ADK image supplies both the agent runtime and sandbox tools. +`TestSessionSendRecoversCrashedRuntime` requires `config.sessions.automaticRecovery: true` +in the test install. CI enables it to exercise recovery through ordinary sends. + `TestSessionIdleExpiration` runs its harness cases in parallel so they share the idle period and expiration sweep. Its parent stays sequential to isolate the controller rollouts from the rest of the suite. It temporarily sets diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index 8a77a0d16c..50909486ee 100644 --- a/go/core/test/e2e/interaction_test.go +++ b/go/core/test/e2e/interaction_test.go @@ -22,7 +22,6 @@ import ( "time" a2atype "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/a2apb/v1/pbconv" "github.com/google/uuid" @@ -196,7 +195,7 @@ func TestSessionAskUserSurvivesSuspension(t *testing.T) { Answers: []kagenta2a.AskUserAnswer{{Answer: []string{"PostgreSQL"}}}, }) reply.TaskID, reply.ContextID = waiting.ID, waiting.ContextID - response, err := a2agrpc.NewGRPCTransportFromClient(fixture.client).SendMessage(fixture.ctx, nil, &a2atype.SendMessageRequest{Tenant: fixture.tenant, Message: reply}) + response, err := sendResume(t, fixture, reply) if err != nil { t.Fatalf("resume A2A task: %v", err) } @@ -207,6 +206,23 @@ func TestSessionAskUserSurvivesSuspension(t *testing.T) { }) } +// sendResume answers a waiting task after its Actor was suspended. Waking the +// Actor can fail before the runtime accepts the reply; the server then answers +// KAGENT_SEND_NOT_ACCEPTED, and a client retries that send as it would any other. +func sendResume(t *testing.T, fixture *interactionFixture, reply *a2atype.Message) (a2atype.SendMessageResult, error) { + t.Helper() + request, err := pbconv.ToProtoSendMessageRequest(&a2atype.SendMessageRequest{Message: reply}) + if err != nil { + t.Fatalf("build resume request: %v", err) + } + request.Tenant = fixture.tenant + response, err := sendMessageWithRetry(fixture.ctx, fixture.client, request) + if err != nil { + return nil, err + } + return pbconv.FromProtoSendMessageResponse(response) +} + func sendToolDecision(t *testing.T, fixture *interactionFixture, prompt, wantTool string, approved bool) *a2atype.Task { t.Helper() fixture.ctx = metadata.AppendToOutgoingContext(fixture.ctx, strings.ToLower(a2atype.SvcParamExtensions), kagenta2a.HITLExtensionURI) @@ -242,7 +258,7 @@ func sendToolDecision(t *testing.T, fixture *interactionFixture, prompt, wantToo }); err != nil { t.Fatalf("attach tool approval response: %v", err) } - response, err := a2agrpc.NewGRPCTransportFromClient(fixture.client).SendMessage(fixture.ctx, nil, &a2atype.SendMessageRequest{Tenant: fixture.tenant, Message: reply}) + response, err := sendResume(t, fixture, reply) if err != nil { t.Fatalf("resume A2A task after tool approval: %v", err) } diff --git a/go/core/test/e2e/session_revert_test.go b/go/core/test/e2e/session_revert_test.go index 349c011612..af1b6ecf4c 100644 --- a/go/core/test/e2e/session_revert_test.go +++ b/go/core/test/e2e/session_revert_test.go @@ -13,6 +13,7 @@ import ( "github.com/a2aproject/a2a-go/v2/a2a" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" "github.com/kagent-dev/mockllm" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" @@ -26,129 +27,184 @@ import ( "github.com/kagent-dev/kagent/go/core/internal/substrate" ) -func TestSessionRevertCrashedRuntime(t *testing.T) { - target := interactionTarget(t) - forEachHarness(t, func(t *testing.T, harness testHarness) { - cfg, err := mockllm.LoadConfigFromFile("mocks/invoke_agent.json", interactionMocks) - require.NoError(t, err) - upstream, err := url.Parse(startMockLLMConfig(t, cfg)) - require.NoError(t, err) - proxy := httputil.NewSingleHostReverseProxy(upstream) - // Only an interaction runs a Session's Actor; quiescence suspends it once - // the turn settles. The gate holds the second turn's model call open so - // the Actor is still running when its worker is deleted. - var blockNext atomic.Bool - blocked := make(chan struct{}, 1) - release := make(chan struct{}) - gate := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost && blockNext.CompareAndSwap(true, false) { - blocked <- struct{}{} - select { - case <-r.Context().Done(): - case <-release: - } - return +// crashedTurn is a Session whose worker was deleted during its second turn. The +// first turn completed and left a durable snapshot; the second is still open. +type crashedTurn struct { + fixture *interactionFixture + namespace, name string + running *ateapipb.Actor + revert *apiv1.RevertSessionRequest + getTask *a2apb.GetTaskRequest + before *a2apb.Task + stopTurn context.CancelFunc + turnDone chan error +} + +// crashDuringTurn holds the second turn's model call open, so the Actor is still +// running when its worker is deleted, and returns once the Actor is CRASHED. +func crashDuringTurn(t *testing.T, harness testHarness, target string) *crashedTurn { + t.Helper() + cfg, err := mockllm.LoadConfigFromFile("mocks/invoke_agent.json", interactionMocks) + require.NoError(t, err) + upstream, err := url.Parse(startMockLLMConfig(t, cfg)) + require.NoError(t, err) + proxy := httputil.NewSingleHostReverseProxy(upstream) + // Only an interaction runs a Session's Actor; quiescence suspends it once the + // turn settles. + var blockNext atomic.Bool + blocked := make(chan struct{}, 1) + release := make(chan struct{}) + gate := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && blockNext.CompareAndSwap(true, false) { + blocked <- struct{}{} + select { + case <-r.Context().Done(): + case <-release: } - proxy.ServeHTTP(w, r) - })) - require.NoError(t, gate.Listener.Close()) - gate.Listener, err = net.Listen("tcp", "0.0.0.0:0") - require.NoError(t, err) - gate.Start() - t.Cleanup(gate.Close) - fixture := newInteractionFixture(t, harness, target, reachableModelURL(t, gate.URL)) - defer func() { - close(release) - gate.CloseClientConnections() - }() - _, _, task := fixture.send(t, "What is 2+2?") - require.Equal(t, a2a.TaskStateCompleted, task.Status.State) - assertActorSuspended(t, fixture) - getTask := &a2apb.GetTaskRequest{Tenant: fixture.tenant, Id: string(task.ID)} - before, err := fixture.client.GetTask(fixture.ctx, getTask) - require.NoError(t, err) - session, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) - require.NoError(t, err) - namespace, name := session.GetSession().GetAgent().GetRef().GetNamespace(), substrate.ActorName(fixture.sessionID) - blockNext.Store(true) - _, turn := newMessageRequest(t, "What is 2+2?") - turn.Tenant, turn.Message.ContextId = fixture.tenant, fixture.sessionID - turnCtx, stopTurn := context.WithCancel(fixture.ctx) - defer stopTurn() - sendDone := make(chan error, 1) - go func() { _, err := sendMessageWithRetry(turnCtx, fixture.client, turn); sendDone <- err }() - select { - case <-blocked: - case <-time.After(time.Minute): - t.Fatal("second turn never reached the blocking model") + return } - running, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) + proxy.ServeHTTP(w, r) + })) + require.NoError(t, gate.Listener.Close()) + gate.Listener, err = net.Listen("tcp", "0.0.0.0:0") + require.NoError(t, err) + gate.Start() + t.Cleanup(gate.Close) + fixture := newInteractionFixture(t, harness, target, reachableModelURL(t, gate.URL)) + t.Cleanup(func() { + close(release) + gate.CloseClientConnections() + }) + _, _, task := fixture.send(t, "What is 2+2?") + require.Equal(t, a2a.TaskStateCompleted, task.Status.State) + assertActorSuspended(t, fixture) + crash := &crashedTurn{fixture: fixture, getTask: &a2apb.GetTaskRequest{Tenant: fixture.tenant, Id: string(task.ID)}, + revert: &apiv1.RevertSessionRequest{SessionId: fixture.sessionID}, turnDone: make(chan error, 1)} + crash.before, err = fixture.client.GetTask(fixture.ctx, crash.getTask) + require.NoError(t, err) + session, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + crash.namespace, crash.name = session.GetSession().GetAgent().GetRef().GetNamespace(), substrate.ActorName(fixture.sessionID) + blockNext.Store(true) + _, turn := newMessageRequest(t, "What is 2+2?") + turn.Tenant, turn.Message.ContextId = fixture.tenant, fixture.sessionID + var turnCtx context.Context + turnCtx, crash.stopTurn = context.WithCancel(fixture.ctx) + t.Cleanup(crash.stopTurn) + go func() { _, err := sendMessageWithRetry(turnCtx, fixture.client, turn); crash.turnDone <- err }() + select { + case <-blocked: + case <-time.After(time.Minute): + t.Fatal("second turn never reached the blocking model") + } + crash.running, err = findSubstrateActor(fixture.ctx, fixture.system, crash.namespace, crash.name) + require.NoError(t, err) + require.NotNil(t, crash.running) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_RUNNING, crash.running.GetStatus().GetState()) + require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(crash.running))) + healthy, err := fixture.sessions.RevertSession(fixture.ctx, crash.revert) + require.NoError(t, err, "a healthy runtime has nothing to recover") + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) + assignment := crash.running.GetStatus().GetWorkerAssignment() + require.Equal(t, crash.namespace, assignment.GetWorkerNamespace()) + require.NotEmpty(t, assignment.GetWorkerPod()) + for token := ""; ; { + page, err := fixture.system.ListSubstrateActors(fixture.ctx, &apiv1.ListSubstrateActorsRequest{ + Page: &apiv1.PageRequest{Limit: 100, PageToken: token}, + }) require.NoError(t, err) - require.NotNil(t, running) - require.Equal(t, ateapipb.ActorState_ACTOR_STATE_RUNNING, running.GetStatus().GetState()) - require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running))) - request := &apiv1.RevertSessionRequest{SessionId: fixture.sessionID} - healthy, err := fixture.sessions.RevertSession(fixture.ctx, request) - require.NoError(t, err, "a healthy runtime has nothing to recover") - require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) - assignment := running.GetStatus().GetWorkerAssignment() - require.Equal(t, namespace, assignment.GetWorkerNamespace()) - require.NotEmpty(t, assignment.GetWorkerPod()) - for token := ""; ; { - page, err := fixture.system.ListSubstrateActors(fixture.ctx, &apiv1.ListSubstrateActorsRequest{ - Page: &apiv1.PageRequest{Limit: 100, PageToken: token}, - }) - require.NoError(t, err) - require.Empty(t, page.GetAteApiError()) - for _, actor := range page.GetActors() { - other := actor.GetStatus().GetWorkerAssignment() - if actor.GetMetadata().GetUid() != running.GetMetadata().GetUid() { - require.False(t, other.GetWorkerNamespace() == assignment.GetWorkerNamespace() && other.GetWorkerPod() == assignment.GetWorkerPod(), - "refuse to delete a worker shared with another Actor") - } - } - token = page.GetPage().GetNextPageToken() - if token == "" { - break + require.Empty(t, page.GetAteApiError()) + for _, actor := range page.GetActors() { + other := actor.GetStatus().GetWorkerAssignment() + if actor.GetMetadata().GetUid() != crash.running.GetMetadata().GetUid() { + require.False(t, other.GetWorkerNamespace() == assignment.GetWorkerNamespace() && other.GetWorkerPod() == assignment.GetWorkerPod(), + "refuse to delete a worker shared with another Actor") } } - kube := interactionKubeClient(t) - pod := &corev1.Pod{} - require.NoError(t, kube.Get(fixture.ctx, ctrlclient.ObjectKey{ - Namespace: assignment.GetWorkerNamespace(), Name: assignment.GetWorkerPod(), - }, pod)) - require.NotEmpty(t, pod.UID) - require.NoError(t, kube.Delete(fixture.ctx, pod, &ctrlclient.DeleteOptions{ - GracePeriodSeconds: new(int64(0)), Preconditions: &metav1.Preconditions{UID: &pod.UID}, - })) - err = wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { - actor, err := findSubstrateActor(ctx, fixture.system, namespace, name) - return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED, err - }) - require.NoError(t, err) - reverted, err := fixture.sessions.RevertSession(fixture.ctx, request) + token = page.GetPage().GetNextPageToken() + if token == "" { + break + } + } + kube := interactionKubeClient(t) + pod := &corev1.Pod{} + require.NoError(t, kube.Get(fixture.ctx, ctrlclient.ObjectKey{ + Namespace: assignment.GetWorkerNamespace(), Name: assignment.GetWorkerPod(), + }, pod)) + require.NotEmpty(t, pod.UID) + require.NoError(t, kube.Delete(fixture.ctx, pod, &ctrlclient.DeleteOptions{ + GracePeriodSeconds: new(int64(0)), Preconditions: &metav1.Preconditions{UID: &pod.UID}, + })) + err = wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + actor, err := findSubstrateActor(ctx, fixture.system, crash.namespace, crash.name) + return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED, err + }) + require.NoError(t, err) + return crash +} + +// assertRecovered checks the Actor was restored in place, the completed turn kept +// its history, and the interrupted caller stopped. +func (c *crashedTurn) assertRecovered(t *testing.T) *ateapipb.Actor { + t.Helper() + recovered, err := findSubstrateActor(c.fixture.ctx, c.fixture.system, c.namespace, c.name) + require.NoError(t, err) + require.Equal(t, c.running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid(), "recovery restores the recorded Actor in place") + retained, err := c.fixture.client.GetTask(c.fixture.ctx, c.getTask) + require.NoError(t, err) + require.True(t, proto.Equal(c.before, retained), "recovery preserves the completed task and its history") + c.stopTurn() + select { + case <-c.turnDone: + case <-time.After(10 * time.Second): + t.Fatal("interrupted caller did not stop") + } + return recovered +} + +// A worker is lost during a turn and the client calls RevertSession, as chat's +// Recover does. The crashed Actor returns to its last durable snapshot. +func TestSessionRevertCrashedRuntime(t *testing.T) { + target := interactionTarget(t) + forEachHarness(t, func(t *testing.T, harness testHarness) { + crash := crashDuringTurn(t, harness, target) + reverted, err := crash.fixture.sessions.RevertSession(crash.fixture.ctx, crash.revert) require.NoError(t, err) require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, reverted.Session.State) // Retry as a client whose success response was lost. - retried, err := fixture.sessions.RevertSession(fixture.ctx, request) + retried, err := crash.fixture.sessions.RevertSession(crash.fixture.ctx, crash.revert) require.NoError(t, err) require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, retried.Session.State) - recovered, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) - require.NoError(t, err) - require.Equal(t, running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid()) + recovered := crash.assertRecovered(t) require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, recovered.GetStatus().GetState()) - require.Equal(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running)), substrate.SnapshotURI(substrate.LatestActorSnapshot(recovered))) - retained, err := fixture.client.GetTask(fixture.ctx, getTask) - require.NoError(t, err) - require.True(t, proto.Equal(before, retained), "recovery preserves the completed task and its history") - stopTurn() - select { - case <-sendDone: - case <-time.After(10 * time.Second): - t.Fatal("interrupted caller did not stop") - } - _, _, next := fixture.send(t, "What is 2+2?") + require.Equal(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(crash.running)), substrate.SnapshotURI(substrate.LatestActorSnapshot(recovered))) + _, _, next := crash.fixture.send(t, "What is 2+2?") require.Equal(t, a2a.TaskStateCompleted, next.Status.State) require.Contains(t, taskText(next), "4") }) } + +// A worker is lost during a turn. The next send recovers the crashed Actor to its +// last durable snapshot and is answered by the restored runtime; the client calls +// nothing but SendMessage. Requires sessions.automaticRecovery. +func TestSessionSendRecoversCrashedRuntime(t *testing.T) { + target := interactionTarget(t) + forEachHarness(t, func(t *testing.T, harness testHarness) { + crash := crashDuringTurn(t, harness, target) + _, next := newMessageRequest(t, "What is 2+2?") + next.Tenant, next.Message.ContextId = crash.fixture.tenant, crash.fixture.sessionID + answered, err := sendMessageWithRetry(crash.fixture.ctx, crash.fixture.client, next) + require.NoError(t, err, "the send recovers the crashed runtime and reaches the restored one") + answer, err := pbconv.FromProtoSendMessageResponse(answered) + require.NoError(t, err) + answerTask, ok := answer.(*a2a.Task) + require.True(t, ok, "response = %T", answer) + require.Equal(t, a2a.TaskStateCompleted, answerTask.Status.State) + require.Contains(t, taskText(answerTask), "4") + crash.assertRecovered(t) + // RevertSession afterwards finds a healthy runtime. + healthy, err := crash.fixture.sessions.RevertSession(crash.fixture.ctx, crash.revert) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) + }) +} diff --git a/helm/kagent/tests/controller-configmap_test.yaml b/helm/kagent/tests/controller-configmap_test.yaml index 15a6c077a6..a2a84e4ba9 100644 --- a/helm/kagent/tests/controller-configmap_test.yaml +++ b/helm/kagent/tests/controller-configmap_test.yaml @@ -91,6 +91,20 @@ tests: path: data["controller.yaml"] pattern: 'apiVersion:|kind:|KAGENT_|OTEL_' + - it: omits default-off feature flags so older releases can read the config + template: controller-configmap.yaml + asserts: + - notMatchRegex: + path: data["controller.yaml"] + pattern: automaticRecovery + - it: renders automatic recovery when enabled + template: controller-configmap.yaml + set: + config.sessions.automaticRecovery: true + asserts: + - matchRegex: + path: data["controller.yaml"] + pattern: 'automaticRecovery: true' - it: preserves explicit false and zero duration template: controller-configmap.yaml set: diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 22528fbaa7..cbc80fb1a0 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -188,6 +188,10 @@ config: # Recovery scans for missed wake-ups and unbound registration cleanup. recoveryInterval: 60s sessions: + # Set automaticRecovery: true to recover crashed runtimes when a client sends + # another message. It is off by default and not rendered, so the default config + # stays readable by releases that predate the setting. + # automaticRecovery: false idleTtl: 604800s shareMaxTtl: 0s sandboxes: diff --git a/proto/kagent/api/v1/sessions.proto b/proto/kagent/api/v1/sessions.proto index 93651ce5c0..f766049dbe 100644 --- a/proto/kagent/api/v1/sessions.proto +++ b/proto/kagent/api/v1/sessions.proto @@ -42,6 +42,9 @@ message Session { string a2a_authority = 6; RuntimeState state = 7; RuntimeOperation operation = 8; + // Output only. With state FAILED, reason RuntimeCrashed and operation REVERT + // mean a crash recovery is pending; reason RuntimeLost and operation NONE mean + // the recorded runtime is gone and the Session can only be deleted. Failure failure = 9; google.protobuf.Timestamp created_at = 10; google.protobuf.Timestamp updated_at = 11; diff --git a/proto/kagent/config/v1/config.proto b/proto/kagent/config/v1/config.proto index fd97827559..ac529c064f 100644 --- a/proto/kagent/config/v1/config.proto +++ b/proto/kagent/config/v1/config.proto @@ -231,6 +231,8 @@ message SubstrateConfig { } message SessionsConfig { + // Recover crashed runtimes during sends. Disabled by default; explicit revert remains available. + bool automatic_recovery = 5; // Zero disables idle expiration. google.protobuf.Duration idle_ttl = 1 [(buf.validate.field).duration = { gte: {} diff --git a/ui/src/api/ApiError.ts b/ui/src/api/ApiError.ts index 558030b2c1..d71d864299 100644 --- a/ui/src/api/ApiError.ts +++ b/ui/src/api/ApiError.ts @@ -84,6 +84,32 @@ export function isNotFound(error: unknown): boolean { return isApiError(error) && error.status === 404; } +/** + * `ApiError.reason` of a send the gateway kept refusing as not accepted until chat's + * retry window ran out. The runtime may have crashed. + */ +export const SEND_NOT_ACCEPTED = "KAGENT_SEND_NOT_ACCEPTED"; + +/** + * `ApiError.reason` of a send refused because an earlier turn is still open. A + * runtime that crashed mid-turn leaves its turn open, so unless the controller + * recovers on send, this refusal is the only sign of the crash. + */ +export const TURN_ACTIVE = "KAGENT_TURN_ACTIVE"; + +/** Whether `error` is a send the gateway never accepted, after every retry. */ +export function isSendNotAccepted(error: unknown): boolean { + return isApiError(error) && error.reason === SEND_NOT_ACCEPTED; +} + +/** + * Whether chat should offer a revert after `error`. `RevertSession` returns a + * healthy runtime unchanged, so a refusal that only might mean a crash is enough. + */ +export function mayNeedRecovery(error: unknown): boolean { + return isApiError(error) && (error.reason === SEND_NOT_ACCEPTED || error.reason === TURN_ACTIVE); +} + /** * The status each gRPC code stands for, from the Connect protocol's own table. * diff --git a/ui/src/api/chat/a2aGrpcChatClient.test.ts b/ui/src/api/chat/a2aGrpcChatClient.test.ts index 278ca76673..3f8ecf3c10 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.test.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.test.ts @@ -23,8 +23,10 @@ import { TaskState, type SendMessageRequest, } from "@/generated/a2a_pb"; +import { ErrorInfoSchema } from "@/generated/google/rpc/error_details_pb"; import { setApiTransport } from "../transport"; import { A2AGrpcChatClient } from "./a2aGrpcChatClient"; +import { isSendNotAccepted, mayNeedRecovery, SEND_NOT_ACCEPTED, TURN_ACTIVE } from "../ApiError"; import { LinkFile } from "./attachments"; import { A2A_METADATA } from "./a2aMetadata"; import type { ChatEvent, ChatMessage } from "./types"; @@ -2101,3 +2103,167 @@ describe("A2AGrpcChatClient.send, when the stream breaks", () => { expect(subscribed).toEqual([]); }); }); + +/* + * A send that reaches a crashed runtime reverts it on the server and is refused with + * proof that the message was not accepted. The client sends the identical request + * again, which is safe because a repeated message ID cannot be accepted twice. + */ +describe("A2AGrpcChatClient.send, when the gateway did not accept the message", () => { + /** The error a2a-go writes for kagent's ErrSendNotAccepted. */ + const notAccepted = (metadata: Record = { reason: "KAGENT_SEND_NOT_ACCEPTED", retryAfterMs: "100" }, code = Code.FailedPrecondition) => + new ConnectError("input was not accepted; retry after the session becomes available", code, undefined, [ + { desc: ErrorInfoSchema, value: { reason: "UNSUPPORTED_OPERATION", domain: "a2a-protocol.org", metadata } }, + ]); + + const completed = statusFrame({ state: TaskState.COMPLETED }); + + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + async function run(attempt: (n: number) => AsyncGenerator, signal?: AbortSignal) { + const requests: SendMessageRequest[] = []; + serve(({ service }) => { + service(A2AService, { + sendStreamingMessage: async function* (request) { + requests.push(request); + for await (const frame of attempt(requests.length)) yield frame as never; + }, + }); + }); + const events: ChatEvent[] = []; + let error: unknown; + let done = false; + const turn = (async () => { + try { + for await (const event of new A2AGrpcChatClient().send({ + conversation: CONVERSATION, + text: "why is checkout crashlooping?", + messageId: "msg-local-1", + signal, + })) { + events.push(event); + } + } catch (cause) { + error = cause; + } finally { + done = true; + } + })(); + for (let tick = 0; !done && tick < 600; tick += 1) { + await vi.advanceTimersByTimeAsync(100); + } + expect(done).toBe(true); + await turn; + return { events, error, requests }; + } + + it("re-sends the identical request until the gateway accepts it", async () => { + const { events, error, requests } = await run(async function* (n) { + if (n < 3) throw notAccepted(); + yield completed; + }); + expect(error).toBeUndefined(); + expect(requests).toHaveLength(3); + expect(requests.map((request) => request.message?.messageId)).toEqual(["msg-local-1", "msg-local-1", "msg-local-1"]); + expect(requests[2]).toEqual(requests[0]); + expect(events.at(-1)).toMatchObject({ type: "status", state: "completed" }); + }); + + it("waits as long as the gateway asks before re-sending", async () => { + const sentAt: number[] = []; + const { error } = await run(async function* (n) { + sentAt.push(Date.now()); + if (n === 1) throw notAccepted({ reason: "KAGENT_SEND_NOT_ACCEPTED", retryAfterMs: "1500" }); + yield completed; + }); + expect(error).toBeUndefined(); + expect(sentAt[1] - sentAt[0]).toBeGreaterThanOrEqual(1_500); + }); + + it("never re-sends once a frame has arrived", async () => { + const { error, requests } = await run(async function* () { + yield statusFrame({}); + throw notAccepted(); + }); + expect(error).toMatchObject({ code: "FailedPrecondition" }); + expect(mayNeedRecovery(error)).toBe(false); + expect(requests).toHaveLength(1); + }); + + it.each([ + ["an earlier turn that is still open", notAccepted({ reason: TURN_ACTIVE }), TURN_ACTIVE], + ["another metadata reason", notAccepted({ reason: "SOMETHING_ELSE" }), undefined], + ["no metadata, as for a runtime that cannot be recovered", notAccepted({}), undefined], + ["another gRPC code", notAccepted(undefined, Code.Unavailable), undefined], + ["no ErrorInfo", new ConnectError("Session runtime cannot be recovered; start another conversation", Code.FailedPrecondition), undefined], + ])("does not re-send for %s", async (_, refusal, reason) => { + const { error, requests } = await run(async function* () { + yield* []; + throw refusal; + }); + expect(error).toBeDefined(); + expect(isSendNotAccepted(error)).toBe(false); + expect((error as { reason?: string }).reason).toBe(reason); + expect(mayNeedRecovery(error)).toBe(reason !== undefined); + expect(requests).toHaveLength(1); + }); + + it("keeps the gateway's words for a busy Session", async () => { + const { error } = await run(async function* () { + yield* []; + throw new ConnectError("resource conflict", Code.FailedPrecondition, undefined, [ + { desc: ErrorInfoSchema, value: { reason: "UNSUPPORTED_OPERATION", domain: "a2a-protocol.org", metadata: { reason: TURN_ACTIVE } } }, + ]); + }); + expect(error).toMatchObject({ message: "resource conflict", code: "FailedPrecondition", reason: TURN_ACTIVE }); + }); + + it("doubles the wait between re-sends, up to two seconds", async () => { + const sentAt: number[] = []; + await run(async function* () { + sentAt.push(Date.now()); + yield* []; + throw notAccepted(); + }); + const waits = sentAt.slice(1).map((at, i) => at - sentAt[i]); + expect(waits.slice(0, 6)).toEqual([100, 200, 400, 800, 1_600, 2_000]); + expect(Math.max(...waits)).toBe(2_000); + }); + + it("surfaces the refusal as not accepted once the retry window is spent", async () => { + const { error, requests } = await run(async function* () { + yield* []; + throw notAccepted(); + }); + expect(error).toMatchObject({ code: "FailedPrecondition", reason: SEND_NOT_ACCEPTED }); + expect(isSendNotAccepted(error)).toBe(true); + // 100 ms doubling to 2 s within 30 s: 19 sends, rather than one every 100 ms. + expect(requests).toHaveLength(19); + }); + + it("keeps re-sending while the runtime recovers, and says so when the window is spent", async () => { + const sentAt: number[] = []; + const { error } = await run(async function* () { + sentAt.push(Date.now()); + yield* []; + throw notAccepted({ reason: "KAGENT_SEND_NOT_ACCEPTED", state: "RECOVERING", retryAfterMs: "1000" }); + }); + expect(error).toMatchObject({ reason: SEND_NOT_ACCEPTED, message: "The agent's runtime is still recovering from a crash." }); + expect(mayNeedRecovery(error)).toBe(true); + // The gateway's longer delay is the floor, so a recovery costs a few sends. + expect(sentAt[1] - sentAt[0]).toBeGreaterThanOrEqual(1_000); + }); + + it("stops re-sending when the reader presses stop", async () => { + const stop = new AbortController(); + // The reader presses stop while the client waits to send a third time. + const { error, requests } = await run(async function* (n) { + if (n === 2) setTimeout(() => stop.abort(), 50); + yield* []; + throw notAccepted({ reason: "KAGENT_SEND_NOT_ACCEPTED", retryAfterMs: "1000" }); + }, stop.signal); + expect(error).toMatchObject({ name: "AbortError" }); + expect(requests).toHaveLength(2); + }); +}); diff --git a/ui/src/api/chat/a2aGrpcChatClient.ts b/ui/src/api/chat/a2aGrpcChatClient.ts index 94b7657940..28ca992b0a 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.ts @@ -59,7 +59,8 @@ import { type Task as A2ATask, type TaskStatus, } from "@/generated/a2a_pb"; -import { ApiError, fromConnectError, rethrowIfAborted } from "../ApiError"; +import { ErrorInfoSchema } from "@/generated/google/rpc/error_details_pb"; +import { ApiError, fromConnectError, rethrowIfAborted, SEND_NOT_ACCEPTED, TURN_ACTIVE } from "../ApiError"; import { A2A_METADATA, metadataString } from "./a2aMetadata"; import { HITL_EXTENSION_HEADER, @@ -115,6 +116,60 @@ const REATTACH_HEALTHY_MS = 10_000; /** The wait before the first re-attach after a failure, doubled for each one after. */ const REATTACH_BACKOFF_MS = 500; +/** + * How long a send keeps re-sending a message the gateway proved it did not accept. + * The window bounds the retries, not how long an accepted turn may run. + */ +const UNACCEPTED_SEND_WINDOW_MS = 30_000; + +/** The wait before re-sending when the gateway does not name one. */ +const UNACCEPTED_SEND_RETRY_MS = 100; + +/** + * The longest wait between re-sends. Each wait doubles from the gateway's, so a + * runtime that stays down costs a handful of sends per window rather than hundreds. + */ +const UNACCEPTED_SEND_MAX_RETRY_MS = 2_000; + +/** `ErrorInfo.domain` of every error the A2A server writes (`a2a.ProtocolDomain`). */ +const A2A_PROTOCOL_DOMAIN = "a2a-protocol.org"; + +/** + * A2A's UNSUPPORTED_OPERATION refusal (gRPC FAILED_PRECONDITION), or undefined. + * The gateway answers an input it cannot take this way, before running anything. + */ +function a2aRefusal(error: unknown) { + if (!(error instanceof ConnectError) || error.code !== Code.FailedPrecondition) return undefined; + return error.findDetails(ErrorInfoSchema).find((detail) => + detail.domain === A2A_PROTOCOL_DOMAIN && detail.reason === "UNSUPPORTED_OPERATION"); +} + +/** + * The wait before re-sending, when `error` is the gateway's proof that it did not + * accept the message; otherwise undefined. + * + * The gateway answers a send that reached a crashed runtime, or a recovery still in + * progress, with A2A's UNSUPPORTED_OPERATION (gRPC FAILED_PRECONDITION) and a + * `google.rpc.ErrorInfo` whose metadata says `reason: KAGENT_SEND_NOT_ACCEPTED`. + * Only that combination is proof: any other refusal, or a broken connection, may + * have accepted the input, and sending it again could start a second turn. + */ +function unacceptedSendRetryMs(error: unknown): number | undefined { + const info = a2aRefusal(error); + if (info?.metadata.reason !== SEND_NOT_ACCEPTED) return undefined; + const wait = Number.parseInt(info.metadata.retryAfterMs ?? "", 10); + return Number.isFinite(wait) && wait >= 0 ? wait : UNACCEPTED_SEND_RETRY_MS; +} + +/** A refused send, tagged with `reason` for the page to offer a revert. */ +function refusedSendError(error: unknown, reason: string, message?: string): ApiError { + const failure = fromConnectError(error, "A2AService/SendStreamingMessage"); + return new ApiError(message ?? failure.message, { + kind: failure.kind, status: failure.status, code: failure.code, + reason, url: failure.url, cause: error, + }); +} + /** Resolves after `ms`, or rejects with the abort as soon as `signal` fires. */ function pause(ms: number, signal: AbortSignal | undefined): Promise { return new Promise((resolve, reject) => { @@ -600,16 +655,26 @@ export class A2AGrpcChatClient implements ChatClient { * not taken as a sign the path works: a path that breaks right after it would * otherwise never run out of attempts. What resets the count is a frame after * it, or a stream that stayed open long enough to have been cut for being quiet. + * + * Before any frame, a send the gateway proves it did not accept is sent again, + * identically, after a wait that doubles up to `UNACCEPTED_SEND_MAX_RETRY_MS`. + * With automatic recovery enabled, that is how a crashed runtime recovers: the + * send that finds it reverts it, and the same message ID can never be accepted + * twice. Once a frame has arrived, the message was accepted and is never re-sent. */ let taskId = ""; let settled = false; const follow = async function* (): AsyncGenerator { + const resendUntil = Date.now() + UNACCEPTED_SEND_WINDOW_MS; + let resendWait = 0; + let received = false; let source = stream; for (let failures = 0; ; ) { const opened = Date.now(); let progressed = false; try { for await (const frame of source) { + received = true; if (frame.payload.case !== "task") progressed = true; const event = frame.payload.value as | { id?: string; taskId?: string; status?: TaskStatus } @@ -623,6 +688,22 @@ export class A2AGrpcChatClient implements ChatClient { } return; } catch (error) { + const resendAfter = received ? undefined : unacceptedSendRetryMs(error); + if (resendAfter !== undefined && !signal?.aborted) { + resendWait = Math.max(resendAfter, Math.min(resendWait * 2, UNACCEPTED_SEND_MAX_RETRY_MS)); + if (Date.now() + resendWait > resendUntil) { + const recovering = a2aRefusal(error)?.metadata.state === "RECOVERING"; + throw refusedSendError(error, SEND_NOT_ACCEPTED, recovering + ? "The agent's runtime is still recovering from a crash." + : "The agent's runtime did not accept the message. It may have crashed."); + } + await pause(resendWait, signal); + source = client.sendStreamingMessage(request, options); + continue; + } + if (!received && !signal?.aborted && a2aRefusal(error)?.metadata.reason === TURN_ACTIVE) { + throw refusedSendError(error, TURN_ACTIVE); + } if (progressed || Date.now() - opened >= REATTACH_HEALTHY_MS) failures = 0; failures += 1; const reattach = !signal?.aborted && taskId !== "" && !settled diff --git a/ui/src/api/client.ts b/ui/src/api/client.ts index b30625483e..46d0ebaf09 100644 --- a/ui/src/api/client.ts +++ b/ui/src/api/client.ts @@ -211,6 +211,9 @@ export interface AgentInstancesApi { ): Promise; get(id: string, options?: ReadOptions): Promise; + /** Restores a crashed runtime to its last retained snapshot. */ + revert(id: string): Promise; + create(input: { agent: ResourceRefInput; /** @@ -390,6 +393,7 @@ export function createApiClient(): KagentApiClient { invoke("agentInstances.rename", { id, name }), get: (id, options) => invoke("agentInstances.get", { id }, options), + revert: (id) => invoke("agentInstances.revert", { id }), create: (input) => invoke("agentInstances.create", input), remove: (id, options) => invoke("agentInstances.delete", { id, ...options }), fork: (id, name) => diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index f7e7916b05..731c59beee 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -677,6 +677,7 @@ const agentInstances: Pick< | "agentInstances.rename" | "agentInstances.fork" | "agentInstances.delete" + | "agentInstances.revert" > = { "agentInstances.list": async (input, options) => { const name = "SessionService/ListSessions"; @@ -952,6 +953,16 @@ const agentInstances: Pick< return toAgentInstance(required(response.session, name, "agent instance")); }, + "agentInstances.revert": async (input, options) => { + const name = "SessionService/RevertSession"; + const response = await rpc(name, options.signal, () => + serviceClient(SessionService).revertSession( + { sessionId: input.id }, + call("agentInstances.revert", options), + ), + ); + return toAgentInstance(required(response.session, name, "agent instance")); + }, }; // endregion diff --git a/ui/src/api/index.ts b/ui/src/api/index.ts index 7e909617e1..eb90f691b4 100644 --- a/ui/src/api/index.ts +++ b/ui/src/api/index.ts @@ -7,7 +7,15 @@ * is still held over HTTP. */ -export { ApiError, isApiError, isNotFound } from "./ApiError"; +export { + ApiError, + isApiError, + isNotFound, + isSendNotAccepted, + mayNeedRecovery, + SEND_NOT_ACCEPTED, + TURN_ACTIVE, +} from "./ApiError"; export type { ApiErrorKind } from "./ApiError"; export { apiBaseUrl, apiMode, isMockMode, MOCK_API_BASE_URL } from "./config"; diff --git a/ui/src/api/operations.test.ts b/ui/src/api/operations.test.ts index b862f8b9c7..6f7dc0703d 100644 --- a/ui/src/api/operations.test.ts +++ b/ui/src/api/operations.test.ts @@ -994,6 +994,21 @@ describe("agent instances", () => { expect(isNotFound(missing)).toBe(true); }); + it("reverts through the existing Session RPC", async () => { + let received: string | undefined; + serve(({ service }) => { + service(SessionService, { + revertSession: (request) => { + received = request.sessionId; + return { session: instanceMessage({ state: RuntimeState.SUSPENDED }) }; + }, + }); + }); + const reverted = await apiClient.agentInstances.revert(INSTANCE_ID); + expect(received).toBe(INSTANCE_ID); + expect(reverted.state).toBe("suspended"); + }); + }); describe("session shares", () => { diff --git a/ui/src/api/operations.ts b/ui/src/api/operations.ts index ee98421d74..f35f66a90e 100644 --- a/ui/src/api/operations.ts +++ b/ui/src/api/operations.ts @@ -291,6 +291,7 @@ export interface OperationMap { output: void; }; + "agentInstances.revert": { input: AgentInstanceRef; output: AgentInstance }; "harnesses.list": { input: { namespace?: string }; output: Harness[] }; /** Creates a harness from a whole custom resource. */ diff --git a/ui/src/components/chat/ChatTranscript.tsx b/ui/src/components/chat/ChatTranscript.tsx index 0ac76403b5..8f2ac2e4e7 100644 --- a/ui/src/components/chat/ChatTranscript.tsx +++ b/ui/src/components/chat/ChatTranscript.tsx @@ -3,7 +3,7 @@ import type { ReactNode } from "react"; import { Alert, Button, Empty, Skeleton, Tag, Tooltip } from "antd"; import { ChevronDown } from "lucide-react"; import { useTheme } from "@emotion/react"; -import type { ChatController, ChatTurnPhase, Checkpoint } from "@/api"; +import { isSendNotAccepted, mayNeedRecovery, type ChatController, type ChatTurnPhase, type Checkpoint } from "@/api"; import { AskUserPrompt } from "./AskUserPrompt"; import { ChatMessageItem } from "./ChatMessageItem"; import { CheckpointDivider } from "./CheckpointDivider"; @@ -62,6 +62,8 @@ export function ChatTranscript({ onFork, onDeleteCheckpoint, onManageTaskPush, + onRecover, + isRecovering = false, checkpointsById, checkpointByMessage, }: { @@ -83,6 +85,14 @@ export function ChatTranscript({ onDeleteCheckpoint?: (checkpointId: string) => void; /** Opens active callbacks for one A2A task. */ onManageTaskPush?: (taskId: string) => void; + /** + * Reverts a runtime that may have crashed, offered when a send is refused as not + * accepted after every retry, or because an open turn holds the Session. Absent when + * read-only. + */ + onRecover?: () => void; + /** A revert from `onRecover` is in flight. */ + isRecovering?: boolean; /** The boundaries the controller has described, for the line to name itself by. */ checkpointsById?: ReadonlyMap; /** Which boundary each message sits inside, for the messages that sit inside one. */ @@ -447,18 +457,40 @@ export function ChatTranscript({ ) : null} {chat.turnError && !chat.pendingQuestion ? ( - void chat.retry()}> - Retry - - )} - /> + mayNeedRecovery(chat.turnError) && onRecover && !readOnly ? ( + + + + + } + /> + ) : ( + void chat.retry()}> + Retry + + )} + /> + ) ) : null} {/* diff --git a/ui/src/components/chat/lifecycleReading.test.ts b/ui/src/components/chat/lifecycleReading.test.ts index 34bcb12996..303a4025f8 100644 --- a/ui/src/components/chat/lifecycleReading.test.ts +++ b/ui/src/components/chat/lifecycleReading.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; import type { AgentInstance } from "@/api"; -import { isLifecycleBusy, lifecycleReading } from "./lifecycleReading"; +import { isLifecycleBusy, isRuntimeLost, lifecycleReading } from "./lifecycleReading"; /** * What the chat page is allowed to say the agent is doing. @@ -42,6 +42,17 @@ describe("lifecycleReading", () => { expect(reading.source).toMatch(/claimed a suspend operation/); }); + it("reads a pending crash recovery from the stored revert operation", () => { + // The Session is FAILED while the recovery is pending; the operation is what says + // the controller is acting on it, so this is a stage rather than a verdict. + const reading = lifecycleReading( + instance({ state: "failed", operation: "revert", failure: { reason: "RuntimeCrashed" } }), + "idle", + ); + expect(reading.active).toBe("resuming"); + expect(reading.summary).toBe("Recovering this conversation after a crash"); + }); + it("lets a claimed operation outrank the turn, because the controller is acting on the instance right now", () => { // Both are true at once — a turn in flight and a suspend being carried out — // and the instance's own lifecycle is the more specific of the two. @@ -144,6 +155,9 @@ describe("isLifecycleBusy", () => { // Two decisions taken separately would drift: a page polling on one rule while // the indicator changes on another either misses a stage or reads for nothing. expect(isLifecycleBusy(instance({ operation: "resume" }), "idle")).toBe(true); + expect(isLifecycleBusy(instance({ state: "failed", operation: "revert" }), "idle")).toBe(true); + // A lost runtime never changes on its own, so it is not polled. + expect(isLifecycleBusy(instance({ state: "failed", failure: { reason: "RuntimeLost" } }), "idle")).toBe(false); expect(isLifecycleBusy(instance(), "streaming")).toBe(true); expect(isLifecycleBusy(instance(), "idle")).toBe(false); expect(isLifecycleBusy(instance({ state: "suspended" }), "idle")).toBe(false); @@ -153,3 +167,14 @@ describe("isLifecycleBusy", () => { expect(isLifecycleBusy(instance(), "idle", true)).toBe(false); }); }); + +describe("isRuntimeLost", () => { + it("is true only for a failed Session with no operation and reason RuntimeLost", () => { + const lost = { state: "failed", operation: "unspecified", failure: { reason: "RuntimeLost" } } as const; + expect(isRuntimeLost(instance(lost))).toBe(true); + expect(isRuntimeLost(instance({ ...lost, operation: "revert" }))).toBe(false); + expect(isRuntimeLost(instance({ ...lost, failure: { reason: "RuntimeCrashed" } }))).toBe(false); + expect(isRuntimeLost(instance({ ...lost, state: "ready" }))).toBe(false); + expect(isRuntimeLost(undefined)).toBe(false); + }); +}); diff --git a/ui/src/components/chat/lifecycleReading.ts b/ui/src/components/chat/lifecycleReading.ts index 35f835f8ae..47f45fc5fe 100644 --- a/ui/src/components/chat/lifecycleReading.ts +++ b/ui/src/components/chat/lifecycleReading.ts @@ -60,6 +60,15 @@ export function lifecycleReading( }; } + if (instance.operation === "revert") { + return { + active: "resuming", + summary: "Recovering this conversation after a crash", + source: + "The agent's runtime crashed and the controller is restoring it from its last snapshot.", + }; + } + if (instance.operation === "suspend") { return { active: "suspending", @@ -167,3 +176,13 @@ export function isLifecycleBusy( ): boolean { return lifecycleReading(instance, turnPhase, isAwaitingReply).active !== undefined; } + +/** + * Whether the conversation's recorded runtime is gone for good, so it can no longer + * answer. It can still be read, forked from a saved snapshot, or deleted. + */ +export function isRuntimeLost(instance: AgentInstance | undefined): boolean { + return instance?.state === "failed" + && instance.operation === "unspecified" + && instance.failure?.reason === "RuntimeLost"; +} diff --git a/ui/src/generated/kagent/api/v1/sessions_pb.ts b/ui/src/generated/kagent/api/v1/sessions_pb.ts index fdc14c14d4..9df755e0fb 100644 --- a/ui/src/generated/kagent/api/v1/sessions_pb.ts +++ b/ui/src/generated/kagent/api/v1/sessions_pb.ts @@ -83,6 +83,10 @@ export type Session = Message<"kagent.api.v1.Session"> & { operation: RuntimeOperation; /** + * Output only. With state FAILED, reason RuntimeCrashed and operation REVERT + * mean a crash recovery is pending; reason RuntimeLost and operation NONE mean + * the recorded runtime is gone and the Session can only be deleted. + * * @generated from field: kagent.api.v1.Failure failure = 9; */ failure?: Failure | undefined; diff --git a/ui/src/mocks/mockBackend.test.ts b/ui/src/mocks/mockBackend.test.ts index e6807cf786..56efadcaf5 100644 --- a/ui/src/mocks/mockBackend.test.ts +++ b/ui/src/mocks/mockBackend.test.ts @@ -108,6 +108,7 @@ const INPUTS = { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44", }, + "agentInstances.revert": { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44" }, "agentInstances.shares.list": { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44", diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index 14a1cc3c91..bfc06bd922 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -838,6 +838,11 @@ on(SessionService.method.getSession, (input, call) => ({ ), })); +// Fixture runtimes never crash, so a revert returns the Session unchanged. +on(SessionService.method.revertSession, (input, call) => ({ + session: agentInstanceMessage(instanceFor(requireInstanceId(input.sessionId), call)), +})); + on(SessionService.method.createSession, (input, call) => { const namespace = requireNamespace(input.agent?.namespace ?? ""); if (!input.agent?.name.trim()) throw new ConnectError("an Agent is required", Code.InvalidArgument); diff --git a/ui/src/pages/AgentChatPage.recovery.test.tsx b/ui/src/pages/AgentChatPage.recovery.test.tsx new file mode 100644 index 0000000000..05dc95eedc --- /dev/null +++ b/ui/src/pages/AgentChatPage.recovery.test.tsx @@ -0,0 +1,247 @@ +import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { ThemeProvider } from "@emotion/react"; +import { MemoryRouter, Route, Routes } from "react-router-dom"; +import { SWRConfig } from "swr"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { createRouterTransport } from "@connectrpc/connect"; +import { SessionService } from "@/generated/kagent/api/v1/sessions_pb"; +import { CheckpointService, CheckpointState } from "@/generated/kagent/api/v1/checkpoints_pb"; +import { RuntimeOperation, RuntimeState } from "@/generated/kagent/api/v1/runtime_pb"; +import { setApiTransport } from "@/api/transport"; +import { ApiError, SEND_NOT_ACCEPTED, TURN_ACTIVE, type ChatController } from "@/api"; +import { themeFor } from "@/theme/theme"; +import { AgentChatPage } from "./AgentChatPage"; + +/* + * Crash recovery as the chat page shows it: read from the stored Session, plus an + * explicit Recover for a recovery that does not finish and for sends the gateway + * keeps refusing. Recover calls RevertSession, so it works whether or not the + * controller recovers on send. + */ + +// Only the A2A chat controller is faked; Session reads go through the real hooks and client. +const chat = { + messages: [], + isLoadingHistory: false, + turnState: "idle", + turnPhase: "idle", + phase: "idle", + refreshTranscript: vi.fn(() => Promise.resolve()), + send: vi.fn(), + cancel: vi.fn(), + dismissQuestion: vi.fn(), + answerQuestion: vi.fn(), + answerToolApproval: vi.fn(), + retry: vi.fn(), +} as unknown as ChatController; + +vi.mock("@/api", async (importOriginal) => ({ + ...(await importOriginal()), + useChat: () => chat, +})); + +const ID = "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44"; +const LIFECYCLE_POLL_MS = 1_000; + +type Lifecycle = { + state: RuntimeState; + operation: RuntimeOperation; + failure?: { reason: string; message: string }; +}; + +const READY: Lifecycle = { state: RuntimeState.READY, operation: RuntimeOperation.NONE }; +const RECOVERING: Lifecycle = { + state: RuntimeState.FAILED, + operation: RuntimeOperation.REVERT, + failure: { reason: "RuntimeCrashed", message: "runtime crashed" }, +}; +const LOST: Lifecycle = { + state: RuntimeState.FAILED, + operation: RuntimeOperation.NONE, + failure: { reason: "RuntimeLost", message: "The Session runtime no longer exists." }, +}; + +let lifecycle: Lifecycle; +let checkpointState: CheckpointState | undefined; +let getSession: ReturnType; +let revertSession: ReturnType; +/** What a revert leaves stored; READY unless a test says otherwise. */ +let afterRevert: Lifecycle; + +function session() { + return { + id: ID, + contextId: ID, + creator: "alice@example.com", + agent: { ref: { namespace: "kagent", name: "assistant" }, uid: "agent-uid" }, + ...lifecycle, + }; +} + +function renderPage() { + setApiTransport(createRouterTransport(({ service }) => { + service(SessionService, { + getSession: (getSession = vi.fn(() => ({ session: session() }))), + revertSession: (revertSession = vi.fn(() => { + lifecycle = afterRevert; + return { session: session() }; + })), + listSessions: () => ({ sessions: [session()] }), + }); + service(CheckpointService, { + listCheckpoints: () => ({ + checkpoints: checkpointState === undefined + ? [] + : [{ id: "checkpoint-1", sessionId: ID, headTaskId: "task-1", state: checkpointState }], + }), + }); + })); + render( + + new Map(), shouldRetryOnError: false, revalidateOnFocus: false }}> + + + } /> + + + + , + ); +} + +const composer = () => screen.getByTestId("chat-input"); + +async function tick(ms: number) { + await act(async () => { + vi.advanceTimersByTime(ms); + }); +} + +beforeEach(() => { + // Only intervals are faked: the page's polls are intervals, and SWR and the + // transport keep their real timeouts and promises. + vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] }); + checkpointState = undefined; + afterRevert = READY; + (chat as { turnError?: Error }).turnError = undefined; +}); + +afterEach(() => { + vi.restoreAllMocks(); + vi.useRealTimers(); + setApiTransport(undefined); +}); + +describe("a crash recovery in progress", () => { + it("says the conversation is recovering and enables the composer once it is ready", async () => { + lifecycle = RECOVERING; + renderPage(); + expect(await screen.findByText("Recovering this conversation after a crash")).toBeInTheDocument(); + expect(composer()).toBeDisabled(); + + // The stored revert operation keeps the lifecycle poll running. + lifecycle = READY; + const reads = getSession.mock.calls.length; + await tick(LIFECYCLE_POLL_MS); + await waitFor(() => expect(getSession.mock.calls.length).toBeGreaterThan(reads)); + + await waitFor(() => expect(composer()).toBeEnabled()); + expect(screen.queryByTestId("chat-not-ready")).not.toBeInTheDocument(); + }); +}); + +describe("a crash recovery that does not finish", () => { + it("offers Recover, which reverts the Session and enables the composer", async () => { + lifecycle = RECOVERING; + renderPage(); + const alert = await screen.findByTestId("chat-not-ready"); + fireEvent.click(within(alert).getByRole("button", { name: "Recover" })); + await waitFor(() => expect(revertSession).toHaveBeenCalledTimes(1)); + expect(revertSession.mock.calls[0][0]).toMatchObject({ sessionId: ID }); + await waitFor(() => expect(composer()).toBeEnabled()); + }); +}); + +describe("a send the gateway never accepted", () => { + const refused = () => new ApiError("The agent's runtime did not accept the message. It may have crashed.", { + kind: "http", code: "FailedPrecondition", reason: SEND_NOT_ACCEPTED, url: "A2AService/SendStreamingMessage", + }); + + it("says the runtime may have crashed and offers Recover", async () => { + lifecycle = READY; + (chat as { turnError?: Error }).turnError = refused(); + renderPage(); + const alert = await screen.findByTestId("chat-turn-error"); + expect(alert).toHaveTextContent("The agent's runtime may have crashed"); + fireEvent.click(within(alert).getByRole("button", { name: "Recover" })); + await waitFor(() => expect(revertSession).toHaveBeenCalledTimes(1)); + }); + + it("shows the recorded loss when Recover finds the runtime gone", async () => { + lifecycle = READY; + afterRevert = LOST; + (chat as { turnError?: Error }).turnError = refused(); + renderPage(); + fireEvent.click(within(await screen.findByTestId("chat-turn-error")).getByRole("button", { name: "Recover" })); + expect(await screen.findByText("This conversation's runtime cannot be recovered")).toBeInTheDocument(); + }); + + it("offers Recover when an earlier turn left the conversation refusing sends", async () => { + // A runtime that crashed mid-turn leaves that turn open, so sends are refused as a + // conflict while the Session still reads READY. + lifecycle = READY; + (chat as { turnError?: Error }).turnError = new ApiError("resource conflict", { + kind: "http", code: "FailedPrecondition", reason: TURN_ACTIVE, url: "A2AService/SendStreamingMessage", + }); + renderPage(); + const alert = await screen.findByTestId("chat-turn-error"); + expect(alert).toHaveTextContent("The agent's runtime may have crashed"); + expect(alert).toHaveTextContent("resource conflict"); + fireEvent.click(within(alert).getByRole("button", { name: "Recover" })); + await waitFor(() => expect(revertSession).toHaveBeenCalledTimes(1)); + }); + + it("leaves Recover to the not-ready alert while a recovery is pending", async () => { + lifecycle = RECOVERING; + (chat as { turnError?: Error }).turnError = refused(); + renderPage(); + await screen.findByTestId("chat-not-ready"); + const alert = await screen.findByTestId("chat-turn-error"); + expect(within(alert).queryByRole("button", { name: "Recover" })).not.toBeInTheDocument(); + expect(screen.getAllByRole("button", { name: "Recover" })).toHaveLength(1); + }); + + it("offers only Retry for any other failure", async () => { + lifecycle = READY; + (chat as { turnError?: Error }).turnError = new Error("model unavailable"); + renderPage(); + const alert = await screen.findByTestId("chat-turn-error"); + expect(within(alert).queryByRole("button", { name: "Recover" })).not.toBeInTheDocument(); + expect(within(alert).getByRole("button", { name: "Retry" })).toBeInTheDocument(); + }); +}); + +describe("a runtime that cannot be recovered", () => { + it("says so, keeps the composer disabled, and does not poll a state that never changes", async () => { + lifecycle = LOST; + renderPage(); + const alert = await screen.findByTestId("chat-not-ready"); + expect(alert).toHaveTextContent("This conversation's runtime cannot be recovered"); + expect(alert).toHaveTextContent("The Session runtime no longer exists."); + expect(alert).toHaveTextContent("Start another conversation. This one can still be read."); + expect(composer()).toBeDisabled(); + + const reads = getSession.mock.calls.length; + await tick(LIFECYCLE_POLL_MS * 3); + expect(getSession.mock.calls.length).toBe(reads); + }); + + it("offers forking only when a saved snapshot can fork", async () => { + lifecycle = LOST; + checkpointState = CheckpointState.READY; + renderPage(); + await waitFor(() => + expect(screen.getByTestId("chat-not-ready")).toHaveTextContent("fork a saved snapshot from the transcript"), + ); + }); +}); diff --git a/ui/src/pages/AgentChatPage.tsx b/ui/src/pages/AgentChatPage.tsx index e2241495de..0b11f68ba5 100644 --- a/ui/src/pages/AgentChatPage.tsx +++ b/ui/src/pages/AgentChatPage.tsx @@ -14,10 +14,12 @@ import { SnapshotDetailsModal } from "@/components/chat/SnapshotDetailsModal"; import { SnapshotRenameDialog } from "@/components/chat/SnapshotRenameDialog"; import { ChatTranscript } from "@/components/chat/ChatTranscript"; import { TaskPushNotificationsModal } from "@/components/chat/TaskPushNotificationsModal"; -import { isLifecycleBusy } from "@/components/chat/lifecycleReading"; +import { isLifecycleBusy, isRuntimeLost } from "@/components/chat/lifecycleReading"; import { paths } from "@/router/routes"; import { + ApiError, apiClient, + canForkFrom, useAgentInstance, useAgentInstances, useAgentTakesFiles, @@ -27,7 +29,7 @@ import { autoTitleFrom } from "@/components/agent-instances/instanceLabels"; import { useLiveTranscript } from "@/api/hooks/useLiveTranscript"; import { useInvalidateConversations } from "@/api/hooks/useInvalidateConversations"; import { useCheckpoints } from "@/api/hooks/useCheckpoints"; -import type { Checkpoint } from "@/api"; +import type { AgentInstance, Checkpoint } from "@/api"; import { useCollapsedBelow } from "@/components/chat/useNarrowViewport"; import { checkpointsByMessage } from "@/components/chat/messageCheckpoints"; import { messageSummary } from "@/components/chat/messageText"; @@ -36,6 +38,42 @@ import { useExtensionAgentLinks } from "@/appExtensions/hooks"; import { agentUrl } from "@/components/agent/agentUrl"; import { isMockMode } from "@/api/config"; +/** + * What the not-ready alert says about a conversation that cannot take a message. + * + * Read from the stored Session alone. A pending crash recovery and a runtime that is + * gone are both FAILED, and only the operation and failure reason tell them apart. + */ +function notReadyCopy( + instance: AgentInstance, + canFork: boolean, +): { title: string; description: string } { + if (isRuntimeLost(instance)) { + return { + title: "This conversation's runtime cannot be recovered", + description: [ + instance.failure?.message, + canFork + ? "Start another conversation, or fork a saved snapshot from the transcript. This one can still be read." + : "Start another conversation. This one can still be read.", + ].filter(Boolean).join(" "), + }; + } + if (instance.operation === "revert") { + return { + title: "Recovering this conversation after a crash", + description: + "The agent's runtime crashed and is being restored from its last snapshot. The agent accepts messages again once that finishes, and does not remember messages sent after that snapshot. If recovery does not finish, select Recover to try again.", + }; + } + return { + title: `This conversation is ${instance.state}, so the agent cannot answer`, + description: + instance.failure?.message || + "The agent accepts messages once this conversation is ready. You can also start another conversation.", + }; +} + /** * How often the instance is re-read while it is doing something. * @@ -452,6 +490,37 @@ export function AgentChatPage() { refreshInstance.current = instance.refresh; }); + /* + * Reverts the runtime on the reader's request. + * + * Offered while a recovery is pending, since a failed one stays FAILED/REVERT until + * something retries it, and when sends keep being refused, since a crashed runtime + * looks READY until something observes it. Either way it is explicit, so it works + * whether or not the controller recovers on send. A healthy runtime is returned + * unchanged, and one that cannot be recovered is recorded as lost. + */ + const [isRecovering, setRecovering] = useState(false); + const recover = useCallback(async () => { + if (!id) return; + setRecovering(true); + try { + const recovered = await apiClient.agentInstances.revert(id); + toast.success(recovered.state === "ready" + ? "The runtime is ready. Retry your message." + : "Recovery finished."); + } catch (cause: unknown) { + console.error("Could not recover the conversation:", cause); + toast.error(cause instanceof ApiError && cause.code === "Aborted" + ? "The runtime is busy with another operation. Try again shortly." + : `Could not recover: ${cause instanceof Error ? cause.message : String(cause)}`); + } finally { + setRecovering(false); + await refreshInstance.current().catch((cause: unknown) => { + console.error("Could not re-read the conversation:", cause); + }); + } + }, [id]); + useEffect(() => { if (!isBusy) return; let inFlight = false; @@ -625,11 +694,12 @@ export function AgentChatPage() { type="warning" showIcon data-testid="chat-not-ready" - title={`This conversation is ${state}, so the agent cannot answer`} - description={ - instance.data.failure?.message || - "The agent accepts messages once this conversation is ready. You can also start another conversation." - } + {...notReadyCopy(instance.data, (checkpoints.data ?? []).some(canForkFrom))} + action={instance.data.operation === "revert" ? ( + + ) : undefined} /> ) : null} @@ -643,6 +713,9 @@ export function AgentChatPage() { onRenameCheckpoint={renameSnapshot} onFork={forkCheckpoint} onDeleteCheckpoint={deleteCheckpoint} + // A conversation that is not ready has its own alert with Recover. + onRecover={canSend ? () => void recover() : undefined} + isRecovering={isRecovering} // The mock chat client has no A2A push store; hide the control rather // than showing a backend error over an otherwise scripted transcript. onManageTaskPush={isMockMode ? undefined : (taskId) => { diff --git a/ui/vite.config.ts b/ui/vite.config.ts index 8b9d1419c1..34b70821dd 100644 --- a/ui/vite.config.ts +++ b/ui/vite.config.ts @@ -1,6 +1,7 @@ import { defineConfig } from "vitest/config"; import { loadEnv, type Plugin } from "vite"; import react from "@vitejs/plugin-react"; +import { availableParallelism } from "node:os"; import path from "node:path"; import { CORE_ENV_VARS, ENV_DEFAULTS } from "./src/env.ts"; @@ -168,6 +169,15 @@ export default defineConfig(({ mode }) => ({ environment: "jsdom", setupFiles: ["./src/testSetup.ts"], css: true, + // At most eight workers, and one less than the cores below that. A worker + // rendering a whole page is CPU-bound and keeps about one and a half cores busy + // (V8's collector runs beside it), so vitest's default of one per core less one + // asks a large machine for more CPU than it has. The heaviest page tests need + // about a second on a free core. Oversubscribed, they stretch to 5-9 s and fail at + // the 5 s timeout or the 1 s findBy, a different few on each run. Fewer workers + // did not lengthen the suite, because the CPU was already the limit. Small CI + // runners keep the default. + maxWorkers: Math.max(1, Math.min(availableParallelism() - 1, 8)), /* Playwright's own specs are `.spec.ts` and must not run here. `.test.ts` under playwright/ is a check *about* that suite rather than part of it, and does. */ exclude: ["**/node_modules/**", "**/playwright/**/*.spec.ts"],