CloudAgentNext - Move cli_sessions_v2 creation to session-ingest RPC

cloud-agent-next/src/execution/orchestrator.ts

@@ -322,7 +322,6 @@ export class ExecutionOrchestrator {
           setupCommands: initContext.setupCommands,
           mcpServers: initContext.mcpServers,
           botId: initContext.botId,
-          skipLinking: true,
           githubAppType: initContext.githubAppType,
           // Note: existingMetadata requires CloudAgentSessionState, not our simplified type
           ...gitSource,

cloud-agent-next/src/persistence/CloudAgentSession.ts

@@ -225,7 +225,6 @@ export class CloudAgentSession extends DurableObject {
       // Create DO context for the ingest handler to call back into the DO
       const doContext: IngestDOContext = {
         updateKiloSessionId: (id: string) => this.updateKiloSessionId(id),
-        linkKiloSessionInBackend: (id: string) => this.linkKiloSessionInBackend(id),
         updateUpstreamBranch: (branch: string) => this.updateUpstreamBranch(branch),
         clearActiveExecution: () => this.clearActiveExecution(),
         getExecution: async (executionId: string) => {
@@ -574,41 +573,6 @@ export class CloudAgentSession extends DurableObject {
     await this.updateMetadata(updated);
   }
 
-  /**
-   * Link the kiloSessionId to the backend for analytics/tracking.
-   * Called when a session_created event is received from the CLI.
-   *
-   * @param kiloSessionId - The kilo CLI session ID to link
-   */
-  async linkKiloSessionInBackend(kiloSessionId: string): Promise<void> {
-    const metadata = await this.getMetadata();
-    if (!metadata?.kilocodeToken) {
-      throw new Error('Cannot link session: missing kilocodeToken');
-    }
-
-    const backendUrl = (this.env as unknown as WorkerEnv).KILOCODE_BACKEND_BASE_URL;
-    if (!backendUrl) {
-      throw new Error('Cannot link session: KILOCODE_BACKEND_BASE_URL not configured');
-    }
-
-    const response = await fetch(`${backendUrl}/api/cloud-sessions/linkSessions`, {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: `Bearer ${metadata.kilocodeToken}`,
-      },
-      body: JSON.stringify({
-        cloudSessionId: this.sessionId,
-        kiloSessionId: kiloSessionId,
-      }),
-    });
-
-    if (!response.ok) {
-      const text = await response.text();
-      throw new Error(`Backend link failed: ${response.status} ${text}`);
-    }
-  }
-
   // ---------------------------------------------------------------------------
   // Wrapper Communication Methods
   // ---------------------------------------------------------------------------

cloud-agent-next/src/persistence/types.ts

@@ -4,6 +4,7 @@ import type { CloudAgentSession } from './CloudAgentSession.js';
 import type { EncryptedSecrets } from '../router/schemas.js';
 import type { CallbackTarget } from '../callbacks/index.js';
 import type { Images } from './schemas.js';
+import type { SessionIngestBinding } from '../session-ingest-binding.js';
 
 /**
  * Base configuration shared by all MCP server types
@@ -170,7 +171,7 @@ export type PersistenceEnv = {
   /** Durable Object namespace for CloudAgentSession metadata (SQLite-backed) with RPC support */
   CLOUD_AGENT_SESSION: DurableObjectNamespace<CloudAgentSession>;
   /** Service binding for the session ingest worker */
-  SESSION_INGEST: Fetcher;
+  SESSION_INGEST: SessionIngestBinding;
   /** Shared secret for JWT token validation */
   NEXTAUTH_SECRET: string;
   /** Comma-separated list of allowed Origins for /stream WebSocket connections */

cloud-agent-next/src/router/handlers/session-prepare.ts

@@ -6,6 +6,7 @@ import {
   SessionService,
   determineBranchName,
   runSetupCommands,
+  writeAuthFile,
   writeMCPSettings,
 } from '../../session-service.js';
 import { InstallationLookupService } from '../../services/installation-lookup-service.js';
@@ -190,35 +191,6 @@ const prepareSessionHandler = internalApiProtectedProcedure
         }
       }
 
-      // NOTE: Backend session creation (createKiloSessionInBackend) is temporarily disabled.
-      // The kiloSessionId will now come from the kilo CLI server's POST /session API.
-      // This can be re-enabled later if backend analytics/tracking is needed.
-      // const gitUrlForBackend = input.githubRepo
-      //   ? `https://github.com/${input.githubRepo}`
-      //   : input.gitUrl;
-      // let backendKiloSessionId: string;
-      // try {
-      //   backendKiloSessionId = await sessionService.createKiloSessionInBackend(
-      //     cloudAgentSessionId,
-      //     ctx.authToken,
-      //     ctx.env,
-      //     input.kilocodeOrganizationId,
-      //     input.mode,
-      //     input.model,
-      //     gitUrlForBackend
-      //   );
-      // } catch (error) {
-      //   logger
-      //     .withFields({ error: error instanceof Error ? error.message : String(error) })
-      //     .error('Failed to create cliSession in backend');
-      //   throw new TRPCError({
-      //     code: 'INTERNAL_SERVER_ERROR',
-      //     message: `Failed to create session in backend: ${
-      //       error instanceof Error ? error.message : String(error)
-      //     }`,
-      //   });
-      // }
-
       // 3. Get sandbox
       logger.info('Getting sandbox');
       const sandbox = getSandbox(ctx.env.Sandbox, sandboxId, { sleepAfter: 900 });
@@ -310,6 +282,9 @@ const prepareSessionHandler = internalApiProtectedProcedure
         await writeMCPSettings(sandbox, sessionHome, input.mcpServers);
       }
 
+      // 9b. Write auth file for session ingest
+      await writeAuthFile(sandbox, sessionHome, ctx.authToken);
+
       // 10. Start kilo server
       logger.info('Starting kilo server');
       const kiloServerPort = await ensureKiloServer(
@@ -327,52 +302,98 @@ const prepareSessionHandler = internalApiProtectedProcedure
       logger.setTags({ kiloSessionId });
       logger.info('Created kilo CLI session');
 
-      // 12. Get DO stub and store metadata
+      // 12. Create cli_sessions_v2 record via session-ingest RPC (blocking)
+      logger.info('Creating cli_sessions_v2 record via session-ingest');
+      try {
+        await sessionService.createCliSessionViaSessionIngest(
+          kiloSessionId,
+          cloudAgentSessionId,
+          ctx.userId,
+          ctx.env,
+          input.kilocodeOrganizationId,
+          'cloud-agent'
+        );
+      } catch (error) {
+        logger
+          .withFields({ error: error instanceof Error ? error.message : String(error) })
+          .error('Failed to create cli_sessions_v2 record');
+        throw new TRPCError({
+          code: 'INTERNAL_SERVER_ERROR',
+          message: `Failed to create session record: ${
+            error instanceof Error ? error.message : String(error)
+          }`,
+        });
+      }
+
+      const rollbackCliSession = async () => {
+        await sessionService
+          .deleteCliSessionViaSessionIngest(kiloSessionId, ctx.userId, ctx.env)
+          .catch((rollbackError: unknown) => {
+            logger
+              .withFields({
+                error:
+                  rollbackError instanceof Error ? rollbackError.message : String(rollbackError),
+              })
+              .error('Failed to rollback cli_sessions_v2 record');
+          });
+      };
+
+      // 13. Get DO stub and store metadata
       const doId = ctx.env.CLOUD_AGENT_SESSION.idFromName(`${ctx.userId}:${cloudAgentSessionId}`);
       const stub = ctx.env.CLOUD_AGENT_SESSION.get(doId);
 
-      const prepareResult = await stub.prepare({
-        sessionId: cloudAgentSessionId,
-        userId: ctx.userId,
-        orgId: input.kilocodeOrganizationId,
-        botId: ctx.botId,
-        kiloSessionId,
-        prompt: input.prompt,
-        mode: input.mode,
-        model: input.model,
-        kilocodeToken: ctx.authToken,
-        githubRepo: input.githubRepo,
-        githubToken: input.githubToken,
-        githubInstallationId: resolvedInstallationId,
-        githubAppType: resolvedGithubAppType,
-        gitUrl: input.gitUrl,
-        gitToken: input.gitToken,
-        envVars: input.envVars,
-        encryptedSecrets: input.encryptedSecrets,
-        setupCommands: input.setupCommands,
-        mcpServers: input.mcpServers,
-        upstreamBranch: input.upstreamBranch,
-        autoCommit: input.autoCommit,
-        condenseOnComplete: input.condenseOnComplete,
-        appendSystemPrompt: input.appendSystemPrompt,
-        callbackTarget: input.callbackTarget,
-        images: input.images,
-        // Workspace metadata
-        workspacePath,
-        sessionHome,
-        branchName,
-        sandboxId,
-      });
+      let prepareResult;
+      try {
+        prepareResult = await stub.prepare({
+          sessionId: cloudAgentSessionId,
+          userId: ctx.userId,
+          orgId: input.kilocodeOrganizationId,
+          botId: ctx.botId,
+          kiloSessionId,
+          prompt: input.prompt,
+          mode: input.mode,
+          model: input.model,
+          kilocodeToken: ctx.authToken,
+          githubRepo: input.githubRepo,
+          githubToken: input.githubToken,
+          githubInstallationId: resolvedInstallationId,
+          githubAppType: resolvedGithubAppType,
+          gitUrl: input.gitUrl,
+          gitToken: input.gitToken,
+          envVars: input.envVars,
+          encryptedSecrets: input.encryptedSecrets,
+          setupCommands: input.setupCommands,
+          mcpServers: input.mcpServers,
+          upstreamBranch: input.upstreamBranch,
+          autoCommit: input.autoCommit,
+          condenseOnComplete: input.condenseOnComplete,
+          appendSystemPrompt: input.appendSystemPrompt,
+          callbackTarget: input.callbackTarget,
+          images: input.images,
+          // Workspace metadata
+          workspacePath,
+          sessionHome,
+          branchName,
+          sandboxId,
+        });
+      } catch (error) {
+        logger
+          .withFields({ error: error instanceof Error ? error.message : String(error) })
+          .error('DO prepare() threw, rolling back cli_sessions_v2 record');
+        await rollbackCliSession();
+        throw error;
+      }
 
       if (!prepareResult.success) {
         logger.withFields({ error: prepareResult.error }).error('Failed to prepare session in DO');
+        await rollbackCliSession();
         throw new TRPCError({
           code: 'BAD_REQUEST',
           message: prepareResult.error ?? 'Failed to prepare session',
         });
       }
 
-      // 13. Record kilo server activity for idle timeout tracking
+      // 14. Record kilo server activity for idle timeout tracking
       try {
         await withDORetry(
           () => ctx.env.CLOUD_AGENT_SESSION.get(doId),
@@ -388,7 +409,7 @@ const prepareSessionHandler = internalApiProtectedProcedure
 
       logger.info('Session prepared successfully');
 
-      // 14. Return both IDs
+      // 15. Return both IDs
       return { cloudAgentSessionId, kiloSessionId };
     });
   });

cloud-agent-next/src/session-ingest-binding.ts

@@ -0,0 +1,33 @@
+/**
+ * RPC method types for the SESSION_INGEST service binding.
+ *
+ * `wrangler types` only sees `Fetcher` for service bindings; the actual RPC
+ * shape comes from the session-ingest worker's WorkerEntrypoint and is
+ * declared here so the generated file can be freely regenerated.
+ *
+ * Keep in sync with: cloudflare-session-ingest/src/session-ingest-rpc.ts
+ */
+
+export type CreateSessionForCloudAgentParams = {
+  sessionId: string;
+  kiloUserId: string;
+  cloudAgentSessionId: string;
+  organizationId?: string;
+  createdOnPlatform: string;
+};
+
+export type DeleteSessionForCloudAgentParams = {
+  sessionId: string;
+  kiloUserId: string;
+};
+
+export type ExportSessionParams = {
+  sessionId: string;
+  kiloUserId: string;
+};
+
+export type SessionIngestBinding = Fetcher & {
+  createSessionForCloudAgent(params: CreateSessionForCloudAgentParams): Promise<void>;
+  deleteSessionForCloudAgent(params: DeleteSessionForCloudAgentParams): Promise<void>;
+  exportSession(params: ExportSessionParams): Promise<string | null>;
+};

cloud-agent-next/src/session-prepare.test.ts

@@ -53,14 +53,15 @@ vi.mock('./kilo/server-manager.js', () => ({
 
 // Define mocks BEFORE vi.mock() to avoid hoisting issues
 // vi.hoisted() ensures these are available when the mock factory runs
-const { generateSessionIdMock, createKiloSessionInBackendMock, deleteKiloSessionInBackendMock } =
-  vi.hoisted(() => ({
-    generateSessionIdMock: vi.fn(() => 'agent_12345678-1234-1234-1234-123456789abc'),
-    createKiloSessionInBackendMock: vi
-      .fn()
-      .mockResolvedValue('123e4567-e89b-12d3-a456-426614174000'),
-    deleteKiloSessionInBackendMock: vi.fn().mockResolvedValue(undefined),
-  }));
+const {
+  generateSessionIdMock,
+  createCliSessionViaSessionIngestMock,
+  deleteCliSessionViaSessionIngestMock,
+} = vi.hoisted(() => ({
+  generateSessionIdMock: vi.fn(() => 'agent_12345678-1234-1234-1234-123456789abc'),
+  createCliSessionViaSessionIngestMock: vi.fn().mockResolvedValue(undefined),
+  deleteCliSessionViaSessionIngestMock: vi.fn().mockResolvedValue(undefined),
+}));
 
 // Mock session-service to isolate router tests
 vi.mock('./session-service.js', () => ({
@@ -70,6 +71,7 @@ vi.mock('./session-service.js', () => ({
     (sessionId: string, upstreamBranch?: string) => upstreamBranch || `session/${sessionId}`
   ),
   runSetupCommands: vi.fn().mockResolvedValue(undefined),
+  writeAuthFile: vi.fn().mockResolvedValue(undefined),
   writeMCPSettings: vi.fn().mockResolvedValue(undefined),
   InvalidSessionMetadataError: class InvalidSessionMetadataError extends Error {
     constructor(
@@ -82,8 +84,8 @@ vi.mock('./session-service.js', () => ({
     }
   },
   SessionService: class SessionService {
-    createKiloSessionInBackend = createKiloSessionInBackendMock;
-    deleteKiloSessionInBackend = deleteKiloSessionInBackendMock;
+    createCliSessionViaSessionIngest = createCliSessionViaSessionIngestMock;
+    deleteCliSessionViaSessionIngest = deleteCliSessionViaSessionIngestMock;
     getOrCreateSession = vi.fn().mockResolvedValue(createMockExecutionSession());
     buildContext = vi.fn().mockReturnValue({
       sandboxId: 'test-sandbox',
@@ -174,6 +176,11 @@ function createInternalApiContext(options: {
         idFromName: vi.fn((id: string) => ({ id })),
         get: vi.fn(() => doStub),
       } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'],
+      SESSION_INGEST: {
+        fetch: vi.fn(),
+        createSessionForCloudAgent: vi.fn().mockResolvedValue(undefined),
+        deleteSessionForCloudAgent: vi.fn().mockResolvedValue(undefined),
+      } as unknown as TRPCContext['env']['SESSION_INGEST'],
       INTERNAL_API_SECRET: effectiveInternalApiSecret,
       NEXTAUTH_SECRET: 'test-secret',
     },
@@ -184,8 +191,8 @@ describe('prepareSession endpoint', () => {
   beforeEach(() => {
     vi.clearAllMocks();
     generateSessionIdMock.mockReturnValue('agent_12345678-1234-1234-1234-123456789abc');
-    createKiloSessionInBackendMock.mockResolvedValue('123e4567-e89b-12d3-a456-426614174000');
-    deleteKiloSessionInBackendMock.mockResolvedValue(undefined);
+    createCliSessionViaSessionIngestMock.mockResolvedValue(undefined);
+    deleteCliSessionViaSessionIngestMock.mockResolvedValue(undefined);
   });
 
   describe('authentication', () => {
@@ -426,7 +433,7 @@ describe('prepareSession endpoint', () => {
       ).rejects.toThrow('Session already prepared');
     });
 
-    // NOTE: Backend session creation (createKiloSessionInBackend) is currently disabled.
+    // NOTE: CLI session creation (createCliSessionViaSessionIngest) is handled via session-ingest.
     // The kiloSessionId now comes from the kilo CLI server's POST /session API.
     // Tests for backend session creation error handling and rollback have been removed.
   });