fix: bounded SDK retries, runnable langgraph example, MCP client shutdown

examples/langgraph/main.py

@@ -1,7 +1,7 @@
 """Use DMA as a retrieval node in a LangGraph StateGraph.
 
-Install the optional framework dependency first:
-``pip install 'dma-langgraph[langgraph]'``.
+Start the API first with ``make api`` from the repository root, and install the
+optional framework dependency: ``pip install 'dma-langgraph[langgraph]'``.
 """
 
 from __future__ import annotations
@@ -20,22 +20,29 @@ class AgentState(TypedDict):
     response: str
 
 
-memory = DMAClient(api_key=os.environ["DMA_API_KEY"], agent_id="langgraph-agent")
-adapter = DMAMemoryAdapter(memory=memory, query_builder=lambda state: state["user_input"])
-
-
 def respond(state: AgentState) -> dict[str, str]:
     # Replace this deterministic placeholder with your model invocation.
     return {"response": f"Relevant memory:\n{state['dma_context']}"}
 
 
-builder = StateGraph(AgentState)
-builder.add_node("recall_memory", adapter.recall_node)
-builder.add_node("respond", respond)
-builder.add_edge(START, "recall_memory")
-builder.add_edge("recall_memory", "respond")
-builder.add_edge("respond", END)
-graph = builder.compile()
+def main() -> None:
+    with DMAClient(
+        api_key=os.environ["DMA_API_KEY"],
+        agent_id="langgraph-agent",
+        base_url=os.getenv("DMA_BASE_URL", "http://127.0.0.1:8000"),
+    ) as memory:
+        adapter = DMAMemoryAdapter(memory=memory, query_builder=lambda state: state["user_input"])
+        builder = StateGraph(AgentState)
+        builder.add_node("recall_memory", adapter.recall_node)
+        builder.add_node("respond", respond)
+        builder.add_edge(START, "recall_memory")
+        builder.add_edge("recall_memory", "respond")
+        builder.add_edge("respond", END)
+        graph = builder.compile()
+        result = graph.invoke({"user_input": "What backend does the user prefer?"})
+
+    print(result["response"])
+
 
-result = graph.invoke({"user_input": "What backend does the user prefer?"})
-print(result["response"])
+if __name__ == "__main__":
+    main()

packages/dma-mcp/src/dma_mcp/server.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import logging
 import os
+import sys
 
 from dma import DMAClient
 
@@ -35,5 +37,9 @@ def main() -> None:
     api_key = os.environ.get("DMA_API_KEY")
     if not api_key:
         raise SystemExit("DMA_API_KEY must be set")
+    logging.basicConfig(stream=sys.stderr, level=os.getenv("DMA_MCP_LOG_LEVEL", "INFO"))
     client = DMAClient(api_key=api_key, agent_id=os.getenv("DMA_MCP_AGENT_ID", "mcp-agent"), base_url=os.getenv("DMA_BASE_URL", "http://127.0.0.1:8000"))
-    create_server(DMATools(client)).run(transport="stdio")
+    try:
+        create_server(DMATools(client)).run(transport="stdio")
+    finally:
+        client.close()

packages/dma-sdk-python/src/dma/client.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import random
+import time
 from collections.abc import Mapping
 from datetime import datetime
 from typing import Any, Self
@@ -20,13 +22,19 @@
 )
 
 _DEFAULT_BASE_URL = "https://api.dma.dev"
+_DEFAULT_MAX_RETRIES = 2
+_RETRY_BACKOFF_SECONDS = 0.2
 
 
 class DMAClient:
     """A small, typed client for storing and recalling agent memory.
 
     The caller owns the client lifecycle. Use it as a context manager in scripts
     and services to ensure its HTTP connection pool is closed.
+
+    Transport failures are retried up to ``max_retries`` times for requests that
+    are safe to replay: reads, and writes carrying an ``Idempotency-Key``.
+    Deletes are never retried.
     """
 
     def __init__(
@@ -36,6 +44,7 @@ def __init__(
         agent_id: str,
         base_url: str = _DEFAULT_BASE_URL,
         timeout: float = 5.0,
+        max_retries: int = _DEFAULT_MAX_RETRIES,
         transport: httpx.BaseTransport | None = None,
     ) -> None:
         if not api_key.strip():
@@ -44,7 +53,10 @@ def __init__(
             raise ValidationError("agent_id must not be blank")
         if timeout <= 0:
             raise ValidationError("timeout must be greater than zero")
+        if max_retries < 0:
+            raise ValidationError("max_retries must not be negative")
         self._agent_id = agent_id
+        self._max_retries = max_retries
         self._client = httpx.Client(
             base_url=base_url.rstrip("/") + "/",
             timeout=timeout,
@@ -75,7 +87,9 @@ def remember(
         }
         if expires_at is not None:
             payload["expires_at"] = expires_at.isoformat()
-        response = self._request("POST", "v1/memories", json=payload, headers={"Idempotency-Key": key})
+        response = self._request(
+            "POST", "v1/memories", json=payload, headers={"Idempotency-Key": key}, retryable=True
+        )
         return _memory_from_payload(response.json())
 
     def recall(
@@ -92,7 +106,7 @@ def recall(
         payload: dict[str, Any] = {"agent_id": self._agent_id, "query": query, "limit": limit}
         if types is not None:
             payload["types"] = [self._memory_type(memory_type).value for memory_type in types]
-        response = self._request("POST", "v1/memories/recall", json=payload)
+        response = self._request("POST", "v1/memories/recall", json=payload, retryable=True)
         return [_recall_result_from_payload(item) for item in response.json()["results"]]
 
     def list(
@@ -110,7 +124,7 @@ def list(
             params["type"] = self._memory_type(type).value
         if cursor is not None:
             params["cursor"] = cursor
-        response = self._request("GET", "v1/memories", params=params)
+        response = self._request("GET", "v1/memories", params=params, retryable=True)
         payload = response.json()
         return MemoryPage(
             items=[_memory_from_payload(item) for item in payload["items"]],
@@ -130,7 +144,9 @@ def explain(self, memory_id: str, *, query: str | None = None) -> MemoryExplanat
         params: dict[str, Any] = {"agent_id": self._agent_id}
         if query is not None:
             params["query"] = query
-        response = self._request("GET", f"v1/memories/{memory_id}/explanation", params=params)
+        response = self._request(
+            "GET", f"v1/memories/{memory_id}/explanation", params=params, retryable=True
+        )
         payload = response.json()
         retrieval = payload["retrieval"]
         return MemoryExplanation(
@@ -155,18 +171,25 @@ def __enter__(self) -> Self:
     def __exit__(self, *_: object) -> None:
         self.close()
 
-    def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
-        try:
-            response = self._client.request(method, path, **kwargs)
-        except httpx.HTTPError as error:
-            raise DMAConnectionError("unable to reach the DMA API") from error
+    def _request(self, method: str, path: str, *, retryable: bool = False, **kwargs: Any) -> httpx.Response:
+        response = self._send(method, path, self._max_retries if retryable else 0, kwargs)
         if response.status_code == 401:
             raise AuthenticationError(401, "DMA API key was rejected")
         if response.is_error:
             message, code = _error_details(response)
             raise DMAApiError(response.status_code, message, code=code)
         return response
 
+    def _send(self, method: str, path: str, retries: int, kwargs: dict[str, Any]) -> httpx.Response:
+        for attempt in range(retries + 1):
+            try:
+                return self._client.request(method, path, **kwargs)
+            except httpx.HTTPError as error:
+                if attempt == retries:
+                    raise DMAConnectionError("unable to reach the DMA API") from error
+                time.sleep(_RETRY_BACKOFF_SECONDS * 2**attempt * (0.5 + random.random()))
+        raise DMAConnectionError("unable to reach the DMA API")
+
     @staticmethod
     def _validate_non_blank(value: str, field: str) -> None:
         if not value or not value.strip():

packages/dma-sdk-python/tests/test_client.py

@@ -5,7 +5,14 @@
 import httpx
 import pytest
 
-from dma import AuthenticationError, DMAApiError, DMAClient, MemoryType, ValidationError
+from dma import (
+    AuthenticationError,
+    DMAApiError,
+    DMAClient,
+    DMAConnectionError,
+    MemoryType,
+    ValidationError,
+)
 
 
 def _memory_payload() -> dict[str, object]:
@@ -90,6 +97,51 @@ def test_client_surfaces_auth_and_api_errors() -> None:
         api_client.forget("mem_missing")
 
 
+def test_retries_replay_safe_requests_with_the_same_idempotency_key() -> None:
+    seen: list[str] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen.append(request.headers["Idempotency-Key"])
+        if len(seen) == 1:
+            raise httpx.ConnectError("boom")
+        return httpx.Response(201, json=_memory_payload())
+
+    client = DMAClient(
+        api_key="test-key",
+        agent_id="coding-agent",
+        base_url="https://dma.test",
+        max_retries=1,
+        transport=httpx.MockTransport(handler),
+    )
+    memory = client.remember(content="User prefers Java Spring Boot.", type=MemoryType.SEMANTIC)
+    assert memory.id == "mem_abc123"
+    assert seen == [seen[0], seen[0]]
+
+
+def test_delete_is_not_retried_and_retries_are_bounded() -> None:
+    attempts: list[str] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        attempts.append(request.method)
+        raise httpx.ConnectError("boom")
+
+    client = DMAClient(
+        api_key="test-key",
+        agent_id="coding-agent",
+        base_url="https://dma.test",
+        max_retries=2,
+        transport=httpx.MockTransport(handler),
+    )
+    with pytest.raises(DMAConnectionError):
+        client.forget("mem_abc123")
+    assert attempts == ["DELETE"]
+
+    attempts.clear()
+    with pytest.raises(DMAConnectionError):
+        client.recall(query="Java")
+    assert attempts == ["POST", "POST", "POST"]
+
+
 def test_client_validates_inputs_before_requests() -> None:
     with pytest.raises(ValidationError, match="api_key"):
         DMAClient(api_key=" ", agent_id="coding-agent")