feat(chat): support local, worktree, and remote branches

apps/desktop/src/main/agent/agent-store.ts

@@ -1,12 +1,17 @@
 import { randomUUID } from "node:crypto";
-import type { AgentSessionInfo, SubagentWorktreeInfo } from "../../shared/contracts";
+import type {
+  AgentSessionInfo,
+  SessionWorktreeInfo,
+  SubagentWorktreeInfo,
+} from "../../shared/contracts";
 import { getDatabase } from "../db/database";
 
 type AgentSessionRow = {
   id: string;
   workspace_id: string;
   title: string;
   cwd: string;
+  branch: string | null;
   status: AgentSessionInfo["status"];
   runtime: "pi-sdk" | "pi-rpc";
   model: string | null;
@@ -16,6 +21,11 @@ type AgentSessionRow = {
   subagent_task: string | null;
   subagent_type: string | null;
   subagent_readonly: number;
+  worktree_path: string | null;
+  worktree_branch: string | null;
+  worktree_base_branch: string | null;
+  worktree_base_sha: string | null;
+  worktree_status: string | null;
   subagent_worktree_path: string | null;
   subagent_worktree_branch: string | null;
   subagent_worktree_base_sha: string | null;
@@ -28,8 +38,9 @@ type AgentSessionRow = {
   updated_at: string;
 };
 
-const SESSION_COLUMNS = `id, workspace_id, title, cwd, status, runtime, model, pi_session_id,
+const SESSION_COLUMNS = `id, workspace_id, title, cwd, branch, status, runtime, model, pi_session_id,
   pi_session_file, parent_session_id, subagent_task, subagent_type, subagent_readonly,
+  worktree_path, worktree_branch, worktree_base_branch, worktree_base_sha, worktree_status,
   subagent_worktree_path, subagent_worktree_branch, subagent_worktree_base_sha,
   subagent_integration_status, subagent_changed_files_json, subagent_conflict_files_json,
   pinned_at, archived_at, created_at, updated_at`;
@@ -58,6 +69,9 @@ function toSession(row: AgentSessionRow): AgentSessionInfo {
     updatedAt: row.updated_at,
   };
 
+  if (row.branch !== null) {
+    session.branch = row.branch;
+  }
   if (row.model !== null) {
     session.model = row.model;
   }
@@ -85,6 +99,21 @@ function toSession(row: AgentSessionRow): AgentSessionInfo {
   if (row.subagent_readonly !== 0) {
     session.subagentReadOnly = true;
   }
+  if (
+    row.worktree_path !== null &&
+    row.worktree_branch !== null &&
+    row.worktree_base_branch !== null &&
+    row.worktree_base_sha !== null
+  ) {
+    const status = row.worktree_status as SessionWorktreeInfo["status"] | null;
+    session.worktree = {
+      path: row.worktree_path,
+      branch: row.worktree_branch,
+      baseBranch: row.worktree_base_branch,
+      baseSha: row.worktree_base_sha,
+      status: status ?? "active",
+    };
+  }
   if (
     row.subagent_worktree_path !== null &&
     row.subagent_worktree_branch !== null &&
@@ -112,6 +141,7 @@ export function createAgentSessionRecord(input: {
   workspaceId: string;
   title: string;
   cwd: string;
+  branch?: string;
   runtime?: "pi-sdk" | "pi-rpc";
   model?: string;
   piSessionId?: string;
@@ -120,6 +150,7 @@ export function createAgentSessionRecord(input: {
   subagentTask?: string;
   subagentType?: string;
   subagentReadOnly?: boolean;
+  worktree?: SessionWorktreeInfo;
   subagentWorktree?: SubagentWorktreeInfo;
 }): AgentSessionInfo {
   const now = new Date().toISOString();
@@ -135,6 +166,9 @@ export function createAgentSessionRecord(input: {
     updatedAt: now,
   };
 
+  if (input.branch !== undefined) {
+    session.branch = input.branch;
+  }
   if (input.model !== undefined) {
     session.model = input.model;
   }
@@ -156,25 +190,30 @@ export function createAgentSessionRecord(input: {
   if (input.subagentReadOnly !== undefined) {
     session.subagentReadOnly = input.subagentReadOnly;
   }
+  if (input.worktree !== undefined) {
+    session.worktree = input.worktree;
+  }
   if (input.subagentWorktree !== undefined) {
     session.subagentWorktree = input.subagentWorktree;
   }
   getDatabase()
     .prepare(
       `insert into agent_sessions (
-        id, workspace_id, title, cwd, status, runtime, model, pi_session_id, pi_session_file,
+        id, workspace_id, title, cwd, branch, status, runtime, model, pi_session_id, pi_session_file,
         parent_session_id, subagent_task, subagent_type, subagent_readonly,
+        worktree_path, worktree_branch, worktree_base_branch, worktree_base_sha, worktree_status,
         subagent_worktree_path, subagent_worktree_branch, subagent_worktree_base_sha,
         subagent_integration_status, subagent_changed_files_json, subagent_conflict_files_json,
         created_at, updated_at
        )
-       values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+       values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
     )
     .run(
       session.id,
       session.workspaceId,
       session.title,
       session.cwd,
+      session.branch ?? null,
       session.status,
       runtime,
       session.model ?? null,
@@ -184,6 +223,11 @@ export function createAgentSessionRecord(input: {
       session.subagentTask ?? null,
       session.subagentType ?? null,
       session.subagentReadOnly ? 1 : 0,
+      session.worktree?.path ?? null,
+      session.worktree?.branch ?? null,
+      session.worktree?.baseBranch ?? null,
+      session.worktree?.baseSha ?? null,
+      session.worktree?.status ?? null,
       session.subagentWorktree?.path ?? null,
       session.subagentWorktree?.branch ?? null,
       session.subagentWorktree?.baseSha ?? null,
@@ -202,6 +246,39 @@ export function createAgentSessionRecord(input: {
 }
 
 export function updateAgentSessionWorktree(
+  sessionId: string,
+  worktree: SessionWorktreeInfo | undefined,
+): AgentSessionInfo | undefined {
+  const existing = getAgentSession(sessionId);
+  if (!existing) {
+    return undefined;
+  }
+
+  getDatabase()
+    .prepare(
+      `update agent_sessions
+       set worktree_path = ?,
+           worktree_branch = ?,
+           worktree_base_branch = ?,
+           worktree_base_sha = ?,
+           worktree_status = ?,
+           updated_at = ?
+       where id = ?`,
+    )
+    .run(
+      worktree?.path ?? null,
+      worktree?.branch ?? null,
+      worktree?.baseBranch ?? null,
+      worktree?.baseSha ?? null,
+      worktree?.status ?? null,
+      new Date().toISOString(),
+      sessionId,
+    );
+
+  return getAgentSession(sessionId);
+}
+
+export function updateAgentSessionSubagentWorktree(
   sessionId: string,
   worktree: SubagentWorktreeInfo | undefined,
 ): AgentSessionInfo | undefined {

apps/desktop/src/main/agent/pi-rpc-service.ts

@@ -27,7 +27,7 @@ function stringifyLine(line: unknown): string {
 
 export function createPiRpcSession(
   window: BrowserWindowType,
-  input: { workspaceId: string; cwd: string; title: string },
+  input: { workspaceId: string; cwd: string; title: string; branch?: string },
 ): AgentSessionInfo {
   const info = createAgentSessionRecord({ ...input, runtime: "pi-rpc" });
   const sessionDir = join(app.getPath("userData"), "pi-sessions");

apps/desktop/src/main/agent/pi-sdk-runtime.ts

@@ -52,8 +52,8 @@ import {
   listSubagentSessions,
   updateAgentSessionMetadata,
   updateAgentSessionStatus,
+  updateAgentSessionSubagentWorktree,
   updateAgentSessionTitle,
-  updateAgentSessionWorktree,
 } from "./agent-store";
 import { createCheckpoint } from "./checkpoint-service";
 import {
@@ -511,12 +511,14 @@ export class PiSdkRuntime implements AgentRuntime {
       workspaceId: input.workspaceId,
       cwd: input.cwd,
       title: input.title,
+      ...(input.branch !== undefined ? { branch: input.branch } : {}),
       runtime: "pi-sdk",
       ...(input.id !== undefined ? { id: input.id } : {}),
       ...(input.parentSessionId !== undefined ? { parentSessionId: input.parentSessionId } : {}),
       ...(input.subagentTask !== undefined ? { subagentTask: input.subagentTask } : {}),
       ...(input.subagentType !== undefined ? { subagentType: input.subagentType } : {}),
       ...(input.subagentReadOnly !== undefined ? { subagentReadOnly: input.subagentReadOnly } : {}),
+      ...(input.worktree !== undefined ? { worktree: input.worktree } : {}),
       ...(input.subagentWorktree !== undefined ? { subagentWorktree: input.subagentWorktree } : {}),
     };
     if (modelId !== undefined) {
@@ -1134,7 +1136,7 @@ export class PiSdkRuntime implements AgentRuntime {
           return;
         }
         const updated = await finishSubagentWorktree(session.subagentWorktree, input.task);
-        updateAgentSessionWorktree(session.id, updated);
+        updateAgentSessionSubagentWorktree(session.id, updated);
         if (!runFailed) {
           emit({
             type: "subagent.updated",

apps/desktop/src/main/agent/runtime.ts

@@ -17,11 +17,13 @@ export type CreateAgentRuntimeInput = {
   workspaceId: string;
   cwd: string;
   title: string;
+  branch?: string;
   model?: string;
   parentSessionId?: string;
   subagentTask?: string;
   subagentType?: string;
   subagentReadOnly?: boolean;
+  worktree?: AgentSessionInfo["worktree"];
   subagentWorktree?: AgentSessionInfo["subagentWorktree"];
 };
 

apps/desktop/src/main/db/database.ts

@@ -170,6 +170,7 @@ function migrate(db: DatabaseSync): void {
 
   addColumn(db, "agent_sessions", "runtime", "text not null default 'pi-sdk'");
   addColumn(db, "agent_sessions", "model", "text");
+  addColumn(db, "agent_sessions", "branch", "text");
   addColumn(db, "agent_sessions", "pi_session_id", "text");
   addColumn(db, "agent_sessions", "pi_session_file", "text");
   addColumn(
@@ -181,6 +182,11 @@ function migrate(db: DatabaseSync): void {
   addColumn(db, "agent_sessions", "subagent_task", "text");
   addColumn(db, "agent_sessions", "subagent_type", "text");
   addColumn(db, "agent_sessions", "subagent_readonly", "integer not null default 0");
+  addColumn(db, "agent_sessions", "worktree_path", "text");
+  addColumn(db, "agent_sessions", "worktree_branch", "text");
+  addColumn(db, "agent_sessions", "worktree_base_branch", "text");
+  addColumn(db, "agent_sessions", "worktree_base_sha", "text");
+  addColumn(db, "agent_sessions", "worktree_status", "text");
   addColumn(db, "agent_sessions", "subagent_worktree_path", "text");
   addColumn(db, "agent_sessions", "subagent_worktree_branch", "text");
   addColumn(db, "agent_sessions", "subagent_worktree_base_sha", "text");
@@ -196,18 +202,6 @@ function migrate(db: DatabaseSync): void {
   // Sidebar project pinning: pinned projects sort to the top (pinned_at breaks ties).
   addColumn(db, "workspaces", "pinned", "integer not null default 0");
   addColumn(db, "workspaces", "pinned_at", "text");
-  if (hasColumn(db, "agent_sessions", "worktree_path")) {
-    db.exec(`
-      update agent_sessions
-      set
-        cwd = coalesce(
-          (select root_path from workspaces where workspaces.id = agent_sessions.workspace_id),
-          cwd
-        ),
-        worktree_path = null
-      where worktree_path is not null
-    `);
-  }
   // PI session-tree leaf id captured right before each prompt — the exact
   // branch point used to rewind the conversation when the message is edited.
   // "root" marks an empty tree (first message); NULL marks legacy runs.

apps/desktop/src/main/git/git-service.test.ts

@@ -9,8 +9,10 @@ import {
   abortSubagentWorktreeApply,
   applySubagentWorktree,
   checkoutBranch,
+  cleanupSessionWorktree,
   cleanupSubagentWorktree,
   commitOrPush,
+  createChatWorktree,
   createSubagentWorktree,
   discardFile,
   finishSubagentWorktree,
@@ -276,6 +278,55 @@ describe("git-service", () => {
     }
   });
 
+  it("creates a chat worktree from a selected local base branch", async () => {
+    const baseBranch = (await git(["branch", "--show-current"])).trim();
+    const worktree = await createChatWorktree(repo, {
+      sessionId: "abcdef12-3456-7890-abcd-ef1234567890",
+      baseBranch,
+    });
+
+    expect(worktree.path.replace(/\\/g, "/")).toContain(
+      `/.modus/worktrees/chat-${baseBranch}-abcdef12`,
+    );
+    expect(worktree.branch).toBe(`modus/chat/${baseBranch}-abcdef12`);
+    expect(worktree.baseBranch).toBe(baseBranch);
+    expect(worktree.status).toBe("active");
+    expect(existsSync(worktree.path)).toBe(true);
+
+    const cleaned = await cleanupSessionWorktree(repo, worktree);
+    expect(cleaned.status).toBe("cleaned");
+    expect(existsSync(worktree.path)).toBe(false);
+  });
+
+  it("creates a chat worktree directly from a remote-tracking base branch", async () => {
+    const baseSha = (await git(["rev-parse", "HEAD"])).trim();
+    await git(["update-ref", "refs/remotes/origin/feature/remote", baseSha]);
+
+    const worktree = await createChatWorktree(repo, {
+      sessionId: "abcdef12-3456-7890-abcd-ef1234567890",
+      baseBranch: "origin/feature/remote",
+      baseBranchRemote: true,
+    });
+
+    expect(worktree.baseBranch).toBe("origin/feature/remote");
+    expect(worktree.baseSha).toBe(baseSha);
+    expect((await execFileAsync("git", ["branch", "--show-current"], { cwd: worktree.path })).stdout.trim()).toBe(
+      worktree.branch,
+    );
+    expect((await git(["branch", "--list", "feature/remote"])).trim()).toBe("");
+
+    await cleanupSessionWorktree(repo, worktree);
+  });
+
+  it("rejects a chat worktree from a missing local base branch", async () => {
+    await expect(
+      createChatWorktree(repo, {
+        sessionId: "abcdef12-3456-7890-abcd-ef1234567890",
+        baseBranch: "origin/main",
+      }),
+    ).rejects.toThrow("not a local branch");
+  });
+
   it("creates, finishes, applies, and cleans up a subagent worktree", async () => {
     const worktree = await createSubagentWorktree(repo, {
       sessionId: "abcdef12-3456-7890-abcd-ef1234567890",