Kill session only on UserScript stage

When the build container is done the session server needs to be
notified, so waitForTerminal is not triggered then the job is finished,
since the user can no longer user the terminal, because the container is
dead.

executors/docker/executor_docker.go

@@ -983,11 +983,6 @@ func (e *executor) waitForContainer(id string) error {
 			continue
 		}
 
-		// If container has finished, kill session.
-		if e.Build.Session != nil {
-			e.Build.Session.Kill()
-		}
-
 		if container.State.ExitCode != 0 {
 			return &common.BuildError{
 				Inner: fmt.Errorf("exit code %d", container.State.ExitCode),

executors/docker/executor_docker_command.go

@@ -86,6 +86,12 @@ func (s *commandExecutor) Run(cmd common.ExecutorCommand) error {
 
 	s.SetCurrentStage(DockerExecutorStageRun)
 
+	defer func() {
+		if cmd.Stage == common.BuildStageUserScript && s.Build.Session != nil {
+			s.Build.Session.Kill()
+		}
+	}()
+
 	return s.watchContainer(cmd.Context, runOn.ID, bytes.NewBufferString(cmd.Script))
 }