Tier 1 fixes for #1954 (variable composition review feedback)

logfire/_internal/main.py

@@ -2622,6 +2622,17 @@ def var(
 
         return variable
 
+    @overload
+    def template_var(
+        self,
+        name: str,
+        *,
+        default: T,
+        inputs_type: type[InputsT],
+        description: str | None = None,
+    ) -> TemplateVariable[T, InputsT]: ...
+
+    @overload
     def template_var(
         self,
         name: str,
@@ -2630,6 +2641,16 @@ def template_var(
         default: T | ResolveFunction[T],
         inputs_type: type[InputsT],
         description: str | None = None,
+    ) -> TemplateVariable[T, InputsT]: ...
+
+    def template_var(
+        self,
+        name: str,
+        *,
+        type: type[T] | None = None,
+        default: T | ResolveFunction[T],
+        inputs_type: type[InputsT],
+        description: str | None = None,
     ) -> TemplateVariable[T, InputsT]:
         """Define a managed template variable with integrated rendering.
 
@@ -2652,7 +2673,6 @@ class PromptInputs(BaseModel):
 
         prompt = logfire.template_var(
             'system_prompt',
-            type=str,
             default='Hello {{user_name}}',
             inputs_type=PromptInputs,
         )
@@ -2663,16 +2683,30 @@ class PromptInputs(BaseModel):
 
         Args:
             name: Unique identifier for the variable.
-            type: Expected type for validation and JSON schema generation.
+            type: Expected type for validation and JSON schema generation. Can be a primitive
+                type or a Pydantic model. If not provided, the type is inferred from `default`.
+                Required when `default` is a resolve function.
             default: Default value used when no remote configuration is found.
-                Can also be a callable with `targeting_key` and `attributes` parameters.
+                When `type` is not provided, the type is inferred from this value.
+                Can also be a callable with `targeting_key` and `attributes` parameters
+                (requires `type` to be set explicitly).
             inputs_type: The type (typically a Pydantic `BaseModel`) describing the expected
                 template inputs. Used for type-safe `get(inputs)` calls and JSON schema generation.
             description: Optional human-readable description of what the variable controls.
         """
         import re
 
-        from logfire.variables.variable import TemplateVariable
+        from logfire.variables.variable import TemplateVariable, is_resolve_function
+
+        if type is None:
+            if is_resolve_function(default):
+                raise TypeError(
+                    'When `default` is a resolve function (callable with targeting_key and attributes parameters), '
+                    '`type` must be provided to specify the variable value type.'
+                )
+            tp = cast(Type[T], default.__class__)  # noqa UP006
+        else:
+            tp = type
 
         if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', name):
             raise ValueError(
@@ -2692,7 +2726,7 @@ class PromptInputs(BaseModel):
 
         variable = TemplateVariable[T, InputsT](
             name,
-            type=type,
+            type=tp,
             default=default,
             inputs_type=inputs_type,
             description=description,

logfire/variables/abstract.py

@@ -536,49 +536,82 @@ def _check_reference_errors(
 ) -> list[str]:
     """Check for reference errors: non-existent refs and cycles.
 
-    Scans local variable defaults and server label values for @{references}@
-    and validates that referenced variables exist and there are no cycles.
+    Walks the full composition graph starting from each locally-registered
+    variable, transitively following `@{ref}@` edges into server-only
+    variables — so a missing reference reachable only through a chain that
+    passes through a server-only node still surfaces, as does a cycle whose
+    midpoints are server-only.
+
+    `VariablesConfig` is treated as self-contained for substitution: any
+    `@{name}@` whose `name` isn't in either the local registration set or
+    `server_config` is reported as a non-existent reference, the same way
+    a registration miss is.
     """
     from logfire.variables.composition import find_references
     from logfire.variables.config import LabeledValue
     from logfire.variables.variable import is_resolve_function
 
     warnings_list: list[str] = []
 
-    # Collect all known variable names (local + server)
     all_names: set[str] = {v.name for v in variables} | set(server_config.variables.keys())
+    locals_by_name = {v.name: v for v in variables}
 
-    # Build a reference graph: variable_name -> set of referenced names
-    ref_graph: dict[str, set[str]] = {}
+    def _refs_of(name: str) -> set[str]:
+        """Collect refs from every serialized value reachable for *name*.
 
-    # Scan local variable defaults for references
-    for variable in variables:
+        That's the local code default (if registered locally) plus every
+        labeled server value plus the `latest_version`. Failures to
+        serialize the local default are tolerated — we want the walker to
+        keep going.
+        """
         refs: set[str] = set()
-        if not is_resolve_function(variable.default):
+        local = locals_by_name.get(name)
+        if local is not None and not is_resolve_function(local.default):
             try:
-                serialized_default = variable.type_adapter.dump_json(variable.default).decode('utf-8')
+                serialized_default = local.type_adapter.dump_json(local.default).decode('utf-8')
                 refs.update(find_references(serialized_default))
             except Exception:
                 pass
-
-        # Also scan server label values for this variable
-        server_var = server_config.variables.get(variable.name)
+        server_var = server_config.variables.get(name)
         if server_var is not None:
-            for _, labeled_value in server_var.labels.items():
-                if isinstance(labeled_value, LabeledValue):
-                    refs.update(find_references(labeled_value.serialized_value))
+            for labeled in server_var.labels.values():
+                if isinstance(labeled, LabeledValue):
+                    refs.update(find_references(labeled.serialized_value))
             if server_var.latest_version is not None:
                 refs.update(find_references(server_var.latest_version.serialized_value))
+        return refs
 
-        if refs:
-            ref_graph[variable.name] = refs
-
-        # Check for non-existent references
-        for ref_name in refs:
-            if ref_name not in all_names:
-                warnings_list.append(f"Variable '{variable.name}' references '@{{{ref_name}}}@' which does not exist.")
+    # BFS the composition graph from every local variable in declaration
+    # order. Each node we visit contributes its outgoing edges to
+    # `ref_graph` and, if any point at an unknown name, a
+    # non-existent-reference warning. Visited names are gated on `seen` so
+    # a shared sub-tree is walked once.
+    from collections import deque
 
-    # Check for cycles using DFS
+    ref_graph: dict[str, set[str]] = {}
+    seen: set[str] = set()
+    missing_reported: set[tuple[str, str]] = set()
+    frontier: deque[str] = deque(v.name for v in variables)
+    while frontier:
+        current = frontier.popleft()
+        if current in seen:
+            continue
+        seen.add(current)
+        refs = _refs_of(current)
+        if refs:
+            ref_graph[current] = refs
+        for ref in refs:
+            if ref not in all_names:
+                key = (current, ref)
+                if key not in missing_reported:
+                    missing_reported.add(key)
+                    warnings_list.append(f"Variable '{current}' references '@{{{ref}}}@' which does not exist.")
+            elif ref not in seen:
+                frontier.append(ref)
+
+    # Cycle detection on the assembled graph. Because the graph includes
+    # nodes reached transitively through server-only variables, cycles
+    # whose midpoints are server-only are now caught too.
     def _detect_cycles(graph: dict[str, set[str]]) -> list[list[str]]:
         cycles: list[list[str]] = []
         visited: set[str] = set()
@@ -587,7 +620,6 @@ def _detect_cycles(graph: dict[str, set[str]]) -> list[list[str]]:
 
         def dfs(node: str) -> None:
             if node in in_stack:
-                # Found a cycle - extract it
                 cycle_start = path.index(node)
                 cycles.append(path[cycle_start:] + [node])
                 return

logfire/variables/variable.py

@@ -19,7 +19,6 @@
     ComposedReference,
     VariableCompositionError,
     expand_references,
-    has_references,
 )
 
 if TYPE_CHECKING:
@@ -211,6 +210,21 @@ def _resolve(
     ) -> ResolvedVariable[T_co]:
         serialized_result: ResolvedVariable[str | None] | None = None
         try:
+            # Top-level context-override fast path: handled here, before
+            # `_lookup_serialized` even sees the name. Overrides do not
+            # participate in `@{ref}@` composition (their value is the user's
+            # literal choice), and the round-trip through dump_json /
+            # validate_json that `_lookup_serialized` would otherwise perform
+            # silently drops any value that isn't JSON-serializable. Restore
+            # the pre-#1951 behaviour: if the override serializes, take the
+            # render path; if it doesn't, return the typed Python value
+            # verbatim.
+            context_overrides = _VARIABLE_OVERRIDES.get()
+            if context_overrides is not None and self.name in context_overrides:
+                return self._resolve_context_override(
+                    context_overrides[self.name], targeting_key, attributes, render_fn
+                )
+
             provider = self.logfire_instance.config.get_variable_provider()
 
             serialized_result = self._lookup_serialized(
@@ -221,16 +235,6 @@ def _resolve(
                 label=label,
             )
 
-            # Context overrides skip composition expansion: the override is the
-            # user's literal choice, so we only optionally render (for
-            # TemplateVariable) and then deserialize.
-            if serialized_result.reason == 'context_override' and serialized_result.value is not None:
-                serialized = serialized_result.value
-                if render_fn is not None:
-                    serialized = render_fn(serialized)
-                value = self.type_adapter.validate_json(serialized)
-                return ResolvedVariable(name=self.name, value=value, reason='context_override')
-
             if serialized_result.value is None:
                 return self._resolve_code_default(
                     targeting_key,
@@ -261,6 +265,38 @@ def _resolve(
                 default = cast('T_co', None)
             return ResolvedVariable(name=self.name, value=default, exception=e, reason='other_error')
 
+    def _resolve_context_override(
+        self,
+        override_value: T_co | ResolveFunction[T_co],
+        targeting_key: str | None,
+        attributes: Mapping[str, Any] | None,
+        render_fn: Callable[[str], str] | None,
+    ) -> ResolvedVariable[T_co]:
+        """Return a resolution for the top-level context override.
+
+        Overrides do not participate in composition. When the override value
+        serializes cleanly, run any provided `render_fn` (template rendering)
+        against the JSON form and revalidate so the user gets the same shape
+        a provider value would yield. When it doesn't serialize — common for
+        custom Python types, arbitrary objects, etc. — return the user's
+        value verbatim under `reason='context_override'`. Returning verbatim
+        is the legacy behaviour Devin / Alex flagged on #1951; the previous
+        implementation silently dropped these values back to the provider /
+        code default.
+        """
+        if is_resolve_function(override_value):
+            resolved_value = cast('T_co', override_value(targeting_key, attributes))
+        else:
+            resolved_value = cast('T_co', override_value)
+        try:
+            serialized = self.type_adapter.dump_json(resolved_value).decode('utf-8')
+        except (ValueError, TypeError, RuntimeError):
+            return ResolvedVariable(name=self.name, value=resolved_value, reason='context_override')
+        if render_fn is not None:
+            serialized = render_fn(serialized)
+        validated = self.type_adapter.validate_json(serialized)
+        return ResolvedVariable(name=self.name, value=validated, reason='context_override')
+
     def _lookup_serialized(
         self,
         name: str,
@@ -339,45 +375,48 @@ def _expand_and_deserialize(
         serialized_value = serialized_result.value
         composed: list[ComposedReference] = []
 
-        # Expand @{references}@ if any are present
-        if has_references(serialized_value):
-
-            def resolve_ref(
-                ref_name: str,
-            ) -> tuple[str | None, str | None, int | None, ResolutionReason]:
-                # Shares the lookup priority with `_resolve` so that composition
-                # respects overrides and registered code defaults rather than
-                # only consulting the provider.
-                ref_result = self._lookup_serialized(
-                    ref_name,
-                    provider=provider,
-                    targeting_key=targeting_key,
-                    attributes=attributes,
-                )
-                return (ref_result.value, ref_result.label, ref_result.version, ref_result.reason)
+        # Always run through `expand_references`, even when no `@{ref}@` tags
+        # are present: it's also responsible for unescaping `\@{...}@` →
+        # `@{...}@`. Gating on `has_references` produced inconsistent
+        # observable behaviour where an escaped-only value (e.g.
+        # `r'\@{baz}@'`) kept its backslash, but the same escape combined
+        # with a real reference correctly produced the literal `@{baz}@`.
+        def resolve_ref(
+            ref_name: str,
+        ) -> tuple[str | None, str | None, int | None, ResolutionReason]:
+            # Shares the lookup priority with `_resolve` so that composition
+            # respects overrides and registered code defaults rather than
+            # only consulting the provider.
+            ref_result = self._lookup_serialized(
+                ref_name,
+                provider=provider,
+                targeting_key=targeting_key,
+                attributes=attributes,
+            )
+            return (ref_result.value, ref_result.label, ref_result.version, ref_result.reason)
 
-            try:
-                serialized_value, composed = expand_references(
-                    serialized_value,
-                    self.name,
-                    resolve_ref,
-                )
-                if composition_error := _first_composition_error(composed):
-                    return self._composition_failure(
-                        exception=VariableCompositionError(composition_error),
-                        targeting_key=targeting_key,
-                        attributes=attributes,
-                        serialized_result=serialized_result,
-                        composed=composed,
-                    )
-            except VariableCompositionError as e:
+        try:
+            serialized_value, composed = expand_references(
+                serialized_value,
+                self.name,
+                resolve_ref,
+            )
+            if composition_error := _first_composition_error(composed):
                 return self._composition_failure(
-                    exception=e,
+                    exception=VariableCompositionError(composition_error),
                     targeting_key=targeting_key,
                     attributes=attributes,
                     serialized_result=serialized_result,
                     composed=composed,
                 )
+        except VariableCompositionError as e:
+            return self._composition_failure(
+                exception=e,
+                targeting_key=targeting_key,
+                attributes=attributes,
+                serialized_result=serialized_result,
+                composed=composed,
+            )
 
         # Apply render_fn (template rendering) if provided
         if render_fn is not None: