diff --git a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream/api.py b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream/api.py
index b3aaeec65..9fe8fa760 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream/api.py
+++ b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream/api.py
@@ -67,7 +67,7 @@ class AgentContext(ABC):
     """The Agent context interface"""
 
     @abstractmethod
-    def get_persistent_state_directory(self):
+    def get_persistent_state_directory(self) -> Optional[str]:
         """Return the path of the agent disk. Return None if not configured."""
         pass
 
diff --git a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/api.py b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/api.py
index da41c151f..855a5cec0 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/api.py
+++ b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/api.py
@@ -67,9 +67,9 @@ class AgentContext(ABC):
     """The Agent context interface"""
 
     @abstractmethod
-    def get_persistent_state_directory(self):
+    def get_persistent_state_directory(self) -> Optional[str]:
         """Return the path of the agent disk. Return None if not configured."""
-        return None
+        pass
 
 
 class Agent(ABC):
diff --git a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py
index a12d5772d..9fca983a1 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py
+++ b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py
@@ -383,11 +383,8 @@ def __init__(self, configuration: dict, context: dict):
         self.configuration = configuration
         self.context = context
 
-    def get_persistent_state_directory(self) -> str:
-        dir = self.context.get("persistentStateDirectory", "")
-        if not dir:
-            return None
-        return dir
+    def get_persistent_state_directory(self) -> Optional[str]:
+        return self.context.get("persistentStateDirectory")
 
 
 class AgentServer(object):
diff --git a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py
index 25ae85b93..c76c26818 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py
+++ b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py
@@ -281,23 +281,27 @@ def process(self, record: Record) -> Future[List[RecordType]]:
         return self.executor.submit(lambda r: [r], record)
 
 
-class ProcessorInitOneParameter(Processor):
+class ProcessorInitOneParameter:
+    def __init__(self):
+        self.myparam = None
+
     def init(self, agent_config):
         self.myparam = agent_config["my-param"]
 
-    def process(self, record: Record) -> List[RecordType]:
+    def process(self, _) -> List[RecordType]:
         return [{"value": self.myparam}]
 
 
 class ProcessorUseContext(Processor):
+    def __init__(self):
+        self.myparam = None
+        self.context = None
+
     def init(self, agent_config, context: AgentContext):
         self.myparam = agent_config["my-param"]
         self.context = context
 
     def process(self, record: Record) -> List[RecordType]:
         return [
-            {
-                "value": "directory is "
-                + str(self.context.get_persistent_state_directory())
-            }
+            {"value": f"directory is {self.context.get_persistent_state_directory()}"}
         ]