diff --git a/apps/backend/internal/agent/handlers/shell_handlers.go b/apps/backend/internal/agent/handlers/shell_handlers.go index d53454cd1b..e7f7598ecf 100644 --- a/apps/backend/internal/agent/handlers/shell_handlers.go +++ b/apps/backend/internal/agent/handlers/shell_handlers.go @@ -35,10 +35,11 @@ func NewShellHandlers(lifecycleMgr *lifecycle.Manager, scriptService scripts.Scr } func (h *ShellHandlers) resolveScopeID(ctx context.Context, sessionID string) (string, error) { - if _, err := h.lifecycleMgr.GetOrEnsureExecution(ctx, sessionID); err != nil { - return "", fmt.Errorf("failed to ensure execution for scope: %w", err) + envID, err := h.lifecycleMgr.ResolveTaskEnvironmentID(ctx, sessionID) + if err != nil { + return "", fmt.Errorf("failed to resolve task environment for shell scope: %w", err) } - return h.lifecycleMgr.ResolveScopeKey(sessionID), nil + return envID, nil } // RegisterHandlers registers shell handlers with the WebSocket dispatcher diff --git a/apps/backend/internal/agent/lifecycle/execution_store.go b/apps/backend/internal/agent/lifecycle/execution_store.go index e1c0377510..600b30a76d 100644 --- a/apps/backend/internal/agent/lifecycle/execution_store.go +++ b/apps/backend/internal/agent/lifecycle/execution_store.go @@ -97,6 +97,19 @@ func (s *ExecutionStore) GetBySessionID(sessionID string) (*AgentExecution, bool return execution, exists } +// GetByTaskEnvironmentID returns any execution associated with a task environment ID. +func (s *ExecutionStore) GetByTaskEnvironmentID(taskEnvironmentID string) (*AgentExecution, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + + for _, execution := range s.executions { + if execution.TaskEnvironmentID == taskEnvironmentID { + return execution, true + } + } + return nil, false +} + // GetByContainerID returns the agent execution associated with a container ID. func (s *ExecutionStore) GetByContainerID(containerID string) (*AgentExecution, bool) { s.mu.RLock() diff --git a/apps/backend/internal/agent/lifecycle/manager_execution.go b/apps/backend/internal/agent/lifecycle/manager_execution.go index fd43f63a79..96f5cf6d78 100644 --- a/apps/backend/internal/agent/lifecycle/manager_execution.go +++ b/apps/backend/internal/agent/lifecycle/manager_execution.go @@ -49,6 +49,52 @@ func (m *Manager) GetOrEnsureExecution(ctx context.Context, sessionID string) (* return v.(*AgentExecution), nil } +// GetOrEnsureExecutionForEnvironment returns an execution for a task environment, +// creating one on-demand from the workspace info provider when needed. +func (m *Manager) GetOrEnsureExecutionForEnvironment(ctx context.Context, taskEnvironmentID string) (*AgentExecution, error) { + if taskEnvironmentID == "" { + return nil, fmt.Errorf("task_environment_id is required") + } + + if execution, exists := m.executionStore.GetByTaskEnvironmentID(taskEnvironmentID); exists { + return execution, nil + } + + groupKey := "env:" + taskEnvironmentID + v, err, _ := m.ensureExecutionGroup.Do(groupKey, func() (interface{}, error) { + if execution, exists := m.executionStore.GetByTaskEnvironmentID(taskEnvironmentID); exists { + return execution, nil + } + if m.workspaceInfoProvider == nil { + return nil, fmt.Errorf("workspace info provider not configured") + } + info, err := m.workspaceInfoProvider.GetWorkspaceInfoForEnvironment(ctx, taskEnvironmentID) + if err != nil { + return nil, fmt.Errorf("failed to get workspace info for environment %s: %w", taskEnvironmentID, err) + } + if info == nil { + return nil, fmt.Errorf("task environment %s not found", taskEnvironmentID) + } + if info.TaskEnvironmentID == "" { + return nil, fmt.Errorf("task environment %s has no task_environment_id", taskEnvironmentID) + } + if info.TaskEnvironmentID != taskEnvironmentID { + return nil, fmt.Errorf("workspace info resolved environment %s, want %s", info.TaskEnvironmentID, taskEnvironmentID) + } + if info.WorkspacePath == "" { + return nil, fmt.Errorf("%w: task environment %s has no workspace path yet", ErrSessionWorkspaceNotReady, taskEnvironmentID) + } + if info.SessionID == "" { + return nil, fmt.Errorf("task environment %s has no task session", taskEnvironmentID) + } + return m.createExecution(ctx, info.TaskID, info) + }) + if err != nil { + return nil, err + } + return v.(*AgentExecution), nil +} + // EnsureWorkspaceExecutionForSession ensures an agentctl execution exists for a specific task session. // This is used when the frontend provides a session ID (e.g., from URL path /task/[id]/[sessionId]). // If an execution already exists for the session, it returns it. Otherwise, it creates a new execution diff --git a/apps/backend/internal/agent/lifecycle/manager_execution_test.go b/apps/backend/internal/agent/lifecycle/manager_execution_test.go index 6318b3c97c..6e0a73a032 100644 --- a/apps/backend/internal/agent/lifecycle/manager_execution_test.go +++ b/apps/backend/internal/agent/lifecycle/manager_execution_test.go @@ -4,11 +4,19 @@ import ( "context" "errors" "fmt" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strconv" "sync" "sync/atomic" "testing" "time" + "github.com/kandev/kandev/internal/agent/executor" + agentctl "github.com/kandev/kandev/internal/agentctl/client" + "github.com/kandev/kandev/internal/common/logger" v1 "github.com/kandev/kandev/pkg/api/v1" ) @@ -38,7 +46,7 @@ func TestErrSessionWorkspaceNotReady_UnrelatedError(t *testing.T) { } } -func TestResolveScopeKey(t *testing.T) { +func TestResolveTaskEnvironmentID(t *testing.T) { t.Run("returns TaskEnvironmentID when execution carries it", func(t *testing.T) { store := NewExecutionStore() store.Add(&AgentExecution{ @@ -50,20 +58,46 @@ func TestResolveScopeKey(t *testing.T) { }) mgr := &Manager{executionStore: store, logger: newTestLogger()} - if got := mgr.ResolveScopeKey("session-A"); got != "env-1" { - t.Errorf("ResolveScopeKey = %q, want %q", got, "env-1") + got, err := mgr.ResolveTaskEnvironmentID(context.Background(), "session-A") + if err != nil { + t.Fatalf("ResolveTaskEnvironmentID returned error: %v", err) + } + if got != "env-1" { + t.Errorf("ResolveTaskEnvironmentID = %q, want %q", got, "env-1") } }) - t.Run("falls back to sessionID when no execution", func(t *testing.T) { + t.Run("returns TaskEnvironmentID from provider when no execution", func(t *testing.T) { + provider := &mockWorkspaceInfoProvider{ + infos: map[string]*WorkspaceInfo{ + "session-X": {SessionID: "session-X", TaskEnvironmentID: "env-X"}, + }, + } + mgr := &Manager{executionStore: NewExecutionStore(), logger: newTestLogger()} + mgr.workspaceInfoProvider = provider + + got, err := mgr.ResolveTaskEnvironmentID(context.Background(), "session-X") + if err != nil { + t.Fatalf("ResolveTaskEnvironmentID returned error: %v", err) + } + if got != "env-X" { + t.Errorf("ResolveTaskEnvironmentID = %q, want %q", got, "env-X") + } + }) + + t.Run("errors when no execution and no provider", func(t *testing.T) { mgr := &Manager{executionStore: NewExecutionStore(), logger: newTestLogger()} - if got := mgr.ResolveScopeKey("session-X"); got != "session-X" { - t.Errorf("ResolveScopeKey fallback = %q, want %q", got, "session-X") + _, err := mgr.ResolveTaskEnvironmentID(context.Background(), "session-X") + if err == nil { + t.Fatal("expected error") + } + if !containsString(err.Error(), "workspace info provider not configured") { + t.Errorf("unexpected error: %v", err) } }) - t.Run("falls back to sessionID when execution has empty env", func(t *testing.T) { + t.Run("errors when execution has empty env", func(t *testing.T) { store := NewExecutionStore() store.Add(&AgentExecution{ ID: "exec-2", @@ -73,8 +107,33 @@ func TestResolveScopeKey(t *testing.T) { }) mgr := &Manager{executionStore: store, logger: newTestLogger()} - if got := mgr.ResolveScopeKey("session-B"); got != "session-B" { - t.Errorf("ResolveScopeKey legacy = %q, want %q", got, "session-B") + _, err := mgr.ResolveTaskEnvironmentID(context.Background(), "session-B") + if err == nil { + t.Fatal("expected error") + } + if !containsString(err.Error(), "no task environment ID") { + t.Errorf("unexpected error: %v", err) + } + }) + + t.Run("errors when provider returns empty env", func(t *testing.T) { + provider := &mockWorkspaceInfoProvider{ + infos: map[string]*WorkspaceInfo{ + "session-C": {SessionID: "session-C"}, + }, + } + mgr := &Manager{ + executionStore: NewExecutionStore(), + workspaceInfoProvider: provider, + logger: newTestLogger(), + } + + _, err := mgr.ResolveTaskEnvironmentID(context.Background(), "session-C") + if err == nil { + t.Fatal("expected error") + } + if !containsString(err.Error(), "no task environment ID") { + t.Errorf("unexpected error: %v", err) } }) @@ -90,7 +149,15 @@ func TestResolveScopeKey(t *testing.T) { }) mgr := &Manager{executionStore: store, logger: newTestLogger()} - if mgr.ResolveScopeKey("sess-A") != mgr.ResolveScopeKey("sess-B") { + envA, err := mgr.ResolveTaskEnvironmentID(context.Background(), "sess-A") + if err != nil { + t.Fatalf("ResolveTaskEnvironmentID(sess-A): %v", err) + } + envB, err := mgr.ResolveTaskEnvironmentID(context.Background(), "sess-B") + if err != nil { + t.Fatalf("ResolveTaskEnvironmentID(sess-B): %v", err) + } + if envA != envB { t.Error("sessions in the same env must resolve to the same scope key") } }) @@ -223,6 +290,191 @@ func TestGetOrEnsureExecution(t *testing.T) { }) } +func TestGetOrEnsureExecutionForEnvironment(t *testing.T) { + t.Run("returns existing execution by environment", func(t *testing.T) { + store := NewExecutionStore() + execution := &AgentExecution{ + ID: "exec-1", + SessionID: "session-1", + TaskID: "task-1", + TaskEnvironmentID: "env-1", + Status: v1.AgentStatusRunning, + } + store.Add(execution) + mgr := &Manager{executionStore: store, logger: newTestLogger()} + + got, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-1") + if err != nil { + t.Fatalf("GetOrEnsureExecutionForEnvironment returned error: %v", err) + } + if got.ID != "exec-1" { + t.Errorf("execution ID = %q, want exec-1", got.ID) + } + }) + + t.Run("creates execution from provider and caches it", func(t *testing.T) { + mgr, backend := newEnvironmentExecutionTestManager(t, &mockWorkspaceInfoProvider{ + envInfos: map[string]*WorkspaceInfo{ + "env-new": { + TaskID: "task-1", + SessionID: "session-1", + TaskEnvironmentID: "env-new", + WorkspacePath: "/workspace/task-1", + AgentID: "auggie", + }, + }, + }) + + got, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-new") + if err != nil { + t.Fatalf("GetOrEnsureExecutionForEnvironment returned error: %v", err) + } + if got.TaskEnvironmentID != "env-new" { + t.Errorf("TaskEnvironmentID = %q, want env-new", got.TaskEnvironmentID) + } + if got.WorkspacePath != "/workspace/task-1" { + t.Errorf("WorkspacePath = %q, want /workspace/task-1", got.WorkspacePath) + } + + got2, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-new") + if err != nil { + t.Fatalf("second GetOrEnsureExecutionForEnvironment returned error: %v", err) + } + if got2.ID != got.ID { + t.Errorf("cached execution ID = %q, want %q", got2.ID, got.ID) + } + if backend.createCount.Load() != 1 { + t.Errorf("CreateInstance calls = %d, want 1", backend.createCount.Load()) + } + }) + + t.Run("concurrent creates use singleflight", func(t *testing.T) { + mgr, backend := newEnvironmentExecutionTestManager(t, &mockWorkspaceInfoProvider{ + envInfos: map[string]*WorkspaceInfo{ + "env-new": { + TaskID: "task-1", + SessionID: "session-1", + TaskEnvironmentID: "env-new", + WorkspacePath: "/workspace/task-1", + AgentID: "auggie", + }, + }, + }) + backend.delay = 50 * time.Millisecond + + var wg sync.WaitGroup + errs := make(chan error, 2) + ids := make(chan string, 2) + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + execution, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-new") + if err != nil { + errs <- err + return + } + ids <- execution.ID + }() + } + wg.Wait() + close(errs) + close(ids) + + for err := range errs { + t.Fatalf("GetOrEnsureExecutionForEnvironment returned error: %v", err) + } + var firstID string + for id := range ids { + if firstID == "" { + firstID = id + continue + } + if id != firstID { + t.Errorf("execution ID = %q, want %q", id, firstID) + } + } + if backend.createCount.Load() != 1 { + t.Errorf("CreateInstance calls = %d, want 1", backend.createCount.Load()) + } + }) + + t.Run("empty environment ID returns error", func(t *testing.T) { + mgr := &Manager{executionStore: NewExecutionStore(), logger: newTestLogger()} + + _, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "") + if err == nil { + t.Fatal("expected error") + } + if err.Error() != "task_environment_id is required" { + t.Errorf("unexpected error: %v", err) + } + }) + + t.Run("missing provider returns error instead of session fallback", func(t *testing.T) { + mgr := &Manager{executionStore: NewExecutionStore(), logger: newTestLogger()} + + _, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-missing") + if err == nil { + t.Fatal("expected error") + } + if !containsString(err.Error(), "workspace info provider not configured") { + t.Errorf("unexpected error: %v", err) + } + }) + + t.Run("provider must return matching environment ID", func(t *testing.T) { + provider := &mockWorkspaceInfoProvider{ + envInfos: map[string]*WorkspaceInfo{ + "env-want": { + TaskID: "task-1", + SessionID: "session-1", + TaskEnvironmentID: "env-other", + WorkspacePath: "/tmp/test", + }, + }, + } + mgr := &Manager{ + executionStore: NewExecutionStore(), + workspaceInfoProvider: provider, + logger: newTestLogger(), + } + + _, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-want") + if err == nil { + t.Fatal("expected error") + } + if !containsString(err.Error(), "workspace info resolved environment env-other") { + t.Errorf("unexpected error: %v", err) + } + }) + + t.Run("provider must return a workspace path", func(t *testing.T) { + provider := &mockWorkspaceInfoProvider{ + envInfos: map[string]*WorkspaceInfo{ + "env-1": { + TaskID: "task-1", + SessionID: "session-1", + TaskEnvironmentID: "env-1", + }, + }, + } + mgr := &Manager{ + executionStore: NewExecutionStore(), + workspaceInfoProvider: provider, + logger: newTestLogger(), + } + + _, err := mgr.GetOrEnsureExecutionForEnvironment(context.Background(), "env-1") + if err == nil { + t.Fatal("expected error") + } + if !errors.Is(err, ErrSessionWorkspaceNotReady) { + t.Errorf("expected ErrSessionWorkspaceNotReady, got %v", err) + } + }) +} + func TestEnsureWorkspaceExecutionForSession_EmptyTaskID(t *testing.T) { t.Run("resolves taskID from provider when empty", func(t *testing.T) { provider := &mockWorkspaceInfoProvider{ @@ -288,6 +540,77 @@ func TestEnsureWorkspaceExecutionForSession_EmptyTaskID(t *testing.T) { // --- test helpers --- +type createInstanceExecutor struct { + MockExecutor + client *agentctl.Client + createCount atomic.Int32 + delay time.Duration +} + +func (e *createInstanceExecutor) CreateInstance(ctx context.Context, req *ExecutorCreateRequest) (*ExecutorInstance, error) { + if e.delay > 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(e.delay): + } + } + e.createCount.Add(1) + return &ExecutorInstance{ + InstanceID: req.InstanceID, + TaskID: req.TaskID, + SessionID: req.SessionID, + RuntimeName: string(e.Name()), + Client: e.client, + WorkspacePath: req.WorkspacePath, + }, nil +} + +func newEnvironmentExecutionTestManager(t *testing.T, provider WorkspaceInfoProvider) (*Manager, *createInstanceExecutor) { + t.Helper() + log := newTestLogger() + execRegistry := NewExecutorRegistry(log) + backend := &createInstanceExecutor{ + MockExecutor: MockExecutor{name: executor.NameStandalone}, + client: newReadyAgentctlClient(t, log), + } + execRegistry.Register(backend) + + mgr := NewManager( + newTestRegistry(), &MockEventBus{}, execRegistry, &MockCredentialsManager{}, &MockProfileResolver{}, nil, + ExecutorFallbackWarn, "", log, + ) + mgr.workspaceInfoProvider = provider + t.Cleanup(func() { close(mgr.stopCh) }) + return mgr, backend +} + +func newReadyAgentctlClient(t *testing.T, log *logger.Logger) *agentctl.Client { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/health" { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusNotFound) + })) + t.Cleanup(server.Close) + + parsed, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("parse test server URL: %v", err) + } + host, portString, err := net.SplitHostPort(parsed.Host) + if err != nil { + t.Fatalf("split test server host: %v", err) + } + port, err := strconv.Atoi(portString) + if err != nil { + t.Fatalf("parse test server port: %v", err) + } + return agentctl.NewClient(host, port, log) +} + func containsString(s, substr string) bool { return len(s) >= len(substr) && searchString(s, substr) } @@ -312,6 +635,11 @@ func (p *trackingWorkspaceInfoProvider) GetWorkspaceInfoForSession(ctx context.C return p.delegate.GetWorkspaceInfoForSession(ctx, taskID, sessionID) } +func (p *trackingWorkspaceInfoProvider) GetWorkspaceInfoForEnvironment(ctx context.Context, taskEnvironmentID string) (*WorkspaceInfo, error) { + *p.called = true + return p.delegate.GetWorkspaceInfoForEnvironment(ctx, taskEnvironmentID) +} + // slowWorkspaceInfoProvider adds a delay to simulate slow DB lookups for concurrency tests. type slowWorkspaceInfoProvider struct { delay time.Duration @@ -328,3 +656,12 @@ func (p *slowWorkspaceInfoProvider) GetWorkspaceInfoForSession(_ context.Context } return p.info, nil } + +func (p *slowWorkspaceInfoProvider) GetWorkspaceInfoForEnvironment(_ context.Context, _ string) (*WorkspaceInfo, error) { + p.callCount.Add(1) + time.Sleep(p.delay) + if p.err != nil { + return nil, p.err + } + return p.info, nil +} diff --git a/apps/backend/internal/agent/lifecycle/manager_interaction.go b/apps/backend/internal/agent/lifecycle/manager_interaction.go index 16b1a76448..9ed4195a1f 100644 --- a/apps/backend/internal/agent/lifecycle/manager_interaction.go +++ b/apps/backend/internal/agent/lifecycle/manager_interaction.go @@ -634,19 +634,30 @@ func (m *Manager) GetExecutionBySessionID(sessionID string) (*AgentExecution, bo return m.executionStore.GetBySessionID(sessionID) } -// ResolveScopeKey returns a stable scope key for env-keyed resources (user -// shells, etc.) given a sessionID. Returns the session's TaskEnvironmentID -// when an active execution carries it, otherwise falls back to the sessionID -// itself so callers always get a non-empty key. -// -// This is the seam that makes terminals task-keyed: two sessions in the -// same task share an env, so they resolve to the same scope key and the -// runner returns one shared shell list instead of duplicating per-session. -func (m *Manager) ResolveScopeKey(sessionID string) string { - if exec, ok := m.executionStore.GetBySessionID(sessionID); ok && exec.TaskEnvironmentID != "" { - return exec.TaskEnvironmentID +// ResolveTaskEnvironmentID returns the task environment ID for a session. +// User shell resources must be environment-scoped; missing mappings are +// lifecycle errors and must not be converted into session-scoped shell state. +func (m *Manager) ResolveTaskEnvironmentID(ctx context.Context, sessionID string) (string, error) { + if sessionID == "" { + return "", fmt.Errorf("session_id is required") } - return sessionID + if exec, ok := m.executionStore.GetBySessionID(sessionID); ok { + if exec.TaskEnvironmentID != "" { + return exec.TaskEnvironmentID, nil + } + return "", fmt.Errorf("session %s has no task environment ID", sessionID) + } + if m.workspaceInfoProvider == nil { + return "", fmt.Errorf("workspace info provider not configured") + } + info, err := m.workspaceInfoProvider.GetWorkspaceInfoForSession(ctx, "", sessionID) + if err != nil { + return "", fmt.Errorf("failed to resolve task environment for session %s: %w", sessionID, err) + } + if info == nil || info.TaskEnvironmentID == "" { + return "", fmt.Errorf("session %s has no task environment ID", sessionID) + } + return info.TaskEnvironmentID, nil } // IsRemoteSession checks whether a session is associated with a remote executor diff --git a/apps/backend/internal/agent/lifecycle/manager_interaction_test.go b/apps/backend/internal/agent/lifecycle/manager_interaction_test.go index 3d995f7c9b..2b28135997 100644 --- a/apps/backend/internal/agent/lifecycle/manager_interaction_test.go +++ b/apps/backend/internal/agent/lifecycle/manager_interaction_test.go @@ -517,8 +517,9 @@ func TestIsAgentRunningForSession(t *testing.T) { // --- IsRemoteSession tests --- type mockWorkspaceInfoProvider struct { - infos map[string]*WorkspaceInfo - err error + infos map[string]*WorkspaceInfo + envInfos map[string]*WorkspaceInfo + err error } func (m *mockWorkspaceInfoProvider) GetWorkspaceInfoForSession(_ context.Context, _, sessionID string) (*WorkspaceInfo, error) { @@ -528,6 +529,21 @@ func (m *mockWorkspaceInfoProvider) GetWorkspaceInfoForSession(_ context.Context return m.infos[sessionID], nil } +func (m *mockWorkspaceInfoProvider) GetWorkspaceInfoForEnvironment(_ context.Context, taskEnvironmentID string) (*WorkspaceInfo, error) { + if m.err != nil { + return nil, m.err + } + if m.envInfos != nil { + return m.envInfos[taskEnvironmentID], nil + } + for _, info := range m.infos { + if info.TaskEnvironmentID == taskEnvironmentID { + return info, nil + } + } + return nil, nil +} + func TestIsRemoteSession(t *testing.T) { t.Run("in-memory execution with sprites runtime", func(t *testing.T) { store := NewExecutionStore() diff --git a/apps/backend/internal/agent/lifecycle/types.go b/apps/backend/internal/agent/lifecycle/types.go index de0f1dc18e..c29fd0c02f 100644 --- a/apps/backend/internal/agent/lifecycle/types.go +++ b/apps/backend/internal/agent/lifecycle/types.go @@ -375,6 +375,8 @@ type WorkspaceInfo struct { type WorkspaceInfoProvider interface { // GetWorkspaceInfoForSession returns workspace info for a specific task session GetWorkspaceInfoForSession(ctx context.Context, taskID, sessionID string) (*WorkspaceInfo, error) + // GetWorkspaceInfoForEnvironment returns workspace info for a task environment. + GetWorkspaceInfoForEnvironment(ctx context.Context, taskEnvironmentID string) (*WorkspaceInfo, error) } // RecoveredExecution contains info about an execution recovered from a runtime. diff --git a/apps/backend/internal/agentctl/server/process/interactive_shells.go b/apps/backend/internal/agentctl/server/process/interactive_shells.go index fdf17ac392..9e2911b924 100644 --- a/apps/backend/internal/agentctl/server/process/interactive_shells.go +++ b/apps/backend/internal/agentctl/server/process/interactive_shells.go @@ -43,8 +43,7 @@ type UserShellInfo struct { // The entry is registered atomically to prevent races with ListUserShells. // // `scopeID` groups shells. Callers should pass `taskEnvironmentID` so sessions -// in the same task share one shell list. Falls back to sessionID is preserved -// at the caller boundary (lifecycle.Manager.ResolveScopeKey). +// in the same task share one shell list. func (r *InteractiveRunner) CreateUserShell(scopeID string) CreateUserShellResult { r.userShellsMu.Lock() defer r.userShellsMu.Unlock() diff --git a/apps/backend/internal/gateway/websocket/setup.go b/apps/backend/internal/gateway/websocket/setup.go index bf4545f795..4490cd9806 100644 --- a/apps/backend/internal/gateway/websocket/setup.go +++ b/apps/backend/internal/gateway/websocket/setup.go @@ -72,7 +72,7 @@ func (g *Gateway) SetupRoutes(router *gin.Engine) { // Add dedicated terminal WebSocket route if terminal handler is configured if g.TerminalHandler != nil { - router.GET("/terminal/:sessionId", g.TerminalHandler.HandleTerminalWS) + router.GET("/terminal/*target", g.TerminalHandler.HandleTerminalWS) } // Add LSP routes if LSP handler is configured diff --git a/apps/backend/internal/gateway/websocket/terminal_handler.go b/apps/backend/internal/gateway/websocket/terminal_handler.go index f962625340..c267aa27b8 100644 --- a/apps/backend/internal/gateway/websocket/terminal_handler.go +++ b/apps/backend/internal/gateway/websocket/terminal_handler.go @@ -112,7 +112,6 @@ func isResizeCommand(data []byte) bool { const ( passthroughReadyTimeout = 30 * time.Second - shellExecutionTimeout = 15 * time.Second pollInterval = 500 * time.Millisecond ) @@ -155,60 +154,89 @@ func (w *wsWriter) Close() error { return nil } -// HandleTerminalWS handles WebSocket connections at /terminal/:sessionId -// This creates a binary WebSocket bridge between xterm.js and the PTY. -// Query parameter "mode" determines the type of terminal: -// - "agent": Agent passthrough terminal (CLI passthrough mode) -// - "shell": Independent user shell terminal (requires terminalId query param) +type terminalRoute struct { + kind string + id string +} + +const ( + terminalRouteEnvironment = "environment" + terminalRouteSession = "session" +) + +func parseTerminalRoute(c *gin.Context) terminalRoute { + target := strings.Trim(c.Param("target"), "/") + if target == "" { + return terminalRoute{} + } + if id, ok := strings.CutPrefix(target, "environment/"); ok { + return terminalRoute{kind: terminalRouteEnvironment, id: id} + } + if id, ok := strings.CutPrefix(target, "session/"); ok { + return terminalRoute{kind: terminalRouteSession, id: id} + } + return terminalRoute{} +} + +// HandleTerminalWS handles terminal WebSocket connections. +// Explicit routes: +// - /terminal/session/:sessionId?mode=agent for agent passthrough terminals +// - /terminal/environment/:environmentId?terminalId=... for user shell terminals func (h *TerminalHandler) HandleTerminalWS(c *gin.Context) { - sessionID := c.Param("sessionId") + route := parseTerminalRoute(c) + if route.id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "terminal route must be /environment/:environmentId for shell terminals or /session/:sessionId?mode=agent for agent terminals"}) + return + } + + switch route.kind { + case terminalRouteEnvironment: + h.handleEnvironmentTerminalRoute(c, route.id) + case terminalRouteSession: + h.handleSessionTerminalRoute(c, route.id) + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid terminal route"}) + } +} + +func (h *TerminalHandler) handleSessionTerminalRoute(c *gin.Context, sessionID string) { + if c.Query("mode") != "agent" { + c.JSON(http.StatusBadRequest, gin.H{"error": "session terminal route requires mode=agent; shell terminals must use /terminal/environment/:environmentId"}) + return + } + h.handleAgentTerminalRoute(c, sessionID) +} + +func (h *TerminalHandler) handleEnvironmentTerminalRoute(c *gin.Context, environmentID string) { + mode := c.Query("mode") + if mode != "" && mode != "shell" { + c.JSON(http.StatusBadRequest, gin.H{"error": "environment terminal route only supports shell mode"}) + return + } + terminalID := c.Query("terminalId") + if terminalID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "terminalId required for shell mode"}) + return + } + h.handleEnvironmentUserShellWS(c, environmentID, terminalID) +} + +func (h *TerminalHandler) handleAgentTerminalRoute(c *gin.Context, sessionID string) { if sessionID == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "sessionId is required"}) return } - - // Route based on explicit mode parameter - mode := c.Query("mode") - switch mode { - case "agent": - interactiveRunner := h.lifecycleMgr.GetInteractiveRunner() - if interactiveRunner == nil { - h.logger.Error("interactive runner not available") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "interactive runner not available"}) - return - } - h.logger.Info("terminal WebSocket connection request", - zap.String("session_id", sessionID), - zap.String("mode", "agent"), - zap.String("remote_addr", c.Request.RemoteAddr)) - h.handleAgentPassthroughWS(c, sessionID, interactiveRunner) - - case "shell": - terminalID := c.Query("terminalId") - if terminalID == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "terminalId required for shell mode"}) - return - } - if h.shouldUseWorkspaceShell(c.Request.Context(), sessionID) { - h.handleRemoteUserShellWS(c, sessionID, terminalID) - return - } - interactiveRunner := h.lifecycleMgr.GetInteractiveRunner() - if interactiveRunner == nil { - h.logger.Error("interactive runner not available") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "interactive runner not available"}) - return - } - h.logger.Info("terminal WebSocket connection request", - zap.String("session_id", sessionID), - zap.String("mode", "shell"), - zap.String("terminal_id", terminalID), - zap.String("remote_addr", c.Request.RemoteAddr)) - h.handleUserShellWS(c, sessionID, terminalID, interactiveRunner) - - default: - c.JSON(http.StatusBadRequest, gin.H{"error": "mode query param required: agent or shell"}) + interactiveRunner := h.lifecycleMgr.GetInteractiveRunner() + if interactiveRunner == nil { + h.logger.Error("interactive runner not available") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "interactive runner not available"}) + return } + h.logger.Info("terminal WebSocket connection request", + zap.String("session_id", sessionID), + zap.String("mode", "agent"), + zap.String("remote_addr", c.Request.RemoteAddr)) + h.handleAgentPassthroughWS(c, sessionID, interactiveRunner) } func (h *TerminalHandler) shouldUseWorkspaceShell(ctx context.Context, sessionID string) bool { @@ -217,56 +245,115 @@ func (h *TerminalHandler) shouldUseWorkspaceShell(ctx context.Context, sessionID return h.lifecycleMgr.ShouldUseContainerShell(ctx, sessionID) } -// waitForRemoteExecution polls the lifecycle manager for the session's execution, -// waiting up to 30 seconds for it to become available. After a backend restart the -// resume flow recreates the execution asynchronously; this bridges the gap so the -// terminal doesn't fall back to a local shell. -func (h *TerminalHandler) waitForRemoteExecution(ctx context.Context, sessionID string) (*lifecycle.AgentExecution, bool) { - return h.waitForRemoteExecutionWithTimeout(ctx, sessionID, 30*time.Second) -} - -func (h *TerminalHandler) waitForRemoteExecutionWithTimeout( +func (h *TerminalHandler) waitForRemoteExecutionReadyWithTimeout( ctx context.Context, sessionID string, timeout time.Duration, ) (*lifecycle.AgentExecution, bool) { - // Fast path: execution already exists. - if execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID); exists { - return execution, true - } + deadline := time.Now().Add(timeout) - h.logger.Info("waiting for execution to become available", + h.logger.Info("waiting for remote execution to become ready", zap.String("session_id", sessionID), zap.Duration("timeout", timeout)) - ticker := time.NewTicker(pollInterval) - defer ticker.Stop() - timer := time.NewTimer(timeout) - defer timer.Stop() - for { - select { - case <-ctx.Done(): - return nil, false - case <-timer.C: - h.logger.Warn("timed out waiting for execution", + if execution, ok := h.remoteExecutionWithClientReady(ctx, sessionID, deadline); ok { + return execution, true + } + + remaining := time.Until(deadline) + if remaining <= 0 { + h.logger.Warn("timed out waiting for remote execution readiness", zap.String("session_id", sessionID), zap.Duration("timeout", timeout)) return nil, false - case <-ticker.C: - if execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID); exists { - h.logger.Info("execution became available", - zap.String("session_id", sessionID), - zap.String("execution_id", execution.ID)) - return execution, true - } + } + + wait := pollInterval + if remaining < wait { + wait = remaining + } + timer := time.NewTimer(wait) + select { + case <-ctx.Done(): + timer.Stop() + return nil, false + case <-timer.C: } } } -func (h *TerminalHandler) handleRemoteUserShellWS(c *gin.Context, sessionID, terminalID string) { - execution, exists := h.waitForRemoteExecution(c.Request.Context(), sessionID) +func (h *TerminalHandler) remoteExecutionWithClientReady( + ctx context.Context, + sessionID string, + deadline time.Time, +) (*lifecycle.AgentExecution, bool) { + execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID) if !exists || execution.GetAgentCtlClient() == nil { + return nil, false + } + remaining := time.Until(deadline) + if remaining <= 0 { + return nil, false + } + if err := execution.GetAgentCtlClient().WaitForReady(ctx, remaining); err != nil { + h.logger.Debug("remote execution agentctl client not ready", + zap.String("session_id", sessionID), + zap.String("execution_id", execution.ID), + zap.Error(err)) + return nil, false + } + h.logger.Info("remote execution became ready", + zap.String("session_id", sessionID), + zap.String("execution_id", execution.ID)) + return execution, true +} + +func (h *TerminalHandler) handleEnvironmentUserShellWS(c *gin.Context, environmentID, terminalID string) { + execution, err := h.lifecycleMgr.GetOrEnsureExecutionForEnvironment(c.Request.Context(), environmentID) + if err != nil { + h.logger.Warn("environment terminal execution not ready", + zap.String("task_environment_id", environmentID), + zap.String("terminal_id", terminalID), + zap.Error(err)) + c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()}) + return + } + if h.shouldUseWorkspaceShell(c.Request.Context(), execution.SessionID) { + remoteExecution, ok := h.waitForRemoteExecutionReadyWithTimeout( + c.Request.Context(), execution.SessionID, passthroughReadyTimeout, + ) + if !ok { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "remote execution not available"}) + return + } + execution = remoteExecution + h.handleRemoteUserShellWS(c, execution, environmentID, terminalID) + return + } + interactiveRunner := h.lifecycleMgr.GetInteractiveRunner() + if interactiveRunner == nil { + h.logger.Error("interactive runner not available") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "interactive runner not available"}) + return + } + h.logger.Info("terminal WebSocket connection request", + zap.String("task_environment_id", environmentID), + zap.String("session_id", execution.SessionID), + zap.String("mode", "shell"), + zap.String("terminal_id", terminalID), + zap.String("remote_addr", c.Request.RemoteAddr)) + h.handleUserShellWS(c, execution, environmentID, terminalID, interactiveRunner) +} + +func (h *TerminalHandler) handleRemoteUserShellWS( + c *gin.Context, + execution *lifecycle.AgentExecution, + scopeID string, + terminalID string, +) { + sessionID := execution.SessionID + if execution.GetAgentCtlClient() == nil { c.JSON(http.StatusNotFound, gin.H{"error": "remote execution not available"}) return } @@ -277,13 +364,10 @@ func (h *TerminalHandler) handleRemoteUserShellWS(c *gin.Context, sessionID, ter c.JSON(http.StatusBadRequest, gin.H{"error": httpErr}) return } - // The per-terminal WS URL only carries terminalId, so fall back to the - // command pre-registered by wsUserShellCreate when the URL didn't include - // scriptId or label-derived command. Without this, script terminals (custom - // scripts and dev_script) would open empty in containerized sessions. + // The per-terminal WS URL carries only env + terminalId, so recover the + // command pre-registered by wsUserShellCreate for script/dev terminals. if initialCommand == "" { if runner := h.lifecycleMgr.GetInteractiveRunner(); runner != nil { - scopeID := scopeIDForExecution(execution, sessionID) initialCommand = runner.LookupShellInitialCommand(scopeID, terminalID) } } @@ -326,6 +410,7 @@ func (h *TerminalHandler) handleRemoteUserShellWS(c *gin.Context, sessionID, ter }() h.logger.Info("remote terminal shell WebSocket connected", + zap.String("task_environment_id", scopeID), zap.String("session_id", sessionID), zap.String("terminal_id", terminalID)) @@ -850,47 +935,20 @@ func (h *TerminalHandler) resolvePreferredShell(ctx context.Context) string { return shell } -// resolveWorkingDir returns the workspace path for the given session. -// It waits briefly for execution/workspace readiness to avoid launching shells in the wrong cwd. -func (h *TerminalHandler) resolveWorkingDir(ctx context.Context, sessionID string) (string, error) { - execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID) - if exists && execution.WorkspacePath != "" { - h.logger.Info("handleUserShellWS: got working directory from execution", - zap.String("working_dir", execution.WorkspacePath)) - return execution.WorkspacePath, nil - } - - h.logger.Info("handleUserShellWS: waiting for execution workspace path", - zap.String("session_id", sessionID), - zap.Duration("timeout", shellExecutionTimeout)) - - execution, exists = h.waitForRemoteExecutionWithTimeout(ctx, sessionID, shellExecutionTimeout) - if !exists || execution == nil { - return "", fmt.Errorf("session %s execution is not ready yet", sessionID) +// resolveWorkingDir returns the workspace path from the environment execution. +func (h *TerminalHandler) resolveWorkingDir(execution *lifecycle.AgentExecution) (string, error) { + if execution == nil { + return "", fmt.Errorf("environment execution is not ready") } if execution.WorkspacePath == "" { - return "", fmt.Errorf("%w: session %s has no workspace path configured", lifecycle.ErrSessionWorkspaceNotReady, sessionID) + return "", fmt.Errorf("%w: task environment %s has no workspace path configured", + lifecycle.ErrSessionWorkspaceNotReady, execution.TaskEnvironmentID) } - - h.logger.Info("handleUserShellWS: got working directory after wait", + h.logger.Info("handleUserShellWS: got working directory from environment execution", zap.String("working_dir", execution.WorkspacePath)) return execution.WorkspacePath, nil } -func (h *TerminalHandler) resolveScopeID(ctx context.Context, sessionID string) (string, error) { - if _, err := h.lifecycleMgr.GetOrEnsureExecution(ctx, sessionID); err != nil { - return "", fmt.Errorf("failed to ensure execution for scope: %w", err) - } - return h.lifecycleMgr.ResolveScopeKey(sessionID), nil -} - -func scopeIDForExecution(execution *lifecycle.AgentExecution, sessionID string) string { - if execution != nil && execution.TaskEnvironmentID != "" { - return execution.TaskEnvironmentID - } - return sessionID -} - // resolveShellLabel returns the label and initial command for a user shell, derived // from either a script ID lookup or a plain label query parameter. // Returns an HTTP error string (non-empty) when the script lookup fails. @@ -915,13 +973,15 @@ func (h *TerminalHandler) resolveShellLabel(c *gin.Context) (label, initialComma // user shell process. Returns the process ID on success, or an HTTP status + message on failure. func (h *TerminalHandler) startUserShellProcess( c *gin.Context, - sessionID string, + execution *lifecycle.AgentExecution, + scopeID string, terminalID string, interactiveRunner *process.InteractiveRunner, -) (processID string, scopeID string, httpStatus int, errMsg string) { +) (string, int, string) { + sessionID := execution.SessionID label, initialCommand, httpErr := h.resolveShellLabel(c) if httpErr != "" { - return "", "", http.StatusBadRequest, httpErr + return "", http.StatusBadRequest, httpErr } h.logger.Info("handleUserShellWS: starting user shell handling", @@ -930,21 +990,13 @@ func (h *TerminalHandler) startUserShellProcess( zap.String("label", label), zap.String("initial_command", initialCommand)) - scopeID, err := h.resolveScopeID(c.Request.Context(), sessionID) - if err != nil { - h.logger.Warn("handleUserShellWS: scope not ready", - zap.String("session_id", sessionID), - zap.Error(err)) - return "", "", http.StatusServiceUnavailable, err.Error() - } - preferredShell := h.resolvePreferredShell(c.Request.Context()) - workingDir, err := h.resolveWorkingDir(c.Request.Context(), sessionID) + workingDir, err := h.resolveWorkingDir(execution) if err != nil { h.logger.Warn("handleUserShellWS: workspace path not ready", zap.String("session_id", sessionID), zap.Error(err)) - return "", "", http.StatusServiceUnavailable, err.Error() + return "", http.StatusServiceUnavailable, err.Error() } opts := &process.UserShellOptions{Label: label, InitialCommand: initialCommand} @@ -965,7 +1017,7 @@ func (h *TerminalHandler) startUserShellProcess( zap.String("session_id", sessionID), zap.String("terminal_id", terminalID), zap.Error(err)) - return "", "", http.StatusServiceUnavailable, err.Error() + return "", http.StatusServiceUnavailable, err.Error() } h.logger.Info("handleUserShellWS: user shell started successfully", @@ -973,18 +1025,22 @@ func (h *TerminalHandler) startUserShellProcess( zap.String("terminal_id", terminalID), zap.String("process_id", info.ID)) - return info.ID, scopeID, 0, "" + return info.ID, 0, "" } // handleUserShellWS handles WebSocket connections for user shell terminals. // Each terminal tab gets its own independent shell process. func (h *TerminalHandler) handleUserShellWS( c *gin.Context, - sessionID string, + execution *lifecycle.AgentExecution, + scopeID string, terminalID string, interactiveRunner *process.InteractiveRunner, ) { - processID, scopeID, httpStatus, errMsg := h.startUserShellProcess(c, sessionID, terminalID, interactiveRunner) + sessionID := execution.SessionID + processID, httpStatus, errMsg := h.startUserShellProcess( + c, execution, scopeID, terminalID, interactiveRunner, + ) if errMsg != "" { c.JSON(httpStatus, gin.H{"error": errMsg}) return diff --git a/apps/backend/internal/gateway/websocket/terminal_handler_test.go b/apps/backend/internal/gateway/websocket/terminal_handler_test.go index ff782d7104..d32fd746a2 100644 --- a/apps/backend/internal/gateway/websocket/terminal_handler_test.go +++ b/apps/backend/internal/gateway/websocket/terminal_handler_test.go @@ -2,9 +2,13 @@ package websocket import ( "bytes" + "encoding/json" "net/http" + "net/http/httptest" "net/url" "testing" + + "github.com/gin-gonic/gin" ) func TestCheckWebSocketOrigin(t *testing.T) { @@ -145,3 +149,58 @@ func TestStripTerminalResponses(t *testing.T) { }) } } + +func TestParseTerminalRoute(t *testing.T) { + tests := []struct { + name string + target string + wantKind string + wantID string + }{ + {name: "environment route", target: "/environment/env-1", wantKind: terminalRouteEnvironment, wantID: "env-1"}, + {name: "session route", target: "/session/session-1", wantKind: terminalRouteSession, wantID: "session-1"}, + {name: "unknown route", target: "/unknown-target", wantKind: "", wantID: ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := parseTerminalRoute(&gin.Context{Params: gin.Params{{Key: "target", Value: tt.target}}}) + if got.kind != tt.wantKind || got.id != tt.wantID { + t.Fatalf("parseTerminalRoute() = {%q %q}, want {%q %q}", + got.kind, got.id, tt.wantKind, tt.wantID) + } + }) + } +} + +func TestSessionTerminalRouteRequiresAgentMode(t *testing.T) { + tests := []struct { + name string + query string + }{ + {name: "missing mode", query: ""}, + {name: "shell mode", query: "?mode=shell"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/terminal/session/session-1"+tt.query, nil) + + handler := &TerminalHandler{} + handler.handleSessionTerminalRoute(c, "session-1") + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + + var body map[string]string + if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil { + t.Fatalf("decode response: %v", err) + } + if body["error"] != "session terminal route requires mode=agent; shell terminals must use /terminal/environment/:environmentId" { + t.Fatalf("error = %q", body["error"]) + } + }) + } +} diff --git a/apps/backend/internal/gateway/websocket/terminal_handler_wait_test.go b/apps/backend/internal/gateway/websocket/terminal_handler_wait_test.go index ca8a10b14d..0d66cfd06a 100644 --- a/apps/backend/internal/gateway/websocket/terminal_handler_wait_test.go +++ b/apps/backend/internal/gateway/websocket/terminal_handler_wait_test.go @@ -62,7 +62,16 @@ func newTestTerminalHandler(t *testing.T) (*TerminalHandler, *lifecycle.Executio return handler, mgr.ExecutionStoreForTesting() } -func TestWaitForRemoteExecutionWithTimeout_FastPath(t *testing.T) { +func TestWaitForRemoteExecutionReadyWithTimeout_Timeout(t *testing.T) { + handler, _ := newTestTerminalHandler(t) + + _, ok := handler.waitForRemoteExecutionReadyWithTimeout(context.Background(), "session-missing", 600*time.Millisecond) + if ok { + t.Fatal("expected timeout, got execution") + } +} + +func TestWaitForRemoteExecutionReadyWithTimeout_NoClient(t *testing.T) { handler, store := newTestTerminalHandler(t) store.Add(&lifecycle.AgentExecution{ @@ -71,53 +80,19 @@ func TestWaitForRemoteExecutionWithTimeout_FastPath(t *testing.T) { Status: v1.AgentStatusRunning, }) - got, ok := handler.waitForRemoteExecutionWithTimeout(context.Background(), "session-1", 1*time.Second) - if !ok { - t.Fatal("expected fast path to find execution") - } - if got.ID != "exec-1" { - t.Errorf("expected exec-1, got %s", got.ID) - } -} - -func TestWaitForRemoteExecutionWithTimeout_PollSuccess(t *testing.T) { - handler, store := newTestTerminalHandler(t) - - // Add execution after a short delay (simulates async workspace setup) - go func() { - time.Sleep(200 * time.Millisecond) - store.Add(&lifecycle.AgentExecution{ - ID: "exec-delayed", - SessionID: "session-1", - Status: v1.AgentStatusRunning, - }) - }() - - got, ok := handler.waitForRemoteExecutionWithTimeout(context.Background(), "session-1", 5*time.Second) - if !ok { - t.Fatal("expected polling to find execution") - } - if got.ID != "exec-delayed" { - t.Errorf("expected exec-delayed, got %s", got.ID) - } -} - -func TestWaitForRemoteExecutionWithTimeout_Timeout(t *testing.T) { - handler, _ := newTestTerminalHandler(t) - - _, ok := handler.waitForRemoteExecutionWithTimeout(context.Background(), "session-missing", 600*time.Millisecond) + _, ok := handler.waitForRemoteExecutionReadyWithTimeout(context.Background(), "session-1", 600*time.Millisecond) if ok { - t.Fatal("expected timeout, got execution") + t.Fatal("expected timeout without agentctl client") } } -func TestWaitForRemoteExecutionWithTimeout_ContextCancelled(t *testing.T) { +func TestWaitForRemoteExecutionReadyWithTimeout_ContextCancelled(t *testing.T) { handler, _ := newTestTerminalHandler(t) ctx, cancel := context.WithCancel(context.Background()) cancel() - _, ok := handler.waitForRemoteExecutionWithTimeout(ctx, "session-1", 5*time.Second) + _, ok := handler.waitForRemoteExecutionReadyWithTimeout(ctx, "session-1", 5*time.Second) if ok { t.Fatal("expected context cancellation to return false") } diff --git a/apps/backend/internal/task/service/service_test.go b/apps/backend/internal/task/service/service_test.go index a427174ae8..3c3181fa72 100644 --- a/apps/backend/internal/task/service/service_test.go +++ b/apps/backend/internal/task/service/service_test.go @@ -98,18 +98,19 @@ func createTestService(t *testing.T) (*Service, *MockEventBus, *sqliterepo.Repos eventBus := NewMockEventBus() log, _ := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "json", OutputPath: "stdout"}) svc := NewService(Repos{ - Workspaces: repo, - Tasks: repo, - TaskRepos: repo, - Workflows: repo, - Messages: repo, - Turns: repo, - Sessions: repo, - GitSnapshots: repo, - RepoEntities: repo, - Executors: repo, - Environments: repo, - Reviews: repo, + Workspaces: repo, + Tasks: repo, + TaskRepos: repo, + Workflows: repo, + Messages: repo, + Turns: repo, + Sessions: repo, + GitSnapshots: repo, + RepoEntities: repo, + Executors: repo, + Environments: repo, + TaskEnvironments: repo, + Reviews: repo, }, eventBus, log, RepositoryDiscoveryConfig{}) return svc, eventBus, repo } diff --git a/apps/backend/internal/task/service/service_turns.go b/apps/backend/internal/task/service/service_turns.go index 8d86af2261..5cca5e15c7 100644 --- a/apps/backend/internal/task/service/service_turns.go +++ b/apps/backend/internal/task/service/service_turns.go @@ -5,6 +5,7 @@ import ( "database/sql" "errors" "fmt" + "sort" "time" "github.com/google/uuid" @@ -287,6 +288,43 @@ func (s *Service) GetWorkspaceInfoForSession(ctx context.Context, taskID, sessio return info, nil } +// GetWorkspaceInfoForEnvironment returns workspace information for a task environment. +func (s *Service) GetWorkspaceInfoForEnvironment(ctx context.Context, taskEnvironmentID string) (*lifecycle.WorkspaceInfo, error) { + if taskEnvironmentID == "" { + return nil, fmt.Errorf("task_environment_id is required") + } + env, err := s.taskEnvironments.GetTaskEnvironment(ctx, taskEnvironmentID) + if err != nil { + return nil, err + } + if env == nil { + return nil, fmt.Errorf("task environment %s not found", taskEnvironmentID) + } + sessions, err := s.sessions.ListTaskSessions(ctx, env.TaskID) + if err != nil { + return nil, fmt.Errorf("failed to list sessions for task environment %s: %w", taskEnvironmentID, err) + } + matching := make([]*models.TaskSession, 0, len(sessions)) + for _, session := range sessions { + if session.TaskEnvironmentID == taskEnvironmentID { + matching = append(matching, session) + } + } + if len(matching) > 0 { + sort.SliceStable(matching, func(i, j int) bool { + if !matching[i].StartedAt.Equal(matching[j].StartedAt) { + return matching[i].StartedAt.After(matching[j].StartedAt) + } + if !matching[i].UpdatedAt.Equal(matching[j].UpdatedAt) { + return matching[i].UpdatedAt.After(matching[j].UpdatedAt) + } + return matching[i].ID > matching[j].ID + }) + return s.GetWorkspaceInfoForSession(ctx, env.TaskID, matching[0].ID) + } + return nil, fmt.Errorf("task environment %s has no linked task session", taskEnvironmentID) +} + func isExecutorRunningNotFoundError(err error) bool { if err == nil { return false diff --git a/apps/backend/internal/task/service/service_turns_test.go b/apps/backend/internal/task/service/service_turns_test.go index bc6c9c42b6..05d1a4e150 100644 --- a/apps/backend/internal/task/service/service_turns_test.go +++ b/apps/backend/internal/task/service/service_turns_test.go @@ -212,3 +212,57 @@ func TestGetWorkspaceInfoForSession_SessionNotFound(t *testing.T) { t.Fatal("expected error for nonexistent session") } } + +func TestGetWorkspaceInfoForEnvironment(t *testing.T) { + svc, _, repo := createTestService(t) + ctx := context.Background() + + setupTestTask(t, repo) + now := time.Now().UTC() + if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{ + ID: "env-123", + TaskID: "task-123", + Status: models.TaskEnvironmentStatusReady, + CreatedAt: now, + UpdatedAt: now, + }); err != nil { + t.Fatalf("failed to create task environment: %v", err) + } + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "session-1", + TaskID: "task-123", + TaskEnvironmentID: "env-123", + State: models.TaskSessionStateCompleted, + AgentProfileSnapshot: map[string]interface{}{ + "agent_name": "auggie", + }, + StartedAt: now, + UpdatedAt: now, + }); err != nil { + t.Fatalf("failed to create session: %v", err) + } + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "session-2", + TaskID: "task-123", + TaskEnvironmentID: "env-123", + State: models.TaskSessionStateCompleted, + AgentProfileSnapshot: map[string]interface{}{ + "agent_name": "auggie", + }, + StartedAt: now.Add(time.Minute), + UpdatedAt: now.Add(time.Minute), + }); err != nil { + t.Fatalf("failed to create newer session: %v", err) + } + + info, err := svc.GetWorkspaceInfoForEnvironment(ctx, "env-123") + if err != nil { + t.Fatalf("GetWorkspaceInfoForEnvironment returned error: %v", err) + } + if info.SessionID != "session-2" { + t.Errorf("SessionID = %q, want session-2", info.SessionID) + } + if info.TaskEnvironmentID != "env-123" { + t.Errorf("TaskEnvironmentID = %q, want env-123", info.TaskEnvironmentID) + } +} diff --git a/apps/web/components/task/bottom-terminal-panel.tsx b/apps/web/components/task/bottom-terminal-panel.tsx index a36fb77e35..ebea52634c 100644 --- a/apps/web/components/task/bottom-terminal-panel.tsx +++ b/apps/web/components/task/bottom-terminal-panel.tsx @@ -3,7 +3,7 @@ import { useRef, useState, useCallback, useEffect } from "react"; import { IconMinus } from "@tabler/icons-react"; import { useAppStore, useAppStoreApi } from "@/components/state-provider"; -import { useEnvironmentSessionId } from "@/hooks/use-environment-session-id"; +import { useEnvironmentId } from "@/hooks/use-environment-session-id"; import { PassthroughTerminal } from "./passthrough-terminal"; import { Button } from "@kandev/ui/button"; import { cn } from "@/lib/utils"; @@ -15,10 +15,7 @@ const STORAGE_KEY_HEIGHT = "bottom-terminal-height"; const STORAGE_KEY_OPEN = "bottom-terminal-open"; export function BottomTerminalPanel() { - // Pin to a stable sessionId across same-task session switches so the bottom - // terminal doesn't reconnect / land on a different shell. Matches the - // dockview terminal panel's behavior in `terminal-panel.tsx`. - const sessionId = useEnvironmentSessionId(); + const environmentId = useEnvironmentId(); const visible = useAppStore((s) => s.bottomTerminal.isOpen); const pendingCommand = useAppStore((s) => s.bottomTerminal.pendingCommand); const storeApi = useAppStoreApi(); @@ -98,8 +95,8 @@ export function BottomTerminalPanel() { {/* Terminal content */}
(["chat", "files", "shell", "git"]); +const VALID_COMPONENTS = new Set(["chat", "files", "shell", "git", "terminal"]); /** * Build a minimal valid SerializedDockview-shaped object — matches what diff --git a/apps/web/components/task/passthrough-terminal.tsx b/apps/web/components/task/passthrough-terminal.tsx index 2aef4747fb..ba9c172e3a 100644 --- a/apps/web/components/task/passthrough-terminal.tsx +++ b/apps/web/components/task/passthrough-terminal.tsx @@ -22,25 +22,30 @@ import { useSendInput, useFitAndResize, } from "./use-passthrough-terminal"; +import { useEnvironmentSessionId } from "@/hooks/use-environment-session-id"; import { useTerminalSearch } from "./use-terminal-search"; import { TerminalSearchBar } from "./terminal-search-bar"; import { usePanelSearch } from "@/hooks/use-panel-search"; type BaseProps = { - sessionId?: string | null; autoFocus?: boolean; pendingCommand?: string | null; onCommandSent?: () => void; }; -type AgentTerminalProps = BaseProps & { mode: "agent"; label?: string }; -type ShellTerminalProps = BaseProps & { mode: "shell"; terminalId: string; label?: string }; +type AgentTerminalProps = BaseProps & { mode: "agent"; sessionId?: string | null; label?: string }; +type ShellTerminalProps = BaseProps & { + mode: "shell"; + environmentId: string | null | undefined; + terminalId: string; + label?: string; +}; type PassthroughTerminalProps = AgentTerminalProps | ShellTerminalProps; /** * PassthroughTerminal provides direct terminal interaction with an agent CLI. * * Design: Dedicated Binary WebSocket + AttachAddon - * - Uses a dedicated WebSocket connection to /terminal/:sessionId + * - Uses dedicated WebSocket routes for session-scoped agent terminals and env-scoped shells * - Raw binary frames bypass JSON encoding/decoding latency * - AttachAddon (official xterm.js addon) handles the bridging * - Unicode11Addon enables proper unicode character support @@ -93,13 +98,15 @@ function useWsBaseUrl() { // eslint-disable-next-line max-lines-per-function -- wires many hooks + refs; each block is already its own hook export function PassthroughTerminal(props: PassthroughTerminalProps) { - const { sessionId: propSessionId, mode, label, autoFocus, pendingCommand, onCommandSent } = props; + const { mode, label, autoFocus, pendingCommand, onCommandSent } = props; const terminalId = mode === "shell" ? props.terminalId : undefined; + const environmentId = mode === "shell" ? props.environmentId : undefined; const refs = useTerminalRefs(); const { terminalRef, xtermRef, fitAddonRef, wsRef, attachAddonRef } = refs; const storeSessionId = useAppStore((state) => state.tasks.activeSessionId); - const sessionId = propSessionId ?? storeSessionId; + const environmentSessionId = useEnvironmentSessionId(); + const sessionId = mode === "agent" ? (props.sessionId ?? storeSessionId) : environmentSessionId; const { session } = useSession(sessionId); const agentctlStatus = useSessionAgentctl(sessionId); @@ -109,20 +116,21 @@ export function PassthroughTerminal(props: PassthroughTerminalProps) { // spamming reconnects burns cycles and confuses the loading state. // Don't require isActive — restored workspaces (terminal-state sessions) // have agentctl running but the session state stays COMPLETED/CANCELLED. - const canConnect = Boolean(sessionId && agentctlStatus.isReady); + const connectionID = mode === "agent" ? sessionId : environmentId; + const canConnect = Boolean(connectionID && sessionId && agentctlStatus.isReady); const wsBaseUrl = useWsBaseUrl(); const [isTerminalReady, setIsTerminalReady] = useState(false); const onTerminalReady = useCallback(() => setIsTerminalReady(true), []); - // Track which session has an active WebSocket connection. The loading overlay - // resets on session switches without needing a separate setState effect. - const [connectedSessionId, setConnectedSessionId] = useState(null); - const isConnected = sessionId != null && connectedSessionId === sessionId; + // Track which terminal target has an active WebSocket connection. The loading + // overlay resets on target switches without needing a separate setState effect. + const [connectedTargetId, setConnectedTargetId] = useState(null); + const isConnected = connectionID != null && connectedTargetId === connectionID; const onConnected = useCallback(() => { - setConnectedSessionId(sessionId ?? null); + setConnectedTargetId(connectionID ?? null); if (autoFocus) refs.xtermRef.current?.textarea?.focus({ preventScroll: true }); - }, [sessionId, autoFocus, refs.xtermRef]); + }, [connectionID, autoFocus, refs.xtermRef]); const linkHandler = useTerminalLinkHandler(); const terminalFontFamily = useAppStore((s) => s.userSettings.terminalFontFamily); @@ -171,6 +179,7 @@ export function PassthroughTerminal(props: PassthroughTerminalProps) { useWebSocketConnection({ taskId, sessionId, + environmentId, canConnect, isTerminalReady, fitAndResize, diff --git a/apps/web/components/task/task-right-panel.tsx b/apps/web/components/task/task-right-panel.tsx index af412c7492..353f2a3b1d 100644 --- a/apps/web/components/task/task-right-panel.tsx +++ b/apps/web/components/task/task-right-panel.tsx @@ -170,7 +170,7 @@ type RightPanelContentProps = { scripts: RepositoryScript[]; handleRunCommand: (script: RepositoryScript) => void; terminals: Terminal[]; - sessionId: string | null; + environmentId: string | null; devProcessId: string | null | undefined; devOutput: string | undefined; isStoppingDev: boolean; @@ -190,7 +190,7 @@ function RightPanelContent({ scripts, handleRunCommand, terminals, - sessionId, + environmentId, devProcessId, devOutput, isStoppingDev, @@ -236,7 +236,7 @@ function RightPanelContent({ state.setRightPanelActiveTab); + const environmentId = useAppStore((state) => + sessionId ? (state.environmentIdBySessionId[sessionId] ?? null) : null, + ); const closeLayoutPreview = useLayoutStore((state) => state.closePreview); // Use the terminals hook @@ -326,7 +329,7 @@ const TaskRightPanel = memo(function TaskRightPanel({ scripts={scripts} handleRunCommand={handleRunCommand} terminals={terminals} - sessionId={sessionId} + environmentId={environmentId} devProcessId={devProcessId} devOutput={devOutput} isStoppingDev={isStoppingDev} @@ -368,13 +371,13 @@ function CommandsTabContent({ /** Terminal tab contents (dev-server and shell terminals) */ function TerminalTabContents({ terminals, - sessionId, + environmentId, devProcessId, devOutput, isStoppingDev, }: { terminals: Terminal[]; - sessionId: string | null; + environmentId: string | null; devProcessId: string | null | undefined; devOutput: string | undefined; isStoppingDev: boolean; @@ -394,8 +397,8 @@ function TerminalTabContents({ ) : ( diff --git a/apps/web/components/task/terminal-panel.tsx b/apps/web/components/task/terminal-panel.tsx index b408141284..2d094eaef2 100644 --- a/apps/web/components/task/terminal-panel.tsx +++ b/apps/web/components/task/terminal-panel.tsx @@ -6,7 +6,7 @@ import { PassthroughTerminal } from "./passthrough-terminal"; import { ShellTerminal } from "./shell-terminal"; import { useAppStore } from "@/components/state-provider"; import { useIsTaskArchived, ArchivedPanelPlaceholder } from "./task-archived-context"; -import { useEnvironmentSessionId } from "@/hooks/use-environment-session-id"; +import { useEnvironmentId } from "@/hooks/use-environment-session-id"; type TerminalPanelProps = { panelId: string; @@ -18,9 +18,7 @@ export const TerminalPanel = memo(function TerminalPanel({ params }: TerminalPan const type = (params.type as string) ?? "shell"; const processId = params.processId as string | undefined; - // Only reconnect the terminal when the environment changes, not on every - // session tab switch within the same environment. - const sessionId = useEnvironmentSessionId(); + const environmentId = useEnvironmentId(); const devOutput = useAppStore((state) => processId ? (state.processes.outputsByProcessId[processId] ?? "") : "", @@ -47,11 +45,7 @@ export const TerminalPanel = memo(function TerminalPanel({ params }: TerminalPan return ( - + ); diff --git a/apps/web/components/task/use-passthrough-terminal.test.ts b/apps/web/components/task/use-passthrough-terminal.test.ts index f9cc063e7a..001fef91f2 100644 --- a/apps/web/components/task/use-passthrough-terminal.test.ts +++ b/apps/web/components/task/use-passthrough-terminal.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import { buildTerminalWsUrl } from "./use-passthrough-terminal"; import { reconnectDelayMs } from "./ws-reconnect"; describe("reconnectDelayMs", () => { @@ -23,3 +24,24 @@ describe("reconnectDelayMs", () => { expect(reconnectDelayMs(100)).toBe(5000); }); }); + +describe("buildTerminalWsUrl", () => { + it("routes shell terminals by task environment ID", () => { + expect( + buildTerminalWsUrl("ws://localhost:38429", { + mode: "shell", + environmentId: "env-1", + terminalId: "terminal with spaces", + }), + ).toBe("ws://localhost:38429/terminal/environment/env-1?terminalId=terminal%20with%20spaces"); + }); + + it("routes agent terminals by session ID", () => { + expect( + buildTerminalWsUrl("ws://localhost:38429", { + mode: "agent", + sessionId: "session-1", + }), + ).toBe("ws://localhost:38429/terminal/session/session-1?mode=agent"); + }); +}); diff --git a/apps/web/components/task/use-passthrough-terminal.ts b/apps/web/components/task/use-passthrough-terminal.ts index 13a9969b5a..e6972c0a5b 100644 --- a/apps/web/components/task/use-passthrough-terminal.ts +++ b/apps/web/components/task/use-passthrough-terminal.ts @@ -311,6 +311,7 @@ export function useTerminalInit({ export type WebSocketConnectionOptions = { taskId: string | null; sessionId: string | null | undefined; + environmentId?: string | null | undefined; canConnect: boolean; isTerminalReady: boolean; fitAndResize: (force?: boolean) => void; @@ -325,23 +326,33 @@ export type WebSocketConnectionOptions = { onConnected: () => void; }; -function buildWsUrl( +export function buildTerminalWsUrl( wsBaseUrl: string, - sessionId: string, - mode: "agent" | "shell", - terminalId: string | undefined, - label?: string, + params: { + mode: "agent" | "shell"; + sessionId?: string; + environmentId?: string; + terminalId?: string; + label?: string; + }, ): string { - let wsUrl = - mode === "agent" - ? `${wsBaseUrl}/terminal/${sessionId}?mode=agent` - : `${wsBaseUrl}/terminal/${sessionId}?mode=shell&terminalId=${encodeURIComponent(terminalId!)}`; + const { mode, sessionId, environmentId, terminalId, label } = params; + let wsUrl: string; + if (mode === "agent") { + if (!sessionId) throw new Error("sessionId is required for agent terminal"); + wsUrl = `${wsBaseUrl}/terminal/session/${encodeURIComponent(sessionId)}?mode=agent`; + } else { + if (!environmentId) throw new Error("environmentId is required for shell terminal"); + if (!terminalId) throw new Error("terminalId is required for shell terminal"); + wsUrl = `${wsBaseUrl}/terminal/environment/${encodeURIComponent(environmentId)}?terminalId=${encodeURIComponent(terminalId)}`; + } if (label) wsUrl += `&label=${encodeURIComponent(label)}`; return wsUrl; } type ConnectWebSocketOptions = { - sessionId: string; + sessionId?: string; + environmentId?: string; wsBaseUrl: string; mode: "agent" | "shell"; terminalId: string | undefined; @@ -358,6 +369,7 @@ type ConnectWebSocketOptions = { function connectWebSocket({ sessionId, + environmentId, wsBaseUrl, mode, terminalId, @@ -384,8 +396,14 @@ function connectWebSocket({ } wsRef.current = null; } - const wsUrl = buildWsUrl(wsBaseUrl, sessionId, mode, terminalId, label); - log("Connecting to", wsUrl, { mode, terminalId, label }); + const wsUrl = buildTerminalWsUrl(wsBaseUrl, { + mode, + sessionId, + environmentId, + terminalId, + label, + }); + log("Connecting to", wsUrl, { mode, sessionId, environmentId, terminalId, label }); const ws = new WebSocket(wsUrl); ws.binaryType = "arraybuffer"; wsRef.current = ws; @@ -438,6 +456,7 @@ export { reconnectDelayMs } from "./ws-reconnect"; export function useWebSocketConnection({ taskId, sessionId, + environmentId, canConnect, isTerminalReady, fitAndResize, @@ -452,17 +471,24 @@ export function useWebSocketConnection({ onConnected, }: WebSocketConnectionOptions) { useEffect(() => { + const connectionKey = mode === "agent" ? sessionId : environmentId; log("WebSocket effect:", { taskId, sessionId, + environmentId, mode, terminalId, canConnect, isTerminalReady, hasTerminal: !!xtermRef.current, }); - if (!taskId || !sessionId || !canConnect || !isTerminalReady) { - log("WebSocket effect: early return", { taskId, sessionId, canConnect, isTerminalReady }); + if (!taskId || !connectionKey || !canConnect || !isTerminalReady) { + log("WebSocket effect: early return", { + taskId, + connectionKey, + canConnect, + isTerminalReady, + }); return; } if (!xtermRef.current || !fitAddonRef.current) { @@ -475,11 +501,12 @@ export function useWebSocketConnection({ // (not session-scoped), so the same xterm instance persists across session // switches. Without this reset, the previous session's PTY output leaks // into the new session's scrollback. - log("Resetting terminal buffer for session", sessionId); + log("Resetting terminal buffer for terminal target", connectionKey); terminal.reset(); const stopReconnectLoop = startReconnectLoop({ - sessionId, + sessionId: sessionId ?? undefined, + environmentId: environmentId ?? undefined, wsBaseUrl, mode, terminalId, @@ -514,6 +541,7 @@ export function useWebSocketConnection({ }, [ taskId, sessionId, + environmentId, canConnect, isTerminalReady, fitAndResize, diff --git a/apps/web/components/task/ws-reconnect.ts b/apps/web/components/task/ws-reconnect.ts index 255a92dc08..f223e4808c 100644 --- a/apps/web/components/task/ws-reconnect.ts +++ b/apps/web/components/task/ws-reconnect.ts @@ -10,7 +10,8 @@ export function reconnectDelayMs(attempt: number): number { } type ConnectWebSocketFn = (opts: { - sessionId: string; + sessionId?: string; + environmentId?: string; wsBaseUrl: string; mode: "agent" | "shell"; terminalId: string | undefined; @@ -26,7 +27,8 @@ type ConnectWebSocketFn = (opts: { }) => void; export type ReconnectLoopOptions = { - sessionId: string; + sessionId?: string; + environmentId?: string; wsBaseUrl: string; mode: "agent" | "shell"; terminalId: string | undefined; @@ -41,6 +43,7 @@ export type ReconnectLoopOptions = { export function startReconnectLoop({ sessionId, + environmentId, wsBaseUrl, mode, terminalId, @@ -65,6 +68,7 @@ export function startReconnectLoop({ if (!isMounted) return; connectWebSocket({ sessionId, + environmentId, wsBaseUrl, mode, terminalId, diff --git a/apps/web/hooks/domains/session/use-user-shells.ts b/apps/web/hooks/domains/session/use-user-shells.ts index 289a7f8ff4..5bb6c35be8 100644 --- a/apps/web/hooks/domains/session/use-user-shells.ts +++ b/apps/web/hooks/domains/session/use-user-shells.ts @@ -23,22 +23,23 @@ const EMPTY_SHELLS: UserShellInfo[] = []; */ export function useUserShells(sessionId: string | null): UseUserShellsReturn { const store = useAppStoreApi(); + const environmentId = useAppStore((state) => + sessionId ? (state.environmentIdBySessionId[sessionId] ?? null) : null, + ); - // Read from store — resolve environmentId for shared workspace state + // Read from store by explicit environment ID. Missing env means lifecycle data + // is incomplete; do not synthesize a session-scoped shell list. const shells = useAppStore((state) => { - if (!sessionId) return EMPTY_SHELLS; - const envKey = state.environmentIdBySessionId[sessionId] ?? sessionId; - return state.userShells.byEnvironmentId[envKey] ?? EMPTY_SHELLS; + if (!environmentId) return EMPTY_SHELLS; + return state.userShells.byEnvironmentId[environmentId] ?? EMPTY_SHELLS; }); const isLoading = useAppStore((state) => { - if (!sessionId) return false; - const envKey = state.environmentIdBySessionId[sessionId] ?? sessionId; - return state.userShells.loading[envKey] ?? false; + if (!environmentId) return false; + return state.userShells.loading[environmentId] ?? false; }); const isLoaded = useAppStore((state) => { - if (!sessionId) return false; - const envKey = state.environmentIdBySessionId[sessionId] ?? sessionId; - return state.userShells.loaded[envKey] ?? false; + if (!environmentId) return false; + return state.userShells.loaded[environmentId] ?? false; }); const connectionStatus = useAppStore((state) => state.connection.status); @@ -55,7 +56,7 @@ export function useUserShells(sessionId: string | null): UseUserShellsReturn { // Fetch user shells from backend useEffect(() => { - if (!sessionId) return; + if (!sessionId || !environmentId) return; if (connectionStatus !== "connected") return; // Detect session change to force refetch @@ -118,7 +119,7 @@ export function useUserShells(sessionId: string | null): UseUserShellsReturn { }; fetchShells(); - }, [sessionId, connectionStatus, isLoaded, store]); + }, [sessionId, environmentId, connectionStatus, isLoaded, store]); // Actions const addShell = (shell: UserShellInfo) => { diff --git a/apps/web/hooks/use-environment-session-id.ts b/apps/web/hooks/use-environment-session-id.ts index dfed910bd3..99671596a8 100644 --- a/apps/web/hooks/use-environment-session-id.ts +++ b/apps/web/hooks/use-environment-session-id.ts @@ -9,10 +9,9 @@ type StoreSlice = { /** * Return a sessionId that only changes when the underlying environment changes. * - * Many backend endpoints (terminal WS, file tree, agentctl) are routed by - * sessionId but resolve to a shared per-environment resource. This hook - * keeps the returned sessionId stable across same-environment tab switches - * so downstream components don't needlessly reconnect or re-fetch. + * Some workspace readiness APIs still need a session lookup handle. This hook + * keeps that handle stable across same-environment tab switches so downstream + * components don't needlessly reconnect or re-fetch. */ export function useEnvironmentSessionId(): string | null { const cacheRef = useRef({ envId: null as string | null, sessionId: null as string | null }); @@ -25,3 +24,11 @@ export function useEnvironmentSessionId(): string | null { }, []); return useAppStore(selector); } + +/** Return the active task environment ID without falling back to session ID. */ +export function useEnvironmentId(): string | null { + return useAppStore((state: StoreSlice) => { + const sid = state.tasks.activeSessionId; + return sid ? (state.environmentIdBySessionId[sid] ?? null) : null; + }); +} diff --git a/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts b/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts index f2fb08b9fa..d987811d16 100644 --- a/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts +++ b/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts @@ -24,6 +24,16 @@ function makeCombinedStore() { } describe("registerSessionEnvironment — migrateEnvKeyedData", () => { + it("does not create user shell state under a session fallback key", () => { + const store = makeStore(); + + store.getState().setUserShells("sess-missing-env", [{ terminalId: "t1" } as never]); + + const state = store.getState(); + expect(state.userShells.byEnvironmentId["sess-missing-env"]).toBeUndefined(); + expect(state.userShells.loaded["sess-missing-env"]).toBeUndefined(); + }); + it("migrates data from sessionId key to environmentId key", () => { const store = makeStore(); @@ -90,7 +100,9 @@ describe("registerSessionEnvironment — migrateEnvKeyedData", () => { // Only put data in shell outputs and userShells, leave others empty store.getState().appendShellOutput("sess-3", "output"); - store.getState().setUserShells("sess-3", [{ terminalId: "t1" } as never]); + store.setState((draft) => { + draft.userShells.byEnvironmentId["sess-3"] = [{ terminalId: "t1" } as never]; + }); store.getState().registerSessionEnvironment("sess-3", "env-3"); diff --git a/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts b/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts index 3f967df40d..533b552776 100644 --- a/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts +++ b/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts @@ -152,25 +152,30 @@ function buildTerminalShellProcessActions(set: ImmerSet) { } function buildUserShellActions(set: ImmerSet) { + const getEnvKey = (draft: SessionRuntimeSliceState, sessionId: string) => + draft.environmentIdBySessionId[sessionId]; return { setUserShells: ( sessionId: string, shells: Parameters[1], ) => set((draft) => { - const envKey = draft.environmentIdBySessionId[sessionId] ?? sessionId; + const envKey = getEnvKey(draft, sessionId); + if (!envKey) return; draft.userShells.byEnvironmentId[envKey] = shells; draft.userShells.loaded[envKey] = true; draft.userShells.loading[envKey] = false; }), setUserShellsLoading: (sessionId: string, loading: boolean) => set((draft) => { - const envKey = draft.environmentIdBySessionId[sessionId] ?? sessionId; + const envKey = getEnvKey(draft, sessionId); + if (!envKey) return; draft.userShells.loading[envKey] = loading; }), addUserShell: (sessionId: string, shell: Parameters[1]) => set((draft) => { - const envKey = draft.environmentIdBySessionId[sessionId] ?? sessionId; + const envKey = getEnvKey(draft, sessionId); + if (!envKey) return; const existing = draft.userShells.byEnvironmentId[envKey] || []; if (!existing.some((s) => s.terminalId === shell.terminalId)) { draft.userShells.byEnvironmentId[envKey] = [...existing, shell]; @@ -178,7 +183,8 @@ function buildUserShellActions(set: ImmerSet) { }), removeUserShell: (sessionId: string, terminalId: string) => set((draft) => { - const envKey = draft.environmentIdBySessionId[sessionId] ?? sessionId; + const envKey = getEnvKey(draft, sessionId); + if (!envKey) return; const existing = draft.userShells.byEnvironmentId[envKey] || []; draft.userShells.byEnvironmentId[envKey] = existing.filter( (s) => s.terminalId !== terminalId, diff --git a/apps/web/lib/state/slices/session-runtime/types.ts b/apps/web/lib/state/slices/session-runtime/types.ts index b3ef89cdf1..317d373843 100644 --- a/apps/web/lib/state/slices/session-runtime/types.ts +++ b/apps/web/lib/state/slices/session-runtime/types.ts @@ -215,8 +215,7 @@ export type UserShellInfo = { }; export type UserShellsState = { - /** User shells keyed by environmentId (shared across sessions in the same environment). - * Falls back to sessionId when no environment mapping exists. */ + /** User shells keyed by environmentId (shared across sessions in the same environment). */ byEnvironmentId: Record; /** Keyed by environmentId (same key strategy as byEnvironmentId). */ loading: Record;