CLI: Add persistent shell mode to CLI with workspace management

cli/README.md

@@ -32,6 +32,17 @@ We've only tested the CLI on Mac and Linux, and are aware that there are some is
 
 ## Usage
 
+### Shell Mode
+
+Kilo Code CLI includes a persistent shell mode that allows you to execute shell commands while maintaining directory changes and environment variables across commands. To enter shell mode, press `Shift+!` or type `!` when the input is empty.
+
+#### Features
+
+- **Persistent Session**: Directory changes (`cd`) and environment variable modifications persist for the duration of the shell session
+- **Workspace Propagation**: When you exit shell mode, the CLI's working directory is updated to match the shell's current directory, and the AI context is reinitialized for the new workspace
+- **Command History**: Navigate through previous commands using the up/down arrow keys
+- **Real-time Output**: Command output is displayed immediately in the CLI interface
+
 ### Interactive Mode
 
 ```bash

cli/src/cli.ts

@@ -6,6 +6,8 @@ import { createExtensionService, ExtensionService } from "./services/extension.j
 import { App } from "./ui/App.js"
 import { logs } from "./services/logs.js"
 import { extensionServiceAtom } from "./state/atoms/service.js"
+import { addMessageAtom } from "./state/atoms/ui.js"
+import { disposeShellSessionAtom } from "./state/atoms/shell.js"
 import { initializeServiceEffectAtom } from "./state/atoms/effects.js"
 import { loadConfigAtom, mappedExtensionStateAtom, providersAtom, saveConfigAtom } from "./state/atoms/config.js"
 import { ciExitReasonAtom } from "./state/atoms/ci.js"
@@ -195,6 +197,7 @@ export class CLI {
 					parallel: this.options.parallel || false,
 					worktreeBranch: this.options.worktreeBranch || undefined,
 					noSplash: this.options.noSplash || false,
+					onWorkspaceChange: (newWorkspace: string) => this.handleWorkspaceChange(newWorkspace),
 				},
 				onExit: () => this.dispose(),
 			}),
@@ -322,6 +325,11 @@ export class CLI {
 				this.service = null
 			}
 
+			// Dispose shell session
+			if (this.store) {
+				this.store.set(disposeShellSessionAtom)
+			}
+
 			// Clear store reference
 			this.store = null
 
@@ -418,6 +426,86 @@ export class CLI {
 		}
 	}
 
+	/**
+	 * Handle workspace changes from shell mode
+	 */
+	public async handleWorkspaceChange(newWorkspace: string): Promise<void> {
+		if (!this.store || !this.service) {
+			logs.warn("Cannot handle workspace change: service or store not available", "CLI")
+			return
+		}
+
+		try {
+			logs.info("Handling workspace change", "CLI", { oldWorkspace: this.options.workspace, newWorkspace })
+
+			// Update CLI options
+			this.options.workspace = newWorkspace
+
+			// Show system message about workspace change
+			const workspaceChangeMessage = {
+				id: `workspace-change-${Date.now()}`,
+				type: "system" as const,
+				ts: Date.now(),
+				content: `📁 Workspace changed to: ${newWorkspace}\nReinitializing AI context...`,
+				partial: false,
+			}
+			this.store.set(addMessageAtom, workspaceChangeMessage)
+
+			// Dispose current extension service
+			await this.service.dispose()
+			this.service = null
+
+			// Create new extension service with updated workspace
+			const telemetryService = getTelemetryService()
+			const identityManager = getIdentityManager()
+			const identity = identityManager.getIdentity()
+
+			const serviceOptions: Parameters<typeof createExtensionService>[0] = {
+				workspace: newWorkspace,
+				mode: this.options.mode || "code",
+			}
+
+			if (identity) {
+				serviceOptions.identity = {
+					machineId: identity.machineId,
+					sessionId: identity.sessionId,
+					cliUserId: identity.cliUserId,
+				}
+			}
+
+			this.service = createExtensionService(serviceOptions)
+
+			// Set service in store
+			this.store.set(extensionServiceAtom, this.service)
+
+			// Reinitialize service through effect atom
+			await this.store.set(initializeServiceEffectAtom, this.store)
+
+			// Re-inject configuration
+			await this.injectConfigurationToExtension()
+
+			// Re-request router models
+			await this.requestRouterModels()
+
+			// Update telemetry context (access private fields directly since there's no public update method)
+			;(telemetryService as unknown as { currentWorkspace: string }).currentWorkspace = newWorkspace
+
+			logs.info("Workspace change handled successfully", "CLI", { newWorkspace })
+		} catch (error) {
+			logs.error("Failed to handle workspace change", "CLI", { error, newWorkspace })
+
+			// Show error message
+			const errorMessage = {
+				id: `workspace-change-error-${Date.now()}`,
+				type: "error" as const,
+				ts: Date.now(),
+				content: `❌ Failed to change workspace: ${error instanceof Error ? error.message : error}`,
+				partial: false,
+			}
+			this.store.set(addMessageAtom, errorMessage)
+		}
+	}
+
 	/**
 	 * Resume the last conversation from the current workspace
 	 */

cli/src/services/__tests__/shellSession.test.ts

@@ -0,0 +1,229 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
+import { ShellSession } from "../shell/session.js"
+import { spawn, ChildProcess } from "child_process"
+import { EventEmitter } from "events"
+
+// Define mock types for testing
+interface MockStdin {
+	write: ReturnType<typeof vi.fn>
+	end: ReturnType<typeof vi.fn>
+}
+
+interface MockChildProcess extends EventEmitter {
+	stdout: EventEmitter
+	stderr: EventEmitter
+	stdin: MockStdin
+	kill: ReturnType<typeof vi.fn>
+}
+
+// Mock the default-shell module
+vi.mock("default-shell", () => ({
+	detectDefaultShell: vi.fn(() => "/bin/bash"),
+}))
+
+// Mock child_process.spawn
+vi.mock("child_process", () => ({
+	spawn: vi.fn(),
+}))
+
+describe("ShellSession", () => {
+	let mockChildProcess: MockChildProcess
+	let stdoutEmitter: EventEmitter
+	let stderrEmitter: EventEmitter
+	let stdin: MockStdin
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		// Use EventEmitters for stdout/stderr to easily simulate data events
+		stdoutEmitter = new EventEmitter()
+		stderrEmitter = new EventEmitter()
+
+		stdin = {
+			write: vi.fn(),
+			end: vi.fn(),
+		}
+
+		mockChildProcess = new EventEmitter()
+		mockChildProcess.stdout = stdoutEmitter
+		mockChildProcess.stderr = stderrEmitter
+		mockChildProcess.stdin = stdin
+		mockChildProcess.kill = vi.fn()
+
+		vi.mocked(spawn).mockReturnValue(mockChildProcess as unknown as ChildProcess)
+	})
+
+	afterEach(() => {
+		vi.restoreAllMocks()
+	})
+
+	describe("initialization", () => {
+		it("should create a new shell session with default shell", async () => {
+			const session = new ShellSession()
+
+			await session.ensureSession("/tmp/test")
+
+			expect(spawn).toHaveBeenCalledWith("/bin/bash", [], {
+				cwd: "/tmp/test",
+				stdio: ["pipe", "pipe", "pipe"],
+				env: expect.objectContaining({
+					PS1: "",
+					// PROMPT_COMMAND is no longer set in env
+				}),
+				detached: false,
+			})
+		})
+
+		it("should handle shell spawn errors", async () => {
+			const session = new ShellSession()
+
+			// ensureSession doesn't reject on spawn error because spawn error is async
+			// and ensureSession returns immediately.
+			// We should verify that the error event is emitted.
+
+			const errorPromise = new Promise<void>((resolve, reject) => {
+				session.on("error", (err) => {
+					try {
+						expect(err.message).toBe("Spawn failed")
+						resolve()
+					} catch (e) {
+						reject(e)
+					}
+				})
+			})
+
+			await session.ensureSession("/tmp/test")
+
+			// Simulate spawn error
+			mockChildProcess.emit("error", new Error("Spawn failed"))
+
+			await errorPromise
+		})
+	})
+
+	describe("command execution", () => {
+		let session: ShellSession
+
+		beforeEach(async () => {
+			session = new ShellSession()
+			await session.ensureSession("/tmp/test")
+		})
+
+		it("should execute commands and parse sentinel output", async () => {
+			const runPromise = session.run("echo hello")
+
+			// Simulate output
+			stdoutEmitter.emit("data", "__KILO_DONE__:0:/tmp/test\n")
+
+			const result = await runPromise
+
+			expect(result).toEqual({
+				stdout: "",
+				stderr: "",
+				exitCode: 0,
+				cwd: "/tmp/test",
+			})
+
+			expect(stdin.write).toHaveBeenCalledWith(
+				expect.stringContaining("echo hello"),
+				"utf8",
+				expect.any(Function),
+			)
+		})
+
+		it("should handle command output", async () => {
+			const runPromise = session.run("echo hello")
+
+			// Simulate output
+			stdoutEmitter.emit("data", "hello world\n__KILO_DONE__:0:/tmp/test\n")
+
+			const result = await runPromise
+
+			expect(result).toEqual({
+				stdout: "hello world",
+				stderr: "",
+				exitCode: 0,
+				cwd: "/tmp/test",
+			})
+		})
+
+		it("should handle stderr output", async () => {
+			const runPromise = session.run("failing command")
+
+			// Simulate output
+			stderrEmitter.emit("data", "error message\n")
+			stdoutEmitter.emit("data", "__KILO_DONE__:1:/tmp/test\n")
+
+			const result = await runPromise
+
+			expect(result).toEqual({
+				stdout: "",
+				stderr: "error message",
+				exitCode: 1,
+				cwd: "/tmp/test",
+			})
+		})
+
+		it("should timeout long-running commands", async () => {
+			const session = new ShellSession({ timeout: 50 })
+			// We need to initialize the session first (which was done in beforeEach but for the other 'session' instance)
+			// Actually beforeEach initializes 'session'. But here we create a new one.
+
+			// Mock spawn again for this new session if needed, but the global mock persists.
+			await session.ensureSession("/tmp/test")
+
+			await expect(session.run("sleep 1")).rejects.toThrow("Command timeout after 50ms")
+		})
+
+		it("should reject concurrent commands", async () => {
+			const promise1 = session.run("cmd1")
+
+			await expect(session.run("cmd2")).rejects.toThrow("Command already in progress")
+
+			// Complete the first command
+			stdoutEmitter.emit("data", "__KILO_DONE__:0:/tmp/test\n")
+
+			await promise1
+		})
+	})
+
+	describe("directory tracking", () => {
+		it("should track current directory from sentinel", async () => {
+			const session = new ShellSession()
+			await session.ensureSession("/tmp/test")
+
+			const runPromise = session.run("cd /home/user")
+
+			stdoutEmitter.emit("data", "__KILO_DONE__:0:/home/user\n")
+
+			await runPromise
+
+			expect(session.getCurrentDirectory()).toBe("/home/user")
+		})
+	})
+
+	describe("lifecycle", () => {
+		it("should dispose properly", async () => {
+			const session = new ShellSession()
+			await session.ensureSession("/tmp/test")
+
+			session.dispose()
+
+			expect(mockChildProcess.kill).toHaveBeenCalled()
+			expect(session.isSessionReady()).toBe(false)
+		})
+
+		it("should handle process exit", async () => {
+			const session = new ShellSession()
+			await session.ensureSession("/tmp/test")
+
+			const runPromise = session.run("test")
+
+			mockChildProcess.emit("close", 1, "SIGTERM")
+
+			await expect(runPromise).rejects.toThrow("Shell process exited unexpectedly with code 1")
+
+			expect(session.isSessionReady()).toBe(false)
+		})
+	})
+})