fix: prevent invalid MCP task sessions (#1494)
This commit is contained in:
Carlos Florêncio
@@ -26,6 +26,7 @@ import (
|
||||
taskrepo "github.com/kandev/kandev/internal/task/repository/sqlite"
|
||||
"github.com/kandev/kandev/internal/task/service"
|
||||
workflowctrl "github.com/kandev/kandev/internal/workflow/controller"
|
||||
workflowmodels "github.com/kandev/kandev/internal/workflow/models"
|
||||
workflowsvc "github.com/kandev/kandev/internal/workflow/service"
|
||||
v1 "github.com/kandev/kandev/pkg/api/v1"
|
||||
ws "github.com/kandev/kandev/pkg/websocket"
|
||||
@@ -459,6 +460,23 @@ func (h *Handlers) handleCreateTask(ctx context.Context, msg *ws.Message) (*ws.M
|
||||
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "workflow_id is required", nil)
|
||||
}
|
||||
|
||||
var autoStartConfig *mcpAutoStartConfig
|
||||
if startAgent && h.sessionLauncher != nil {
|
||||
pendingTask := &models.Task{
|
||||
ParentID: req.ParentID,
|
||||
WorkspaceID: req.WorkspaceID,
|
||||
WorkflowID: req.WorkflowID,
|
||||
WorkflowStepID: req.WorkflowStepID,
|
||||
}
|
||||
config := h.resolveMCPAutoStartConfig(ctx, pendingTask, req.AgentProfileID, req.ExecutorProfileID, req.SourceTaskID)
|
||||
if config.AgentProfileID == "" {
|
||||
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation,
|
||||
"agent_profile_id is required when start_agent is true and no agent profile can be resolved from the parent task, source task, workflow, or workspace defaults; pass agent_profile_id or set start_agent=false",
|
||||
nil)
|
||||
}
|
||||
autoStartConfig = &config
|
||||
}
|
||||
|
||||
task, err := h.taskSvc.CreateTask(ctx, &service.CreateTaskRequest{
|
||||
ParentID: req.ParentID,
|
||||
WorkspaceID: req.WorkspaceID,
|
||||
@@ -481,8 +499,12 @@ func (h *Handlers) handleCreateTask(ctx context.Context, msg *ws.Message) (*ws.M
|
||||
}
|
||||
|
||||
// Auto-start agent session asynchronously only if requested
|
||||
if startAgent {
|
||||
h.autoStartTask(task, req.AgentProfileID, req.ExecutorProfileID, req.SourceTaskID)
|
||||
if startAgent && h.sessionLauncher != nil {
|
||||
if autoStartConfig == nil {
|
||||
config := h.resolveMCPAutoStartConfig(ctx, task, req.AgentProfileID, req.ExecutorProfileID, req.SourceTaskID)
|
||||
autoStartConfig = &config
|
||||
}
|
||||
h.launchAutoStartTask(ctx, task, *autoStartConfig)
|
||||
}
|
||||
|
||||
return ws.NewResponse(msg.ID, msg.Action, dto.FromTask(task))
|
||||
@@ -625,28 +647,43 @@ func inheritedRepoInputs(src []*models.TaskRepository) []service.TaskRepositoryI
|
||||
return repos
|
||||
}
|
||||
|
||||
type mcpAutoStartConfig struct {
|
||||
AgentProfileID string
|
||||
ExecutorID string
|
||||
ExecutorProfileID string
|
||||
}
|
||||
|
||||
// autoStartTask launches an agent session for a newly created task in the background.
|
||||
// It resolves the agent profile: explicit > parent's session > source task's session > workspace default.
|
||||
// It resolves the executor: explicit executor_profile_id > parent's executor_profile_id >
|
||||
// source task's executor_profile_id > parent's executor_id > "exec-worktree" (default for MCP-created tasks).
|
||||
// It is kept as a small compatibility wrapper for direct tests; handleCreateTask
|
||||
// uses resolveMCPAutoStartConfig before persisting so invalid auto-start
|
||||
// requests fail synchronously.
|
||||
func (h *Handlers) autoStartTask(task *models.Task, agentProfileID, executorProfileID, sourceTaskID string) {
|
||||
if h.sessionLauncher == nil {
|
||||
return
|
||||
}
|
||||
ctx := context.Background()
|
||||
config := h.resolveMCPAutoStartConfig(ctx, task, agentProfileID, executorProfileID, sourceTaskID)
|
||||
h.launchAutoStartTask(ctx, task, config)
|
||||
}
|
||||
|
||||
executorID := h.inheritFromParentSession(task.ParentID, &agentProfileID, &executorProfileID)
|
||||
// resolveMCPAutoStartConfig resolves the agent profile and executor for
|
||||
// create_task_kandev auto-start. Agent profile resolution follows the MCP
|
||||
// ergonomics first (explicit > parent/source session), then the same durable
|
||||
// defaults used by task opening where this handler can see them (workflow step
|
||||
// when a workflow controller is wired, workflow default, workspace default).
|
||||
func (h *Handlers) resolveMCPAutoStartConfig(ctx context.Context, task *models.Task, agentProfileID, executorProfileID, sourceTaskID string) mcpAutoStartConfig {
|
||||
executorID := h.inheritFromTaskSession(ctx, task.ParentID, &agentProfileID, &executorProfileID)
|
||||
|
||||
// For top-level tasks, inherit from the source task (the calling agent's task)
|
||||
// For top-level tasks, inherit from the source task (the calling agent's task).
|
||||
if task.ParentID == "" && sourceTaskID != "" {
|
||||
sourceExecutorID := h.inheritFromParentSession(sourceTaskID, &agentProfileID, &executorProfileID)
|
||||
sourceExecutorID := h.inheritFromTaskSession(ctx, sourceTaskID, &agentProfileID, &executorProfileID)
|
||||
if executorID == "" {
|
||||
executorID = sourceExecutorID
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to workspace defaults for agent profile and worktree executor
|
||||
if agentProfileID == "" {
|
||||
workspace, err := h.taskSvc.GetWorkspace(context.Background(), task.WorkspaceID)
|
||||
agentProfileID = h.resolveWorkflowAgentProfile(ctx, task.WorkflowStepID, task.WorkflowID)
|
||||
}
|
||||
if agentProfileID == "" && h.taskSvc != nil {
|
||||
workspace, err := h.taskSvc.GetWorkspace(ctx, task.WorkspaceID)
|
||||
if err == nil && workspace.DefaultAgentProfileID != nil {
|
||||
agentProfileID = *workspace.DefaultAgentProfileID
|
||||
}
|
||||
@@ -655,22 +692,110 @@ func (h *Handlers) autoStartTask(task *models.Task, agentProfileID, executorProf
|
||||
executorID = models.ExecutorIDWorktree
|
||||
}
|
||||
|
||||
if agentProfileID == "" {
|
||||
return mcpAutoStartConfig{
|
||||
AgentProfileID: agentProfileID,
|
||||
ExecutorID: executorID,
|
||||
ExecutorProfileID: executorProfileID,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handlers) resolveWorkflowAgentProfile(ctx context.Context, workflowStepID, workflowID string) string {
|
||||
profileID, resolvedWorkflowID := h.resolveWorkflowControllerAgentProfile(ctx, workflowStepID, workflowID)
|
||||
if profileID != "" {
|
||||
return profileID
|
||||
}
|
||||
if workflowID == "" {
|
||||
workflowID = resolvedWorkflowID
|
||||
}
|
||||
return h.workflowDefaultAgentProfile(ctx, workflowID)
|
||||
}
|
||||
|
||||
func (h *Handlers) resolveWorkflowControllerAgentProfile(ctx context.Context, workflowStepID, workflowID string) (string, string) {
|
||||
if h.workflowCtrl == nil {
|
||||
return "", workflowID
|
||||
}
|
||||
if workflowStepID != "" {
|
||||
return h.resolveWorkflowStepAgentProfile(ctx, workflowStepID, workflowID)
|
||||
}
|
||||
if workflowID == "" {
|
||||
return "", ""
|
||||
}
|
||||
return h.resolveWorkflowStartStepAgentProfile(ctx, workflowID), workflowID
|
||||
}
|
||||
|
||||
func (h *Handlers) resolveWorkflowStepAgentProfile(ctx context.Context, workflowStepID, workflowID string) (string, string) {
|
||||
resp, err := h.workflowCtrl.GetStep(ctx, workflowStepID)
|
||||
if err != nil || resp == nil || resp.Step == nil {
|
||||
return "", workflowID
|
||||
}
|
||||
if workflowID == "" {
|
||||
workflowID = resp.Step.WorkflowID
|
||||
}
|
||||
return resp.Step.AgentProfileID, workflowID
|
||||
}
|
||||
|
||||
func (h *Handlers) resolveWorkflowStartStepAgentProfile(ctx context.Context, workflowID string) string {
|
||||
resp, err := h.workflowCtrl.ListStepsByWorkflow(ctx, workflowctrl.ListStepsRequest{WorkflowID: workflowID})
|
||||
if err != nil || resp == nil {
|
||||
return ""
|
||||
}
|
||||
return startStepAgentProfile(resp.Steps)
|
||||
}
|
||||
|
||||
func (h *Handlers) workflowDefaultAgentProfile(ctx context.Context, workflowID string) string {
|
||||
if workflowID == "" || h.taskSvc == nil {
|
||||
return ""
|
||||
}
|
||||
workflow, err := h.taskSvc.GetWorkflow(ctx, workflowID)
|
||||
if err != nil || workflow == nil {
|
||||
return ""
|
||||
}
|
||||
return workflow.AgentProfileID
|
||||
}
|
||||
|
||||
func startStepAgentProfile(steps []*workflowmodels.WorkflowStep) string {
|
||||
if len(steps) == 0 {
|
||||
return ""
|
||||
}
|
||||
var firstByPosition *workflowmodels.WorkflowStep
|
||||
for _, step := range steps {
|
||||
if step == nil {
|
||||
continue
|
||||
}
|
||||
if firstByPosition == nil || step.Position < firstByPosition.Position {
|
||||
firstByPosition = step
|
||||
}
|
||||
if step.IsStartStep {
|
||||
return step.AgentProfileID
|
||||
}
|
||||
}
|
||||
if firstByPosition == nil {
|
||||
return ""
|
||||
}
|
||||
return firstByPosition.AgentProfileID
|
||||
}
|
||||
|
||||
func (h *Handlers) launchAutoStartTask(ctx context.Context, task *models.Task, config mcpAutoStartConfig) {
|
||||
if h.sessionLauncher == nil {
|
||||
return
|
||||
}
|
||||
if config.AgentProfileID == "" {
|
||||
h.logger.Warn("no agent profile available, skipping auto-start",
|
||||
zap.String("task_id", task.ID))
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), constants.AgentLaunchTimeout)
|
||||
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), constants.AgentLaunchTimeout)
|
||||
defer cancel()
|
||||
|
||||
resp, err := h.sessionLauncher.LaunchSession(ctx, &orchestrator.LaunchSessionRequest{
|
||||
TaskID: task.ID,
|
||||
Intent: orchestrator.IntentStart,
|
||||
AgentProfileID: agentProfileID,
|
||||
ExecutorID: executorID,
|
||||
ExecutorProfileID: executorProfileID,
|
||||
AgentProfileID: config.AgentProfileID,
|
||||
ExecutorID: config.ExecutorID,
|
||||
ExecutorProfileID: config.ExecutorProfileID,
|
||||
WorkflowStepID: task.WorkflowStepID,
|
||||
Prompt: task.Description,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -684,16 +809,19 @@ func (h *Handlers) autoStartTask(task *models.Task, agentProfileID, executorProf
|
||||
}()
|
||||
}
|
||||
|
||||
// inheritFromParentSession fills agentProfileID and executorProfileID from the parent
|
||||
// task's primary session when not explicitly provided. It returns the parent's ExecutorID
|
||||
// inheritFromTaskSession fills agentProfileID and executorProfileID from another
|
||||
// task's primary session when not explicitly provided. It returns that session's ExecutorID
|
||||
// as a fallback for when the parent session has no executor profile (common for
|
||||
// UI-created sessions). If ExecutorProfileID is resolved, ExecutorID is redundant
|
||||
// since the profile already encodes the executor reference.
|
||||
func (h *Handlers) inheritFromParentSession(parentID string, agentProfileID, executorProfileID *string) string {
|
||||
func (h *Handlers) inheritFromTaskSession(ctx context.Context, parentID string, agentProfileID, executorProfileID *string) string {
|
||||
if parentID == "" {
|
||||
return ""
|
||||
}
|
||||
parent, err := h.taskSvc.GetPrimarySession(context.Background(), parentID)
|
||||
if h.taskSvc == nil {
|
||||
return ""
|
||||
}
|
||||
parent, err := h.taskSvc.GetPrimarySession(ctx, parentID)
|
||||
if err != nil || parent == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -21,6 +21,10 @@ import (
|
||||
"github.com/kandev/kandev/internal/task/repository"
|
||||
sqliterepo "github.com/kandev/kandev/internal/task/repository/sqlite"
|
||||
"github.com/kandev/kandev/internal/task/service"
|
||||
workflowcontroller "github.com/kandev/kandev/internal/workflow/controller"
|
||||
workflowmodels "github.com/kandev/kandev/internal/workflow/models"
|
||||
workflowrepo "github.com/kandev/kandev/internal/workflow/repository"
|
||||
workflowservice "github.com/kandev/kandev/internal/workflow/service"
|
||||
"github.com/kandev/kandev/internal/worktree"
|
||||
v1 "github.com/kandev/kandev/pkg/api/v1"
|
||||
ws "github.com/kandev/kandev/pkg/websocket"
|
||||
@@ -64,6 +68,45 @@ func newTestTaskService(t *testing.T) (*service.Service, *sqliterepo.Repository)
|
||||
return svc, repo
|
||||
}
|
||||
|
||||
func newTestTaskServiceWithWorkflow(t *testing.T) (*service.Service, *sqliterepo.Repository, *workflowcontroller.Controller, *workflowrepo.Repository) {
|
||||
t.Helper()
|
||||
dbConn, err := db.OpenSQLite(filepath.Join(t.TempDir(), "test.db"))
|
||||
require.NoError(t, err)
|
||||
sqlxDB := sqlx.NewDb(dbConn, "sqlite3")
|
||||
t.Cleanup(func() {
|
||||
_ = sqlxDB.Close()
|
||||
})
|
||||
repo, cleanup, err := repository.Provide(sqlxDB, sqlxDB, nil)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_ = cleanup()
|
||||
})
|
||||
workflowRepo, err := workflowrepo.NewWithDB(sqlxDB, sqlxDB, testLogger(t))
|
||||
require.NoError(t, err)
|
||||
_, err = worktree.NewSQLiteStore(sqlxDB, sqlxDB)
|
||||
require.NoError(t, err)
|
||||
|
||||
log, _ := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "json"})
|
||||
eventBus := bus.NewMemoryEventBus(log)
|
||||
t.Cleanup(func() { eventBus.Close() })
|
||||
svc := service.NewService(service.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,
|
||||
}, eventBus, log, service.RepositoryDiscoveryConfig{})
|
||||
workflowSvc := workflowservice.NewService(workflowRepo, log)
|
||||
return svc, repo, workflowcontroller.NewController(workflowSvc), workflowRepo
|
||||
}
|
||||
|
||||
func seedMCPHandlerSession(t *testing.T, repo *sqliterepo.Repository, taskID, sessionID string, state models.TaskSessionState) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
@@ -264,21 +307,231 @@ func TestHandleCreateTask_TopLevel_MissingWorkflowID(t *testing.T) {
|
||||
assertWSError(t, resp, ws.ErrorCodeValidation)
|
||||
}
|
||||
|
||||
func TestHandleCreateTask_StartAgentRequiresResolvableAgentProfile(t *testing.T) {
|
||||
svc, _ := newTestTaskService(t)
|
||||
ctx := context.Background()
|
||||
workspaces, err := svc.ListWorkspaces(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, workspaces, 1)
|
||||
workflows, err := svc.ListWorkflows(ctx, workspaces[0].ID, false)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, workflows, 1)
|
||||
|
||||
h := &Handlers{
|
||||
taskSvc: svc,
|
||||
sessionLauncher: newMockSessionLauncher(),
|
||||
logger: testLogger(t).WithFields(),
|
||||
}
|
||||
msg := makeWSMessage(t, ws.ActionMCPCreateTask, map[string]interface{}{
|
||||
"workspace_id": workspaces[0].ID,
|
||||
"workflow_id": workflows[0].ID,
|
||||
"title": "Invalid auto-start task",
|
||||
"description": "This should not be persisted without a profile.",
|
||||
})
|
||||
|
||||
resp, err := h.handleCreateTask(ctx, msg)
|
||||
require.NoError(t, err)
|
||||
assertWSError(t, resp, ws.ErrorCodeValidation)
|
||||
|
||||
tasks, err := svc.ListTasks(ctx, workflows[0].ID)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, tasks, "failed auto-start preflight must not leave a broken task behind")
|
||||
}
|
||||
|
||||
func TestHandleCreateTask_StartAgentUsesWorkspaceDefaultAgentProfile(t *testing.T) {
|
||||
svc, _ := newTestTaskService(t)
|
||||
ctx := context.Background()
|
||||
workspaces, err := svc.ListWorkspaces(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, workspaces, 1)
|
||||
workflows, err := svc.ListWorkflows(ctx, workspaces[0].ID, false)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, workflows, 1)
|
||||
|
||||
defaultProfileID := "workspace-default-profile"
|
||||
_, err = svc.UpdateWorkspace(ctx, workspaces[0].ID, &service.UpdateWorkspaceRequest{
|
||||
DefaultAgentProfileID: &defaultProfileID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
launcher := newMockSessionLauncher()
|
||||
h := &Handlers{
|
||||
taskSvc: svc,
|
||||
sessionLauncher: launcher,
|
||||
logger: testLogger(t).WithFields(),
|
||||
}
|
||||
msg := makeWSMessage(t, ws.ActionMCPCreateTask, map[string]interface{}{
|
||||
"workspace_id": workspaces[0].ID,
|
||||
"workflow_id": workflows[0].ID,
|
||||
"title": "Valid auto-start task",
|
||||
"description": "This should launch with the workspace default profile.",
|
||||
})
|
||||
|
||||
resp, err := h.handleCreateTask(ctx, msg)
|
||||
require.NoError(t, err)
|
||||
require.Equalf(t, ws.MessageTypeResponse, resp.Type, "create_task should succeed; payload: %s", string(resp.Payload))
|
||||
|
||||
select {
|
||||
case <-launcher.called:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("LaunchSession was not called within timeout")
|
||||
}
|
||||
req := launcher.getRequest()
|
||||
assert.Equal(t, defaultProfileID, req.AgentProfileID)
|
||||
}
|
||||
|
||||
func TestResolveMCPAutoStartConfig_UsesWorkflowStepAgentProfile(t *testing.T) {
|
||||
svc, _, workflowCtrl, workflowRepo := newTestTaskServiceWithWorkflow(t)
|
||||
ctx := context.Background()
|
||||
workspace, workflow := defaultWorkspaceAndWorkflow(t, ctx, svc)
|
||||
step := seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "Implement",
|
||||
Position: 1,
|
||||
AgentProfileID: "step-profile",
|
||||
AllowManualMove: true,
|
||||
})
|
||||
h := &Handlers{taskSvc: svc, workflowCtrl: workflowCtrl, logger: testLogger(t).WithFields()}
|
||||
|
||||
config := h.resolveMCPAutoStartConfig(ctx, &models.Task{
|
||||
WorkspaceID: workspace.ID,
|
||||
WorkflowID: workflow.ID,
|
||||
WorkflowStepID: step.ID,
|
||||
}, "", "", "")
|
||||
|
||||
assert.Equal(t, "step-profile", config.AgentProfileID)
|
||||
}
|
||||
|
||||
func TestResolveMCPAutoStartConfig_UsesLowestPositionStepAgentProfileWhenNoStartStep(t *testing.T) {
|
||||
svc, _, workflowCtrl, workflowRepo := newTestTaskServiceWithWorkflow(t)
|
||||
ctx := context.Background()
|
||||
workspace, workflow := defaultWorkspaceAndWorkflow(t, ctx, svc)
|
||||
seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "Later",
|
||||
Position: 10,
|
||||
AgentProfileID: "later-profile",
|
||||
AllowManualMove: true,
|
||||
})
|
||||
seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "First",
|
||||
Position: 2,
|
||||
AgentProfileID: "first-profile",
|
||||
AllowManualMove: true,
|
||||
})
|
||||
h := &Handlers{taskSvc: svc, workflowCtrl: workflowCtrl, logger: testLogger(t).WithFields()}
|
||||
|
||||
config := h.resolveMCPAutoStartConfig(ctx, &models.Task{
|
||||
WorkspaceID: workspace.ID,
|
||||
WorkflowID: workflow.ID,
|
||||
}, "", "", "")
|
||||
|
||||
assert.Equal(t, "first-profile", config.AgentProfileID)
|
||||
}
|
||||
|
||||
func TestStartStepAgentProfile_UsesLowestPositionWhenNoStartStep(t *testing.T) {
|
||||
profileID := startStepAgentProfile([]*workflowmodels.WorkflowStep{
|
||||
{Position: 10, AgentProfileID: "later-profile"},
|
||||
nil,
|
||||
{Position: 2, AgentProfileID: "first-profile"},
|
||||
})
|
||||
|
||||
assert.Equal(t, "first-profile", profileID)
|
||||
}
|
||||
|
||||
func TestResolveMCPAutoStartConfig_UsesMarkedStartStepAgentProfile(t *testing.T) {
|
||||
svc, _, workflowCtrl, workflowRepo := newTestTaskServiceWithWorkflow(t)
|
||||
ctx := context.Background()
|
||||
workspace, workflow := defaultWorkspaceAndWorkflow(t, ctx, svc)
|
||||
seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "First",
|
||||
Position: 1,
|
||||
AgentProfileID: "first-profile",
|
||||
AllowManualMove: true,
|
||||
})
|
||||
seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "Start Here",
|
||||
Position: 2,
|
||||
IsStartStep: true,
|
||||
AgentProfileID: "start-profile",
|
||||
AllowManualMove: true,
|
||||
})
|
||||
h := &Handlers{taskSvc: svc, workflowCtrl: workflowCtrl, logger: testLogger(t).WithFields()}
|
||||
|
||||
config := h.resolveMCPAutoStartConfig(ctx, &models.Task{
|
||||
WorkspaceID: workspace.ID,
|
||||
WorkflowID: workflow.ID,
|
||||
}, "", "", "")
|
||||
|
||||
assert.Equal(t, "start-profile", config.AgentProfileID)
|
||||
}
|
||||
|
||||
func TestResolveMCPAutoStartConfig_FallsBackToWorkflowAgentProfile(t *testing.T) {
|
||||
svc, _, workflowCtrl, workflowRepo := newTestTaskServiceWithWorkflow(t)
|
||||
ctx := context.Background()
|
||||
workspace, workflow := defaultWorkspaceAndWorkflow(t, ctx, svc)
|
||||
workflowProfileID := "workflow-profile"
|
||||
_, err := svc.UpdateWorkflow(ctx, workflow.ID, &service.UpdateWorkflowRequest{
|
||||
AgentProfileID: &workflowProfileID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
seedWorkflowStep(t, ctx, workflowRepo, &workflowmodels.WorkflowStep{
|
||||
WorkflowID: workflow.ID,
|
||||
Name: "Unassigned",
|
||||
Position: 1,
|
||||
AllowManualMove: true,
|
||||
})
|
||||
h := &Handlers{taskSvc: svc, workflowCtrl: workflowCtrl, logger: testLogger(t).WithFields()}
|
||||
|
||||
config := h.resolveMCPAutoStartConfig(ctx, &models.Task{
|
||||
WorkspaceID: workspace.ID,
|
||||
WorkflowID: workflow.ID,
|
||||
}, "", "", "")
|
||||
|
||||
assert.Equal(t, workflowProfileID, config.AgentProfileID)
|
||||
}
|
||||
|
||||
func defaultWorkspaceAndWorkflow(t *testing.T, ctx context.Context, svc *service.Service) (*models.Workspace, *models.Workflow) {
|
||||
t.Helper()
|
||||
workspaces, err := svc.ListWorkspaces(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, workspaces, 1)
|
||||
workflow, err := svc.CreateWorkflow(ctx, &service.CreateWorkflowRequest{
|
||||
WorkspaceID: workspaces[0].ID,
|
||||
Name: "Workflow under test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return workspaces[0], workflow
|
||||
}
|
||||
|
||||
func seedWorkflowStep(t *testing.T, ctx context.Context, repo *workflowrepo.Repository, step *workflowmodels.WorkflowStep) *workflowmodels.WorkflowStep {
|
||||
t.Helper()
|
||||
require.NoError(t, repo.CreateStep(ctx, step))
|
||||
return step
|
||||
}
|
||||
|
||||
// mockSessionLauncher captures LaunchSession calls for testing autoStartTask.
|
||||
type mockSessionLauncher struct {
|
||||
mu sync.Mutex
|
||||
req *orchestrator.LaunchSessionRequest
|
||||
called chan struct{}
|
||||
mu sync.Mutex
|
||||
req *orchestrator.LaunchSessionRequest
|
||||
called chan struct{}
|
||||
inspect func(context.Context)
|
||||
}
|
||||
|
||||
func newMockSessionLauncher() *mockSessionLauncher {
|
||||
return &mockSessionLauncher{called: make(chan struct{})}
|
||||
}
|
||||
|
||||
func (m *mockSessionLauncher) LaunchSession(_ context.Context, req *orchestrator.LaunchSessionRequest) (*orchestrator.LaunchSessionResponse, error) {
|
||||
func (m *mockSessionLauncher) LaunchSession(ctx context.Context, req *orchestrator.LaunchSessionRequest) (*orchestrator.LaunchSessionResponse, error) {
|
||||
m.mu.Lock()
|
||||
m.req = req
|
||||
m.mu.Unlock()
|
||||
if m.inspect != nil {
|
||||
m.inspect(ctx)
|
||||
}
|
||||
close(m.called)
|
||||
return &orchestrator.LaunchSessionResponse{
|
||||
Success: true,
|
||||
@@ -316,8 +569,9 @@ func TestAutoStartTask_DefaultsToWorktreeExecutor(t *testing.T) {
|
||||
}
|
||||
|
||||
task := &models.Task{
|
||||
ID: "task-1",
|
||||
WorkspaceID: "ws-1",
|
||||
ID: "task-1",
|
||||
WorkspaceID: "ws-1",
|
||||
WorkflowStepID: "step-1",
|
||||
}
|
||||
|
||||
// Call with agent profile but no executor info
|
||||
@@ -333,6 +587,7 @@ func TestAutoStartTask_DefaultsToWorktreeExecutor(t *testing.T) {
|
||||
assert.Equal(t, models.ExecutorIDWorktree, req.ExecutorID, "should default to exec-worktree")
|
||||
assert.Equal(t, "", req.ExecutorProfileID)
|
||||
assert.Equal(t, "agent-profile-1", req.AgentProfileID)
|
||||
assert.Equal(t, "step-1", req.WorkflowStepID)
|
||||
}
|
||||
|
||||
func TestAutoStartTask_ExplicitExecutorProfilePreserved(t *testing.T) {
|
||||
@@ -362,6 +617,45 @@ func TestAutoStartTask_ExplicitExecutorProfilePreserved(t *testing.T) {
|
||||
assert.Equal(t, "", req.ExecutorID, "executorID should be empty when profile is set")
|
||||
}
|
||||
|
||||
func TestLaunchAutoStartTask_PreservesContextValuesWithoutCallerCancellation(t *testing.T) {
|
||||
type contextKey string
|
||||
|
||||
launcher := newMockSessionLauncher()
|
||||
log := testLogger(t)
|
||||
h := &Handlers{
|
||||
sessionLauncher: launcher,
|
||||
logger: log.WithFields(),
|
||||
}
|
||||
|
||||
key := contextKey("request-id")
|
||||
var capturedValue any
|
||||
var capturedErr error
|
||||
launcher.inspect = func(ctx context.Context) {
|
||||
capturedValue = ctx.Value(key)
|
||||
capturedErr = ctx.Err()
|
||||
}
|
||||
|
||||
parentCtx, cancel := context.WithCancel(context.WithValue(context.Background(), key, "request-1"))
|
||||
cancel()
|
||||
|
||||
h.launchAutoStartTask(parentCtx, &models.Task{
|
||||
ID: "task-1",
|
||||
WorkspaceID: "ws-1",
|
||||
}, mcpAutoStartConfig{
|
||||
AgentProfileID: "agent-profile-1",
|
||||
ExecutorID: models.ExecutorIDWorktree,
|
||||
})
|
||||
|
||||
select {
|
||||
case <-launcher.called:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("LaunchSession was not called within timeout")
|
||||
}
|
||||
|
||||
assert.Equal(t, "request-1", capturedValue)
|
||||
assert.NoError(t, capturedErr)
|
||||
}
|
||||
|
||||
func TestResolveTaskRepositories_ExplicitRepos(t *testing.T) {
|
||||
log := testLogger(t)
|
||||
h := &Handlers{logger: log.WithFields()}
|
||||
|
||||
@@ -25,8 +25,10 @@ vi.mock("@/lib/links", () => ({
|
||||
|
||||
import { launchSession, type LaunchSessionResponse } from "@/lib/services/session-launch-service";
|
||||
import { performLayoutSwitch, releaseLayoutToDefault } from "@/lib/state/dockview-store";
|
||||
import { replaceTaskUrl } from "@/lib/links";
|
||||
import type { StoreApi } from "zustand";
|
||||
import type { AppState } from "@/lib/state/store";
|
||||
import type { TaskSession } from "@/lib/types/http";
|
||||
|
||||
const NEW_TASK_ID = "task-new";
|
||||
const OLD_SESSION_ID = "old-session";
|
||||
@@ -50,6 +52,123 @@ function makeEnvStore(envIds: Record<string, string>): StoreApi<AppState> {
|
||||
} as unknown as StoreApi<AppState>;
|
||||
}
|
||||
|
||||
const TASK_ID = "task-A";
|
||||
const PRIMARY = "sess-primary";
|
||||
const ORIGINAL_TASK_ID = "task-original";
|
||||
const PENDING_TASK_ID = "task-pending";
|
||||
const ORIGINAL_SESSION_ID = "sess-original";
|
||||
const PENDING_SESSION_ID = "sess-pending";
|
||||
const ORIGINAL_ENV_ID = "env-original";
|
||||
|
||||
function makeKanbanStore(args: {
|
||||
activeTaskId?: string | null;
|
||||
activeSessionId: string | null;
|
||||
envIds: Record<string, string>;
|
||||
lastSessionByTaskId?: Record<string, string>;
|
||||
sessionTaskIds?: Record<string, string>;
|
||||
}): StoreApi<AppState> {
|
||||
const items: Record<string, { id: string; task_id: string }> = {};
|
||||
for (const [sid, tid] of Object.entries(args.sessionTaskIds ?? {})) {
|
||||
items[sid] = { id: sid, task_id: tid };
|
||||
}
|
||||
const state = {
|
||||
tasks: {
|
||||
activeTaskId: args.activeTaskId ?? null,
|
||||
activeSessionId: args.activeSessionId,
|
||||
lastSessionByTaskId: args.lastSessionByTaskId ?? {},
|
||||
},
|
||||
taskPRs: { byTaskId: {} as Record<string, unknown[]> },
|
||||
environmentIdBySessionId: args.envIds,
|
||||
taskSessions: { items },
|
||||
};
|
||||
return {
|
||||
getState: () => state as unknown as AppState,
|
||||
setState: vi.fn(),
|
||||
subscribe: vi.fn(),
|
||||
} as unknown as StoreApi<AppState>;
|
||||
}
|
||||
|
||||
function runSelect(store: StoreApi<AppState>) {
|
||||
const switchToSession = vi.fn();
|
||||
selectTaskWithLayout({
|
||||
taskId: TASK_ID,
|
||||
task: { primarySessionId: PRIMARY },
|
||||
store,
|
||||
switchToSession,
|
||||
loadTaskSessionsForTask: vi.fn(async () => []),
|
||||
setActiveTask: vi.fn(),
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
return switchToSession;
|
||||
}
|
||||
|
||||
async function flushTaskSelection() {
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
}
|
||||
|
||||
function makeSelectionHarness(args: {
|
||||
activeTaskId: string;
|
||||
activeSessionId: string | null;
|
||||
envIds?: Record<string, string>;
|
||||
sessions?: Record<string, { id: string; task_id: string }>;
|
||||
}) {
|
||||
const listeners: Array<(state: AppState, previousState: AppState) => void> = [];
|
||||
const state = {
|
||||
tasks: {
|
||||
activeTaskId: args.activeTaskId,
|
||||
activeSessionId: args.activeSessionId,
|
||||
lastSessionByTaskId: {},
|
||||
},
|
||||
taskPRs: { byTaskId: {} as Record<string, unknown[]> },
|
||||
environmentIdBySessionId: args.envIds ?? {},
|
||||
taskSessions: { items: (args.sessions ?? {}) as Record<string, TaskSession> },
|
||||
};
|
||||
const snapshot = () =>
|
||||
({
|
||||
...state,
|
||||
tasks: { ...state.tasks },
|
||||
}) as unknown as AppState;
|
||||
const notify = (previousState: AppState) => {
|
||||
for (const listener of listeners) {
|
||||
listener(state as unknown as AppState, previousState);
|
||||
}
|
||||
};
|
||||
const store = {
|
||||
getState: () => state as unknown as AppState,
|
||||
setState: vi.fn(),
|
||||
subscribe: vi.fn((listener: (state: AppState, previousState: AppState) => void) => {
|
||||
listeners.push(listener);
|
||||
return () => {
|
||||
const index = listeners.indexOf(listener);
|
||||
if (index >= 0) listeners.splice(index, 1);
|
||||
};
|
||||
}),
|
||||
} as unknown as StoreApi<AppState>;
|
||||
const setActiveTask = vi.fn((taskId: string) => {
|
||||
const previousState = snapshot();
|
||||
state.tasks.activeTaskId = taskId;
|
||||
state.tasks.activeSessionId = null;
|
||||
notify(previousState);
|
||||
});
|
||||
return { state, store, setActiveTask, getListenerCount: () => listeners.length };
|
||||
}
|
||||
|
||||
function makeDeferredSessionLoader() {
|
||||
let resolveLoad: (sessions: TaskSession[]) => void = () => {};
|
||||
const loadTaskSessionsForTask = vi.fn(
|
||||
() =>
|
||||
new Promise<TaskSession[]>((resolve) => {
|
||||
resolveLoad = resolve;
|
||||
}),
|
||||
);
|
||||
return {
|
||||
loadTaskSessionsForTask,
|
||||
resolveLoad: (sessions: TaskSession[]) => resolveLoad(sessions),
|
||||
};
|
||||
}
|
||||
|
||||
describe("prepareAndSwitchTask — outgoing-env panel cleanup", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
@@ -185,53 +304,10 @@ describe("buildSwitchToSession", () => {
|
||||
* re-entry, as long as the session still has a known environment.
|
||||
*/
|
||||
describe("selectTaskWithLayout — last-selected session preference", () => {
|
||||
const TASK_ID = "task-A";
|
||||
const PRIMARY = "sess-primary";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
function makeKanbanStore(args: {
|
||||
activeSessionId: string | null;
|
||||
envIds: Record<string, string>;
|
||||
lastSessionByTaskId?: Record<string, string>;
|
||||
sessionTaskIds?: Record<string, string>;
|
||||
}): StoreApi<AppState> {
|
||||
const items: Record<string, { id: string; task_id: string }> = {};
|
||||
for (const [sid, tid] of Object.entries(args.sessionTaskIds ?? {})) {
|
||||
items[sid] = { id: sid, task_id: tid };
|
||||
}
|
||||
const state = {
|
||||
tasks: {
|
||||
activeSessionId: args.activeSessionId,
|
||||
lastSessionByTaskId: args.lastSessionByTaskId ?? {},
|
||||
},
|
||||
taskPRs: { byTaskId: {} as Record<string, unknown[]> },
|
||||
environmentIdBySessionId: args.envIds,
|
||||
taskSessions: { items },
|
||||
};
|
||||
return {
|
||||
getState: () => state as unknown as AppState,
|
||||
setState: vi.fn(),
|
||||
subscribe: vi.fn(),
|
||||
} as unknown as StoreApi<AppState>;
|
||||
}
|
||||
|
||||
function runSelect(store: StoreApi<AppState>) {
|
||||
const switchToSession = vi.fn();
|
||||
selectTaskWithLayout({
|
||||
taskId: TASK_ID,
|
||||
task: { primarySessionId: PRIMARY },
|
||||
store,
|
||||
switchToSession,
|
||||
loadTaskSessionsForTask: vi.fn(async () => []),
|
||||
setActiveTask: vi.fn(),
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
return switchToSession;
|
||||
}
|
||||
|
||||
it("prefers the user's last-selected session over primarySessionId on re-entry", () => {
|
||||
const LAST = "sess-gpt";
|
||||
const switchToSession = runSelect(
|
||||
@@ -292,3 +368,252 @@ describe("selectTaskWithLayout — last-selected session preference", () => {
|
||||
expect(switchToSession).toHaveBeenCalledWith(TASK_ID, PRIMARY, null);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectTaskWithLayout — pending selection races", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("does not switch back to a sessionless task when its prepare failure resolves after another task was selected", async () => {
|
||||
const SESSIONLESS = "task-sessionless";
|
||||
const OTHER = "task-other";
|
||||
const { state, store, setActiveTask } = makeSelectionHarness({
|
||||
activeTaskId: SESSIONLESS,
|
||||
activeSessionId: null,
|
||||
envIds: { "sess-other": "env-other" },
|
||||
sessions: { "sess-other": { id: "sess-other", task_id: OTHER } },
|
||||
});
|
||||
const { loadTaskSessionsForTask, resolveLoad } = makeDeferredSessionLoader();
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: SESSIONLESS,
|
||||
task: undefined,
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask,
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
state.tasks.activeTaskId = OTHER;
|
||||
state.tasks.activeSessionId = "sess-other";
|
||||
selectTaskWithLayout({
|
||||
taskId: OTHER,
|
||||
task: { primarySessionId: "sess-other" },
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask: vi.fn(async () => []),
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
resolveLoad([]);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(setActiveTask).not.toHaveBeenCalledWith(SESSIONLESS);
|
||||
expect(releaseLayoutToDefault).not.toHaveBeenCalled();
|
||||
expect(replaceTaskUrl).not.toHaveBeenCalledWith(SESSIONLESS);
|
||||
});
|
||||
|
||||
it("does not switch to an old pending selection after the user clicks back to the original task", async () => {
|
||||
const { store, setActiveTask } = makeSelectionHarness({
|
||||
activeTaskId: ORIGINAL_TASK_ID,
|
||||
activeSessionId: ORIGINAL_SESSION_ID,
|
||||
envIds: { [ORIGINAL_SESSION_ID]: ORIGINAL_ENV_ID },
|
||||
sessions: {
|
||||
[ORIGINAL_SESSION_ID]: { id: ORIGINAL_SESSION_ID, task_id: ORIGINAL_TASK_ID },
|
||||
},
|
||||
});
|
||||
const { loadTaskSessionsForTask, resolveLoad } = makeDeferredSessionLoader();
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: PENDING_TASK_ID,
|
||||
task: undefined,
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask,
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: ORIGINAL_TASK_ID,
|
||||
task: { primarySessionId: ORIGINAL_SESSION_ID },
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask: vi.fn(async () => []),
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
resolveLoad([]);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(setActiveTask).not.toHaveBeenCalledWith(PENDING_TASK_ID);
|
||||
expect(replaceTaskUrl).not.toHaveBeenCalledWith(PENDING_TASK_ID);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectTaskWithLayout — external active-task changes", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("does not apply a pending selection after an external active-task change", async () => {
|
||||
const OTHER = "task-other";
|
||||
const switchToSession = vi.fn();
|
||||
const { state, store, setActiveTask } = makeSelectionHarness({
|
||||
activeTaskId: ORIGINAL_TASK_ID,
|
||||
activeSessionId: ORIGINAL_SESSION_ID,
|
||||
envIds: { [ORIGINAL_SESSION_ID]: ORIGINAL_ENV_ID },
|
||||
sessions: {
|
||||
[PENDING_SESSION_ID]: { id: PENDING_SESSION_ID, task_id: PENDING_TASK_ID },
|
||||
[ORIGINAL_SESSION_ID]: { id: ORIGINAL_SESSION_ID, task_id: ORIGINAL_TASK_ID },
|
||||
},
|
||||
});
|
||||
const { loadTaskSessionsForTask, resolveLoad } = makeDeferredSessionLoader();
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: PENDING_TASK_ID,
|
||||
task: { primarySessionId: PENDING_SESSION_ID },
|
||||
store,
|
||||
switchToSession,
|
||||
loadTaskSessionsForTask,
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
state.tasks.activeTaskId = OTHER;
|
||||
state.tasks.activeSessionId = null;
|
||||
resolveLoad([
|
||||
{ id: PENDING_SESSION_ID, task_id: PENDING_TASK_ID, is_primary: true } as TaskSession,
|
||||
]);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(switchToSession).not.toHaveBeenCalled();
|
||||
expect(replaceTaskUrl).not.toHaveBeenCalledWith(PENDING_TASK_ID);
|
||||
});
|
||||
|
||||
it("does not apply a pending selection after an external task switch returns to the original task", async () => {
|
||||
const OTHER = "task-other";
|
||||
const switchToSession = vi.fn();
|
||||
const { store, setActiveTask } = makeSelectionHarness({
|
||||
activeTaskId: ORIGINAL_TASK_ID,
|
||||
activeSessionId: ORIGINAL_SESSION_ID,
|
||||
envIds: { [ORIGINAL_SESSION_ID]: ORIGINAL_ENV_ID },
|
||||
sessions: {
|
||||
[PENDING_SESSION_ID]: { id: PENDING_SESSION_ID, task_id: PENDING_TASK_ID },
|
||||
[ORIGINAL_SESSION_ID]: { id: ORIGINAL_SESSION_ID, task_id: ORIGINAL_TASK_ID },
|
||||
},
|
||||
});
|
||||
const { loadTaskSessionsForTask, resolveLoad } = makeDeferredSessionLoader();
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: PENDING_TASK_ID,
|
||||
task: { primarySessionId: PENDING_SESSION_ID },
|
||||
store,
|
||||
switchToSession,
|
||||
loadTaskSessionsForTask,
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
setActiveTask(OTHER);
|
||||
setActiveTask(ORIGINAL_TASK_ID);
|
||||
resolveLoad([
|
||||
{ id: PENDING_SESSION_ID, task_id: PENDING_TASK_ID, is_primary: true } as TaskSession,
|
||||
]);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(switchToSession).not.toHaveBeenCalled();
|
||||
expect(replaceTaskUrl).not.toHaveBeenCalledWith(PENDING_TASK_ID);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectTaskWithLayout — selection guard cleanup", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("disposes the selection guard when loading a primary task's sessions rejects", async () => {
|
||||
const { store, setActiveTask, getListenerCount } = makeSelectionHarness({
|
||||
activeTaskId: ORIGINAL_TASK_ID,
|
||||
activeSessionId: ORIGINAL_SESSION_ID,
|
||||
envIds: { [ORIGINAL_SESSION_ID]: ORIGINAL_ENV_ID },
|
||||
});
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: PENDING_TASK_ID,
|
||||
task: { primarySessionId: PENDING_SESSION_ID },
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask: vi.fn(async () => {
|
||||
throw new Error("load failed");
|
||||
}),
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
expect(getListenerCount()).toBe(1);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(getListenerCount()).toBe(0);
|
||||
});
|
||||
|
||||
it("disposes the selection guard when loading a sessionless task rejects", async () => {
|
||||
const { store, setActiveTask, getListenerCount } = makeSelectionHarness({
|
||||
activeTaskId: ORIGINAL_TASK_ID,
|
||||
activeSessionId: ORIGINAL_SESSION_ID,
|
||||
envIds: { [ORIGINAL_SESSION_ID]: ORIGINAL_ENV_ID },
|
||||
});
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: PENDING_TASK_ID,
|
||||
task: undefined,
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask: vi.fn(async () => {
|
||||
throw new Error("load failed");
|
||||
}),
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
expect(getListenerCount()).toBe(1);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(getListenerCount()).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectTaskWithLayout — pending old-session changes", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("keeps a pending selection alive when only the old task's session id changes", async () => {
|
||||
const OLD = "task-old";
|
||||
const SESSIONLESS = "task-sessionless";
|
||||
vi.mocked(launchSession).mockResolvedValue({} as never);
|
||||
const { state, store, setActiveTask } = makeSelectionHarness({
|
||||
activeTaskId: OLD,
|
||||
activeSessionId: "sess-old",
|
||||
});
|
||||
const { loadTaskSessionsForTask, resolveLoad } = makeDeferredSessionLoader();
|
||||
|
||||
selectTaskWithLayout({
|
||||
taskId: SESSIONLESS,
|
||||
task: undefined,
|
||||
store,
|
||||
switchToSession: vi.fn(),
|
||||
loadTaskSessionsForTask,
|
||||
setActiveTask,
|
||||
setPreparingTaskId: vi.fn(),
|
||||
});
|
||||
|
||||
state.tasks.activeSessionId = "sess-old-replaced";
|
||||
resolveLoad([]);
|
||||
await flushTaskSelection();
|
||||
|
||||
expect(setActiveTask).toHaveBeenCalledWith(SESSIONLESS);
|
||||
expect(replaceTaskUrl).toHaveBeenCalledWith(SESSIONLESS);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -19,6 +19,7 @@ import { buildPrepareRequest } from "@/lib/services/session-launch-helpers";
|
||||
import { createDebugLogger, isDebug } from "@/lib/debug/log";
|
||||
|
||||
const debug = createDebugLogger("dockview:task-select");
|
||||
let taskSelectionSequence = 0;
|
||||
|
||||
export type SwitchToSessionFn = (
|
||||
taskId: string,
|
||||
@@ -109,11 +110,21 @@ export function buildSwitchToSession(
|
||||
};
|
||||
}
|
||||
|
||||
function nextTaskSelectionToken(): number {
|
||||
taskSelectionSequence += 1;
|
||||
return taskSelectionSequence;
|
||||
}
|
||||
|
||||
function taskSelectionWasSuperseded(selectionToken: number): boolean {
|
||||
return selectionToken !== taskSelectionSequence;
|
||||
}
|
||||
|
||||
export async function prepareAndSwitchTask(
|
||||
taskId: string,
|
||||
store: StoreApi<AppState>,
|
||||
switchToSession: SwitchToSessionFn,
|
||||
setPreparingTaskId: (id: string | null) => void,
|
||||
shouldContinue: () => boolean = () => true,
|
||||
): Promise<boolean> {
|
||||
setPreparingTaskId(taskId);
|
||||
// Capture before the async launch; WS events may update activeSessionId
|
||||
@@ -131,6 +142,7 @@ export async function prepareAndSwitchTask(
|
||||
try {
|
||||
const { request } = buildPrepareRequest(taskId);
|
||||
const resp = await launchSession(request);
|
||||
if (!shouldContinue()) return false;
|
||||
if (resp.session_id) {
|
||||
// Pass `null` instead of the original oldSessionId — releaseLayoutToDefault
|
||||
// already saved + released the outgoing env, and the dockview now holds the
|
||||
@@ -163,7 +175,26 @@ export function selectTaskWithLayout(params: {
|
||||
}): void {
|
||||
const { taskId, task, store, switchToSession, loadTaskSessionsForTask } = params;
|
||||
const state = store.getState();
|
||||
const selectionToken = nextTaskSelectionToken();
|
||||
const startActiveTaskId = state.tasks.activeTaskId ?? null;
|
||||
const oldSessionId = state.tasks.activeSessionId;
|
||||
let activeTaskChangedExternally = false;
|
||||
const unsubscribeSelectionGuard = store.subscribe((current, previous) => {
|
||||
const currentTaskId = current.tasks.activeTaskId ?? null;
|
||||
const previousTaskId = previous.tasks.activeTaskId ?? null;
|
||||
if (currentTaskId !== previousTaskId && currentTaskId !== taskId) {
|
||||
activeTaskChangedExternally = true;
|
||||
}
|
||||
});
|
||||
const disposeSelectionGuard = () => {
|
||||
if (typeof unsubscribeSelectionGuard === "function") unsubscribeSelectionGuard();
|
||||
};
|
||||
const selectionWasSuperseded = () => {
|
||||
if (taskSelectionWasSuperseded(selectionToken)) return true;
|
||||
if (activeTaskChangedExternally) return true;
|
||||
const activeTaskId = store.getState().tasks.activeTaskId ?? null;
|
||||
return activeTaskId !== startActiveTaskId && activeTaskId !== taskId;
|
||||
};
|
||||
if (isDebug()) {
|
||||
debug("selectTaskWithLayout: entry", {
|
||||
taskId,
|
||||
@@ -182,45 +213,56 @@ export function selectTaskWithLayout(params: {
|
||||
});
|
||||
const hasEnvId = !!state.environmentIdBySessionId[targetSessionId];
|
||||
if (hasEnvId) {
|
||||
disposeSelectionGuard();
|
||||
switchToSession(taskId, targetSessionId, oldSessionId);
|
||||
loadTaskSessionsForTask(taskId);
|
||||
replaceTaskUrl(taskId);
|
||||
return;
|
||||
}
|
||||
loadTaskSessionsForTask(taskId).then((sessions) => {
|
||||
switchToSession(taskId, resolveLoadedSessionId(sessions, targetSessionId), oldSessionId);
|
||||
replaceTaskUrl(taskId);
|
||||
});
|
||||
void loadTaskSessionsForTask(taskId)
|
||||
.then((sessions) => {
|
||||
if (selectionWasSuperseded()) return;
|
||||
switchToSession(taskId, resolveLoadedSessionId(sessions, targetSessionId), oldSessionId);
|
||||
replaceTaskUrl(taskId);
|
||||
})
|
||||
.finally(disposeSelectionGuard)
|
||||
.catch(() => undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
loadTaskSessionsForTask(taskId).then(async (sessions) => {
|
||||
const currentOldSessionId = store.getState().tasks.activeSessionId;
|
||||
const primary = sessions.find((s) => s.is_primary);
|
||||
const sessionId = primary?.id ?? sessions[0]?.id ?? null;
|
||||
if (sessionId) {
|
||||
switchToSession(taskId, sessionId, currentOldSessionId);
|
||||
replaceTaskUrl(taskId);
|
||||
return;
|
||||
}
|
||||
void loadTaskSessionsForTask(taskId)
|
||||
.then(async (sessions) => {
|
||||
if (selectionWasSuperseded()) return;
|
||||
const currentOldSessionId = store.getState().tasks.activeSessionId;
|
||||
const primary = sessions.find((s) => s.is_primary);
|
||||
const sessionId = primary?.id ?? sessions[0]?.id ?? null;
|
||||
if (sessionId) {
|
||||
switchToSession(taskId, sessionId, currentOldSessionId);
|
||||
replaceTaskUrl(taskId);
|
||||
return;
|
||||
}
|
||||
|
||||
const switched = await prepareAndSwitchTask(
|
||||
taskId,
|
||||
store,
|
||||
switchToSession,
|
||||
params.setPreparingTaskId,
|
||||
);
|
||||
if (switched) {
|
||||
replaceTaskUrl(taskId);
|
||||
return;
|
||||
}
|
||||
const switched = await prepareAndSwitchTask(
|
||||
taskId,
|
||||
store,
|
||||
switchToSession,
|
||||
params.setPreparingTaskId,
|
||||
() => !selectionWasSuperseded(),
|
||||
);
|
||||
if (switched) {
|
||||
replaceTaskUrl(taskId);
|
||||
return;
|
||||
}
|
||||
if (selectionWasSuperseded()) return;
|
||||
|
||||
// Failure path: prepareAndSwitchTask already called releaseLayoutToDefault
|
||||
// before awaiting, so the outgoing env's layout is already saved and the
|
||||
// dockview is showing the default layout. A second release here would
|
||||
// overwrite the just-saved env layout with `api.toJSON()` (the default),
|
||||
// losing the user's real layout for the originating task.
|
||||
params.setActiveTask(taskId);
|
||||
replaceTaskUrl(taskId);
|
||||
});
|
||||
// Failure path: prepareAndSwitchTask already called releaseLayoutToDefault
|
||||
// before awaiting, so the outgoing env's layout is already saved and the
|
||||
// dockview is showing the default layout. A second release here would
|
||||
// overwrite the just-saved env layout with `api.toJSON()` (the default),
|
||||
// losing the user's real layout for the originating task.
|
||||
params.setActiveTask(taskId);
|
||||
replaceTaskUrl(taskId);
|
||||
})
|
||||
.finally(disposeSelectionGuard)
|
||||
.catch(() => undefined);
|
||||
}
|
||||
|
||||
@@ -29,20 +29,22 @@ function flushMicrotasks() {
|
||||
return act(() => Promise.resolve());
|
||||
}
|
||||
|
||||
describe("useEnsureTaskSession", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockLoadSessions.mockResolvedValue(undefined);
|
||||
mockSessionsResult = { sessions: [], isLoaded: true, loadSessions: mockLoadSessions };
|
||||
mockEnsureTaskSession.mockResolvedValue({
|
||||
success: true,
|
||||
task_id: "task-1",
|
||||
session_id: "sess-new",
|
||||
state: "CREATED",
|
||||
source: "created_prepare",
|
||||
newly_created: true,
|
||||
});
|
||||
function resetEnsureTaskSessionMocks() {
|
||||
vi.clearAllMocks();
|
||||
mockLoadSessions.mockResolvedValue(undefined);
|
||||
mockSessionsResult = { sessions: [], isLoaded: true, loadSessions: mockLoadSessions };
|
||||
mockEnsureTaskSession.mockResolvedValue({
|
||||
success: true,
|
||||
task_id: "task-1",
|
||||
session_id: "sess-new",
|
||||
state: "CREATED",
|
||||
source: "created_prepare",
|
||||
newly_created: true,
|
||||
});
|
||||
}
|
||||
|
||||
describe("useEnsureTaskSession", () => {
|
||||
beforeEach(resetEnsureTaskSessionMocks);
|
||||
|
||||
it("calls the backend ensure endpoint once when the task has no sessions", async () => {
|
||||
const { result } = renderHook(() => useEnsureTaskSession(TASK));
|
||||
@@ -117,6 +119,33 @@ describe("useEnsureTaskSession", () => {
|
||||
await flushMicrotasks();
|
||||
expect(result.current.status).toBe("idle");
|
||||
});
|
||||
});
|
||||
|
||||
describe("useEnsureTaskSession — task changes", () => {
|
||||
beforeEach(resetEnsureTaskSessionMocks);
|
||||
|
||||
it("clears a stale error when switching to a task that already has a session", async () => {
|
||||
mockEnsureTaskSession.mockRejectedValueOnce(new Error("task one failed"));
|
||||
const { result, rerender } = renderHook(
|
||||
({ task }: { task: { id: string } }) => useEnsureTaskSession(task),
|
||||
{ initialProps: { task: TASK } },
|
||||
);
|
||||
|
||||
await flushMicrotasks();
|
||||
expect(result.current.status).toBe("error");
|
||||
expect(result.current.error?.message).toBe("task one failed");
|
||||
|
||||
mockSessionsResult = {
|
||||
sessions: [{ id: "sess-2" }],
|
||||
isLoaded: true,
|
||||
loadSessions: mockLoadSessions,
|
||||
};
|
||||
rerender({ task: { id: "task-2" } });
|
||||
|
||||
expect(result.current.status).toBe("idle");
|
||||
expect(result.current.error).toBeNull();
|
||||
expect(mockEnsureTaskSession).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("calls ensure again when the task id changes", () => {
|
||||
const { rerender } = renderHook(
|
||||
|
||||
@@ -43,6 +43,17 @@ export function useEnsureTaskSession(
|
||||
// Latch keyed by `${taskId}:${retryToken}` so a re-mount on the same task
|
||||
// doesn't refire, but switching tasks or calling retry() does.
|
||||
const launchedKeyRef = useRef<string | null>(null);
|
||||
const previousTaskIdRef = useRef<string | null>(taskId);
|
||||
|
||||
/* eslint-disable react-hooks/set-state-in-effect -- task changes must clear stale ensure-session errors before early returns */
|
||||
useEffect(() => {
|
||||
if (previousTaskIdRef.current === taskId) return;
|
||||
previousTaskIdRef.current = taskId;
|
||||
launchedKeyRef.current = null;
|
||||
setStatus("idle");
|
||||
setError(null);
|
||||
}, [taskId]);
|
||||
/* eslint-enable react-hooks/set-state-in-effect */
|
||||
|
||||
/* eslint-disable react-hooks/set-state-in-effect -- ensuring a session is a side effect; status mirrors that external work */
|
||||
useEffect(() => {
|
||||
|
||||
Reference in New Issue
Block a user