diff --git a/llama_stack/distribution/utils/context.py b/llama_stack/distribution/utils/context.py
index 3fcd3315f..24b249890 100644
--- a/llama_stack/distribution/utils/context.py
+++ b/llama_stack/distribution/utils/context.py
@@ -6,12 +6,9 @@
 
 from collections.abc import AsyncGenerator
 from contextvars import ContextVar
-from typing import TypeVar
-
-T = TypeVar("T")
 
 
-def preserve_contexts_async_generator(
+def preserve_contexts_async_generator[T](
     gen: AsyncGenerator[T, None], context_vars: list[ContextVar]
 ) -> AsyncGenerator[T, None]:
     """
diff --git a/llama_stack/providers/utils/telemetry/trace_protocol.py b/llama_stack/providers/utils/telemetry/trace_protocol.py
index eb6d8b331..916f7622a 100644
--- a/llama_stack/providers/utils/telemetry/trace_protocol.py
+++ b/llama_stack/providers/utils/telemetry/trace_protocol.py
@@ -9,14 +9,12 @@ import inspect
 import json
 from collections.abc import AsyncGenerator, Callable
 from functools import wraps
-from typing import Any, TypeVar
+from typing import Any
 
 from pydantic import BaseModel
 
 from llama_stack.models.llama.datatypes import Primitive
 
-T = TypeVar("T")
-
 
 def serialize_value(value: Any) -> Primitive:
     return str(_prepare_for_json(value))
@@ -44,7 +42,7 @@ def _prepare_for_json(value: Any) -> str:
             return str(value)
 
 
-def trace_protocol(cls: type[T]) -> type[T]:
+def trace_protocol[T](cls: type[T]) -> type[T]:
     """
     A class decorator that automatically traces all methods in a protocol/base class
     and its inheriting classes.