From a6a873dfe8276833ae72aa47c12adb473e534954 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 09:47:08 +0200 Subject: [PATCH 1/2] ref(langgraph): Drop send_default_pii support --- sentry_sdk/integrations/langgraph.py | 66 +- .../integrations/langgraph/test_langgraph.py | 567 +++++++----------- 2 files changed, 241 insertions(+), 392 deletions(-) diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 07708dd007..114d835cd6 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -11,9 +11,7 @@ # This is fine because langgraph depends on langchain-base, and LangchainIntegration only imports from langchain-base. from sentry_sdk.integrations.langchain import LangchainIntegration -from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import ( - has_data_collection_enabled, package_version, safe_serialize, ) @@ -51,24 +49,6 @@ def setup_once() -> None: Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke) -def _should_record_inputs(integration: "LanggraphIntegration") -> bool: - client = sentry_sdk.get_client() - if has_data_collection_enabled(client.options): - return bool(client.options["data_collection"]["gen_ai"]["inputs"]) - - # To remove once data collection has been fully rolled out - return should_send_default_pii() - - -def _should_record_outputs(integration: "LanggraphIntegration") -> bool: - client = sentry_sdk.get_client() - if has_data_collection_enabled(client.options): - return bool(client.options["data_collection"]["gen_ai"]["outputs"]) - - # To remove once data collection has been fully rolled out - return should_send_default_pii() - - def _get_graph_name(graph_obj: "Any") -> "Optional[str]": for attr in ["name", "graph_name", "__name__", "_name"]: if hasattr(graph_obj, attr): @@ -156,7 +136,10 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": input_messages = None if len(args) > 0: input_messages = _parse_langgraph_messages(args[0]) - if input_messages and _should_record_inputs(integration): + if ( + input_messages + and client.options["data_collection"]["gen_ai"]["inputs"] + ): normalized_input_messages = normalize_message_roles(input_messages) set_data_normalized( span, @@ -199,7 +182,10 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": input_messages = None if len(args) > 0: input_messages = _parse_langgraph_messages(args[0]) - if input_messages and _should_record_inputs(integration): + if ( + input_messages + and client.options["data_collection"]["gen_ai"]["inputs"] + ): normalized_input_messages = normalize_message_roles(input_messages) set_data_normalized( span, @@ -277,20 +263,22 @@ def _set_response_attributes( if new_messages is None: return - if _should_record_outputs(integration): - llm_response_text = _extract_llm_response_text(new_messages) - if llm_response_text: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text) - elif new_messages: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages) - else: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result) - - tool_calls = _extract_tool_calls(new_messages) - if tool_calls: - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - safe_serialize(tool_calls), - unpack=False, - ) + if not sentry_sdk.get_client().options["data_collection"]["gen_ai"]["outputs"]: + return + + llm_response_text = _extract_llm_response_text(new_messages) + if llm_response_text: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text) + elif new_messages: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages) + else: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result) + + tool_calls = _extract_tool_calls(new_messages) + if tool_calls: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + safe_serialize(tool_calls), + unpack=False, + ) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index a65b9f39aa..3acab77ebc 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -137,20 +137,20 @@ def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResu raise GraphBubbleUp("interrupt") -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_pregel_invoke( sentry_init, capture_items, - send_default_pii, ): """Test Pregel.invoke() wrapper creates proper invoke_agent span.""" sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) test_state = { @@ -205,52 +205,118 @@ def original_invoke(self, *args, **kwargs): assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" - if send_default_pii: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] - request_messages = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + request_messages = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - if isinstance(request_messages, str): - request_messages = json.loads(request_messages) - assert len(request_messages) == 2 - assert request_messages[0]["content"] == "Hello, can you help me?" - assert request_messages[1]["content"] == "Of course! How can I assist you?" + if isinstance(request_messages, str): + request_messages = json.loads(request_messages) + assert len(request_messages) == 2 + assert request_messages[0]["content"] == "Hello, can you help me?" + assert request_messages[1]["content"] == "Of course! How can I assist you?" - response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response + response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] + tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) + if isinstance(tool_calls_data, str): + tool_calls_data = json.loads(tool_calls_data) - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_test_123" - assert tool_calls_data[0]["function"]["name"] == "search_tool" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "attributes", {} - ) + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_test_123" + assert tool_calls_data[0]["function"]["name"] == "search_tool" + + +def test_pregel_invoke_no_sensitive_data( + sentry_init, + capture_items, +): + """Test Pregel.invoke() wrapper creates proper invoke_agent span.""" + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + test_state = { + "messages": [ + MockMessage("Hello, can you help me?", name="user"), + MockMessage("Of course! How can I assist you?", name="assistant"), + ] + } + + pregel = MockPregelInstance("test_graph") + + expected_assistant_response = "I'll help you with that task!" + expected_tool_calls = [ + { + "id": "call_test_123", + "type": "function", + "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, + } + ] + + def original_invoke(self, *args, **kwargs): + input_messages = args[0].get("messages", []) + new_messages = input_messages + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + return {"messages": new_messages} + + items = capture_items("span") + + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) + + assert result is not None + + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span["name"] == "invoke_agent test_graph" + assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get("attributes", {}) -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_pregel_ainvoke( sentry_init, capture_items, - send_default_pii, ): """Test Pregel.ainvoke() async wrapper creates proper invoke_agent span.""" sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) test_state = {"messages": [MockMessage("What's the weather like?", name="user")]} @@ -302,28 +368,91 @@ async def run_test(): assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" - if send_default_pii: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + + response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response + + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] + tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response + if isinstance(tool_calls_data, str): + tool_calls_data = json.loads(tool_calls_data) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_weather_456" + assert tool_calls_data[0]["function"]["name"] == "get_weather" - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_weather_456" - assert tool_calls_data[0]["function"]["name"] == "get_weather" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "attributes", {} - ) +def test_pregel_ainvoke_no_sensitive_data( + sentry_init, + capture_items, +): + """Test Pregel.ainvoke() async wrapper creates proper invoke_agent span.""" + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + test_state = {"messages": [MockMessage("What's the weather like?", name="user")]} + pregel = MockPregelInstance("async_graph") + + expected_assistant_response = "It's sunny and 72°F today!" + expected_tool_calls = [ + { + "id": "call_weather_456", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "current"}'}, + } + ] + + async def original_ainvoke(self, *args, **kwargs): + input_messages = args[0].get("messages", []) + new_messages = input_messages + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + return {"messages": new_messages} + + async def run_test(): + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result + + items = capture_items("span") + + result = asyncio.run(run_test()) + assert result is not None + + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span["name"] == "invoke_agent async_graph" + assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get("attributes", {}) def test_pregel_invoke_error( @@ -334,7 +463,12 @@ def test_pregel_invoke_error( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) test_state = {"messages": [MockMessage("This will fail")]} @@ -370,7 +504,12 @@ def test_pregel_ainvoke_error( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) test_state = {"messages": [MockMessage("This will fail async")]} @@ -411,7 +550,12 @@ def test_pregel_invoke_with_different_graph_names( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) pregel = MockPregelInstance(graph_name) if graph_name else MockPregelInstance() @@ -503,7 +647,12 @@ def test_extraction_functions_complex_scenario( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) pregel = MockPregelInstance("complex_graph") @@ -582,7 +731,12 @@ def test_langgraph_message_role_mapping( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock a langgraph message with mixed roles @@ -648,6 +802,7 @@ def __init__(self, content, message_type="human"): def test_graph_bubble_up_ignored(sentry_init, capture_items): sentry_init( integrations=[LanggraphIntegration()], + data_collection={}, ) events = capture_items("event") @@ -657,297 +812,3 @@ def test_graph_bubble_up_ignored(sentry_init, capture_items): model.invoke([HumanMessage(content="hi")]) assert len(events) == 0 - - -def _invoke_span_data(items_or_events): - sentry_sdk.flush() - spans = [item.payload for item in items_or_events] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - return invoke_spans[0]["attributes"] - - -@pytest.mark.parametrize( - "data_collection, send_default_pii, expect_inputs", - [ - pytest.param( - {"gen_ai": {"inputs": True}}, - False, - True, - id="gen-ai-inputs-enabled-overrides-pii-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - True, - False, - id="gen-ai-inputs-disabled-overrides-pii-enabled", - ), - pytest.param( - {}, - False, - True, - id="gen-ai-omitted-defaults-to-enabled", - ), - pytest.param( - {"gen_ai": {"outputs": False}}, - False, - True, - id="gen-ai-outputs-disabled-does-not-affect-inputs", - ), - pytest.param( - None, - False, - False, - id="no-data-collection-falls-back-to-send-default-pii", - ), - pytest.param( - None, - True, - True, - id="no-data-collection-pii-enabled-collects", - ), - ], -) -def test_pregel_invoke_gates_request_messages_on_inputs_setting( - sentry_init, - capture_items, - data_collection, - send_default_pii, - expect_inputs, -): - init_kwargs = { - "integrations": [LanggraphIntegration()], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - } - if data_collection is not None: - init_kwargs["data_collection"] = data_collection - - sentry_init(**init_kwargs) - - test_state = {"messages": [MockMessage("Hello, can you help me?", name="user")]} - pregel = MockPregelInstance("test_graph") - - def original_invoke(self, *args, **kwargs): - return { - "messages": args[0].get("messages", []) - + [ - MockMessage( - content="I'll help you with that task!", - name="assistant", - ) - ] - } - - captured = capture_items("span") - - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) - - data = _invoke_span_data(captured) - - if expect_inputs: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data - - -@pytest.mark.parametrize( - "data_collection, send_default_pii, expect_outputs", - [ - pytest.param( - {"gen_ai": {"outputs": True}}, - False, - True, - id="gen-ai-outputs-enabled-overrides-pii-disabled", - ), - pytest.param( - {"gen_ai": {"outputs": False}}, - True, - False, - id="gen-ai-outputs-disabled-overrides-pii-enabled", - ), - pytest.param( - {}, - False, - True, - id="gen-ai-omitted-defaults-to-enabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - False, - True, - id="gen-ai-inputs-disabled-does-not-affect-outputs", - ), - pytest.param( - None, - False, - False, - id="no-data-collection-falls-back-to-send-default-pii", - ), - pytest.param( - None, - True, - True, - id="no-data-collection-pii-enabled-collects", - ), - ], -) -def test_pregel_invoke_gates_response_text_and_tool_calls_on_outputs_setting( - sentry_init, - capture_items, - data_collection, - send_default_pii, - expect_outputs, -): - init_kwargs = { - "integrations": [LanggraphIntegration()], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - } - if data_collection is not None: - init_kwargs["data_collection"] = data_collection - - sentry_init(**init_kwargs) - - test_state = {"messages": [MockMessage("Hello, can you help me?", name="user")]} - pregel = MockPregelInstance("test_graph") - expected_assistant_response = "I'll help you with that task!" - expected_tool_calls = [ - { - "id": "call_test_123", - "type": "function", - "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, - } - ] - - def original_invoke(self, *args, **kwargs): - return { - "messages": args[0].get("messages", []) - + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - ) - ] - } - - captured = capture_items("span") - - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) - - data = _invoke_span_data(captured) - - if expect_outputs: - assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response - assert ( - json.loads(data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == expected_tool_calls - ) - else: - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data - - -@pytest.mark.parametrize( - "data_collection, send_default_pii, expect_inputs, expect_outputs", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - True, - False, - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - True, - False, - True, - id="gen-ai-inputs-disabled-outputs-enabled", - ), - pytest.param( - None, - False, - False, - False, - id="no-data-collection-falls-back-to-send-default-pii", - ), - pytest.param( - None, - True, - True, - True, - id="no-data-collection-pii-enabled-collects", - ), - ], -) -def test_pregel_ainvoke_gates_inputs_and_outputs_independently( - sentry_init, - capture_items, - data_collection, - send_default_pii, - expect_inputs, - expect_outputs, -): - init_kwargs = { - "integrations": [LanggraphIntegration()], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - } - if data_collection is not None: - init_kwargs["data_collection"] = data_collection - - sentry_init(**init_kwargs) - - test_state = {"messages": [MockMessage("What is the weather?", name="user")]} - pregel = MockPregelInstance("async_graph") - expected_assistant_response = "Let me check the weather for you!" - expected_tool_calls = [ - { - "id": "call_weather_456", - "type": "function", - "function": {"name": "get_weather", "arguments": '{"location": "Berlin"}'}, - } - ] - - async def original_ainvoke(self, *args, **kwargs): - return { - "messages": args[0].get("messages", []) - + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - ) - ] - } - - async def run_test(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - return await wrapped_ainvoke(pregel, test_state) - - captured = capture_items("span") - - asyncio.run(run_test()) - - data = _invoke_span_data(captured) - - if expect_inputs: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data - - if expect_outputs: - assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response - assert ( - json.loads(data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == expected_tool_calls - ) - else: - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data From 755244e416fadc6718680aaf220417b75633830b Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 11:00:12 +0200 Subject: [PATCH 2/2] restore cases of test_pregel_ainvoke_gates_inputs_and_outputs_independently --- .../integrations/langgraph/test_langgraph.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index 3acab77ebc..6964bdba7f 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -812,3 +812,96 @@ def test_graph_bubble_up_ignored(sentry_init, capture_items): model.invoke([HumanMessage(content="hi")]) assert len(events) == 0 + + +def _invoke_span_data(items_or_events): + sentry_sdk.flush() + spans = [item.payload for item in items_or_events] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + return invoke_spans[0]["attributes"] + + +@pytest.mark.parametrize( + "data_collection, expect_inputs, expect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + True, + id="gen-ai-inputs-disabled-outputs-enabled", + ), + ], +) +def test_pregel_ainvoke_gates_inputs_and_outputs_independently( + sentry_init, + capture_items, + data_collection, + expect_inputs, + expect_outputs, +): + init_kwargs = { + "integrations": [LanggraphIntegration()], + "traces_sample_rate": 1.0, + } + if data_collection is not None: + init_kwargs["data_collection"] = data_collection + + sentry_init(**init_kwargs) + + test_state = {"messages": [MockMessage("What is the weather?", name="user")]} + pregel = MockPregelInstance("async_graph") + expected_assistant_response = "Let me check the weather for you!" + expected_tool_calls = [ + { + "id": "call_weather_456", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "Berlin"}'}, + } + ] + + async def original_ainvoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + } + + async def run_test(): + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + return await wrapped_ainvoke(pregel, test_state) + + captured = capture_items("span") + + asyncio.run(run_test()) + + data = _invoke_span_data(captured) + + if expect_inputs: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data + + if expect_outputs: + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response + assert ( + json.loads(data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == expected_tool_calls + ) + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data