from contextvars import ContextVar

import pytest
from llama_index.core.tools import FunctionTool
from llama_index.core.workflow import Context


# The LLM being used shouldn't matter much: please try it on multiple.
def _make_nova_micro_llm() -> "BedrockConverse":  # noqa: F821
    """Construct a BedrockConverse LLM for amazon.nova-micro-v1:0.
    """
    # import llama_index.llms.bedrock_converse.utils as _bedrock_utils  # noqa: F401  (ensure patches applied)
    from llama_index.llms.bedrock_converse import BedrockConverse

    AWS_BEDROCK_GEO = "us"

    model = f"{AWS_BEDROCK_GEO}.amazon.nova-micro-v1:0"
    return BedrockConverse(model=model, max_tokens=5_120, **{"region_name": "us-east-1", "timeout": 600})



@pytest.mark.asyncio
@pytest.mark.parametrize("use_async_fn", [False, True], ids=["sync_fn", "async_fn"])
async def test_call_tool_reads_context_var_with_nova_micro(use_async_fn: bool) -> None:
    """FunctionTool (sync and async) should be able to read a ContextVar.

    When a FunctionAgent drives tool execution the Python contextvars context
    must be propagated into the tool call so that ContextVar.get() returns the
    value set on the calling context, not the default.

    Args:
        use_async_fn: When True the tool is registered as an async function;
            when False as a sync function.
    """
    from llama_index.core.agent.workflow import FunctionAgent

    some_context_var: ContextVar[list[str]] = ContextVar("some_context_var", default=["starting value"])

    # Particularly, it seems like llama_index has some custom event loop handling
    # that causes non-async tool calls to fail to correctly read `ContextVar`s. This
    # problem can be worked around by consumers who have ownership of those `ContextVar`s,
    # but cannot be if those ContextVars are owned by another package.
    # This becomes a problem particularly in OpenTelemetry, which uses a ContextVar to 
    # correctly parent a span. Currently, all non-async tool calls have spans that appear
    # as orphans in Otel.
    def read_context_var(ctx: Context) -> str:
        """Read the secret value stored in the some_context_var and return it."""
        return str(some_context_var.get())

    async def read_context_var_async(ctx: Context) -> str:
        """Read the secret value stored in the some_context_var and return it."""
        return str(some_context_var.get())

    if use_async_fn:
        tool = FunctionTool.from_defaults(
            async_fn=read_context_var_async,
            name="read_context_var",
            description="Always call this, just once.",
        )
    else:
        tool = FunctionTool.from_defaults(
            fn=read_context_var,
            name="read_context_var",
            description="Always call this, just once.",
        )

    llm = _make_nova_micro_llm()

    updated_value = f"updated value: {use_async_fn}"
    some_context_var.set([updated_value])

    agent = FunctionAgent(
        llm=llm,
        tools=[tool],
        system_prompt="You MUST call the given tool; you are testing its behavior.",
    )

    handler = agent.run(user_msg="What is the context var?")
    output = await handler

    assert updated_value in str(output), (
        f"Expected '{updated_value}' to appear in the agent output, but got: {output!r}\n"
        "This may be the llama-index contextvars propagation bug."
    )

---