← vibesdk
import { Connection } from 'agents';
import { createLogger } from '../../logger';
import { WebSocketMessageRequests, WebSocketMessageResponses } from '../constants';
import { WebSocketMessage, WebSocketMessageData, WebSocketMessageType } from '../../api/websocketTypes';
import { MAX_IMAGES_PER_MESSAGE, MAX_IMAGE_SIZE_BYTES, type ImageAttachment } from '../../types/image-attachment';
import { type CredentialsPayload } from '../inferutils/config.types';
import { checkUsageAndBalance } from '../../services/rate-limit';
import type { CodeGeneratorAgent } from './codingAgent';
import type { DeploymentTarget } from './types';
// Type for incoming WebSocket messages
interface IncomingWebSocketMessage {
type: string;
message?: string;
images?: ImageAttachment[];
credentials?: CredentialsPayload;
commitHash?: string;
target?: DeploymentTarget;
data?: {
url?: string;
viewport?: unknown;
[key: string]: unknown;
};
}
const logger = createLogger('CodeGeneratorWebSocket');
export async function handleWebSocketMessage(
agent: CodeGeneratorAgent,
connection: Connection,
message: string
): Promise<void> {
try {
logger.info(`Received WebSocket message from ${connection.id}: ${message}`);
const parsedMessage = JSON.parse(message) as IncomingWebSocketMessage;
switch (parsedMessage.type) {
case WebSocketMessageRequests.SESSION_INIT: {
// const credentials = parsedMessage.credentials as CredentialsPayload | undefined;
// agent.getBehavior().setRuntimeOverrides(credentialsToRuntimeOverrides(credentials));
logger.info(`Received session init message from ${connection.id}, Disable for now`);
break;
}
case WebSocketMessageRequests.GENERATE_ALL:
// Set shouldBeGenerating flag to indicate persistent intent
agent.setState({
...agent.state,
shouldBeGenerating: true
});
// Check if generation is already active to avoid duplicate processes
if (agent.getBehavior().isCodeGenerating()) {
logger.info('Generation already in progress, skipping duplicate request');
// sendToConnection(connection, WebSocketMessageResponses.GENERATION_STARTED, {
// message: 'Code generation is already in progress'
// });
return;
}
// Start generation process
logger.info('Starting code generation process');
agent.getBehavior().generateAllFiles().catch(error => {
logger.error('Error during code generation:', error);
sendError(connection, `Error generating files: ${error instanceof Error ? error.message : String(error)}`);
}).finally(() => {
// Only clear shouldBeGenerating on successful completion
// (errors might want to retry, so this could be handled differently)
if (!agent.getBehavior().isCodeGenerating()) {
agent.setState({
...agent.state,
shouldBeGenerating: false
});
}
});
break;
case WebSocketMessageRequests.DEPLOY: {
const target = parsedMessage.target ??
(agent.state.behaviorType === 'think' ? 'user' : 'platform');
if (target !== 'platform' && target !== 'user') {
sendError(connection, `Unsupported deployment target: ${String(target)}`);
return;
}
agent.deployToCloudflare(target).then((deploymentResult) => {
if (!deploymentResult?.deploymentUrl) {
logger.error('Deployment failed', { target });
return;
}
logger.info('Deployment completed', deploymentResult);
}).catch((error: unknown) => {
logger.error('Error during deployment:', error);
sendError(connection, error instanceof Error ? error.message : String(error));
});
break;
}
case WebSocketMessageRequests.PREVIEW:
// Deploy current state for preview
logger.info('Deploying for preview');
agent.getBehavior().deployToSandbox().then((deploymentResult) => {
logger.info(`Preview deployed successfully!, deploymentResult:`, deploymentResult);
}).catch((error: unknown) => {
logger.error('Error during preview deployment:', error);
});
break;
case WebSocketMessageRequests.ROLLBACK_TO_COMMIT: {
const commitHash = parsedMessage.commitHash;
if (!commitHash) {
sendError(connection, 'Missing commitHash for rollback');
return;
}
const behavior = agent.getBehavior() as unknown as {
rollbackToCommit?: (commitHash: string) => Promise<void>;
};
if (typeof behavior.rollbackToCommit !== 'function') {
sendError(connection, 'Rollback is not supported for this app');
return;
}
logger.info('Rolling back to commit', { commitHash });
behavior.rollbackToCommit(commitHash).catch((error: unknown) => {
logger.error('Error during rollback:', error);
sendError(connection, `Error during rollback: ${error instanceof Error ? error.message : String(error)}`);
});
break;
}
case WebSocketMessageRequests.CAPTURE_SCREENSHOT:
if (!parsedMessage.data?.url) {
sendError(connection, 'Missing url for screenshot capture');
return;
}
agent.getBehavior().captureScreenshot(parsedMessage.data.url, parsedMessage.data.viewport as { width: number; height: number } | undefined).then((screenshotResult) => {
if (!screenshotResult) {
logger.error('Failed to capture screenshot');
return;
}
logger.info('Screenshot captured successfully!', screenshotResult);
}).catch((error: unknown) => {
logger.error('Error during screenshot capture:', error);
});
break;
case WebSocketMessageRequests.STOP_GENERATION: {
logger.info('User requested to stop generation');
// Cancel current inference operation
const wasCancelled = agent.getBehavior().cancelCurrentInference();
// Clear shouldBeGenerating flag
agent.setState({
...agent.state,
shouldBeGenerating: false
});
sendToConnection(connection, WebSocketMessageResponses.GENERATION_STOPPED, {
message: wasCancelled
? 'Inference operation cancelled successfully'
: 'No active inference to cancel'
});
break;
}
case WebSocketMessageRequests.RESUME_GENERATION:
// Set shouldBeGenerating and restart generation
logger.info('Resuming code generation');
agent.setState({
...agent.state,
shouldBeGenerating: true
});
if (!agent.getBehavior().isCodeGenerating()) {
sendToConnection(connection, WebSocketMessageResponses.GENERATION_RESUMED, {
message: 'Code generation resumed'
});
agent.getBehavior().generateAllFiles().catch(error => {
logger.error('Error resuming code generation:', error);
sendError(connection, `Error resuming generation: ${error instanceof Error ? error.message : String(error)}`);
});
} else {
// sendToConnection(connection, WebSocketMessageResponses.GENERATION_STARTED, {
// message: 'Code generation is already in progress'
// });
}
break;
case WebSocketMessageRequests.GITHUB_EXPORT:
// DEPRECATED: WebSocket-based GitHub export replaced with OAuth flow
// GitHub Apps require OAuth user access tokens for user repository creation
sendToConnection(connection, WebSocketMessageResponses.GITHUB_EXPORT_ERROR, {
message: 'GitHub export via WebSocket is deprecated',
error: 'Please use the GitHub export button which will redirect you to authorize with GitHub OAuth'
});
break;
case WebSocketMessageRequests.USER_SUGGESTION:
// Handle user suggestion for conversational AI
logger.info('Received user suggestion', {
messageLength: parsedMessage.message?.length || 0,
hasImages: !!parsedMessage.images && parsedMessage.images.length > 0,
imageCount: parsedMessage.images?.length || 0
});
if (!parsedMessage.message) {
sendError(connection, 'No message provided in user suggestion');
return;
}
// Validate image count and size
if (parsedMessage.images && parsedMessage.images.length > 0) {
if (parsedMessage.images.length > MAX_IMAGES_PER_MESSAGE) {
sendError(connection, `Maximum ${MAX_IMAGES_PER_MESSAGE} images allowed per message. Received ${parsedMessage.images.length} images.`);
return;
}
// Validate each image size
for (const image of parsedMessage.images) {
if (image.size && image.size > MAX_IMAGE_SIZE_BYTES) {
sendError(connection, `Image "${image.filename}" exceeds maximum size of ${MAX_IMAGE_SIZE_BYTES / (1024 * 1024)}MB`);
return;
}
}
}
// Check usage limits before processing user suggestion
try {
const env = agent.env;
const userId = agent.state.metadata.userId;
// The encrypted blob was captured from the HttpOnly cookie at WS
// upgrade time (see codingAgent.onConnect) and stored in DO state.
// WS frames do not carry cookies, so we rely on that snapshot.
const userToken = agent.state.cloudflareToken || null;
// Check limits and balance (this may transparently refresh the token).
const wsOrigin = agent.state.wsOrigin || undefined;
const limitResult = await checkUsageAndBalance(env, userId, undefined, userToken, wsOrigin);
// If a refresh occurred, keep the DO-cached blob fresh so subsequent
// user_suggestion messages pick up the new access token.
if (limitResult.refreshedBlob) {
agent.setState({ ...agent.state, cloudflareToken: limitResult.refreshedBlob });
}
if (!limitResult.allowed) {
logger.warn('User suggestion blocked by usage check', {
userId,
reason: limitResult.reason,
withinLimits: limitResult.withinLimits,
remaining: limitResult.remaining,
hasUserToken: limitResult.hasUserToken,
balance: limitResult.balance,
});
// Send structured error for frontend to show as popup
sendToConnection(connection, WebSocketMessageResponses.ERROR, {
error: limitResult.reason,
code: 'USAGE_LIMIT_EXCEEDED',
showAsPopup: true,
});
return;
}
} catch (error) {
logger.error('Failed to check usage:', error);
sendToConnection(connection, WebSocketMessageResponses.ERROR, {
error: `Error processing request: ${error instanceof Error ? error.message : String(error)}`,
showAsPopup: true,
});
return;
}
agent.handleUserInput(parsedMessage.message, parsedMessage.images).catch((error: unknown) => {
logger.error('Error handling user suggestion:', error);
sendError(connection, `Error processing user suggestion: ${error instanceof Error ? error.message : String(error)}`);
});
break;
case WebSocketMessageRequests.GET_MODEL_CONFIGS:
logger.info('Fetching model configurations');
agent.getBehavior().getModelConfigsInfo().then(configsInfo => {
sendToConnection(connection, WebSocketMessageResponses.MODEL_CONFIGS_INFO, {
message: 'Model configurations retrieved',
configs: configsInfo
});
}).catch((error: unknown) => {
logger.error('Error fetching model configs:', error);
sendError(connection, `Error fetching model configurations: ${error instanceof Error ? error.message : String(error)}`);
});
break;
case WebSocketMessageRequests.CLEAR_CONVERSATION:
logger.info('Clearing conversation history');
agent.clearConversation();
break;
case WebSocketMessageRequests.VAULT_UNLOCKED:
agent.handleVaultUnlocked();
break;
case WebSocketMessageRequests.VAULT_LOCKED:
agent.handleVaultLocked();
break;
case WebSocketMessageRequests.GET_CONVERSATION_STATE:
try {
const loader = agent.getConversationMessageLoader();
const state = await loader.load();
const debugState = agent
.getBehavior()
.getDeepDebugSessionState();
logger.info(`Conversation state retrieved for ${state.id}`);
sendToConnection(connection, WebSocketMessageResponses.CONVERSATION_STATE, {
state,
deepDebugSession: debugState
});
} catch (error) {
logger.error('Error fetching conversation state:', error);
sendError(connection, `Error fetching conversation state: ${error instanceof Error ? error.message : String(error)}`);
}
break;
// Disabled it for now
// case WebSocketMessageRequests.TERMINAL_COMMAND:
// // Handle terminal command execution
// logger.info('Received terminal command', {
// command: parsedMessage.command,
// timestamp: parsedMessage.timestamp
// });
// if (!parsedMessage.command) {
// sendError(connection, 'No command provided');
// return;
// }
// // Execute terminal command
// agent.executeTerminalCommand(parsedMessage.command, connection as any)
// .catch((error: unknown) => {
// logger.error('Error executing terminal command:', error);
// sendToConnection(connection, WebSocketMessageResponses.TERMINAL_OUTPUT, {
// output: `Error: ${error instanceof Error ? error.message : String(error)}`,
// outputType: 'stderr' as const,
// timestamp: Date.now()
// });
// });
// break;
default:
sendError(connection, `Unknown message type: ${parsedMessage.type}`);
}
} catch (error) {
logger.error('Error processing WebSocket message:', error);
sendError(connection, `Error processing message: ${error instanceof Error ? error.message : String(error)}`);
}
}
export function handleWebSocketClose(agent: CodeGeneratorAgent, connection: Connection): void {
logger.info(`WebSocket connection closed: ${connection.id}`);
// Clear vault session on disconnect for security
agent.handleVaultLocked();
}
export function broadcastToConnections<T extends WebSocketMessageType>(
agent: { getWebSockets(): WebSocket[] },
type: T,
data: WebSocketMessageData<T>
): void {
const connections = agent.getWebSockets();
for (const connection of connections) {
sendToConnection(connection, type, data);
}
}
export function sendToConnection<T extends WebSocketMessageType>(
connection: WebSocket,
type: T,
data: WebSocketMessageData<T>
): void {
try {
const message: WebSocketMessage = { type, ...data } as WebSocketMessage;
connection.send(JSON.stringify(message));
} catch (error) {
console.error(`Error sending message to connection ${connection.url}:`, error);
}
}
export function sendError(connection: WebSocket, errorMessage: string): void {
sendToConnection(connection, 'error', { error: errorMessage });
}