chore: fix tracing for text chat mode

api/tests/test_workflow_run_cost.py

@@ -0,0 +1,87 @@
+from datetime import UTC, datetime
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from api.services.pricing import workflow_run_cost as workflow_run_cost_mod
+from api.services.pricing.workflow_run_cost import (
+    build_workflow_run_cost_info,
+    calculate_workflow_run_cost,
+)
+
+
+def _make_workflow_run():
+    return SimpleNamespace(
+        id=123,
+        workflow_id=456,
+        mode="textchat",
+        created_at=datetime.now(UTC),
+        usage_info={
+            "llm": {},
+            "tts": {},
+            "stt": {},
+            "call_duration_seconds": 7,
+        },
+        cost_info={},
+        workflow=SimpleNamespace(
+            organization_id=42,
+            user=SimpleNamespace(selected_organization_id=42),
+        ),
+    )
+
+
+@pytest.mark.asyncio
+async def test_build_workflow_run_cost_info_does_not_update_org_usage(monkeypatch):
+    workflow_run = _make_workflow_run()
+    get_org = AsyncMock(return_value=SimpleNamespace(id=42, price_per_second_usd=1.5))
+    update_usage = AsyncMock()
+
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client, "get_organization_by_id", get_org
+    )
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client, "update_usage_after_run", update_usage
+    )
+
+    cost_info = await build_workflow_run_cost_info(workflow_run)
+
+    assert cost_info is not None
+    assert cost_info["call_duration_seconds"] == 7
+    assert "cost_breakdown" in cost_info
+    assert "dograh_token_usage" in cost_info
+    assert cost_info["charge_usd"] == 10.5
+    update_usage.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_calculate_workflow_run_cost_keeps_org_usage_side_effect_in_wrapper(
+    monkeypatch,
+):
+    workflow_run = _make_workflow_run()
+    get_org = AsyncMock(return_value=SimpleNamespace(id=42, price_per_second_usd=None))
+    update_run = AsyncMock()
+    update_usage = AsyncMock()
+
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client,
+        "get_workflow_run_by_id",
+        AsyncMock(return_value=workflow_run),
+    )
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client, "get_organization_by_id", get_org
+    )
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client, "update_workflow_run", update_run
+    )
+    monkeypatch.setattr(
+        workflow_run_cost_mod.db_client, "update_usage_after_run", update_usage
+    )
+
+    await calculate_workflow_run_cost(workflow_run.id)
+
+    update_run.assert_awaited_once()
+    saved_kwargs = update_run.await_args.kwargs
+    assert saved_kwargs["run_id"] == workflow_run.id
+    assert "cost_breakdown" in saved_kwargs["cost_info"]
+    update_usage.assert_awaited_once()