From 942ecddc100e157b16e79e2801a04cb4ff001822 Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 08:54:27 +0000 Subject: [PATCH 1/2] feat(cloud-agent-next): expose worktree ownership on getSession --- .../cloud-agent-next/cloud-agent-client.ts | 5 ++ services/cloud-agent-next/src/router.test.ts | 71 +++++++++++++++++++ .../src/router/handlers/session-management.ts | 55 +++++++++++++- .../cloud-agent-next/src/router/schemas.ts | 17 +++++ 4 files changed, 146 insertions(+), 2 deletions(-) diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 1d2167d2289..992a8cb1801 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/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; diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 205902de60c..a96d24c3f90 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/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): 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'; diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index 04f575c3524..02c48a4a42a 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/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 { + 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, diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index 63c5b60641f..941f1378597 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/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'), From a04fb3e9ecd8491870de62c823b36c692d8c402b Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:02:05 +0000 Subject: [PATCH 2/2] fix(web): preserve worktree ownership in cloud-agent-next getSession schemas --- .../routers/cloud-agent-next-schemas.test.ts | 45 +++++++++++++++++++ .../src/routers/cloud-agent-next-schemas.ts | 5 +++ 2 files changed, 50 insertions(+) diff --git a/apps/web/src/routers/cloud-agent-next-schemas.test.ts b/apps/web/src/routers/cloud-agent-next-schemas.test.ts index 87bc2d1f788..556cfebdfab 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/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}`; diff --git a/apps/web/src/routers/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index 25d1e665018..4bf539ed5f7 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/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(),