feat(cloud-agent-next): expose worktree ownership on getSession

apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts

@@ -348,6 +348,11 @@ export type GetSessionOutput = {
   /** Sandbox ID (hashed format like usr-abc123...) for correlating with Cloudflare logs */
   sandboxId?: string;
 
+  // Worktree ownership (present only for worktree sessions)
+  worktreeId?: string | null;
+  parentSessionId?: string | null;
+  cloudAgentSessionScopeId?: string | null;
+
   // Repository info (no tokens)
   githubRepo?: string;
   gitUrl?: string;

apps/web/src/routers/cloud-agent-next-schemas.test.ts

@@ -9,6 +9,7 @@ import {
   baseCreateWorktreeChatNextSchema,
   baseGetSandboxStatusNextOutputSchema,
   baseGetSandboxStatusNextSchema,
+  baseGetSessionNextOutputSchema,
   basePrepareSessionNextSchema,
   organizationPrepareSessionNextSchema,
   personalPrepareSessionNextSchema,
@@ -223,6 +224,50 @@ describe('baseGetSandboxStatusNextSchema', () => {
   });
 });
 
+describe('baseGetSessionNextOutputSchema', () => {
+  const workspaceSessionId = `workspace_${MESSAGE_UUID}`;
+  const worktreeId = `worktree_${MESSAGE_UUID}`;
+  const baseSession = {
+    sessionId: workspaceSessionId,
+    userId: 'user_123',
+    execution: null,
+    timestamp: 1,
+    version: 1,
+  };
+
+  it('preserves worktree ownership for a grouped worktree session', () => {
+    const response = {
+      ...baseSession,
+      worktreeId,
+      parentSessionId: KILO_SESSION_ID,
+      cloudAgentSessionScopeId: workspaceSessionId,
+    };
+
+    expect(baseGetSessionNextOutputSchema.parse(response)).toEqual(response);
+  });
+
+  it('preserves explicit null ownership fields when no ownership row exists', () => {
+    const response = baseGetSessionNextOutputSchema.parse({
+      ...baseSession,
+      worktreeId,
+      parentSessionId: null,
+      cloudAgentSessionScopeId: null,
+    });
+
+    expect(response.worktreeId).toBe(worktreeId);
+    expect(response.parentSessionId).toBeNull();
+    expect(response.cloudAgentSessionScopeId).toBeNull();
+  });
+
+  it('omits ownership fields for an ordinary session', () => {
+    const response = baseGetSessionNextOutputSchema.parse(baseSession);
+
+    expect(response).not.toHaveProperty('worktreeId');
+    expect(response).not.toHaveProperty('parentSessionId');
+    expect(response).not.toHaveProperty('cloudAgentSessionScopeId');
+  });
+});
+
 describe('baseWorktreeChangesNextSchema', () => {
   const cloudAgentSessionId = `workspace_${MESSAGE_UUID}`;
 

apps/web/src/routers/cloud-agent-next-schemas.ts

@@ -693,6 +693,11 @@ export const baseGetSessionNextOutputSchema = z.object({
   orgId: z.string().optional(),
   sandboxId: z.string().optional(),
 
+  // Worktree ownership (present only for worktree sessions)
+  worktreeId: z.string().nullable().optional(),
+  parentSessionId: z.string().nullable().optional(),
+  cloudAgentSessionScopeId: z.string().nullable().optional(),
+
   // Repository info (no tokens)
   githubRepo: z.string().optional(),
   gitUrl: z.string().optional(),

services/cloud-agent-next/src/router.test.ts

@@ -36,6 +36,10 @@ vi.mock('./session-access.js', () => ({
   requireCurrentSessionAccess: requireCurrentSessionAccessMock,
 }));
 
+const { getPgDbMock } = vi.hoisted(() => ({ getPgDbMock: vi.fn() }));
+
+vi.mock('./db/pg.js', () => ({ getPgDb: getPgDbMock }));
+
 vi.mock('./session/model-preflight.js', () => ({
   preflightExistingPromptModel: preflightExistingPromptModelMock,
   preflightPreparedInitialPromptModel: preflightPreparedInitialPromptModelMock,
@@ -121,6 +125,15 @@ function legacySessionMetadata(input: Record<string, unknown>): CloudAgentSessio
   return parseSessionMetadata(input);
 }
 
+function mockWorktreeOwnershipRow(
+  row: { parentSessionId: string | null; cloudAgentSessionScopeId: string | null } | null
+): void {
+  const limit = vi.fn().mockResolvedValue(row ? [row] : []);
+  const where = vi.fn(() => ({ limit }));
+  const from = vi.fn(() => ({ where }));
+  getPgDbMock.mockReturnValue({ select: vi.fn(() => ({ from })) });
+}
+
 // Note: Balance validation is now handled in the worker entry point (index.ts)
 // via pre-flight validation before the tRPC handler is called.
 // This returns proper HTTP status codes (401, 402) instead of SSE error events.
@@ -1186,6 +1199,64 @@ describe('router sessionId validation', () => {
         });
       });
 
+      describe('worktree ownership', () => {
+        const WORKTREE_ID = 'worktree_aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa';
+        const SCOPE_ID = 'workspace_bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb';
+        const PARENT_KILO_SESSION_ID = 'ses_abcdefghijklmnopqrstuvwxyz';
+
+        function worktreeSessionMetadata(sessionId: SessionId): CloudAgentSessionState {
+          return parseSessionMetadata({
+            metadataSchemaVersion: 2,
+            identity: { sessionId, userId: 'test-user-123' },
+            auth: { kiloSessionId: PARENT_KILO_SESSION_ID },
+            workspace: { worktreeId: WORKTREE_ID },
+            lifecycle: { version: 1, timestamp: 123456789 },
+          });
+        }
+
+        it('exposes worktree ownership from the PostgreSQL row for a worktree session', async () => {
+          const sessionId: SessionId = 'agent_77777777-7777-4777-8777-777777777771';
+          mockGetMetadata.mockResolvedValue(worktreeSessionMetadata(sessionId));
+          mockWorktreeOwnershipRow({
+            parentSessionId: PARENT_KILO_SESSION_ID,
+            cloudAgentSessionScopeId: SCOPE_ID,
+          });
+
+          const result = await caller.getSession({ cloudAgentSessionId: sessionId });
+
+          expect(result.worktreeId).toBe(WORKTREE_ID);
+          expect(result.parentSessionId).toBe(PARENT_KILO_SESSION_ID);
+          expect(result.cloudAgentSessionScopeId).toBe(SCOPE_ID);
+          expect(getPgDbMock).toHaveBeenCalledTimes(1);
+        });
+
+        it('returns null ownership fields when a worktree session has no ownership row', async () => {
+          const sessionId: SessionId = 'agent_77777777-7777-4777-8777-777777777772';
+          mockGetMetadata.mockResolvedValue(worktreeSessionMetadata(sessionId));
+          mockWorktreeOwnershipRow(null);
+
+          const result = await caller.getSession({ cloudAgentSessionId: sessionId });
+
+          expect(result.worktreeId).toBe(WORKTREE_ID);
+          expect(result.parentSessionId).toBeNull();
+          expect(result.cloudAgentSessionScopeId).toBeNull();
+        });
+
+        it('omits ownership fields and skips PostgreSQL for an ordinary session', async () => {
+          const sessionId: SessionId = 'agent_77777777-7777-4777-8777-777777777773';
+          mockGetMetadata.mockResolvedValue(
+            legacySessionMetadata({ version: 1, sessionId, userId: 'test-user-123', timestamp: 1 })
+          );
+
+          const result = await caller.getSession({ cloudAgentSessionId: sessionId });
+
+          expect(result).not.toHaveProperty('worktreeId');
+          expect(result).not.toHaveProperty('parentSessionId');
+          expect(result).not.toHaveProperty('cloudAgentSessionScopeId');
+          expect(getPgDbMock).not.toHaveBeenCalled();
+        });
+      });
+
       describe('not found', () => {
         it('should return NOT_FOUND for non-existent session', async () => {
           const sessionId: SessionId = 'agent_00000000-0000-0000-0000-000000000000';

services/cloud-agent-next/src/router/handlers/session-management.ts

@@ -1,5 +1,6 @@
 import { TRPCError } from '@trpc/server';
 import * as z from 'zod';
+import type { WorkerDb } from '@kilocode/db/client';
 import { AgentSandboxUnavailableError } from '../../agent-sandbox/protocol.js';
 import { createAgentSandbox } from '../../agent-sandbox/factory.js';
 import { logger, withLogTags } from '../../logger.js';
@@ -37,7 +38,7 @@ import type { CloudAgentSessionState } from '../../persistence/types.js';
 import type { MessageResultRPCResponse } from '../../session/message-result.js';
 import { requireCurrentSessionAccess } from '../../session-access.js';
 import { getPgDb } from '../../db/pg.js';
-import { cloud_billing_sku, container_usage_interval } from '@kilocode/db/schema';
+import { cloud_billing_sku, cli_sessions_v2, container_usage_interval } from '@kilocode/db/schema';
 import { and, desc, eq, like } from 'drizzle-orm';
 import { SANDBOX_USAGE_SKUS } from '../../container-usage-context.js';
 
@@ -64,6 +65,34 @@ function toIso(value: string): string {
   return new Date(value).toISOString();
 }
 
+type WorktreeOwnership = {
+  parentSessionId: string | null;
+  cloudAgentSessionScopeId: string | null;
+};
+
+async function findWorktreeOwnership(
+  db: WorkerDb,
+  userId: string,
+  kiloSessionId: string,
+  cloudAgentSessionId: string
+): Promise<WorktreeOwnership | null> {
+  const [row] = await db
+    .select({
+      parentSessionId: cli_sessions_v2.parent_session_id,
+      cloudAgentSessionScopeId: cli_sessions_v2.cloud_agent_session_scope_id,
+    })
+    .from(cli_sessions_v2)
+    .where(
+      and(
+        eq(cli_sessions_v2.kilo_user_id, userId),
+        eq(cli_sessions_v2.session_id, kiloSessionId),
+        eq(cli_sessions_v2.cloud_agent_session_id, cloudAgentSessionId)
+      )
+    )
+    .limit(1);
+  return row ?? null;
+}
+
 function microdollarsForSeconds(seconds: number, rateCentsPerSecond: string): number {
   const [whole, fraction = ''] = rateCentsPerSecond.split('.');
   const scale = 10n ** BigInt(fraction.length);
@@ -326,7 +355,7 @@ export function createSessionManagementHandlers() {
 
           logger.setTags({ userId, sessionId });
           logger.info('Fetching session metadata');
-          await requireCurrentSessionAccess({
+          const sessionAccess = await requireCurrentSessionAccess({
             env,
             kiloUserId: userId,
             cloudAgentSessionId: sessionId,
@@ -395,6 +424,20 @@ export function createSessionManagementHandlers() {
           logger.setTags({ sandboxId, orgId: sessionMetadata.identity.orgId ?? '(personal)' });
           logger.info('Session metadata retrieved successfully');
 
+          // Worktree ownership lives on the external ownership row, not in DO
+          // metadata. Only worktree sessions need that read, so ordinary
+          // sessions never touch PostgreSQL here.
+          const worktreeId = sessionMetadata.workspace?.worktreeId;
+          let worktreeOwnership: WorktreeOwnership | null = null;
+          if (worktreeId) {
+            worktreeOwnership = await findWorktreeOwnership(
+              getPgDb(env),
+              userId,
+              sessionAccess.kiloSessionId,
+              sessionId
+            );
+          }
+
           // Sanitize and return safe fields only (no tokens/secrets)
           const repositoryFields = publicRepositoryFields(sessionMetadata);
           return {
@@ -404,6 +447,14 @@ export function createSessionManagementHandlers() {
             orgId: sessionMetadata.identity.orgId,
             sandboxId,
 
+            ...(worktreeId
+              ? {
+                  worktreeId,
+                  parentSessionId: worktreeOwnership?.parentSessionId ?? null,
+                  cloudAgentSessionScopeId: worktreeOwnership?.cloudAgentSessionScopeId ?? null,
+                }
+              : {}),
+
             githubRepo: repositoryFields.githubRepo,
             gitUrl: repositoryFields.gitUrl,
             platform: repositoryFields.platform,

services/cloud-agent-next/src/router/schemas.ts

@@ -1124,6 +1124,23 @@ export const GetSessionOutput = z.object({
     .optional()
     .describe('Sandbox ID (hashed format like usr-abc123...) for correlating with Cloudflare logs'),
 
+  // Worktree ownership (present only for worktree sessions)
+  worktreeId: z
+    .string()
+    .nullable()
+    .optional()
+    .describe('Worktree ID when this session belongs to a shared worktree'),
+  parentSessionId: z
+    .string()
+    .nullable()
+    .optional()
+    .describe('Parent Kilo session ID for a grouped worktree session (null when no group)'),
+  cloudAgentSessionScopeId: z
+    .string()
+    .nullable()
+    .optional()
+    .describe('Cloud-agent session scope ID for a grouped worktree session (null when no group)'),
+
   // Repository info (no tokens)
   githubRepo: z.string().optional().describe('GitHub repository in org/repo format'),
   gitUrl: z.string().optional().describe('Generic git URL'),