feat: allow specifying shell in /prepare command

README.md

@@ -190,6 +190,12 @@ To activate Prepare Mode, simply use:
 TmuxAI » /prepare
 ```
 
+By default, TmuxAI will attempt to detect the shell running in the execution pane. If you need to specify the shell manually, you can provide it as an argument:
+
+```
+TmuxAI » /prepare bash
+```
+
 **Prepared Fish Example:**
 
 ```shell
@@ -291,7 +297,7 @@ TmuxAI » /squash
 | `/config`                   | View current configuration settings                              |
 | `/config set <key> <value>` | Override configuration for current session                       |
 | `/squash`                   | Manually trigger context summarization                           |
-| `/prepare`                  | Initialize Prepared Mode for the Exec Pane                       |
+| `/prepare [shell]`          | Initialize Prepared Mode for the Exec Pane (e.g., bash, zsh)     |
 | `/watch <description>`      | Enable Watch Mode with specified goal                            |
 | `/exit`                     | Exit TmuxAI                                                      |
 

internal/chat.go

@@ -177,6 +177,13 @@ func (c *CLIInterface) newCompleter() *completion.CmdCompletionOrList2 {
 					return AllowedConfigKeys, AllowedConfigKeys
 				}
 			}
+
+			// Handle /prepare subcommands
+			if len(field) > 0 && field[0] == "/prepare" {
+				if len(field) == 1 || (len(field) == 2 && !strings.HasSuffix(field[1], " ")) {
+					return []string{"bash", "zsh", "fish"}, []string{"bash", "zsh", "fish"}
+				}
+			}
 			return nil, nil
 		},
 	}

internal/chat_command.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/alvinunreal/tmuxai/config"
 	"github.com/alvinunreal/tmuxai/logger"
@@ -64,12 +65,56 @@ func (m *Manager) ProcessSubCommand(command string) {
 		return
 
 	case prefixMatch(commandPrefix, "/prepare"):
+		supportedShells := []string{"bash", "zsh", "fish"}
 		m.InitExecPane()
-		m.PrepareExecPane()
-		m.Messages = []ChatMessage{}
-		if m.ExecPane.IsPrepared {
-			m.Println("Exec pane prepared successfully")
+
+		// Check if exec pane is a subshell
+		if m.ExecPane.IsSubShell {
+			if len(parts) > 1 {
+				shell := parts[1]
+				isSupported := false
+				for _, supportedShell := range supportedShells {
+					if shell == supportedShell {
+						isSupported = true
+						break
+					}
+				}
+				if !isSupported {
+					m.Println(fmt.Sprintf("Shell '%s' is not supported. Supported shells are: %s", shell, strings.Join(supportedShells, ", ")))
+					return
+				}
+				m.PrepareExecPaneWithShell(shell)
+			} else {
+				m.Println("Shell detection is not supported on subshells.")
+				m.Println("Please specify the shell manually: /prepare bash, /prepare zsh, or /prepare fish")
+				return
+			}
+		} else {
+			if len(parts) > 1 {
+				shell := parts[1]
+				isSupported := false
+				for _, supportedShell := range supportedShells {
+					if shell == supportedShell {
+						isSupported = true
+						break
+					}
+				}
+
+				if !isSupported {
+					m.Println(fmt.Sprintf("Shell '%s' is not supported. Supported shells are: %s", shell, strings.Join(supportedShells, ", ")))
+					return
+				}
+				m.PrepareExecPaneWithShell(shell)
+			} else {
+				m.PrepareExecPane()
+			}
 		}
+
+		// for latency over ssh connections
+		time.Sleep(500 * time.Millisecond)
+		m.ExecPane.Refresh(m.GetMaxCaptureLines())
+		m.Messages = []ChatMessage{}
+
 		fmt.Println(m.ExecPane.String())
 		m.parseExecPaneCommandHistory()
 

internal/chat_command_test.go

@@ -0,0 +1,169 @@
+package internal
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/alvinunreal/tmuxai/config"
+	"github.com/alvinunreal/tmuxai/system"
+	"github.com/stretchr/testify/assert"
+)
+
+// Test /prepare command behavior with subshell
+func TestProcessSubCommand_PrepareSubshell(t *testing.T) {
+	manager := &Manager{
+		Config:           &config.Config{MaxCaptureLines: 1000},
+		SessionOverrides: make(map[string]any),
+		Messages:         []ChatMessage{},
+		ExecPane: &system.TmuxPaneDetails{
+			Id:         "test-pane",
+			IsSubShell: true, // This is a subshell
+		},
+	}
+
+	// Mock system functions to prevent actual tmux calls
+	originalTmuxSend := system.TmuxSendCommandToPane
+	originalTmuxCapture := system.TmuxCapturePane
+	originalTmuxCurrentPaneId := system.TmuxCurrentPaneId
+	originalTmuxPanesDetails := system.TmuxPanesDetails
+	defer func() {
+		system.TmuxSendCommandToPane = originalTmuxSend
+		system.TmuxCapturePane = originalTmuxCapture
+		system.TmuxCurrentPaneId = originalTmuxCurrentPaneId
+		system.TmuxPanesDetails = originalTmuxPanesDetails
+	}()
+
+	commandsSent := []string{}
+	system.TmuxSendCommandToPane = func(paneId string, command string, enter bool) error {
+		commandsSent = append(commandsSent, command)
+		return nil
+	}
+
+	system.TmuxCapturePane = func(paneId string, maxLines int) (string, error) {
+		return "", nil
+	}
+
+	// Mock the system functions used by GetTmuxPanes to return the test pane
+	system.TmuxCurrentPaneId = func() (string, error) {
+		return "main-pane", nil
+	}
+	system.TmuxPanesDetails = func(windowTarget string) ([]system.TmuxPaneDetails, error) {
+		// Return the test pane as the only available pane
+		return []system.TmuxPaneDetails{*manager.ExecPane}, nil
+	}
+
+	// Test case 1: /prepare with valid shell on subshell (should work and send commands)
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare bash")
+
+	assert.Len(t, commandsSent, 2, "Should send PS1 command and clear command for bash")
+	assert.Contains(t, commandsSent[0], "PS1=", "Should send bash PS1 command")
+	assert.Equal(t, "C-l", commandsSent[1], "Should send clear screen command")
+
+	// Test case 2: /prepare with zsh on subshell
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare zsh")
+
+	assert.Len(t, commandsSent, 2, "Should send PROMPT command and clear command for zsh")
+	assert.Contains(t, commandsSent[0], "PROMPT=", "Should send zsh PROMPT command")
+	assert.Equal(t, "C-l", commandsSent[1], "Should send clear screen command")
+
+	// Test case 3: /prepare with fish on subshell
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare fish")
+
+	assert.Len(t, commandsSent, 2, "Should send fish_prompt function and clear command for fish")
+	assert.Contains(t, commandsSent[0], "fish_prompt", "Should send fish prompt function")
+	assert.Equal(t, "C-l", commandsSent[1], "Should send clear screen command")
+
+	// Test case 4: /prepare without shell specification on subshell (should not send commands, just print warning)
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare")
+
+	fmt.Println(commandsSent)
+	assert.Len(t, commandsSent, 0, "Should not send commands when no shell specified on subshell (should show warning instead)")
+}
+
+// Test /prepare command behavior with normal shell (not subshell)
+func TestProcessSubCommand_PrepareNormalShell(t *testing.T) {
+	manager := &Manager{
+		Config:           &config.Config{MaxCaptureLines: 1000},
+		SessionOverrides: make(map[string]any),
+		Messages:         []ChatMessage{},
+		ExecPane: &system.TmuxPaneDetails{
+			Id:             "test-pane",
+			IsSubShell:     false,    // This is NOT a subshell
+			CurrentCommand: "unknown", // Unsupported shell - should not send commands
+		},
+	}
+
+	// Mock system functions to prevent actual tmux calls
+	originalTmuxSend := system.TmuxSendCommandToPane
+	originalTmuxCapture := system.TmuxCapturePane
+	originalTmuxCurrentPaneId := system.TmuxCurrentPaneId
+	originalTmuxPanesDetails := system.TmuxPanesDetails
+	defer func() {
+		system.TmuxSendCommandToPane = originalTmuxSend
+		system.TmuxCapturePane = originalTmuxCapture
+		system.TmuxCurrentPaneId = originalTmuxCurrentPaneId
+		system.TmuxPanesDetails = originalTmuxPanesDetails
+	}()
+
+	commandsSent := []string{}
+	system.TmuxSendCommandToPane = func(paneId string, command string, enter bool) error {
+		commandsSent = append(commandsSent, command)
+		return nil
+	}
+
+	system.TmuxCapturePane = func(paneId string, maxLines int) (string, error) {
+		return "", nil
+	}
+
+	// Mock the system functions used by GetTmuxPanes to return the test pane
+	system.TmuxCurrentPaneId = func() (string, error) {
+		return "main-pane", nil
+	}
+	system.TmuxPanesDetails = func(windowTarget string) ([]system.TmuxPaneDetails, error) {
+		// Return the test pane as the only available pane
+		return []system.TmuxPaneDetails{*manager.ExecPane}, nil
+	}
+
+	// Test case 1: /prepare without shell specification when CurrentCommand is not a shell (should not send commands)
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare")
+
+	assert.Len(t, commandsSent, 0, "Should not send commands when CurrentCommand is not a supported shell")
+
+	// Test case 2: /prepare with explicit shell on normal shell (should work)
+	commandsSent = []string{} // Reset
+	manager.ProcessSubCommand("/prepare zsh")
+
+	assert.Len(t, commandsSent, 2, "Should send commands when explicitly specifying shell on normal pane")
+	assert.Contains(t, commandsSent[0], "PROMPT=", "Should send zsh PROMPT command")
+	assert.Equal(t, "C-l", commandsSent[1], "Should send clear screen command")
+}
+
+// Test IsMessageSubcommand function
+func TestIsMessageSubcommand(t *testing.T) {
+	manager := &Manager{}
+
+	// Test cases for command detection
+	testCases := []struct {
+		input    string
+		expected bool
+		desc     string
+	}{
+		{"/help", true, "Simple command should be detected"},
+		{"/prepare bash", true, "Command with arguments should be detected"},
+		{"  /info  ", true, "Command with whitespace should be detected"},
+		{"/PREPARE", true, "Uppercase command should be detected"},
+		{"hello world", false, "Regular message should not be detected as command"},
+		{"", false, "Empty string should not be detected as command"},
+		{"/ invalid", true, "Any string starting with / should be detected as command"},
+	}
+
+	for _, tc := range testCases {
+		result := manager.IsMessageSubcommand(tc.input)
+		assert.Equal(t, tc.expected, result, tc.desc)
+	}
+}

internal/exec_pane.go

@@ -34,23 +34,22 @@ func (m *Manager) InitExecPane() {
 	m.ExecPane = &availablePane
 }
 
-func (m *Manager) PrepareExecPane() {
+func (m *Manager) PrepareExecPaneWithShell(shell string) {
 	m.ExecPane.Refresh(m.GetMaxCaptureLines())
 	if m.ExecPane.IsPrepared && m.ExecPane.Shell != "" {
 		return
 	}
 
-	shellCommand := m.ExecPane.CurrentCommand
 	var ps1Command string
-	switch shellCommand {
+	switch shell {
 	case "zsh":
 		ps1Command = `export PROMPT='%n@%m:%~[%T][%?]» '`
 	case "bash":
 		ps1Command = `export PS1='\u@\h:\w[\A][$?]» '`
 	case "fish":
 		ps1Command = `function fish_prompt; set -l s $status; printf '%s@%s:%s[%s][%d]» ' $USER (hostname -s) (prompt_pwd) (date +"%H:%M") $s; end`
 	default:
-		errMsg := fmt.Sprintf("Shell '%s' in pane %s is recognized but not yet supported for PS1 modification.", shellCommand, m.ExecPane.Id)
+		errMsg := fmt.Sprintf("Shell '%s' in pane %s is recognized but not yet supported for PS1 modification.", shell, m.ExecPane.Id)
 		logger.Info(errMsg)
 		return
 	}
@@ -59,11 +58,17 @@ func (m *Manager) PrepareExecPane() {
 	_ = system.TmuxSendCommandToPane(m.ExecPane.Id, "C-l", false)
 }
 
+func (m *Manager) PrepareExecPane() {
+	m.PrepareExecPaneWithShell(m.ExecPane.CurrentCommand)
+}
+
 func (m *Manager) ExecWaitCapture(command string) (CommandExecHistory, error) {
 	_ = system.TmuxSendCommandToPane(m.ExecPane.Id, command, true)
-	m.ExecPane.Refresh(m.GetMaxCaptureLines())
 
-	m.Println("")
+	// wait for keys to be sent, duo to sometimes ssh latency
+	time.Sleep(500 * time.Millisecond)
+
+	m.ExecPane.Refresh(m.GetMaxCaptureLines())
 
 	animChars := []string{"⋯", "⋱", "⋮", "⋰"}
 	animIndex := 0
@@ -76,13 +81,25 @@ func (m *Manager) ExecWaitCapture(command string) (CommandExecHistory, error) {
 	fmt.Print("\r\033[K")
 
 	m.parseExecPaneCommandHistory()
+	if len(m.ExecHistory) == 0 {
+		logger.Error("Failed to parse command history from exec pane")
+		return CommandExecHistory{}, fmt.Errorf("failed to parse command history from exec pane")
+	}
 	cmd := m.ExecHistory[len(m.ExecHistory)-1]
 	logger.Debug("Command: %s\nOutput: %s\nCode: %d\n", cmd.Command, cmd.Output, cmd.Code)
 	return cmd, nil
 }
 
 func (m *Manager) parseExecPaneCommandHistory() {
-	m.ExecPane.Refresh(m.GetMaxCaptureLines())
+	m.parseExecPaneCommandHistoryWithContent("")
+}
+
+func (m *Manager) parseExecPaneCommandHistoryWithContent(testContent string) {
+	if testContent == "" {
+		m.ExecPane.Refresh(m.GetMaxCaptureLines())
+	} else {
+		m.ExecPane.Content = testContent
+	}
 
 	var history []CommandExecHistory