diff --git a/src/forge/integrations/agents/agent.py b/src/forge/integrations/agents/agent.py index 2497bbda..5ac52fd6 100644 --- a/src/forge/integrations/agents/agent.py +++ b/src/forge/integrations/agents/agent.py @@ -71,15 +71,30 @@ "retry_count", "repo", "pr_number", + "is_question", + "is_revision", + "revision_requested", + "task", } ) def _forward_trace_fields(context: dict[str, Any] | None) -> dict[str, Any]: """Extract trace-relevant fields from an incoming context dict.""" - if not context: - return {} - return {k: v for k, v in context.items() if k in _TRACE_FIELD_KEYS} + extracted = {k: v for k, v in context.items() if k in _TRACE_FIELD_KEYS} if context else {} + + # Resolve executing node name from LangGraph config if available + try: + from langchain_core.runnables.config import ensure_config + + config = ensure_config() + langgraph_node = config.get("metadata", {}).get("langgraph_node") + if langgraph_node: + extracted["current_node"] = langgraph_node + except Exception: + pass + + return extracted def _prompt_context_fields( @@ -755,6 +770,16 @@ async def run_task( logger.info(f"Running task '{task}' using Deep Agents") record_agent_invocation(task_type=task) _start = time.monotonic() + # Resolve executing node name from LangGraph config if available + langgraph_node = None + try: + from langchain_core.runnables.config import ensure_config + + config = ensure_config() + langgraph_node = config.get("metadata", {}).get("langgraph_node") + except Exception: + pass + # Merge prompt context + trace-only fields for Langfuse resolution. # trace_context fields are intentionally excluded from system_prompt above. trace_state: dict[str, Any] = { @@ -763,6 +788,10 @@ async def run_task( "system_prompt_length": len(system_prompt), "llm_model": self.settings.claude_model, } + + if langgraph_node: + trace_state["current_node"] = langgraph_node + trace_tags, trace_metadata = resolve_trace_fields(trace_state) result = await self._run_agent( diff --git a/src/forge/integrations/langfuse/fields.py b/src/forge/integrations/langfuse/fields.py index fe121a30..be929262 100644 --- a/src/forge/integrations/langfuse/fields.py +++ b/src/forge/integrations/langfuse/fields.py @@ -70,8 +70,84 @@ def _resolve_project_id(state: dict[str, Any]) -> str | None: def _resolve_workflow_step(state: dict[str, Any]) -> str | None: - val = state.get("current_node") - return str(val) if val is not None else None + current_node = state.get("current_node") + if current_node is None: + return None + + node_str = str(current_node) + + # Determine workflow_stage_or_node + # If the active operation is decompose_epics or generate_tasks, categorize under breakdown flow + if node_str in ("decompose_epics", "generate_tasks"): + stage_or_node = "decompose_epics" if node_str == "decompose_epics" else "generate_tasks" + else: + stage_or_node = node_str + + # Determine operation_type + is_question = state.get("is_question", False) + is_revision = state.get("is_revision", False) + revision_requested = state.get("revision_requested", False) + + if is_question or "qa" in node_str or "question" in node_str: + operation_type = "question_asking" + elif ( + is_revision + or revision_requested + or node_str.startswith("regenerate_") + or "revise" in node_str + ): + operation_type = "revision" + elif node_str in ("decompose_epics", "generate_tasks"): + operation_type = "breakdown" + elif "gate" in node_str or "approval" in node_str: + operation_type = "approval_gate" + else: + operation_type = "initial_generation" + + # Determine artifact_type + # Map based on node name, task, or state properties + # Target artifacts: spec, prd, plan, tasks, epic_breakdown, etc. + if state.get("artifact_type") is not None: + artifact_type = str(state["artifact_type"]) + else: + artifact_type = "unknown" + task = state.get("task", "") or "" + task_str = str(task).lower() + + # Analyze state/node to map artifact type + if "spec" in node_str or "spec" in task_str: + artifact_type = "spec" + elif "prd" in node_str or "prd" in task_str: + artifact_type = "prd" + elif "plan" in node_str or "plan" in task_str: + artifact_type = "plan" + elif "task" in node_str or "task" in task_str: + artifact_type = "tasks" + elif ( + "epic" in node_str + or "epic" in task_str + or "breakdown" in node_str + or "breakdown" in task_str + or node_str == "decompose_epics" + ): + artifact_type = "epic_breakdown" + elif node_str == "generate_tasks": + artifact_type = "tasks" + + # Format base label: [workflow_stage_or_node]:[operation_type]:[artifact_type] + label = f"{stage_or_node}:{operation_type}:{artifact_type}" + + # Appended when retry_count is greater than 0 + retry_count = state.get("retry_count", 0) + try: + retry_int = int(retry_count) if retry_count is not None else 0 + except (ValueError, TypeError): + retry_int = 0 + + if retry_int > 0: + label += f":attempt-{retry_int}" + + return label def _resolve_repo(state: dict[str, Any]) -> str | None: diff --git a/tests/unit/integrations/agents/test_run_task_tracing.py b/tests/unit/integrations/agents/test_run_task_tracing.py index ae1857bb..d0b98803 100644 --- a/tests/unit/integrations/agents/test_run_task_tracing.py +++ b/tests/unit/integrations/agents/test_run_task_tracing.py @@ -36,9 +36,7 @@ async def test_builds_trace_state_from_context_and_system_prompt( with ( patch.object(agent, "_run_agent", new_callable=AsyncMock) as mock_run, - patch( - "forge.integrations.agents.agent.resolve_trace_fields" - ) as mock_resolve, + patch("forge.integrations.agents.agent.resolve_trace_fields") as mock_resolve, patch("forge.integrations.agents.agent.load_prompt", return_value="prompt"), ): mock_run.return_value = "result" @@ -98,7 +96,26 @@ async def test_uses_trace_context_ticket_key_for_session_when_context_omits_it( call_kwargs = mock_run.call_args.kwargs assert call_kwargs["session_id"] == "PROJ-42" assert call_kwargs["ticket_key"] == "PROJ-42" - assert "PROJ-42" not in call_kwargs["system_prompt"] + + @pytest.mark.asyncio + async def test_run_task_resolves_node_from_langgraph_config(self, agent: ForgeAgent) -> None: + context = {"ticket_key": "PROJ-42"} + config = {"metadata": {"langgraph_node": "generate_tasks"}} + + with ( + patch.object(agent, "_run_agent", new_callable=AsyncMock) as mock_run, + patch( + "forge.integrations.agents.agent.resolve_trace_fields", + return_value=([], {}), + ) as mock_resolve, + patch("forge.integrations.agents.agent.load_prompt", return_value="prompt"), + patch("langchain_core.runnables.config.ensure_config", return_value=config), + ): + mock_run.return_value = "result" + await agent.run_task(task="generate-tasks", prompt="test", context=context) + + resolve_call_state = mock_resolve.call_args[0][0] + assert resolve_call_state["current_node"] == "generate_tasks" @pytest.mark.asyncio async def test_empty_tags_passed_as_none(self, agent: ForgeAgent) -> None: @@ -123,9 +140,7 @@ async def test_none_context_produces_trace_state_with_prompt_and_model( ) -> None: with ( patch.object(agent, "_run_agent", new_callable=AsyncMock) as mock_run, - patch( - "forge.integrations.agents.agent.resolve_trace_fields" - ) as mock_resolve, + patch("forge.integrations.agents.agent.resolve_trace_fields") as mock_resolve, patch("forge.integrations.agents.agent.load_prompt", return_value="prompt"), ): mock_run.return_value = "result" @@ -164,8 +179,6 @@ async def test_session_id_from_ticket_key(self, agent: ForgeAgent) -> None: patch("forge.integrations.agents.agent.load_prompt", return_value="prompt"), ): mock_run.return_value = "result" - await agent.run_task( - task="test", prompt="test", context={"ticket_key": "PROJ-42"} - ) + await agent.run_task(task="test", prompt="test", context={"ticket_key": "PROJ-42"}) assert mock_run.call_args.kwargs["session_id"] == "PROJ-42" diff --git a/tests/unit/integrations/agents/test_trace_forwarding.py b/tests/unit/integrations/agents/test_trace_forwarding.py index f28764d5..e3a9bd11 100644 --- a/tests/unit/integrations/agents/test_trace_forwarding.py +++ b/tests/unit/integrations/agents/test_trace_forwarding.py @@ -71,6 +71,14 @@ def test_does_not_mutate_input(self) -> None: _forward_trace_fields(context) assert context == original + def test_extracts_node_from_langgraph_config(self) -> None: + context = {"ticket_key": "PROJ-42"} + config = {"metadata": {"langgraph_node": "actual_executing_node"}} + with patch("langchain_core.runnables.config.ensure_config", return_value=config): + result = _forward_trace_fields(context) + assert result["current_node"] == "actual_executing_node" + assert result["ticket_key"] == "PROJ-42" + class TestGeneratePrdTraceForwarding: """generate_prd() uses _forward_trace_fields() and adds project_key.""" diff --git a/tests/unit/integrations/langfuse/test_fields.py b/tests/unit/integrations/langfuse/test_fields.py index 6623e20e..aa19ea7b 100644 --- a/tests/unit/integrations/langfuse/test_fields.py +++ b/tests/unit/integrations/langfuse/test_fields.py @@ -63,7 +63,49 @@ def test_project_id_no_dash(self) -> None: assert resolve_field(TracingField.PROJECT_ID, _make_state(ticket_key="NODASH")) is None def test_workflow_step(self) -> None: - assert resolve_field(TracingField.WORKFLOW_STEP, _make_state()) == "analyze_bug" + assert ( + resolve_field(TracingField.WORKFLOW_STEP, _make_state(retry_count=0)) + == "analyze_bug:initial_generation:unknown" + ) + + def test_workflow_step_epic_breakdown(self) -> None: + state = _make_state(current_node="decompose_epics", retry_count=0) + assert ( + resolve_field(TracingField.WORKFLOW_STEP, state) + == "decompose_epics:breakdown:epic_breakdown" + ) + + def test_workflow_step_generate_tasks(self) -> None: + state = _make_state(current_node="generate_tasks", retry_count=0) + assert resolve_field(TracingField.WORKFLOW_STEP, state) == "generate_tasks:breakdown:tasks" + + def test_workflow_step_spec_qa(self) -> None: + state = _make_state(current_node="spec_approval_gate", is_question=True, retry_count=0) + assert ( + resolve_field(TracingField.WORKFLOW_STEP, state) + == "spec_approval_gate:question_asking:spec" + ) + + def test_workflow_step_spec_revision(self) -> None: + state = _make_state(current_node="regenerate_spec", is_revision=True, retry_count=2) + assert ( + resolve_field(TracingField.WORKFLOW_STEP, state) + == "regenerate_spec:revision:spec:attempt-2" + ) + + def test_workflow_step_prd_approval_gate(self) -> None: + state = _make_state(current_node="prd_approval_gate", retry_count=0) + assert ( + resolve_field(TracingField.WORKFLOW_STEP, state) + == "prd_approval_gate:approval_gate:prd" + ) + + def test_workflow_step_explicit_artifact_type(self) -> None: + state = _make_state(current_node="answer_question", artifact_type="custom_type", retry_count=0) + assert ( + resolve_field(TracingField.WORKFLOW_STEP, state) + == "answer_question:question_asking:custom_type" + ) def test_workflow_step_missing(self) -> None: state = _make_state() @@ -313,14 +355,12 @@ class TestResolveTraceFields: """Integration: resolve configured fields from workflow state.""" def test_resolves_tags_and_metadata(self) -> None: - state = _make_state() + state = _make_state(retry_count=0) tag_fields = [TracingField.TICKET_TYPE, TracingField.PROJECT_ID, TracingField.WORKFLOW_STEP] metadata_fields = [TracingField.TICKET_KEY, TracingField.RETRY_COUNT] with ( - patch( - "forge.config.get_settings" - ) as mock_get_settings, + patch("forge.config.get_settings") as mock_get_settings, ): mock_settings = mock_get_settings.return_value type(mock_settings).trace_tag_fields = PropertyMock(return_value=tag_fields) @@ -328,8 +368,8 @@ def test_resolves_tags_and_metadata(self) -> None: tags, metadata = resolve_trace_fields(state) - assert tags == ["Bug", "PROJ", "analyze_bug"] - assert metadata == {"ticket_key": "PROJ-42", "retry_count": "3"} + assert tags == ["Bug", "PROJ", "analyze_bug:initial_generation:unknown"] + assert metadata == {"ticket_key": "PROJ-42", "retry_count": "0"} def test_skips_missing_fields(self) -> None: state = _make_state() @@ -337,9 +377,7 @@ def test_skips_missing_fields(self) -> None: tag_fields = [TracingField.TICKET_TYPE, TracingField.REPO] metadata_fields = [TracingField.PR_NUMBER] - with patch( - "forge.config.get_settings" - ) as mock_get_settings: + with patch("forge.config.get_settings") as mock_get_settings: mock_settings = mock_get_settings.return_value type(mock_settings).trace_tag_fields = PropertyMock(return_value=tag_fields) type(mock_settings).trace_metadata_fields = PropertyMock(return_value=metadata_fields) @@ -350,9 +388,7 @@ def test_skips_missing_fields(self) -> None: assert metadata == {"pr_number": "99"} def test_empty_config_returns_empty(self) -> None: - with patch( - "forge.config.get_settings" - ) as mock_get_settings: + with patch("forge.config.get_settings") as mock_get_settings: mock_settings = mock_get_settings.return_value type(mock_settings).trace_tag_fields = PropertyMock(return_value=[]) type(mock_settings).trace_metadata_fields = PropertyMock(return_value=[]) @@ -366,9 +402,7 @@ def test_system_prompt_length_in_metadata(self) -> None: state = _make_state(system_prompt_length=4523) metadata_fields = [TracingField.SYSTEM_PROMPT_LENGTH] - with patch( - "forge.config.get_settings" - ) as mock_get_settings: + with patch("forge.config.get_settings") as mock_get_settings: mock_settings = mock_get_settings.return_value type(mock_settings).trace_tag_fields = PropertyMock(return_value=[]) type(mock_settings).trace_metadata_fields = PropertyMock(return_value=metadata_fields) @@ -382,9 +416,7 @@ def test_llm_model_in_tags(self) -> None: state = _make_state(llm_model="claude-sonnet-4-6-20250514") tag_fields = [TracingField.LLM_MODEL] - with patch( - "forge.config.get_settings" - ) as mock_get_settings: + with patch("forge.config.get_settings") as mock_get_settings: mock_settings = mock_get_settings.return_value type(mock_settings).trace_tag_fields = PropertyMock(return_value=tag_fields) type(mock_settings).trace_metadata_fields = PropertyMock(return_value=[])