@motiadev/core 0.1.0-beta.3 → 0.1.0-beta.30

package/dist/index.d.ts

@@ -9,3 +9,4 @@ export { isApiStep, isCronStep, isEventStep, isNoopStep } from './src/guards';
 export { LockedData } from './src/locked-data';
 export { getStepConfig } from './src/get-step-config';
 export { StateAdapter } from './src/state/state-adapter';
+export { createMermaidGenerator } from './src/mermaid-generator';

package/dist/index.js

@@ -14,7 +14,7 @@ var __exportStar = (this && this.__exportStar) || function(m, exports) {
     for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.getStepConfig = exports.LockedData = exports.isNoopStep = exports.isEventStep = exports.isCronStep = exports.isApiStep = exports.setupCronHandlers = exports.createStateAdapter = exports.Logger = exports.globalLogger = exports.createEventManager = exports.createStepHandlers = exports.createServer = void 0;
+exports.createMermaidGenerator = exports.getStepConfig = exports.LockedData = exports.isNoopStep = exports.isEventStep = exports.isCronStep = exports.isApiStep = exports.setupCronHandlers = exports.createStateAdapter = exports.Logger = exports.globalLogger = exports.createEventManager = exports.createStepHandlers = exports.createServer = void 0;
 __exportStar(require("./src/types"), exports);
 var server_1 = require("./src/server");
 Object.defineProperty(exports, "createServer", { enumerable: true, get: function () { return server_1.createServer; } });
@@ -38,3 +38,5 @@ var locked_data_1 = require("./src/locked-data");
 Object.defineProperty(exports, "LockedData", { enumerable: true, get: function () { return locked_data_1.LockedData; } });
 var get_step_config_1 = require("./src/get-step-config");
 Object.defineProperty(exports, "getStepConfig", { enumerable: true, get: function () { return get_step_config_1.getStepConfig; } });
+var mermaid_generator_1 = require("./src/mermaid-generator");
+Object.defineProperty(exports, "createMermaidGenerator", { enumerable: true, get: function () { return mermaid_generator_1.createMermaidGenerator; } });

package/dist/src/flows-config-endpoint.js

@@ -8,13 +8,16 @@ const fs_1 = __importDefault(require("fs"));
 const path_1 = __importDefault(require("path"));
 const flowsConfigEndpoint = (app, baseDir) => {
     const configPath = path_1.default.join(baseDir, 'motia-workbench.json');
+    const getConfig = () => {
+        if (fs_1.default.existsSync(configPath)) {
+            return JSON.parse(fs_1.default.readFileSync(configPath, 'utf8'));
+        }
+        return {};
+    };
     app.post('/flows/:id/config', (req, res) => {
         const newFlowConfig = req.body;
         try {
-            if (!fs_1.default.existsSync(configPath)) {
-                fs_1.default.writeFileSync(configPath, JSON.stringify({}, null, 2));
-            }
-            const existingConfig = JSON.parse(fs_1.default.readFileSync(configPath, 'utf8'));
+            const existingConfig = getConfig();
             const updatedConfig = {
                 ...existingConfig,
             };
@@ -28,21 +31,20 @@ const flowsConfigEndpoint = (app, baseDir) => {
             res.status(200).send({ message: 'Flow config saved successfully' });
         }
         catch (error) {
-            console.error('Error saving flow config:', error);
+            console.error('Error saving flow config:', error.message);
             res.status(500).json({ error: 'Failed to save flow config' });
         }
     });
     app.get('/flows/:id/config', (req, res) => {
         const { id } = req.params;
         try {
-            const file = fs_1.default.readFileSync(configPath, 'utf8');
-            const allFlowsConfig = JSON.parse(file);
+            const allFlowsConfig = getConfig();
             const flowConfig = allFlowsConfig[id] || {};
             res.status(200).send(flowConfig);
         }
         catch (error) {
-            console.error('Error reading flow config:', error);
-            res.status(400).send({});
+            console.error('Error reading flow config:', error.message);
+            res.status(400).send({ error: 'Failed to read flow config' });
         }
     });
 };

package/dist/src/mermaid-generator.d.ts

@@ -0,0 +1,4 @@
+import { LockedData } from './locked-data';
+export declare const createMermaidGenerator: (baseDir: string) => {
+    initialize: (lockedData: LockedData) => void;
+};

package/dist/src/mermaid-generator.js

@@ -0,0 +1,203 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.createMermaidGenerator = void 0;
+const fs_1 = __importDefault(require("fs"));
+const path_1 = __importDefault(require("path"));
+const guards_1 = require("./guards");
+// Pure function to ensure diagrams directory exists
+const ensureDiagramsDirectory = (diagramsDir) => {
+    if (!fs_1.default.existsSync(diagramsDir)) {
+        fs_1.default.mkdirSync(diagramsDir, { recursive: true });
+    }
+};
+// Pure function to get node ID
+const getNodeId = (step, baseDir) => {
+    // Get relative path from the base directory
+    const relativePath = path_1.default.relative(baseDir, step.filePath);
+    // Remove common file extensions
+    const pathWithoutExtension = relativePath.replace(/\.(ts|js|tsx|jsx)$/, '');
+    // Replace slashes with underscores and dots with underscores
+    // Only keep alphanumeric characters and underscores
+    return pathWithoutExtension.replace(/[^a-zA-Z0-9]/g, '_');
+};
+// Pure function to get node label
+const getNodeLabel = (step) => {
+    // Get display name for node
+    const displayName = step.config.name || path_1.default.basename(step.filePath, path_1.default.extname(step.filePath));
+    // Add node type prefix to help distinguish types
+    let prefix = '';
+    if ((0, guards_1.isApiStep)(step))
+        prefix = '🌐 ';
+    else if ((0, guards_1.isEventStep)(step))
+        prefix = '⚡ ';
+    else if ((0, guards_1.isCronStep)(step))
+        prefix = '⏰ ';
+    else if ((0, guards_1.isNoopStep)(step))
+        prefix = '⚙️ ';
+    // Create a node label with the step name
+    return `["${prefix}${displayName}"]`;
+};
+// Pure function to get node style
+const getNodeStyle = (step) => {
+    // Apply style class based on step type
+    if ((0, guards_1.isApiStep)(step))
+        return ':::apiStyle';
+    if ((0, guards_1.isEventStep)(step))
+        return ':::eventStyle';
+    if ((0, guards_1.isCronStep)(step))
+        return ':::cronStyle';
+    if ((0, guards_1.isNoopStep)(step))
+        return ':::noopStyle';
+    return '';
+};
+// Pure function to generate connections
+const generateConnections = (emits, sourceStep, steps, sourceId, baseDir) => {
+    const connections = [];
+    if (!emits || !Array.isArray(emits) || emits.length === 0) {
+        return '';
+    }
+    // Helper function to check if a step subscribes to a topic
+    const stepSubscribesToTopic = (step, topic) => {
+        // Event steps use regular subscribes
+        if ((0, guards_1.isEventStep)(step) &&
+            step.config.subscribes &&
+            Array.isArray(step.config.subscribes) &&
+            step.config.subscribes.includes(topic)) {
+            return true;
+        }
+        // Noop and API steps use virtualSubscribes
+        if (((0, guards_1.isNoopStep)(step) || (0, guards_1.isApiStep)(step)) &&
+            step.config.virtualSubscribes &&
+            Array.isArray(step.config.virtualSubscribes) &&
+            step.config.virtualSubscribes.includes(topic)) {
+            return true;
+        }
+        return false;
+    };
+    emits.forEach((emit) => {
+        const topic = typeof emit === 'string' ? emit : emit.topic;
+        const label = typeof emit === 'string' ? topic : emit.label || topic;
+        steps.forEach((targetStep) => {
+            if (stepSubscribesToTopic(targetStep, topic)) {
+                const targetId = getNodeId(targetStep, baseDir);
+                connections.push(`    ${sourceId} -->|${label}| ${targetId}`);
+            }
+        });
+    });
+    return connections.join('\n');
+};
+// Pure function to generate flow diagram
+const generateFlowDiagram = (flowName, steps, baseDir) => {
+    // Start mermaid flowchart with top-down direction
+    let diagram = `flowchart TD\n`;
+    // Add class definitions for styling with explicit text color
+    const classDefinitions = [
+        `    classDef apiStyle fill:#f96,stroke:#333,stroke-width:2px,color:#fff`,
+        `    classDef eventStyle fill:#69f,stroke:#333,stroke-width:2px,color:#fff`,
+        `    classDef cronStyle fill:#9c6,stroke:#333,stroke-width:2px,color:#fff`,
+        `    classDef noopStyle fill:#3f3a50,stroke:#333,stroke-width:2px,color:#fff`,
+    ];
+    diagram += classDefinitions.join('\n') + '\n';
+    // Check if we have any steps
+    if (!steps || steps.length === 0) {
+        return diagram + '    empty[No steps in this flow]';
+    }
+    // Create node definitions with proper format
+    steps.forEach((step) => {
+        const nodeId = getNodeId(step, baseDir);
+        const nodeLabel = getNodeLabel(step);
+        const nodeStyle = getNodeStyle(step);
+        diagram += `    ${nodeId}${nodeLabel}${nodeStyle}\n`;
+    });
+    // Create connections between nodes
+    let connectionsStr = '';
+    steps.forEach((sourceStep) => {
+        const sourceId = getNodeId(sourceStep, baseDir);
+        // Helper function to process emissions if they exist
+        function processEmissions(emissionsArray, stepSource, stepsCollection, sourceIdentifier) {
+            if (emissionsArray && Array.isArray(emissionsArray)) {
+                return generateConnections(emissionsArray, stepSource, stepsCollection, sourceIdentifier, baseDir);
+            }
+            return '';
+        }
+        // Semantic variables to clarify which step types support which emission types
+        const supportsEmits = (0, guards_1.isApiStep)(sourceStep) || (0, guards_1.isEventStep)(sourceStep) || (0, guards_1.isCronStep)(sourceStep);
+        const supportsVirtualEmits = supportsEmits || (0, guards_1.isNoopStep)(sourceStep);
+        // Process regular emissions if supported
+        if (supportsEmits) {
+            const emitConnections = processEmissions(sourceStep.config.emits, sourceStep, steps, sourceId);
+            if (emitConnections) {
+                connectionsStr += emitConnections + '\n';
+            }
+        }
+        // Process virtual emissions if supported
+        if (supportsVirtualEmits) {
+            const virtualEmitConnections = processEmissions(sourceStep.config.virtualEmits, sourceStep, steps, sourceId);
+            if (virtualEmitConnections) {
+                connectionsStr += virtualEmitConnections + '\n';
+            }
+        }
+    });
+    // Add connections to the diagram
+    diagram += connectionsStr;
+    return diagram;
+};
+// Function to save a diagram to a file
+const saveDiagram = (diagramsDir, flowName, diagram) => {
+    const filePath = path_1.default.join(diagramsDir, `${flowName}.mmd`);
+    fs_1.default.writeFileSync(filePath, diagram);
+};
+// Function to remove a diagram file
+const removeDiagram = (diagramsDir, flowName) => {
+    const filePath = path_1.default.join(diagramsDir, `${flowName}.mmd`);
+    if (fs_1.default.existsSync(filePath)) {
+        fs_1.default.unlinkSync(filePath);
+    }
+};
+// Function to generate and save a diagram
+const generateAndSaveDiagram = (diagramsDir, flowName, flow, baseDir) => {
+    const diagram = generateFlowDiagram(flowName, flow.steps, baseDir);
+    saveDiagram(diagramsDir, flowName, diagram);
+};
+// Main exported function that creates the mermaid generator
+const createMermaidGenerator = (baseDir) => {
+    const diagramsDir = path_1.default.join(baseDir, '.mermaid');
+    ensureDiagramsDirectory(diagramsDir);
+    // Event handlers
+    const handleFlowCreated = (flowName, flow) => {
+        generateAndSaveDiagram(diagramsDir, flowName, flow, baseDir);
+    };
+    const handleFlowUpdated = (flowName, flow) => {
+        generateAndSaveDiagram(diagramsDir, flowName, flow, baseDir);
+    };
+    const handleFlowRemoved = (flowName) => {
+        removeDiagram(diagramsDir, flowName);
+    };
+    // Initialize function to hook into LockedData events
+    const initialize = (lockedData) => {
+        // Hook into flow events
+        lockedData.on('flow-created', (flowName) => {
+            handleFlowCreated(flowName, lockedData.flows[flowName]);
+        });
+        lockedData.on('flow-updated', (flowName) => {
+            handleFlowUpdated(flowName, lockedData.flows[flowName]);
+        });
+        lockedData.on('flow-removed', (flowName) => {
+            handleFlowRemoved(flowName);
+        });
+        // Generate diagrams for all existing flows
+        if (lockedData.flows && typeof lockedData.flows === 'object') {
+            Object.entries(lockedData.flows).forEach(([flowName, flow]) => {
+                generateAndSaveDiagram(diagramsDir, flowName, flow, baseDir);
+            });
+        }
+    };
+    // Return the public API
+    return {
+        initialize,
+    };
+};
+exports.createMermaidGenerator = createMermaidGenerator;

package/dist/src/node/middleware-compose.d.ts

@@ -0,0 +1 @@
+export declare const composeMiddleware: (...middlewares: any[]) => (req: any, ctx: any, handler: () => Promise<any>) => Promise<any>;

package/dist/src/node/middleware-compose.js

@@ -0,0 +1,12 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.composeMiddleware = void 0;
+/* eslint-disable @typescript-eslint/no-explicit-any */
+const composeMiddleware = (...middlewares) => {
+    return async (req, ctx, handler) => {
+        const composedHandler = middlewares.reduceRight((nextHandler, middleware) => () => middleware(req, ctx, nextHandler), handler);
+        return composedHandler();
+    };
+};
+exports.composeMiddleware = composeMiddleware;
+/* eslint-enable @typescript-eslint/no-explicit-any */

package/dist/src/node/node-runner.js

@@ -7,6 +7,7 @@ const path_1 = __importDefault(require("path"));
 const logger_1 = require("./logger");
 const rpc_state_manager_1 = require("./rpc-state-manager");
 const rpc_1 = require("./rpc");
+const middleware_compose_1 = require("./middleware-compose");
 // eslint-disable-next-line @typescript-eslint/no-require-imports
 require('dotenv').config();
 // Add ts-node registration before dynamic imports
@@ -38,8 +39,12 @@ async function runTypescriptModule(filePath, event) {
         const emit = async (data) => sender.send('emit', data);
         const context = { traceId, flows, logger, state, emit };
         sender.init();
-        // Call the function with provided arguments
-        const result = contextInFirstArg ? await module.handler(context) : await module.handler(event.data, context);
+        const middlewares = Array.isArray(module.config.middleware) ? module.config.middleware : [];
+        const composedMiddleware = (0, middleware_compose_1.composeMiddleware)(...middlewares);
+        const handlerFn = () => {
+            return contextInFirstArg ? module.handler(context) : module.handler(event.data, context);
+        };
+        const result = await composedMiddleware(event.data, context, handlerFn);
         await sender.send('result', result);
         await sender.close();
         process.exit(0);

package/dist/src/python/context.py

@@ -0,0 +1,21 @@
+from typing import Any, List, Optional
+from type_definitions import HandlerResult
+from rpc import RpcSender
+from rpc_state_manager import RpcStateManager
+from logger import Logger
+
+class Context:
+    def __init__(
+        self,
+        trace_id: str,
+        flows: List[str],
+        rpc: RpcSender,
+    ):
+        self.trace_id = trace_id
+        self.flows = flows
+        self.rpc = rpc
+        self.state = RpcStateManager(rpc)
+        self.logger = Logger(self.trace_id, self.flows, rpc)
+
+    async def emit(self, event: Any) -> Optional[HandlerResult]:
+        return await self.rpc.send('emit', event)

package/dist/src/python/get-config.py

@@ -2,7 +2,6 @@ import sys
 import json
 import importlib.util
 import os
-from typing import Any, Callable
 
 # get the FD from ENV
 NODEIPCFD = int(os.environ["NODE_CHANNEL_FD"])
@@ -39,11 +38,14 @@ async def run_python_module(file_path: str) -> None:
         spec.loader.exec_module(module)
 
         # Print the config if it exists
-        if hasattr(module, 'config'):
-            sendMessage(module.config)
-        else:
+        if not hasattr(module, 'config'):
             raise AttributeError(f"No 'config' found in module {file_path}")
 
+        if 'middleware' in module.config:
+            del module.config['middleware']
+
+        sendMessage(module.config)
+
     except Exception as error:
         print('Error running Python module:', str(error), file=sys.stderr)
         sys.exit(1)

package/dist/src/python/logger.py

@@ -18,7 +18,7 @@ class Logger:
         }
 
         if args:
-            # Convert SimpleNamespace to dict if needed
+            # Use our serializer to ensure args are JSON-serializable
             if hasattr(args, '__dict__'):
                 args = vars(args)
             elif not isinstance(args, dict):

package/dist/src/python/middleware.py

@@ -0,0 +1,17 @@
+from typing import Any, Callable
+from functools import reduce
+from context import Context
+
+def compose_middleware(*middlewares):
+    """Compose multiple middleware functions into a single middleware"""
+    def compose_two(f: Callable, g: Callable) -> Callable:
+        async def composed(data: Any, context: Context, next_fn: Callable):
+            async def wrapped_next(d=data):
+                return await g(d, context, next_fn)
+            return await f(data, context, wrapped_next)
+        return composed
+
+    if not middlewares:
+        return lambda data, context, next_fn: next_fn()
+    
+    return reduce(compose_two, middlewares)

package/dist/src/python/python-runner.py

@@ -4,74 +4,115 @@ import importlib.util
 import os
 import asyncio
 import traceback
-from logger import Logger
+from typing import Optional, Any, Callable, List, Dict
+
 from rpc import RpcSender
-from rpc_state_manager import RpcStateManager
-from typing import Any
+from type_definitions import FlowConfig, ApiResponse
+from context import Context
+from validation import validate_with_jsonschema
+from middleware import compose_middleware
 
-def parse_args(arg: str) -> Any:
-    from types import SimpleNamespace
-    
+def parse_args(arg: str) -> Dict:
+    """Parse command line arguments into HandlerArgs"""
     try:
-        return json.loads(arg, object_hook=lambda d: SimpleNamespace(**d))
+        return json.loads(arg)
     except json.JSONDecodeError:
         print('Error parsing args:', arg)
         return arg
 
-class Context:
-    def __init__(self, args: Any, rpc: RpcSender):
-        self.trace_id = args.traceId
-        self.traceId = args.traceId
-        self.flows = args.flows
-        self.rpc = rpc
-        self.state = RpcStateManager(rpc)
-        self.logger = Logger(self.trace_id, self.flows, rpc)
+async def validate_handler_input(
+    module: Any,
+    args: Dict,
+    context: Context,
+    is_api_handler: bool
+) -> Optional[ApiResponse]:
+    """Validate handler input based on module configuration"""
+    if not hasattr(args, "contextInFirstArg") or hasattr(module, "config"):
+        return None
+
+    config: FlowConfig = module.config
+    input_data = args.get("data")
+
+    if is_api_handler and hasattr(input_data, "body") and hasattr(config, "bodySchema"):
+        body = input_data.get("body", {})
+        body_schema = config.get("bodySchema", {})
+        validation_result = validate_with_jsonschema(body, body_schema)
+
+        if not validation_result["success"]:
+            return ApiResponse(
+                status=400,
+                body={"message": f'Input validation error: {validation_result["error"]}'}
+            )
+        input_data.body = validation_result['data']
+    
+    elif not is_api_handler and config.input is not None:
+        validation_result = validate_with_jsonschema(input_data, config.get("input", {}))
+
+        if not validation_result["success"]:
+            context.logger.info(f'Input Validation Error: {validation_result["error"]}', validation_result['details'])
+            return None
+        
+        input_data = validation_result["data"]
 
-    async def emit(self, event: Any):
-        return await self.rpc.send('emit', event)
+    return None
 
-async def run_python_module(file_path: str, rpc: RpcSender, args: Any) -> None:
+async def run_python_module(file_path: str, rpc: RpcSender, args: Dict) -> None:
+    """Execute a Python module with the given arguments"""
     try:
-        # Get the directory containing the module file
         module_dir = os.path.dirname(os.path.abspath(file_path))
-        
-        # Add module directory to Python path
-        if module_dir not in sys.path:
-            sys.path.insert(0, module_dir)
-            
-        # Get the flows directory (parent of steps)
         flows_dir = os.path.dirname(module_dir)
-        if flows_dir not in sys.path:
-            sys.path.insert(0, flows_dir)
-
-        contextInFirstArg = args.contextInFirstArg
+        
+        for path in [module_dir, flows_dir]:
+            if path not in sys.path:
+                sys.path.insert(0, path)
 
-        # Load the module dynamically
         spec = importlib.util.spec_from_file_location("dynamic_module", file_path)
         if spec is None or spec.loader is None:
             raise ImportError(f"Could not load module from {file_path}")
             
         module = importlib.util.module_from_spec(spec)
-        # Add module's directory as its __package__
         module.__package__ = os.path.basename(module_dir)
         spec.loader.exec_module(module)
 
-        # Check if the handler function exists
-        if not hasattr(module, 'handler'):
+        if not hasattr(module, "handler"):
             raise AttributeError(f"Function 'handler' not found in module {file_path}")
 
-        context = Context(args, rpc)
+        config = module.config
+        is_api_handler = (config and config.get("type") == "api")
+
+        trace_id = args.get("traceId")
+        flows = args.get("flows") or []
+        data = args.get("data")
+        context_in_first_arg = args.get("contextInFirstArg")
+        context = Context(trace_id, flows, rpc)
+
+        validation_result = await validate_handler_input(module, args, context, is_api_handler)
+        if validation_result:
+            await rpc.send('result', validation_result)
+            return
+
+    
+        middlewares: List[Callable] = config.get("middleware", [])
+        composed_middleware = compose_middleware(*middlewares)
+        
+        async def handler_fn():
+            if context_in_first_arg:
+                return await module.handler(context)
+            else:
+                return await module.handler(data, context)
+
+        result = await composed_middleware(data, context, handler_fn)
 
-        if contextInFirstArg:
-            result = await module.handler(context)
-        else:
-            result = await module.handler(args.data, context)
+        if not is_api_handler:
+            pending = asyncio.all_tasks() - {asyncio.current_task()}
+            if pending:
+                await asyncio.gather(*pending)
 
         if result:
             await rpc.send('result', result)
 
         rpc.close()
-        rpc.send_no_wait('close', None)
+        rpc.send_no_wait("close", None)
         
     except Exception as error:
         print(f'Error running Python module: {error}', file=sys.stderr)
@@ -80,13 +121,19 @@ async def run_python_module(file_path: str, rpc: RpcSender, args: Any) -> None:
 
 if __name__ == "__main__":
     if len(sys.argv) < 2:
-        print('Usage: python pythonRunner.py <file-path> <arg>', file=sys.stderr)
+        print("Usage: python pythonRunner.py <file-path> <arg>", file=sys.stderr)
         sys.exit(1)
 
     file_path = sys.argv[1]
     arg = sys.argv[2] if len(sys.argv) > 2 else None
 
     rpc = RpcSender()
-    loop = asyncio.get_event_loop()
-    tasks = asyncio.gather(rpc.init(), run_python_module(file_path, rpc, parse_args(arg)))
-    loop.run_until_complete(tasks)
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    args = parse_args(arg) if arg else None
+    tasks = asyncio.gather(rpc.init(), run_python_module(file_path, rpc, args))
+    loop.run_until_complete(tasks)

package/dist/src/python/rpc.py

@@ -8,6 +8,20 @@ from typing import Any, Dict, Tuple
 # get the FD from ENV
 NODEIPCFD = int(os.environ["NODE_CHANNEL_FD"])
 
+def serialize_for_json(obj: Any) -> Any:
+    """Convert Python objects to JSON-serializable types"""
+    if hasattr(obj, '__dict__'):
+        return obj.__dict__
+    elif hasattr(obj, '_asdict'):  # For namedtuples
+        return obj._asdict()
+    elif isinstance(obj, (list, tuple)):
+        return [serialize_for_json(item) for item in obj]
+    elif isinstance(obj, dict):
+        return {k: serialize_for_json(v) for k, v in obj.items()}
+    else:
+        # Try to return the object as is, letting json handle basic types
+        return obj
+
 class RpcSender:
     def __init__(self):
         self.executing = True

package/dist/src/python/rpc_state_manager.py

@@ -1,19 +1,56 @@
 from typing import Any
 import asyncio
 from rpc import RpcSender
+import functools
+import sys
 
 class RpcStateManager:
     def __init__(self, rpc: RpcSender):
         self.rpc = rpc
+        self._loop = asyncio.get_event_loop()
 
     async def get(self, trace_id: str, key: str) -> asyncio.Future[Any]:
-        return await self.rpc.send('state.get', {'traceId': trace_id, 'key': key})
+        result = await self.rpc.send('state.get', {'traceId': trace_id, 'key': key})
+        
+        if result is None:
+            return {'data': None}
+        elif isinstance(result, dict):
+            if 'data' not in result:
+                return {'data': result}
+        
+        return result
 
     async def set(self, trace_id: str, key: str, value: Any) -> asyncio.Future[None]:
-        return await self.rpc.send('state.set', {'traceId': trace_id, 'key': key, 'value': value})
+        future = await self.rpc.send('state.set', {'traceId': trace_id, 'key': key, 'value': value})
+        return future
 
     async def delete(self, trace_id: str, key: str) -> asyncio.Future[None]:
         return await self.rpc.send('state.delete', {'traceId': trace_id, 'key': key})
 
     async def clear(self, trace_id: str) -> asyncio.Future[None]:
         return await self.rpc.send('state.clear', {'traceId': trace_id})
+
+    # Add wrappers to handle non-awaited coroutines
+    def __getattribute__(self, name):
+        attr = super().__getattribute__(name)
+        if name in ('get', 'set', 'delete', 'clear') and asyncio.iscoroutinefunction(attr):
+            @functools.wraps(attr)
+            def wrapper(*args, **kwargs):
+                coro = attr(*args, **kwargs)
+                # Check if this is being awaited
+                frame = sys._getframe(1)
+                if frame.f_code.co_name != '__await__':
+                    # Not being awaited, schedule in background
+                    # But we need to make sure this task completes before the handler returns
+                    # So we'll return the task for the caller to await if needed
+                    task = asyncio.create_task(coro)
+                    # Add error handling for the background task
+                    def handle_exception(t):
+                        if t.done() and not t.cancelled() and t.exception():
+                            print(f"Unhandled exception in background task: {t.exception()}", file=sys.stderr)
+                    task.add_done_callback(handle_exception)
+                    return task
+                # Being awaited, return coroutine as normal
+                return coro
+            return wrapper
+        return attr