RPC Rework: Part 5: Add ModuleProxy type

[jeff-hykin]

No description provided.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR continues the RPC rework series by adding the ModuleProxy type hint and refactoring blueprint internals to support inter-module dependencies through Spec protocols.

Key Changes

• Added ModuleProxy type: A TYPE_CHECKING-only class that helps type checkers understand that RPCClient instances have all Module methods available
• Renamed ModuleBlueprintSet to Blueprint: Cleaner, more concise naming throughout the codebase
• Renamed ModuleConnection to StreamRef: Better reflects that it represents a reference to a stream
• Added ModuleRef tracking: New _BlueprintAtom.create() logic detects inter-module dependencies by examining type annotations for Spec protocols and direct Module references
• New dimos.spec.utils module: Provides utilities for working with Spec protocols, including structural/annotation compliance checking and method signature extraction
• Comprehensive updates: All imports, tests, and documentation updated to reflect the new naming

Critical Issues Found

Two logic bugs were identified that will cause runtime failures:

1. Line 98 of dimos/core/blueprints.py: Attempts to access .rpc_calls on a class instead of using .rpcs.keys()
2. Line 426 of dimos/core/blueprints.py: model_copy() result not assigned, so rerun_server_addr is never actually updated

Type Safety Improvements

The addition of ModuleProxy significantly improves the development experience by giving IDE autocomplete and type checkers proper visibility into the dynamic RPC methods available on deployed modules.

Confidence Score: 2/5

• This PR contains two critical runtime bugs that must be fixed before merging
• While the refactoring is well-structured with comprehensive test and documentation updates, two logic errors in blueprints.py will cause immediate runtime failures when the new ModuleRef tracking code is exercised or when Rerun initialization occurs
• Pay close attention to dimos/core/blueprints.py - fix the bugs on lines 98 and 426 before merging

Important Files Changed

Filename	Overview
dimos/core/blueprints.py	Major refactoring: renamed ModuleBlueprintSet to Blueprint, added ModuleRef tracking for inter-module dependencies, but contains a critical attribute access bug on line 98
dimos/core/rpc_client.py	Added ModuleProxy type hint class to help type checkers understand that RPCClient instances have Module methods
dimos/core/init.py	Added type annotations to deploy function, properly typed return as ModuleProxy using TYPE_CHECKING import
dimos/core/test_blueprints.py	Updated all test cases to use new Blueprint and StreamRef naming, tests properly updated to include module_refs field
dimos/spec/utils.py	New utility file for Spec protocol support, includes structural and annotation compliance checking

Sequence Diagram

sequenceDiagram
    participant Client as Client Code
    participant MC as ModuleCoordinator
    participant Dask as DimosCluster (Dask)
    participant RPC as RPCClient
    participant Module as Module Instance

    Client->>MC: deploy(ModuleClass, args, kwargs)
    MC->>Dask: deploy(ModuleClass, args, kwargs)
    
    Note over Dask: Submits module construction<br/>to dask worker
    
    Dask->>Module: ModuleClass(*args, **kwargs)
    activate Module
    Module-->>Dask: module instance
    deactivate Module
    
    Dask->>RPC: RPCClient(actor, ModuleClass)
    activate RPC
    Note over RPC: Creates RPC client with<br/>LCMRPC transport
    RPC-->>Dask: rpc_client instance
    deactivate RPC
    
    Dask-->>MC: cast to ModuleProxy
    Note over MC: ModuleProxy is TYPE_CHECKING only<br/>Tells type checker RPCClient has<br/>Module methods
    MC-->>Client: ModuleProxy (actually RPCClient)
    
    Note over Client,Module: Module method calls via proxy
    
    Client->>RPC: proxy.some_rpc_method()
    Note over RPC: __getattr__ intercepts call<br/>Creates RpcCall wrapper
    RPC->>Module: LCMRPC.call_sync(method, args)
    Module-->>RPC: result
    RPC-->>Client: result

Loading

[greptile-apps]

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

rpc_calls is a list attribute on Module instances, not a class attribute. This line will fail at runtime when isinstance(getattr(module, name, None), Module) evaluates to True because module here is a class, not an instance.

Suggested change

	module_refs.append(ModuleRef(name=name, rpc_method_names=other_module.rpc_calls))
	module_refs.append(ModuleRef(name=name, rpc_method_names=tuple(module.rpcs.keys())))

[greptile-apps]

model_copy returns a new instance but the result is not assigned to anything. This means rerun_server_addr is never actually updated in the global config.

Suggested change

	global_config.model_copy(update={"rerun_server_addr": server_addr})
	global_config = global_config.model_copy(update={"rerun_server_addr": server_addr})