PROTOBUS: Streaming, State, Service Auto-Config (RooCodeInc#3253)

* round 1

* round 2 - searchFiles integration attempt

* undo streaming search experiments

* Start state.proto and related migrations

* state subscription

* get the main state flow using it

* correct stream ending early, debug statements

* clean up build-proto service config

* autogenerate index.tses

* auto-generate grpc-client service exports

* rename web-content -> web to make codegen work

* cleaned up streaming flow & cancels

* v3.14.0 Release Notes

v3.14.0 Release Notes

* prettier

* uhh prettier ?

* rename GrpcRequestRegistry file

* auto-generate directory for new services in the config

* generate template proto if it doesn't exist and provide instructions

* format fix

* add models service back to new system

---------

Co-authored-by: Andrei Edell <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: 3 people

proto/build-proto.js

@@ -12,11 +12,27 @@ const require = createRequire(import.meta.url)
 const protoc = path.join(require.resolve("grpc-tools"), "../bin/protoc")
 const tsProtoPlugin = require.resolve("ts-proto/protoc-gen-ts_proto")
 
-// Get script directory and root directory
 const __filename = fileURLToPath(import.meta.url)
 const SCRIPT_DIR = path.dirname(__filename)
 const ROOT_DIR = path.resolve(SCRIPT_DIR, "..")
 
+// List of gRPC services
+// To add a new service, simply add it to this map and run this script
+// The service handler will be automatically discovered and used by grpc-handler.ts
+const serviceNameMap = {
+	account: "cline.AccountService",
+	browser: "cline.BrowserService",
+	checkpoints: "cline.CheckpointsService",
+	file: "cline.FileService",
+	mcp: "cline.McpService",
+	state: "cline.StateService",
+	task: "cline.TaskService",
+	web: "cline.WebService",
+	models: "cline.ModelsService",
+	// Add new services here - no other code changes needed!
+}
+const serviceDirs = Object.keys(serviceNameMap).map((serviceKey) => path.join(ROOT_DIR, "src", "core", "controller", serviceKey))
+
 async function main() {
 	console.log(chalk.bold.blue("Starting Protocol Buffer code generation..."))
 
@@ -33,6 +49,9 @@ async function main() {
 		await fs.unlink(path.join(TS_OUT_DIR, file))
 	}
 
+	// Check for missing proto files for services in serviceNameMap
+	await ensureProtoFilesExist()
+
 	// Process all proto files
 	console.log(chalk.cyan("Processing proto files from"), SCRIPT_DIR)
 	const protoFiles = await globby("*.proto", { cwd: SCRIPT_DIR })
@@ -64,42 +83,135 @@ async function main() {
 	console.log(chalk.green("Protocol Buffer code generation completed successfully."))
 	console.log(chalk.green(`TypeScript files generated in: ${TS_OUT_DIR}`))
 
-	// Generate method registration files
 	await generateMethodRegistrations()
+	await generateServiceConfig()
+	await generateGrpcClientConfig()
+}
+
+/**
+ * Generate a gRPC client configuration file for the webview
+ * This eliminates the need for manual imports and client creation in grpc-client.ts
+ */
+async function generateGrpcClientConfig() {
+	console.log(chalk.cyan("Generating gRPC client configuration..."))
+
+	const serviceImports = []
+	const serviceClientCreations = []
+	const serviceExports = []
+
+	// Process each service in the serviceNameMap
+	for (const [dirName, fullServiceName] of Object.entries(serviceNameMap)) {
+		const capitalizedName = dirName.charAt(0).toUpperCase() + dirName.slice(1)
+
+		// Add import statement
+		serviceImports.push(`import { ${capitalizedName}ServiceDefinition } from "@shared/proto/${dirName}"`)
+
+		// Add client creation
+		serviceClientCreations.push(
+			`const ${capitalizedName}ServiceClient = createGrpcClient(${capitalizedName}ServiceDefinition)`,
+		)
+
+		// Add to exports
+		serviceExports.push(`${capitalizedName}ServiceClient`)
+	}
+
+	// Generate the file content
+	const content = `// AUTO-GENERATED FILE - DO NOT MODIFY DIRECTLY
+// Generated by proto/build-proto.js
+
+import { createGrpcClient } from "./grpc-client-base"
+${serviceImports.join("\n")}
+
+${serviceClientCreations.join("\n")}
+
+export {
+	${serviceExports.join(",\n\t")}
+}`
+
+	const configPath = path.join(ROOT_DIR, "webview-ui", "src", "services", "grpc-client.ts")
+	await fs.writeFile(configPath, content)
+	console.log(chalk.green(`Generated gRPC client at ${configPath}`))
+}
 
-	// Make the script executable
-	try {
-		await fs.chmod(path.join(SCRIPT_DIR, "build-proto.js"), 0o755)
-	} catch (error) {
-		console.warn(chalk.yellow("Warning: Could not make script executable:"), error)
+/**
+ * Parse proto files to extract streaming method information
+ * @param protoFiles Array of proto file names
+ * @param scriptDir Directory containing proto files
+ * @returns Map of service names to their streaming methods
+ */
+async function parseProtoForStreamingMethods(protoFiles, scriptDir) {
+	console.log(chalk.cyan("Parsing proto files for streaming methods..."))
+
+	// Map of service name to array of streaming method names
+	const streamingMethodsMap = new Map()
+
+	for (const protoFile of protoFiles) {
+		const content = await fs.readFile(path.join(scriptDir, protoFile), "utf8")
+
+		// Extract package name
+		const packageMatch = content.match(/package\s+([^;]+);/)
+		const packageName = packageMatch ? packageMatch[1].trim() : "unknown"
+
+		// Extract service definitions
+		const serviceMatches = Array.from(content.matchAll(/service\s+(\w+)\s*\{([^}]+)\}/g))
+		for (const serviceMatch of serviceMatches) {
+			const serviceName = serviceMatch[1]
+			const serviceBody = serviceMatch[2]
+			const fullServiceName = `${packageName}.${serviceName}`
+
+			// Extract method definitions with streaming
+			const methodMatches = Array.from(
+				serviceBody.matchAll(/rpc\s+(\w+)\s*\(\s*(stream\s+)?(\w+)\s*\)\s*returns\s*\(\s*(stream\s+)?(\w+)\s*\)/g),
+			)
+
+			const streamingMethods = []
+			for (const methodMatch of methodMatches) {
+				const methodName = methodMatch[1]
+				const isRequestStreaming = !!methodMatch[2]
+				const requestType = methodMatch[3]
+				const isResponseStreaming = !!methodMatch[4]
+				const responseType = methodMatch[5]
+
+				if (isResponseStreaming) {
+					streamingMethods.push({
+						name: methodName,
+						requestType,
+						responseType,
+						isRequestStreaming,
+					})
+				}
+			}
+
+			if (streamingMethods.length > 0) {
+				streamingMethodsMap.set(fullServiceName, streamingMethods)
+			}
+		}
 	}
+
+	return streamingMethodsMap
 }
 
 async function generateMethodRegistrations() {
 	console.log(chalk.cyan("Generating method registration files..."))
 
-	const serviceDirs = [
-		path.join(ROOT_DIR, "src", "core", "controller", "account"),
-		path.join(ROOT_DIR, "src", "core", "controller", "browser"),
-		path.join(ROOT_DIR, "src", "core", "controller", "checkpoints"),
-		path.join(ROOT_DIR, "src", "core", "controller", "file"),
-		path.join(ROOT_DIR, "src", "core", "controller", "mcp"),
-		path.join(ROOT_DIR, "src", "core", "controller", "models"),
-		path.join(ROOT_DIR, "src", "core", "controller", "task"),
-		path.join(ROOT_DIR, "src", "core", "controller", "web-content"),
-		// Add more service directories here as needed
-	]
+	// Parse proto files for streaming methods
+	const protoFiles = await globby("*.proto", { cwd: SCRIPT_DIR })
+	const streamingMethodsMap = await parseProtoForStreamingMethods(protoFiles, SCRIPT_DIR)
 
 	for (const serviceDir of serviceDirs) {
 		try {
 			await fs.access(serviceDir)
 		} catch (error) {
-			console.log(chalk.gray(`Skipping ${serviceDir} - directory does not exist`))
-			continue
+			console.log(chalk.cyan(`Creating directory ${serviceDir} for new service`))
+			await fs.mkdir(serviceDir, { recursive: true })
 		}
 
 		const serviceName = path.basename(serviceDir)
 		const registryFile = path.join(serviceDir, "methods.ts")
+		const indexFile = path.join(serviceDir, "index.ts")
+
+		const fullServiceName = serviceNameMap[serviceName]
+		const streamingMethods = streamingMethodsMap.get(fullServiceName) || []
 
 		console.log(chalk.cyan(`Generating method registrations for ${serviceName}...`))
 
@@ -109,8 +221,8 @@ async function generateMethodRegistrations() {
 		// Filter out index.ts and methods.ts
 		const implementationFiles = files.filter((file) => file !== "index.ts" && file !== "methods.ts")
 
-		// Create the output file with header
-		let content = `// AUTO-GENERATED FILE - DO NOT MODIFY DIRECTLY
+		// Create the methods.ts file with header
+		let methodsContent = `// AUTO-GENERATED FILE - DO NOT MODIFY DIRECTLY
 // Generated by proto/build-proto.js
 
 // Import all method implementations
@@ -119,31 +231,177 @@ import { registerMethod } from "./index"\n`
 		// Add imports for all implementation files
 		for (const file of implementationFiles) {
 			const baseName = path.basename(file, ".ts")
-			content += `import { ${baseName} } from "./${baseName}"\n`
+			methodsContent += `import { ${baseName} } from "./${baseName}"\n`
+		}
+
+		// Add streaming methods information
+		if (streamingMethods.length > 0) {
+			methodsContent += `\n// Streaming methods for this service
+export const streamingMethods = ${JSON.stringify(
+				streamingMethods.map((m) => m.name),
+				null,
+				2,
+			)}\n`
 		}
 
 		// Add registration function
-		content += `\n// Register all ${serviceName} service methods
+		methodsContent += `\n// Register all ${serviceName} service methods
 export function registerAllMethods(): void {
 \t// Register each method with the registry\n`
 
 		// Add registration statements
 		for (const file of implementationFiles) {
 			const baseName = path.basename(file, ".ts")
-			content += `\tregisterMethod("${baseName}", ${baseName})\n`
+			const isStreaming = streamingMethods.some((m) => m.name === baseName)
+
+			if (isStreaming) {
+				methodsContent += `\tregisterMethod("${baseName}", ${baseName}, { isStreaming: true })\n`
+			} else {
+				methodsContent += `\tregisterMethod("${baseName}", ${baseName})\n`
+			}
 		}
 
 		// Close the function
-		content += `}`
+		methodsContent += `}`
 
-		// Write the file
-		await fs.writeFile(registryFile, content)
+		// Write the methods.ts file
+		await fs.writeFile(registryFile, methodsContent)
 		console.log(chalk.green(`Generated ${registryFile}`))
+
+		// Generate index.ts file
+		const capitalizedServiceName = serviceName.charAt(0).toUpperCase() + serviceName.slice(1)
+		const indexContent = `// AUTO-GENERATED FILE - DO NOT MODIFY DIRECTLY
+// Generated by proto/build-proto.js
+
+import { createServiceRegistry, ServiceMethodHandler, StreamingMethodHandler } from "../grpc-service"
+import { StreamingResponseHandler } from "../grpc-handler"
+import { registerAllMethods } from "./methods"
+
+// Create ${serviceName} service registry
+const ${serviceName}Service = createServiceRegistry("${serviceName}")
+
+// Export the method handler types and registration function
+export type ${capitalizedServiceName}MethodHandler = ServiceMethodHandler
+export type ${capitalizedServiceName}StreamingMethodHandler = StreamingMethodHandler
+export const registerMethod = ${serviceName}Service.registerMethod
+
+// Export the request handlers
+export const handle${capitalizedServiceName}ServiceRequest = ${serviceName}Service.handleRequest
+export const handle${capitalizedServiceName}ServiceStreamingRequest = ${serviceName}Service.handleStreamingRequest
+export const isStreamingMethod = ${serviceName}Service.isStreamingMethod
+
+// Register all ${serviceName} methods
+registerAllMethods()`
+
+		// Write the index.ts file
+		await fs.writeFile(indexFile, indexContent)
+		console.log(chalk.green(`Generated ${indexFile}`))
 	}
 
 	console.log(chalk.green("Method registration files generated successfully."))
 }
 
+/**
+ * Generate a service configuration file that maps service names to their handlers
+ * This eliminates the need for manual switch/case statements in grpc-handler.ts
+ */
+async function generateServiceConfig() {
+	console.log(chalk.cyan("Generating service configuration file..."))
+
+	const serviceImports = []
+	const serviceConfigs = []
+
+	// Add all services from the serviceNameMap
+	for (const [dirName, fullServiceName] of Object.entries(serviceNameMap)) {
+		const capitalizedName = dirName.charAt(0).toUpperCase() + dirName.slice(1)
+		serviceImports.push(
+			`import { handle${capitalizedName}ServiceRequest, handle${capitalizedName}ServiceStreamingRequest } from "./${dirName}/index"`,
+		)
+		serviceConfigs.push(`
+  "${fullServiceName}": {
+    requestHandler: handle${capitalizedName}ServiceRequest,
+    streamingHandler: handle${capitalizedName}ServiceStreamingRequest
+  }`)
+	}
+
+	const content = `// AUTO-GENERATED FILE - DO NOT MODIFY DIRECTLY
+// Generated by proto/build-proto.js
+
+import { Controller } from "./index"
+import { StreamingResponseHandler } from "./grpc-handler"
+${serviceImports.join("\n")}
+
+/**
+ * Configuration for a service handler
+ */
+export interface ServiceHandlerConfig {
+  requestHandler: (controller: Controller, method: string, message: any) => Promise<any>;
+  streamingHandler: (controller: Controller, method: string, message: any, responseStream: StreamingResponseHandler, requestId?: string) => Promise<void>;
+}
+
+/**
+ * Map of service names to their handler configurations
+ */
+export const serviceHandlers: Record<string, ServiceHandlerConfig> = {${serviceConfigs.join(",")}
+};`
+
+	const configPath = path.join(ROOT_DIR, "src", "core", "controller", "grpc-service-config.ts")
+	await fs.writeFile(configPath, content)
+	console.log(chalk.green(`Generated service configuration at ${configPath}`))
+}
+
+/**
+ * Ensure that a .proto file exists for each service in the serviceNameMap
+ * If a .proto file doesn't exist, create a template file
+ */
+async function ensureProtoFilesExist() {
+	console.log(chalk.cyan("Checking for missing proto files..."))
+
+	// Get existing proto files
+	const existingProtoFiles = await globby("*.proto", { cwd: SCRIPT_DIR })
+	const existingProtoServices = existingProtoFiles.map((file) => path.basename(file, ".proto"))
+
+	// Check each service in serviceNameMap
+	for (const [serviceName, fullServiceName] of Object.entries(serviceNameMap)) {
+		if (!existingProtoServices.includes(serviceName)) {
+			console.log(chalk.yellow(`Creating template proto file for ${serviceName}...`))
+
+			// Extract service class name from full name (e.g., "cline.ModelsService" -> "ModelsService")
+			const serviceClassName = fullServiceName.split(".").pop()
+
+			// Create template proto file
+			const protoContent = `syntax = "proto3";
+
+package cline;
+option java_package = "bot.cline.proto";
+option java_multiple_files = true;
+
+import "common.proto";
+
+// ${serviceClassName} provides methods for managing ${serviceName}
+service ${serviceClassName} {
+  // Add your RPC methods here
+  // Example (String is from common.proto, responses should be generic types):
+  // rpc YourMethod(YourRequest) returns (String);
+}
+
+// Add your message definitions here
+// Example (Requests must always start with Metadata):
+// message YourRequest {
+//   Metadata metadata = 1;
+//   string stringField = 2;
+//   int32 int32Field = 3;
+// }
+`
+
+			// Write the template proto file
+			const protoFilePath = path.join(SCRIPT_DIR, `${serviceName}.proto`)
+			await fs.writeFile(protoFilePath, protoContent)
+			console.log(chalk.green(`Created template proto file at ${protoFilePath}`))
+		}
+	}
+}
+
 // Run the main function
 main().catch((error) => {
 	console.error(chalk.red("Error:"), error)