diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a7f7d861bb..a35dd256a9 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 b5e7d375de..7ecd8c1cae 100644 --- a/docs/architecture/process-configuration.md +++ b/docs/architecture/process-configuration.md @@ -99,6 +99,7 @@ push: signingPrivateKeyFile: secrets/push-private-key.pem sessions: idleTtl: 604800s + automaticRecovery: false runtimeDefaults: controlPlaneUrl: http://kagent-controller.kagent.svc:8083 gatewayUrl: http://kagent-controller.kagent.svc:8083 @@ -106,6 +107,10 @@ runtimeDefaults: level: info ``` +`sessions.automaticRecovery` enables recovery of crashed runtimes during sends. +It defaults to false. 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..87d4178343 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -265,25 +265,47 @@ 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`. + +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 active, and active work refuses every +send, so a send refused for active work 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. A send while recovery is pending +joins it. Reverting needs Session update access, as `RevertSession` does; a +caller who may only send gets the original refusal. The revert continues if the +caller disconnects. No Session read calls Substrate. + +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..b325ee31c6 100644 --- a/go/core/internal/a2agateway/gateway_test.go +++ b/go/core/internal/a2agateway/gateway_test.go @@ -55,6 +55,8 @@ type gatewayTestStore struct { initialID string listedIDs []string revoked bool + reverts int + onRevert func() session *apiv1.Session revision *database.RuntimeRevision err error @@ -808,6 +810,74 @@ 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 active, so a busy send probes the runtime once. +func TestGatewayBusySendRecoversACrashedTurn(t *testing.T) { + for _, crashed := range []bool{false, true} { + store := &gatewayTestStore{session: gatewayTestSession(), reserveErr: database.ErrConflict, taskErr: database.ErrNotFound} + if crashed { + store.onRevert = func() { store.reserveErr = nil } + } + runtime := &gatewayTestRuntime{} + gateway := newTestGateway(store, &gatewayTestAuthorizer{}, &gatewayTestDialer{client: gatewayTestClient(t, runtime)}, gatewayTestURL) + _, err := gateway.SendMessage(gatewayTestContext(), gatewayTestRequest()) + if store.reverts != 1 || store.reserveCalls != 2 { + t.Fatalf("crashed = %t: reverts = %d, reservations = %d", crashed, store.reverts, store.reserveCalls) + } + if crashed { + if runtime.sendCalls != 1 { + t.Fatalf("a restored runtime must receive the send: %d sends, err = %v", runtime.sendCalls, err) + } + continue + } + var protocolError *a2atype.Error + if !errors.As(err, &protocolError) || !errors.Is(err, a2atype.ErrUnsupportedOperation) || runtime.sendCalls != 0 { + t.Fatalf("a busy healthy runtime must refuse the send as before: %v, %d sends", err, runtime.sendCalls) + } + if meta, _ := protocolError.ErrorInfo().Value["metadata"].(map[string]string); meta["reason"] == "KAGENT_SEND_NOT_ACCEPTED" { + t.Fatalf("a busy healthy runtime must not invite a retry loop: %v", err) + } + } +} + +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..69f1f70c4f 100644 --- a/go/core/internal/a2agateway/routing_test.go +++ b/go/core/internal/a2agateway/routing_test.go @@ -43,18 +43,29 @@ 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.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 +142,42 @@ 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, busy, refused bool + want error + }{ + {name: "an admitted recovery", pending: true, want: a2a.ErrUnsupportedOperation}, + {name: "an active turn", busy: true, want: a2a.ErrUnsupportedOperation}, + {name: "a runtime refusal", refused: true, want: sessionsvc.ErrSendNotAccepted}, + } { + 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} + if scenario.busy { + store.reserveErr = database.ErrConflict + } + 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) + if !scenario.refused { + require.NotErrorIs(t, err, sessionsvc.ErrSendNotAccepted) + } + require.Zero(t, store.reverts) + if scenario.busy { + require.Equal(t, 1, store.reserveCalls) + } + }) + } +} 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/session_recovery.go b/go/core/internal/database/session_recovery.go index d591f33b5b..a739a9fb96 100644 --- a/go/core/internal/database/session_recovery.go +++ b/go/core/internal/database/session_recovery.go @@ -18,6 +18,17 @@ 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 +} + // 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 +130,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 +139,67 @@ 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 +} + +func lifecycleIdle(session *apiv1.Session) bool { + return session.State == apiv1.RuntimeState_RUNTIME_STATE_READY && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE +} + // 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 +236,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..43709159bb 100644 --- a/go/core/internal/database/session_recovery_test.go +++ b/go/core/internal/database/session_recovery_test.go @@ -327,3 +327,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/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..d67e23a9c7 100644 --- a/go/core/internal/service/session/interactions.go +++ b/go/core/internal/service/session/interactions.go @@ -48,14 +48,15 @@ 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 } -func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP bool) *InteractionService { - return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP} +func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP, automaticRecovery bool) *InteractionService { + return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP, automaticRecovery: automaticRecovery} } // PreparedSend either grants a new dispatch or returns an already accepted task. @@ -163,6 +164,21 @@ func (s *InteractionService) PrepareSend(ctx context.Context, agent types.Namesp req.Config = &config } id, err := s.reserveDispatch(ctx, session.Id, initialID) + if s.automaticRecovery && errors.Is(err, database.ErrConflict) { + // A crash during a turn leaves that turn active, and active work 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. A caller who may + // send but not revert gets the original refusal. + if recoverErr := s.recoverRuntime(ctx, session.Id); recoverErr == nil { + id, err = s.reserveDispatch(ctx, session.Id, initialID) + } else if !errors.Is(recoverErr, a2atype.ErrUnauthorized) { + return nil, recoverErr + } + } + if errors.Is(err, database.ErrConflict) { + err = interactionStoreError(ctx, err) + } if errors.Is(err, database.ErrMessageAccepted) { var historyLength *int if req.Config != nil { @@ -222,7 +238,7 @@ func (s *InteractionService) reserveDispatch(ctx context.Context, sessionID, ini 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) { @@ -260,6 +276,47 @@ 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.recoverRuntime(ctx, session.Id); err != nil { + 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") + +// recoverRuntime reverts a crashed runtime for a send that found it, and returns +// nil once the runtime is healthy or restored. Reverting needs the same update +// access as RevertSession, which sending alone does not grant. Contention and +// outages answer ErrSendNotAccepted, so the caller retries. The revert outlives a +// caller that disconnects, so a retry finds it finished. +func (s *InteractionService) recoverRuntime(ctx context.Context, sessionID string) error { + _, err := s.sessions.Revert(context.WithoutCancel(ctx), sessionID) + switch { + case err == nil: + return nil + case errors.Is(err, errRuntimeUnrecoverable): + return ErrRuntimeUnrecoverable + case serviceerrors.CodeOf(err) == serviceerrors.CodeAborted || serviceerrors.CodeOf(err) == serviceerrors.CodeUnavailable: + logging.FromContext(ctx).WarnContext(ctx, "recover crashed runtime", "session_id", sessionID, "error", err) + return ErrSendNotAccepted + default: + return interactionError(ctx, err) + } +} + // 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 +494,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.recoverRuntime(ctx, session.Id); err != nil { + 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..2f7ee1a9ad 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" @@ -25,14 +27,16 @@ func lifecycleIdle(session *apiv1.Session) bool { type crashObservation struct { generation *database.SessionGeneration actor *ateapipb.Actor + // lost replaces actor with why the recorded Actor cannot be recovered. + lost string } 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 lost, 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 +54,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{generation: generation, lost: "its Actor is missing"}, nil } if err != nil { return crashObservation{}, err } if !binding.Matches(actor) { - return crashObservation{}, fmt.Errorf("runtime Actor identity or template changed: %w", errRuntimeUnrecoverable) + return crashObservation{generation: generation, lost: "its Actor identity or template changed"}, nil } if actor.GetMetadata().GetUid() != generation.ActorUID { - return crashObservation{}, fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable) + return crashObservation{generation: generation, lost: "its Actor was replaced"}, nil } 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) } @@ -86,6 +111,9 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap if err != nil { return nil, err } + if observed.lost != "" { + return nil, w.recordLost(ctx, session.Id, observed.generation.OperationID, observed.lost) + } if !observed.crashed() { return session, nil } diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 2f0d05dcc5..168e69fe47 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,70 @@ func TestDeleteSupersedesCrashRecovery(t *testing.T) { }) } } + +func TestSendRecoveryAnswersTheSend(t *testing.T) { + for _, test := range []struct { + name string + err error + authErr error + want error + }{ + {name: "healthy or reverted"}, + {name: "contention", err: database.ErrConflict, want: ErrSendNotAccepted}, + {name: "substrate unavailable", err: status.Error(codes.Unavailable, "control plane unavailable"), want: ErrSendNotAccepted}, + {name: "unrecoverable", err: fmt.Errorf("runtime Actor is missing: %w", errRuntimeUnrecoverable), want: ErrRuntimeUnrecoverable}, + // Sending needs only create access; reverting needs RevertSession's update access. + {name: "no update access", authErr: serviceerrors.NewPermissionDenied("denied", nil), want: a2a.ErrUnauthorized}, + } { + t.Run(test.name, func(t *testing.T) { + session := &apiv1.Session{Id: uuid.NewString(), State: apiv1.RuntimeState_RUNTIME_STATE_READY} + 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.recoverRuntime(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") + }) + } +} + +// 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) + }) + } +} 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 2b3458dbd5..6671d9c182 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -376,7 +376,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..b64bfdd582 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,7 +27,10 @@ import ( "github.com/kagent-dev/kagent/go/core/internal/substrate" ) -func TestSessionRevertCrashedRuntime(t *testing.T) { +// 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. +func TestSessionSendRecoversCrashedRuntime(t *testing.T) { target := interactionTarget(t) forEachHarness(t, func(t *testing.T, harness testHarness) { cfg, err := mockllm.LoadConfigFromFile("mocks/invoke_agent.json", interactionMocks) @@ -126,18 +130,23 @@ func TestSessionRevertCrashedRuntime(t *testing.T) { return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED, err }) require.NoError(t, err) - reverted, err := fixture.sessions.RevertSession(fixture.ctx, request) + _, next := newMessageRequest(t, "What is 2+2?") + next.Tenant, next.Message.ContextId = fixture.tenant, fixture.sessionID + answered, err := sendMessageWithRetry(fixture.ctx, 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) - 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) - require.NoError(t, err) - require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, retried.Session.State) + 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") recovered, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) require.NoError(t, err) - require.Equal(t, running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid()) - require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, recovered.GetStatus().GetState()) - require.Equal(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running)), substrate.SnapshotURI(substrate.LatestActorSnapshot(recovered))) + require.Equal(t, running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid(), "recovery restores the recorded Actor in place") + // RevertSession afterwards finds a healthy runtime. + healthy, err = fixture.sessions.RevertSession(fixture.ctx, request) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) 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") @@ -147,8 +156,5 @@ func TestSessionRevertCrashedRuntime(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("interrupted caller did not stop") } - _, _, next := fixture.send(t, "What is 2+2?") - require.Equal(t, a2a.TaskStateCompleted, next.Status.State) - require.Contains(t, taskText(next), "4") }) } diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 69c5b77417..499b936b70 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -186,6 +186,8 @@ config: # Recovery scans for missed wake-ups and unbound registration cleanup. recoveryInterval: 60s sessions: + # Recover crashed runtimes when a client sends another message. + 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/chat/a2aGrpcChatClient.test.ts b/ui/src/api/chat/a2aGrpcChatClient.test.ts index 278ca76673..193a29a18e 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.test.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.test.ts @@ -23,6 +23,7 @@ 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 { LinkFile } from "./attachments"; @@ -2101,3 +2102,127 @@ 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(requests).toHaveLength(1); + }); + + it.each([ + ["another metadata reason", notAccepted({ reason: "SOMETHING_ELSE" })], + ["no metadata, as for a runtime that cannot be recovered", notAccepted({})], + ["another gRPC code", notAccepted(undefined, Code.Unavailable)], + ["no ErrorInfo", new ConnectError("Session runtime cannot be recovered; start another conversation", Code.FailedPrecondition)], + ])("does not re-send for %s", async (_, refusal) => { + const { error, requests } = await run(async function* () { + yield* []; + throw refusal; + }); + expect(error).toBeDefined(); + expect(requests).toHaveLength(1); + }); + + it("surfaces the refusal once the retry window is spent", async () => { + const { error, requests } = await run(async function* () { + yield* []; + throw notAccepted(); + }); + expect(error).toMatchObject({ code: "FailedPrecondition" }); + // One send every 100 ms for 30 s, and not one more. + expect(requests.length).toBeGreaterThan(250); + expect(requests.length).toBeLessThanOrEqual(301); + }); + + 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..c6635cdf7a 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.ts @@ -59,6 +59,7 @@ import { type Task as A2ATask, type TaskStatus, } from "@/generated/a2a_pb"; +import { ErrorInfoSchema } from "@/generated/google/rpc/error_details_pb"; import { ApiError, fromConnectError, rethrowIfAborted } from "../ApiError"; import { A2A_METADATA, metadataString } from "./a2aMetadata"; import { @@ -115,6 +116,39 @@ 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; + +/** `ErrorInfo.domain` of every error the A2A server writes (`a2a.ProtocolDomain`). */ +const A2A_PROTOCOL_DOMAIN = "a2a-protocol.org"; + +/** + * 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 { + if (!(error instanceof ConnectError) || error.code !== Code.FailedPrecondition) return undefined; + const info = error.findDetails(ErrorInfoSchema).find((detail) => + detail.domain === A2A_PROTOCOL_DOMAIN + && detail.reason === "UNSUPPORTED_OPERATION" + && detail.metadata.reason === "KAGENT_SEND_NOT_ACCEPTED"); + if (!info) return undefined; + const wait = Number.parseInt(info.metadata.retryAfterMs ?? "", 10); + return Number.isFinite(wait) && wait >= 0 ? wait : UNACCEPTED_SEND_RETRY_MS; +} + /** 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 +634,24 @@ 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. 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 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 +665,12 @@ export class A2AGrpcChatClient implements ChatClient { } return; } catch (error) { + const resendAfter = received ? undefined : unacceptedSendRetryMs(error); + if (resendAfter !== undefined && !signal?.aborted && Date.now() + resendAfter <= resendUntil) { + await pause(resendAfter, signal); + source = client.sendStreamingMessage(request, options); + continue; + } if (progressed || Date.now() - opened >= REATTACH_HEALTHY_MS) failures = 0; failures += 1; const reattach = !signal?.aborted && taskId !== "" && !settled 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/pages/AgentChatPage.recovery.test.tsx b/ui/src/pages/AgentChatPage.recovery.test.tsx new file mode 100644 index 0000000000..680c679d5d --- /dev/null +++ b/ui/src/pages/AgentChatPage.recovery.test.tsx @@ -0,0 +1,166 @@ +import { act, render, screen, waitFor } 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 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 alone. + * A send that finds a crashed runtime recovers it on the server; the page only says + * what the stored state already records. + */ + +// 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; + +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() }))), + 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; +}); + +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 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..819d78c15f 100644 --- a/ui/src/pages/AgentChatPage.tsx +++ b/ui/src/pages/AgentChatPage.tsx @@ -14,10 +14,11 @@ 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 { apiClient, + canForkFrom, useAgentInstance, useAgentInstances, useAgentTakesFiles, @@ -27,7 +28,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 +37,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.", + }; + } + 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. * @@ -625,11 +662,7 @@ 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))} /> ) : null}