claude-code/packages/remote-control-server/src/services/work-dispatch.ts

import { log, error as logError } from "../logger";
import {
  storeCreateWorkItem,
  storeGetWorkItem,
  storeGetPendingWorkItem,
  storeUpdateWorkItem,
  storeListSessionsByEnvironment,
  storeGetEnvironment,
} from "../store";
import { config } from "../config";
import { getBaseUrl } from "../config";
import type { WorkResponse } from "../types/api";

/** Encode work secret as base64 JSON (no JWT — just API key as token) */
function encodeWorkSecret(): string {
  const payload = {
    version: 1,
    session_ingress_token: config.apiKeys[0] || "",
    api_base_url: getBaseUrl(),
    sources: [] as string[],
    auth: [] as string[],
    use_code_sessions: false,
  };
  return Buffer.from(JSON.stringify(payload)).toString("base64url");
}

export async function createWorkItem(environmentId: string, sessionId: string): Promise<string> {
  // Validate environment exists and is active
  const env = storeGetEnvironment(environmentId);
  if (!env) {
    throw new Error(`Environment ${environmentId} not found`);
  }
  if (env.status !== "active") {
    throw new Error(`Environment ${environmentId} is not active (status: ${env.status})`);
  }

  const secret = encodeWorkSecret();
  const record = storeCreateWorkItem({ environmentId, sessionId, secret });
  log(`[RCS] Work item created: ${record.id} for env=${environmentId} session=${sessionId}`);
  return record.id;
}

/** Long-poll for work — blocks until work is available or timeout.
 *  Returns null when no work is available, matching the CLI bridge client protocol. */
export async function pollWork(environmentId: string, timeoutSeconds = config.pollTimeout): Promise<WorkResponse | null> {
  const deadline = Date.now() + timeoutSeconds * 1000;

  while (Date.now() < deadline) {
    const item = storeGetPendingWorkItem(environmentId);

    if (item) {
      storeUpdateWorkItem(item.id, { state: "dispatched" });

      return {
        id: item.id,
        type: "work",
        environment_id: environmentId,
        state: "dispatched",
        data: {
          type: "session",
          id: item.sessionId,
        },
        secret: item.secret,
        created_at: item.createdAt.toISOString(),
      };
    }

    await new Promise((r) => setTimeout(r, 500));
  }

  return null;
}

export function ackWork(workId: string) {
  storeUpdateWorkItem(workId, { state: "acked" });
}

export function stopWork(workId: string) {
  storeUpdateWorkItem(workId, { state: "completed" });
}

export function heartbeatWork(workId: string): { lease_extended: boolean; state: string; last_heartbeat: string; ttl_seconds: number } {
  storeUpdateWorkItem(workId, {} as any); // just bump updatedAt
  const item = storeGetWorkItem(workId);
  const now = new Date();
  return {
    lease_extended: true,
    state: item?.state ?? "acked",
    last_heartbeat: now.toISOString(),
    ttl_seconds: config.heartbeatInterval * 2,
  };
}

/** Reconnect: re-queue sessions associated with an environment */
export function reconnectWorkForEnvironment(envId: string) {
  const activeSessions = storeListSessionsByEnvironment(envId).filter((s) => s.status === "idle");
  const promises = activeSessions.map((s) => createWorkItem(envId, s.id));
  return Promise.all(promises);
}