Fix exit code for run command

internal/command/command_exec_normalizer.go

@@ -50,33 +50,18 @@ func (n *argsNormalizer) Normalize(cfg *Config) (*Config, error) {
 		var buf strings.Builder
 
 		if isShellLanguage(filepath.Base(cfg.ProgramName)) {
-			if options := shellOptionsFromProgram(cfg.ProgramName); options != "" {
-				_, _ = buf.WriteString(options)
-				_, _ = buf.WriteString("\n\n")
+			if err := n.inlineShell(cfg, &buf); err != nil {
+				return nil, err
 			}
-
-			if n.session != nil {
-				if err := n.createTempDir(); err != nil {
-					return nil, err
+		} else {
+			// Write the script from the commands or the script.
+			if commands := cfg.GetCommands(); commands != nil {
+				for _, cmd := range commands.Items {
+					_, _ = buf.WriteString(cmd)
+					_, _ = buf.WriteRune('\n')
 				}
-				_, _ = buf.WriteString(fmt.Sprintf("%s > %s\n", EnvDumpCommand, filepath.Join(n.tempDir, envStartFileName)))
-			}
-		}
-
-		if commands := cfg.GetCommands(); commands != nil {
-			for _, cmd := range commands.Items {
-				_, _ = buf.WriteString(cmd)
-				_, _ = buf.WriteRune('\n')
-			}
-		} else if script := cfg.GetScript(); script != "" {
-			_, _ = buf.WriteString(script)
-		}
-
-		if isShellLanguage(filepath.Base(cfg.ProgramName)) {
-			if n.session != nil {
-				_, _ = buf.WriteString(fmt.Sprintf("trap \"%s > %s\" EXIT\n", EnvDumpCommand, filepath.Join(n.tempDir, envEndFileName)))
-
-				n.isEnvCollectable = true
+			} else if script := cfg.GetScript(); script != "" {
+				_, _ = buf.WriteString(script)
 			}
 		}
 
@@ -107,6 +92,39 @@ func (n *argsNormalizer) Normalize(cfg *Config) (*Config, error) {
 	return result, nil
 }
 
+func (n *argsNormalizer) inlineShell(cfg *Config, buf *strings.Builder) error {
+	if options := shellOptionsFromProgram(cfg.ProgramName); options != "" {
+		_, _ = buf.WriteString(options)
+		_, _ = buf.WriteString("\n\n")
+	}
+
+	// If the session is provided, we need to collect the environment before and after the script execution.
+	// Here, we dump env before the script execution and use trap on EXIT to collect the env after the script execution.
+	if n.session != nil {
+		if err := n.createTempDir(); err != nil {
+			return err
+		}
+
+		_, _ = buf.WriteString(fmt.Sprintf("%s > %s\n", EnvDumpCommand, filepath.Join(n.tempDir, envStartFileName)))
+		_, _ = buf.WriteString(fmt.Sprintf("__cleanup() {\nrv=$?\n%s > %s\nexit $rv\n}\n", EnvDumpCommand, filepath.Join(n.tempDir, envEndFileName)))
+		_, _ = buf.WriteString("trap -- \"__cleanup\" EXIT\n")
+
+		n.isEnvCollectable = true
+	}
+
+	// Write the script from the commands or the script.
+	if commands := cfg.GetCommands(); commands != nil {
+		for _, cmd := range commands.Items {
+			_, _ = buf.WriteString(cmd)
+			_, _ = buf.WriteRune('\n')
+		}
+	} else if script := cfg.GetScript(); script != "" {
+		_, _ = buf.WriteString(script)
+	}
+
+	return nil
+}
+
 func (n *argsNormalizer) Cleanup() error {
 	if n.tempDir == "" {
 		return nil

internal/command/command_test.go

@@ -128,7 +128,7 @@ func TestExecutionCommandFromCodeBlocks(t *testing.T) {
 		})
 
 		t.Run("VirtualCommand", func(t *testing.T) {
-			tc := tc
+			t.Parallel()
 
 			t.Run(tc.name, func(t *testing.T) {
 				t.Parallel()

internal/runner/command.go

@@ -161,6 +161,8 @@ func newCommand(cfg *commandConfig) (*command, error) {
 			tmpEnvDir = envStorePath
 
 			_, _ = builder.WriteString(fmt.Sprintf("%s > %s\n", dumpCmd, filepath.Join(envStorePath, envStartFileName)))
+			_, _ = builder.WriteString(fmt.Sprintf("__cleanup() {\nrv=$?\n%s > %s\nexit $rv\n}\n", dumpCmd, filepath.Join(envStorePath, envEndFileName)))
+			_, _ = builder.WriteString("trap -- \"__cleanup\" EXIT\n")
 
 			if len(cfg.Commands) > 0 {
 				_, _ = builder.WriteString(
@@ -172,8 +174,6 @@ func newCommand(cfg *commandConfig) (*command, error) {
 				)
 			}
 
-			_, _ = builder.WriteString(fmt.Sprintf("%s > %s\n", dumpCmd, filepath.Join(envStorePath, envEndFileName)))
-
 			extraArgs = []string{}
 
 			if cfg.Tty {

internal/runnerv2service/service_execute.go

@@ -96,7 +96,7 @@ func (r *runnerService) Execute(srv runnerv2alpha1.RunnerService_ExecuteServer)
 			case err == io.EOF:
 				logger.Info("client closed its send direction; stopping the program")
 				if err := exec.Cmd.StopWithSignal(os.Interrupt); err != nil {
-					logger.Info("failed to stop the command with signal", zap.Error(err))
+					logger.Info("failed to stop the command with interrupt signal", zap.Error(err))
 				}
 				return
 			case status.Convert(err).Code() == codes.Canceled || status.Convert(err).Code() == codes.DeadlineExceeded:
@@ -105,7 +105,7 @@ func (r *runnerService) Execute(srv runnerv2alpha1.RunnerService_ExecuteServer)
 				} else {
 					logger.Info("stream canceled while the process is still running; program will be stopped if non-background")
 					if err := exec.Cmd.StopWithSignal(os.Kill); err != nil {
-						logger.Info("failed to stop program", zap.Error(err))
+						logger.Info("failed to stop program with kill signal", zap.Error(err))
 					}
 				}
 				return

internal/runnerv2service/service_execute_test.go

@@ -322,14 +322,9 @@ func TestRunnerServiceServerExecuteWithInput(t *testing.T) {
 
 		err = stream.Send(&runnerv2alpha1.ExecuteRequest{
 			Config: &runnerv2alpha1.ProgramConfig{
-				ProgramName: "bash",
-				Source: &runnerv2alpha1.ProgramConfig_Commands{
-					Commands: &runnerv2alpha1.ProgramConfig_CommandList{
-						Items: []string{
-							"sleep 30",
-						},
-					},
-				},
+				ProgramName: "sleep",
+				Arguments:   []string{"30"},
+				Mode:        runnerv2alpha1.CommandMode_COMMAND_MODE_INLINE,
 			},
 		})
 		require.NoError(t, err)

testdata/script/exitcode.txtar

@@ -0,0 +1,24 @@
+! exec runme run test-singleline
+stdout 'failhereplease: (command )?not found'
+
+! exec runme run test-multiline
+stdout 'failhereplease: (command )?not found'
+stderr 'could not execute command: failed to run command "test-multiline": exit code: 127'
+
+! exec runme run-locally test-singleline
+! stdout .
+stderr 'failhereplease: (command )?not found'
+
+! exec runme run-locally test-multiline
+! stdout .
+stderr 'failhereplease: (command )?not found'
+
+-- README.md --
+```sh { "name": "test-singleline" }
+failhereplease && echo single
+```
+
+```sh { "name": "test-multiline" }
+failhereplease
+echo multi
+```