Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Comment thread
kilo-code-bot[bot] marked this conversation as resolved.
parentSessionId?: string | null;
cloudAgentSessionScopeId?: string | null;

// Repository info (no tokens)
githubRepo?: string;
gitUrl?: string;
Expand Down
71 changes: 71 additions & 0 deletions services/cloud-agent-next/src/router.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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';
Expand Down
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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';

Expand All @@ -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);
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
Expand Down
17 changes: 17 additions & 0 deletions services/cloud-agent-next/src/router/schemas.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'),
Expand Down
Loading