feat: integrated with sandbox

backend/src/sandbox/middleware.py

@@ -0,0 +1,33 @@
+from typing import override
+
+from langchain.agents import AgentState
+from langchain.agents.middleware import AgentMiddleware
+from langgraph.runtime import Runtime
+
+from src.agents.thread_state import SandboxState
+from src.sandbox import get_sandbox_provider
+
+
+class SandboxMiddlewareState(AgentState):
+    """Compatible with the `ThreadState` schema."""
+
+    sandbox: SandboxState | None = None
+
+
+class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
+    """Create a sandbox environment and assign it to an agent."""
+
+    state_schema = SandboxMiddlewareState
+
+    def _acquire_sandbox(self) -> str:
+        provider = get_sandbox_provider()
+        sandbox_id = provider.acquire()
+        print(f"Acquiring sandbox {sandbox_id}")
+        return sandbox_id
+
+    @override
+    def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
+        if "sandbox" not in state or state["sandbox"] is None:
+            sandbox_id = self._acquire_sandbox()
+            return {"sandbox": {"sandbox_id": sandbox_id}}
+        return super().before_agent(state, runtime)