@@ -14,12 +14,17 @@
 DEFAULT_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 DEFAULT_OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 DEFAULT_VERCEL_API_KEY = os.getenv("AI_GATEWAY_API_KEY")
+DEFAULT_PRIME_API_KEY = os.getenv("PRIME_API_KEY")
 DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/"
 
 
 class OpenAIClient(BaseLM):
     """
     LM Client for running models with the OpenAI API. Works with vLLM as well.
+
+    Any additional keyword arguments (e.g. default_headers, default_query, max_retries)
+    are passed through to the underlying openai.OpenAI and openai.AsyncOpenAI constructors.
+    Only model_name is excluded, since it is not a client constructor argument.
     """
 
     def __init__(
@@ -38,12 +43,19 @@ def __init__(
                 api_key = DEFAULT_OPENROUTER_API_KEY
             elif base_url == "https://ai-gateway.vercel.sh/v1":
                 api_key = DEFAULT_VERCEL_API_KEY
-
-        # For vLLM, set base_url to local vLLM server address.
-        self.client = openai.OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout)
-        self.async_client = openai.AsyncOpenAI(
-            api_key=api_key, base_url=base_url, timeout=self.timeout
-        )
+            elif base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
+                api_key = DEFAULT_PRIME_API_KEY
+
+        # Pass through arbitrary kwargs to the OpenAI client (e.g. default_headers, default_query, max_retries).
+        # Exclude model_name since it is not an OpenAI client constructor argument.
+        client_kwargs = {
+            "api_key": api_key,
+            "base_url": base_url,
+            "timeout": self.timeout,
+            **{k: v for k, v in self.kwargs.items() if k != "model_name"},
+        }
+        self.client = openai.OpenAI(**client_kwargs)
+        self.async_client = openai.AsyncOpenAI(**client_kwargs)
         self.model_name = model_name
         self.base_url = base_url  # Track for cost extraction
 
 
@@ -85,8 +85,14 @@ def _handle_batched(self, request: LMRequest, handler: "LMHandler") -> LMRespons
 
         start_time = time.perf_counter()
 
+        sem = asyncio.Semaphore(handler.batch_max_concurrent)
+
+        async def run_one(prompt: str):
+            async with sem:
+                return await client.acompletion(prompt)
+
         async def run_all():
-            tasks = [client.acompletion(prompt) for prompt in request.prompts]
+            tasks = [run_one(prompt) for prompt in request.prompts]
             return await asyncio.gather(*tasks)
 
         results = asyncio.run(run_all())
@@ -132,6 +138,7 @@ def __init__(
         host: str = "127.0.0.1",
         port: int = 0,  # auto-assign available port
         other_backend_client: BaseLM | None = None,
+        batch_max_concurrent: int = 16,
     ):
         self.default_client = client
         self.other_backend_client = other_backend_client
@@ -140,6 +147,7 @@ def __init__(
         self._server: ThreadingLMServer | None = None
         self._thread: Thread | None = None
         self._port = port
+        self.batch_max_concurrent = batch_max_concurrent
 
         self.register_client(client.model_name, client)
 
 
@@ -1,27 +1,61 @@
+"""Mock LM for testing without a real model."""
+
+from collections.abc import Callable
+from typing import Any
+
 from rlm.clients.base_lm import BaseLM
 from rlm.core.types import ModelUsageSummary, UsageSummary
 
 
 class MockLM(BaseLM):
-    """Simple mock LM that echoes prompts."""
+    """
+    In-memory mock LM that implements BaseLM for testing.
 
-    def __init__(self):
-        super().__init__(model_name="mock-model")
+    - completion() / acompletion(): echo-style response by default, or use
+      optional responses list / response_fn if provided.
+    - get_usage_summary() / get_last_usage(): return usage suitable for tests.
+    """
 
-    def completion(self, prompt):
-        return f"Mock response to: {prompt[:50]}"
+    def __init__(
+        self,
+        model_name: str = "mock-model",
+        responses: list[str] | None = None,
+        response_fn: Callable[[str | dict[str, Any]], str] | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self._responses = list(responses) if responses is not None else None
+        self._response_fn = response_fn
+        self._call_count = 0
 
-    async def acompletion(self, prompt):
+    def completion(self, prompt: str | dict[str, Any]) -> str:
+        self._call_count += 1
+        if self._responses is not None:
+            if not self._responses:
+                raise IndexError("MockLM: no more responses in list")
+            return self._responses.pop(0)
+        if self._response_fn is not None:
+            return self._response_fn(prompt)
+        prompt_str = prompt if isinstance(prompt, str) else str(prompt)[:80]
+        return f"Mock response to: {prompt_str}"
+
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
         return self.completion(prompt)
 
-    def get_usage_summary(self):
+    def get_usage_summary(self) -> UsageSummary:
         return UsageSummary(
             model_usage_summaries={
-                "mock-model": ModelUsageSummary(
-                    total_calls=1, total_input_tokens=10, total_output_tokens=10
+                self.model_name: ModelUsageSummary(
+                    total_calls=self._call_count,
+                    total_input_tokens=self._call_count * 10,
+                    total_output_tokens=self._call_count * 10,
                 )
             }
         )
 
-    def get_last_usage(self):
-        return self.get_usage_summary()
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=10,
+            total_output_tokens=10,
+        )
@@ -0,0 +1,45 @@
+"""Tests for LMHandler using MockLM (no real LM required)."""
+
+from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
+from rlm.core.lm_handler import LMHandler
+from tests.mock_lm import MockLM
+
+
+def test_lm_handler_single_request():
+    """Single prompt request returns success and echo-style content."""
+    mock = MockLM(responses=["hello back"])
+    with LMHandler(client=mock) as handler:
+        request = LMRequest(prompt="hello")
+        response = send_lm_request(handler.address, request)
+    assert response.success
+    assert response.chat_completion is not None
+    assert response.chat_completion.response == "hello back"
+
+
+def test_lm_handler_batched_request():
+    """Batched prompts return one response per prompt in order."""
+    responses = [f"r{i}" for i in range(5)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=3) as handler:
+        prompts = [f"prompt-{i}" for i in range(5)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == 5
+    for i, resp in enumerate(result):
+        assert resp.success, resp.error
+        assert resp.chat_completion is not None
+        assert resp.chat_completion.response == f"r{i}"
+
+
+def test_lm_handler_batched_many_prompts_semaphore_cap():
+    """Many prompts complete successfully with semaphore limiting concurrency."""
+    # 50 prompts, max 4 concurrent: should still all complete
+    count = 50
+    responses = [f"resp-{i}" for i in range(count)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=4) as handler:
+        prompts = [f"p-{i}" for i in range(count)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == count
+    for i, resp in enumerate(result):
+        assert resp.success, (i, resp.error)
+        assert resp.chat_completion.response == f"resp-{i}"