Merge pull request #4145 from Kilo-Org/add-session-commands

Author: iscekic

.changeset/late-paths-kick.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+update shared session url

.changeset/twenty-rocks-joke.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": minor
+---
+
+add session sharing and forking

cli/src/commands/__tests__/session.test.ts

@@ -100,7 +100,7 @@ describe("sessionCommand", () => {
 			expect(sessionCommand.examples).toContain("/session search <query>")
 			expect(sessionCommand.examples).toContain("/session select <sessionId>")
 			expect(sessionCommand.examples).toContain("/session share")
-			expect(sessionCommand.examples).toContain("/session fork <shareId>")
+			expect(sessionCommand.examples).toContain("/session fork <id>")
 			expect(sessionCommand.examples).toContain("/session delete <sessionId>")
 			expect(sessionCommand.examples).toContain("/session rename <new name>")
 		})
@@ -621,7 +621,7 @@ describe("sessionCommand", () => {
 
 			const replacedMessages = (mockContext.replaceMessages as ReturnType<typeof vi.fn>).mock.calls[0][0]
 			expect(replacedMessages).toHaveLength(2)
-			expect(replacedMessages[1].content).toContain("Forking session from share ID")
+			expect(replacedMessages[1].content).toContain("Forking session from ID")
 			expect(replacedMessages[1].content).toContain("share-123")
 		})
 
@@ -633,7 +633,7 @@ describe("sessionCommand", () => {
 			expect(mockContext.addMessage).toHaveBeenCalledTimes(1)
 			const message = (mockContext.addMessage as ReturnType<typeof vi.fn>).mock.calls[0][0]
 			expect(message.type).toBe("error")
-			expect(message.content).toContain("Usage: /session fork <shareId>")
+			expect(message.content).toContain("Usage: /session fork <id>")
 			expect(mockSessionManager.forkSession).not.toHaveBeenCalled()
 		})
 
@@ -645,7 +645,7 @@ describe("sessionCommand", () => {
 			expect(mockContext.addMessage).toHaveBeenCalledTimes(1)
 			const message = (mockContext.addMessage as ReturnType<typeof vi.fn>).mock.calls[0][0]
 			expect(message.type).toBe("error")
-			expect(message.content).toContain("Usage: /session fork <shareId>")
+			expect(message.content).toContain("Usage: /session fork <id>")
 		})
 
 		it("should handle fork error gracefully", async () => {

cli/src/commands/session.ts

@@ -41,10 +41,8 @@ async function listSessions(context: CommandContext): Promise<void> {
 	const sessionClient = sessionService.sessionClient
 
 	try {
-		const result = await sessionClient.list({ limit: 50 })
-		const { cliSessions } = result
-
-		if (cliSessions.length === 0) {
+		const result = await sessionClient?.list({ limit: 50 })
+		if (!result || result.cliSessions.length === 0) {
 			addMessage({
 				...generateMessage(),
 				type: "system",
@@ -53,6 +51,8 @@ async function listSessions(context: CommandContext): Promise<void> {
 			return
 		}
 
+		const { cliSessions } = result
+
 		// Format and display sessions
 		let content = `**Available Sessions:**\n\n`
 		cliSessions.forEach((session, index) => {
@@ -148,10 +148,9 @@ async function searchSessions(context: CommandContext, query: string): Promise<v
 	}
 
 	try {
-		const result = await sessionClient.search({ search_string: query, limit: 20 })
-		const { results, total } = result
+		const result = await sessionClient?.search({ search_string: query, limit: 20 })
 
-		if (results.length === 0) {
+		if (!result || result.results.length === 0) {
 			addMessage({
 				...generateMessage(),
 				type: "system",
@@ -160,6 +159,8 @@ async function searchSessions(context: CommandContext, query: string): Promise<v
 			return
 		}
 
+		const { results, total } = result
+
 		let content = `**Search Results** (${results.length} of ${total}):\n\n`
 		results.forEach((session, index) => {
 			const isActive = session.session_id === sessionService.sessionId ? " * [Active]" : ""
@@ -198,7 +199,7 @@ async function shareSession(context: CommandContext): Promise<void> {
 		addMessage({
 			...generateMessage(),
 			type: "system",
-			content: `✅ Session shared successfully!\n\n\`https://kilo.ai/share/${result.share_id}\``,
+			content: `✅ Session shared successfully!\n\n\`https://app.kilo.ai/share/${result.share_id}\``,
 		})
 	} catch (error) {
 		addMessage({
@@ -212,15 +213,15 @@ async function shareSession(context: CommandContext): Promise<void> {
 /**
  * Fork a shared session by share ID
  */
-async function forkSession(context: CommandContext, shareId: string): Promise<void> {
+async function forkSession(context: CommandContext, id: string): Promise<void> {
 	const { addMessage, replaceMessages, refreshTerminal } = context
 	const sessionService = SessionManager.init()
 
-	if (!shareId) {
+	if (!id) {
 		addMessage({
 			...generateMessage(),
 			type: "error",
-			content: "Usage: /session fork <shareId>",
+			content: "Usage: /session fork <id>",
 		})
 		return
 	}
@@ -238,14 +239,14 @@ async function forkSession(context: CommandContext, shareId: string): Promise<vo
 			{
 				id: `system-${now + 1}`,
 				type: "system",
-				content: `Forking session from share ID \`${shareId}\`...`,
+				content: `Forking session from ID \`${id}\`...`,
 				ts: 2,
 			},
 		])
 
 		await refreshTerminal()
 
-		await sessionService.forkSession(shareId, true)
+		await sessionService.forkSession(id, true)
 
 		// Success message handled by restoreSession via extension messages
 	} catch (error) {
@@ -275,6 +276,10 @@ async function deleteSession(context: CommandContext, sessionId: string): Promis
 	}
 
 	try {
+		if (!sessionClient) {
+			throw new Error("SessionManager used before initialization")
+		}
+
 		await sessionClient.delete({ session_id: sessionId })
 
 		addMessage({
@@ -340,7 +345,11 @@ async function sessionIdAutocompleteProvider(context: ArgumentProviderContext):
 	}
 
 	try {
-		const response = await sessionClient.search({ search_string: prefix, limit: 20 })
+		const response = await sessionClient?.search({ search_string: prefix, limit: 20 })
+
+		if (!response) {
+			return []
+		}
 
 		return response.results.map((session, index) => {
 			const title = session.title || "Untitled"
@@ -373,7 +382,7 @@ export const sessionCommand: Command = {
 		"/session search <query>",
 		"/session select <sessionId>",
 		"/session share",
-		"/session fork <shareId>",
+		"/session fork <id>",
 		"/session delete <sessionId>",
 		"/session rename <new name>",
 	],
@@ -390,7 +399,7 @@ export const sessionCommand: Command = {
 				{ value: "search", description: "Search sessions by title or ID" },
 				{ value: "select", description: "Restore a session" },
 				{ value: "share", description: "Share current session publicly" },
-				{ value: "fork", description: "Fork a shared session" },
+				{ value: "fork", description: "Fork a session" },
 				{ value: "delete", description: "Delete a session" },
 				{ value: "rename", description: "Rename the current session" },
 			],

cli/src/index.ts

@@ -44,7 +44,7 @@ program
 	.option("-pv, --provider <id>", "Select provider by ID (e.g., 'kilocode-1')")
 	.option("-mo, --model <model>", "Override model for the selected provider")
 	.option("-s, --session <sessionId>", "Restore a session by ID")
-	.option("-f, --fork <shareId>", "Fork a shared session by share ID")
+	.option("-f, --fork <shareId>", "Fork a session by ID")
 	.option("--nosplash", "Disable the welcome message and update notifications", false)
 	.argument("[prompt]", "The prompt or command to execute")
 	.action(async (prompt, options) => {

src/activate/handleUri.ts

@@ -46,6 +46,21 @@ export const handleUri = async (uri: vscode.Uri) => {
 			})
 			break
 		}
+		case "/kilo/fork": {
+			const id = query.get("id")
+			if (id) {
+				await visibleProvider.postMessageToWebview({
+					type: "invoke",
+					invoke: "setChatBoxMessage",
+					text: `/session fork ${id}`,
+				})
+				await visibleProvider.postMessageToWebview({
+					type: "action",
+					action: "focusInput",
+				})
+			}
+			break
+		}
 		// kilocode_change end
 		case "/requesty": {
 			const code = query.get("code")

src/core/webview/webviewMessageHandler.ts

@@ -93,6 +93,7 @@ import { getSapAiCoreDeployments } from "../../api/providers/fetchers/sap-ai-cor
 import { AutoPurgeScheduler } from "../../services/auto-purge" // kilocode_change
 import { setPendingTodoList } from "../tools/UpdateTodoListTool"
 import { ManagedIndexer } from "../../services/code-index/managed/ManagedIndexer"
+import { SessionManager } from "../../shared/kilocode/cli-sessions/core/SessionManager" // kilocode_change
 
 export const webviewMessageHandler = async (
 	provider: ClineProvider,
@@ -3842,6 +3843,72 @@ export const webviewMessageHandler = async (
 			}
 			break
 		}
+		case "sessionShare": {
+			try {
+				const sessionService = SessionManager.init()
+
+				if (!sessionService.sessionId) {
+					vscode.window.showErrorMessage("No active session. Start a new task to create a session.")
+					break
+				}
+
+				const result = await sessionService.shareSession()
+
+				const shareUrl = `https://app.kilo.ai/share/${result.share_id}`
+
+				// Copy URL to clipboard and show success notification
+				await vscode.env.clipboard.writeText(shareUrl)
+				vscode.window.showInformationMessage(`Session shared! Link copied to clipboard: ${shareUrl}`)
+			} catch (error) {
+				const errorMessage = error instanceof Error ? error.message : String(error)
+				vscode.window.showErrorMessage(`Failed to share session: ${errorMessage}`)
+			}
+			break
+		}
+		case "shareTaskSession": {
+			try {
+				if (!message.text) {
+					vscode.window.showErrorMessage("Task ID is required for sharing a task session")
+					break
+				}
+
+				const taskId = message.text
+				const sessionService = SessionManager.init()
+
+				const sessionId = await sessionService.getSessionFromTask(taskId, provider)
+
+				const result = await sessionService.shareSession(sessionId)
+
+				const shareUrl = `https://app.kilo.ai/share/${result.share_id}`
+
+				await vscode.env.clipboard.writeText(shareUrl)
+				vscode.window.showInformationMessage(`Session shared! Link copied to clipboard.`)
+			} catch (error) {
+				const errorMessage = error instanceof Error ? error.message : String(error)
+				vscode.window.showErrorMessage(`Failed to share task session: ${errorMessage}`)
+			}
+			break
+		}
+		case "sessionFork": {
+			try {
+				if (!message.shareId) {
+					vscode.window.showErrorMessage("ID is required for forking a session")
+					break
+				}
+
+				const sessionService = SessionManager.init()
+
+				await sessionService.forkSession(message.shareId, true)
+
+				await provider.postStateToWebview()
+
+				vscode.window.showInformationMessage(`Session forked successfully from ${message.shareId}`)
+			} catch (error) {
+				const errorMessage = error instanceof Error ? error.message : String(error)
+				vscode.window.showErrorMessage(`Failed to fork session: ${errorMessage}`)
+			}
+			break
+		}
 		case "singleCompletion": {
 			try {
 				const { text, completionRequestId } = message

src/extension.ts

@@ -51,6 +51,30 @@ import { flushModels, getModels } from "./api/providers/fetchers/modelCache"
 import { ManagedIndexer } from "./services/code-index/managed/ManagedIndexer" // kilocode_change
 import { kilo_initializeSessionManager } from "./shared/kilocode/cli-sessions/extension/session-manager-utils"
 
+// kilocode_change start
+async function findKilocodeTokenFromAnyProfile(provider: ClineProvider): Promise<string | undefined> {
+	const { apiConfiguration } = await provider.getState()
+	if (apiConfiguration.kilocodeToken) {
+		return apiConfiguration.kilocodeToken
+	}
+
+	const profiles = await provider.providerSettingsManager.listConfig()
+
+	for (const profile of profiles) {
+		try {
+			const fullProfile = await provider.providerSettingsManager.getProfile({ name: profile.name })
+			if (fullProfile.kilocodeToken) {
+				return fullProfile.kilocodeToken
+			}
+		} catch {
+			continue
+		}
+	}
+
+	return undefined
+}
+// kilocode_change end
+
 /**
  * Built using https://github.com/microsoft/vscode-webview-ui-toolkit
  *
@@ -270,11 +294,11 @@ export async function activate(context: vscode.ExtensionContext) {
 
 	// kilocode_change start
 	try {
-		const { apiConfiguration } = await provider.getState()
+		const kiloToken = await findKilocodeTokenFromAnyProfile(provider)
 
 		await kilo_initializeSessionManager({
 			context: context,
-			kiloToken: apiConfiguration.kilocodeToken,
+			kiloToken,
 			log: provider.log.bind(provider),
 			outputChannel,
 			provider,

src/services/kilo-session/ExtensionMessengerImpl.ts

@@ -3,12 +3,14 @@ import type { WebviewMessage } from "../../shared/kilocode/cli-sessions/types/IE
 import type { ExtensionMessage } from "../../shared/ExtensionMessage"
 import type { ClineProvider } from "../../core/webview/ClineProvider"
 import { singleCompletionHandler } from "../../utils/single-completion-handler"
+import { webviewMessageHandler } from "../../core/webview/webviewMessageHandler"
 
 export class ExtensionMessengerImpl implements IExtensionMessenger {
 	constructor(private readonly provider: ClineProvider) {}
 
+	// we can directly handle whatever is sent
 	async sendWebviewMessage(message: WebviewMessage): Promise<void> {
-		await this.provider.postMessageToWebview(message as unknown as ExtensionMessage)
+		return webviewMessageHandler(this.provider, message)
 	}
 
 	async requestSingleCompletion(prompt: string, timeoutMs: number): Promise<string> {

src/services/kilo-session/ExtensionPathProvider.ts

@@ -9,34 +9,20 @@ export class ExtensionPathProvider implements IPathProvider {
 
 	constructor(context: vscode.ExtensionContext) {
 		this.globalStoragePath = context.globalStorageUri.fsPath
-		this.ensureDirectories()
-	}
-
-	private ensureDirectories(): void {
-		const sessionsDir = path.join(this.globalStoragePath, "sessions")
-		const tasksDir = this.getTasksDir()
-		const workspacesDir = path.join(sessionsDir, "workspaces")
-
-		for (const dir of [sessionsDir, tasksDir, workspacesDir]) {
-			if (!existsSync(dir)) {
-				mkdirSync(dir, { recursive: true })
-			}
-		}
 	}
 
 	getTasksDir(): string {
-		return path.join(this.globalStoragePath, "sessions", "tasks")
+		return path.join(this.globalStoragePath, "tasks")
 	}
 
-	getSessionFilePath(workspaceDir: string): string {
-		const hash = createHash("sha256").update(workspaceDir).digest("hex").substring(0, 16)
-		const workspacesDir = path.join(this.globalStoragePath, "sessions", "workspaces")
-		const workspaceSessionDir = path.join(workspacesDir, hash)
+	getSessionFilePath(workspaceName: string): string {
+		const hash = createHash("sha256").update(workspaceName).digest("hex").substring(0, 16)
+		const workspaceDir = path.join(this.globalStoragePath, "sessions", hash)
 
-		if (!existsSync(workspaceSessionDir)) {
-			mkdirSync(workspaceSessionDir, { recursive: true })
+		if (!existsSync(workspaceDir)) {
+			mkdirSync(workspaceDir, { recursive: true })
 		}
 
-		return path.join(workspaceSessionDir, "session.json")
+		return path.join(workspaceDir, "session.json")
 	}
 }