onebox/ts/opsserver/handlers/workspace.handler.ts

import * as plugins from '../../plugins.ts';
import { logger } from '../../logging.ts';
import type { OpsServer } from '../classes.opsserver.ts';
import * as interfaces from '../../../ts_interfaces/index.ts';
import { requireAdminIdentity } from '../helpers/guards.ts';
import { getErrorMessage } from '../../utils/error.ts';

interface IWorkspaceProcessSession {
  processId: string;
  serviceName: string;
  userId: string;
  stream: plugins.nodeStream.Duplex;
  close: () => Promise<void>;
  inspect: () => Promise<{ ExitCode?: number | null; Running?: boolean }>;
  finalized: boolean;
}

const getWorkspaceProcessTag = (processIdArg: string) => `workspaceProcess:${processIdArg}`;

export class WorkspaceHandler {
  public typedrouter = new plugins.typedrequest.TypedRouter();
  private workspaceProcesses = new Map<string, IWorkspaceProcessSession>();

  constructor(private opsServerRef: OpsServer) {
    this.opsServerRef.typedrouter.addTypedRouter(this.typedrouter);
    this.registerHandlers();
  }

  /**
   * Resolve a service name to a container ID (handling Swarm service IDs)
   */
  private async resolveContainerId(serviceName: string): Promise<string> {
    const service = this.opsServerRef.oneboxRef.services.getService(serviceName);
    if (!service || !service.containerID) {
      throw new plugins.typedrequest.TypedResponseError(`Service not found or has no container: ${serviceName}`);
    }
    return service.containerID;
  }

  private validateProcessId(processIdArg: string): void {
    if (!/^[a-zA-Z0-9_-]{8,80}$/.test(processIdArg)) {
      throw new plugins.typedrequest.TypedResponseError('Invalid workspace process id');
    }
  }

  private async getShellCommandForContainer(
    containerIdArg: string,
  ): Promise<interfaces.requests.IWorkspaceShellCommand> {
    const candidates: interfaces.requests.IWorkspaceShellCommand[] = [
      { command: '/bin/bash', args: ['-il'], label: 'bash', prompt: '# ' },
      { command: 'bash', args: ['-il'], label: 'bash', prompt: '# ' },
      { command: '/bin/sh', args: ['-i'], label: 'sh', prompt: '# ' },
      { command: 'sh', args: ['-i'], label: 'sh', prompt: '# ' },
      { command: '/bin/ash', args: ['-i'], label: 'ash', prompt: '# ' },
      { command: 'ash', args: ['-i'], label: 'ash', prompt: '# ' },
      { command: '/usr/bin/zsh', args: ['-il'], label: 'zsh', prompt: '# ' },
      { command: 'zsh', args: ['-il'], label: 'zsh', prompt: '# ' },
    ];

    for (const candidate of candidates) {
      const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
        containerIdArg,
        [candidate.command, '-c', 'printf onebox-shell'],
      );
      if (result.exitCode === 0 && result.stdout.includes('onebox-shell')) {
        return candidate;
      }
    }

    throw new plugins.typedrequest.TypedResponseError(
      'No supported interactive shell found in the target container',
    );
  }

  private async getProcessSession(
    dataArg: { identity: interfaces.data.IIdentity; processId: string },
  ): Promise<IWorkspaceProcessSession> {
    const identity = await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
    this.validateProcessId(dataArg.processId);
    const session = this.workspaceProcesses.get(dataArg.processId);
    if (!session) {
      throw new plugins.typedrequest.TypedResponseError(`Workspace process not found: ${dataArg.processId}`);
    }
    if (session.userId !== identity.userId) {
      throw new plugins.typedrequest.TypedResponseError('Workspace process belongs to another session');
    }
    return session;
  }

  private async pushWorkspaceProcessOutput(processIdArg: string, outputArg: string): Promise<void> {
    const typedsocket = (this.opsServerRef.server as any)?.typedserver?.typedsocket;
    if (!typedsocket) return;

    const connections = await typedsocket.findAllTargetConnectionsByTag(getWorkspaceProcessTag(processIdArg));
    await Promise.allSettled(
      connections.map((connection: any) => typedsocket
        .createTypedRequest(
          'pushWorkspaceProcessOutput',
          connection,
        )
        .fire({ processId: processIdArg, output: outputArg })),
    );
  }

  private async pushWorkspaceProcessExit(processIdArg: string, exitCodeArg: number): Promise<void> {
    const typedsocket = (this.opsServerRef.server as any)?.typedserver?.typedsocket;
    if (!typedsocket) return;

    const connections = await typedsocket.findAllTargetConnectionsByTag(getWorkspaceProcessTag(processIdArg));
    await Promise.allSettled(
      connections.map((connection: any) => typedsocket
        .createTypedRequest(
          'pushWorkspaceProcessExit',
          connection,
        )
        .fire({ processId: processIdArg, exitCode: exitCodeArg })),
    );
  }

  private async finalizeWorkspaceProcess(processIdArg: string, fallbackExitCodeArg = -1): Promise<void> {
    const session = this.workspaceProcesses.get(processIdArg);
    if (!session || session.finalized) return;
    session.finalized = true;

    let exitCode = fallbackExitCodeArg;
    try {
      await new Promise((resolve) => setTimeout(resolve, 50));
      const inspectResult = await session.inspect();
      if (typeof inspectResult.ExitCode === 'number') {
        exitCode = inspectResult.ExitCode;
      }
    } catch (error) {
      logger.debug(`Failed to inspect workspace process ${processIdArg}: ${getErrorMessage(error)}`);
    }

    this.workspaceProcesses.delete(processIdArg);
    await this.pushWorkspaceProcessExit(processIdArg, exitCode);
    try {
      await session.close();
    } catch {
      // The hijacked connection may already be closed by Docker.
    }
  }

  private registerHandlers(): void {
    // Read file from container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceReadFile>(
        'workspaceReadFile',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            ['cat', dataArg.path],
          );
          if (result.exitCode !== 0) {
            throw new plugins.typedrequest.TypedResponseError(`Failed to read file: ${result.stderr || 'File not found'}`);
          }
          return { content: result.stdout };
        },
      ),
    );

    // Write file to container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceWriteFile>(
        'workspaceWriteFile',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          // Use sh -c with printf to write content (handles special characters)
          const escaped = dataArg.content.replace(/'/g, "'\\''");
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            ['sh', '-c', `printf '%s' '${escaped}' > ${dataArg.path}`],
          );
          if (result.exitCode !== 0) {
            throw new plugins.typedrequest.TypedResponseError(`Failed to write file: ${result.stderr}`);
          }
          return {};
        },
      ),
    );

    // Read directory from container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceReadDir>(
        'workspaceReadDir',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          // Use ls with -1 -F to get entries with type indicators (/ for dirs)
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            ['ls', '-1', '-F', dataArg.path],
          );
          if (result.exitCode !== 0) {
            throw new plugins.typedrequest.TypedResponseError(`Failed to read directory: ${result.stderr}`);
          }
          const entries = result.stdout
            .split('\n')
            .filter((line) => line.trim())
            .map((line) => {
              const isDir = line.endsWith('/');
              const name = isDir ? line.slice(0, -1) : line.replace(/[*@=|]$/, '');
              const basePath = dataArg.path.endsWith('/') ? dataArg.path : dataArg.path + '/';
              return {
                type: (isDir ? 'directory' : 'file') as 'file' | 'directory',
                name,
                path: basePath + name,
              };
            });
          return { entries };
        },
      ),
    );

    // Create directory in container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceMkdir>(
        'workspaceMkdir',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            ['mkdir', '-p', dataArg.path],
          );
          if (result.exitCode !== 0) {
            throw new plugins.typedrequest.TypedResponseError(`Failed to create directory: ${result.stderr}`);
          }
          return {};
        },
      ),
    );

    // Remove file/directory from container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceRm>(
        'workspaceRm',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const args = dataArg.recursive ? ['rm', '-rf', dataArg.path] : ['rm', '-f', dataArg.path];
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            args,
          );
          if (result.exitCode !== 0) {
            throw new plugins.typedrequest.TypedResponseError(`Failed to remove: ${result.stderr}`);
          }
          return {};
        },
      ),
    );

    // Check if path exists in container
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceExists>(
        'workspaceExists',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            ['test', '-e', dataArg.path],
          );
          return { exists: result.exitCode === 0 };
        },
      ),
    );

    // Execute a command in the container (non-interactive)
    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceExec>(
        'workspaceExec',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const cmd = dataArg.args
            ? [dataArg.command, ...dataArg.args]
            : [dataArg.command];
          const result = await this.opsServerRef.oneboxRef.docker.execInContainer(
            containerId,
            cmd,
          );
          return {
            stdout: result.stdout,
            stderr: result.stderr,
            exitCode: result.exitCode,
          };
        },
      ),
    );

    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceGetShellCommand>(
        'workspaceGetShellCommand',
        async (dataArg) => {
          await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const shellCommand = await this.getShellCommandForContainer(containerId);
          return { shellCommand };
        },
      ),
    );

    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceStartProcess>(
        'workspaceStartProcess',
        async (dataArg) => {
          const identity = await requireAdminIdentity(this.opsServerRef.adminHandler, dataArg);
          this.validateProcessId(dataArg.processId);
          if (this.workspaceProcesses.has(dataArg.processId)) {
            throw new plugins.typedrequest.TypedResponseError(`Workspace process already exists: ${dataArg.processId}`);
          }

          const containerId = await this.resolveContainerId(dataArg.serviceName);
          const command = dataArg.args ? [dataArg.command, ...dataArg.args] : [dataArg.command];
          const interactiveExec = await this.opsServerRef.oneboxRef.docker.startInteractiveExecInContainer(
            containerId,
            command,
          );

          const session: IWorkspaceProcessSession = {
            processId: dataArg.processId,
            serviceName: dataArg.serviceName,
            userId: identity.userId,
            stream: interactiveExec.stream,
            close: interactiveExec.close,
            inspect: interactiveExec.inspect,
            finalized: false,
          };
          this.workspaceProcesses.set(dataArg.processId, session);

          interactiveExec.stream.on('data', (chunk: Uint8Array | string) => {
            const output = typeof chunk === 'string' ? chunk : new TextDecoder().decode(chunk);
            void this.pushWorkspaceProcessOutput(dataArg.processId, output);
          });
          interactiveExec.stream.on('error', (error: Error) => {
            void this.pushWorkspaceProcessOutput(
              dataArg.processId,
              `\r\n[workspace process error: ${getErrorMessage(error)}]\r\n`,
            );
            void this.finalizeWorkspaceProcess(dataArg.processId, -1);
          });
          interactiveExec.stream.on('end', () => {
            void this.finalizeWorkspaceProcess(dataArg.processId, -1);
          });
          interactiveExec.stream.on('close', () => {
            void this.finalizeWorkspaceProcess(dataArg.processId, -1);
          });

          return { processId: dataArg.processId };
        },
      ),
    );

    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceProcessInput>(
        'workspaceProcessInput',
        async (dataArg) => {
          const session = await this.getProcessSession(dataArg);
          if (session.finalized || session.stream.writableEnded) {
            return {};
          }
          await new Promise<void>((resolve, reject) => {
            session.stream.write(dataArg.input, (error?: Error | null) => {
              if (error) {
                reject(error);
              } else {
                resolve();
              }
            });
          });
          return {};
        },
      ),
    );

    this.typedrouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<interfaces.requests.IReq_WorkspaceKillProcess>(
        'workspaceKillProcess',
        async (dataArg) => {
          const session = await this.getProcessSession(dataArg);
          session.stream.destroy();
          try {
            await session.close();
          } catch {
            // The stream may already be closed.
          }
          await this.finalizeWorkspaceProcess(dataArg.processId, -1);
          return {};
        },
      ),
    );

    logger.info('Workspace handler registered');
  }
}