feat(code): session forking

apps/code/src/main/db/migrations/0006_fork_relationships.sql

@@ -0,0 +1,12 @@
+CREATE TABLE `fork_relationships` (
+	`id` text PRIMARY KEY NOT NULL,
+	`forked_task_id` text NOT NULL,
+	`source_task_id` text NOT NULL,
+	`source_task_run_id` text NOT NULL,
+	`source_task_title` text NOT NULL,
+	`fork_at_message_index` integer NOT NULL,
+	`forked_at` text NOT NULL,
+	`created_at` text DEFAULT (CURRENT_TIMESTAMP) NOT NULL
+);
+--> statement-breakpoint
+CREATE UNIQUE INDEX `fork_relationships_forked_task_id_unique` ON `fork_relationships` (`forked_task_id`);

apps/code/src/main/db/migrations/meta/_journal.json

@@ -43,6 +43,13 @@
       "when": 1775755977659,
       "tag": "0005_youthful_scarlet_spider",
       "breakpoints": true
+    },
+    {
+      "idx": 6,
+      "version": "6",
+      "when": 1748131200000,
+      "tag": "0006_fork_relationships",
+      "breakpoints": true
     }
   ]
 }

apps/code/src/main/db/repositories/fork-relationship-repository.ts

@@ -0,0 +1,54 @@
+import { eq } from "drizzle-orm";
+import { inject, injectable } from "inversify";
+import { MAIN_TOKENS } from "../../di/tokens";
+import { forkRelationships } from "../schema";
+import type { DatabaseService } from "../service";
+
+export type ForkRelationship = typeof forkRelationships.$inferSelect;
+
+export interface CreateForkRelationshipData {
+  forkedTaskId: string;
+  sourceTaskId: string;
+  sourceTaskRunId: string;
+  sourceTaskTitle: string;
+  forkAtMessageIndex: number;
+  forkedAt: string;
+}
+
+export interface IForkRelationshipRepository {
+  create(data: CreateForkRelationshipData): ForkRelationship;
+  findByForkedTaskId(forkedTaskId: string): ForkRelationship | null;
+}
+
+@injectable()
+export class ForkRelationshipRepository implements IForkRelationshipRepository {
+  constructor(
+    @inject(MAIN_TOKENS.DatabaseService)
+    private readonly databaseService: DatabaseService,
+  ) {}
+
+  private get db() {
+    return this.databaseService.db;
+  }
+
+  create(data: CreateForkRelationshipData): ForkRelationship {
+    return this.db
+      .insert(forkRelationships)
+      .values({
+        id: crypto.randomUUID(),
+        ...data,
+      })
+      .returning()
+      .get();
+  }
+
+  findByForkedTaskId(forkedTaskId: string): ForkRelationship | null {
+    return (
+      this.db
+        .select()
+        .from(forkRelationships)
+        .where(eq(forkRelationships.forkedTaskId, forkedTaskId))
+        .get() ?? null
+    );
+  }
+}

apps/code/src/main/db/schema.ts

@@ -88,6 +88,18 @@ export const authSessions = sqliteTable("auth_sessions", {
   updatedAt: updatedAt(),
 });
 
+export const forkRelationships = sqliteTable("fork_relationships", {
+  id: id(),
+  forkedTaskId: text().notNull().unique(),
+  sourceTaskId: text().notNull(),
+  sourceTaskRunId: text().notNull(),
+  /** Title of the source task captured at fork time, shown if parent is later deleted. */
+  sourceTaskTitle: text().notNull(),
+  forkAtMessageIndex: integer().notNull(),
+  forkedAt: text().notNull(),
+  createdAt: createdAt(),
+});
+
 export const authPreferences = sqliteTable(
   "auth_preferences",
   {

apps/code/src/main/di/container.ts

@@ -4,6 +4,7 @@ import { Container } from "inversify";
 import { ArchiveRepository } from "../db/repositories/archive-repository";
 import { AuthPreferenceRepository } from "../db/repositories/auth-preference-repository";
 import { AuthSessionRepository } from "../db/repositories/auth-session-repository";
+import { ForkRelationshipRepository } from "../db/repositories/fork-relationship-repository";
 import { RepositoryRepository } from "../db/repositories/repository-repository";
 import { SuspensionRepositoryImpl } from "../db/repositories/suspension-repository";
 import { WorkspaceRepository } from "../db/repositories/workspace-repository";
@@ -41,6 +42,7 @@ import { FileWatcherService } from "../services/file-watcher/service";
 import { FocusService } from "../services/focus/service";
 import { FocusSyncService } from "../services/focus/sync-service";
 import { FoldersService } from "../services/folders/service";
+import { ForkService } from "../services/fork/service";
 import { FsService } from "../services/fs/service";
 import { GitService } from "../services/git/service";
 import { GitHubIntegrationService } from "../services/github-integration/service";
@@ -95,6 +97,9 @@ container.bind(MAIN_TOKENS.AuthSessionRepository).to(AuthSessionRepository);
 container.bind(MAIN_TOKENS.RepositoryRepository).to(RepositoryRepository);
 container.bind(MAIN_TOKENS.WorkspaceRepository).to(WorkspaceRepository);
 container.bind(MAIN_TOKENS.WorktreeRepository).to(WorktreeRepository);
+container
+  .bind(MAIN_TOKENS.ForkRelationshipRepository)
+  .to(ForkRelationshipRepository);
 container.bind(MAIN_TOKENS.ArchiveRepository).to(ArchiveRepository);
 container.bind(MAIN_TOKENS.SuspensionRepository).to(SuspensionRepositoryImpl);
 container.bind(MAIN_TOKENS.AgentAuthAdapter).to(AgentAuthAdapter);
@@ -142,5 +147,6 @@ container.bind(MAIN_TOKENS.TaskLinkService).to(TaskLinkService);
 container.bind(MAIN_TOKENS.InboxLinkService).to(InboxLinkService);
 container.bind(MAIN_TOKENS.WatcherRegistryService).to(WatcherRegistryService);
 container.bind(MAIN_TOKENS.WorkspaceService).to(WorkspaceService);
+container.bind(MAIN_TOKENS.ForkService).to(ForkService);
 
 container.bind(MAIN_TOKENS.SettingsStore).toConstantValue(settingsStore);

apps/code/src/main/di/tokens.ts

@@ -32,6 +32,7 @@ export const MAIN_TOKENS = Object.freeze({
   RepositoryRepository: Symbol.for("Main.RepositoryRepository"),
   WorkspaceRepository: Symbol.for("Main.WorkspaceRepository"),
   WorktreeRepository: Symbol.for("Main.WorktreeRepository"),
+  ForkRelationshipRepository: Symbol.for("Main.ForkRelationshipRepository"),
   ArchiveRepository: Symbol.for("Main.ArchiveRepository"),
   SuspensionRepository: Symbol.for("Main.SuspensionRepository"),
 
@@ -76,5 +77,6 @@ export const MAIN_TOKENS = Object.freeze({
   EnvironmentService: Symbol.for("Main.EnvironmentService"),
   ProvisioningService: Symbol.for("Main.ProvisioningService"),
   WorkspaceService: Symbol.for("Main.WorkspaceService"),
+  ForkService: Symbol.for("Main.ForkService"),
   EnrichmentService: Symbol.for("Main.EnrichmentService"),
 });

apps/code/src/main/services/agent/service.ts

@@ -37,6 +37,7 @@ import {
 import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
 import { extractCreatedPrUrl } from "@posthog/agent/pr-url-detector";
 import type * as AgentTypes from "@posthog/agent/types";
+import { createGitClient } from "@posthog/git/client";
 import { getCurrentBranch } from "@posthog/git/queries";
 import type { IAppMeta } from "@posthog/platform/app-meta";
 import type { IBundledResources } from "@posthog/platform/bundled-resources";
@@ -890,6 +891,44 @@ When creating pull requests, add the following footer at the end of the PR descr
       if (!this.hasActiveSessions()) {
         this.emit(AgentServiceEvent.SessionsIdle, undefined);
       }
+
+      void this.captureAgentCheckpoint(
+        sessionId,
+        session.taskId,
+        session.repoPath,
+        session.agent,
+      );
+    }
+  }
+
+  private async captureAgentCheckpoint(
+    taskRunId: string,
+    taskId: string,
+    repoPath: string,
+    agent: Agent,
+  ): Promise<void> {
+    if (taskId === "__preview__") return;
+
+    try {
+      const git = createGitClient(repoPath);
+      const headSha = await git.revparse(["HEAD"]);
+
+      const posthogAPI = agent.getPosthogAPI();
+      if (!posthogAPI) return;
+
+      await posthogAPI.appendTaskRunLog(taskId, taskRunId, [
+        {
+          type: "notification",
+          timestamp: new Date().toISOString(),
+          notification: {
+            jsonrpc: "2.0",
+            method: "_posthog/agent_checkpoint",
+            params: { headSha },
+          },
+        },
+      ]);
+    } catch (err) {
+      log.debug("Failed to capture agent checkpoint", { taskRunId, err });
     }
   }