refactor: Add ProviderContext for a flexible storage directory

- Introduce ProviderContext class to decouple provider storage paths from absolute paths
- Add storage_dir attribute to StackRunConfig to accept CLI options
- Implement storage directory resolution with prioritized fallbacks:
  1. CLI option (--state-directory)
  2. Environment variable (LLAMA_STACK_STATE_DIR)
  3. Default distribution directory
- Standardize provider signatures to follow context, config, deps pattern
- Update provider implementations to use the new context-based approach
- Add comprehensive tests to verify state directory resolution

tests/unit/providers/agent/test_meta_reference_agent.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 from datetime import datetime
+from pathlib import Path
 from unittest.mock import AsyncMock
 
 import pytest
@@ -20,6 +21,7 @@ from llama_stack.apis.inference import Inference
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
+from llama_stack.providers.datatypes import ProviderContext
 from llama_stack.providers.inline.agents.meta_reference.agents import MetaReferenceAgentsImpl
 from llama_stack.providers.inline.agents.meta_reference.config import MetaReferenceAgentsImplConfig
 from llama_stack.providers.inline.agents.meta_reference.persistence import AgentInfo
@@ -48,7 +50,9 @@ def config(tmp_path):
 
 @pytest_asyncio.fixture
 async def agents_impl(config, mock_apis):
+    context = ProviderContext(storage_dir=Path("/tmp"))
     impl = MetaReferenceAgentsImpl(
+        context,
         config,
         mock_apis["inference_api"],
         mock_apis["vector_io_api"],

tests/unit/providers/inference/test_remote_vllm.py

@@ -62,9 +62,13 @@ class MockInferenceAdapterWithSleep:
             # ruff: noqa: N802
             def do_POST(self):
                 time.sleep(sleep_time)
+                response_json = json.dumps(response).encode("utf-8")
                 self.send_response(code=200)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(response_json)))
                 self.end_headers()
-                self.wfile.write(json.dumps(response).encode("utf-8"))
+                self.wfile.write(response_json)
+                self.wfile.flush()
 
         self.request_handler = DelayedRequestHandler
 

tests/unit/test_state_dir_resolution.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from pathlib import Path
+
+from llama_stack.distribution.resolver import resolve_storage_dir
+
+
+class DummyConfig:
+    pass
+
+
+def test_storage_dir_cli(monkeypatch):
+    config = DummyConfig()
+    config.storage_dir = "/cli/dir"
+    monkeypatch.delenv("LLAMA_STACK_STORAGE_DIR", raising=False)
+    result = resolve_storage_dir(config, "distro")
+    assert result == Path("/cli/dir")
+
+
+def test_storage_dir_env(monkeypatch):
+    config = DummyConfig()
+    if hasattr(config, "storage_dir"):
+        delattr(config, "storage_dir")
+    monkeypatch.setenv("LLAMA_STACK_STORAGE_DIR", "/env/dir")
+    result = resolve_storage_dir(config, "distro")
+    assert result == Path("/env/dir")
+
+
+def test_storage_dir_fallback(monkeypatch):
+    # Mock the DISTRIBS_BASE_DIR
+    monkeypatch.setattr("llama_stack.distribution.utils.config_dirs.DISTRIBS_BASE_DIR", Path("/mock/distribs"))
+
+    config = DummyConfig()
+    if hasattr(config, "storage_dir"):
+        delattr(config, "storage_dir")
+    monkeypatch.delenv("LLAMA_STACK_STORAGE_DIR", raising=False)
+
+    result = resolve_storage_dir(config, "distro")
+    assert result == Path("/mock/distribs/distro")