zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-13 16:38:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): defend against @ in var names at all template-split sites (#16469)

Browse Source 
## Summary

While fixing #16467 (IterationItem crash on `@` in user-defined output
keys), an audit of `agent/**/*.py` revealed **three additional sites**
with the same vulnerability. This PR hardens all of them with
`maxsplit=1` and adds regression tests.

This is **defense-in-depth hardening**, not a behavior change. The
current `variable_ref_patt` regex constrains `var_nm` to
`[A-Za-z0-9_.-]+`, so single-`@` templates resolve exactly as before.
The `maxsplit=1` only kicks in if the trailing side itself contains `@`
— currently unreachable from the public DSL surface, but trivially
exploitable the moment a user-defined output key happens to contain `@`
(e.g. `user@email`) or the regex is ever relaxed.

> **Note on issue scope**: The primary fix for #16467 (the
`list_tenant_added_models` `ValueError` crash on `@` in model names) is
in PR #16468. This PR is a **follow-up hardening sweep** of the same
vulnerability class found in `agent/` during that audit; it does not
duplicate or replace #16468.

## Sites hardened

| File | Line | Method |
|------|------|--------|
| `agent/canvas.py` | 206 | `Graph.get_variable_value` |
| `agent/canvas.py` | 256 | `Graph.set_variable_value` |
| `agent/component/base.py` | 533 |
`ComponentBase.get_input_elements_from_text` |
| `agent/component/iterationitem.py` | 88 |
`IterationItem.output_collation` |

All now use `split("@", 1)` with an inline comment explaining the
rationale. The trailing side keeps any embedded `@`.

## Sites already safe (audited but left alone)

| File | Reason safe |
|------|------------|
| `agent/canvas.py:708` (`is_reff`) | Pre-checks `len(arr) != 2` |
| `agent/component/categorize.py` | Uses `rsplit` |
| `agent/component/iteration.py` | Pre-validates via regex |
| Other call sites | `rsplit` or regex pre-validation |

## Regression tests

9 new tests across 2 files, all `pytest.mark.p2`:

| File | Tests |
|------|-------|
| `test/unit_test/agent/test_canvas_at_split.py` | 6 —
`get_variable_value`, `set_variable_value`, round-trip, single-`@`,
missing-component |
| `test/unit_test/agent/component/test_iterationitem_at_split.py` | 3 —
`output_collation` with `@` in var, single-`@`, non-matching cid |

Each test was **verified to fail with `ValueError: too many values to
unpack (expected 2)`** when the corresponding fix is temporarily
reverted, confirming the tests actually catch the bug rather than just
exercising the happy path.

## Test results

```
9 passed in 0.04s
```

Full agent unit suite also clean (38 passed, 3 skipped; 6 unrelated
pre-existing collection errors from missing `peewee`/`requests` in local
venv — not caused by this PR).

## Related

- Issue: #16467
- Primary fix PR: #16468 (closes the issue)
- This PR: defense-in-depth follow-up, intentionally non-blocking on
#16467

---------

Co-authored-by: skbs-eng <skbs-eng@users.noreply.github.com>
This commit is contained in:
S
2026-07-03 16:56:27 +05:30
committed by GitHub
parent fd7fb6669a
commit 1861087787
5 changed files with 665 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
agent/canvas.py
View File

@@ -25,18 +25,18 @@ import time
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from functools import partial
from typing import Any, Union, Tuple
from typing import Any, Tuple, Union
from agent.component import component_class
from agent.component.base import ComponentBase
from agent.dsl_migration import normalize_chunker_dsl
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
from api.db.services.file_service import FileService
from api.db.services.llm_service import LLMBundle
from api.db.services.task_service import has_canceled
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
from common.constants import LLMType
from common.misc_utils import get_uuid, hash_str2int
from common.exceptions import TaskCanceledException
from common.misc_utils import get_uuid, hash_str2int
from common.token_utils import token_usage_sink, langfuse_run_attrs
from rag.prompts.generator import chunks_format
from rag.utils.redis_conn import REDIS_CONN
@@ -229,7 +229,13 @@ class Graph:
exp = exp.strip("{").strip("}").strip(" ").strip("{").strip("}")
if exp.find("@") < 0:
return self.globals[exp]
cpn_id, var_nm = exp.split("@")
# Split from the left with maxsplit=1 so the trailing var_nm can
# legitimately contain '@' characters (defensive: although the
# upstream regex in `get_value_with_variable` constrains `var_nm`
# to `[A-Za-z0-9_.-]+`, direct callers of this method may pass
# any string and should not raise `ValueError: too many values
# to unpack`). `cpn_id` is system-generated and never contains '@'.
cpn_id, var_nm = exp.split("@", 1)
cpn = self.get_component(cpn_id)
if not cpn:
raise Exception(f"Can't find variable: '{cpn_id}@{var_nm}'")
@@ -276,7 +282,10 @@ class Graph:
if exp.find("@") < 0:
self.globals[exp] = value
return
cpn_id, var_nm = exp.split("@")
# See `get_variable_value` above for rationale on `maxsplit=1`.
# Without it, a var_nm containing '@' would raise
# `ValueError: too many values to unpack` instead of being preserved.
cpn_id, var_nm = exp.split("@", 1)
cpn = self.get_component(cpn_id)
if not cpn:
raise Exception(f"Can't find variable: '{cpn_id}@{var_nm}'")

17
agent/component/base.py
View File

@@ -15,15 +15,17 @@
#
import asyncio
import builtins
import json
import logging
import os
import re
import time
from abc import ABC
import builtins
import json
import os
import logging
from typing import Any, List, Union
import pandas as pd
from agent import settings
from common.connection_utils import timeout
@@ -517,7 +519,12 @@ class ComponentBase(ABC):
res = {}
for r in re.finditer(self.variable_ref_patt, txt, flags=re.IGNORECASE | re.DOTALL):
exp = r.group(1)
cpn_id, var_nm = exp.split("@") if exp.find("@") > 0 else ("", exp)
# Use maxsplit=1 to be defensive: although `exp` here comes
# from `variable_ref_patt` (which constrains `var_nm` to
# `[A-Za-z0-9_.-]+`), a future regex relaxation or a non-
# pattern caller should not raise `ValueError: too many values
# to unpack` if the trailing part happens to contain '@'.
cpn_id, var_nm = exp.split("@", 1) if exp.find("@") > 0 else ("", exp)
res[exp] = {
"name": (self._canvas.get_component_name(cpn_id) + f"@{var_nm}") if cpn_id else exp,
"value": self._canvas.get_variable_value(exp),

7
agent/component/iterationitem.py
View File

@@ -14,6 +14,7 @@
# limitations under the License.
#
from abc import ABC
from agent.component.base import ComponentBase, ComponentParamBase
@@ -80,7 +81,11 @@ class IterationItem(ComponentBase, ABC):
for k, o in p._param.outputs.items():
if "ref" not in o:
continue
_cid, var = o["ref"].split("@")
# Use maxsplit=1 so an `@` legitimately embedded in `var`
# (e.g. a user-defined output key that happens to contain
# '@') does not raise `ValueError: too many values to unpack`.
# `_cid` is system-generated and never contains '@'.
_cid, var = o["ref"].split("@", 1)
if _cid != cid:
continue
res = p.output(k)

353
test/unit_test/agent/component/test_iterationitem_at_split.py Normal file
View File

@@ -0,0 +1,353 @@
#
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Hardening regression test for `@`-splitting in
`agent/component/iterationitem.py:82` (inside `output_collation`).
Before the fix, `o["ref"].split("@")` raised `ValueError: too many values
to unpack` whenever a user-defined output key contained `@` (e.g. an
output named `foo@bar` referenced from another component). After the
fix, the call uses `split("@", 1)`, preserving the trailing `@`-bearing
key in `var` so the collation append can proceed normally.
The module is loaded via `importlib` after stubbing its heavyweight
imports — same isolation strategy as
`test/testcases/test_web_api/test_canvas_app/test_iterationitem_unit.py`.
"""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import MagicMock
import pytest
# ─── Module loader ────────────────────────────────────────────────────
def _load_iterationitem_module(monkeypatch):
"""Load `agent.component.iterationitem` in isolation with stubs.
Unlike `test_canvas_at_split.py`, the iterationitem module's import
graph is small enough that we can run the real `common.*` modules
(after stubbing their `quart` dependency). That gives us high
fidelity on the `_param.outputs` iteration logic without dragging
in the canvas-level graph (which the tests don't need).
Returns:
The loaded `agent.component.iterationitem` module object.
"""
repo_root = Path(__file__).resolve().parents[4]
# Stub `quart` because `common.connection_utils` imports from it.
quart_stub = ModuleType("quart")
quart_stub.make_response = MagicMock()
quart_stub.jsonify = MagicMock()
monkeypatch.setitem(sys.modules, "quart", quart_stub)
# Stub the heavy transitive imports so we never touch pandas/jinja2/etc.
common_pkg = ModuleType("common")
common_pkg.__path__ = [str(repo_root / "common")]
monkeypatch.setitem(sys.modules, "common", common_pkg)
conn_spec = importlib.util.spec_from_file_location("common.connection_utils", repo_root / "common" / "connection_utils.py")
conn_mod = importlib.util.module_from_spec(conn_spec)
monkeypatch.setitem(sys.modules, "common.connection_utils", conn_mod)
conn_spec.loader.exec_module(conn_mod)
misc_spec = importlib.util.spec_from_file_location("common.misc_utils", repo_root / "common" / "misc_utils.py")
misc_mod = importlib.util.module_from_spec(misc_spec)
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_mod)
misc_spec.loader.exec_module(misc_mod)
constants_mod = ModuleType("common.constants")
class _RetCode:
"""Minimal stand-in for `common.constants.RetCode`.
Only the two integer sentinels actually referenced by
`iterationitem.py` are defined; everything else would otherwise
require loading the entire constants module.
"""
SUCCESS = 0
EXCEPTION_ERROR = 100
constants_mod.RetCode = _RetCode
monkeypatch.setitem(sys.modules, "common.constants", constants_mod)
agent_pkg = ModuleType("agent")
agent_pkg.__path__ = [str(repo_root / "agent")]
monkeypatch.setitem(sys.modules, "agent", agent_pkg)
agent_settings = ModuleType("agent.settings")
agent_settings.FLOAT_ZERO = 1e-8
agent_settings.PARAM_MAXDEPTH = 5
monkeypatch.setitem(sys.modules, "agent.settings", agent_settings)
component_pkg = ModuleType("agent.component")
component_pkg.__path__ = [str(repo_root / "agent" / "component")]
monkeypatch.setitem(sys.modules, "agent.component", component_pkg)
# Provide a minimal `ComponentBase` for the iterationitem module to
# inherit from. Only the surface used by `output_collation` matters:
# `output(var_nm)`, `set_output(key, value)`, plus the parent lookup
# `get_parent()` is defined on the test instance directly.
class _ComponentBaseStub:
"""No-op stand-in for the real `ComponentBase`.
`iterationitem.IterationItem` extends this class, so any
instantiation path that isn't bypassed via `__new__` would
require it to be importable. We deliberately no-op every
method so that nothing incidental fires during loading.
"""
def __init__(self, *a, **kw):
"""Accept and discard all args; the stub is purely nominal."""
pass
def output(self, var_nm=None):
"""Return an empty string for any variable lookup."""
return ""
def set_output(self, key, value):
"""Discard writes; iterationitem test paths verify via
`parent._param.outputs`, not via the base class."""
pass
class _ComponentParamBaseStub:
"""Stand-in for `ComponentParamBase` with mutable `outputs`/`inputs`."""
def __init__(self):
"""Initialize empty `outputs` and `inputs` dicts."""
self.outputs = {}
self.inputs = {}
def check(self):
"""Always succeed; validation is not exercised in these tests."""
return True
base_mod = ModuleType("agent.component.base")
base_mod.ComponentBase = _ComponentBaseStub
base_mod.ComponentParamBase = _ComponentParamBaseStub
monkeypatch.setitem(sys.modules, "agent.component.base", base_mod)
iterationitem_spec = importlib.util.spec_from_file_location(
"agent.component.iterationitem",
repo_root / "agent" / "component" / "iterationitem.py",
)
iterationitem_mod = importlib.util.module_from_spec(iterationitem_spec)
monkeypatch.setitem(sys.modules, "agent.component.iterationitem", iterationitem_mod)
iterationitem_spec.loader.exec_module(iterationitem_mod)
return iterationitem_mod
# ─── Fixtures ─────────────────────────────────────────────────────────
class _ParentStub:
"""Parent iteration component.
`_param.outputs` is keyed by output name; each entry has a `"ref"`
pointing to a child component (`"<cid>@<var_name>"`).
"""
def __init__(self, _id, outputs, component_name="Iteration"):
"""Store `_id`, default each output's `"value"` to `[]`, expose
the canonical `outputs` dict on a `_param` SimpleNamespace."""
self._id = _id
self.component_name = component_name
# Normalize: every output entry must carry a `"value"` key so
# `p.output(k)` and the post-append `p.set_output(k, res)` work
# the same way they would on a real `ComponentBase`.
normalized = {}
for name, payload in outputs.items():
entry = dict(payload)
entry.setdefault("value", [])
normalized[name] = entry
self._param = SimpleNamespace(outputs=normalized)
def output(self, var_nm):
"""Return the `"value"` of the named output, or `""` if absent."""
return self._param.outputs.get(var_nm, {}).get("value", "")
def set_output(self, key, value):
"""Write `value` into `_param.outputs[key]["value"]`, creating
the entry with the appropriate `type` sentinel on first write."""
if key not in self._param.outputs:
self._param.outputs[key] = {"value": None, "type": str(type(value))}
self._param.outputs[key]["value"] = value
class _ChildStub:
"""Child component inside an iteration.
`_param.outputs` stores the child component's own outputs in the
`ComponentBase` `{"value": ...}` shape.
"""
def __init__(self, _id, parent, output_values):
"""Bind to a parent and seed `_param.outputs` from a flat dict.
Args:
_id: Child component id used as the `<cid>` half of refs.
parent: The parent iteration component (for `get_parent`).
output_values: `{output_name: value}` flat mapping, wrapped
into the `ComponentBase` `{"value": ..., "type": ...}`
shape before storage.
"""
self._id = _id
self._parent = parent
self.component_name = "Generate"
self._param = SimpleNamespace(outputs={k: {"value": v, "type": str(type(v))} for k, v in output_values.items()})
def get_parent(self):
"""Return the parent iteration component passed at construction."""
return self._parent
def output(self, var_nm):
"""Return the `"value"` of the named output, or `""` if absent."""
return self._param.outputs.get(var_nm, {}).get("value", "")
def set_output(self, key, value):
"""Write `value` into `_param.outputs[key]["value"]`, creating
the entry with the appropriate `type` sentinel on first write."""
if key not in self._param.outputs:
self._param.outputs[key] = {"value": None, "type": str(type(value))}
self._param.outputs[key]["value"] = value
class _CanvasStub:
"""Minimal canvas: only `components` (dict) and `get_component_obj`
are touched by `output_collation`."""
def __init__(self, components):
"""Store the `{cid: component}` lookup used by `get_component_obj`."""
self.components = components
def get_component_obj(self, cid):
"""Look up a component by id; the stub never returns `None`."""
return self.components[cid]
def _make_iteration_item(module, canvas, parent):
"""Construct an `IterationItem` bypassing its normal `__init__`.
`output_collation` only touches `self._canvas`, `self._id`, and
`self.get_parent()`, so a stripped-down instance is sufficient.
"""
inst = module.IterationItem.__new__(module.IterationItem)
inst._canvas = canvas
inst._id = "IterationItem:test"
inst._param = SimpleNamespace(outputs={}, inputs={})
inst._idx = 0
inst.get_parent = MagicMock(return_value=parent)
return inst
# ─── Tests ────────────────────────────────────────────────────────────
@pytest.mark.p2
def test_output_collation_accepts_at_in_var_name(monkeypatch):
"""The parent iteration's output `result` references the child's
output `foo@bar` (a user-defined key containing `@`). Before the
fix, `o["ref"].split("@")` raised `ValueError: too many values to
unpack`. After the fix, `output_collation` must:
1. Not raise.
2. Append the child's `foo@bar` value into the parent's `result`
list (i.e. treat `var = "foo@bar"` as a literal output name).
"""
module = _load_iterationitem_module(monkeypatch)
parent_outputs = {
"result": {"ref": "child-1@foo@bar"},
}
parent = _ParentStub("iter-pid", parent_outputs)
child = _ChildStub(
_id="child-1",
parent=parent,
output_values={"foo@bar": "value-A"},
)
canvas = _CanvasStub({"child-1": child})
item = _make_iteration_item(module, canvas, parent)
# Should NOT raise ValueError.
item.output_collation()
# The child's `foo@bar` value should have been appended to the
# parent's `result` list.
assert parent._param.outputs["result"]["value"] == ["value-A"]
@pytest.mark.p2
def test_output_collation_single_at_still_works(monkeypatch):
"""Sanity check: the existing single-`@` ref path is unaffected."""
module = _load_iterationitem_module(monkeypatch)
parent_outputs = {
"result": {"ref": "child-1@normal_var"},
}
parent = _ParentStub("iter-pid", parent_outputs)
child = _ChildStub(
_id="child-1",
parent=parent,
output_values={"normal_var": "value-B"},
)
canvas = _CanvasStub({"child-1": child})
item = _make_iteration_item(module, canvas, parent)
item.output_collation()
assert parent._param.outputs["result"]["value"] == ["value-B"]
@pytest.mark.p2
def test_output_collation_skips_non_matching_cid(monkeypatch):
"""When the ref's `cid` doesn't match the current `cid`, the loop
must skip silently even if `var` contains `@`. This guards against
an over-eager fix that might, e.g., re-split on the wrong side."""
module = _load_iterationitem_module(monkeypatch)
# The ref points at `other-cid`, not `child-1`.
parent_outputs = {
"result": {"ref": "other-cid@foo@bar"},
}
parent = _ParentStub("iter-pid", parent_outputs)
child = _ChildStub(
_id="child-1",
parent=parent,
output_values={"foo@bar": "ignored"},
)
canvas = _CanvasStub({"child-1": child})
item = _make_iteration_item(module, canvas, parent)
# Must not raise.
item.output_collation()
# Nothing should have been appended (the ref didn't match this cid).
assert parent._param.outputs["result"]["value"] == []

280
test/unit_test/agent/test_canvas_at_split.py Normal file
View File

@@ -0,0 +1,280 @@
#
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Hardening regression tests for `@`-splitting in agent/canvas.py.
These exercise `get_variable_value` and `set_variable_value` on
`agent/canvas.py` (lines 201 and 248 of the Graph class, inherited by
`Canvas`) with references whose variable name contains `@`. Before the
fix, `exp.split("@")` would raise `ValueError: too many values to unpack`
because Python splits on every `@` by default. After the fix, the calls
use `split("@", 1)`, so the trailing `@`-bearing variable name is
preserved verbatim.
We load `agent/canvas.py` via `importlib` after stubbing its heavyweight
imports (LLM service, task service, etc.) with MagicMock modules — same
isolation strategy as
`test/testcases/test_web_api/test_canvas_app/test_iterationitem_unit.py`,
which lets these tests run without a live DB / Redis / MinIO stack.
"""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import MagicMock
import pytest
# ─── Module loader ────────────────────────────────────────────────────
def _load_canvas_module(monkeypatch):
"""Load `agent.canvas` in isolation with all heavy deps stubbed out.
The canvas import graph pulls in `pandas`, `quart`, `jinja2`, the
real `ComponentBase`, ORM models, Redis, TTS cache, etc. None of
that is needed for the `get_variable_value` / `set_variable_value`
paths under test, so we register lightweight `ModuleType` stubs in
`sys.modules` first and then `exec_module` the real `canvas.py`
against that fake module table.
Returns:
The loaded `agent.canvas` module object.
"""
repo_root = Path(__file__).resolve().parents[3]
# `agent.component.base` pulls in pandas, quart, jinja2, etc.
# Stub the modules it (and the rest of the canvas import graph) needs
# before loading, so we never touch the real implementations.
def _stub_module(name, **attrs):
"""Create a fresh `ModuleType`, attach attrs, register in `sys.modules`.
Used to short-circuit every transitive import that `canvas.py`
performs at module load time. The returned module is also the
actual object that `import` statements resolve to during
`exec_module`.
Args:
name: Fully-qualified module name (e.g. ``"common.constants"``).
**attrs: Attributes to set on the new module before it is
registered (typically classes, callables, or sentinel
objects expected by importers).
Returns:
The newly created `ModuleType` instance.
"""
mod = ModuleType(name)
for k, v in attrs.items():
setattr(mod, k, v)
monkeypatch.setitem(sys.modules, name, mod)
return mod
# Parent packages must exist as packages with `__path__` so submodule
# imports resolve correctly.
for pkg_name, pkg_path in [
("common", repo_root / "common"),
("agent", repo_root / "agent"),
("agent.component", repo_root / "agent" / "component"),
]:
pkg = ModuleType(pkg_name)
pkg.__path__ = [str(pkg_path)]
monkeypatch.setitem(sys.modules, pkg_name, pkg)
_stub_module("agent.settings", FLOAT_ZERO=1e-8, PARAM_MAXDEPTH=5)
# `agent.canvas` and `agent.component.base` import each other
# indirectly. Provide a minimal `ComponentBase` so the canvas module
# can be loaded without dragging the real one in.
class _ComponentBaseStub:
"""Minimal stand-in for `agent.component.base.ComponentBase`.
The real class would drag in the entire component registry and
ORM-layer initialization; we only need an object that
`canvas.py`'s imports can bind to without side effects.
"""
def __init__(self, *a, **kw):
"""Accept and discard all args; the stub is purely nominal."""
pass
base_stub_mod = _stub_module("agent.component.base")
base_stub_mod.ComponentBase = _ComponentBaseStub
base_stub_mod.ComponentParamBase = type("ComponentParamBase", (), {"outputs": {}, "inputs": {}})
# `agent.component.component_class` is a registry factory used at
# canvas load time. The tests below never call `Canvas.load()`, so
# any callable suffices.
_stub_module("agent.component", component_class=lambda *_a, **_kw: MagicMock())
_stub_module("agent.dsl_migration", normalize_chunker_dsl=lambda dsl: dsl)
# `api.*` services imported at the top of canvas.py
_stub_module("api.db.services.file_service", FileService=MagicMock())
_stub_module("api.db.services.llm_service", LLMBundle=MagicMock())
_stub_module("api.db.services.task_service", has_canceled=MagicMock(return_value=False))
_stub_module(
"api.db.joint_services.tenant_model_service",
get_tenant_default_model_by_type=MagicMock(return_value=None),
)
_stub_module("common.constants", LLMType=MagicMock())
_stub_module("common.misc_utils", get_uuid=lambda: "test-uuid", hash_str2int=lambda _s: 0, thread_pool_exec=lambda fn, *a, **kw: fn(*a, **kw))
_stub_module(
"common.token_utils",
token_usage_sink=lambda *_a, **_kw: None,
langfuse_run_attrs=lambda *_a, **_kw: {},
)
_stub_module("common.connection_utils", timeout=lambda *_a, **_kw: lambda fn: fn)
_stub_module("common.exceptions", TaskCanceledException=type("TaskCanceledException", (Exception,), {}))
_stub_module("rag.prompts.generator", chunks_format=MagicMock())
_stub_module("rag.utils.redis_conn", REDIS_CONN=MagicMock())
_stub_module("rag.utils.tts_cache", synthesize_with_cache=MagicMock())
spec = importlib.util.spec_from_file_location("agent.canvas", repo_root / "agent" / "canvas.py")
canvas_mod = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, "agent.canvas", canvas_mod)
spec.loader.exec_module(canvas_mod)
return canvas_mod
# ─── Fixtures ─────────────────────────────────────────────────────────
class _ComponentObjStub:
"""Minimal stub emulating the `cpn["obj"]` interface used by
`get_variable_value` / `set_variable_value`.
Only `.output(var_nm)` and `.set_output(var_nm, value)` are needed,
mirroring `ComponentBase.output` / `ComponentBase.set_output`.
"""
def __init__(self):
"""Initialize an empty per-instance output store."""
self._store: dict = {}
def output(self, var_nm):
return self._store.get(var_nm)
def set_output(self, var_nm, value):
self._store[var_nm] = value
def _make_canvas(canvas_mod, components=None, globals_=None):
"""Construct a `Graph` instance that bypasses `Canvas.__init__`.
The methods under test (`get_variable_value`, `set_variable_value`)
only touch `self.globals` and `self.components` (via `get_component`),
so a stripped-down instance is sufficient.
"""
inst = canvas_mod.Graph.__new__(canvas_mod.Graph)
inst.globals = dict(globals_ or {})
inst.components = dict(components or {})
return inst
# ─── Tests ────────────────────────────────────────────────────────────
@pytest.mark.p2
def test_get_variable_value_preserves_at_in_var_name(monkeypatch):
"""`get_variable_value("{cpn@foo@bar}")` must not raise; it should
return the component's stored value for the literal key `foo@bar`."""
canvas_mod = _load_canvas_module(monkeypatch)
cpn_obj = _ComponentObjStub()
cpn_obj.set_output("foo@bar", "value-1")
canvas = _make_canvas(canvas_mod, components={"cpn-1": {"obj": cpn_obj}})
result = canvas.get_variable_value("{cpn-1@foo@bar}")
assert result == "value-1"
@pytest.mark.p2
def test_set_variable_value_preserves_at_in_var_name(monkeypatch):
"""`set_variable_value` must not raise on a `@`-bearing var name."""
canvas_mod = _load_canvas_module(monkeypatch)
cpn_obj = _ComponentObjStub()
canvas = _make_canvas(canvas_mod, components={"cpn-1": {"obj": cpn_obj}})
canvas.set_variable_value("{cpn-1@foo@bar}", "value-2")
assert cpn_obj.output("foo@bar") == "value-2"
@pytest.mark.p2
def test_get_variable_value_legacy_split_would_raise(monkeypatch):
"""Regression: under the previous `exp.split("@")`, this same
expression raised `ValueError: too many values to unpack`. Asserting
that the call returns cleanly (and resolves to the stored value)
proves the `split("@", 1)` hardening is in effect."""
canvas_mod = _load_canvas_module(monkeypatch)
cpn_obj = _ComponentObjStub()
cpn_obj.set_output("nested@with@ats", {"nested": True})
canvas = _make_canvas(canvas_mod, components={"cpn-1": {"obj": cpn_obj}})
# If the legacy `split("@")` were in place this would raise
# `ValueError: too many values to unpack (expected 2)`.
result = canvas.get_variable_value("{cpn-1@nested@with@ats}")
assert result == {"nested": True}
@pytest.mark.p2
def test_set_variable_value_then_get_round_trip(monkeypatch):
"""Round-trip: a write through `set_variable_value` with a multi-`@`
key should be observable through `get_variable_value`."""
canvas_mod = _load_canvas_module(monkeypatch)
cpn_obj = _ComponentObjStub()
canvas = _make_canvas(canvas_mod, components={"cpn-1": {"obj": cpn_obj}})
canvas.set_variable_value("{cpn-1@user@email}", "alice@example.com")
assert canvas.get_variable_value("{cpn-1@user@email}") == "alice@example.com"
@pytest.mark.p2
def test_get_variable_value_missing_component_still_raises(monkeypatch):
"""The hardening must NOT swallow the existing `Can't find variable`
exception for unknown component IDs — that's a legitimate error
path that callers depend on."""
canvas_mod = _load_canvas_module(monkeypatch)
canvas = _make_canvas(canvas_mod, components={})
with pytest.raises(Exception, match="Can't find variable"):
canvas.get_variable_value("{missing-cpn@foo@bar}")
@pytest.mark.p2
def test_get_variable_value_single_at_still_works(monkeypatch):
"""Sanity check that the existing single-`@` path is unaffected by
the hardening (no regression in the common case)."""
canvas_mod = _load_canvas_module(monkeypatch)
cpn_obj = _ComponentObjStub()
cpn_obj.set_output("normal_var", "normal-value")
canvas = _make_canvas(canvas_mod, components={"cpn-1": {"obj": cpn_obj}})
assert canvas.get_variable_value("{cpn-1@normal_var}") == "normal-value"