mesh changed sdk

Author: randilfernando

node/mesh_peer.ts

@@ -22,12 +22,11 @@ async function main() {
     await mesh.start();
 
     // 3. Register Local Service (Just to be a good citizen, though we mainly act as client here for this test layout)
-    await mesh.registerService(MY_SERVICE_NAME, {
-        onRequestChannel: async (channel: IChannel) => {
+    mesh.registerService(MY_SERVICE_NAME)
+        .onRequestChannel(async (channel: IChannel) => {
             console.log(`[Node] Accepted inbound connection (Unexpected for this test flow, but handling...)`);
             channel.close();
-        }
-    });
+        });
     console.log(`[Node] Registered '${MY_SERVICE_NAME}'`);
 
     // 4. Connect to Python Peer and Start Streaming

node/src/abstract.ts

@@ -66,9 +66,14 @@ export interface IServiceHandlers {
     onRequestReply?: (msg: any) => Promise<any>;
 }
 
+export interface IServiceRegistration {
+    onRequestChannel(handler: (channel: IChannel) => void): IServiceRegistration;
+    onRequestReply(handler: (msg: any) => Promise<any>): IServiceRegistration;
+}
+
 export interface IMesh {
     start(): Promise<void>;
-    registerService(serviceName: string, handlers: IServiceHandlers): Promise<void>;
+    registerService(serviceName: string): IServiceRegistration;
     service(name: string): IServiceClient;
     topic(name: string): ITopic;
     close(): Promise<void>;

node/src/mesh.ts

@@ -4,10 +4,10 @@ import { MeshConfig, Node } from './shared/types';
 import { RedisServiceDiscovery } from './discovery/redis-discovery';
 import { H2Communication } from './communication/h2-communication';
 import { ServiceClient } from './service-client';
-import { IMesh, IServiceClient, ITopic, IServiceHandlers, IChannel, IServiceDiscovery, ICommunication, IPubSub } from './abstract';
+import { IMesh, IServiceClient, ITopic, IServiceHandlers, IChannel, IServiceDiscovery, ICommunication, IPubSub, IServiceRegistration } from './abstract';
 import { RedisPubSub, Topic } from './communication/pubsub';
 
-export { MeshConfig, Node, IChannel, IServiceHandlers, IChannel as Channel, IServiceHandlers as ServiceHandlers, Topic }; // Export both for compat
+export { MeshConfig, Node, IChannel, IServiceHandlers, IServiceRegistration, IChannel as Channel, IServiceHandlers as ServiceHandlers, Topic }; // Export both for compat
 // User asked to standardize naming. I'll export renamed IChannel and IServiceHandlers.
 // IF old code relies on Channel/ServiceHandlers export from Mesh, I should alias them or update consumer code.
 // Given "standardize abstraction", ideally I use IChannel everywhere.
@@ -46,25 +46,31 @@ export class Mesh implements IMesh {
         await this.pubsub.start();
     }
 
-    public async registerService(serviceName: string, handlers: IServiceHandlers) {
-        // Setup listeners
-        if (handlers.onRequestChannel) {
-            this.communication.listenChannel(serviceName, handlers.onRequestChannel);
-        }
-        if (handlers.onRequestReply) {
-            this.communication.listenUnary(serviceName, handlers.onRequestReply);
-        }
+    public registerService(serviceName: string): IServiceRegistration {
+        const registration = {
+            onRequestChannel: (handler: (channel: IChannel) => void) => {
+                this.communication.listenChannel(serviceName, handler);
+                return registration;
+            },
+            onRequestReply: (handler: (msg: any) => Promise<any>) => {
+                this.communication.listenUnary(serviceName, handler);
+                return registration;
+            }
+        };
 
-        // Register with SD
+        // Register with SD immediately
         const node: Node = {
             id: this.nodeId,
-            host: this.config.host!, // Detected in constructor
+            host: this.config.host!,
             port: this.config.port,
             service_name: serviceName
         };
 
-        await this.discovery.registerService(node);
-        console.log(`[Mesh] Registered service '${serviceName}' on ${node.host}:${node.port}`);
+        this.discovery.registerService(node)
+            .then(() => console.log(`[Mesh] Registered service '${serviceName}' on ${node.host}:${node.port}`))
+            .catch(err => console.error(`[Mesh] Failed to register service '${serviceName}':`, err));
+
+        return registration;
     }
 
     private getLocalIP(): string {

python/mesh_peer.py

@@ -67,7 +67,7 @@ async def receive_loop():
 
         logger.info("Connection handler finished.")
 
-    await mesh.register_service("echo-service", {"on_request_channel": incoming_handler})
+    mesh.register_service("echo-service").on_request_channel(incoming_handler)
     logger.info(f"Registered 'echo-service' and waiting for connections...")
 
     # Register signal handlers

python/src/mesh/abstract.py

@@ -141,6 +141,15 @@ def get_service(self, service_name: str) -> List[Node]:
         """
         pass
 
+class IServiceRegistration(ABC):
+    @abstractmethod
+    def on_request_channel(self, handler: Callable[['IChannel'], Awaitable[None]]) -> 'IServiceRegistration':
+        pass
+    
+    @abstractmethod
+    def on_request_reply(self, handler: Callable[[Any], Awaitable[Any]]) -> 'IServiceRegistration':
+        pass
+
 class IMesh(ABC):
     @abstractmethod
     async def start(self):
@@ -150,9 +159,9 @@ async def start(self):
         pass
 
     @abstractmethod
-    async def register_service(self, service_name: str, handlers: IServiceHandlers):
+    def register_service(self, service_name: str) -> IServiceRegistration:
         """
-        Register a service provider with handlers.
+        Register a service provider. Returns a registration object to attach handlers.
         """
         pass
 

python/src/mesh/mesh.py

@@ -1,13 +1,13 @@
 import asyncio
 import uuid
 import socket
-from typing import Optional, List, Callable, Awaitable
+from typing import Optional, List, Callable, Awaitable, Any
 from mesh.shared.types import MeshConfig, Node
 from mesh.service_discovery import RedisServiceDiscovery
 from mesh.communication import H2Communication
 from mesh.communication import H2Communication
 from mesh.communication.pubsub import RedisPubSub
-from mesh.abstract import IMesh, IChannel, ITopic, IServiceHandlers
+from mesh.abstract import IMesh, IChannel, ITopic, IServiceHandlers, IServiceRegistration
 from mesh.service_client import ServiceClient
 
 class Mesh(IMesh):    
@@ -46,18 +46,24 @@ async def start(self):
         # Start a background task to keep the process alive or handle signals? 
         # For now, just starts the asyncio server.
 
-    async def register_service(
+    def register_service(
         self, 
-        service_name: str, 
-        handlers: IServiceHandlers
-    ):
-        # 1. Setup listeners
-        if handlers.get('on_request_channel'):
-            self.communication.listen_channel(service_name, handlers['on_request_channel'])
-        
-        if handlers.get('on_request_reply'):
-            self.communication.listen_unary(service_name, handlers['on_request_reply'])
-        
+        service_name: str
+    ) -> IServiceRegistration:
+        # Internal registration class
+        class ServiceRegistration(IServiceRegistration):
+            def __init__(self, mesh, service_name):
+                self.mesh = mesh
+                self.service_name = service_name
+                
+            def on_request_channel(self, handler: Callable[['IChannel'], Awaitable[None]]) -> 'IServiceRegistration':
+                self.mesh.communication.listen_channel(self.service_name, handler)
+                return self
+                
+            def on_request_reply(self, handler: Callable[[Any], Awaitable[Any]]) -> 'IServiceRegistration':
+                self.mesh.communication.listen_unary(self.service_name, handler)
+                return self
+
         # 2. Register with SD
         node = Node(
             id=self._node_id,
@@ -69,6 +75,8 @@ async def register_service(
         # Register with SD
         self.service_discovery.register_service(node)
         print(f"Mesh: Registered service '{service_name}' on {node.host}:{node.port} (ID: {node.id})")
+        
+        return ServiceRegistration(self, service_name)
 
 
     def service(self, service_name: str) -> ServiceClient:

python/tests/test_mesh.py

@@ -66,8 +66,8 @@ async def handler2(channel):
     await mesh1.start()
     await mesh2.start()
 
-    await mesh1.register_service(service_name, {"on_request_channel": handler1})
-    await mesh2.register_service(service_name, {"on_request_channel": handler2})
+    mesh1.register_service(service_name).on_request_channel(handler1)
+    mesh2.register_service(service_name).on_request_channel(handler2)
 
     # Wait for binding
     await asyncio.sleep(0.1)
@@ -112,8 +112,8 @@ def get_service(self, service_name):
     # Re-register (since we swapped the class, previous registers might be lost or we just restart flow)
     # Actually simpler to just restart flow with correct mock
 
-    await mesh1.register_service(service_name, {})
-    await mesh2.register_service(service_name, {})
+    mesh1.register_service(service_name)
+    mesh2.register_service(service_name)
 
     # Verify SD
     nodes = SharedMockSD(None).get_service(service_name)