diff --git a/docs/architecture/mcp.md b/docs/architecture/mcp.md index 661026b4c3..0ec7af648b 100644 --- a/docs/architecture/mcp.md +++ b/docs/architecture/mcp.md @@ -75,3 +75,30 @@ not protected: registering a name or URI core already uses replaces it. Authorization inside an added tool stays with the program that adds it. The implementation is in [`go/core/internal/mcp`](../../go/core/internal/mcp). + +## Go ADK runtime discovery + +The Go ADK connects to MCP servers and discovers tools during invocation +preparation. Startup and Actor resume do not wait for HTTP or SSE connections +or start stdio MCP processes. This also applies to subagents. + +Discovery filters app-only tools and applies configured allowlists before tools +reach the model. The MCP SDK caches filtered responses and owns TTL expiry and +catalog-notification invalidation. Cached tools remain filtered. Credentials, +approvals, and tool-call authorization continue to apply. + +If discovery fails, that server's toolset stays hidden for the current A2A +execution, including subsequent model steps and subagent calls. Healthy toolsets +remain available. The next execution retries, including an approval continuation. +There is no fallback to unclassified tools. Discovery failures are logged as +warnings. Cancellation and tool-call errors retain their existing handling. + +MCP App results are compacted for the model while stored results remain intact +for UI rendering. Removed app names remain recognized for historical compaction +during the process lifetime. Tools reclassified as ordinary keep their full +output, including historical results. A new process rediscovers app metadata +during invocation preparation. + +Standalone runners can use `mcp.WithDiscoveryScope(ctx)` for each `runner.Run` +to get the same omission behavior. Without it, discovery errors propagate. +Controller-side catalog discovery used by the UI is independent. diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index 7dfd1f0e8f..a76cc94697 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -223,6 +223,67 @@ claims still block new work, but completed results remain readable. The recorded actor UID is checked before lifecycle calls; a same-name replacement cannot be adopted implicitly. +`SessionService.RevertSession` lets an authorized caller restore a confirmed +crashed runtime. It accepts active and idle Sessions, but verifies Substrate's +CRASHED state for the recorded Actor UID and pinned template before admitting +recovery. Quiet turns, transport failures, missing Actors and same-name +replacements cannot authorize a revert. A runtime that is not crashed, including +a Session at rest outside READY, has nothing to recover, so RevertSession +returns the Session unchanged. That is also what a retry sees after a revert +whose response was lost. A missing or replaced Actor returns +`FAILED_PRECONDITION`, because no retry restores it. A transport failure returns `UNAVAILABLE`. Only an active +lifecycle operation returns `ABORTED`. + +Admission records FAILED/REVERT with a RuntimeCrashed failure under the Session +lock. REVERT is its own `RuntimeOperation`: only a confirmed crash starts one, and +`BeginSessionOperation` can join a REVERT but never start one. The same transaction +publishes a durable terminal result saved before the crash, or fails nonterminal +work absent from the retained snapshot. Waiting tasks captured by that snapshot +remain available. Checkpoint creation excludes competing recovery. + +Admission settles waits against the Actor's latest completed durable snapshot, +which is the one Substrate's revert restores. Every suspend kagent issues uses +synchronous durability, so its upload completes before suspension returns, and +a crashed Actor cannot be suspended. No upload can complete between admission +and the revert, so settling at admission settles against what is restored, and +the interrupted turn fails as soon as the crash is confirmed. An asynchronous +suspend would break this and require settling when the revert finishes. + +The caller reads the Session's generation before it observes Substrate: its +lifecycle generation, recorded Actor UID, and Actor template, in one statement. +The observed Actor must match that template and UID. Admission takes the read +back, rejects the crash when another lifecycle operation began after it, and +re-checks the Actor UID under the Session lock. A concurrent revert keeps the +Actor UID and returns the Session to READY, so only the generation exposes a +stale observation. The recorded UID changes only through a lifecycle operation, +which also replaces the generation. + +An issued idle suspension does not exclude recovery. It cannot finish against the +crashed Actor, and its worker keeps the claim when suspension fails. Admission +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. + +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. + +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. + ```mermaid sequenceDiagram participant Client diff --git a/docs/lifecycle-retries.md b/docs/lifecycle-retries.md index cbb2cf71ec..452b09ad1b 100644 --- a/docs/lifecycle-retries.md +++ b/docs/lifecycle-retries.md @@ -2,8 +2,8 @@ Session and Sandbox lifecycle mutations run within the API request. The server records the intended operation before calling the runtime and records completion -before returning success. Sessions expose Create, Fork, and Delete; explicit Suspend -and Resume are available only for Sandboxes. Session runtime suspension and resumption +before returning success. Sessions expose Create, Fork, Revert, and Delete; explicit +Suspend and Resume are available only for Sandboxes. Session runtime suspension and resumption are platform-managed. A failed request can have changed the runtime: a timeout does not mean that nothing happened, and canceling a request does not undo it. @@ -39,6 +39,29 @@ for subsequent invocations. Omitting it generates a new ID on each invocation, so rerunning that command is a new creation rather than a retry. MCP callers 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. +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. +It discards runtime changes since that snapshot and does not replay input. +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. +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. +Read the reason rather than the message text, which is not stable. + ## Errors | gRPC status | Client action | diff --git a/go/adk/examples/oneshot/main.go b/go/adk/examples/oneshot/main.go index 4e7709ef1d..ff4fb37f6f 100644 --- a/go/adk/examples/oneshot/main.go +++ b/go/adk/examples/oneshot/main.go @@ -27,6 +27,7 @@ import ( "github.com/kagent-dev/kagent/go/adk/pkg/agent" "github.com/kagent-dev/kagent/go/adk/pkg/config" + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" "github.com/kagent-dev/kagent/go/pkg/logging" ) @@ -106,7 +107,7 @@ func main() { streamMode = adkagent.StreamingModeSSE } - for ev, err := range r.Run(ctx, "user", sess.Session.ID(), content, adkagent.RunConfig{ + for ev, err := range r.Run(mcp.WithDiscoveryScope(ctx), "user", sess.Session.ID(), content, adkagent.RunConfig{ StreamingMode: streamMode, }) { if err != nil { diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index dd6469b7cf..03b4c845b3 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -22,6 +22,7 @@ import ( "google.golang.org/genai" "github.com/kagent-dev/kagent/go/adk/pkg/auth" + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" "github.com/kagent-dev/kagent/go/adk/pkg/models" "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" apia2a "github.com/kagent-dev/kagent/go/api/a2a" @@ -264,6 +265,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon ctx = context.WithValue(ctx, publicContextIDKey{}, reqCtx.ContextID) ctx = withBearerToken(ctx) + ctx = mcp.WithDiscoveryScope(ctx) // The synthetic ADK user ID is only a native session lookup key. Memory // and outgoing credentials must receive only the passed-through caller. ctx = auth.WithUserID(ctx, trustedUserID) diff --git a/go/adk/pkg/a2a/executor_mcp_test.go b/go/adk/pkg/a2a/executor_mcp_test.go new file mode 100644 index 0000000000..2e23f70fb7 --- /dev/null +++ b/go/adk/pkg/a2a/executor_mcp_test.go @@ -0,0 +1,118 @@ +package a2a + +import ( + "context" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "slices" + "sync/atomic" + "testing" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + adkagent "google.golang.org/adk/v2/agent" + + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestKAgentExecutorScopesMCPDiscoveryFailures(t *testing.T) { + // Given a temporarily unavailable MCP catalog, when execution continues and + // later resumes with reused IDs, assert that omission lasts only one execution. + var unavailable atomic.Bool + unavailable.Store(true) + var listings atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "execution-discovery", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2025-11-25"}, + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "available_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "ok"}, nil + }) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" { + listings.Add(1) + if unavailable.Load() { + return nil, errors.New("catalog temporarily unavailable") + } + } + return next(ctx, method, req) + } + }) + httpServer := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(httpServer.Close) + healthy := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "healthy-discovery"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(healthy, &mcpsdk.Tool{Name: "healthy_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "healthy"}, nil + }) + healthyHTTP := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return healthy }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(healthyHTTP.Close) + toolsets := mcp.CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{ + Params: &configv1.ADKHTTPConnectionParams{Url: httpServer.URL}, + }, { + Params: &configv1.ADKHTTPConnectionParams{Url: healthyHTTP.URL}, + }}, nil, nil, false, nil) + require.Len(t, toolsets, 2) + + var visible [][]string + builtin := &recordingExecutor{onExecute: func(ctx context.Context) { + readonly := &executionDiscoveryContext{StrictContextMock: adkagent.NewStrictContextMock(ctx)} + for range 2 { + var names []string + for _, toolset := range toolsets { + tools, err := toolset.Tools(readonly) + require.NoError(t, err, "discovery failures must leave the execution usable") + for _, tool := range tools { + names = append(names, tool.Name()) + } + } + slices.Sort(names) + visible = append(visible, names) + unavailable.Store(false) + } + }} + executor := &KAgentExecutor{builtin: builtin, logger: slog.New(slog.DiscardHandler)} + request := &a2asrv.ExecutorContext{ + TaskID: "reused-task", ContextID: "reused-session", + Message: a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")), + } + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"healthy_tool"}, {"healthy_tool"}}, visible, "recovery during an execution must not retry its omitted server") + require.EqualValues(t, 1, listings.Load()) + + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"available_tool", "healthy_tool"}, {"available_tool", "healthy_tool"}}, visible, "a later execution with the same IDs must retry discovery") + require.Greater(t, listings.Load(), int32(1)) + + // Given prior successful discovery, when the server goes down again, assert + // that healthy tools remain available and a fresh execution retries recovery. + unavailable.Store(true) + before := listings.Load() + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"healthy_tool"}, {"healthy_tool"}}, visible) + require.Equal(t, before+1, listings.Load(), "the failed catalog must stay omitted across model steps") + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"available_tool", "healthy_tool"}, {"available_tool", "healthy_tool"}}, visible) +} + +type executionDiscoveryContext struct { + adkagent.StrictContextMock +} + +func (*executionDiscoveryContext) InvocationID() string { return "reused-invocation" } +func (*executionDiscoveryContext) SessionID() string { return "reused-session" } diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go index 4c557f9b4c..3d8e8f8251 100644 --- a/go/adk/pkg/a2a/executor_test.go +++ b/go/adk/pkg/a2a/executor_test.go @@ -40,6 +40,7 @@ type recordingExecutor struct { userID string cleanupCalled bool events []a2atype.Event + onExecute func(context.Context) // span, when set, is a span the turn emits, the way ADK does. span string } @@ -48,6 +49,9 @@ func (e *recordingExecutor) Execute(ctx context.Context, reqCtx *a2asrv.Executor return func(yield func(a2atype.Event, error) bool) { e.message = reqCtx.Message e.userID = auth.UserIDFromContext(ctx) + if e.onExecute != nil { + e.onExecute(ctx) + } if e.span != "" { _, span := otel.Tracer("adk-test").Start(ctx, e.span) span.End() diff --git a/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go b/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go new file mode 100644 index 0000000000..e0a9dd8160 --- /dev/null +++ b/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go @@ -0,0 +1,76 @@ +package server + +import ( + "context" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/durationpb" + + kagentagent "github.com/kagent-dev/kagent/go/adk/pkg/agent" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestReadinessDoesNotWaitForMCPDiscovery(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "test-key") + t.Run("Given slow MCP servers and MCP-bound subagents, when the runtime starts, assert that readiness precedes discovery", func(t *testing.T) { + var requests atomic.Int32 + mcpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + select { + case <-time.After(500 * time.Millisecond): + case <-r.Context().Done(): + } + http.Error(w, "slow server", http.StatusServiceUnavailable) + })) + t.Cleanup(mcpServer.Close) + marker := filepath.Join(t.TempDir(), "stdio-started") + config := func(name string) *configv1.ADKAgentConfig { + return &configv1.ADKAgentConfig{ + Name: name, Instruction: "Use the configured tools.", SkillsDirectory: t.TempDir(), + Model: &configv1.ADKModelConfig{Model: "test-model", Provider: &configv1.ADKModelConfig_Openai{Openai: &configv1.ADKOpenAI{}}}, + HttpTools: []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: mcpServer.URL, Timeout: durationpb.New(time.Second)}}}, + SseTools: []*configv1.ADKSSEMCPServerConfig{{Params: &configv1.ADKSSEConnectionParams{Url: mcpServer.URL, Timeout: durationpb.New(time.Second)}}}, + } + } + rootConfig := config("root") + rootConfig.StdioTools = []*configv1.ADKStdioMCPServerConfig{{Command: "touch", Args: []string{marker}}} + rootConfig.SubAgents = []*configv1.ADKAgentConfig{config("child")} + _, err := kagentagent.CreateGoogleADKAgent(t.Context(), &configv1.ADKConfig{Agent: rootConfig}, "root", nil) + require.NoError(t, err) + srv, err := NewA2AServer(a2atype.AgentCard{}, substrateExecutor{}, slog.New(slog.DiscardHandler), ServerConfig{Host: "127.0.0.1", Port: "0"}) + require.NoError(t, err) + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv.readyServer.Addr = listener.Addr().String() + require.NoError(t, listener.Close()) + require.NoError(t, srv.Start()) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, srv.httpServer.Shutdown(ctx)) + require.NoError(t, srv.readyServer.Shutdown(ctx)) + srv.grpcServer.Stop() + }) + client := &http.Client{Timeout: time.Second} + resp, err := client.Get("http://" + srv.readyServer.Addr + "/readyz") + require.NoError(t, err) + _, err = io.Copy(io.Discard, resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Zero(t, requests.Load(), "root and subagent MCP connections must be deferred") + _, err = os.Stat(marker) + require.ErrorIs(t, err, os.ErrNotExist, "stdio process must not start before readiness") + }) +} diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 64413eacba..b5363b6288 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -98,7 +98,6 @@ func createGoogleADKAgent(ctx context.Context, agentConfig *configv1.ADKAgentCon log.InfoContext(ctx, "wired local skills", "skills_directory", skillsDirectory, "skill_count", len(skills)) } } - mcpAppToolNames := mcp.MCPAppToolNamesFromToolsets(toolsets) var remoteAgentTools []tool.Tool for _, remoteAgent := range agentConfig.RemoteAgents { @@ -150,12 +149,11 @@ func createGoogleADKAgent(ctx context.Context, agentConfig *configv1.ADKAgentCon beforeToolCallbacks := []llmagent.BeforeToolCallback{} beforeModelCallbacks := []llmagent.BeforeModelCallback{} - if len(mcpAppToolNames) > 0 { - // For MCP App-capable tools, keep rich tool payloads in chat history for UI rendering, - // but compact what is sent back to the model to avoid redundant polling/tool churn. - log.InfoContext(ctx, "wiring MCP App model result callback", "tool_count", len(mcpAppToolNames)) - beforeModelCallbacks = append(beforeModelCallbacks, MakeMCPAppModelResultCallback(mcpAppToolNames)) - } + // Tools are classified before each model request is prepared. Resolve app + // names here rather than capturing an empty snapshot during construction. + beforeModelCallbacks = append(beforeModelCallbacks, makeMCPAppModelResultCallback(func(ctx agent.Context) map[string]bool { + return mcp.MCPAppToolNamesFromToolsets(ctx, toolsets) + })) beforeToolCallbacks = append(beforeToolCallbacks, makeBeforeToolCallback(log)) llmAgentConfig := llmagent.Config{ diff --git a/go/adk/pkg/agent/mcp_apps.go b/go/adk/pkg/agent/mcp_apps.go index 38d10c4ff5..50617b6b8d 100644 --- a/go/adk/pkg/agent/mcp_apps.go +++ b/go/adk/pkg/agent/mcp_apps.go @@ -19,20 +19,22 @@ import ( // tool carrying a UI resourceUri, independent of the tool's name or payload keys. const mcpAppRenderedNotice = "The interactive UI for this tool has been rendered to the user and now updates live inside the app. Treat this as complete and do not call this tool again unless the user explicitly asks for it." -// MakeMCPAppModelResultCallback replaces what the model sees for MCP App +// makeMCPAppModelResultCallback replaces what the model sees for MCP App // (UI-rendering) tool results: instead of the heavy render payload it receives a // terminal directive (see mcpAppRenderedNotice). The full result is still // streamed to the UI separately, so this only changes the model's view and // prevents the model from looping on the rendering tool. Errors are passed // through so the model can still react to and recover from failures. -func MakeMCPAppModelResultCallback(appToolNames map[string]bool) llmagent.BeforeModelCallback { - return func(_ agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { +// Tool names are resolved during model preparation so discovery stays lazy. +func makeMCPAppModelResultCallback(appToolNames func(agent.Context) map[string]bool) llmagent.BeforeModelCallback { + return func(ctx agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { + names := appToolNames(ctx) for _, content := range req.Contents { if content == nil { continue } for _, part := range content.Parts { - if part == nil || part.FunctionResponse == nil || !appToolNames[part.FunctionResponse.Name] { + if part == nil || part.FunctionResponse == nil || !names[part.FunctionResponse.Name] { continue } part.FunctionResponse.Response = compactMCPAppModelResponse(part.FunctionResponse.Response) diff --git a/go/adk/pkg/agent/mcp_apps_test.go b/go/adk/pkg/agent/mcp_apps_test.go index 65b8e2ffd0..23d1e0a242 100644 --- a/go/adk/pkg/agent/mcp_apps_test.go +++ b/go/adk/pkg/agent/mcp_apps_test.go @@ -1,13 +1,26 @@ package agent import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" "testing" + "time" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + adkagent "google.golang.org/adk/v2/agent" adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/tool" + "google.golang.org/adk/v2/tool/toolconfirmation" "google.golang.org/genai" + + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" ) -func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing.T) { +func TestMCPAppModelResultReplacesRenderPayloadWithNotice(t *testing.T) { t.Parallel() req := &adkmodel.LLMRequest{ @@ -37,7 +50,9 @@ func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"jenkins_monitor_build": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"jenkins_monitor_build": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -70,7 +85,7 @@ func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing } } -func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { +func TestMCPAppModelResultPreservesIsError(t *testing.T) { t.Parallel() req := &adkmodel.LLMRequest{ @@ -91,7 +106,9 @@ func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"jenkins_monitor_build": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"jenkins_monitor_build": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -110,7 +127,7 @@ func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { } } -func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { +func TestMCPAppModelResultLeavesNonAppToolsAlone(t *testing.T) { t.Parallel() original := map[string]any{ @@ -131,7 +148,9 @@ func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"some_app_tool": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"some_app_tool": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -143,3 +162,128 @@ func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { t.Fatalf("non-app tool response modified: %#v", got) } } + +func TestMCPAppCompactionPreservesCacheApprovalAndHistoricalResults(t *testing.T) { + // Given real MCP tools with approvals, when fresh and cached catalogs + // prepare model results, assert that app compaction preserves rich results + // and approval checks, including history after the app tool is removed. + var lists, calls atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "app-compaction", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + for _, descriptor := range []*mcpsdk.Tool{ + {Name: "ordinary"}, + {Name: "panel", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://panel"}}}, + {Name: "app_only", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, + } { + mcpsdk.AddTool(server, descriptor, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + calls.Add(1) + return nil, map[string]any{"payload": "rich result"}, nil + }) + } + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + } + }) + srv := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(func() { srv.CloseClientConnections(); srv.Close() }) + sets := mcp.CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: srv.URL}, RequireApproval: true}}, nil, nil, false, nil) + require.Len(t, sets, 1) + ctx := &appCompactionContext{StrictContextMock: adkagent.NewStrictContextMock(t.Context()), confirmed: true} + callback := makeMCPAppModelResultCallback(func(ctx adkagent.Context) map[string]bool { + return mcp.MCPAppToolNamesFromToolsets(ctx, sets) + }) + var historical map[string]any + for range 2 { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + require.Len(t, listed, 2) + req := &adkmodel.LLMRequest{} + original := make(map[string]map[string]any) + for _, candidate := range listed { + fn := candidate.(interface { + ProcessRequest(adkagent.Context, *adkmodel.LLMRequest) error + Run(adkagent.Context, any) (map[string]any, error) + }) + require.NoError(t, fn.ProcessRequest(ctx, req)) + response, err := fn.Run(ctx, map[string]any{}) + require.NoError(t, err) + original[candidate.Name()] = response + req.Contents = append(req.Contents, &genai.Content{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: candidate.Name(), Response: response}}}}) + } + require.NotContains(t, req.Tools, "app_only") + _, err = callback(ctx, req) + require.NoError(t, err) + for _, content := range req.Contents { + response := content.Parts[0].FunctionResponse + if response.Name == "panel" { + require.Equal(t, mcpAppRenderedNotice, response.Response["content"].([]any)[0].(map[string]any)["text"]) + } else { + require.Equal(t, original["ordinary"], response.Response) + } + } + historical = original["panel"] + require.Equal(t, map[string]any{"output": map[string]any{"payload": "rich result"}}, historical, "history retains the original rich result") + ctx.confirmed = false + _, err = req.Tools["panel"].(interface { + Run(adkagent.Context, any) (map[string]any, error) + }).Run(ctx, map[string]any{}) + require.ErrorIs(t, err, tool.ErrConfirmationRejected) + ctx.confirmed = true + } + require.EqualValues(t, 1, lists.Load(), "cached model preparation makes no extra tools/list call") + require.EqualValues(t, 4, calls.Load(), "rejected approvals must not execute MCP tools") + server.RemoveTools("panel") + require.Eventually(t, func() bool { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + return len(listed) == 1 && listed[0].Name() == "ordinary" + }, 2*time.Second, 10*time.Millisecond) + req := &adkmodel.LLMRequest{Contents: []*genai.Content{{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: "panel", Response: historical}}}}}} + _, err := callback(ctx, req) + require.NoError(t, err) + require.Equal(t, mcpAppRenderedNotice, req.Contents[0].Parts[0].FunctionResponse.Response["content"].([]any)[0].(map[string]any)["text"]) + require.Equal(t, map[string]any{"output": map[string]any{"payload": "rich result"}}, historical) + + // Given the retained app name, when the tool returns as ordinary, assert + // that its new output reaches the model without the rendered notice. + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "panel"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"payload": "ordinary result"}, nil + }) + var ordinary tool.Tool + require.Eventually(t, func() bool { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + for _, candidate := range listed { + if candidate.Name() == "panel" { + ordinary = candidate + return !mcp.MCPAppToolNamesFromToolsets(ctx, sets)["panel"] + } + } + return false + }, time.Second, 10*time.Millisecond) + response, err := ordinary.(interface { + Run(adkagent.Context, any) (map[string]any, error) + }).Run(ctx, map[string]any{}) + require.NoError(t, err) + req = &adkmodel.LLMRequest{Contents: []*genai.Content{{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: "panel", Response: response}}}}}} + _, err = callback(ctx, req) + require.NoError(t, err) + require.Equal(t, map[string]any{"output": map[string]any{"payload": "ordinary result"}}, req.Contents[0].Parts[0].FunctionResponse.Response) +} + +type appCompactionContext struct { + adkagent.StrictContextMock + confirmed bool +} + +func (c *appCompactionContext) ToolConfirmation() *toolconfirmation.ToolConfirmation { + return &toolconfirmation.ToolConfirmation{Confirmed: c.confirmed} +} + +func (*appCompactionContext) FunctionCallID() string { return "app-compaction-call" } diff --git a/go/adk/pkg/mcp/discovery.go b/go/adk/pkg/mcp/discovery.go new file mode 100644 index 0000000000..d79692f483 --- /dev/null +++ b/go/adk/pkg/mcp/discovery.go @@ -0,0 +1,160 @@ +package mcp + +import ( + "context" + "fmt" + "sync" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + adkagent "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/tool" + + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +type discoveryScopeKey struct{} + +// WithDiscoveryScope remembers unavailable MCP toolsets for one runner.Run or +// A2A execution, including all its model steps and sub-agent calls. Start a new +// scope for each execution, even when an approval resume reuses InvocationID. +// The scope contains no connections and needs no explicit cleanup. +func WithDiscoveryScope(ctx context.Context) context.Context { + return context.WithValue(ctx, discoveryScopeKey{}, &discoveryScope{omitted: make(map[*mcpAppToolset]bool)}) +} + +type discoveryScope struct { + mu sync.Mutex + omitted map[*mcpAppToolset]bool +} + +func (d *discoveryScope) isOmitted(ts *mcpAppToolset) bool { + d.mu.Lock() + defer d.mu.Unlock() + return d.omitted[ts] +} + +func (d *discoveryScope) omit(ts *mcpAppToolset) { + d.mu.Lock() + defer d.mu.Unlock() + d.omitted[ts] = true +} + +// Raw metadata survives SDK cache hits, including pages from incomplete lists. +// Published app names change only after successful discovery and are retained +// after removal for historical result compaction. Visibility is enforced on +// each raw tools/list response, never by these maps. +type mcpAppDiscovery struct { + mu sync.Mutex + listedAppNames map[string]bool + appToolNames map[string]bool +} + +type mcpAppToolset struct { + inner tool.Toolset + discovery *mcpAppDiscovery +} + +var _ tool.Toolset = (*mcpAppToolset)(nil) + +func (m *mcpAppToolset) Name() string { return m.inner.Name() } + +func (m *mcpAppToolset) Tools(ctx adkagent.ReadonlyContext) ([]tool.Tool, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + scope, _ := ctx.Value(discoveryScopeKey{}).(*discoveryScope) + if scope != nil && scope.isOmitted(m) { + return nil, nil + } + tools, err := m.inner.Tools(ctx) + if ctx.Err() != nil { + return nil, ctx.Err() + } + if err != nil { + if scope == nil { + return nil, err + } + scope.omit(m) + logging.FromContext(ctx).WarnContext(ctx, "failed to discover MCP tools; omitting toolset for this invocation", "error", err) + return nil, nil + } + m.discovery.mu.Lock() + for _, t := range tools { + if m.discovery.listedAppNames[t.Name()] { + m.discovery.appToolNames[t.Name()] = true + } else { + delete(m.discovery.appToolNames, t.Name()) + } + } + m.discovery.mu.Unlock() + if scope != nil && scope.isOmitted(m) { + return nil, nil + } + return tools, nil +} + +// MCPAppToolNamesFromToolsets returns successfully discovered MCP App-capable +// tool names, including names retained for historical model-result compaction. +// Call it during model preparation, after Tools has run. +func MCPAppToolNamesFromToolsets(ctx context.Context, toolsets []tool.Toolset) map[string]bool { + out := make(map[string]bool) + scope, _ := ctx.Value(discoveryScopeKey{}).(*discoveryScope) + for _, ts := range toolsets { + aware, ok := ts.(*mcpAppToolset) + if !ok || scope != nil && scope.isOmitted(aware) { + continue + } + aware.discovery.mu.Lock() + for name := range aware.discovery.appToolNames { + out[name] = true + } + aware.discovery.mu.Unlock() + } + return out +} + +func newMCPAppClient(configuredFilter map[string]bool, discovery *mcpAppDiscovery) *mcpsdk.Client { + client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "kagent-adk"}, &mcpsdk.ClientOptions{ + Capabilities: mcpUIClientCapabilities(), + // New-protocol servers subscribe only when a handler is configured. + // The SDK itself invalidates its list cache on these notifications. + ToolListChangedHandler: func(context.Context, *mcpsdk.ToolListChangedRequest) {}, + }) + client.AddSendingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method != "tools/list" { + return next(ctx, method, req) + } + result, err := next(ctx, method, req) + if err != nil { + return nil, err + } + page, ok := result.(*mcpsdk.ListToolsResult) + if !ok || page == nil { + return nil, fmt.Errorf("invalid MCP tools/list result") + } + visible := make([]*mcpsdk.Tool, 0, len(page.Tools)) + discovery.mu.Lock() + for _, t := range page.Tools { + if t == nil || t.Name == "" || len(configuredFilter) > 0 && !configuredFilter[t.Name] { + continue + } + switch mcpToolKindOf(t.Meta) { + case mcpToolKindAppInternal: + continue + case mcpToolKindApp: + discovery.listedAppNames[t.Name] = true + default: + delete(discovery.listedAppNames, t.Name) + } + visible = append(visible, t) + } + discovery.mu.Unlock() + page.Tools = visible + // ListTools caches this already-filtered response after middleware + // returns. Preserve its TTL and all other response metadata. + return page, nil + } + }) + return client +} diff --git a/go/adk/pkg/mcp/discovery_test.go b/go/adk/pkg/mcp/discovery_test.go new file mode 100644 index 0000000000..35fbc5019a --- /dev/null +++ b/go/adk/pkg/mcp/discovery_test.go @@ -0,0 +1,675 @@ +package mcp + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + "google.golang.org/adk/v2/tool" + + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestCreateToolsetsDoesNotContactSlowServers(t *testing.T) { + t.Run("Given a slow MCP server, when toolsets are constructed, assert that no request delays startup", func(t *testing.T) { + var requests atomic.Int32 + server := newDiscoveryServer(t, 0, nil, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + select { + case <-time.After(500 * time.Millisecond): + case <-r.Context().Done(): + } + })) + sets := CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: server.URL}}}, nil, nil, false, nil) + require.Len(t, sets, 1) + require.Zero(t, requests.Load()) + }) +} + +func TestFailedClassificationIsOmittedUntilNextRun(t *testing.T) { + // Given a transient discovery failure, when multiple model steps run, assert + // that the server stays omitted until a new execution scope retries it. + var unavailable atomic.Bool + unavailable.Store(true) + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + if unavailable.Load() { + return nil, errors.New("catalog unavailable") + } + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load(), "ADK must not discover after classification fails") + unavailable.Store(false) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load()) + + // Same InvocationID, new run (including approval resumes). + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.Contains(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts}), "c_app") + // Another run's successful classification must not change the omitted run. + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts}), "another run's classification must not tag omitted tools in this run") +} + +func TestClassificationIncludesEveryPageAndPreservesAllowlist(t *testing.T) { + // Given one tool per page and a configured allowlist, when discovery runs, + // assert that later app tools survive and app-only tools remain excluded. + var lists atomic.Int32 + server := newDiscoveryServer(t, 1, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + require.Contains(t, req.GetSession().(*mcpsdk.ServerSession).InitializeParams().Capabilities.Extensions, mcpUIExtensionName) + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, map[string]bool{"b_internal": true, "c_app": true}) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load(), "classification must use ADK's own paginated discovery") +} + +func TestPartialClassificationDoesNotPublishTools(t *testing.T) { + // Given a later page fails, when discovery runs, assert that no partial + // classification or unclassified ADK fallback becomes model-visible. + var pages atomic.Int32 + server := newDiscoveryServer(t, 1, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && pages.Add(1) == 2 { + return nil, errors.New("second page unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + require.EqualValues(t, 2, pages.Load()) +} + +func TestClassificationWithoutScopeReturnsFailure(t *testing.T) { + // Given a standalone caller, when classification fails, assert that an error + // is returned without exposing ADK's unclassified catalog. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + return nil, errors.New("catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.ErrorContains(t, err, "catalog unavailable") + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load()) +} + +func TestDiscoveryFailureAfterSuccessfulDiscoveryIsOmittedUntilNextRun(t *testing.T) { + // Given one successful discovery, when the next listing fails, assert that + // omission lasts for that execution and a new execution can recover. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && lists.Add(1) == 2 { + return nil, errors.New("ADK catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 2, lists.Load(), "later model steps must not retry the omitted toolset") + + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load()) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed, "another execution's recovery must not undo omission") + require.EqualValues(t, 3, lists.Load()) + canceled, cancel := context.WithCancel(ctx) + cancel() + _, err = ts.Tools(testReadonlyContext{Context: canceled}) + require.ErrorIs(t, err, context.Canceled, "omission must not swallow cancellation") +} + +func TestDiscoveryFailureWithoutScopePropagatesAfterSuccess(t *testing.T) { + // Given a previously discovered standalone toolset, when its catalog fails, + // assert that discovery returns the error and the next call can retry. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && lists.Add(1) == 2 { + return nil, errors.New("standalone catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx := testReadonlyContext{Context: t.Context()} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + listed, err = ts.Tools(ctx) + require.ErrorContains(t, err, "standalone catalog unavailable") + require.Empty(t, listed) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load()) +} + +func TestEmptyLegacyCatalogRemainsClassified(t *testing.T) { + // Given an empty legacy catalog, when successive runs discover, assert that + // each model step lists once and a successful empty catalog remains valid. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + lists.Add(1) + result.(*mcpsdk.ListToolsResult).Tools = nil + } + return result, err + }, nil) + ts := discoveryToolset(t, server.URL, nil) + for range 2 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + } + require.EqualValues(t, 2, lists.Load(), "one discovery per legacy-protocol model step") +} + +func TestClassificationWaiterCanBeCanceled(t *testing.T) { + // Given discovery is already in flight, when another invocation is canceled, + // assert that it exits promptly and leaves the original discovery usable. + entered, release := make(chan struct{}), make(chan struct{}) + var once sync.Once + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + once.Do(func() { close(entered); <-release }) + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + done := make(chan error, 1) + go func() { + _, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + done <- err + }() + <-entered + ctx, cancel := context.WithCancel(WithDiscoveryScope(t.Context())) + waiter := &observedReadonlyContext{testReadonlyContext: testReadonlyContext{Context: ctx}, waiting: make(chan struct{})} + waiterDone := make(chan error, 1) + go func() { _, err := ts.Tools(waiter); waiterDone <- err }() + <-waiter.waiting + cancel() + require.ErrorIs(t, <-waiterDone, context.Canceled) + close(release) + require.NoError(t, <-done) +} + +func TestCanceledClassificationCanRecover(t *testing.T) { + // Given listing waits for cancellation, when the caller cancels discovery, + // assert that cancellation propagates and a later run can recover. + var canceled atomic.Bool + entered := make(chan struct{}) + release := make(chan struct{}) + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && !canceled.Load() { + close(entered) + <-release + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx, cancel := context.WithCancel(WithDiscoveryScope(t.Context())) + done := make(chan error, 1) + go func() { _, err := ts.Tools(testReadonlyContext{Context: ctx}); done <- err }() + <-entered + canceled.Store(true) + cancel() + close(release) + require.ErrorIs(t, <-done, context.Canceled) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) +} + +func TestSDKCachePreservesFilteredToolsAndAppMetadata(t *testing.T) { + // Given a cacheable paginated catalog, when multiple executions prepare + // tools, assert that cached pages remain filtered and keep app metadata. + var lists, discoveries atomic.Int32 + server := newCacheableDiscoveryServer(1, 60000) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + if method == "server/discover" { + discoveries.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, map[string]bool{"b_internal": true, "c_app": true}) + for range 3 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + } + require.EqualValues(t, 3, lists.Load(), "cached pages must not repeat discovery or bypass filtering") + require.EqualValues(t, 1, discoveries.Load(), "classification must share ADK's connection") +} + +func TestSDKCachePreservesEmptyFilteredCatalog(t *testing.T) { + // Given only an app-only tool is allowed, when cached tools are requested, + // assert that an empty filtered response remains empty without rediscovery. + var lists atomic.Int32 + server := newCacheableDiscoveryServer(0, 60000) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, map[string]bool{"b_internal": true}) + for range 2 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + } + require.EqualValues(t, 1, lists.Load()) +} + +func TestSDKCacheExpiryDiscoversNewTools(t *testing.T) { + // Given a cacheable catalog without notifications, when its TTL expires, + // assert that a newly advertised tool becomes visible through ADK discovery. + var updated atomic.Bool + var lists atomic.Int32 + server := newCacheableDiscoveryServer(0, 40) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + lists.Add(1) + if updated.Load() { + page := result.(*mcpsdk.ListToolsResult) + page.Tools = append(page.Tools, &mcpsdk.Tool{Name: "d_new", InputSchema: map[string]any{"type": "object"}}) + } + } + return result, err + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + updated.Store(true) + require.Eventually(t, func() bool { + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + return slices.Contains(discoveryNames(listed), "d_new") + }, time.Second, 10*time.Millisecond) + require.GreaterOrEqual(t, lists.Load(), int32(2)) +} + +func TestSDKNotificationRefreshesVisibilityAndRetainsHistoricalAppNames(t *testing.T) { + // Given a long-lived cache, when the catalog changes, assert that the SDK + // refreshes visibility while retaining app names used to compact old results. + server := newCacheableDiscoveryServer(0, 3600000) + url := serveDiscoveryServer(t, server, nil, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_plain", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, discoveryToolHandler) + server.RemoveTools("c_app") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "d_new"}, discoveryToolHandler) + require.Eventually(t, func() bool { + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + return slices.Equal(discoveryNames(listed), []string{"d_new"}) + }, 2*time.Second, 10*time.Millisecond) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) +} + +func TestPartialSDKCacheRetainsMetadataForLaterRecovery(t *testing.T) { + // Given an app page is cached before a later page fails, when a new + // execution recovers, assert that cached app metadata is still available. + var unavailable atomic.Bool + unavailable.Store(true) + var firstPages atomic.Int32 + server := newCacheableDiscoveryServer(1, 60000) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://first"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + if req.GetParams().(*mcpsdk.ListToolsParams).Cursor == "" { + firstPages.Add(1) + } else if unavailable.Load() { + return nil, errors.New("later page unavailable") + } + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_app", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"a_app": true, "c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + require.EqualValues(t, 1, firstPages.Load(), "the SDK must reuse its cached first page") + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed, "recovery cannot undo another execution's omission") +} + +func TestFailedDiscoveryDoesNotPublishStaleAppNamesAfterRecovery(t *testing.T) { + // Given an app page followed by a failure, when the server changes that + // tool to ordinary before recovery, assert that partial metadata is not published. + var unavailable atomic.Bool + unavailable.Store(true) + server := newCacheableDiscoveryServer(1, 0) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://first"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetParams().(*mcpsdk.ListToolsParams).Cursor != "" && unavailable.Load() { + return nil, errors.New("later page unavailable") + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app"}, discoveryToolHandler) + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_app", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) +} + +func TestFailedDiscoveryDoesNotRetainRemovedUnpublishedAppNames(t *testing.T) { + // Given a published ordinary tool, when an app reclassification listing + // fails and recovery removes it, assert that its app name is never published. + var unavailable atomic.Bool + server := newCacheableDiscoveryServer(1, 0) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_changing"}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetParams().(*mcpsdk.ListToolsParams).Cursor != "" && unavailable.Load() { + return nil, errors.New("later page unavailable") + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_changing", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_changing", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://changing"}}}, discoveryToolHandler) + unavailable.Store(true) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + server.RemoveTools("a_changing") + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts}), "failed partial metadata must not compact old ordinary results") +} + +func TestSuccessfulDiscoveryClearsReclassifiedAppNames(t *testing.T) { + // Given a published app tool, when a complete listing changes it to an + // ordinary tool, assert that only removed app names remain for compaction. + for _, protocol := range []string{"2025-11-25", "2026-07-28"} { + t.Run(protocol, func(t *testing.T) { + var lists atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "reclassification"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{protocol}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + addDiscoveryTools(server) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "removed_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://removed"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + _, err := ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, map[string]bool{"c_app": true, "removed_app": true}, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "c_app"}, discoveryToolHandler) + server.RemoveTools("removed_app") + require.Eventually(t, func() bool { + listed, err := ts.Tools(ctx) + require.NoError(t, err) + return slices.Equal(discoveryNames(listed), []string{"a_plain", "c_app"}) && + !MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})["c_app"] + }, time.Second, 10*time.Millisecond) + require.Equal(t, map[string]bool{"removed_app": true}, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + before := lists.Load() + _, err = ts.Tools(ctx) + require.NoError(t, err) + require.NotContains(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts}), "c_app") + if protocol == "2026-07-28" { + require.Equal(t, before, lists.Load(), "cached ordinary tools must not regain app classification") + } + }) + } +} + +func TestConcurrentDiscoveryKeepsEachResponseFiltered(t *testing.T) { + // Given concurrent uncached listings with different visibility, when both + // finish, assert that one response's classification never filters the other. + entered, release := make(chan struct{}), make(chan struct{}) + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + page := result.(*mcpsdk.ListToolsResult) + if lists.Add(1) == 1 { + close(entered) + <-release + page.Tools = []*mcpsdk.Tool{{Name: "same_name", InputSchema: map[string]any{"type": "object"}}} + } else { + page.Tools = []*mcpsdk.Tool{{Name: "same_name", InputSchema: map[string]any{"type": "object"}, Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}} + } + } + return result, err + }, nil) + ts := discoveryToolset(t, server.URL, nil) + var releaseOnce sync.Once + releaseFirst := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(releaseFirst) + done := make(chan []tool.Tool, 1) + go func() { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + if err != nil { + t.Error(err) + } + done <- listed + }() + <-entered + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + releaseFirst() + require.Equal(t, []string{"same_name"}, discoveryNames(<-done)) +} + +func TestStdioDiscoveryStartsOnlyOneProcess(t *testing.T) { + // Given a stdio MCP server, when construction and discovery run, assert + // that startup starts no process and classification uses ADK's one process. + marker := filepath.Join(t.TempDir(), "process-starts") + ts, err := initializeToolSet(t.Context(), mcpServerParams{URL: "helper", ServerType: "stdio", Command: os.Args[0], Args: []string{"-test.run=^TestStdioDiscoveryHelper$"}, Env: map[string]string{"KAGENT_MCP_DISCOVERY_CHILD": marker}}, nil) + require.NoError(t, err) + t.Cleanup(func() { + if err := ts.inner.(io.Closer).Close(); err != nil && !strings.Contains(err.Error(), "signal: killed") { + t.Errorf("close stdio: %v", err) + } + }) + _, err = os.Stat(marker) + require.ErrorIs(t, err, os.ErrNotExist) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"stdio_tool"}, discoveryNames(listed)) + content, err := os.ReadFile(marker) + require.NoError(t, err) + require.Equal(t, "started\n", string(content)) +} + +func TestStdioDiscoveryHelper(t *testing.T) { + marker := os.Getenv("KAGENT_MCP_DISCOVERY_CHILD") + if marker == "" { + t.Skip("subprocess helper") + } + f, err := os.OpenFile(marker, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600) + if err != nil { + os.Exit(2) + } + if _, err := f.WriteString("started\n"); err != nil { + os.Exit(3) + } + if err := f.Close(); err != nil { + os.Exit(4) + } + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "stdio-discovery"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "stdio_tool"}, discoveryToolHandler) + if err := server.Run(context.Background(), &mcpsdk.StdioTransport{}); err != nil { + os.Exit(5) + } + os.Exit(0) +} + +type discoveryMiddleware func(context.Context, string, mcpsdk.Request, mcpsdk.MethodHandler) (mcpsdk.Result, error) + +type observedReadonlyContext struct { + testReadonlyContext + waiting chan struct{} + once sync.Once +} + +func (o *observedReadonlyContext) Done() <-chan struct{} { + o.once.Do(func() { close(o.waiting) }) + return o.Context.Done() +} + +func newDiscoveryServer(t *testing.T, pageSize int, middleware discoveryMiddleware, intercept http.Handler) *httptest.Server { + t.Helper() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "discovery-test", Version: "1"}, &mcpsdk.ServerOptions{PageSize: pageSize, SupportedProtocolVersions: []string{"2025-11-25"}}) + addDiscoveryTools(server) + return serveDiscoveryServer(t, server, middleware, intercept, false) +} + +func newCacheableDiscoveryServer(pageSize, ttl int) *mcpsdk.Server { + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "cached-discovery-test", Version: "1"}, &mcpsdk.ServerOptions{ + PageSize: pageSize, SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = ttl }, + }) + addDiscoveryTools(server) + return server +} + +func addDiscoveryTools(server *mcpsdk.Server) { + for _, descriptor := range []*mcpsdk.Tool{ + {Name: "a_plain"}, + {Name: "b_internal", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, + {Name: "c_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://discovery"}}}, + } { + mcpsdk.AddTool(server, descriptor, discoveryToolHandler) + } +} + +func discoveryToolHandler(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "ok"}, nil +} + +func serveDiscoveryServer(t *testing.T, server *mcpsdk.Server, middleware discoveryMiddleware, intercept http.Handler, stateless bool) *httptest.Server { + t.Helper() + if middleware != nil { + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + return middleware(ctx, method, req, next) + } + }) + } + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: stateless}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if intercept != nil { + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(body)) + intercept.ServeHTTP(w, r) + } + handler.ServeHTTP(w, r) + })) + t.Cleanup(func() { srv.CloseClientConnections(); srv.Close() }) + return srv +} + +func discoveryToolset(t *testing.T, url string, filter map[string]bool) *mcpAppToolset { + t.Helper() + ts, err := initializeToolSet(t.Context(), mcpServerParams{URL: url, ServerType: "http"}, filter) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, ts.inner.(io.Closer).Close()) }) + return ts +} + +func discoveryNames(tools []tool.Tool) []string { + var names []string + for _, t := range tools { + names = append(names, t.Name()) + } + slices.Sort(names) + return names +} diff --git a/go/adk/pkg/mcp/mcp_ui.go b/go/adk/pkg/mcp/mcp_ui.go index 404a9831e0..515ef3bd2f 100644 --- a/go/adk/pkg/mcp/mcp_ui.go +++ b/go/adk/pkg/mcp/mcp_ui.go @@ -17,12 +17,7 @@ package mcp // in one place and is easy to track as the extension evolves. import ( - "context" - "fmt" - mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" - adkagent "google.golang.org/adk/v2/agent" - "google.golang.org/adk/v2/tool" ) const ( @@ -43,44 +38,6 @@ func mcpUIClientCapabilities() *mcpsdk.ClientCapabilities { return caps } -// mcpAppToolset wraps an MCP toolset and records which model-visible tools -// render as MCP App widgets. Classification happens during ListTools inside -// agentVisibleToolFilter because mcpsdk.Tool.Meta is not preserved on the ADK -// tool.Tool values the toolset exposes later. -type mcpAppToolset struct { - inner tool.Toolset - // appToolNames is this server's set of model-visible tool names whose - // results render as interactive MCP App (UI) widgets (the tool declares a - // `_meta.ui.resourceUri`). Used as a set, so only key presence is meaningful. - appToolNames map[string]bool -} - -func (m *mcpAppToolset) Name() string { - return m.inner.Name() -} - -func (m *mcpAppToolset) Tools(ctx adkagent.ReadonlyContext) ([]tool.Tool, error) { - return m.inner.Tools(ctx) -} - -// MCPAppToolNamesFromToolsets returns the union of MCP App-capable tool names -// recorded on toolsets built by CreateToolsets. The agent attaches the -// model-result compaction callback (see agent.MakeMCPAppModelResultCallback) -// only to these tools. -func MCPAppToolNamesFromToolsets(toolsets []tool.Toolset) map[string]bool { - out := make(map[string]bool) - for _, ts := range toolsets { - aware, ok := ts.(*mcpAppToolset) - if !ok { - continue - } - for name := range aware.appToolNames { - out[name] = true - } - } - return out -} - // mcpToolKind classifies how an MCP tool is surfaced, following the MCP Apps // extension. It is derived from the tool's `_meta.ui` block: presence of a // `resourceUri` means the tool renders an app, and the `visibility` field @@ -191,52 +148,3 @@ func normalizeVisibility(value any) []string { } return nil } - -// agentVisibleToolFilter lists tools from the MCP server, filters out app-only -// tools and any not in the configured allow-list, and returns a predicate the -// toolset can apply plus the MCP App-capable tool names discovered on this -// server. -// -// Classification must happen here because MCP Apps metadata lives on -// mcpsdk.Tool.Meta, which ADK mcptoolset drops when converting to tool.Tool. -func agentVisibleToolFilter(ctx context.Context, params mcpServerParams, configuredFilter map[string]bool) (tool.Predicate, map[string]bool, error) { - mcpTransport, err := createTransport(ctx, params) - if err != nil { - return nil, nil, fmt.Errorf("failed to create transport for %s: %w", params.URL, err) - } - - client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "kagent-adk"}, &mcpsdk.ClientOptions{Capabilities: mcpUIClientCapabilities()}) - session, err := client.Connect(ctx, mcpTransport, nil) - if err != nil { - return nil, nil, fmt.Errorf("failed to connect MCP client for %s: %w", params.URL, err) - } - defer session.Close() - - result, err := session.ListTools(ctx, &mcpsdk.ListToolsParams{}) - if err != nil { - return nil, nil, fmt.Errorf("failed to list MCP tools for %s: %w", params.URL, err) - } - - allowedTools := make([]string, 0, len(result.Tools)) - appToolNames := make(map[string]bool) - for _, t := range result.Tools { - if t == nil || t.Name == "" { - continue - } - if len(configuredFilter) > 0 && !configuredFilter[t.Name] { - continue - } - switch mcpToolKindOf(t.Meta) { - case mcpToolKindAppInternal: - // Hidden from the model; only the rendered MCP App calls it. - continue - case mcpToolKindApp: - allowedTools = append(allowedTools, t.Name) - appToolNames[t.Name] = true - default: // mcpToolKindModel - allowedTools = append(allowedTools, t.Name) - } - } - - return tool.StringPredicate(allowedTools), appToolNames, nil -} diff --git a/go/adk/pkg/mcp/passthrough_test.go b/go/adk/pkg/mcp/passthrough_test.go index 6cb1199ba6..b945a01c32 100644 --- a/go/adk/pkg/mcp/passthrough_test.go +++ b/go/adk/pkg/mcp/passthrough_test.go @@ -87,11 +87,13 @@ func TestTokenPassthroughToolsetDiscoversWithCallerToken(t *testing.T) { require.Empty(t, seen, "startup must not dial without a caller token") mu.Unlock() - _, err = toolset.Tools(testReadonlyContext{Context: t.Context()}) - require.Error(t, err) + // Given no caller token, when discovery runs, assert that no tools are exposed. + tools, err := toolset.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, tools) ctx := context.WithValue(t.Context(), models.BearerTokenKey, "alice") - tools, err := toolset.Tools(testReadonlyContext{Context: ctx}) + tools, err = toolset.Tools(testReadonlyContext{Context: WithDiscoveryScope(ctx)}) require.NoError(t, err) require.Len(t, tools, 1) require.Equal(t, "whoami", tools[0].Name()) diff --git a/go/adk/pkg/mcp/registry.go b/go/adk/pkg/mcp/registry.go index 600fcc401b..e7fd478941 100644 --- a/go/adk/pkg/mcp/registry.go +++ b/go/adk/pkg/mcp/registry.go @@ -371,31 +371,18 @@ func (rt *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, erro return rt.base.RoundTrip(req) } -// initializeToolSet fetches tools from an MCP server using Google ADK's -// mcptoolset and wraps the result with any MCP App-capable tool names found -// during classification. +// initializeToolSet prepares a lazy MCP toolset. Neither classification nor +// ADK discovery may connect before the runtime reports ready. func initializeToolSet(ctx context.Context, params mcpServerParams, toolFilter map[string]bool) (*mcpAppToolset, error) { mcpTransport, err := createTransport(ctx, params) if err != nil { return nil, fmt.Errorf("failed to create transport for %s: %w", params.URL, err) } - var toolPredicate tool.Predicate - var appToolNames map[string]bool - if params.TokenPassthrough { - // Startup has no caller token, so discover lazily on the first invocation. - toolPredicate = configuredToolPredicate(toolFilter) - } else { - toolPredicate, appToolNames, err = agentVisibleToolFilter(ctx, params, toolFilter) - if err != nil { - logging.FromContext(ctx).ErrorContext(ctx, "failed to classify MCP tools; falling back to lazy tool discovery", "error", err, "url", params.URL) - appToolNames = nil - toolPredicate = configuredToolPredicate(toolFilter) - } - } - + discovery := &mcpAppDiscovery{listedAppNames: make(map[string]bool), appToolNames: make(map[string]bool)} cfg := mcptoolset.Config{ Transport: mcpTransport, + Client: newMCPAppClient(toolFilter, discovery), } toolset, err := mcptoolset.New(cfg) @@ -403,24 +390,5 @@ func initializeToolSet(ctx context.Context, params mcpServerParams, toolFilter m return nil, fmt.Errorf("failed to create MCP toolset for %s: %w", params.URL, err) } - visibleTools := tool.Toolset(toolset) - if toolPredicate != nil { - visibleTools = tool.FilterToolset(toolset, toolPredicate) - } - return &mcpAppToolset{inner: visibleTools, appToolNames: appToolNames}, nil -} - -// configuredToolPredicate selects the configured tool names, or every tool -// when none are configured. -func configuredToolPredicate(toolFilter map[string]bool) tool.Predicate { - if len(toolFilter) == 0 { - return nil - } - allowedTools := make([]string, 0, len(toolFilter)) - for name, enabled := range toolFilter { - if enabled { - allowedTools = append(allowedTools, name) - } - } - return tool.StringPredicate(allowedTools) + return &mcpAppToolset{inner: toolset, discovery: discovery}, nil } diff --git a/go/adk/pkg/mcp/registry_test.go b/go/adk/pkg/mcp/registry_test.go index 2620816588..fc04f3aa7c 100644 --- a/go/adk/pkg/mcp/registry_test.go +++ b/go/adk/pkg/mcp/registry_test.go @@ -218,12 +218,12 @@ func TestMCPAppToolNamesFromToolsets(t *testing.T) { inner := &stubToolset{name: "mcp-server"} toolsets := []tool.Toolset{ - &mcpAppToolset{inner: inner, appToolNames: map[string]bool{"show_board": true}}, - &mcpAppToolset{inner: inner, appToolNames: map[string]bool{"refresh": true}}, + &mcpAppToolset{inner: inner, discovery: &mcpAppDiscovery{appToolNames: map[string]bool{"show_board": true}}}, + &mcpAppToolset{inner: inner, discovery: &mcpAppDiscovery{appToolNames: map[string]bool{"refresh": true}}}, inner, } - got := MCPAppToolNamesFromToolsets(toolsets) + got := MCPAppToolNamesFromToolsets(t.Context(), toolsets) if len(got) != 2 || !got["show_board"] || !got["refresh"] { t.Fatalf("MCPAppToolNamesFromToolsets() = %#v, want show_board and refresh", got) } diff --git a/go/api/gen/kagent/api/v1/runtime.pb.go b/go/api/gen/kagent/api/v1/runtime.pb.go index 88c28e575d..77fe6e91f0 100644 --- a/go/api/gen/kagent/api/v1/runtime.pb.go +++ b/go/api/gen/kagent/api/v1/runtime.pb.go @@ -95,6 +95,9 @@ const ( RuntimeOperation_RUNTIME_OPERATION_DELETE RuntimeOperation = 4 // No lifecycle operation is active. RuntimeOperation_RUNTIME_OPERATION_NONE RuntimeOperation = 5 + // Restores a crashed runtime to its last durable snapshot. Only a confirmed + // crash admits it; the Session stays FAILED until it finishes. + RuntimeOperation_RUNTIME_OPERATION_REVERT RuntimeOperation = 6 ) // Enum value maps for RuntimeOperation. @@ -106,6 +109,7 @@ var ( 3: "RUNTIME_OPERATION_RESUME", 4: "RUNTIME_OPERATION_DELETE", 5: "RUNTIME_OPERATION_NONE", + 6: "RUNTIME_OPERATION_REVERT", } RuntimeOperation_value = map[string]int32{ "RUNTIME_OPERATION_UNSPECIFIED": 0, @@ -114,6 +118,7 @@ var ( "RUNTIME_OPERATION_RESUME": 3, "RUNTIME_OPERATION_DELETE": 4, "RUNTIME_OPERATION_NONE": 5, + "RUNTIME_OPERATION_REVERT": 6, } ) @@ -156,14 +161,15 @@ const file_kagent_api_v1_runtime_proto_rawDesc = "" + "\x17RUNTIME_STATE_SUSPENDED\x10\x03\x12\x18\n" + "\x14RUNTIME_STATE_FAILED\x10\x04\x12\x1a\n" + "\x16RUNTIME_STATE_DELETING\x10\x05\x12\x19\n" + - "\x15RUNTIME_STATE_DELETED\x10\x06*\xca\x01\n" + + "\x15RUNTIME_STATE_DELETED\x10\x06*\xe8\x01\n" + "\x10RuntimeOperation\x12!\n" + "\x1dRUNTIME_OPERATION_UNSPECIFIED\x10\x00\x12\x1c\n" + "\x18RUNTIME_OPERATION_CREATE\x10\x01\x12\x1d\n" + "\x19RUNTIME_OPERATION_SUSPEND\x10\x02\x12\x1c\n" + "\x18RUNTIME_OPERATION_RESUME\x10\x03\x12\x1c\n" + "\x18RUNTIME_OPERATION_DELETE\x10\x04\x12\x1a\n" + - "\x16RUNTIME_OPERATION_NONE\x10\x05B=Z;github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1;apiv1b\x06proto3" + "\x16RUNTIME_OPERATION_NONE\x10\x05\x12\x1c\n" + + "\x18RUNTIME_OPERATION_REVERT\x10\x06B=Z;github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1;apiv1b\x06proto3" var ( file_kagent_api_v1_runtime_proto_rawDescOnce sync.Once diff --git a/go/api/gen/kagent/api/v1/sessions.pb.go b/go/api/gen/kagent/api/v1/sessions.pb.go index 5657cda54f..ab556d6f9b 100644 --- a/go/api/gen/kagent/api/v1/sessions.pb.go +++ b/go/api/gen/kagent/api/v1/sessions.pb.go @@ -669,6 +669,94 @@ func (x *UpdateSessionNameResponse) GetSession() *Session { return nil } +type RevertSessionRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RevertSessionRequest) Reset() { + *x = RevertSessionRequest{} + mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RevertSessionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RevertSessionRequest) ProtoMessage() {} + +func (x *RevertSessionRequest) ProtoReflect() protoreflect.Message { + mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RevertSessionRequest.ProtoReflect.Descriptor instead. +func (*RevertSessionRequest) Descriptor() ([]byte, []int) { + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{11} +} + +func (x *RevertSessionRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +type RevertSessionResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Session *Session `protobuf:"bytes,1,opt,name=session,proto3" json:"session,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RevertSessionResponse) Reset() { + *x = RevertSessionResponse{} + mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RevertSessionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RevertSessionResponse) ProtoMessage() {} + +func (x *RevertSessionResponse) ProtoReflect() protoreflect.Message { + mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RevertSessionResponse.ProtoReflect.Descriptor instead. +func (*RevertSessionResponse) Descriptor() ([]byte, []int) { + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{12} +} + +func (x *RevertSessionResponse) GetSession() *Session { + if x != nil { + return x.Session + } + return nil +} + type DeleteSessionRequest struct { state protoimpl.MessageState `protogen:"open.v1"` SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` @@ -681,7 +769,7 @@ type DeleteSessionRequest struct { func (x *DeleteSessionRequest) Reset() { *x = DeleteSessionRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -693,7 +781,7 @@ func (x *DeleteSessionRequest) String() string { func (*DeleteSessionRequest) ProtoMessage() {} func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -706,7 +794,7 @@ func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionRequest.ProtoReflect.Descriptor instead. func (*DeleteSessionRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{11} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{13} } func (x *DeleteSessionRequest) GetSessionId() string { @@ -732,7 +820,7 @@ type DeleteSessionResponse struct { func (x *DeleteSessionResponse) Reset() { *x = DeleteSessionResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -744,7 +832,7 @@ func (x *DeleteSessionResponse) String() string { func (*DeleteSessionResponse) ProtoMessage() {} func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -757,7 +845,7 @@ func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionResponse.ProtoReflect.Descriptor instead. func (*DeleteSessionResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{12} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{14} } func (x *DeleteSessionResponse) GetSession() *Session { @@ -781,7 +869,7 @@ type SessionShare struct { func (x *SessionShare) Reset() { *x = SessionShare{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -793,7 +881,7 @@ func (x *SessionShare) String() string { func (*SessionShare) ProtoMessage() {} func (x *SessionShare) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -806,7 +894,7 @@ func (x *SessionShare) ProtoReflect() protoreflect.Message { // Deprecated: Use SessionShare.ProtoReflect.Descriptor instead. func (*SessionShare) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{13} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{15} } func (x *SessionShare) GetId() string { @@ -850,7 +938,7 @@ type CreateSessionShareRequest struct { func (x *CreateSessionShareRequest) Reset() { *x = CreateSessionShareRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -862,7 +950,7 @@ func (x *CreateSessionShareRequest) String() string { func (*CreateSessionShareRequest) ProtoMessage() {} func (x *CreateSessionShareRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -875,7 +963,7 @@ func (x *CreateSessionShareRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateSessionShareRequest.ProtoReflect.Descriptor instead. func (*CreateSessionShareRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{14} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{16} } func (x *CreateSessionShareRequest) GetSessionId() string { @@ -903,7 +991,7 @@ type CreateSessionShareResponse struct { func (x *CreateSessionShareResponse) Reset() { *x = CreateSessionShareResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -915,7 +1003,7 @@ func (x *CreateSessionShareResponse) String() string { func (*CreateSessionShareResponse) ProtoMessage() {} func (x *CreateSessionShareResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -928,7 +1016,7 @@ func (x *CreateSessionShareResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateSessionShareResponse.ProtoReflect.Descriptor instead. func (*CreateSessionShareResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{15} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{17} } func (x *CreateSessionShareResponse) GetShare() *SessionShare { @@ -955,7 +1043,7 @@ type ListSessionSharesRequest struct { func (x *ListSessionSharesRequest) Reset() { *x = ListSessionSharesRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -967,7 +1055,7 @@ func (x *ListSessionSharesRequest) String() string { func (*ListSessionSharesRequest) ProtoMessage() {} func (x *ListSessionSharesRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -980,7 +1068,7 @@ func (x *ListSessionSharesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionSharesRequest.ProtoReflect.Descriptor instead. func (*ListSessionSharesRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{16} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{18} } func (x *ListSessionSharesRequest) GetSessionId() string { @@ -1007,7 +1095,7 @@ type ListSessionSharesResponse struct { func (x *ListSessionSharesResponse) Reset() { *x = ListSessionSharesResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1019,7 +1107,7 @@ func (x *ListSessionSharesResponse) String() string { func (*ListSessionSharesResponse) ProtoMessage() {} func (x *ListSessionSharesResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1032,7 +1120,7 @@ func (x *ListSessionSharesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionSharesResponse.ProtoReflect.Descriptor instead. func (*ListSessionSharesResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{17} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{19} } func (x *ListSessionSharesResponse) GetShares() []*SessionShare { @@ -1058,7 +1146,7 @@ type RevokeSessionShareRequest struct { func (x *RevokeSessionShareRequest) Reset() { *x = RevokeSessionShareRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1070,7 +1158,7 @@ func (x *RevokeSessionShareRequest) String() string { func (*RevokeSessionShareRequest) ProtoMessage() {} func (x *RevokeSessionShareRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1083,7 +1171,7 @@ func (x *RevokeSessionShareRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RevokeSessionShareRequest.ProtoReflect.Descriptor instead. func (*RevokeSessionShareRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{18} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{20} } func (x *RevokeSessionShareRequest) GetShareId() string { @@ -1101,7 +1189,7 @@ type RevokeSessionShareResponse struct { func (x *RevokeSessionShareResponse) Reset() { *x = RevokeSessionShareResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1113,7 +1201,7 @@ func (x *RevokeSessionShareResponse) String() string { func (*RevokeSessionShareResponse) ProtoMessage() {} func (x *RevokeSessionShareResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[21] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1126,7 +1214,7 @@ func (x *RevokeSessionShareResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RevokeSessionShareResponse.ProtoReflect.Descriptor instead. func (*RevokeSessionShareResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{19} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{21} } var File_kagent_api_v1_sessions_proto protoreflect.FileDescriptor @@ -1182,6 +1270,11 @@ const file_kagent_api_v1_sessions_proto_rawDesc = "" + "session_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\tsessionId\x12Q\n" + "\x04name\x18\x02 \x01(\tB=\xbaH:r8\x18\xc8\x0123^(?:$|[^\\p{Z}\\p{Cc}](?:[^\\p{Cc}]*[^\\p{Z}\\p{Cc}])?)$R\x04name\"M\n" + "\x19UpdateSessionNameResponse\x120\n" + + "\asession\x18\x01 \x01(\v2\x16.kagent.api.v1.SessionR\asession\"?\n" + + "\x14RevertSessionRequest\x12'\n" + + "\n" + + "session_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\tsessionId\"I\n" + + "\x15RevertSessionResponse\x120\n" + "\asession\x18\x01 \x01(\v2\x16.kagent.api.v1.SessionR\asession\"U\n" + "\x14DeleteSessionRequest\x12'\n" + "\n" + @@ -1213,13 +1306,14 @@ const file_kagent_api_v1_sessions_proto_rawDesc = "" + "\x04page\x18\x02 \x01(\v2\x1b.kagent.api.v1.PageResponseR\x04page\"@\n" + "\x19RevokeSessionShareRequest\x12#\n" + "\bshare_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\ashareId\"\x1c\n" + - "\x1aRevokeSessionShareResponse2\x9a\x06\n" + + "\x1aRevokeSessionShareResponse2\xf6\x06\n" + "\x0eSessionService\x12Z\n" + "\rCreateSession\x12#.kagent.api.v1.CreateSessionRequest\x1a$.kagent.api.v1.CreateSessionResponse\x12Q\n" + "\n" + "GetSession\x12 .kagent.api.v1.GetSessionRequest\x1a!.kagent.api.v1.GetSessionResponse\x12W\n" + "\fListSessions\x12\".kagent.api.v1.ListSessionsRequest\x1a#.kagent.api.v1.ListSessionsResponse\x12f\n" + "\x11UpdateSessionName\x12'.kagent.api.v1.UpdateSessionNameRequest\x1a(.kagent.api.v1.UpdateSessionNameResponse\x12Z\n" + + "\rRevertSession\x12#.kagent.api.v1.RevertSessionRequest\x1a$.kagent.api.v1.RevertSessionResponse\x12Z\n" + "\rDeleteSession\x12#.kagent.api.v1.DeleteSessionRequest\x1a$.kagent.api.v1.DeleteSessionResponse\x12i\n" + "\x12CreateSessionShare\x12(.kagent.api.v1.CreateSessionShareRequest\x1a).kagent.api.v1.CreateSessionShareResponse\x12f\n" + "\x11ListSessionShares\x12'.kagent.api.v1.ListSessionSharesRequest\x1a(.kagent.api.v1.ListSessionSharesResponse\x12i\n" + @@ -1237,7 +1331,7 @@ func file_kagent_api_v1_sessions_proto_rawDescGZIP() []byte { return file_kagent_api_v1_sessions_proto_rawDescData } -var file_kagent_api_v1_sessions_proto_msgTypes = make([]protoimpl.MessageInfo, 20) +var file_kagent_api_v1_sessions_proto_msgTypes = make([]protoimpl.MessageInfo, 22) var file_kagent_api_v1_sessions_proto_goTypes = []any{ (*Failure)(nil), // 0: kagent.api.v1.Failure (*Session)(nil), // 1: kagent.api.v1.Session @@ -1250,68 +1344,73 @@ var file_kagent_api_v1_sessions_proto_goTypes = []any{ (*ListSessionsResponse)(nil), // 8: kagent.api.v1.ListSessionsResponse (*UpdateSessionNameRequest)(nil), // 9: kagent.api.v1.UpdateSessionNameRequest (*UpdateSessionNameResponse)(nil), // 10: kagent.api.v1.UpdateSessionNameResponse - (*DeleteSessionRequest)(nil), // 11: kagent.api.v1.DeleteSessionRequest - (*DeleteSessionResponse)(nil), // 12: kagent.api.v1.DeleteSessionResponse - (*SessionShare)(nil), // 13: kagent.api.v1.SessionShare - (*CreateSessionShareRequest)(nil), // 14: kagent.api.v1.CreateSessionShareRequest - (*CreateSessionShareResponse)(nil), // 15: kagent.api.v1.CreateSessionShareResponse - (*ListSessionSharesRequest)(nil), // 16: kagent.api.v1.ListSessionSharesRequest - (*ListSessionSharesResponse)(nil), // 17: kagent.api.v1.ListSessionSharesResponse - (*RevokeSessionShareRequest)(nil), // 18: kagent.api.v1.RevokeSessionShareRequest - (*RevokeSessionShareResponse)(nil), // 19: kagent.api.v1.RevokeSessionShareResponse - (RuntimeState)(0), // 20: kagent.api.v1.RuntimeState - (RuntimeOperation)(0), // 21: kagent.api.v1.RuntimeOperation - (*timestamppb.Timestamp)(nil), // 22: google.protobuf.Timestamp - (*ResourceReference)(nil), // 23: kagent.api.v1.ResourceReference - (*PageRequest)(nil), // 24: kagent.api.v1.PageRequest - (*PageResponse)(nil), // 25: kagent.api.v1.PageResponse - (*durationpb.Duration)(nil), // 26: google.protobuf.Duration + (*RevertSessionRequest)(nil), // 11: kagent.api.v1.RevertSessionRequest + (*RevertSessionResponse)(nil), // 12: kagent.api.v1.RevertSessionResponse + (*DeleteSessionRequest)(nil), // 13: kagent.api.v1.DeleteSessionRequest + (*DeleteSessionResponse)(nil), // 14: kagent.api.v1.DeleteSessionResponse + (*SessionShare)(nil), // 15: kagent.api.v1.SessionShare + (*CreateSessionShareRequest)(nil), // 16: kagent.api.v1.CreateSessionShareRequest + (*CreateSessionShareResponse)(nil), // 17: kagent.api.v1.CreateSessionShareResponse + (*ListSessionSharesRequest)(nil), // 18: kagent.api.v1.ListSessionSharesRequest + (*ListSessionSharesResponse)(nil), // 19: kagent.api.v1.ListSessionSharesResponse + (*RevokeSessionShareRequest)(nil), // 20: kagent.api.v1.RevokeSessionShareRequest + (*RevokeSessionShareResponse)(nil), // 21: kagent.api.v1.RevokeSessionShareResponse + (RuntimeState)(0), // 22: kagent.api.v1.RuntimeState + (RuntimeOperation)(0), // 23: kagent.api.v1.RuntimeOperation + (*timestamppb.Timestamp)(nil), // 24: google.protobuf.Timestamp + (*ResourceReference)(nil), // 25: kagent.api.v1.ResourceReference + (*PageRequest)(nil), // 26: kagent.api.v1.PageRequest + (*PageResponse)(nil), // 27: kagent.api.v1.PageResponse + (*durationpb.Duration)(nil), // 28: google.protobuf.Duration } var file_kagent_api_v1_sessions_proto_depIdxs = []int32{ 2, // 0: kagent.api.v1.Session.agent:type_name -> kagent.api.v1.AgentProvenance - 20, // 1: kagent.api.v1.Session.state:type_name -> kagent.api.v1.RuntimeState - 21, // 2: kagent.api.v1.Session.operation:type_name -> kagent.api.v1.RuntimeOperation + 22, // 1: kagent.api.v1.Session.state:type_name -> kagent.api.v1.RuntimeState + 23, // 2: kagent.api.v1.Session.operation:type_name -> kagent.api.v1.RuntimeOperation 0, // 3: kagent.api.v1.Session.failure:type_name -> kagent.api.v1.Failure - 22, // 4: kagent.api.v1.Session.created_at:type_name -> google.protobuf.Timestamp - 22, // 5: kagent.api.v1.Session.updated_at:type_name -> google.protobuf.Timestamp - 23, // 6: kagent.api.v1.AgentProvenance.ref:type_name -> kagent.api.v1.ResourceReference - 23, // 7: kagent.api.v1.CreateSessionRequest.agent:type_name -> kagent.api.v1.ResourceReference + 24, // 4: kagent.api.v1.Session.created_at:type_name -> google.protobuf.Timestamp + 24, // 5: kagent.api.v1.Session.updated_at:type_name -> google.protobuf.Timestamp + 25, // 6: kagent.api.v1.AgentProvenance.ref:type_name -> kagent.api.v1.ResourceReference + 25, // 7: kagent.api.v1.CreateSessionRequest.agent:type_name -> kagent.api.v1.ResourceReference 1, // 8: kagent.api.v1.CreateSessionResponse.session:type_name -> kagent.api.v1.Session 1, // 9: kagent.api.v1.GetSessionResponse.session:type_name -> kagent.api.v1.Session - 24, // 10: kagent.api.v1.ListSessionsRequest.page:type_name -> kagent.api.v1.PageRequest - 23, // 11: kagent.api.v1.ListSessionsRequest.agent:type_name -> kagent.api.v1.ResourceReference + 26, // 10: kagent.api.v1.ListSessionsRequest.page:type_name -> kagent.api.v1.PageRequest + 25, // 11: kagent.api.v1.ListSessionsRequest.agent:type_name -> kagent.api.v1.ResourceReference 1, // 12: kagent.api.v1.ListSessionsResponse.sessions:type_name -> kagent.api.v1.Session - 25, // 13: kagent.api.v1.ListSessionsResponse.page:type_name -> kagent.api.v1.PageResponse + 27, // 13: kagent.api.v1.ListSessionsResponse.page:type_name -> kagent.api.v1.PageResponse 1, // 14: kagent.api.v1.UpdateSessionNameResponse.session:type_name -> kagent.api.v1.Session - 1, // 15: kagent.api.v1.DeleteSessionResponse.session:type_name -> kagent.api.v1.Session - 22, // 16: kagent.api.v1.SessionShare.created_at:type_name -> google.protobuf.Timestamp - 22, // 17: kagent.api.v1.SessionShare.expires_at:type_name -> google.protobuf.Timestamp - 26, // 18: kagent.api.v1.CreateSessionShareRequest.ttl:type_name -> google.protobuf.Duration - 13, // 19: kagent.api.v1.CreateSessionShareResponse.share:type_name -> kagent.api.v1.SessionShare - 24, // 20: kagent.api.v1.ListSessionSharesRequest.page:type_name -> kagent.api.v1.PageRequest - 13, // 21: kagent.api.v1.ListSessionSharesResponse.shares:type_name -> kagent.api.v1.SessionShare - 25, // 22: kagent.api.v1.ListSessionSharesResponse.page:type_name -> kagent.api.v1.PageResponse - 3, // 23: kagent.api.v1.SessionService.CreateSession:input_type -> kagent.api.v1.CreateSessionRequest - 5, // 24: kagent.api.v1.SessionService.GetSession:input_type -> kagent.api.v1.GetSessionRequest - 7, // 25: kagent.api.v1.SessionService.ListSessions:input_type -> kagent.api.v1.ListSessionsRequest - 9, // 26: kagent.api.v1.SessionService.UpdateSessionName:input_type -> kagent.api.v1.UpdateSessionNameRequest - 11, // 27: kagent.api.v1.SessionService.DeleteSession:input_type -> kagent.api.v1.DeleteSessionRequest - 14, // 28: kagent.api.v1.SessionService.CreateSessionShare:input_type -> kagent.api.v1.CreateSessionShareRequest - 16, // 29: kagent.api.v1.SessionService.ListSessionShares:input_type -> kagent.api.v1.ListSessionSharesRequest - 18, // 30: kagent.api.v1.SessionService.RevokeSessionShare:input_type -> kagent.api.v1.RevokeSessionShareRequest - 4, // 31: kagent.api.v1.SessionService.CreateSession:output_type -> kagent.api.v1.CreateSessionResponse - 6, // 32: kagent.api.v1.SessionService.GetSession:output_type -> kagent.api.v1.GetSessionResponse - 8, // 33: kagent.api.v1.SessionService.ListSessions:output_type -> kagent.api.v1.ListSessionsResponse - 10, // 34: kagent.api.v1.SessionService.UpdateSessionName:output_type -> kagent.api.v1.UpdateSessionNameResponse - 12, // 35: kagent.api.v1.SessionService.DeleteSession:output_type -> kagent.api.v1.DeleteSessionResponse - 15, // 36: kagent.api.v1.SessionService.CreateSessionShare:output_type -> kagent.api.v1.CreateSessionShareResponse - 17, // 37: kagent.api.v1.SessionService.ListSessionShares:output_type -> kagent.api.v1.ListSessionSharesResponse - 19, // 38: kagent.api.v1.SessionService.RevokeSessionShare:output_type -> kagent.api.v1.RevokeSessionShareResponse - 31, // [31:39] is the sub-list for method output_type - 23, // [23:31] is the sub-list for method input_type - 23, // [23:23] is the sub-list for extension type_name - 23, // [23:23] is the sub-list for extension extendee - 0, // [0:23] is the sub-list for field type_name + 1, // 15: kagent.api.v1.RevertSessionResponse.session:type_name -> kagent.api.v1.Session + 1, // 16: kagent.api.v1.DeleteSessionResponse.session:type_name -> kagent.api.v1.Session + 24, // 17: kagent.api.v1.SessionShare.created_at:type_name -> google.protobuf.Timestamp + 24, // 18: kagent.api.v1.SessionShare.expires_at:type_name -> google.protobuf.Timestamp + 28, // 19: kagent.api.v1.CreateSessionShareRequest.ttl:type_name -> google.protobuf.Duration + 15, // 20: kagent.api.v1.CreateSessionShareResponse.share:type_name -> kagent.api.v1.SessionShare + 26, // 21: kagent.api.v1.ListSessionSharesRequest.page:type_name -> kagent.api.v1.PageRequest + 15, // 22: kagent.api.v1.ListSessionSharesResponse.shares:type_name -> kagent.api.v1.SessionShare + 27, // 23: kagent.api.v1.ListSessionSharesResponse.page:type_name -> kagent.api.v1.PageResponse + 3, // 24: kagent.api.v1.SessionService.CreateSession:input_type -> kagent.api.v1.CreateSessionRequest + 5, // 25: kagent.api.v1.SessionService.GetSession:input_type -> kagent.api.v1.GetSessionRequest + 7, // 26: kagent.api.v1.SessionService.ListSessions:input_type -> kagent.api.v1.ListSessionsRequest + 9, // 27: kagent.api.v1.SessionService.UpdateSessionName:input_type -> kagent.api.v1.UpdateSessionNameRequest + 11, // 28: kagent.api.v1.SessionService.RevertSession:input_type -> kagent.api.v1.RevertSessionRequest + 13, // 29: kagent.api.v1.SessionService.DeleteSession:input_type -> kagent.api.v1.DeleteSessionRequest + 16, // 30: kagent.api.v1.SessionService.CreateSessionShare:input_type -> kagent.api.v1.CreateSessionShareRequest + 18, // 31: kagent.api.v1.SessionService.ListSessionShares:input_type -> kagent.api.v1.ListSessionSharesRequest + 20, // 32: kagent.api.v1.SessionService.RevokeSessionShare:input_type -> kagent.api.v1.RevokeSessionShareRequest + 4, // 33: kagent.api.v1.SessionService.CreateSession:output_type -> kagent.api.v1.CreateSessionResponse + 6, // 34: kagent.api.v1.SessionService.GetSession:output_type -> kagent.api.v1.GetSessionResponse + 8, // 35: kagent.api.v1.SessionService.ListSessions:output_type -> kagent.api.v1.ListSessionsResponse + 10, // 36: kagent.api.v1.SessionService.UpdateSessionName:output_type -> kagent.api.v1.UpdateSessionNameResponse + 12, // 37: kagent.api.v1.SessionService.RevertSession:output_type -> kagent.api.v1.RevertSessionResponse + 14, // 38: kagent.api.v1.SessionService.DeleteSession:output_type -> kagent.api.v1.DeleteSessionResponse + 17, // 39: kagent.api.v1.SessionService.CreateSessionShare:output_type -> kagent.api.v1.CreateSessionShareResponse + 19, // 40: kagent.api.v1.SessionService.ListSessionShares:output_type -> kagent.api.v1.ListSessionSharesResponse + 21, // 41: kagent.api.v1.SessionService.RevokeSessionShare:output_type -> kagent.api.v1.RevokeSessionShareResponse + 33, // [33:42] is the sub-list for method output_type + 24, // [24:33] is the sub-list for method input_type + 24, // [24:24] is the sub-list for extension type_name + 24, // [24:24] is the sub-list for extension extendee + 0, // [0:24] is the sub-list for field type_name } func init() { file_kagent_api_v1_sessions_proto_init() } @@ -1327,7 +1426,7 @@ func file_kagent_api_v1_sessions_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_kagent_api_v1_sessions_proto_rawDesc), len(file_kagent_api_v1_sessions_proto_rawDesc)), NumEnums: 0, - NumMessages: 20, + NumMessages: 22, NumExtensions: 0, NumServices: 1, }, diff --git a/go/api/gen/kagent/api/v1/sessions_grpc.pb.go b/go/api/gen/kagent/api/v1/sessions_grpc.pb.go index f5bfe3112d..4c262b9426 100644 --- a/go/api/gen/kagent/api/v1/sessions_grpc.pb.go +++ b/go/api/gen/kagent/api/v1/sessions_grpc.pb.go @@ -23,6 +23,7 @@ const ( SessionService_GetSession_FullMethodName = "/kagent.api.v1.SessionService/GetSession" SessionService_ListSessions_FullMethodName = "/kagent.api.v1.SessionService/ListSessions" SessionService_UpdateSessionName_FullMethodName = "/kagent.api.v1.SessionService/UpdateSessionName" + SessionService_RevertSession_FullMethodName = "/kagent.api.v1.SessionService/RevertSession" SessionService_DeleteSession_FullMethodName = "/kagent.api.v1.SessionService/DeleteSession" SessionService_CreateSessionShare_FullMethodName = "/kagent.api.v1.SessionService/CreateSessionShare" SessionService_ListSessionShares_FullMethodName = "/kagent.api.v1.SessionService/ListSessionShares" @@ -43,6 +44,8 @@ type SessionServiceClient interface { GetSession(ctx context.Context, in *GetSessionRequest, opts ...grpc.CallOption) (*GetSessionResponse, error) ListSessions(ctx context.Context, in *ListSessionsRequest, opts ...grpc.CallOption) (*ListSessionsResponse, error) UpdateSessionName(ctx context.Context, in *UpdateSessionNameRequest, opts ...grpc.CallOption) (*UpdateSessionNameResponse, error) + // Restores a crashed runtime to its last snapshot without replaying its turn. + RevertSession(ctx context.Context, in *RevertSessionRequest, opts ...grpc.CallOption) (*RevertSessionResponse, error) DeleteSession(ctx context.Context, in *DeleteSessionRequest, opts ...grpc.CallOption) (*DeleteSessionResponse, error) CreateSessionShare(ctx context.Context, in *CreateSessionShareRequest, opts ...grpc.CallOption) (*CreateSessionShareResponse, error) ListSessionShares(ctx context.Context, in *ListSessionSharesRequest, opts ...grpc.CallOption) (*ListSessionSharesResponse, error) @@ -97,6 +100,16 @@ func (c *sessionServiceClient) UpdateSessionName(ctx context.Context, in *Update return out, nil } +func (c *sessionServiceClient) RevertSession(ctx context.Context, in *RevertSessionRequest, opts ...grpc.CallOption) (*RevertSessionResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RevertSessionResponse) + err := c.cc.Invoke(ctx, SessionService_RevertSession_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *sessionServiceClient) DeleteSession(ctx context.Context, in *DeleteSessionRequest, opts ...grpc.CallOption) (*DeleteSessionResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(DeleteSessionResponse) @@ -151,6 +164,8 @@ type SessionServiceServer interface { GetSession(context.Context, *GetSessionRequest) (*GetSessionResponse, error) ListSessions(context.Context, *ListSessionsRequest) (*ListSessionsResponse, error) UpdateSessionName(context.Context, *UpdateSessionNameRequest) (*UpdateSessionNameResponse, error) + // Restores a crashed runtime to its last snapshot without replaying its turn. + RevertSession(context.Context, *RevertSessionRequest) (*RevertSessionResponse, error) DeleteSession(context.Context, *DeleteSessionRequest) (*DeleteSessionResponse, error) CreateSessionShare(context.Context, *CreateSessionShareRequest) (*CreateSessionShareResponse, error) ListSessionShares(context.Context, *ListSessionSharesRequest) (*ListSessionSharesResponse, error) @@ -177,6 +192,9 @@ func (UnimplementedSessionServiceServer) ListSessions(context.Context, *ListSess func (UnimplementedSessionServiceServer) UpdateSessionName(context.Context, *UpdateSessionNameRequest) (*UpdateSessionNameResponse, error) { return nil, status.Error(codes.Unimplemented, "method UpdateSessionName not implemented") } +func (UnimplementedSessionServiceServer) RevertSession(context.Context, *RevertSessionRequest) (*RevertSessionResponse, error) { + return nil, status.Error(codes.Unimplemented, "method RevertSession not implemented") +} func (UnimplementedSessionServiceServer) DeleteSession(context.Context, *DeleteSessionRequest) (*DeleteSessionResponse, error) { return nil, status.Error(codes.Unimplemented, "method DeleteSession not implemented") } @@ -282,6 +300,24 @@ func _SessionService_UpdateSessionName_Handler(srv interface{}, ctx context.Cont return interceptor(ctx, in, info, handler) } +func _SessionService_RevertSession_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RevertSessionRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(SessionServiceServer).RevertSession(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: SessionService_RevertSession_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(SessionServiceServer).RevertSession(ctx, req.(*RevertSessionRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _SessionService_DeleteSession_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(DeleteSessionRequest) if err := dec(in); err != nil { @@ -377,6 +413,10 @@ var SessionService_ServiceDesc = grpc.ServiceDesc{ MethodName: "UpdateSessionName", Handler: _SessionService_UpdateSessionName_Handler, }, + { + MethodName: "RevertSession", + Handler: _SessionService_RevertSession_Handler, + }, { MethodName: "DeleteSession", Handler: _SessionService_DeleteSession_Handler, diff --git a/go/api/utils/utils.go b/go/api/utils/utils.go deleted file mode 100644 index 1225656e89..0000000000 --- a/go/api/utils/utils.go +++ /dev/null @@ -1,39 +0,0 @@ -package utils - -import ( - "context" - "fmt" - - corev1 "k8s.io/api/core/v1" - "sigs.k8s.io/controller-runtime/pkg/client" -) - -// GetSecretValue fetches a value from a Secret -func GetSecretValue(ctx context.Context, c client.Client, ref client.ObjectKey, key string) (string, error) { - secret := &corev1.Secret{} - err := c.Get(ctx, ref, secret) - if err != nil { - return "", fmt.Errorf("failed to find Secret %s: %v", ref.String(), err) - } - - value, exists := secret.Data[key] - if !exists { - return "", fmt.Errorf("key %s not found in Secret %s", key, ref.String()) - } - return string(value), nil -} - -// GetConfigMapValue fetches a value from a ConfigMap -func GetConfigMapValue(ctx context.Context, c client.Client, ref client.ObjectKey, key string) (string, error) { - configMap := &corev1.ConfigMap{} - err := c.Get(ctx, ref, configMap) - if err != nil { - return "", fmt.Errorf("failed to find ConfigMap for %s: %v", ref.String(), err) - } - - value, exists := configMap.Data[key] - if !exists { - return "", fmt.Errorf("key %s not found in ConfigMap %s", key, ref) - } - return value, nil -} diff --git a/go/api/v1alpha3/common_types.go b/go/api/v1alpha3/common_types.go index 98379f776c..fd1d7d6051 100644 --- a/go/api/v1alpha3/common_types.go +++ b/go/api/v1alpha3/common_types.go @@ -17,16 +17,12 @@ limitations under the License. package v1alpha3 import ( - "context" "fmt" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/types" - "sigs.k8s.io/controller-runtime/pkg/client" - - "github.com/kagent-dev/kagent/go/api/utils" ) // ObjectReference identifies a resource; an omitted namespace means the owner's namespace. @@ -147,21 +143,6 @@ type ValueSource struct { Key string `json:"key"` } -func (s *ValueSource) Resolve(ctx context.Context, client client.Client, namespace string) (string, error) { - if s == nil { - return "", fmt.Errorf("ValueSource cannot be nil") - } - - switch s.Type { - case ConfigMapValueSource: - return utils.GetConfigMapValue(ctx, client, types.NamespacedName{Namespace: namespace, Name: s.Name}, s.Key) - case SecretValueSource: - return utils.GetSecretValue(ctx, client, types.NamespacedName{Namespace: namespace, Name: s.Name}, s.Key) - default: - return "", fmt.Errorf("unknown value source type: %s", s.Type) - } -} - // ValueRef represents a configuration value // +kubebuilder:validation:ExactlyOneOf=value;valueFrom type ValueRef struct { @@ -173,23 +154,3 @@ type ValueRef struct { // +optional ValueFrom *ValueSource `json:"valueFrom,omitempty"` } - -func (r *ValueRef) Resolve(ctx context.Context, client client.Client, namespace string) (string, string, error) { - if r == nil { - return "", "", fmt.Errorf("ValueRef cannot be nil") - } - - switch { - case r.Value != "": - return r.Name, r.Value, nil - case r.ValueFrom != nil: - value, err := r.ValueFrom.Resolve(ctx, client, namespace) - if err != nil { - return "", "", fmt.Errorf("failed to resolve value for ref %s: %v", r.Name, err) - } - - return r.Name, value, nil - default: - return r.Name, "", nil - } -} diff --git a/go/api/v1alpha3/remotemcpserver_types.go b/go/api/v1alpha3/remotemcpserver_types.go index d3453ed56e..b28fa3be0c 100644 --- a/go/api/v1alpha3/remotemcpserver_types.go +++ b/go/api/v1alpha3/remotemcpserver_types.go @@ -17,15 +17,8 @@ limitations under the License. package v1alpha3 import ( - "context" - "database/sql" - "database/sql/driver" - "encoding/json" - "fmt" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "sigs.k8s.io/controller-runtime/pkg/client" ) // +kubebuilder:validation:Enum=SSE;STREAMABLE_HTTP @@ -95,22 +88,6 @@ type RemoteMCPServerSpec struct { TLS *TLSConfig `json:"tls,omitempty"` } -var _ sql.Scanner = (*RemoteMCPServerSpec)(nil) - -func (t *RemoteMCPServerSpec) Scan(src any) error { - switch v := src.(type) { - case []uint8: - return json.Unmarshal(v, t) - } - return nil -} - -var _ driver.Valuer = (*RemoteMCPServerSpec)(nil) - -func (t RemoteMCPServerSpec) Value() (driver.Value, error) { - return json.Marshal(t) -} - // RemoteMCPServerStatus defines the observed state of RemoteMCPServer. type RemoteMCPServerStatus struct { // INSERT ADDITIONAL STATUS FIELD - define observed state of cluster @@ -153,22 +130,6 @@ type RemoteMCPServer struct { Status RemoteMCPServerStatus `json:"status,omitempty"` } -// ResolveHeaders resolves all HeadersFrom entries using the object's namespace. -func (r *RemoteMCPServer) ResolveHeaders(ctx context.Context, client client.Client) (map[string]string, error) { - result := map[string]string{} - - for _, h := range r.Spec.HeadersFrom { - k, v, err := h.Resolve(ctx, client, r.Namespace) - if err != nil { - return nil, fmt.Errorf("failed to resolve header: %v", err) - } - - result[k] = v - } - - return result, nil -} - // +kubebuilder:object:root=true // RemoteMCPServerList contains a list of RemoteMCPServer. type RemoteMCPServerList struct { diff --git a/go/core/cli/internal/common/exec/docker.go b/go/core/cli/internal/common/exec/docker.go deleted file mode 100644 index 65c4087052..0000000000 --- a/go/core/cli/internal/common/exec/docker.go +++ /dev/null @@ -1,95 +0,0 @@ -// Package exec provides utilities for executing external commands. -package exec - -import ( - "fmt" - "os" - "os/exec" - "strings" -) - -// DockerExecutor provides utilities for running Docker commands with consistent -// error handling, logging, and configuration across the CLI. -type DockerExecutor struct { - // Verbose enables detailed logging of Docker commands - Verbose bool - - // WorkDir sets the working directory for Docker commands - WorkDir string -} - -// NewDockerExecutor creates a new Docker executor with the specified configuration. -func NewDockerExecutor(verbose bool, workDir string) *DockerExecutor { - return &DockerExecutor{ - Verbose: verbose, - WorkDir: workDir, - } -} - -// CheckAvailability checks if Docker is installed and the daemon is running. -// Returns an error if Docker is not available or not properly configured. -func (d *DockerExecutor) CheckAvailability() error { - if _, err := exec.LookPath("docker"); err != nil { - return fmt.Errorf("docker command not found in PATH. Please install Docker") - } - - // Check if Docker daemon is running by querying the server version - cmd := exec.Command("docker", "version", "--format", "{{.Server.Version}}") - output, err := cmd.Output() - if err != nil { - return fmt.Errorf("docker daemon is not running or not accessible. Please start Docker Desktop or Docker daemon") - } - - // Verify we got a valid version string - version := strings.TrimSpace(string(output)) - if version == "" { - return fmt.Errorf("docker daemon returned empty version. Docker may not be properly installed") - } - - return nil -} - -// Run executes a docker command with the given arguments. -// It automatically handles verbose logging and working directory configuration. -func (d *DockerExecutor) Run(args ...string) error { - if d.Verbose { - fmt.Printf("Running: docker %s\n", strings.Join(args, " ")) - if d.WorkDir != "" { - fmt.Printf("Working directory: %s\n", d.WorkDir) - } - } - - cmd := exec.Command("docker", args...) - if d.WorkDir != "" { - cmd.Dir = d.WorkDir - } - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - - return cmd.Run() -} - -// Build builds a Docker image with the specified name and context. -// Additional build arguments can be provided (e.g., "--platform", "linux/amd64"). -func (d *DockerExecutor) Build(imageName, context string, extraArgs ...string) error { - args := []string{"build", "-t", imageName} - args = append(args, extraArgs...) - args = append(args, context) - - if err := d.Run(args...); err != nil { - return fmt.Errorf("docker build failed: %w", err) - } - - fmt.Printf("✅ Successfully built Docker image: %s\n", imageName) - return nil -} - -// Push pushes a Docker image to a registry. -func (d *DockerExecutor) Push(imageName string) error { - if err := d.Run("push", imageName); err != nil { - return fmt.Errorf("docker push failed: %w", err) - } - - fmt.Printf("✅ Successfully pushed Docker image: %s\n", imageName) - return nil -} diff --git a/go/core/cli/internal/common/exec/kubectl.go b/go/core/cli/internal/common/exec/kubectl.go index d640a9f089..6e92166dca 100644 --- a/go/core/cli/internal/common/exec/kubectl.go +++ b/go/core/cli/internal/common/exec/kubectl.go @@ -1,14 +1,9 @@ package exec import ( - "bytes" - "context" "fmt" - "io" - "os" "os/exec" "strings" - "time" ) // KubectlExecutor provides utilities for running kubectl commands with consistent @@ -29,34 +24,6 @@ func NewKubectlExecutor(verbose bool, namespace string) *KubectlExecutor { } } -// CheckAvailability checks if kubectl is installed and accessible. -func (k *KubectlExecutor) CheckAvailability() error { - cmd := exec.Command("kubectl", "version", "--client") - if err := cmd.Run(); err != nil { - return fmt.Errorf("kubectl not found or not working: %w", err) - } - return nil -} - -// Run executes a kubectl command with the given arguments. -// It automatically handles verbose logging and captures stderr for better error messages. -func (k *KubectlExecutor) Run(args ...string) error { - if k.Verbose { - fmt.Printf("Running: kubectl %s\n", strings.Join(args, " ")) - } - - cmd := exec.Command("kubectl", args...) - var stderr bytes.Buffer - cmd.Stdout = os.Stdout - cmd.Stderr = io.MultiWriter(os.Stderr, &stderr) - - if err := cmd.Run(); err != nil { - return fmt.Errorf("`kubectl %s` failed: %w\n%s", strings.Join(args, " "), err, stderr.String()) - } - - return nil -} - // RunWithOutput executes a kubectl command and returns the combined output. // This is useful for capturing command output for further processing. func (k *KubectlExecutor) RunWithOutput(args ...string) ([]byte, error) { @@ -72,88 +39,3 @@ func (k *KubectlExecutor) RunWithOutput(args ...string) ([]byte, error) { return output, nil } - -// Apply applies YAML resources to the Kubernetes cluster. -// Multiple YAML documents can be provided and will be combined with separators. -func (k *KubectlExecutor) Apply(yamls ...[]byte) error { - fmt.Printf("🚀 Applying resources to cluster...\n") - - // Check if kubectl is available - if err := k.CheckAvailability(); err != nil { - return fmt.Errorf("kubectl is required for cluster deployment: %w", err) - } - - // Create temporary file for kubectl apply - tmpFile, err := os.CreateTemp("", "resources-*.yaml") - if err != nil { - return fmt.Errorf("failed to create temp file: %w", err) - } - defer func() { - if err := os.Remove(tmpFile.Name()); err != nil { - fmt.Printf("Warning: failed to remove temp file: %v\n", err) - } - }() - - // Combine all YAML resources with separators - var combinedYAML []byte - for i, yaml := range yamls { - if i > 0 { - combinedYAML = append(combinedYAML, []byte("\n---\n")...) - } - combinedYAML = append(combinedYAML, yaml...) - } - - if _, err := tmpFile.Write(combinedYAML); err != nil { - return fmt.Errorf("failed to write to temp file: %w", err) - } - if err := tmpFile.Close(); err != nil { - return fmt.Errorf("failed to close temp file: %w", err) - } - - // Apply using kubectl - if err := k.Run("apply", "-f", tmpFile.Name()); err != nil { - return err - } - - fmt.Printf("✅ Resources applied successfully\n") - return nil -} - -// WaitForDeployment waits for a Kubernetes deployment to be ready. -// It uses kubectl rollout status to monitor the deployment progress. -func (k *KubectlExecutor) WaitForDeployment(name string, timeout time.Duration) error { - ctx, cancel := context.WithTimeout(context.Background(), timeout) - defer cancel() - - namespace := k.Namespace - if namespace == "" { - namespace = "default" - } - - args := []string{ - "rollout", "status", "deployment", name, - "-n", namespace, - "--timeout", timeout.String(), - } - - cmd := exec.CommandContext(ctx, "kubectl", args...) - if k.Verbose { - fmt.Printf("Running: kubectl %s\n", strings.Join(args, " ")) - } - - var stderr bytes.Buffer - cmd.Stdout = os.Stdout - cmd.Stderr = io.MultiWriter(os.Stderr, &stderr) - - // Sleep briefly to allow controller to create the deployment - time.Sleep(1 * time.Second) - - if err := cmd.Run(); err != nil { - if ctx.Err() == context.DeadlineExceeded { - return fmt.Errorf("timed out waiting for deployment to be ready") - } - return fmt.Errorf("`kubectl rollout status` failed: %w\n%s", err, stderr.String()) - } - - return nil -} diff --git a/go/core/cli/internal/common/fs/project.go b/go/core/cli/internal/common/fs/project.go deleted file mode 100644 index 412b9bc28c..0000000000 --- a/go/core/cli/internal/common/fs/project.go +++ /dev/null @@ -1,10 +0,0 @@ -// Package fs provides filesystem utilities. -package fs - -import "os" - -// FileExists checks if a file exists at the given path. -func FileExists(path string) bool { - _, err := os.Stat(path) - return err == nil -} diff --git a/go/core/cli/internal/common/generator/base.go b/go/core/cli/internal/common/generator/base.go deleted file mode 100644 index f1d05c24b8..0000000000 --- a/go/core/cli/internal/common/generator/base.go +++ /dev/null @@ -1,163 +0,0 @@ -// Package generator provides shared utilities for project code generation. -package generator - -import ( - "fmt" - "io/fs" - "os" - "os/exec" - "path/filepath" - "strings" - "text/template" -) - -// ProjectConfig defines the interface that project configuration must implement. -// This allows both agent and MCP configs to use the same base generator. -type ProjectConfig interface { - // GetDirectory returns the target directory for project generation - GetDirectory() string - - // IsVerbose returns whether verbose logging is enabled - IsVerbose() bool - - // ShouldInitGit returns whether to initialize a git repository - ShouldInitGit() bool - - // ShouldSkipPath returns whether to skip a specific template path - ShouldSkipPath(path string) bool -} - -// BaseGenerator provides common functionality for all project generators. -// It handles template walking, rendering, and file creation with consistent -// error handling and logging across both agent and MCP generators. -type BaseGenerator struct { - // TemplateFiles is the embedded filesystem containing template files - TemplateFiles fs.FS - - // TemplateRoot is the root directory within TemplateFiles (usually "templates") - TemplateRoot string -} - -// NewBaseGenerator creates a new base generator with the specified configuration. -func NewBaseGenerator(templateFiles fs.FS, templateRoot string) *BaseGenerator { - if templateRoot == "" { - templateRoot = "templates" - } - return &BaseGenerator{ - TemplateFiles: templateFiles, - TemplateRoot: templateRoot, - } -} - -// GenerateProject generates a new project using the provided templates. -// It walks through all template files, renders them with the config data, -// and writes the results to the target directory. -func (g *BaseGenerator) GenerateProject(config ProjectConfig) error { - // Get templates subdirectory - templateRoot, err := fs.Sub(g.TemplateFiles, g.TemplateRoot) - if err != nil { - return fmt.Errorf("failed to get templates subdirectory: %w", err) - } - - // Walk through all template files - err = fs.WalkDir(templateRoot, ".", func(path string, d fs.DirEntry, err error) error { - if err != nil { - return err - } - - // Allow config to skip specific paths (e.g., tool templates, mcp_server dirs) - if config.ShouldSkipPath(path) { - if d.IsDir() { - return fs.SkipDir - } - return nil - } - - // Determine destination path by removing .tmpl extension - destPath := filepath.Join(config.GetDirectory(), strings.TrimSuffix(path, ".tmpl")) - - // Handle directories - if d.IsDir() { - return os.MkdirAll(destPath, 0755) - } - - // Read template file - templateContent, err := fs.ReadFile(templateRoot, path) - if err != nil { - return fmt.Errorf("failed to read template file %s: %w", path, err) - } - - // Render template content - renderedContent, err := g.RenderTemplate(string(templateContent), config) - if err != nil { - return fmt.Errorf("failed to render template for %s: %w", path, err) - } - - // Create file - if err := os.WriteFile(destPath, []byte(renderedContent), 0644); err != nil { - return fmt.Errorf("failed to write file %s: %w", destPath, err) - } - - // Log if verbose - if config.IsVerbose() { - fmt.Printf(" Generated: %s\n", destPath) - } - - return nil - }) - - if err != nil { - return fmt.Errorf("failed to walk templates: %w", err) - } - - // Initialize git repository if requested - if config.ShouldInitGit() { - if err := g.initGitRepo(config.GetDirectory(), config.IsVerbose()); err != nil { - // Don't fail the whole operation if git init fails - if config.IsVerbose() { - fmt.Printf("Warning: failed to initialize git repository: %v\n", err) - } - } - } - - return nil -} - -// RenderTemplate renders a template string with the provided data. -// This is the core template rendering logic used by all generators. -func (g *BaseGenerator) RenderTemplate(tmplContent string, data any) (string, error) { - tmpl, err := template.New("template").Parse(tmplContent) - if err != nil { - return "", fmt.Errorf("failed to parse template: %w", err) - } - - var result strings.Builder - if err := tmpl.Execute(&result, data); err != nil { - return "", fmt.Errorf("failed to execute template: %w", err) - } - - return result.String(), nil -} - -// initGitRepo initializes a git repository in the specified directory. -func (g *BaseGenerator) initGitRepo(dir string, verbose bool) error { - cmd := exec.Command("git", "init") - cmd.Dir = dir - - if verbose { - fmt.Printf(" Initializing git repository...\n") - } - - if err := cmd.Run(); err != nil { - return fmt.Errorf("failed to run git init: %w", err) - } - - return nil -} - -// ReadTemplateFile reads a template file from the embedded filesystem. -// This is a convenience method for reading individual template files. -func (g *BaseGenerator) ReadTemplateFile(templatePath string) ([]byte, error) { - fullPath := filepath.Join(g.TemplateRoot, templatePath) - return fs.ReadFile(g.TemplateFiles, fullPath) -} diff --git a/go/core/cli/internal/common/k8s/config.go b/go/core/cli/internal/common/k8s/config.go deleted file mode 100644 index 76911d775f..0000000000 --- a/go/core/cli/internal/common/k8s/config.go +++ /dev/null @@ -1,70 +0,0 @@ -// Package k8s provides utilities for working with Kubernetes configurations. -package k8s - -import ( - "fmt" - "strings" - - "k8s.io/client-go/tools/clientcmd" - - clientset "github.com/kagent-dev/kagent/go/api/clientset/versioned" -) - -// NewKagentClientset builds the generated kagent clientset from the ambient -// kubeconfig. Harness and AgentTemplate are CRDs, so every reader goes through here. -func NewKagentClientset() (clientset.Interface, error) { - restConfig, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( - clientcmd.NewDefaultClientConfigLoadingRules(), - &clientcmd.ConfigOverrides{}, - ).ClientConfig() - if err != nil { - return nil, fmt.Errorf("load Kubernetes client config: %w", err) - } - clients, err := clientset.NewForConfig(restConfig) - if err != nil { - return nil, fmt.Errorf("create Kubernetes client: %w", err) - } - return clients, nil -} - -// GetCurrentNamespace returns the current namespace from the kubeconfig. -// If no namespace is configured, it returns an error. -func GetCurrentNamespace() (string, error) { - config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( - clientcmd.NewDefaultClientConfigLoadingRules(), - &clientcmd.ConfigOverrides{}, - ) - - namespace, _, err := config.Namespace() - if err != nil { - return "", fmt.Errorf("failed to get current namespace: %w", err) - } - - return namespace, nil -} - -// GetCurrentKindClusterName extracts the kind cluster name from the current kubeconfig context. -// Returns an error if the current context is not a kind cluster. -func GetCurrentKindClusterName() (string, error) { - config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( - clientcmd.NewDefaultClientConfigLoadingRules(), - &clientcmd.ConfigOverrides{}, - ) - - rawConfig, err := config.RawConfig() - if err != nil { - return "", fmt.Errorf("failed to get raw kubeconfig: %w", err) - } - - currentContext, ok := rawConfig.Contexts[rawConfig.CurrentContext] - if !ok { - return "", fmt.Errorf("current context %q not found in kubeconfig", rawConfig.CurrentContext) - } - - const kindPrefix = "kind-" - if after, ok0 := strings.CutPrefix(currentContext.Cluster, kindPrefix); ok0 { - return after, nil - } - - return "", fmt.Errorf("current cluster %q is not a kind cluster", currentContext.Cluster) -} diff --git a/go/core/cli/internal/common/prompt/input.go b/go/core/cli/internal/common/prompt/input.go deleted file mode 100644 index d4e31ae630..0000000000 --- a/go/core/cli/internal/common/prompt/input.go +++ /dev/null @@ -1,21 +0,0 @@ -// Package prompt provides utilities for interactive user input. -package prompt - -import ( - "bufio" - "fmt" - "os" - "strings" -) - -// PromptForInput displays a prompt and reads user input from stdin. -// Returns the trimmed input string or an error if reading fails. -func PromptForInput(promptText string) (string, error) { - fmt.Print(promptText) - reader := bufio.NewReader(os.Stdin) - input, err := reader.ReadString('\n') - if err != nil { - return "", fmt.Errorf("failed to read input: %w", err) - } - return strings.TrimSpace(input), nil -} diff --git a/go/core/internal/controller/reconciler.go b/go/core/internal/controller/reconciler.go index 0a958b9875..82b5e19a98 100644 --- a/go/core/internal/controller/reconciler.go +++ b/go/core/internal/controller/reconciler.go @@ -540,26 +540,24 @@ func (r *Reconciler) reconcileModelConfigStatus(ctx context.Context, key string) } func statusWithTransitionTimes(desired, current kagentv1alpha3.AgentStatus) kagentv1alpha3.AgentStatus { - desired.Conditions = append([]metav1.Condition(nil), desired.Conditions...) - for i := range desired.Conditions { - condition := &desired.Conditions[i] - previous := apimeta.FindStatusCondition(current.Conditions, condition.Type) - if previous != nil && previous.Status == condition.Status { - condition.LastTransitionTime = previous.LastTransitionTime - } else { - condition.LastTransitionTime = metav1.Now() - } - } + desired.Conditions = conditionsWithTransitionTimes(desired.Conditions, current.Conditions) return desired } func modelConfigStatusWithTransitionTimes(desired, current kagentv1alpha3.ModelConfigStatus) kagentv1alpha3.ModelConfigStatus { - desired.Conditions = append([]metav1.Condition(nil), desired.Conditions...) - for conditionIndex := range desired.Conditions { - condition := &desired.Conditions[conditionIndex] - if previous := apimeta.FindStatusCondition(current.Conditions, condition.Type); previous != nil && - previous.Status == condition.Status && previous.Reason == condition.Reason && - previous.Message == condition.Message && previous.ObservedGeneration == condition.ObservedGeneration { + desired.Conditions = conditionsWithTransitionTimes(desired.Conditions, current.Conditions) + return desired +} + +// conditionsWithTransitionTimes keeps the transition time a condition's +// status holds and stamps the current time when the status changes, as +// apimeta.SetStatusCondition does; a new reason or message alone is no +// transition. +func conditionsWithTransitionTimes(desired, current []metav1.Condition) []metav1.Condition { + desired = append([]metav1.Condition(nil), desired...) + for conditionIndex := range desired { + condition := &desired[conditionIndex] + if previous := apimeta.FindStatusCondition(current, condition.Type); previous != nil && previous.Status == condition.Status { condition.LastTransitionTime = previous.LastTransitionTime continue } diff --git a/go/core/internal/controller/reconciler_test.go b/go/core/internal/controller/reconciler_test.go index 669e3f44c3..5b135b01af 100644 --- a/go/core/internal/controller/reconciler_test.go +++ b/go/core/internal/controller/reconciler_test.go @@ -613,3 +613,21 @@ func reconciliationTestCard(name string) *a2apb.AgentCard { SupportedInterfaces: []*a2apb.AgentInterface{{Url: "http://127.0.0.1:80", ProtocolBinding: "GRPC", ProtocolVersion: "1.0"}}, } } + +func TestModelConfigConditionTransitionTimesMoveOnlyOnStatusChange(t *testing.T) { + then := metav1.NewTime(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)) + current := kagentv1alpha3.ModelConfigStatus{Conditions: []metav1.Condition{ + {Type: "Accepted", Status: metav1.ConditionFalse, Reason: "SecretNotFound", Message: "secret team-a/a not found", ObservedGeneration: 1, LastTransitionTime: then}, + {Type: "Ready", Status: metav1.ConditionTrue, Reason: "Resolved", ObservedGeneration: 1, LastTransitionTime: then}, + }} + desired := kagentv1alpha3.ModelConfigStatus{Conditions: []metav1.Condition{ + {Type: "Accepted", Status: metav1.ConditionFalse, Reason: "SecretKeyNotFound", Message: "key missing in secret team-a/b", ObservedGeneration: 2}, + {Type: "Ready", Status: metav1.ConditionFalse, Reason: "Unresolved", ObservedGeneration: 2}, + }} + + got := modelConfigStatusWithTransitionTimes(desired, current) + + require.Equal(t, then, got.Conditions[0].LastTransitionTime, "a new reason, message or generation alone is no transition") + require.True(t, got.Conditions[1].LastTransitionTime.After(then.Time), "a status change stamps a new transition time") + require.True(t, desired.Conditions[0].LastTransitionTime.IsZero(), "the desired conditions are not mutated") +} diff --git a/go/core/internal/database/README.md b/go/core/internal/database/README.md index a9fa246c34..abc4b18cf8 100644 --- a/go/core/internal/database/README.md +++ b/go/core/internal/database/README.md @@ -55,9 +55,10 @@ The repository [style guide](../../../../STYLE.md) also applies. - Do not use `uuid.MustParse` on runtime inputs. Pass validated values as SQL parameters; use checked parsing when Go needs a UUID for comparison. Store calls must return errors instead of panicking on malformed IDs. -- [Migration 1](../../pkg/migrations/core/000001_initial.sql) is currently the - unreleased baseline: fold approved schema changes into it and recreate local - databases. After release, use new migrations. Keep both Up and Down valid. +- Merged migrations are immutable, including + [migration 1](../../pkg/migrations/core/000001_initial.sql); the required + `migration-immutability` check rejects any change to them. Add each schema + change as a new migration. Keep both Up and Down valid. ## Atomicity and performance diff --git a/go/core/internal/database/runtime_instances.go b/go/core/internal/database/runtime_instances.go index f57ded3066..8afd19d22e 100644 --- a/go/core/internal/database/runtime_instances.go +++ b/go/core/internal/database/runtime_instances.go @@ -68,7 +68,11 @@ func runtimeOperationNeeded(row runtimeInstanceRow, requested apiv1.RuntimeOpera if err != nil { return false, err } - if requested < apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE || requested > apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE { + switch requested { + case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, + apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, + apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT: + default: return false, fmt.Errorf("invalid runtime operation %s: %w", requested, ErrFailedPrecondition) } if state == apiv1.RuntimeState_RUNTIME_STATE_DELETED { @@ -81,11 +85,20 @@ func runtimeOperationNeeded(row runtimeInstanceRow, requested apiv1.RuntimeOpera if operation == requested { return false, nil } - if requested != apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE || row.ExecutorID != nil { + // A revert restores the recorded Actor in place and creates nothing, so + // Delete may supersede one even after it was issued: Substrate serializes + // the two calls, and a late revert fails against the deleted Actor. + issued := row.ExecutorID != nil && operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT + if requested != apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE || issued { return false, ErrConflict } operation = apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE } + if requested == apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT { + // Only BeginCrashedSessionRecovery admits a revert, after Substrate confirms + // the crash. This path can join one but never start one. + return false, ErrConflict + } var expected, target apiv1.RuntimeState expectedOperation := apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE switch requested { @@ -178,7 +191,8 @@ func completeRuntimeOperation(row runtimeInstanceRow, operationID, executorID uu return row, ErrConflict } switch operation { - case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME: + case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, + apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT: row.State = apiv1.RuntimeState_RUNTIME_STATE_READY.String() case apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND: row.State = apiv1.RuntimeState_RUNTIME_STATE_SUSPENDED.String() diff --git a/go/core/internal/database/session_deletion.go b/go/core/internal/database/session_deletion.go index 2658c70d0f..7b7b22dd46 100644 --- a/go/core/internal/database/session_deletion.go +++ b/go/core/internal/database/session_deletion.go @@ -82,26 +82,9 @@ func admitForcedSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow) // deletion or actor absence. The caller holds the session lock and commits // these writes with the session's DELETED state. func finishDeletedSessionTasks(ctx context.Context, tx pgx.Tx, row sessionRow) error { - type boundary struct { - TaskID string - Sequence int64 - } - // A saved result may be waiting for a lost runtime cleanup acknowledgement. - // Publish it first so deletion preserves the result instead of canceling it. - boundaries, err := queryMany(ctx, tx, ` - SELECT task_id, sequence FROM session_task_event - WHERE history_id = $1 AND NOT published - AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL - ORDER BY sequence - `, pgx.RowToStructByName[boundary], row.HistoryID) - if err != nil { + if err := publishSavedSessionResults(ctx, tx, row); err != nil { return err } - for _, saved := range boundaries { - if err := settleSessionTask(ctx, tx, row, saved.TaskID, saved.Sequence); err != nil { - return err - } - } // Cancel only tasks still recorded as active after publishing saved results. // Waiting and terminal tasks keep their historical state. tasks, err := queryMany(ctx, tx, ` @@ -125,8 +108,7 @@ func finishDeletedSessionTasks(ctx context.Context, tx pgx.Tx, row sessionRow) e return err } } - // Retire idle-work markers; deletion no longer needs a later suspension. - return execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, row.HistoryID) + return retireSessionQuiescence(ctx, tx, row) } // finishSessionDeletion removes idle-expired sessions and their creation receipts. diff --git a/go/core/internal/database/session_operations.go b/go/core/internal/database/session_operations.go index ca53013880..b629034aee 100644 --- a/go/core/internal/database/session_operations.go +++ b/go/core/internal/database/session_operations.go @@ -26,7 +26,8 @@ type SessionOperation struct { // turns, pending dispatch, and unfinished runtime cleanup, while checkpoint and // lifecycle guards remain. Accepted deletion stays retryable without force. // Creation retries and already-at-target requests return current state. Delete -// may supersede unclaimed work; uncertain lifecycle work blocks it. +// may supersede unclaimed work and any crash recovery; other uncertain lifecycle +// work blocks it. // Callers authorize access. Deleted sessions return ErrNotFound, except Delete // can observe the tombstone for a caller already authorized before deletion. func (c *Client) BeginSessionOperation(ctx context.Context, sessionID string, kind apiv1.RuntimeOperation, force bool) (*SessionOperation, error) { diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go new file mode 100644 index 0000000000..d591f33b5b --- /dev/null +++ b/go/core/internal/database/session_recovery.go @@ -0,0 +1,185 @@ +package database + +import ( + "context" + "fmt" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "google.golang.org/protobuf/types/known/timestamppb" + + apia2a "github.com/kagent-dev/kagent/go/api/a2a" + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// SessionRuntimeCrashed is the failure reason set once Substrate confirms the +// recorded Actor crashed. Transport failures never set it. +const SessionRuntimeCrashed = "RuntimeCrashed" + +// CrashRecoveryAdmitted reports an admitted, unfinished crash recovery. Only +// BeginCrashedSessionRecovery starts a REVERT, so the operation alone identifies it. +func CrashRecoveryAdmitted(session *apiv1.Session) bool { + return session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT +} + +// SessionGeneration is read before crash recovery observes Substrate. OperationID +// is uuid.Nil when none is recorded; every operation replaces it and completion +// keeps it. ActorUID is "" when none is recorded. +type SessionGeneration struct { + SessionID string + OperationID uuid.UUID + ActorUID string + TemplateAtespace string + TemplateName string +} + +// GetSessionGeneration returns ErrNotFound for deleted Sessions. Callers authorize +// access. The read takes no lock; admission re-checks it under the Session lock. +func (c *Client) GetSessionGeneration(ctx context.Context, id string) (*SessionGeneration, error) { + type generationRow struct { + OperationID *uuid.UUID + ActorUID *string + ActorTemplateAtespace string + ActorTemplateName string + } + row, err := queryOne(ctx, c.db, ` + SELECT i.operation_id, i.actor_uid, r.actor_template_atespace, r.actor_template_name + FROM session_record i JOIN runtime_revision r ON r.revision = i.prepared_revision + WHERE i.id = $1 AND i.state <> 'RUNTIME_STATE_DELETED' + `, pgx.RowToStructByName[generationRow], id) + if err != nil { + return nil, fmt.Errorf("read Session %s generation: %w", id, notFoundOr(err)) + } + generation := &SessionGeneration{SessionID: id, TemplateAtespace: row.ActorTemplateAtespace, TemplateName: row.ActorTemplateName} + if row.OperationID != nil { + generation.OperationID = *row.OperationID + } + if row.ActorUID != nil { + generation.ActorUID = *row.ActorUID + } + return generation, nil +} + +// BeginCrashedSessionRecovery admits recovery of a confirmed crash and preserves +// durable results. generation must be read before the crash was observed: a +// concurrent revert keeps the Actor UID and returns to READY, so only a newer +// lifecycle operation shows that the observation is stale. +func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *SessionGeneration, snapshot SessionTaskSnapshot) (*SessionOperation, error) { + var operation *SessionOperation + err := c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockSession(ctx, tx, generation.SessionID) + if err != nil { + return notFoundOr(err) + } + session, err := toSession(row) + if err != nil { + return err + } + if CrashRecoveryAdmitted(session) { + operation, err = toSessionOperation(row) + return err + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return ErrConflict + } + current := uuid.Nil + if row.OperationID != nil { + current = *row.OperationID + } + if current != generation.OperationID { + return ErrConflict + } + // An issued idle suspension does not block recovery. It cannot finish against + // the crashed Actor and its worker keeps the claim when suspension fails, so + // waiting would block forever. Settlement below retires the claim; a late + // finish changes nothing. + busy, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + `, pgx.RowTo[bool], row.ID) + if err != nil { + return err + } + if busy { + return ErrConflict + } + row.State = apiv1.RuntimeState_RUNTIME_STATE_FAILED.String() + row.runtimeInstanceRow, err = beginRuntimeOperation(ctx, tx, row.runtimeInstanceRow, runtimeKindAgent, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT) + if err != nil { + return err + } + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT + session.Failure = &apiv1.Failure{Reason: SessionRuntimeCrashed, Message: "Runtime crashed; recovery to its last snapshot is pending"} + 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 + } + if err := settleCrashedSessionTasks(ctx, tx, row, snapshot); err != nil { + return err + } + operation, err = toSessionOperation(row) + return err + }) + return operation, err +} + +// 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 { + if err := publishSavedSessionResults(ctx, tx, session); err != nil { + return err + } + type interruptedTask struct { + ID string + Data []byte + Version *int64 + } + // A wait is captured when an event recording the snapshot follows the task's + // latest event. + tasks, err := queryMany(ctx, tx, ` + SELECT t.id, t.data, latest.sequence AS version + FROM session_task t + LEFT JOIN LATERAL (SELECT MAX(sequence) AS sequence FROM session_task_event e + WHERE e.history_id = t.history_id AND e.task_id = t.id) latest ON TRUE + WHERE t.history_id = $1 + AND (t.state IN ('TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING') + OR (t.state IN ('TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED') + AND NOT EXISTS (SELECT 1 FROM session_task_event captured + WHERE captured.history_id = t.history_id AND $2::text <> '' + AND captured.snapshot_uri = $2 AND captured.snapshot_fidelity = $3 + AND captured.snapshot_atespace = $4 + AND captured.sequence >= latest.sequence))) + `, pgx.RowToStructByName[interruptedTask], session.HistoryID, snapshot.URI, snapshot.Fidelity, snapshot.Atespace) + if err != nil { + return fmt.Errorf("read interrupted Session tasks: %w", err) + } + for _, interrupted := range tasks { + if interrupted.Version == nil { + return fmt.Errorf("interrupted Session task %s has no events", interrupted.ID) + } + task, err := unmarshalSessionTask(interrupted.Data) + if err != nil { + 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."))} + 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 { + return err + } + version, err := taskVersion(ctx, tx, session.HistoryID, interrupted.ID) + if err != nil { + return err + } + if err := publishPushBoundary(ctx, tx, session.HistoryID, interrupted.ID, version, task); err != nil { + return err + } + } + return retireSessionQuiescence(ctx, tx, session) +} diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go new file mode 100644 index 0000000000..0412f02c07 --- /dev/null +++ b/go/core/internal/database/session_recovery_test.go @@ -0,0 +1,329 @@ +package database + +import ( + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + apia2a "github.com/kagent-dev/kagent/go/api/a2a" + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +func TestCrashedRecoverySettlesTasksAtomically(t *testing.T) { + for _, test := range []struct { + name string + state, want a2a.TaskState + published, captured bool + }{ + {name: "submitted", state: a2a.TaskStateSubmitted, want: a2a.TaskStateFailed}, + {name: "working", state: a2a.TaskStateWorking, want: a2a.TaskStateFailed}, + {name: "completed before publication", state: a2a.TaskStateCompleted, want: a2a.TaskStateCompleted}, + {name: "unpublished input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed}, + {name: "unpublished auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed}, + {name: "published input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed, published: true}, + {name: "published auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed, published: true}, + {name: "captured input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateInputRequired, published: true, captured: true}, + {name: "captured auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateAuthRequired, published: true, captured: 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, _, err := client.CreateSession(t.Context(), newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + session, err = markSessionReady(t.Context(), client, session.Id, "runtime.example") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), "original input") + task.ContextID, task.Status.State = session.ContextId, test.state + version, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("create"), task, "") + require.NoError(t, err) + if test.published { + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + snapshot := SessionTaskSnapshot{} + if test.captured { + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + snapshot = SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/captured", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &snapshot)) + } + before, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + operation, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, "actor-"+session.Id), snapshot) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, operation.Instance.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, operation.Instance.Operation) + require.Equal(t, SessionRuntimeCrashed, operation.Instance.Failure.Reason) + stored, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, test.want, stored.Status.State) + require.Equal(t, before.History, stored.History) + require.Equal(t, before.Artifacts, stored.Artifacts) + if test.want == test.state { + require.Equal(t, before.Status, stored.Status) + } + _, settledVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + replayed, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, "actor-"+session.Id), snapshot) + require.NoError(t, err) + require.Equal(t, operation.ID, replayed.ID) + _, replayedVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, settledVersion, replayedVersion) + _, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash("late"), task, task, "") + if test.state.Terminal() { + require.ErrorIs(t, err, ErrFailedPrecondition) + } else { + require.ErrorIs(t, err, ErrConflict) + } + require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) + _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, ErrConflict) + }) + } +} + +func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { + for _, test := range []struct { + name string + checkpoint, superseded, lifecycle bool + }{ + {name: "checkpoint", checkpoint: true}, + {name: "revert after the crash was observed", superseded: true}, + {name: "lifecycle", lifecycle: 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 + // Callers read the generation before they observe the crash. + generation := sessionGeneration(t, client, session.Id) + switch { + case test.checkpoint: + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES", + })) + _, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.NoError(t, err) + case test.superseded: + // Another revert finishes on the same Actor UID, and a new turn + // starts once ingress wakes the restored Actor. + completeCrashedRecovery(t, client, session.Id, uid) + task = newSessionTask(uuid.NewString(), "next input") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateWorking + _, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), task, "") + require.NoError(t, err) + case test.lifecycle: + _, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, false) + require.NoError(t, err) + } + before, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + generation.ActorUID = uid + _, err = client.BeginCrashedSessionRecovery(t.Context(), generation, SessionTaskSnapshot{}) + require.ErrorIs(t, err, ErrConflict) + after, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.True(t, proto.Equal(before, after)) + if test.superseded { + running, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "a stale observation must not fail the new turn") + } + }) + } +} + +// A crash during post-turn suspension leaves the worker's claim issued, because +// the worker cannot tell whether its suspension took effect. The confirmed crash +// settles that: admission retires the claim instead of waiting for it. +func TestCrashedRecoveryRetiresIssuedIdleSuspension(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + uid := "actor-" + session.Id + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, uid), SessionTaskSnapshot{}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, operation.Instance.Operation) + pending, err := queryOne(t.Context(), client.db, `SELECT quiescence_pending FROM session_task_event WHERE sequence = $1`, pgx.RowTo[bool], work.Version) + require.NoError(t, err) + require.False(t, pending, "admission retires the claim in its own transaction") + finishCrashedRecovery(t, client, session.Id, operation, uid) + + late := SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &late)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + recorded, err := queryOne(t.Context(), client.db, ` + SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND snapshot_uri = $2) + OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND snapshot_uri = $2) + `, pgx.RowTo[bool], row.HistoryID, late.URI) + require.NoError(t, err) + require.False(t, recorded, "a late finish of the retired claim records nothing") + _, err = client.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, ErrNotFound) + next := newSessionTask(uuid.NewString(), "next input") + next.ContextID, next.Status.State = session.ContextId, a2a.TaskStateWorking + _, err = client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), next, "") + require.NoError(t, err, "the retired claim no longer blocks a new turn") +} + +func sessionGeneration(t *testing.T, client *Client, sessionID string) *SessionGeneration { + t.Helper() + generation, err := client.GetSessionGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +func observedGeneration(t *testing.T, client *Client, sessionID, actorUID string) *SessionGeneration { + t.Helper() + generation := sessionGeneration(t, client, sessionID) + generation.ActorUID = actorUID + return generation +} + +func completeCrashedRecovery(t *testing.T, client *Client, sessionID, actorUID string) { + t.Helper() + operation, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, sessionID, actorUID), SessionTaskSnapshot{}) + require.NoError(t, err) + finishCrashedRecovery(t, client, sessionID, operation, actorUID) +} + +func finishCrashedRecovery(t *testing.T, client *Client, sessionID string, operation *SessionOperation, actorUID string) { + t.Helper() + executor := uuid.New() + claimed, err := client.ClaimSessionOperation(t.Context(), sessionID, operation.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + recovered, err := client.FinishSessionOperation(t.Context(), sessionID, operation.ID, executor, "", actorUID, "") + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) +} + +func TestCrashedRecoveryAdmissionRollsBackOnMalformedTask(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + before, err := readSessionTask(t.Context(), client.db, row.HistoryID, string(task.ID)) + require.NoError(t, err) + // A corrupt projection makes settlement fail after the lifecycle write. + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, []byte{0xff}, string(task.ID)) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, "actor-"+session.Id), SessionTaskSnapshot{}) + require.Error(t, err) + current, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + require.Nil(t, current.Failure) + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, before.Data, string(task.ID)) + require.NoError(t, err) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.Equal(t, string(task.ID), work.TaskID) +} + +// The crash note carries a timeline position. Without one, chat orders it by +// an older key and shows it above the partial reply it ends. +func TestCrashedRecoveryNotePinsItsTimelinePosition(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + task.Status = a2a.TaskStatus{State: a2a.TaskStateWorking} + _, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash("working"), task, task, "") + require.NoError(t, err) + completeCrashedRecovery(t, client, session.Id, "actor-"+session.Id) + + failed, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, failed.Status.State) + require.NotNil(t, failed.Status.Timestamp) + position, ok := apia2a.TimelinePosition(failed.Status.Message) + require.True(t, ok, "the crash note has a timeline position") + require.True(t, position.Equal(*failed.Status.Timestamp), "position %s, status timestamp %s", position, failed.Status.Timestamp) +} + +func TestSessionGenerationReadsRecordedActor(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _ := waitingTaskFixture(t, client) + + generation, err := client.GetSessionGeneration(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, session.Id, generation.SessionID) + require.Equal(t, "actor-"+session.Id, generation.ActorUID) + require.Equal(t, "team-a", generation.TemplateAtespace) + require.Equal(t, "revision-actor-template", generation.TemplateName) + require.NotEqual(t, uuid.Nil, generation.OperationID, "creation recorded a lifecycle operation") + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), generation, SessionTaskSnapshot{}) + require.NoError(t, err) + admitted, err := client.GetSessionGeneration(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, operation.ID, admitted.OperationID) + require.Equal(t, generation.ActorUID, admitted.ActorUID) + + _, err = client.GetSessionGeneration(t.Context(), uuid.NewString()) + require.ErrorIs(t, err, ErrNotFound) +} + +func TestDeleteSupersedesClaimedRecovery(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _ := waitingTaskFixture(t, client) + uid := "actor-" + session.Id + 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) + + deletion, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err, "a revert creates nothing, so even an issued one cannot block Delete") + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, deletion.Instance.Operation) + require.Equal(t, uuid.Nil, deletion.ExecutorID) + _, err = client.FinishSessionOperation(t.Context(), session.Id, operation.ID, executor, "", uid, "") + require.ErrorIs(t, err, ErrConflict, "the superseded revert cannot finish") +} + +func TestRevertCanOnlyJoinAnAdmittedRecovery(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _ := waitingTaskFixture(t, client) + _, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, false) + require.ErrorIs(t, err, ErrConflict, "only a confirmed crash admits a revert") + admitted, err := client.BeginCrashedSessionRecovery(t.Context(), observedGeneration(t, client, session.Id, "actor-"+session.Id), SessionTaskSnapshot{}) + require.NoError(t, err) + joined, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, false) + require.NoError(t, err) + require.Equal(t, admitted.ID, joined.ID) +} diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index 8ffbd36457..2f80fb7c79 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -27,8 +27,39 @@ func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string }) } +// publishSavedSessionResults publishes every result the runtime saved whose +// cleanup acknowledgement was lost, so ending the runtime preserves them. +func publishSavedSessionResults(ctx context.Context, tx pgx.Tx, session sessionRow) error { + type boundary struct { + TaskID string + Sequence int64 + } + boundaries, err := queryMany(ctx, tx, ` + SELECT task_id, sequence FROM session_task_event + WHERE history_id = $1 AND NOT published + AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL + ORDER BY sequence + `, pgx.RowToStructByName[boundary], session.HistoryID) + if err != nil { + return err + } + for _, saved := range boundaries { + if err := settleSessionTask(ctx, tx, session, saved.TaskID, saved.Sequence); err != nil { + return err + } + } + return nil +} + +// retireSessionQuiescence drops idle-work markers for a runtime that no longer +// needs a later suspension. +func retireSessionQuiescence(ctx context.Context, tx pgx.Tx, session sessionRow) error { + return execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, session.HistoryID) +} + // settleSessionTask publishes one saved boundary under the session lock. Runtime -// acknowledgement or confirmed actor deletion must establish cleanup first. +// acknowledgement, confirmed actor deletion, or a confirmed crash must establish +// cleanup first. func settleSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, taskID string, version int64) error { type boundary struct { Data []byte diff --git a/go/core/internal/grpcserver/policy.go b/go/core/internal/grpcserver/policy.go index 74b9a66dd8..e6a249d6f5 100644 --- a/go/core/internal/grpcserver/policy.go +++ b/go/core/internal/grpcserver/policy.go @@ -81,6 +81,7 @@ func DefaultMethodPolicies() MethodPolicies { // A rename is the only write on this service that is not a lifecycle // operation, and it must not inherit the read mode its neighbours carry. policies[apiv1.SessionService_UpdateSessionName_FullMethodName] = auth.AccessUpdate + policies[apiv1.SessionService_RevertSession_FullMethodName] = auth.AccessUpdate policies[apiv1.SessionService_DeleteSession_FullMethodName] = auth.AccessDelete policies[apiv1.SessionService_CreateSessionShare_FullMethodName] = auth.AccessCreate policies[apiv1.SessionService_ListSessionShares_FullMethodName] = auth.AccessRead diff --git a/go/core/internal/grpcserver/policy_test.go b/go/core/internal/grpcserver/policy_test.go index 712e2e243a..655c8e926c 100644 --- a/go/core/internal/grpcserver/policy_test.go +++ b/go/core/internal/grpcserver/policy_test.go @@ -38,6 +38,7 @@ func TestSessionServicePoliciesMatchTheirEffect(t *testing.T) { {name: "get", method: apiv1.SessionService_GetSession_FullMethodName, want: pkgauth.AccessRead}, {name: "list", method: apiv1.SessionService_ListSessions_FullMethodName, want: pkgauth.AccessRead}, {name: "rename is a write", method: apiv1.SessionService_UpdateSessionName_FullMethodName, want: pkgauth.AccessUpdate}, + {name: "revert", method: apiv1.SessionService_RevertSession_FullMethodName, want: pkgauth.AccessUpdate}, {name: "delete", method: apiv1.SessionService_DeleteSession_FullMethodName, want: pkgauth.AccessDelete}, {name: "create share", method: apiv1.SessionService_CreateSessionShare_FullMethodName, want: pkgauth.AccessCreate}, {name: "list shares", method: apiv1.SessionService_ListSessionShares_FullMethodName, want: pkgauth.AccessRead}, diff --git a/go/core/internal/grpcserver/protovalidate_test.go b/go/core/internal/grpcserver/protovalidate_test.go index 740ce014fe..64e1877b20 100644 --- a/go/core/internal/grpcserver/protovalidate_test.go +++ b/go/core/internal/grpcserver/protovalidate_test.go @@ -88,6 +88,7 @@ func TestInvalidSessionAndCheckpointIDsNeverReachHandlers(t *testing.T) { requests := []proto.Message{ &apiv1.GetSessionRequest{SessionId: "invalid"}, &apiv1.UpdateSessionNameRequest{SessionId: "invalid"}, + &apiv1.RevertSessionRequest{SessionId: "invalid"}, &apiv1.DeleteSessionRequest{SessionId: "invalid"}, &apiv1.CreateSessionShareRequest{SessionId: "invalid"}, &apiv1.ListSessionSharesRequest{SessionId: "invalid"}, diff --git a/go/core/internal/grpcserver/session.go b/go/core/internal/grpcserver/session.go index e84263bbf9..ed64db9edc 100644 --- a/go/core/internal/grpcserver/session.go +++ b/go/core/internal/grpcserver/session.go @@ -51,6 +51,14 @@ func (s *sessionServer) UpdateSessionName(ctx context.Context, request *apiv1.Up return &apiv1.UpdateSessionNameResponse{Session: session}, nil } +func (s *sessionServer) RevertSession(ctx context.Context, request *apiv1.RevertSessionRequest) (*apiv1.RevertSessionResponse, error) { + session, err := s.service.Revert(ctx, request.GetSessionId()) + if err != nil { + return nil, err + } + return &apiv1.RevertSessionResponse{Session: session}, nil +} + func (s *sessionServer) DeleteSession(ctx context.Context, request *apiv1.DeleteSessionRequest) (*apiv1.DeleteSessionResponse, error) { session, err := s.service.Delete(ctx, request.GetSessionId(), request.GetForce()) if err != nil { diff --git a/go/core/internal/grpcserver/session_recovery_test.go b/go/core/internal/grpcserver/session_recovery_test.go new file mode 100644 index 0000000000..5971703f9b --- /dev/null +++ b/go/core/internal/grpcserver/session_recovery_test.go @@ -0,0 +1,205 @@ +package grpcserver + +import ( + "context" + "crypto/sha256" + "net" + "sync" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/proto" + + 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/dbtest" + authimpl "github.com/kagent-dev/kagent/go/core/internal/httpserver/auth" + sessionsvc "github.com/kagent-dev/kagent/go/core/internal/service/session" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +type revertTestControl struct { + ateapipb.ControlClient + mu sync.Mutex + actors map[string]*ateapipb.Actor + revert func(*ateapipb.Actor) error +} + +func (c *revertTestControl) GetActor(_ context.Context, request *ateapipb.GetActorRequest, _ ...grpc.CallOption) (*ateapipb.Actor, error) { + c.mu.Lock() + defer c.mu.Unlock() + actor := c.actors[request.GetActor().GetName()] + if actor == nil { + return nil, status.Error(codes.NotFound, "actor not found") + } + return proto.CloneOf(actor), nil +} + +func (c *revertTestControl) RevertActor(_ context.Context, request *ateapipb.RevertActorRequest, _ ...grpc.CallOption) (*ateapipb.RevertActorResponse, error) { + c.mu.Lock() + defer c.mu.Unlock() + actor := c.actors[request.GetActor().GetName()] + if c.revert != nil { + if err := c.revert(actor); err != nil { + return nil, err + } + } + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + return &ateapipb.RevertActorResponse{Actor: proto.CloneOf(actor)}, nil +} + +func TestRevertSessionThroughGRPC(t *testing.T) { + if testing.Short() { + t.Skip("requires PostgreSQL") + } + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + dbtest.MigrateT(t, dsn, false) + db, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn)}) + require.NoError(t, err) + t.Cleanup(db.Close) + store := database.NewClient(db, "public") + control := &revertTestControl{actors: map[string]*ateapipb.Actor{}} + actors := &substrate.Client{ControlClient: control} + workflow := sessionsvc.NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + listener := bufconn.Listen(DefaultMaxMessageSize) + server, err := New(Config{ + Listener: listener, Authenticator: &authimpl.InsecureAuthenticator{}, + SystemService: testSystemService(), ShareStore: store, + SessionService: sessionsvc.NewService(store, &auth.NoopAuthorizer{}, workflow), + }) + require.NoError(t, err) + ctx, stop := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- server.Start(ctx) }() + t.Cleanup(func() { stop(); require.NoError(t, <-done) }) + connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() })) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + client := apiv1.NewSessionServiceClient(connection) + owner := metadata.NewOutgoingContext(t.Context(), metadata.Pairs("x-user-id", "alice")) + visitor := metadata.NewOutgoingContext(t.Context(), metadata.Pairs("x-user-id", "bob")) + _, err = client.RevertSession(owner, &apiv1.RevertSessionRequest{SessionId: "invalid"}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + for _, test := range []struct { + name string + task a2a.TaskState + noSnapshot bool + lostResponse bool + }{ + {name: "interrupted turn", task: a2a.TaskStateWorking}, + {name: "completed turn", task: a2a.TaskStateCompleted}, + {name: "idle session"}, + {name: "published input wait", task: a2a.TaskStateInputRequired}, + {name: "published auth wait", task: a2a.TaskStateAuthRequired}, + {name: "first-turn crash", task: a2a.TaskStateWorking, noSnapshot: true}, + {name: "revert response lost", task: a2a.TaskStateWorking, lostResponse: true}, + } { + t.Run(test.name, func(t *testing.T) { + session := createTaskStoreSession(t, store) + name := substrate.ActorName(session.Id) + control.actors[name] = &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_RUNNING, + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/last-good"}}}}}, + } + var task *a2a.Task + if test.task != "" { + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task = a2a.NewSubmittedTask(message, message) + task.Status.State = test.task + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + if test.task == a2a.TaskStateInputRequired || test.task == a2a.TaskStateAuthRequired { + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + } + request := &apiv1.RevertSessionRequest{SessionId: session.Id} + if test.noSnapshot { + control.actors[name].Status.Snapshots = nil + } + share, err := client.CreateSessionShare(owner, &apiv1.CreateSessionShareRequest{ + SessionId: session.Id, + }) + require.NoError(t, err) + readOnly := metadata.AppendToOutgoingContext(visitor, "x-share-token", share.Token) + _, err = client.RevertSession(readOnly, request) + require.Equal(t, codes.PermissionDenied, status.Code(err)) + _, err = client.RevertSession(visitor, request) + require.Equal(t, codes.NotFound, status.Code(err)) + healthy, err := client.RevertSession(owner, request) + require.NoError(t, err, "a healthy runtime has nothing to recover") + 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 { + calls++ + if failOnce { + failOnce = false + if test.lostResponse { + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + } + return status.Error(codes.Unavailable, "temporary revert outage") + } + return nil + } + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.Unavailable, status.Code(err)) + failed, err := client.GetSession(owner, &apiv1.GetSessionRequest{SessionId: session.Id}) + require.NoError(t, err) + require.Equal(t, database.SessionRuntimeCrashed, failed.Session.Failure.Reason) + ready, err := client.RevertSession(owner, request) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.Session.State) + require.Nil(t, ready.Session.Failure) + // Retry as a client whose success response was lost. + retried, err := client.RevertSession(owner, request) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, retried.Session.State) + if test.lostResponse { + require.Equal(t, 1, calls, "a lost response cannot repeat the revert") + } else { + require.Equal(t, 2, calls) + } + after := control.actors[name] + require.Equal(t, before.Metadata.Uid, after.Metadata.Uid) + require.True(t, proto.Equal(substrate.LatestDurableActorSnapshot(before), substrate.LatestDurableActorSnapshot(after))) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, after.Status.State) + if task != nil { + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + want := test.task + if want == a2a.TaskStateWorking || want == a2a.TaskStateInputRequired || want == a2a.TaskStateAuthRequired { + want = a2a.TaskStateFailed + } + require.Equal(t, want, stored.Status.State) + } + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) + }) + } +} diff --git a/go/core/internal/mcp/tasks_test.go b/go/core/internal/mcp/tasks_test.go index 5e3f1549e4..ad562ee138 100644 --- a/go/core/internal/mcp/tasks_test.go +++ b/go/core/internal/mcp/tasks_test.go @@ -568,6 +568,10 @@ func (*fakeSessionWorkflow) Create(_ context.Context, session *apiv1.Session) (* return session, nil } +func (*fakeSessionWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, nil +} + func (*fakeSessionWorkflow) Delete(_ context.Context, session *apiv1.Session, _ bool) (*apiv1.Session, error) { return session, nil } diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go new file mode 100644 index 0000000000..7d348775bf --- /dev/null +++ b/go/core/internal/service/session/recovery.go @@ -0,0 +1,102 @@ +package session + +import ( + "context" + "errors" + "fmt" + + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +// errRuntimeUnrecoverable tells the client to delete the Session. Transient +// Substrate and database failures must not wrap it, so they stay retryable. +var errRuntimeUnrecoverable = fmt.Errorf("runtime cannot be recovered: %w", database.ErrFailedPrecondition) + +func lifecycleIdle(session *apiv1.Session) bool { + return session.State == apiv1.RuntimeState_RUNTIME_STATE_READY && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE +} + +type crashObservation struct { + generation *database.SessionGeneration + actor *ateapipb.Actor +} + +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. +// +// 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 +// it. The UID changes only through a lifecycle operation, so admission's locked +// generation check also rejects a stale UID. +func (w *ActorWorkflow) observeCrash(ctx context.Context, session *apiv1.Session) (crashObservation, error) { + generation, err := w.store.GetSessionGeneration(ctx, session.Id) + if errors.Is(err, database.ErrNotFound) { + return crashObservation{}, fmt.Errorf("runtime Session was deleted: %w", errRuntimeUnrecoverable) + } + if err != nil { + return crashObservation{}, err + } + binding := substrate.ActorBinding{Atespace: generation.TemplateAtespace, Name: substrate.ActorName(session.Id), + 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) + } + if err != nil { + return crashObservation{}, err + } + if !binding.Matches(actor) { + return crashObservation{}, fmt.Errorf("runtime Actor identity or template changed: %w", errRuntimeUnrecoverable) + } + if actor.GetMetadata().GetUid() != generation.ActorUID { + return crashObservation{}, fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable) + } + return crashObservation{generation: generation, actor: actor}, nil +} + +// 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. +func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*apiv1.Session, error) { + if database.CrashRecoveryAdmitted(session) { + return w.run(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, false) + } + // Another lifecycle operation owns the Actor; a retry converges once it finishes. + if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return nil, database.ErrConflict + } + if !lifecycleIdle(session) { + return session, nil + } + // run bounds a joined attempt. Admission bounds its own reads and execution. + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + observed, err := w.observeCrash(ctx, session) + if err != nil { + return nil, err + } + if !observed.crashed() { + return session, nil + } + actor := observed.actor + snapshot := substrate.LatestDurableActorSnapshot(actor) + operation, err := w.store.BeginCrashedSessionRecovery(ctx, observed.generation, database.SessionTaskSnapshot{ + Atespace: actor.GetMetadata().GetAtespace(), URI: substrate.SnapshotURI(snapshot), + Fidelity: snapshot.GetFidelity().String(), + }) + if err != nil { + return nil, err + } + return w.execute(ctx, operation) +} diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go new file mode 100644 index 0000000000..2f0d05dcc5 --- /dev/null +++ b/go/core/internal/service/session/recovery_test.go @@ -0,0 +1,574 @@ +package session + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +func (a *lifecycleTestActors) RevertActor(_ context.Context, atespace, name string) (*ateapipb.Actor, error) { + a.mu.Lock() + defer a.mu.Unlock() + actor := a.actors[actorKey(atespace, name)] + if actor.Status.State != ateapipb.ActorState_ACTOR_STATE_CRASHED && actor.Status.State != ateapipb.ActorState_ACTOR_STATE_REVERTING { + return nil, status.Error(codes.FailedPrecondition, "actor is not crashed or reverting") + } + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + actor.Status.Crash = nil + snapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) + actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + return proto.CloneOf(actor), nil +} + +type recoveryTestActors struct { + *lifecycleTestActors + revert func(context.Context, string, string) (*ateapipb.Actor, error) + readErr error + // afterRead runs once, after the first successful read returns its copy. + afterRead func(context.Context) +} + +func (a *recoveryTestActors) RevertActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + return a.revert(ctx, atespace, name) +} + +func (a *recoveryTestActors) GetActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + if a.readErr != nil { + return nil, a.readErr + } + actor, err := a.lifecycleTestActors.GetActor(ctx, atespace, name) + if hook := a.afterRead; hook != nil && err == nil { + a.afterRead = nil + hook(ctx) + } + return actor, err +} + +func revertGeneration(t *testing.T, store *lifecycleTestStore, sessionID, actorUID string) *database.SessionGeneration { + t.Helper() + generation, err := store.GetSessionGeneration(t.Context(), sessionID) + require.NoError(t, err) + generation.ActorUID = actorUID + return generation +} + +func crashedTurnFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestActors, *apiv1.Session, *a2a.Task, int64) { + t.Helper() + store, session := lifecycleFixture(t) + actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + var err error + session, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Create(t.Context(), session) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("working")) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + return store, actors, session, task, version +} + +func TestCrashedSessionRecoveryRetriesWithoutLosingState(t *testing.T) { + for _, test := range []struct { + name string + lostResponse bool + wantReverts int + }{ + {name: "revert failed", wantReverts: 2}, + {name: "revert response lost", lostResponse: true, wantReverts: 1}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + calls := 0 + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + calls++ + if calls == 1 { + if test.lostResponse { + _, err := base.RevertActor(ctx, atespace, name) + require.NoError(t, err) + } + return nil, status.Error(codes.Unavailable, "response unavailable") + } + return base.RevertActor(ctx, atespace, name) + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err := workflow.Revert(t.Context(), session) + require.Error(t, err) + failed, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, failed.Operation) + require.Equal(t, database.SessionRuntimeCrashed, failed.Failure.Reason) + interrupted, interruptedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, interrupted.Status.State) + published, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, published.Status.State) + require.Len(t, published.History, 1) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString()), database.ErrConflict) + // A retried RevertSession on another workflow instance joins the admitted + // operation from PostgreSQL. + restarted := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = restarted.Revert(t.Context(), failed) + require.NoError(t, err) + ready, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.State) + require.Nil(t, ready.Failure) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, stored.Status.State) + require.Equal(t, test.wantReverts, calls) + require.Len(t, stored.History, 1, "the original input remains archived once") + _, recoveredVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, interruptedVersion, recoveredVersion, "retries do not append another terminal event") + retried, err := restarted.Revert(t.Context(), ready) + require.NoError(t, err, "a retry after a completed revert returns the Session") + require.True(t, proto.Equal(ready, retried)) + require.Equal(t, test.wantReverts, calls, "completed recovery cannot discard another execution") + }) + } +} + +func TestCrashedSessionRecoveryFencesConcurrentWork(t *testing.T) { + store, base, session, task, version := crashedTurnFixture(t) + entered, release := make(chan struct{}), make(chan struct{}) + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + close(entered) + select { + case <-release: + return base.RevertActor(ctx, atespace, name) + case <-ctx.Done(): + return nil, ctx.Err() + } + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + done := make(chan error, 1) + go func() { _, err := workflow.Revert(ctx, session); done <- err }() + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("recovery was not admitted") + } + failed, err := store.GetSessionByID(ctx, session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + _, err = workflow.Revert(ctx, failed) + require.ErrorIs(t, err, database.ErrConflict) + hash := sha256.Sum256([]byte("late write")) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) + close(release) + require.NoError(t, <-done) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) +} + +func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { + for _, test := range []struct { + name string + state ateapipb.ActorState + uid string + template string + readErr error + // want is the refusal a client sees; nil means the error must stay + // retryable. healthy means Revert returns the Session unchanged. + want error + healthy bool + }{ + {name: "quiet running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, healthy: true}, + {name: "transient read failure", readErr: status.Error(codes.Unavailable, "unreachable")}, + {name: "missing actor", readErr: status.Error(codes.NotFound, "actor not found"), want: errRuntimeUnrecoverable}, + {name: "same-name replacement", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, uid: "another-actor", want: errRuntimeUnrecoverable}, + {name: "template changed", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, template: "another-template", want: errRuntimeUnrecoverable}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + actor := base.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.State = test.state + if test.uid != "" { + actor.Metadata.Uid = test.uid + } + if test.template != "" { + actor.ActorTemplate.Name = test.template + } + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("unconfirmed crash issued a revert") + return nil, nil + } + returned, err := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Revert(t.Context(), session) + switch { + case test.healthy: + require.NoError(t, err) + require.True(t, proto.Equal(session, returned), "a healthy runtime has nothing to recover") + case test.want != nil: + require.ErrorIs(t, err, test.want) + require.ErrorIs(t, err, database.ErrFailedPrecondition) + default: + require.Error(t, err) + require.NotErrorIs(t, err, database.ErrFailedPrecondition, "a transient failure must stay retryable") + } + 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) + require.Equal(t, a2a.TaskStateWorking, stored.Status.State) + }) + } +} + +// A Session at rest outside READY has nothing to recover, so Revert returns it +// rather than ABORTED forever. Only an active operation is contention. +func TestExplicitRevertSeparatesContentionFromRestingState(t *testing.T) { + for _, test := range []struct { + name string + state apiv1.RuntimeState + operation apiv1.RuntimeOperation + want error + }{ + {name: "suspended by deadline cleanup", state: apiv1.RuntimeState_RUNTIME_STATE_SUSPENDED, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE}, + {name: "suspension in progress", state: apiv1.RuntimeState_RUNTIME_STATE_READY, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, want: database.ErrConflict}, + {name: "resume in progress", state: apiv1.RuntimeState_RUNTIME_STATE_SUSPENDED, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, want: database.ErrConflict}, + {name: "creation in progress", state: apiv1.RuntimeState_RUNTIME_STATE_CREATING, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, want: database.ErrConflict}, + } { + t.Run(test.name, func(t *testing.T) { + // Neither check reaches the store or Substrate. + workflow := NewActorWorkflow(nil, nil, make(chan struct{}, 1), time.Second) + session := &apiv1.Session{Id: uuid.NewString(), State: test.state, Operation: test.operation} + returned, err := workflow.Revert(t.Context(), session) + if test.want == nil { + require.NoError(t, err) + require.Same(t, session, returned) + return + } + require.ErrorIs(t, err, test.want) + require.NotErrorIs(t, err, database.ErrFailedPrecondition) + }) + } +} + +func TestCrashedSessionRecoversAnOlderResumedTask(t *testing.T) { + store, actors, session, older, version := crashedTurnFixture(t) + older.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + waitingVersion, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], older, older, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(older.ID), waitingVersion)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + newer := a2a.NewSubmittedTask(message, message) + newer.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("newer completed task")) + newerVersion, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], newer, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(newer.ID), newerVersion)) + work, err = store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/newer-completed-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + older.Status.State = a2a.TaskStateWorking + hash = sha256.Sum256([]byte("resume older task")) + _, err = store.UpdateSessionTask(t.Context(), session.Id, waitingVersion, hash[:], older, older, "") + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + failed, err := store.GetSettledSessionTask(t.Context(), session.Id, string(older.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, failed.Status.State) + retained, retainedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(newer.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) + require.Equal(t, newerVersion, retainedVersion, "recovery cannot rewrite a later completed task") + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) +} + +func TestExplicitRevertRespectsCheckpointCreation(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_CRASHED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +// The Actor crashed during post-turn suspension, and the worker kept its claim +// because it cannot tell whether the suspension took effect. +func TestExplicitRevertRetiresIdleSuspensionOfCrashedActor(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + recovered, err := workflow.Revert(t.Context(), session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + }), "a late finish of the retired claim is a no-op") + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound) +} + +// Revert A observes the crash, then pauses while Revert B restores the same Actor +// and ingress wakes it for a new turn. A then sees the same Actor UID and a READY +// Session, so only the lifecycle generation shows that its observation is stale. +func TestExplicitRevertRejectsCrashObservedBeforeAnotherRevert(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + var next *a2a.Task + paused := &recoveryTestActors{lifecycleTestActors: base, revert: base.RevertActor, afterRead: func(ctx context.Context) { + recovered, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Revert(ctx, session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + base.mu.Lock() + base.actors[key].Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING + base.mu.Unlock() + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next input")) + message.ContextID = session.ContextId + next = a2a.NewSubmittedTask(message, message) + next.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("next")) + _, err = store.CreateRuntimeTask(ctx, session.Id, hash[:], next, "") + require.NoError(t, err) + }} + _, err := NewActorWorkflow(store, paused, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + running, _, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(next.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "the stale revert must not fail the new turn") +} + +func TestExplicitRevertCannotCheckpointDiscardedTerminalState(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed before snapshot")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound, "the reverted snapshot cannot represent discarded terminal state") + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, database.ErrSnapshotPending) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +func TestExplicitRevertPreservesOnlyCapturedWaits(t *testing.T) { + for _, state := range []a2a.TaskState{a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired} { + for _, captured := range []bool{false, true} { + t.Run(fmt.Sprint(state, "/captured=", captured), func(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = state + hash := sha256.Sum256([]byte(state)) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + if captured { + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} + actor := actors.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.Snapshots[0].Locality = "lost-node" + actor.Status.Snapshots = append(actor.Status.Snapshots, &ateapipb.Snapshot{ + Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}, + }) + } + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + want := a2a.TaskStateFailed + if captured { + want = state + } + require.Equal(t, want, stored.Status.State) + }) + } + } +} + +func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { + store, actors, session, waiting, version := crashedTurnFixture(t) + waiting.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], waiting, waiting, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(waiting.ID), version)) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + later := a2a.NewSubmittedTask(message, message) + later.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("later completed turn")) + version, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], later, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(later.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/later-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + retained, err := store.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateInputRequired, retained.Status.State) + retained, err = store.GetSettledSessionTask(t.Context(), session.Id, string(later.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) +} + +// 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. +func TestDeleteSupersedesCrashRecovery(t *testing.T) { + for _, test := range []struct { + name string + claimed bool + gone bool + replaced bool + readErr error + }{ + {name: "actor read fails", readErr: status.Error(codes.Unavailable, "control plane unavailable")}, + {name: "actor is gone", gone: true}, + {name: "after an issued attempt", claimed: true}, + {name: "actor is gone after an issued attempt", claimed: true, gone: true}, + {name: "running replacement after an issued attempt", claimed: true, replaced: true}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + if test.claimed { + unavailable := &recoveryTestActors{lifecycleTestActors: base, revert: func(context.Context, string, string) (*ateapipb.Actor, error) { + return nil, status.Error(codes.Unavailable, "response unavailable") + }} + _, err := NewActorWorkflow(store, unavailable, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.Error(t, err) + } else { + _, err := store.BeginCrashedSessionRecovery(t.Context(), revertGeneration(t, store, session.Id, base.actors[key].GetMetadata().GetUid()), database.SessionTaskSnapshot{}) + require.NoError(t, err) + } + if test.gone { + delete(base.actors, key) + } + if test.replaced { + base.actors[key].Metadata.Uid = "replacement-uid" + base.actors[key].Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING + } + if unrecoverable := test.gone || test.replaced; unrecoverable || test.readErr != nil { + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("an unread or replaced Actor was reverted") + return nil, nil + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + for range 2 { + admitted, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = workflow.Revert(t.Context(), admitted) + require.Error(t, err) + require.Equal(t, unrecoverable, errors.Is(err, errRuntimeUnrecoverable)) + } + } + + 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") + 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) + deleting, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETING, deleting.State, "Delete closed admission") + replacement, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, "replacement-uid", replacement.GetMetadata().GetUid()) + return + } + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.State) + _, err = base.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + }) + } +} diff --git a/go/core/internal/service/session/service.go b/go/core/internal/service/session/service.go index b7549a9304..bb2578a4de 100644 --- a/go/core/internal/service/session/service.go +++ b/go/core/internal/service/session/service.go @@ -39,6 +39,7 @@ type store interface { type sessionWorkflow interface { Create(context.Context, *apiv1.Session) (*apiv1.Session, error) + Revert(context.Context, *apiv1.Session) (*apiv1.Session, error) Delete(context.Context, *apiv1.Session, bool) (*apiv1.Session, error) } @@ -280,6 +281,24 @@ func (s *Service) Delete(ctx context.Context, id string, force bool) (*apiv1.Ses return session, nil } +func (s *Service) Revert(ctx context.Context, id string) (*apiv1.Session, error) { + session, err := s.getAuthorized(ctx, id, auth.VerbUpdate) + if err != nil { + return nil, err + } + session, err = s.workflow.Revert(ctx, session) + if errors.Is(err, database.ErrFailedPrecondition) { + return nil, serviceerrors.NewFailedPrecondition(err.Error(), err) + } + if errors.Is(err, database.ErrConflict) { + return nil, serviceerrors.NewAborted(err.Error(), err) + } + if err != nil { + return nil, serviceerrors.NewUnavailable("Failed to revert Session", err) + } + return session, nil +} + // CreateShare mints a read-only share of the caller's session. A positive ttl bounds how long its // token grants access and must not exceed the configured maximum; zero takes the maximum, // or leaves the share valid until it is revoked or the session deleted when none is set. diff --git a/go/core/internal/service/session/service_test.go b/go/core/internal/service/session/service_test.go index 69e0f97ce0..8e312517f1 100644 --- a/go/core/internal/service/session/service_test.go +++ b/go/core/internal/service/session/service_test.go @@ -129,6 +129,10 @@ func (w serviceTestWorkflow) Create(_ context.Context, session *apiv1.Session) ( return session, w.err } +func (w serviceTestWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, w.err +} + func (w serviceTestWorkflow) Delete(_ context.Context, session *apiv1.Session, _ bool) (*apiv1.Session, error) { return session, w.err } @@ -201,23 +205,43 @@ func TestServiceCreateRejectsInvalidOrUnauthorizedRequests(t *testing.T) { } } -func TestServiceDeletePreservesContentionErrors(t *testing.T) { - for _, failure := range []struct { - cause error - code serviceerrors.Code +func TestServiceRevertRefusesUnrecoverableRuntime(t *testing.T) { + cause := fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable) + service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: cause}) + _, err := service.Revert(serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab") + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeFailedPrecondition), "error = %v", err) + require.Equal(t, cause.Error(), serviceerrors.MessageOf(err)) +} + +func TestServiceLifecycleMethodsPreserveContentionErrors(t *testing.T) { + for _, test := range []struct { + name string + call func(*Service, context.Context, string) (*apiv1.Session, error) }{ - {database.ErrConflict, serviceerrors.CodeAborted}, - {database.ErrFailedPrecondition, serviceerrors.CodeFailedPrecondition}, + {name: "revert", call: (*Service).Revert}, + {name: "delete", call: func(service *Service, ctx context.Context, id string) (*apiv1.Session, error) { + return service.Delete(ctx, id, false) + }}, } { - conflict := fmt.Errorf("Session has an operation in progress: %w", failure.cause) - service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: conflict}) - _, err := service.Delete(serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab", false) - if !serviceerrors.IsCode(err, failure.code) { - t.Fatalf("error = %v, want code %s", err, failure.code) - } - if serviceerrors.MessageOf(err) != conflict.Error() || !errors.Is(err, failure.cause) { - t.Fatalf("error = %v, want preserved contention reason", err) - } + t.Run(test.name, func(t *testing.T) { + for _, failure := range []struct { + cause error + code serviceerrors.Code + }{ + {database.ErrConflict, serviceerrors.CodeAborted}, + {database.ErrFailedPrecondition, serviceerrors.CodeFailedPrecondition}, + } { + conflict := fmt.Errorf("Session has an operation in progress: %w", failure.cause) + service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: conflict}) + _, err := test.call(service, serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab") + if !serviceerrors.IsCode(err, failure.code) { + t.Fatalf("error = %v, want code %s", err, failure.code) + } + if serviceerrors.MessageOf(err) != conflict.Error() || !errors.Is(err, failure.cause) { + t.Fatalf("error = %v, want preserved contention reason", err) + } + } + }) } } @@ -553,6 +577,7 @@ func TestServiceRejectsReadOnlyShareMutationsWithoutTransport(t *testing.T) { name string call func() error }{ + {name: "revert", call: func() error { _, err := service.Revert(ctx, id); return err }}, {name: "delete", call: func() error { _, err := service.Delete(ctx, id, false); return err }}, {name: "rename", call: func() error { _, err := service.Rename(ctx, id, "title"); return err }}, {name: "share listing", call: func() error { _, err := service.ListShares(ctx, id, 0, ""); return err }}, diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index ac4a66f107..14e3c61783 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -15,6 +15,8 @@ import ( ) type workflowStore interface { + GetSessionGeneration(context.Context, string) (*database.SessionGeneration, error) + BeginCrashedSessionRecovery(context.Context, *database.SessionGeneration, database.SessionTaskSnapshot) (*database.SessionOperation, error) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) HasPendingSessionQuiescence(context.Context) (bool, error) FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error @@ -104,7 +106,8 @@ func (w *ActorWorkflow) Suspend(ctx context.Context, session *apiv1.Session) (*a } // Delete closes admission before stopping and deleting compute. It can supersede -// unissued creation, but never deletes a session while a prior call is uncertain. +// unissued work and any crash recovery, but never deletes a session while another +// issued call is uncertain. func (w *ActorWorkflow) Delete(ctx context.Context, session *apiv1.Session, force bool) (*apiv1.Session, error) { return w.run(ctx, session.GetId(), apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, force) } @@ -157,17 +160,24 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session creation.Snapshot = &substrate.ActorSnapshot{Tag: &ateapipb.ObjectRef{Atespace: snapshot.Atespace, Name: tagName}, URI: snapshot.URI, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES} } } - prepare := substrate.PrepareActorTransition - if operation.ExecutorID != uuid.Nil { - prepare = substrate.PrepareActorRetry + var transition *substrate.ActorTransition + switch { + case database.CrashRecoveryAdmitted(session): + transition, err = substrate.PrepareActorRevert(ctx, w.actors, binding) + case operation.ExecutorID != uuid.Nil: + transition, err = substrate.PrepareActorRetry(ctx, w.actors, binding, kind, creation) + default: + transition, err = substrate.PrepareActorTransition(ctx, w.actors, binding, kind, creation) } - transition, err := prepare(ctx, w.actors, binding, kind, creation) if err != nil { return w.failPreparation(ctx, operation, err) } if uid := transition.ActorUID(); uid != "" && kind != apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE { if _, err := w.store.GetSessionForRuntime(ctx, session.Id, uid); err != nil { + if database.CrashRecoveryAdmitted(session) && errors.Is(err, database.ErrNotFound) { + err = fmt.Errorf("runtime Actor was replaced: %w", substrate.ErrActorUnrecoverable) + } return w.failPreparation(ctx, operation, fmt.Errorf("verify runtime actor UID: %w", err)) } } @@ -211,6 +221,15 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session // the same generation instead. Completion returns current state; supersession // returns a conflict. A local preparation error cannot clear a newer operation. func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database.SessionOperation, cause error) (*apiv1.Session, error) { + 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. + if errors.Is(cause, substrate.ErrActorUnrecoverable) { + return nil, fmt.Errorf("%w: %w", cause, errRuntimeUnrecoverable) + } + return nil, cause + } if admitted.ExecutorID != uuid.Nil { return nil, cause // An earlier attempt may have issued work; retain its intent. } diff --git a/go/core/internal/service/tool/mcp.go b/go/core/internal/service/tool/mcp.go index 1f4a067e54..999d5169c3 100644 --- a/go/core/internal/service/tool/mcp.go +++ b/go/core/internal/service/tool/mcp.go @@ -12,7 +12,9 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.opentelemetry.io/otel/trace" + corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "github.com/kagent-dev/kagent/go/api/v1alpha3" @@ -135,7 +137,7 @@ func (c *RuntimeMCPClient) connect(ctx context.Context, ref MCPServerRef) (*mcp. } connectCtx, cancel := context.WithTimeout(ctx, timeout) - headers, err := server.ResolveHeaders(connectCtx, c.kubeClient) + headers, err := resolveHeaders(connectCtx, c.kubeClient, server.Namespace, server.Spec.HeadersFrom) if err != nil { cancel() return nil, nil, fmt.Errorf("failed to resolve RemoteMCPServer headers: %w", err) @@ -276,3 +278,48 @@ func (t *mcpAppsHeaderTransport) RoundTrip(request *http.Request) (*http.Respons } return t.base.RoundTrip(request) } + +// resolveHeaders reads RemoteMCPServer headers, loading ValueFrom entries from +// Secrets or ConfigMaps in the server's namespace. +func resolveHeaders(ctx context.Context, kube client.Reader, namespace string, refs []v1alpha3.ValueRef) (map[string]string, error) { + headers := make(map[string]string, len(refs)) + for _, ref := range refs { + value, err := resolveHeaderValue(ctx, kube, namespace, ref) + if err != nil { + return nil, fmt.Errorf("failed to resolve header %s: %w", ref.Name, err) + } + headers[ref.Name] = value + } + return headers, nil +} + +func resolveHeaderValue(ctx context.Context, kube client.Reader, namespace string, ref v1alpha3.ValueRef) (string, error) { + if ref.ValueFrom == nil { + return ref.Value, nil + } + key := types.NamespacedName{Namespace: namespace, Name: ref.ValueFrom.Name} + switch ref.ValueFrom.Type { + case v1alpha3.SecretValueSource: + secret := &corev1.Secret{} + if err := kube.Get(ctx, key, secret); err != nil { + return "", fmt.Errorf("failed to get Secret %s: %w", key, err) + } + value, ok := secret.Data[ref.ValueFrom.Key] + if !ok { + return "", fmt.Errorf("key %s not found in Secret %s", ref.ValueFrom.Key, key) + } + return string(value), nil + case v1alpha3.ConfigMapValueSource: + configMap := &corev1.ConfigMap{} + if err := kube.Get(ctx, key, configMap); err != nil { + return "", fmt.Errorf("failed to get ConfigMap %s: %w", key, err) + } + value, ok := configMap.Data[ref.ValueFrom.Key] + if !ok { + return "", fmt.Errorf("key %s not found in ConfigMap %s", ref.ValueFrom.Key, key) + } + return value, nil + default: + return "", fmt.Errorf("unknown value source type %q", ref.ValueFrom.Type) + } +} diff --git a/go/core/internal/service/tool/mcp_test.go b/go/core/internal/service/tool/mcp_test.go index 9f5b347849..402b56a9bf 100644 --- a/go/core/internal/service/tool/mcp_test.go +++ b/go/core/internal/service/tool/mcp_test.go @@ -7,6 +7,7 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" @@ -117,3 +118,72 @@ func TestRuntimeMCPClientResolveServerMatrix(t *testing.T) { }) } } + +func TestResolveHeaders(t *testing.T) { + secret := &corev1.Secret{ + ObjectMeta: metav1.ObjectMeta{Name: "creds", Namespace: "team"}, + Data: map[string][]byte{"token": []byte("secret-token")}, + } + configMap := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{Name: "settings", Namespace: "team"}, + Data: map[string]string{"tenant": "acme"}, + } + + t.Run("Given inline, Secret, and ConfigMap headers, when the server headers are resolved, assert that every value is loaded from the server namespace", func(t *testing.T) { + kube := toolTestKube(t, secret, configMap) + + headers, err := resolveHeaders(t.Context(), kube, "team", []v1alpha3.ValueRef{ + {Name: "X-Static", Value: "static"}, + {Name: "Authorization", ValueFrom: &v1alpha3.ValueSource{Type: v1alpha3.SecretValueSource, Name: "creds", Key: "token"}}, + {Name: "X-Tenant", ValueFrom: &v1alpha3.ValueSource{Type: v1alpha3.ConfigMapValueSource, Name: "settings", Key: "tenant"}}, + }) + + require.NoError(t, err) + assert.Equal(t, map[string]string{ + "X-Static": "static", + "Authorization": "secret-token", + "X-Tenant": "acme", + }, headers) + }) + + scenarios := []struct { + name string + objects []client.Object + ref v1alpha3.ValueRef + wantError string + }{ + { + name: "Given the Secret lives in another namespace, when a Secret header is resolved, assert that resolution fails", + objects: []client.Object{&corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: "creds", Namespace: "other"}, Data: secret.Data}}, + ref: v1alpha3.ValueRef{Name: "Authorization", ValueFrom: &v1alpha3.ValueSource{Type: v1alpha3.SecretValueSource, Name: "creds", Key: "token"}}, + wantError: "failed to get Secret team/creds", + }, + { + name: "Given the Secret lacks the key, when a Secret header is resolved, assert that resolution fails", + objects: []client.Object{secret}, + ref: v1alpha3.ValueRef{Name: "Authorization", ValueFrom: &v1alpha3.ValueSource{Type: v1alpha3.SecretValueSource, Name: "creds", Key: "missing"}}, + wantError: "key missing not found in Secret team/creds", + }, + { + name: "Given the ConfigMap lacks the key, when a ConfigMap header is resolved, assert that resolution fails", + objects: []client.Object{configMap}, + ref: v1alpha3.ValueRef{Name: "X-Tenant", ValueFrom: &v1alpha3.ValueSource{Type: v1alpha3.ConfigMapValueSource, Name: "settings", Key: "missing"}}, + wantError: "key missing not found in ConfigMap team/settings", + }, + { + name: "Given an unknown source type, when the header is resolved, assert that resolution fails", + ref: v1alpha3.ValueRef{Name: "X-Other", ValueFrom: &v1alpha3.ValueSource{Type: "Vault", Name: "x", Key: "y"}}, + wantError: `unknown value source type "Vault"`, + }, + } + for _, scenario := range scenarios { + t.Run(scenario.name, func(t *testing.T) { + kube := toolTestKube(t, scenario.objects...) + + headers, err := resolveHeaders(t.Context(), kube, "team", []v1alpha3.ValueRef{scenario.ref}) + + require.ErrorContains(t, err, scenario.wantError) + assert.Nil(t, headers) + }) + } +} diff --git a/go/core/internal/substrate/client.go b/go/core/internal/substrate/client.go index 1db776de1a..002916413d 100644 --- a/go/core/internal/substrate/client.go +++ b/go/core/internal/substrate/client.go @@ -203,6 +203,16 @@ func (c *Client) ResumeActor(ctx context.Context, atespace, actorID string) (*at return resp.GetActor(), nil } +func (c *Client) RevertActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { + ctx, cancel := c.callCtx(ctx) + defer cancel() + resp, err := c.ControlClient.RevertActor(ctx, &ateapipb.RevertActorRequest{Actor: actorRef(atespace, actorID)}) + if err != nil { + return nil, fmt.Errorf("failed to revert Actor %s/%s: %w", atespace, actorID, err) + } + return resp.GetActor(), nil +} + // SuspendActor waits for durable storage, including an upload left by an earlier // attempt. Substrate's synchronous mode returns already-suspended Actors as-is. func (c *Client) SuspendActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { diff --git a/go/core/internal/substrate/client_test.go b/go/core/internal/substrate/client_test.go index d709a34a7b..d063e77b6c 100644 --- a/go/core/internal/substrate/client_test.go +++ b/go/core/internal/substrate/client_test.go @@ -430,6 +430,8 @@ func TestSuspendActorWaitsForLatestDurableSnapshot(t *testing.T) { fake := &suspendActorFake{actor: actor, get: func(context.Context) (*ateapipb.Actor, error) { reads++; return next, nil }} client := &Client{ControlClient: fake, cfg: Config{CallTimeout: time.Second}} got, err := client.SuspendActor(t.Context(), "team", "actor") + // Crash recovery depends on this: with no upload left in progress, a + // revert restores the snapshot admission settled against (applyRevert). require.Equal(t, ateapipb.DurabilityMode_DURABILITY_MODE_SYNCHRONOUS, fake.request.GetOptions().GetMode()) require.Equal(t, 1, reads) if test.wantError != "" { diff --git a/go/core/internal/substrate/lifecycle.go b/go/core/internal/substrate/lifecycle.go index b5f2e6b436..40c040d59a 100644 --- a/go/core/internal/substrate/lifecycle.go +++ b/go/core/internal/substrate/lifecycle.go @@ -2,6 +2,7 @@ package substrate import ( "context" + "errors" "fmt" ateapipb "go.ate.dev/substrate/api/ateapi/v1" @@ -19,6 +20,7 @@ type LifecycleClient interface { CheckActorEgressPolicy(context.Context, string, string) error EnsureActorEgressPolicy(context.Context, string, string, string) error ResumeActor(context.Context, string, string) (*ateapipb.Actor, error) + RevertActor(context.Context, string, string) (*ateapipb.Actor, error) SuspendActor(context.Context, string, string) (*ateapipb.Actor, error) DeleteActor(context.Context, string, string, string) error } @@ -53,6 +55,10 @@ type ActorTransition struct { creation *ActorCreation } +// ErrActorUnrecoverable reports that a crash recovery's Actor is missing or was +// replaced. No retry changes that, and no revert of it can still be in progress. +var ErrActorUnrecoverable = errors.New("actor is missing or was replaced") + // ActorUID returns the observed runtime identity, including after creation. func (t *ActorTransition) ActorUID() string { return t.actor.GetMetadata().GetUid() @@ -118,6 +124,23 @@ func prepareActorTransition(ctx context.Context, actors LifecycleClient, binding return &ActorTransition{binding: binding, operation: operation, actor: actor, creation: creation}, nil } +// PrepareActorRevert reads the Actor before the store's durable claim, so a +// failed read leaves no claim. A missing or retemplated Actor is unrecoverable. +// The caller then checks the recorded UID; applyRevert checks the state. +func PrepareActorRevert(ctx context.Context, actors LifecycleClient, binding ActorBinding) (*ActorTransition, error) { + actor, err := actors.GetActor(ctx, binding.Atespace, binding.Name) + if status.Code(err) == codes.NotFound { + return nil, fmt.Errorf("actor %s/%s is missing: %w", binding.Atespace, binding.Name, ErrActorUnrecoverable) + } + if err != nil { + return nil, fmt.Errorf("get Actor %s/%s: %w", binding.Atespace, binding.Name, err) + } + if !binding.Matches(actor) { + return nil, fmt.Errorf("actor %s/%s identity or template changed: %w", binding.Atespace, binding.Name, ErrActorUnrecoverable) + } + return &ActorTransition{binding: binding, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT, actor: actor}, nil +} + // ApplyActorTransition requires durable execution authority. Errors retain the // claim's durable intent and pins. Another bounded attempt can retry that intent. func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transition *ActorTransition) error { @@ -178,6 +201,8 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio err = fmt.Errorf("suspended Actor requires a durable VOLUMES snapshot") } } + case apiv1.RuntimeOperation_RUNTIME_OPERATION_REVERT: + err = transition.applyRevert(ctx, actors) case apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE: if actor != nil { // An interrupted resume must not make deletion depend on a successful @@ -194,6 +219,21 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio return nil } +// applyRevert requires the same Actor back, SUSPENDED. Substrate restores the +// last completed durable snapshot, which is the one admission settled against: +// every suspend kagent issues is synchronous, so no upload completes in between. +func (t *ActorTransition) applyRevert(ctx context.Context, actors LifecycleClient) error { + actor := t.actor + if state := actor.GetStatus().GetState(); state == ateapipb.ActorState_ACTOR_STATE_CRASHED || state == ateapipb.ActorState_ACTOR_STATE_REVERTING { + reverted, err := actors.RevertActor(ctx, t.binding.Atespace, t.binding.Name) + if err != nil { + return err + } + actor = reverted + } + return t.checkResult(actor, ateapipb.ActorState_ACTOR_STATE_SUSPENDED) +} + // Matches checks that the actor belongs to the expected runtime and template. func (b ActorBinding) Matches(actor *ateapipb.Actor) bool { return actor.GetMetadata().GetName() == b.Name && actor.GetMetadata().GetAtespace() == b.Atespace && actor.GetMetadata().GetUid() != "" && diff --git a/go/core/internal/substrate/lifecycle_test.go b/go/core/internal/substrate/lifecycle_test.go index 7bd221c199..1cf9f33b2e 100644 --- a/go/core/internal/substrate/lifecycle_test.go +++ b/go/core/internal/substrate/lifecycle_test.go @@ -23,6 +23,7 @@ type transitionActors struct { checkErr error resumeUID string suspendSnapshot *ateapipb.Snapshot + revertUID string } func (a *transitionActors) GetActor(context.Context, string, string) (*ateapipb.Actor, error) { @@ -70,6 +71,21 @@ func (a *transitionActors) SuspendActor(context.Context, string, string) (*ateap return proto.CloneOf(a.actor), nil } +func (a *transitionActors) RevertActor(context.Context, string, string) (*ateapipb.Actor, error) { + a.calls = append(a.calls, "revert") + a.actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + snapshot := proto.CloneOf(LatestDurableActorSnapshot(a.actor)) + a.actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + a.actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + if a.revertUID != "" { + a.actor.Metadata.Uid = a.revertUID + } + return proto.CloneOf(a.actor), nil +} + func (a *transitionActors) DeleteActor(context.Context, string, string, string) error { a.calls = append(a.calls, "delete") a.actor = nil @@ -242,3 +258,58 @@ func TestActorSuspendRequiresDurableVolumesSnapshot(t *testing.T) { }) } } + +func TestActorRevertKeepsTheRecordedActor(t *testing.T) { + durable := func(generation int32, uri string) *ateapipb.Snapshot { + return &ateapipb.Snapshot{Generation: generation, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, Locality: "lost-node", + DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} + } + for _, test := range []struct { + name string + missing bool + state ateapipb.ActorState + revertUID string + wantPrepareErr bool + unrecoverable bool + wantApplyErr bool + wantCalls []string + }{ + {name: "crashed", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, wantCalls: []string{"revert"}}, + {name: "revert in progress", state: ateapipb.ActorState_ACTOR_STATE_REVERTING, wantCalls: []string{"revert"}}, + {name: "earlier revert response lost", state: ateapipb.ActorState_ACTOR_STATE_SUSPENDED}, + {name: "running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantApplyErr: true}, + {name: "actor gone", missing: true, wantPrepareErr: true, unrecoverable: true}, + {name: "replaced actor", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, revertUID: "another-actor", wantApplyErr: true, wantCalls: []string{"revert"}}, + } { + t.Run(test.name, func(t *testing.T) { + actors := &transitionActors{revertUID: test.revertUID} + if !test.missing { + actors.actor = &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team", Name: "instance", Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team", Name: "revision"}, + Status: &ateapipb.ActorStatus{State: test.state, Snapshots: []*ateapipb.Snapshot{durable(2, "s3://snapshots/latest")}}, + } + if test.state == ateapipb.ActorState_ACTOR_STATE_SUSPENDED { + actors.actor.Status.Snapshots[0].Locality = "" + } + } + binding := ActorBinding{Revision: "revision", Atespace: "team", Name: "instance", TemplateAtespace: "team", TemplateName: "revision"} + transition, err := PrepareActorRevert(t.Context(), actors, binding) + if test.wantPrepareErr { + require.Error(t, err) + require.Equal(t, test.unrecoverable, errors.Is(err, ErrActorUnrecoverable)) + require.Empty(t, actors.calls, "a failed read issues nothing") + return + } + require.NoError(t, err) + require.Equal(t, "actor-uid", transition.ActorUID()) + err = ApplyActorTransition(t.Context(), actors, transition) + if test.wantApplyErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.Equal(t, test.wantCalls, actors.calls) + }) + } +} diff --git a/go/core/internal/substrate/snapshot.go b/go/core/internal/substrate/snapshot.go index 3d1163c28e..302c07e975 100644 --- a/go/core/internal/substrate/snapshot.go +++ b/go/core/internal/substrate/snapshot.go @@ -22,3 +22,14 @@ func SnapshotURI(snapshot *ateapipb.Snapshot) string { } return storage.GetObject().GetSnapshotUri() } + +// Checkpoints require the newest generation; revert retains the newest durable copy. +func LatestDurableActorSnapshot(actor *ateapipb.Actor) *ateapipb.Snapshot { + var latest *ateapipb.Snapshot + for _, snapshot := range actor.GetStatus().GetSnapshots() { + if SnapshotURI(snapshot) != "" && snapshot.GetGeneration() > latest.GetGeneration() { + latest = snapshot + } + } + return latest +} diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index faee8a936c..4e923131a8 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -40,3 +40,23 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { require.Nil(t, LatestActorSnapshot(nil)) require.Empty(t, SnapshotURI(nil)) } + +func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { + durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ + State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, + Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, + }} + for _, latest := range []*ateapipb.Snapshot{ + {Generation: 3, Locality: "node-1"}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED}}, + } { + actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} + require.Same(t, durable, LatestDurableActorSnapshot(actor)) + require.Same(t, latest, LatestActorSnapshot(actor)) + } + require.Nil(t, LatestDurableActorSnapshot(nil)) + require.Nil(t, LatestDurableActorSnapshot(&ateapipb.Actor{Status: &ateapipb.ActorStatus{ + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Locality: "node-1"}}, + }})) +} diff --git a/go/core/pkg/migrations/core/000003_session_revert_operation.sql b/go/core/pkg/migrations/core/000003_session_revert_operation.sql new file mode 100644 index 0000000000..2fe3fc1493 --- /dev/null +++ b/go/core/pkg/migrations/core/000003_session_revert_operation.sql @@ -0,0 +1,17 @@ +-- +goose Up + +-- Crash recovery is its own lifecycle operation, started only after Substrate +-- confirms that the recorded Actor crashed. +ALTER TABLE runtime_instance DROP CONSTRAINT runtime_instance_operation_check; +ALTER TABLE runtime_instance ADD CONSTRAINT runtime_instance_operation_check + CHECK (operation IN ('RUNTIME_OPERATION_NONE', 'RUNTIME_OPERATION_CREATE', 'RUNTIME_OPERATION_SUSPEND', + 'RUNTIME_OPERATION_RESUME', 'RUNTIME_OPERATION_DELETE', 'RUNTIME_OPERATION_REVERT')); + +-- +goose Down + +-- Fails while any recovery is still in flight, rather than rewriting its +-- operation into one the older binary would execute differently. +ALTER TABLE runtime_instance DROP CONSTRAINT runtime_instance_operation_check; +ALTER TABLE runtime_instance ADD CONSTRAINT runtime_instance_operation_check + CHECK (operation IN ('RUNTIME_OPERATION_NONE', 'RUNTIME_OPERATION_CREATE', 'RUNTIME_OPERATION_SUSPEND', + 'RUNTIME_OPERATION_RESUME', 'RUNTIME_OPERATION_DELETE')); diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 8c15d7f17a..10fae8fbd3 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -320,7 +320,7 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { for _, source := range sources { expected := []int64{0, 1} if source.Name == "core" { - expected = append(expected, 2) + expected = append(expected, 2, 3) } if versions := testVersions(t, dsn, source.TrackingTable); !slices.Equal(versions, expected) { t.Fatalf("%s versions = %v, want %v", source.Name, versions, expected) diff --git a/go/core/test/e2e/mcp_discovery_test.go b/go/core/test/e2e/mcp_discovery_test.go new file mode 100644 index 0000000000..bd4369a2a3 --- /dev/null +++ b/go/core/test/e2e/mcp_discovery_test.go @@ -0,0 +1,440 @@ +package e2e_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" + "github.com/google/uuid" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + kagenta2a "github.com/kagent-dev/kagent/go/api/a2a" + "github.com/kagent-dev/kagent/go/api/v1alpha3" +) + +func TestMCPClassificationFailurePreservesHealthyToolsAndRecovers(t *testing.T) { + // Given healthy and unavailable catalogs, when an execution continues and + // another execution resumes, assert that failed tools stay hidden until recovery. + target := interactionTarget(t) + var unavailable atomic.Bool + unavailable.Store(true) + var runtimeLists, healthyCalls, recoveredCalls, hiddenCalls atomic.Int32 + healthy := newMCPDiscoveryServer() + mcpsdk.AddTool(healthy, &mcpsdk.Tool{Name: "healthy_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + healthyCalls.Add(1) + // Recovery between model steps must not undo this execution's omission. + unavailable.Store(false) + return nil, map[string]any{"result": "healthy"}, nil + }) + recovering := newMCPDiscoveryServer() + mcpsdk.AddTool(recovering, &mcpsdk.Tool{Name: "recovered_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + recoveredCalls.Add(1) + return nil, map[string]any{"result": "recovered"}, nil + }) + mcpsdk.AddTool(recovering, &mcpsdk.Tool{Name: "app_only_tool", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + hiddenCalls.Add(1) + return nil, map[string]any{"result": "hidden"}, nil + }) + recovering.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetSession().(*mcpsdk.ServerSession).InitializeParams().ClientInfo.Name == "kagent-adk" { + runtimeLists.Add(1) + if unavailable.Load() { + return nil, errors.New("runtime catalog temporarily unavailable") + } + } + return next(ctx, method, req) + } + }) + var modelRequests atomic.Int32 + modelURL := startMCPDiscoveryModel(t, func(request discoveryModelRequest) (string, string, error) { + modelRequests.Add(1) + recovered := strings.Contains(request.lastUser(), "recovered") + names := request.toolNames() + if !slices.Contains(names, "healthy_tool") || slices.Contains(names, "app_only_tool") || slices.Contains(names, "recovered_tool") != recovered { + return "", "", fmt.Errorf("unexpected tool visibility for recovered=%t: %v", recovered, names) + } + toolName, result := "healthy_tool", "HEALTHY_DISCOVERY_OK" + if recovered { + toolName, result = "recovered_tool", "RECOVERED_DISCOVERY_OK" + } + if _, ok := request.toolResults(true)[toolName]; !ok { + return toolName, "", nil + } + return "", result, nil + }) + template := createMCPDiscoveryTemplate(t, modelURL, startMCPDiscoveryServer(t, healthy), startMCPDiscoveryServer(t, recovering)) + require.Zero(t, runtimeLists.Load(), "runtime startup must not classify remote tools") + fixture := newInteractionFixtureForTemplate(t, testHarness{name: "kagent"}, target, template) + _, _, first := fixture.send(t, "Use the healthy tool.") + require.Equal(t, a2atype.TaskStateCompleted, first.Status.State, taskText(first)) + require.Contains(t, taskText(first), "HEALTHY_DISCOVERY_OK") + require.EqualValues(t, 1, healthyCalls.Load()) + require.Zero(t, recoveredCalls.Load()) + require.EqualValues(t, 1, runtimeLists.Load(), "an omitted catalog must not be rediscovered in later model steps") + + assertActorSuspended(t, fixture) + _, _, second := fixture.send(t, "Use the recovered tool.") + require.Equal(t, a2atype.TaskStateCompleted, second.Status.State, taskText(second)) + require.Contains(t, taskText(second), "RECOVERED_DISCOVERY_OK") + require.EqualValues(t, 1, recoveredCalls.Load()) + require.EqualValues(t, 1, healthyCalls.Load()) + require.Zero(t, hiddenCalls.Load()) + require.Greater(t, runtimeLists.Load(), int32(1)) + require.EqualValues(t, 4, modelRequests.Load(), "visibility must be checked before and after both tool calls") +} + +func TestMCPAppCompactionPreservesHistoryAcrossResume(t *testing.T) { + // Given an app tool and an ordinary tool, when a restored execution prepares + // model history, assert that only app results are compacted and public history stays rich. + target := interactionTarget(t) + renderPayload := map[string]any{"panel": map[string]any{"title": "Discovery test panel", "rows": []any{"RICH_MCP_RENDER_PAYLOAD", strings.Repeat("render data ", 128)}}} + ordinaryPayload := map[string]any{"note": "ORDINARY_MCP_RESULT"} + var appCalls, ordinaryCalls, hiddenCalls, runtimeLists atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "cached-app-discovery", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" && req.(*mcpsdk.ListToolsRequest).ClientInfo().Name == "kagent-adk" { + runtimeLists.Add(1) + } + return next(ctx, method, req) + } + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "show_panel", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://discovery/panel"}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + appCalls.Add(1) + return nil, renderPayload, nil + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "ordinary_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + ordinaryCalls.Add(1) + return nil, ordinaryPayload, nil + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "app_only_tool", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + hiddenCalls.Add(1) + return nil, map[string]any{"result": "hidden"}, nil + }) + var restoredRequests atomic.Int32 + modelURL := startMCPDiscoveryModel(t, func(request discoveryModelRequest) (string, string, error) { + names := request.toolNames() + if !slices.Contains(names, "show_panel") || !slices.Contains(names, "ordinary_tool") || slices.Contains(names, "app_only_tool") { + return "", "", fmt.Errorf("unexpected model tools: %v", names) + } + results := request.toolResults(false) + appResult, rendered := results["show_panel"] + if !rendered { + return "show_panel", "", nil + } + if appResult != "The interactive UI for this tool has been rendered to the user and now updates live inside the app. Treat this as complete and do not call this tool again unless the user explicitly asks for it." { + return "", "", fmt.Errorf("app result was not compacted: %s", appResult) + } + ordinaryResult, called := results["ordinary_tool"] + if !called { + return "ordinary_tool", "", nil + } + var output struct { + Output map[string]any `json:"output"` + } + if err := json.Unmarshal([]byte(ordinaryResult), &output); err != nil || len(output.Output) != len(ordinaryPayload) || output.Output["note"] != ordinaryPayload["note"] { + return "", "", fmt.Errorf("ordinary result was modified: %s", ordinaryResult) + } + if strings.Contains(request.lastUser(), "without calling tools") { + restoredRequests.Add(1) + return "", "MCP_APP_HISTORY_RESTORED", nil + } + return "", "MCP_APP_RENDERED", nil + }) + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true}) + serverURL := reachableServerURL(t, startMCPDiscoveryHTTP(t, handler), "/mcp") + template := createMCPDiscoveryTemplate(t, modelURL, serverURL) + fixture := newInteractionFixtureForTemplate(t, testHarness{name: "kagent"}, target, template) + _, _, first := fixture.send(t, "Render the panel and call the ordinary tool.") + require.Equal(t, a2atype.TaskStateCompleted, first.Status.State, taskText(first)) + require.Contains(t, taskText(first), "MCP_APP_RENDERED") + require.EqualValues(t, 1, runtimeLists.Load(), "SDK caching must preserve filtering and app metadata across model steps") + assertPersistedMCPResult(t, fixture, first.ID, "show_panel", renderPayload) + assertPersistedMCPResult(t, fixture, first.ID, "ordinary_tool", ordinaryPayload) + + assertActorSuspended(t, fixture) + _, _, second := fixture.send(t, "Summarize the previous results without calling tools.") + require.Equal(t, a2atype.TaskStateCompleted, second.Status.State, taskText(second)) + require.Contains(t, taskText(second), "MCP_APP_HISTORY_RESTORED") + require.EqualValues(t, 1, restoredRequests.Load(), "the first restored model request must already contain compacted history") + require.EqualValues(t, 1, appCalls.Load(), "restoring history must not invoke the app again") + require.EqualValues(t, 1, ordinaryCalls.Load()) + require.Zero(t, hiddenCalls.Load()) + require.EqualValues(t, 2, runtimeLists.Load(), "a restored process must discover once after readiness") + assertPersistedMCPResult(t, fixture, first.ID, "show_panel", renderPayload) + assertPersistedMCPResult(t, fixture, first.ID, "ordinary_tool", ordinaryPayload) +} + +func TestSlowMCPDiscoveryDoesNotBlockResume(t *testing.T) { + for _, slow := range []string{"tools/list", "GET"} { + t.Run("Given slow "+strings.ReplaceAll(slow, "/", " ")+", when an agent becomes ready and resumes repeatedly, assert that discovery waits until invocation", func(t *testing.T) { + target := interactionTarget(t) + harness := testHarness{name: "kagent", runtimeLabel: "kagent"} + url, runtimeDiscovery := startSlowMCPServer(t, slow) + template, _ := createMCPInteractionTemplate(t, harness, url, false, nil) + require.Zero(t, runtimeDiscovery.Load(), "golden runtime readiness must not depend on MCP discovery") + fixture := newInteractionFixtureForTemplate(t, harness, target, template) + for turn := range 2 { + if turn > 0 { + assertActorSuspended(t, fixture) + } + // Use one SendMessage call: retries would hide a parking timeout. + _, request := newMessageRequest(t, "Add 3 and 5 using the configured MCP server.") + request.Message.ContextId = fixture.sessionID + request.Tenant = fixture.tenant + response, err := fixture.client.SendMessage(fixture.ctx, request) + require.NoError(t, err, "slow MCP discovery must not make the router reject a restore") + result, err := pbconv.FromProtoSendMessageResponse(response) + require.NoError(t, err) + task, ok := result.(*a2atype.Task) + require.True(t, ok, "A2A response = %T, want Task", result) + require.Equal(t, a2atype.TaskStateCompleted, task.Status.State, taskText(task)) + require.Contains(t, taskText(task), "result is 8") + } + require.Positive(t, runtimeDiscovery.Load(), "successful invocations must actually discover the slow server") + }) + } +} + +func startSlowMCPServer(t *testing.T, slow string) (string, *atomic.Int32) { + t.Helper() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "slow-startup", Version: "1"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "add_numbers", Description: "Add two numbers"}, func(_ context.Context, _ *mcpsdk.CallToolRequest, input struct { + A int `json:"a"` + B int `json:"b"` + }) (*mcpsdk.CallToolResult, map[string]any, error) { + return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "The result is 8."}}}, map[string]any{"result": input.A + input.B}, nil + }) + var clients sync.Map + runtimeDiscovery := new(atomic.Int32) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + session := req.GetSession().(*mcpsdk.ServerSession) + if method == "initialize" { + clients.Store(session.ID(), req.GetParams().(*mcpsdk.InitializeParams).ClientInfo.Name) + } + if method == "tools/list" && session.InitializeParams().ClientInfo.Name == "kagent-adk" { + runtimeDiscovery.Add(1) + if slow == method { + if err := waitForSlowDiscovery(ctx); err != nil { + return nil, err + } + } + } + return next(ctx, method, req) + } + }) + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, nil) + baseURL := startMCPDiscoveryHTTP(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + if client, ok := clients.Load(r.Header.Get("Mcp-Session-Id")); ok && client == "kagent-adk" { + runtimeDiscovery.Add(1) + if slow == r.Method { + if err := waitForSlowDiscovery(r.Context()); err != nil { + return + } + } + } + } + handler.ServeHTTP(w, r) + })) + return reachableServerURL(t, baseURL, "/mcp"), runtimeDiscovery +} + +func waitForSlowDiscovery(ctx context.Context) error { + // Longer than the E2E router's five-second parking budget, while fitting + // within the configured MCP request timeout. Workload requests may take longer. + timer := time.NewTimer(8 * time.Second) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func newMCPDiscoveryServer() *mcpsdk.Server { + return mcpsdk.NewServer(&mcpsdk.Implementation{Name: "discovery-e2e", Version: "1"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) +} + +func startMCPDiscoveryServer(t *testing.T, server *mcpsdk.Server) string { + t.Helper() + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, nil) + return reachableServerURL(t, startMCPDiscoveryHTTP(t, handler), "/mcp") +} + +func startMCPDiscoveryHTTP(t *testing.T, handler http.Handler) string { + t.Helper() + server := httptest.NewUnstartedServer(handler) + require.NoError(t, server.Listener.Close()) + var err error + server.Listener, err = net.Listen("tcp", "0.0.0.0:0") + require.NoError(t, err) + server.Start() + t.Cleanup(func() { server.CloseClientConnections(); server.Close() }) + return server.URL +} + +func createMCPDiscoveryTemplate(t *testing.T, modelURL string, serverURLs ...string) string { + t.Helper() + kube := interactionKubeClient(t) + harness := testHarness{name: "kagent", runtimeLabel: "kagent"} + model := &v1alpha3.ModelConfig{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent"}, + Spec: v1alpha3.ModelConfigSpec{Provider: v1alpha3.ModelProviderOpenAI, Model: "gpt-4.1-mini", Stream: new(false), OpenAI: &v1alpha3.OpenAIConfig{BaseURL: modelURL}}, + } + require.NoError(t, kube.Create(t.Context(), model)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), model); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete discovery model: %v", err) + } + }) + template := &v1alpha3.AgentTemplate{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent", Labels: harness.labels()}, + Spec: v1alpha3.AgentTemplateSpec{Description: "MCP discovery E2E fixture", SystemPrompt: "Use the configured MCP tools.", ModelConfig: &v1alpha3.ObjectReference{Name: model.Name}}, + } + for _, serverURL := range serverURLs { + server := &v1alpha3.RemoteMCPServer{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent"}, + Spec: v1alpha3.RemoteMCPServerSpec{Description: "MCP discovery E2E fixture", Protocol: v1alpha3.RemoteMCPServerProtocolStreamableHttp, URL: serverURL}, + } + require.NoError(t, kube.Create(t.Context(), server)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), server); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete MCP server: %v", err) + } + }) + template.Spec.Tools = append(template.Spec.Tools, v1alpha3.ToolBinding{MCP: &v1alpha3.MCPToolBinding{Server: corev1.TypedObjectReference{Kind: "RemoteMCPServer", Name: server.Name}}}) + } + createAndWaitInteractionTemplateForHarness(t, kube, template, harness.name) + return template.Name +} + +func assertPersistedMCPResult(t *testing.T, fixture *interactionFixture, taskID a2atype.TaskID, toolName string, payload map[string]any) { + t.Helper() + request, err := pbconv.ToProtoGetTaskRequest(&a2atype.GetTaskRequest{Tenant: fixture.tenant, ID: taskID}) + require.NoError(t, err) + response, err := fixture.client.GetTask(fixture.ctx, request) + require.NoError(t, err) + task, err := pbconv.FromProtoTask(response) + require.NoError(t, err) + var responses []any + // The runtime publishes tool results as A2A artifacts. These are the + // persisted chat timeline payloads; Task.History contains input messages. + for _, artifact := range task.Artifacts { + for _, part := range artifact.Parts { + activity, ok := kagenta2a.ParseToolActivity(part) + if ok && activity.Kind == kagenta2a.ToolResultKind && activity.Name == toolName { + responses = append(responses, activity.Response) + } + } + } + require.Len(t, responses, 1, "persisted artifacts must contain the original %s result", toolName) + require.Equal(t, map[string]any{"output": payload}, responses[0]) +} + +type discoveryModelRequest struct { + Messages []struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCallID string `json:"tool_call_id"` + ToolCalls []struct { + ID string `json:"id"` + Function struct { + Name string `json:"name"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"messages"` + Tools []struct { + Function struct { + Name string `json:"name"` + } `json:"function"` + } `json:"tools"` +} + +func (r discoveryModelRequest) lastUser() string { + for _, message := range slices.Backward(r.Messages) { + if message.Role == "user" { + return message.Content + } + } + return "" +} + +func (r discoveryModelRequest) toolNames() []string { + var names []string + for _, tool := range r.Tools { + names = append(names, tool.Function.Name) + } + return names +} + +func (r discoveryModelRequest) toolResults(currentTurnOnly bool) map[string]string { + names := make(map[string]string) + start := 0 + for i, message := range r.Messages { + for _, call := range message.ToolCalls { + names[call.ID] = call.Function.Name + } + if currentTurnOnly && message.Role == "user" { + start = i + } + } + results := make(map[string]string) + for _, message := range r.Messages[start:] { + if message.Role == "tool" { + results[names[message.ToolCallID]] = message.Content + } + } + return results +} + +func startMCPDiscoveryModel(t *testing.T, respond func(discoveryModelRequest) (toolName, text string, err error)) string { + t.Helper() + baseURL := startMCPDiscoveryHTTP(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request discoveryModelRequest + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Logf("model mock could not decode request: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + toolName, text, err := respond(request) + if err != nil { + t.Logf("model mock rejected request: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + message := map[string]any{"role": "assistant", "content": text} + finishReason := "stop" + if toolName != "" { + message["tool_calls"] = []any{map[string]any{"id": uuid.NewString(), "type": "function", "function": map[string]any{"name": toolName, "arguments": "{}"}}} + finishReason = "tool_calls" + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{"id": uuid.NewString(), "object": "chat.completion", "model": "test-model", "choices": []any{map[string]any{"index": 0, "message": message, "finish_reason": finishReason}}}); err != nil { + t.Errorf("write model response: %v", err) + } + })) + return reachableModelURL(t, baseURL) +} diff --git a/go/core/test/e2e/session_revert_test.go b/go/core/test/e2e/session_revert_test.go new file mode 100644 index 0000000000..349c011612 --- /dev/null +++ b/go/core/test/e2e/session_revert_test.go @@ -0,0 +1,154 @@ +package e2e_test + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "sync/atomic" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/kagent-dev/mockllm" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/protobuf/proto" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/wait" + ctrlclient "sigs.k8s.io/controller-runtime/pkg/client" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +func TestSessionRevertCrashedRuntime(t *testing.T) { + target := interactionTarget(t) + forEachHarness(t, func(t *testing.T, harness testHarness) { + cfg, err := mockllm.LoadConfigFromFile("mocks/invoke_agent.json", interactionMocks) + require.NoError(t, err) + upstream, err := url.Parse(startMockLLMConfig(t, cfg)) + require.NoError(t, err) + proxy := httputil.NewSingleHostReverseProxy(upstream) + // Only an interaction runs a Session's Actor; quiescence suspends it once + // the turn settles. The gate holds the second turn's model call open so + // the Actor is still running when its worker is deleted. + var blockNext atomic.Bool + blocked := make(chan struct{}, 1) + release := make(chan struct{}) + gate := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && blockNext.CompareAndSwap(true, false) { + blocked <- struct{}{} + select { + case <-r.Context().Done(): + case <-release: + } + return + } + proxy.ServeHTTP(w, r) + })) + require.NoError(t, gate.Listener.Close()) + gate.Listener, err = net.Listen("tcp", "0.0.0.0:0") + require.NoError(t, err) + gate.Start() + t.Cleanup(gate.Close) + fixture := newInteractionFixture(t, harness, target, reachableModelURL(t, gate.URL)) + defer func() { + close(release) + gate.CloseClientConnections() + }() + _, _, task := fixture.send(t, "What is 2+2?") + require.Equal(t, a2a.TaskStateCompleted, task.Status.State) + assertActorSuspended(t, fixture) + getTask := &a2apb.GetTaskRequest{Tenant: fixture.tenant, Id: string(task.ID)} + before, err := fixture.client.GetTask(fixture.ctx, getTask) + require.NoError(t, err) + session, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + namespace, name := session.GetSession().GetAgent().GetRef().GetNamespace(), substrate.ActorName(fixture.sessionID) + blockNext.Store(true) + _, turn := newMessageRequest(t, "What is 2+2?") + turn.Tenant, turn.Message.ContextId = fixture.tenant, fixture.sessionID + turnCtx, stopTurn := context.WithCancel(fixture.ctx) + defer stopTurn() + sendDone := make(chan error, 1) + go func() { _, err := sendMessageWithRetry(turnCtx, fixture.client, turn); sendDone <- err }() + select { + case <-blocked: + case <-time.After(time.Minute): + t.Fatal("second turn never reached the blocking model") + } + running, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) + require.NoError(t, err) + require.NotNil(t, running) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_RUNNING, running.GetStatus().GetState()) + require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running))) + request := &apiv1.RevertSessionRequest{SessionId: fixture.sessionID} + healthy, err := fixture.sessions.RevertSession(fixture.ctx, request) + require.NoError(t, err, "a healthy runtime has nothing to recover") + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, healthy.Session.State) + assignment := running.GetStatus().GetWorkerAssignment() + require.Equal(t, namespace, assignment.GetWorkerNamespace()) + require.NotEmpty(t, assignment.GetWorkerPod()) + for token := ""; ; { + page, err := fixture.system.ListSubstrateActors(fixture.ctx, &apiv1.ListSubstrateActorsRequest{ + Page: &apiv1.PageRequest{Limit: 100, PageToken: token}, + }) + require.NoError(t, err) + require.Empty(t, page.GetAteApiError()) + for _, actor := range page.GetActors() { + other := actor.GetStatus().GetWorkerAssignment() + if actor.GetMetadata().GetUid() != running.GetMetadata().GetUid() { + require.False(t, other.GetWorkerNamespace() == assignment.GetWorkerNamespace() && other.GetWorkerPod() == assignment.GetWorkerPod(), + "refuse to delete a worker shared with another Actor") + } + } + token = page.GetPage().GetNextPageToken() + if token == "" { + break + } + } + kube := interactionKubeClient(t) + pod := &corev1.Pod{} + require.NoError(t, kube.Get(fixture.ctx, ctrlclient.ObjectKey{ + Namespace: assignment.GetWorkerNamespace(), Name: assignment.GetWorkerPod(), + }, pod)) + require.NotEmpty(t, pod.UID) + require.NoError(t, kube.Delete(fixture.ctx, pod, &ctrlclient.DeleteOptions{ + GracePeriodSeconds: new(int64(0)), Preconditions: &metav1.Preconditions{UID: &pod.UID}, + })) + err = wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + actor, err := findSubstrateActor(ctx, fixture.system, namespace, name) + return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED, err + }) + require.NoError(t, err) + reverted, err := fixture.sessions.RevertSession(fixture.ctx, request) + 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) + 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))) + retained, err := fixture.client.GetTask(fixture.ctx, getTask) + require.NoError(t, err) + require.True(t, proto.Equal(before, retained), "recovery preserves the completed task and its history") + stopTurn() + select { + case <-sendDone: + case <-time.After(10 * time.Second): + t.Fatal("interrupted caller did not stop") + } + _, _, next := fixture.send(t, "What is 2+2?") + require.Equal(t, a2a.TaskStateCompleted, next.Status.State) + require.Contains(t, taskText(next), "4") + }) +} diff --git a/go/core/test/upgrade/config_test.go b/go/core/test/upgrade/config_test.go index a08f391597..db5ee83d14 100644 --- a/go/core/test/upgrade/config_test.go +++ b/go/core/test/upgrade/config_test.go @@ -1,15 +1,62 @@ package upgrade import ( + "encoding/json" "os" "path/filepath" + "strings" "testing" "github.com/stretchr/testify/require" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" + "github.com/kagent-dev/kagent/go/api/v1alpha3" kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" ) +func TestInvokeHarnesses(t *testing.T) { + manifest, err := os.ReadFile("../e2e/manifests/lifecycle.yaml.tmpl") + require.NoError(t, err) + for _, scenario := range []struct { + name string + repository string + }{ + {"Given a released controller, when preparing invocation fixtures, assert that they use its published runtimes", "ghcr.io/kagent-dev/kagent"}, + {"Given an upgraded controller, when preparing invocation fixtures, assert that they use its locally built runtimes", "localhost:5001/kagent-dev/kagent"}, + } { + t.Run(scenario.name, func(t *testing.T) { + images := &configv1.HarnessImagesConfig{ + Kagent: scenario.repository + "/golang-adk@sha256:" + strings.Repeat("a", 64), + Claude: scenario.repository + "/claude-harness@sha256:" + strings.Repeat("b", 64), + Codex: scenario.repository + "/codex-harness@sha256:" + strings.Repeat("c", 64), + } + encoded, err := invokeHarnesses(manifest, images) + require.NoError(t, err) + var fixtures v1alpha3.HarnessList + require.NoError(t, json.Unmarshal(encoded, &fixtures)) + require.Equal(t, "List", fixtures.Kind) + require.Len(t, fixtures.Items, 4) + want := map[string]string{ + "kagent": images.Kagent, "byo-adk-e2e": images.Kagent, + "claude-e2e": images.Claude, "codex-e2e": images.Codex, + } + for _, fixture := range fixtures.Items { + require.Equal(t, "Harness", fixture.Kind) + require.Contains(t, want, fixture.Name) + require.NotNil(t, fixture.Spec.Workload.Image) + require.Equal(t, want[fixture.Name], *fixture.Spec.Workload.Image) + require.Equal(t, "kagent-default", fixture.Spec.Substrate.WorkerPoolRef.Name) + delete(want, fixture.Name) + } + require.Empty(t, want) + }) + } + t.Run("Given no runtime catalog, when preparing invocation fixtures, assert that setup reports the missing images", func(t *testing.T) { + _, err := invokeHarnesses(manifest, nil) + require.ErrorContains(t, err, "installed controller must provide kagent, claude, and codex runtime images") + }) +} + func TestLoadUpgradeEnv(t *testing.T) { repoRoot := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(repoRoot, "Makefile"), nil, 0o600)) diff --git a/go/core/test/upgrade/invoke_e2e_test.go b/go/core/test/upgrade/invoke_e2e_test.go index bedff1dee3..1d1a55627a 100644 --- a/go/core/test/upgrade/invoke_e2e_test.go +++ b/go/core/test/upgrade/invoke_e2e_test.go @@ -1,8 +1,11 @@ package upgrade import ( + "bytes" "context" + "encoding/json" "fmt" + "io" "os" "os/exec" "path/filepath" @@ -11,7 +14,12 @@ import ( "time" "github.com/stretchr/testify/require" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + kubeyaml "k8s.io/apimachinery/pkg/util/yaml" + "sigs.k8s.io/yaml" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" ) @@ -52,6 +60,7 @@ func runInvokeE2E(t *testing.T, env upgradeEnv, treeGoDir, label string) { t.Helper() requireInvokeEnvironment(t, treeGoDir, label) + installInvokeHarnesses(t, env, treeGoDir) port, stop := startPortForward(t, env, controllerServiceName, controllerAPIPort) defer stop() @@ -68,6 +77,60 @@ func runInvokeE2E(t *testing.T, env upgradeEnv, treeGoDir, label string) { "[%s] selected invoke test did not run:\n%s", label, string(out)) } +// Reapply the matching release's fixtures at each invocation stage so upgrades +// and application rollbacks exercise that installation's runtime images too. +func installInvokeHarnesses(t *testing.T, env upgradeEnv, treeGoDir string) { + t.Helper() + + manifest, err := os.ReadFile(filepath.Join(treeGoDir, "core/test/e2e/manifests/lifecycle.yaml.tmpl")) + require.NoError(t, err) + controllerYAML := kubectl(t, env, time.Minute, "get", "configmap", controllerServiceName, + "-n", env.namespace, "-o", `jsonpath={.data.controller\.yaml}`) + var config struct { + HarnessImages *configv1.HarnessImagesConfig `json:"harnessImages"` + } + require.NoError(t, yaml.Unmarshal([]byte(controllerYAML), &config)) + fixtures, err := invokeHarnesses(manifest, config.HarnessImages) + require.NoError(t, err) + path := filepath.Join(t.TempDir(), "harnesses.json") + require.NoError(t, os.WriteFile(path, fixtures, 0o600)) + kubectl(t, env, time.Minute, "apply", "-f", path) +} + +func invokeHarnesses(manifest []byte, images *configv1.HarnessImagesConfig) ([]byte, error) { + if images.GetKagent() == "" || images.GetClaude() == "" || images.GetCodex() == "" { + return nil, fmt.Errorf("installed controller must provide kagent, claude, and codex runtime images") + } + rendered := strings.NewReplacer( + "${KAGENT_E2E_RUNTIME_IMAGE}", images.GetKagent(), + "${KAGENT_E2E_CLAUDE_IMAGE}", images.GetClaude(), + "${KAGENT_E2E_CODEX_IMAGE}", images.GetCodex(), + ).Replace(string(manifest)) + decoder := kubeyaml.NewYAMLOrJSONDecoder(bytes.NewBufferString(rendered), 4096) + fixtures := metav1.List{TypeMeta: metav1.TypeMeta{APIVersion: "v1", Kind: "List"}} + for { + var resource runtime.RawExtension + if err := decoder.Decode(&resource); err != nil { + if err == io.EOF { + break + } + return nil, fmt.Errorf("decode lifecycle fixture: %w", err) + } + var metadata metav1.PartialObjectMetadata + if err := json.Unmarshal(resource.Raw, &metadata); err != nil { + return nil, fmt.Errorf("decode lifecycle fixture metadata: %w", err) + } + if metadata.Kind != "Harness" { + continue + } + switch metadata.Name { + case "kagent", "codex-e2e", "claude-e2e", "byo-adk-e2e": + fixtures.Items = append(fixtures.Items, resource) + } + } + return json.Marshal(fixtures) +} + func upgradeE2EEnv(t *testing.T, env upgradeEnv, port int) []string { t.Helper() diff --git a/go/core/test/upgrade/roundtrip_test.go b/go/core/test/upgrade/roundtrip_test.go index 7b1e043d79..28d56e7ba5 100644 --- a/go/core/test/upgrade/roundtrip_test.go +++ b/go/core/test/upgrade/roundtrip_test.go @@ -39,12 +39,20 @@ const ( postgresRole = "kagent_owner" ) -// pgTrackVersion returns the current applied migration version. +// pgTrackVersion returns the current applied migration version, or zero when +// the migration source is disabled and its tracking table does not exist. func pgTrackVersion(t *testing.T, env upgradeEnv, table string) int { t.Helper() + // PostgreSQL resolves table references before evaluating CASE, so an absent + // optional tracking table must be checked in a separate statement. + exists := pgQuery(t, env, fmt.Sprintf("SELECT to_regclass(%s) IS NOT NULL", pgQuote(kagentSchema+"."+table))) + if exists == "f" { + return 0 + } + raw := pgQuery(t, env, fmt.Sprintf( - "SELECT CASE WHEN to_regclass('%[1]s.%[2]s') IS NULL THEN 0 ELSE (SELECT COALESCE(MAX(version_id), 0) FROM %[1]s.%[2]s WHERE is_applied) END", + "SELECT COALESCE(MAX(version_id), 0) FROM %s.%s WHERE is_applied", kagentSchema, table)) return parseInt(t, raw, table+" version") } diff --git a/go/core/test/upgrade/upgrade_test.go b/go/core/test/upgrade/upgrade_test.go index df6a40c03c..4cfd1c7864 100644 --- a/go/core/test/upgrade/upgrade_test.go +++ b/go/core/test/upgrade/upgrade_test.go @@ -430,7 +430,9 @@ func pgMigrationState(t *testing.T, env upgradeEnv) postgresMigrationState { // pgMigrationStateE is the error-returning core of pgMigrationState, for use // inside require.Eventually conditions (see pgQueryE). func pgMigrationStateE(t *testing.T, env upgradeEnv) (postgresMigrationState, error) { - raw, err := pgQueryE(t, env, fmt.Sprintf("SELECT CASE WHEN to_regclass('%[1]s.schema_migrations') IS NULL THEN 0 ELSE (SELECT COALESCE(MAX(version_id), 0) FROM %[1]s.schema_migrations WHERE is_applied) END", kagentSchema)) + // Core migrations are required. Before the tracking table exists, return the + // query error so readiness polling retries until the controller creates it. + raw, err := pgQueryE(t, env, fmt.Sprintf("SELECT COALESCE(MAX(version_id), 0) FROM %s.schema_migrations WHERE is_applied", kagentSchema)) if err != nil { return postgresMigrationState{}, err } diff --git a/go/pkg/logging/logging.go b/go/pkg/logging/logging.go index da8d8794a9..3fef7f7baa 100644 --- a/go/pkg/logging/logging.go +++ b/go/pkg/logging/logging.go @@ -56,11 +56,6 @@ func FromContext(ctx context.Context) *slog.Logger { return slog.New(logr.ToSlogHandler(logger)) } -// FromLogr adapts a logger supplied by controller-runtime. -func FromLogr(logger logr.Logger) *slog.Logger { - return slog.New(logr.ToSlogHandler(logger)) -} - // AsLogr adapts logger for dependencies which still require logr. func AsLogr(logger *slog.Logger) logr.Logger { return logr.FromSlogHandler(logger.Handler()) diff --git a/proto/kagent/api/v1/runtime.proto b/proto/kagent/api/v1/runtime.proto index 1b22ea1feb..312777d4ff 100644 --- a/proto/kagent/api/v1/runtime.proto +++ b/proto/kagent/api/v1/runtime.proto @@ -25,4 +25,7 @@ enum RuntimeOperation { RUNTIME_OPERATION_DELETE = 4; // No lifecycle operation is active. RUNTIME_OPERATION_NONE = 5; + // Restores a crashed runtime to its last durable snapshot. Only a confirmed + // crash admits it; the Session stays FAILED until it finishes. + RUNTIME_OPERATION_REVERT = 6; } diff --git a/proto/kagent/api/v1/sessions.proto b/proto/kagent/api/v1/sessions.proto index e000d269d5..93651ce5c0 100644 --- a/proto/kagent/api/v1/sessions.proto +++ b/proto/kagent/api/v1/sessions.proto @@ -20,6 +20,8 @@ service SessionService { rpc GetSession(GetSessionRequest) returns (GetSessionResponse); rpc ListSessions(ListSessionsRequest) returns (ListSessionsResponse); rpc UpdateSessionName(UpdateSessionNameRequest) returns (UpdateSessionNameResponse); + // Restores a crashed runtime to its last snapshot without replaying its turn. + rpc RevertSession(RevertSessionRequest) returns (RevertSessionResponse); rpc DeleteSession(DeleteSessionRequest) returns (DeleteSessionResponse); rpc CreateSessionShare(CreateSessionShareRequest) returns (CreateSessionShareResponse); rpc ListSessionShares(ListSessionSharesRequest) returns (ListSessionSharesResponse); @@ -107,6 +109,14 @@ message UpdateSessionNameResponse { Session session = 1; } +message RevertSessionRequest { + string session_id = 1 [(buf.validate.field).string.uuid = true]; +} + +message RevertSessionResponse { + Session session = 1; +} + message DeleteSessionRequest { string session_id = 1 [(buf.validate.field).string.uuid = true]; // Terminate ongoing turns instead of rejecting deletion while they are active. diff --git a/ui/src/api/domain/agentInstances.ts b/ui/src/api/domain/agentInstances.ts index 68568298fe..6c85dde927 100644 --- a/ui/src/api/domain/agentInstances.ts +++ b/ui/src/api/domain/agentInstances.ts @@ -56,6 +56,7 @@ export type AgentInstanceOperation = | "suspend" | "resume" | "delete" + | "revert" | "unknown"; /** diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index 202ed7cdb4..2ae8b2d8bc 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -565,6 +565,7 @@ const INSTANCE_OPERATION_BY_ENUM: Record< [RuntimeOperation.SUSPEND]: "suspend", [RuntimeOperation.RESUME]: "resume", [RuntimeOperation.DELETE]: "delete", + [RuntimeOperation.REVERT]: "revert", }; /** diff --git a/ui/src/components/agent-instances/instanceLabels.ts b/ui/src/components/agent-instances/instanceLabels.ts index 200e906f03..c1bcb306e8 100644 --- a/ui/src/components/agent-instances/instanceLabels.ts +++ b/ui/src/components/agent-instances/instanceLabels.ts @@ -109,6 +109,7 @@ const OPERATIONS: Record = { suspend: { label: "Suspending", tone: "processing", inProgress: true }, resume: { label: "Resuming", tone: "processing", inProgress: true }, delete: { label: "Deleting", tone: "warning", inProgress: true }, + revert: { label: "Reverting", tone: "warning", inProgress: true }, unknown: { label: "Unrecognised operation", tone: "warning", inProgress: true }, }; diff --git a/ui/src/generated/kagent/api/v1/runtime_pb.ts b/ui/src/generated/kagent/api/v1/runtime_pb.ts index 5177311c97..901a61509d 100644 --- a/ui/src/generated/kagent/api/v1/runtime_pb.ts +++ b/ui/src/generated/kagent/api/v1/runtime_pb.ts @@ -9,7 +9,7 @@ import { enumDesc, fileDesc } from "@bufbuild/protobuf/codegenv2"; * Describes the file kagent/api/v1/runtime.proto. */ export const file_kagent_api_v1_runtime: GenFile = /*@__PURE__*/ - fileDesc("ChtrYWdlbnQvYXBpL3YxL3J1bnRpbWUucHJvdG8SDWthZ2VudC5hcGkudjEq0AEKDFJ1bnRpbWVTdGF0ZRIdChlSVU5USU1FX1NUQVRFX1VOU1BFQ0lGSUVEEAASGgoWUlVOVElNRV9TVEFURV9DUkVBVElORxABEhcKE1JVTlRJTUVfU1RBVEVfUkVBRFkQAhIbChdSVU5USU1FX1NUQVRFX1NVU1BFTkRFRBADEhgKFFJVTlRJTUVfU1RBVEVfRkFJTEVEEAQSGgoWUlVOVElNRV9TVEFURV9ERUxFVElORxAFEhkKFVJVTlRJTUVfU1RBVEVfREVMRVRFRBAGKsoBChBSdW50aW1lT3BlcmF0aW9uEiEKHVJVTlRJTUVfT1BFUkFUSU9OX1VOU1BFQ0lGSUVEEAASHAoYUlVOVElNRV9PUEVSQVRJT05fQ1JFQVRFEAESHQoZUlVOVElNRV9PUEVSQVRJT05fU1VTUEVORBACEhwKGFJVTlRJTUVfT1BFUkFUSU9OX1JFU1VNRRADEhwKGFJVTlRJTUVfT1BFUkFUSU9OX0RFTEVURRAEEhoKFlJVTlRJTUVfT1BFUkFUSU9OX05PTkUQBUI9WjtnaXRodWIuY29tL2thZ2VudC1kZXYva2FnZW50L2dvL2FwaS9nZW4va2FnZW50L2FwaS92MTthcGl2MWIGcHJvdG8z"); + fileDesc("ChtrYWdlbnQvYXBpL3YxL3J1bnRpbWUucHJvdG8SDWthZ2VudC5hcGkudjEq0AEKDFJ1bnRpbWVTdGF0ZRIdChlSVU5USU1FX1NUQVRFX1VOU1BFQ0lGSUVEEAASGgoWUlVOVElNRV9TVEFURV9DUkVBVElORxABEhcKE1JVTlRJTUVfU1RBVEVfUkVBRFkQAhIbChdSVU5USU1FX1NUQVRFX1NVU1BFTkRFRBADEhgKFFJVTlRJTUVfU1RBVEVfRkFJTEVEEAQSGgoWUlVOVElNRV9TVEFURV9ERUxFVElORxAFEhkKFVJVTlRJTUVfU1RBVEVfREVMRVRFRBAGKugBChBSdW50aW1lT3BlcmF0aW9uEiEKHVJVTlRJTUVfT1BFUkFUSU9OX1VOU1BFQ0lGSUVEEAASHAoYUlVOVElNRV9PUEVSQVRJT05fQ1JFQVRFEAESHQoZUlVOVElNRV9PUEVSQVRJT05fU1VTUEVORBACEhwKGFJVTlRJTUVfT1BFUkFUSU9OX1JFU1VNRRADEhwKGFJVTlRJTUVfT1BFUkFUSU9OX0RFTEVURRAEEhoKFlJVTlRJTUVfT1BFUkFUSU9OX05PTkUQBRIcChhSVU5USU1FX09QRVJBVElPTl9SRVZFUlQQBkI9WjtnaXRodWIuY29tL2thZ2VudC1kZXYva2FnZW50L2dvL2FwaS9nZW4va2FnZW50L2FwaS92MTthcGl2MWIGcHJvdG8z"); /** * RuntimeState is shared by Sessions and Sandboxes. Its names are also the @@ -97,6 +97,14 @@ export enum RuntimeOperation { * @generated from enum value: RUNTIME_OPERATION_NONE = 5; */ NONE = 5, + + /** + * Restores a crashed runtime to its last durable snapshot. Only a confirmed + * crash admits it; the Session stays FAILED until it finishes. + * + * @generated from enum value: RUNTIME_OPERATION_REVERT = 6; + */ + REVERT = 6, } /** diff --git a/ui/src/generated/kagent/api/v1/sessions_pb.ts b/ui/src/generated/kagent/api/v1/sessions_pb.ts index 9c7abc8506..fdc14c14d4 100644 --- a/ui/src/generated/kagent/api/v1/sessions_pb.ts +++ b/ui/src/generated/kagent/api/v1/sessions_pb.ts @@ -17,7 +17,7 @@ import type { Message } from "@bufbuild/protobuf"; * Describes the file kagent/api/v1/sessions.proto. */ export const file_kagent_api_v1_sessions: GenFile = /*@__PURE__*/ - fileDesc("ChxrYWdlbnQvYXBpL3YxL3Nlc3Npb25zLnByb3RvEg1rYWdlbnQuYXBpLnYxIioKB0ZhaWx1cmUSDgoGcmVhc29uGAEgASgJEg8KB21lc3NhZ2UYAiABKAkikgMKB1Nlc3Npb24SCgoCaWQYASABKAkSDwoHY3JlYXRvchgCIAEoCRItCgVhZ2VudBgPIAEoCzIeLmthZ2VudC5hcGkudjEuQWdlbnRQcm92ZW5hbmNlEhkKEXByZXBhcmVkX3JldmlzaW9uGAUgASgJEhUKDWEyYV9hdXRob3JpdHkYBiABKAkSKgoFc3RhdGUYByABKA4yGy5rYWdlbnQuYXBpLnYxLlJ1bnRpbWVTdGF0ZRIyCglvcGVyYXRpb24YCCABKA4yHy5rYWdlbnQuYXBpLnYxLlJ1bnRpbWVPcGVyYXRpb24SJwoHZmFpbHVyZRgJIAEoCzIWLmthZ2VudC5hcGkudjEuRmFpbHVyZRIuCgpjcmVhdGVkX2F0GAogASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIuCgp1cGRhdGVkX2F0GAsgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIMCgRuYW1lGA0gASgJEhIKCmNvbnRleHRfaWQYDiABKAkiTQoPQWdlbnRQcm92ZW5hbmNlEi0KA3JlZhgBIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2USCwoDdWlkGAIgASgJIrwBChRDcmVhdGVTZXNzaW9uUmVxdWVzdBI3CgVhZ2VudBgFIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2VCBrpIA8gBARIeCgpyZXF1ZXN0X2lkGAMgASgJQgq6SAdyBRABGIABEksKBG5hbWUYBCABKAlCPbpIOnI4GMgBMjNeKD86JHxbXlxwe1p9XHB7Q2N9XSg/OlteXHB7Q2N9XSpbXlxwe1p9XHB7Q2N9XSk/KSQiQAoVQ3JlYXRlU2Vzc2lvblJlc3BvbnNlEicKB3Nlc3Npb24YASABKAsyFi5rYWdlbnQuYXBpLnYxLlNlc3Npb24iMQoRR2V0U2Vzc2lvblJlcXVlc3QSHAoKc2Vzc2lvbl9pZBgBIAEoCUIIukgFcgOwAQEiPQoSR2V0U2Vzc2lvblJlc3BvbnNlEicKB3Nlc3Npb24YASABKAsyFi5rYWdlbnQuYXBpLnYxLlNlc3Npb24ihgEKE0xpc3RTZXNzaW9uc1JlcXVlc3QSFAoMYWxsX2NyZWF0b3JzGAIgASgIEigKBHBhZ2UYAyABKAsyGi5rYWdlbnQuYXBpLnYxLlBhZ2VSZXF1ZXN0Ei8KBWFnZW50GAYgASgLMiAua2FnZW50LmFwaS52MS5SZXNvdXJjZVJlZmVyZW5jZSJrChRMaXN0U2Vzc2lvbnNSZXNwb25zZRIoCghzZXNzaW9ucxgBIAMoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbhIpCgRwYWdlGAIgASgLMhsua2FnZW50LmFwaS52MS5QYWdlUmVzcG9uc2UihQEKGFVwZGF0ZVNlc3Npb25OYW1lUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARJLCgRuYW1lGAIgASgJQj26SDpyOBjIATIzXig/OiR8W15ccHtafVxwe0NjfV0oPzpbXlxwe0NjfV0qW15ccHtafVxwe0NjfV0pPykkIkQKGVVwZGF0ZVNlc3Npb25OYW1lUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiJDChREZWxldGVTZXNzaW9uUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARINCgVmb3JjZRgCIAEoCCJAChVEZWxldGVTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiKOAQoMU2Vzc2lvblNoYXJlEgoKAmlkGAEgASgJEhIKCnNlc3Npb25faWQYAiABKAkSLgoKY3JlYXRlZF9hdBgDIAEoCzIaLmdvb2dsZS5wcm90b2J1Zi5UaW1lc3RhbXASLgoKZXhwaXJlc19hdBgEIAEoCzIaLmdvb2dsZS5wcm90b2J1Zi5UaW1lc3RhbXAiawoZQ3JlYXRlU2Vzc2lvblNoYXJlUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARIwCgN0dGwYAiABKAsyGS5nb29nbGUucHJvdG9idWYuRHVyYXRpb25CCLpIBaoBAioAIlcKGkNyZWF0ZVNlc3Npb25TaGFyZVJlc3BvbnNlEioKBXNoYXJlGAEgASgLMhsua2FnZW50LmFwaS52MS5TZXNzaW9uU2hhcmUSDQoFdG9rZW4YAiABKAkiYgoYTGlzdFNlc3Npb25TaGFyZXNSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEigKBHBhZ2UYAiABKAsyGi5rYWdlbnQuYXBpLnYxLlBhZ2VSZXF1ZXN0InMKGUxpc3RTZXNzaW9uU2hhcmVzUmVzcG9uc2USKwoGc2hhcmVzGAEgAygLMhsua2FnZW50LmFwaS52MS5TZXNzaW9uU2hhcmUSKQoEcGFnZRgCIAEoCzIbLmthZ2VudC5hcGkudjEuUGFnZVJlc3BvbnNlIjcKGVJldm9rZVNlc3Npb25TaGFyZVJlcXVlc3QSGgoIc2hhcmVfaWQYASABKAlCCLpIBXIDsAEBIhwKGlJldm9rZVNlc3Npb25TaGFyZVJlc3BvbnNlMpoGCg5TZXNzaW9uU2VydmljZRJaCg1DcmVhdGVTZXNzaW9uEiMua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uUmVxdWVzdBokLmthZ2VudC5hcGkudjEuQ3JlYXRlU2Vzc2lvblJlc3BvbnNlElEKCkdldFNlc3Npb24SIC5rYWdlbnQuYXBpLnYxLkdldFNlc3Npb25SZXF1ZXN0GiEua2FnZW50LmFwaS52MS5HZXRTZXNzaW9uUmVzcG9uc2USVwoMTGlzdFNlc3Npb25zEiIua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvbnNSZXF1ZXN0GiMua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvbnNSZXNwb25zZRJmChFVcGRhdGVTZXNzaW9uTmFtZRInLmthZ2VudC5hcGkudjEuVXBkYXRlU2Vzc2lvbk5hbWVSZXF1ZXN0Gigua2FnZW50LmFwaS52MS5VcGRhdGVTZXNzaW9uTmFtZVJlc3BvbnNlEloKDURlbGV0ZVNlc3Npb24SIy5rYWdlbnQuYXBpLnYxLkRlbGV0ZVNlc3Npb25SZXF1ZXN0GiQua2FnZW50LmFwaS52MS5EZWxldGVTZXNzaW9uUmVzcG9uc2USaQoSQ3JlYXRlU2Vzc2lvblNoYXJlEigua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uU2hhcmVSZXF1ZXN0Gikua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uU2hhcmVSZXNwb25zZRJmChFMaXN0U2Vzc2lvblNoYXJlcxInLmthZ2VudC5hcGkudjEuTGlzdFNlc3Npb25TaGFyZXNSZXF1ZXN0Gigua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvblNoYXJlc1Jlc3BvbnNlEmkKElJldm9rZVNlc3Npb25TaGFyZRIoLmthZ2VudC5hcGkudjEuUmV2b2tlU2Vzc2lvblNoYXJlUmVxdWVzdBopLmthZ2VudC5hcGkudjEuUmV2b2tlU2Vzc2lvblNoYXJlUmVzcG9uc2VCPVo7Z2l0aHViLmNvbS9rYWdlbnQtZGV2L2thZ2VudC9nby9hcGkvZ2VuL2thZ2VudC9hcGkvdjE7YXBpdjFiBnByb3RvMw", [file_buf_validate_validate, file_google_protobuf_duration, file_google_protobuf_timestamp, file_kagent_api_v1_common, file_kagent_api_v1_runtime]); + fileDesc("ChxrYWdlbnQvYXBpL3YxL3Nlc3Npb25zLnByb3RvEg1rYWdlbnQuYXBpLnYxIioKB0ZhaWx1cmUSDgoGcmVhc29uGAEgASgJEg8KB21lc3NhZ2UYAiABKAkikgMKB1Nlc3Npb24SCgoCaWQYASABKAkSDwoHY3JlYXRvchgCIAEoCRItCgVhZ2VudBgPIAEoCzIeLmthZ2VudC5hcGkudjEuQWdlbnRQcm92ZW5hbmNlEhkKEXByZXBhcmVkX3JldmlzaW9uGAUgASgJEhUKDWEyYV9hdXRob3JpdHkYBiABKAkSKgoFc3RhdGUYByABKA4yGy5rYWdlbnQuYXBpLnYxLlJ1bnRpbWVTdGF0ZRIyCglvcGVyYXRpb24YCCABKA4yHy5rYWdlbnQuYXBpLnYxLlJ1bnRpbWVPcGVyYXRpb24SJwoHZmFpbHVyZRgJIAEoCzIWLmthZ2VudC5hcGkudjEuRmFpbHVyZRIuCgpjcmVhdGVkX2F0GAogASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIuCgp1cGRhdGVkX2F0GAsgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIMCgRuYW1lGA0gASgJEhIKCmNvbnRleHRfaWQYDiABKAkiTQoPQWdlbnRQcm92ZW5hbmNlEi0KA3JlZhgBIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2USCwoDdWlkGAIgASgJIrwBChRDcmVhdGVTZXNzaW9uUmVxdWVzdBI3CgVhZ2VudBgFIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2VCBrpIA8gBARIeCgpyZXF1ZXN0X2lkGAMgASgJQgq6SAdyBRABGIABEksKBG5hbWUYBCABKAlCPbpIOnI4GMgBMjNeKD86JHxbXlxwe1p9XHB7Q2N9XSg/OlteXHB7Q2N9XSpbXlxwe1p9XHB7Q2N9XSk/KSQiQAoVQ3JlYXRlU2Vzc2lvblJlc3BvbnNlEicKB3Nlc3Npb24YASABKAsyFi5rYWdlbnQuYXBpLnYxLlNlc3Npb24iMQoRR2V0U2Vzc2lvblJlcXVlc3QSHAoKc2Vzc2lvbl9pZBgBIAEoCUIIukgFcgOwAQEiPQoSR2V0U2Vzc2lvblJlc3BvbnNlEicKB3Nlc3Npb24YASABKAsyFi5rYWdlbnQuYXBpLnYxLlNlc3Npb24ihgEKE0xpc3RTZXNzaW9uc1JlcXVlc3QSFAoMYWxsX2NyZWF0b3JzGAIgASgIEigKBHBhZ2UYAyABKAsyGi5rYWdlbnQuYXBpLnYxLlBhZ2VSZXF1ZXN0Ei8KBWFnZW50GAYgASgLMiAua2FnZW50LmFwaS52MS5SZXNvdXJjZVJlZmVyZW5jZSJrChRMaXN0U2Vzc2lvbnNSZXNwb25zZRIoCghzZXNzaW9ucxgBIAMoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbhIpCgRwYWdlGAIgASgLMhsua2FnZW50LmFwaS52MS5QYWdlUmVzcG9uc2UihQEKGFVwZGF0ZVNlc3Npb25OYW1lUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARJLCgRuYW1lGAIgASgJQj26SDpyOBjIATIzXig/OiR8W15ccHtafVxwe0NjfV0oPzpbXlxwe0NjfV0qW15ccHtafVxwe0NjfV0pPykkIkQKGVVwZGF0ZVNlc3Npb25OYW1lUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiI0ChRSZXZlcnRTZXNzaW9uUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABASJAChVSZXZlcnRTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiJDChREZWxldGVTZXNzaW9uUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARINCgVmb3JjZRgCIAEoCCJAChVEZWxldGVTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiKOAQoMU2Vzc2lvblNoYXJlEgoKAmlkGAEgASgJEhIKCnNlc3Npb25faWQYAiABKAkSLgoKY3JlYXRlZF9hdBgDIAEoCzIaLmdvb2dsZS5wcm90b2J1Zi5UaW1lc3RhbXASLgoKZXhwaXJlc19hdBgEIAEoCzIaLmdvb2dsZS5wcm90b2J1Zi5UaW1lc3RhbXAiawoZQ3JlYXRlU2Vzc2lvblNoYXJlUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABARIwCgN0dGwYAiABKAsyGS5nb29nbGUucHJvdG9idWYuRHVyYXRpb25CCLpIBaoBAioAIlcKGkNyZWF0ZVNlc3Npb25TaGFyZVJlc3BvbnNlEioKBXNoYXJlGAEgASgLMhsua2FnZW50LmFwaS52MS5TZXNzaW9uU2hhcmUSDQoFdG9rZW4YAiABKAkiYgoYTGlzdFNlc3Npb25TaGFyZXNSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEigKBHBhZ2UYAiABKAsyGi5rYWdlbnQuYXBpLnYxLlBhZ2VSZXF1ZXN0InMKGUxpc3RTZXNzaW9uU2hhcmVzUmVzcG9uc2USKwoGc2hhcmVzGAEgAygLMhsua2FnZW50LmFwaS52MS5TZXNzaW9uU2hhcmUSKQoEcGFnZRgCIAEoCzIbLmthZ2VudC5hcGkudjEuUGFnZVJlc3BvbnNlIjcKGVJldm9rZVNlc3Npb25TaGFyZVJlcXVlc3QSGgoIc2hhcmVfaWQYASABKAlCCLpIBXIDsAEBIhwKGlJldm9rZVNlc3Npb25TaGFyZVJlc3BvbnNlMvYGCg5TZXNzaW9uU2VydmljZRJaCg1DcmVhdGVTZXNzaW9uEiMua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uUmVxdWVzdBokLmthZ2VudC5hcGkudjEuQ3JlYXRlU2Vzc2lvblJlc3BvbnNlElEKCkdldFNlc3Npb24SIC5rYWdlbnQuYXBpLnYxLkdldFNlc3Npb25SZXF1ZXN0GiEua2FnZW50LmFwaS52MS5HZXRTZXNzaW9uUmVzcG9uc2USVwoMTGlzdFNlc3Npb25zEiIua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvbnNSZXF1ZXN0GiMua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvbnNSZXNwb25zZRJmChFVcGRhdGVTZXNzaW9uTmFtZRInLmthZ2VudC5hcGkudjEuVXBkYXRlU2Vzc2lvbk5hbWVSZXF1ZXN0Gigua2FnZW50LmFwaS52MS5VcGRhdGVTZXNzaW9uTmFtZVJlc3BvbnNlEloKDVJldmVydFNlc3Npb24SIy5rYWdlbnQuYXBpLnYxLlJldmVydFNlc3Npb25SZXF1ZXN0GiQua2FnZW50LmFwaS52MS5SZXZlcnRTZXNzaW9uUmVzcG9uc2USWgoNRGVsZXRlU2Vzc2lvbhIjLmthZ2VudC5hcGkudjEuRGVsZXRlU2Vzc2lvblJlcXVlc3QaJC5rYWdlbnQuYXBpLnYxLkRlbGV0ZVNlc3Npb25SZXNwb25zZRJpChJDcmVhdGVTZXNzaW9uU2hhcmUSKC5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25TaGFyZVJlcXVlc3QaKS5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25TaGFyZVJlc3BvbnNlEmYKEUxpc3RTZXNzaW9uU2hhcmVzEicua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvblNoYXJlc1JlcXVlc3QaKC5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uU2hhcmVzUmVzcG9uc2USaQoSUmV2b2tlU2Vzc2lvblNoYXJlEigua2FnZW50LmFwaS52MS5SZXZva2VTZXNzaW9uU2hhcmVSZXF1ZXN0Gikua2FnZW50LmFwaS52MS5SZXZva2VTZXNzaW9uU2hhcmVSZXNwb25zZUI9WjtnaXRodWIuY29tL2thZ2VudC1kZXYva2FnZW50L2dvL2FwaS9nZW4va2FnZW50L2FwaS92MTthcGl2MWIGcHJvdG8z", [file_buf_validate_validate, file_google_protobuf_duration, file_google_protobuf_timestamp, file_kagent_api_v1_common, file_kagent_api_v1_runtime]); /** * @generated from message kagent.api.v1.Failure @@ -317,6 +317,40 @@ export type UpdateSessionNameResponse = Message<"kagent.api.v1.UpdateSessionName export const UpdateSessionNameResponseSchema: GenMessage = /*@__PURE__*/ messageDesc(file_kagent_api_v1_sessions, 10); +/** + * @generated from message kagent.api.v1.RevertSessionRequest + */ +export type RevertSessionRequest = Message<"kagent.api.v1.RevertSessionRequest"> & { + /** + * @generated from field: string session_id = 1; + */ + sessionId: string; +}; + +/** + * Describes the message kagent.api.v1.RevertSessionRequest. + * Use `create(RevertSessionRequestSchema)` to create a new message. + */ +export const RevertSessionRequestSchema: GenMessage = /*@__PURE__*/ + messageDesc(file_kagent_api_v1_sessions, 11); + +/** + * @generated from message kagent.api.v1.RevertSessionResponse + */ +export type RevertSessionResponse = Message<"kagent.api.v1.RevertSessionResponse"> & { + /** + * @generated from field: kagent.api.v1.Session session = 1; + */ + session?: Session | undefined; +}; + +/** + * Describes the message kagent.api.v1.RevertSessionResponse. + * Use `create(RevertSessionResponseSchema)` to create a new message. + */ +export const RevertSessionResponseSchema: GenMessage = /*@__PURE__*/ + messageDesc(file_kagent_api_v1_sessions, 12); + /** * @generated from message kagent.api.v1.DeleteSessionRequest */ @@ -340,7 +374,7 @@ export type DeleteSessionRequest = Message<"kagent.api.v1.DeleteSessionRequest"> * Use `create(DeleteSessionRequestSchema)` to create a new message. */ export const DeleteSessionRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 11); + messageDesc(file_kagent_api_v1_sessions, 13); /** * @generated from message kagent.api.v1.DeleteSessionResponse @@ -357,7 +391,7 @@ export type DeleteSessionResponse = Message<"kagent.api.v1.DeleteSessionResponse * Use `create(DeleteSessionResponseSchema)` to create a new message. */ export const DeleteSessionResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 12); + messageDesc(file_kagent_api_v1_sessions, 14); /** * Grants read-only access to a session, including A2A get, list, and subscribe. @@ -393,7 +427,7 @@ export type SessionShare = Message<"kagent.api.v1.SessionShare"> & { * Use `create(SessionShareSchema)` to create a new message. */ export const SessionShareSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 13); + messageDesc(file_kagent_api_v1_sessions, 15); /** * @generated from message kagent.api.v1.CreateSessionShareRequest @@ -419,7 +453,7 @@ export type CreateSessionShareRequest = Message<"kagent.api.v1.CreateSessionShar * Use `create(CreateSessionShareRequestSchema)` to create a new message. */ export const CreateSessionShareRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 14); + messageDesc(file_kagent_api_v1_sessions, 16); /** * @generated from message kagent.api.v1.CreateSessionShareResponse @@ -443,7 +477,7 @@ export type CreateSessionShareResponse = Message<"kagent.api.v1.CreateSessionSha * Use `create(CreateSessionShareResponseSchema)` to create a new message. */ export const CreateSessionShareResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 15); + messageDesc(file_kagent_api_v1_sessions, 17); /** * @generated from message kagent.api.v1.ListSessionSharesRequest @@ -465,7 +499,7 @@ export type ListSessionSharesRequest = Message<"kagent.api.v1.ListSessionSharesR * Use `create(ListSessionSharesRequestSchema)` to create a new message. */ export const ListSessionSharesRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 16); + messageDesc(file_kagent_api_v1_sessions, 18); /** * @generated from message kagent.api.v1.ListSessionSharesResponse @@ -487,7 +521,7 @@ export type ListSessionSharesResponse = Message<"kagent.api.v1.ListSessionShares * Use `create(ListSessionSharesResponseSchema)` to create a new message. */ export const ListSessionSharesResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 17); + messageDesc(file_kagent_api_v1_sessions, 19); /** * @generated from message kagent.api.v1.RevokeSessionShareRequest @@ -504,7 +538,7 @@ export type RevokeSessionShareRequest = Message<"kagent.api.v1.RevokeSessionShar * Use `create(RevokeSessionShareRequestSchema)` to create a new message. */ export const RevokeSessionShareRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 18); + messageDesc(file_kagent_api_v1_sessions, 20); /** * @generated from message kagent.api.v1.RevokeSessionShareResponse @@ -517,7 +551,7 @@ export type RevokeSessionShareResponse = Message<"kagent.api.v1.RevokeSessionSha * Use `create(RevokeSessionShareResponseSchema)` to create a new message. */ export const RevokeSessionShareResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 19); + messageDesc(file_kagent_api_v1_sessions, 21); /** * Lifecycle mutations execute inline. After a retryable error, repeat the same @@ -561,6 +595,16 @@ export const SessionService: GenService<{ input: typeof UpdateSessionNameRequestSchema; output: typeof UpdateSessionNameResponseSchema; }, + /** + * Restores a crashed runtime to its last snapshot without replaying its turn. + * + * @generated from rpc kagent.api.v1.SessionService.RevertSession + */ + revertSession: { + methodKind: "unary"; + input: typeof RevertSessionRequestSchema; + output: typeof RevertSessionResponseSchema; + }, /** * @generated from rpc kagent.api.v1.SessionService.DeleteSession */ diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index ea5c593157..a944067b96 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -635,6 +635,7 @@ const PB_OPERATION_BY_NAME: Record = { suspend: RuntimeOperation.SUSPEND, resume: RuntimeOperation.RESUME, delete: RuntimeOperation.DELETE, + revert: RuntimeOperation.REVERT, unknown: RuntimeOperation.NONE, }; diff --git a/ui/src/pages/SubstratePage.tsx b/ui/src/pages/SubstratePage.tsx index d3e8dc5c74..d9f822c075 100644 --- a/ui/src/pages/SubstratePage.tsx +++ b/ui/src/pages/SubstratePage.tsx @@ -692,25 +692,74 @@ function SectionSearch({ const PAGE_SIZE = 25; -function PageAge({ computedAt }: { computedAt?: string }) { +/** + * When the list was read, as a clock time rather than a running age: an age that ticks + * twice a second pulls the eye off the rows while saying nothing new. + * + * While polling, a list with no new read for `STALLED_AFTER_TICKS` intervals adds how + * long it has gone without one, in the warning colour. That is the stall the old age + * showed, now shown only when it is happening. + */ +function PageAge({ computedAt, pollMs }: { computedAt?: string; pollMs?: number }) { const theme = useTheme(); - const age = useDataAge(computedAt); - return {age}; + const stalledMs = useStalledFor(computedAt, pollMs); + const at = readAt(computedAt); + if (stalledMs === undefined) { + return {at}; + } + return ( + + {`${at}, ${elapsed(stalledMs)} ago`} + + ); } -function useDataAge(computedAt: string | undefined): string { - const [now, setNow] = useState(() => Date.now()); +const STALLED_AFTER_TICKS = 3; +// A floor, so a half-second interval does not call one slow read a stall. +const STALLED_FLOOR_MS = 5_000; + +/** + * How long a polled list has gone without a new read, once that is a stall. + * + * Measured on the browser's clock from when `computedAt` last changed, not from + * `computedAt` itself, so a skewed server clock cannot raise it. + */ +function useStalledFor(computedAt: string | undefined, pollMs: number | undefined) { + const [stall, setStall] = useState<{ key?: string; ms: number }>(); + useEffect(() => { - const timer = window.setInterval(() => setNow(Date.now()), 500); + if (pollMs === undefined) return; + const sinceMs = Date.now(); + const limitMs = Math.max(STALLED_AFTER_TICKS * pollMs, STALLED_FLOOR_MS); + const timer = window.setInterval(() => { + const ms = Date.now() - sinceMs; + setStall(ms >= limitMs ? { key: computedAt, ms } : undefined); + }, 1_000); return () => window.clearInterval(timer); - }, []); + }, [computedAt, pollMs]); + // A stall belongs to the read it was measured against; a new read clears it at once. + return pollMs !== undefined && stall !== undefined && stall.key === computedAt + ? stall.ms + : undefined; +} + +function readAt(computedAt: string | undefined): string { if (!computedAt) return ""; - const at = new Date(computedAt).getTime(); - if (Number.isNaN(at)) return ""; - const seconds = Math.max(0, (now - at) / 1000); - if (seconds < 1) return "read just now"; - return `read ${seconds.toFixed(1)}s ago`; + const at = new Date(computedAt); + if (Number.isNaN(at.getTime())) return ""; + return `Read at ${at.toLocaleTimeString([], { + hour: "numeric", + minute: "2-digit", + second: "2-digit", + })}`; +} + +/** `12s`, `3m 05s`. */ +function elapsed(ms: number): string { + const whole = Math.floor(ms / 1_000); + if (whole < 60) return `${whole}s`; + return `${Math.floor(whole / 60)}m ${String(whole % 60).padStart(2, "0")}s`; } function PageWarning({ message, testId }: { message: string; testId: string }) { @@ -1484,7 +1533,7 @@ export function SubstratePage() { "& > [data-testid$='-pages']": { marginTop: 0 }, }} > - + [data-testid$='-pages']": { marginTop: 0 }, }} > - + { + setApiTransport(undefined); + vi.useRealTimers(); +}); + +// The lists answer once, then hang: polling keeps asking and no new read lands. +it("shows how long a polled list has gone without a new read", async () => { + vi.useFakeTimers({ toFake: ["setInterval", "clearInterval", "Date"] }); + const hang = () => new Promise(() => {}); + let actorReads = 0; + setApiTransport(createRouterTransport(({ service }) => service(SystemService, { + listNamespaces: () => ({}), + getSubstrateSummary: () => ({}), + listSubstrateActors: () => + actorReads++ === 0 ? { computedAt: timestampFromDate(new Date()) } : hang(), + listSubstrateWorkers: () => ({ computedAt: timestampFromDate(new Date()) }), + }))); + render( + new Map(), shouldRetryOnError: false }}> + + + ); + + const [actorsAge] = await screen.findAllByText(/^Read at /); + expect(actorsAge).not.toHaveTextContent(/ago$/); + + fireEvent.click(screen.getByTestId("substrate-poll-toggle")); + await act(async () => { + vi.advanceTimersByTime(8_000); + }); + + expect(actorsAge).toHaveTextContent(/^Read at .+, \d+s ago$/); +});