From fc4d40fb1e138b0823f1c8b39c2a21bb6733cc28 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaan@berri.ai>
Date: Wed, 6 May 2026 14:52:21 -0700
Subject: [PATCH] feat(agent_session_endpoints): add VM provider registry

---
 .../vm_providers/registry.py                  | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)
 create mode 100644 litellm/proxy/agent_session_endpoints/vm_providers/registry.py

diff --git a/litellm/proxy/agent_session_endpoints/vm_providers/registry.py b/litellm/proxy/agent_session_endpoints/vm_providers/registry.py
new file mode 100644
index 00000000000..d878ef25eeb
--- /dev/null
+++ b/litellm/proxy/agent_session_endpoints/vm_providers/registry.py
@@ -0,0 +1,33 @@
+"""Process-wide registry of VM providers, keyed by provider name."""
+
+from typing import Dict
+
+from litellm.proxy.agent_session_endpoints.vm_providers.base import AgentVMProvider
+from litellm.proxy.agent_session_endpoints.vm_providers.noop import NoopVMProvider
+
+_REGISTRY: Dict[str, AgentVMProvider] = {}
+
+
+def register_vm_provider(provider: AgentVMProvider) -> None:
+    """Register a provider by ``provider.name``. Last-write-wins; tests
+    use this to swap in a fresh ``NoopVMProvider`` between cases."""
+    _REGISTRY[provider.name] = provider
+
+
+def get_vm_provider(name: str) -> AgentVMProvider:
+    """Return the registered provider for ``name``.
+
+    Lazily instantiates a default ``NoopVMProvider`` on first access so
+    tests don't need a setup hook just to use the noop.
+    """
+    if name not in _REGISTRY:
+        if name == "noop":
+            _REGISTRY[name] = NoopVMProvider()
+        else:
+            raise KeyError(f"No VM provider registered for '{name}'")
+    return _REGISTRY[name]
+
+
+def reset_vm_provider_registry() -> None:
+    """Test helper: drop all registered providers."""
+    _REGISTRY.clear()