fix: storage scoped to runtime id

Author: cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.1.8"
+version = "0.2.0"
 description = "UiPath LlamaIndex SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
@@ -10,7 +10,8 @@ dependencies = [
     "llama-index-embeddings-azure-openai>=0.4.1",
     "llama-index-llms-azure-openai>=0.4.2",
     "openinference-instrumentation-llama-index>=4.3.9",
-    "uipath>=2.2.26, <2.3.0",
+    "uipath==2.3.0.dev1010343497",
+    "uipath-runtime>=0.3.0, <0.4.0",
 ]
 classifiers = [
     "Intended Audience :: Developers",
@@ -58,6 +59,7 @@ dev = [
     "pytest-cov>=4.1.0",
     "pytest-mock>=3.11.1",
     "pre-commit>=4.1.0",
+    "pytest-asyncio>=1.0.0",
     "numpy>=1.24.0",
 ]
 
@@ -98,10 +100,19 @@ disallow_untyped_defs = false
 testpaths = ["tests"]
 python_files = "test_*.py"
 addopts = "-ra -q"
+asyncio_default_fixture_loop_scope = "function"
+asyncio_mode = "auto"
 
 [[tool.uv.index]]
 name = "testpypi"
 url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
 
+[tool.uv.sources]
+uipath = { index = "testpypi" }
+
+[tool.uv]
+override-dependencies = [
+    "uipath>=2.3.0.dev1010340000,<2.3.0.dev1010350000",
+]

src/uipath_llamaindex/runtime/factory.py

@@ -260,6 +260,7 @@ async def _create_runtime_instance(
             delegate=base_runtime,
             storage=storage,
             trigger_manager=trigger_manager,
+            runtime_id=runtime_id,
         )
 
     async def new_runtime(

src/uipath_llamaindex/runtime/storage.py

@@ -3,9 +3,10 @@
 import json
 import os
 import pickle
-from typing import Any
+from typing import Any, cast
 
 import aiosqlite
+from pydantic import BaseModel
 from uipath.core.errors import ErrorCategory, UiPathFaultedTriggerError
 from uipath.runtime import (
     UiPathApiTrigger,
@@ -35,6 +36,9 @@ async def setup(self) -> None:
 
         try:
             async with aiosqlite.connect(self.storage_path) as conn:
+                # Enable WAL mode for better concurrent write handling
+                await conn.execute("PRAGMA journal_mode=WAL")
+
                 # Table for workflow contexts
                 await conn.execute("""
                     CREATE TABLE IF NOT EXISTS workflow_contexts (
@@ -47,26 +51,40 @@ async def setup(self) -> None:
                 await conn.execute("""
                     CREATE TABLE IF NOT EXISTS resume_triggers (
                         id INTEGER PRIMARY KEY AUTOINCREMENT,
+                        runtime_id TEXT NOT NULL,
                         trigger_data TEXT NOT NULL,
                         created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                     )
                 """)
 
+                await conn.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS runtime_kv (
+                        runtime_id TEXT NOT NULL,
+                        namespace TEXT NOT NULL,
+                        key TEXT NOT NULL,
+                        value TEXT,
+                        timestamp DATETIME DEFAULT (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc')),
+                        PRIMARY KEY (runtime_id, namespace, key)
+                    )
+                    """
+                )
+
                 await conn.commit()
         except aiosqlite.Error as exc:
             msg = f"Failed to initialize SQLite storage at {self.storage_path!r}: {exc.sqlite_errorname} {exc.sqlite_errorcode}"
             raise UiPathFaultedTriggerError(ErrorCategory.SYSTEM, msg) from exc
 
-    async def save_trigger(self, trigger: UiPathResumeTrigger) -> None:
+    async def save_trigger(self, runtime_id: str, trigger: UiPathResumeTrigger) -> None:
         """Save resume trigger to SQLite database."""
         trigger_dict = self._serialize_trigger(trigger)
         trigger_json = json.dumps(trigger_dict)
 
         try:
             async with aiosqlite.connect(self.storage_path) as conn:
                 await conn.execute(
-                    "INSERT INTO resume_triggers (trigger_data) VALUES (?)",
-                    (trigger_json,),
+                    "INSERT INTO resume_triggers (runtime_id, trigger_data) VALUES (?, ?)",
+                    (runtime_id, trigger_json),
                 )
                 await conn.commit()
         except aiosqlite.Error as exc:
@@ -78,12 +96,13 @@ async def save_trigger(self, trigger: UiPathResumeTrigger) -> None:
             )
             raise UiPathFaultedTriggerError(ErrorCategory.SYSTEM, msg) from exc
 
-    async def get_latest_trigger(self) -> UiPathResumeTrigger | None:
+    async def get_latest_trigger(self, runtime_id: str) -> UiPathResumeTrigger | None:
         """Get most recent trigger from SQLite database."""
         try:
             async with aiosqlite.connect(self.storage_path) as conn:
                 cursor = await conn.execute(
-                    "SELECT trigger_data FROM resume_triggers ORDER BY created_at DESC LIMIT 1"
+                    "SELECT trigger_data FROM resume_triggers WHERE runtime_id = ? ORDER BY id DESC LIMIT 1",
+                    (runtime_id,),
                 )
                 row = await cursor.fetchone()
         except aiosqlite.Error as exc:
@@ -148,6 +167,58 @@ async def load_context(self, runtime_id: str) -> dict[str, Any] | None:
 
         return pickle.loads(row[0])
 
+    async def set_value(
+        self,
+        runtime_id: str,
+        namespace: str,
+        key: str,
+        value: Any,
+    ) -> None:
+        """Save arbitrary key-value pair to database."""
+        if not (
+            isinstance(value, str)
+            or isinstance(value, dict)
+            or isinstance(value, BaseModel)
+            or value is None
+        ):
+            raise TypeError("Value must be str, dict, BaseModel or None.")
+
+        value_text = self._dump_value(value)
+
+        async with aiosqlite.connect(self.storage_path) as conn:
+            await conn.execute(
+                """
+                INSERT INTO runtime_kv (runtime_id, namespace, key, value)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(runtime_id, namespace, key)
+                DO UPDATE SET
+                    value = excluded.value,
+                    timestamp = (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc'))
+                """,
+                (runtime_id, namespace, key, value_text),
+            )
+            await conn.commit()
+
+    async def get_value(self, runtime_id: str, namespace: str, key: str) -> Any:
+        """Get arbitrary key-value pair from database (scoped by runtime_id + namespace)."""
+
+        async with aiosqlite.connect(self.storage_path) as conn:
+            cur = await conn.execute(
+                """
+                SELECT value
+                FROM runtime_kv
+                WHERE runtime_id = ? AND namespace = ? AND key = ?
+                LIMIT 1
+                """,
+                (runtime_id, namespace, key),
+            )
+            row = await cur.fetchone()
+
+        if not row:
+            return None
+
+        return self._load_value(cast(str | None, row[0]))
+
     def _serialize_trigger(self, trigger: UiPathResumeTrigger) -> dict[str, Any]:
         """Serialize a resume trigger to a dictionary."""
         trigger_key = (
@@ -194,3 +265,21 @@ def _deserialize_trigger(self, trigger_data: dict[str, Any]) -> UiPathResumeTrig
             )
 
         return resume_trigger
+
+    def _dump_value(self, value: str | dict[str, Any] | BaseModel | None) -> str | None:
+        if value is None:
+            return None
+        if isinstance(value, BaseModel):
+            return "j:" + json.dumps(value.model_dump())
+        if isinstance(value, dict):
+            return "j:" + json.dumps(value)
+        return "s:" + value
+
+    def _load_value(self, raw: str | None) -> Any:
+        if raw is None:
+            return None
+        if raw.startswith("s:"):
+            return raw[2:]
+        if raw.startswith("j:"):
+            return json.loads(raw[2:])
+        return raw