feat(SKY-9206): copilot attribution columns on workflows + workflow_runs (#5669)

tests/unit/forge/sdk/db/test_workflow_attribution_columns.py

@@ -0,0 +1,237 @@
+"""Regression tests for copilot attribution columns."""
+
+from __future__ import annotations
+
+from typing import Any, AsyncGenerator
+
+import pytest
+import pytest_asyncio
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from skyvern.forge.sdk.core import skyvern_context
+from skyvern.forge.sdk.db.agent_db import AgentDB
+from skyvern.forge.sdk.db.models import Base
+
+
+@pytest_asyncio.fixture
+async def db_engine() -> AsyncGenerator[Any]:
+    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    yield engine
+    await engine.dispose()
+
+
+@pytest_asyncio.fixture
+async def agent_db(db_engine: Any) -> AsyncGenerator[AgentDB]:
+    yield AgentDB(database_string="sqlite+aiosqlite:///:memory:", debug_enabled=True, db_engine=db_engine)
+
+
+@pytest_asyncio.fixture
+async def org_id(agent_db: AgentDB) -> str:
+    org = await agent_db.organizations.create_organization(
+        organization_name="Attribution Org",
+        domain="attribution.test",
+    )
+    return org.organization_id
+
+
+@pytest.mark.asyncio
+async def test_create_workflow_without_attribution_defaults_to_none(agent_db: AgentDB, org_id: str) -> None:
+    workflow = await agent_db.workflows.create_workflow(
+        title="plain-create",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+    )
+    assert workflow.created_by is None
+    assert workflow.edited_by is None
+
+
+@pytest.mark.asyncio
+async def test_create_workflow_stamps_attribution_when_passed(agent_db: AgentDB, org_id: str) -> None:
+    workflow = await agent_db.workflows.create_workflow(
+        title="copilot-create",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+        created_by="copilot",
+        edited_by="copilot",
+    )
+    assert workflow.created_by == "copilot"
+    assert workflow.edited_by == "copilot"
+
+
+@pytest.mark.asyncio
+async def test_update_workflow_omit_attribution_preserves_stamps(agent_db: AgentDB, org_id: str) -> None:
+    workflow = await agent_db.workflows.create_workflow(
+        title="seed",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+        created_by="copilot",
+        edited_by="copilot",
+    )
+    # Omit created_by / edited_by — the repo must NOT touch either column.
+    await agent_db.workflows.update_workflow(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+        title="renamed",
+    )
+    reread = await agent_db.workflows.get_workflow(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+    )
+    assert reread is not None
+    assert reread.created_by == "copilot"
+    assert reread.edited_by == "copilot"
+
+
+@pytest.mark.asyncio
+async def test_update_workflow_explicit_none_clears_attribution(agent_db: AgentDB, org_id: str) -> None:
+    # _UNSET sentinel distinguishes omit (preserve) from None (clear); rollback relies on this.
+    workflow = await agent_db.workflows.create_workflow(
+        title="seed",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+        created_by="copilot",
+        edited_by="copilot",
+    )
+    await agent_db.workflows.update_workflow(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+        created_by=None,
+        edited_by=None,
+    )
+    reread = await agent_db.workflows.get_workflow(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+    )
+    assert reread is not None
+    assert reread.created_by is None
+    assert reread.edited_by is None
+
+
+@pytest.mark.asyncio
+async def test_update_workflow_and_reconcile_explicit_none_clears_attribution(agent_db: AgentDB, org_id: str) -> None:
+    # Reconcile path must honor the same omit/None semantics as update_workflow.
+    from skyvern.forge.sdk.workflow.models.workflow import WorkflowDefinition
+
+    workflow = await agent_db.workflows.create_workflow(
+        title="seed",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+        created_by="copilot",
+        edited_by="copilot",
+    )
+    await agent_db.workflows.update_workflow_and_reconcile_definition_params(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+        workflow_definition=WorkflowDefinition(parameters=[], blocks=[]),
+        created_by=None,
+        edited_by=None,
+    )
+    reread = await agent_db.workflows.get_workflow(
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+    )
+    assert reread is not None
+    assert reread.created_by is None
+    assert reread.edited_by is None
+
+
+@pytest.mark.asyncio
+async def test_create_workflow_run_without_session_id_defaults_to_none(agent_db: AgentDB, org_id: str) -> None:
+    # No ambient skyvern_context; no explicit param — copilot_session_id stays NULL.
+    workflow = await agent_db.workflows.create_workflow(
+        title="wf",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+    )
+    run = await agent_db.workflow_runs.create_workflow_run(
+        workflow_permanent_id=workflow.workflow_permanent_id,
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+    )
+    assert run.copilot_session_id is None
+
+
+@pytest.mark.asyncio
+async def test_create_workflow_run_explicit_session_id_persists(agent_db: AgentDB, org_id: str) -> None:
+    workflow = await agent_db.workflows.create_workflow(
+        title="wf",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+    )
+    run = await agent_db.workflow_runs.create_workflow_run(
+        workflow_permanent_id=workflow.workflow_permanent_id,
+        workflow_id=workflow.workflow_id,
+        organization_id=org_id,
+        copilot_session_id="chat_abc123",
+    )
+    assert run.copilot_session_id == "chat_abc123"
+
+
+@pytest.mark.asyncio
+async def test_create_workflow_run_ignores_ambient_context(agent_db: AgentDB, org_id: str) -> None:
+    # Ambient-context resolution lives in the service layer, not the repo. Repo trusts the param.
+    workflow = await agent_db.workflows.create_workflow(
+        title="wf",
+        workflow_definition={"parameters": [], "blocks": []},
+        organization_id=org_id,
+    )
+    ambient = skyvern_context.SkyvernContext(copilot_session_id="chat_from_ctx")
+    with skyvern_context.scoped(ambient):
+        run = await agent_db.workflow_runs.create_workflow_run(
+            workflow_permanent_id=workflow.workflow_permanent_id,
+            workflow_id=workflow.workflow_id,
+            organization_id=org_id,
+        )
+    assert run.copilot_session_id is None
+
+
+# ---------------------------------------------------------------------------
+# Stub-heuristic regression coverage
+# ---------------------------------------------------------------------------
+
+
+def _make_workflow_stub(*, version: int, created_by: str | None, block_count: int) -> Any:
+    blocks = [object()] * block_count
+    definition = type("D", (), {"blocks": blocks})()
+    return type(
+        "W",
+        (),
+        {"version": version, "created_by": created_by, "workflow_definition": definition},
+    )()
+
+
+def test_is_copilot_born_stub_true_on_version_one_empty_unstamped() -> None:
+    from skyvern.forge.sdk.copilot.attribution import is_copilot_born_initial_write
+
+    wf = _make_workflow_stub(version=1, created_by=None, block_count=0)
+    assert is_copilot_born_initial_write(wf) is True
+
+
+def test_is_copilot_born_stub_false_on_later_version() -> None:
+    # v1 is the only version that can be copilot-born; cleared v2+ would otherwise false-positive.
+    from skyvern.forge.sdk.copilot.attribution import is_copilot_born_initial_write
+
+    wf = _make_workflow_stub(version=2, created_by=None, block_count=0)
+    assert is_copilot_born_initial_write(wf) is False
+
+
+def test_is_copilot_born_stub_false_on_already_stamped() -> None:
+    from skyvern.forge.sdk.copilot.attribution import is_copilot_born_initial_write
+
+    wf = _make_workflow_stub(version=1, created_by="copilot", block_count=0)
+    assert is_copilot_born_initial_write(wf) is False
+
+
+def test_is_copilot_born_stub_false_on_non_empty_definition() -> None:
+    from skyvern.forge.sdk.copilot.attribution import is_copilot_born_initial_write
+
+    wf = _make_workflow_stub(version=1, created_by=None, block_count=3)
+    assert is_copilot_born_initial_write(wf) is False
+
+
+def test_is_copilot_born_stub_false_on_none() -> None:
+    from skyvern.forge.sdk.copilot.attribution import is_copilot_born_initial_write
+
+    assert is_copilot_born_initial_write(None) is False

tests/unit/test_copilot_session_injection.py

@@ -138,11 +138,7 @@ def test_mcp_to_copilot_error() -> None:
 
 
 class TestMcpBrowserContextBridge:
-    """Bridge-specific behavior of mcp_browser_context (not scoped_session).
-
-    Covers: copilot session registry, API-key override install/reset, and the
-    teardown guarantees that must hold under every failure mode.
-    """
+    """Bridge-specific behavior of mcp_browser_context."""
 
     def _install_happy_path_mocks(
         self, monkeypatch: pytest.MonkeyPatch
@@ -374,6 +370,7 @@ class TestUpdateWorkflowDirect:
 
         mock_wf_service = MagicMock()
         mock_wf_service.update_workflow_definition = AsyncMock()
+        mock_wf_service.get_workflow = AsyncMock(return_value=None)
         monkeypatch.setattr("skyvern.forge.sdk.copilot.tools.app.WORKFLOW_SERVICE", mock_wf_service)
 
         yaml_str = "title: Test\nworkflow_definition:\n  blocks: []"
@@ -405,6 +402,7 @@ class TestUpdateWorkflowDirect:
 
         mock_wf_service = MagicMock()
         mock_wf_service.update_workflow_definition = AsyncMock()
+        mock_wf_service.get_workflow = AsyncMock(return_value=None)
         monkeypatch.setattr("skyvern.forge.sdk.copilot.tools.app.WORKFLOW_SERVICE", mock_wf_service)
 
         result = await _update_workflow({"workflow_yaml": "title: Test"}, ctx)

tests/unit/test_copilot_session_span_tag.py

@@ -0,0 +1,146 @@
+"""Tests for the copilot.session_id span attribute on LLM spans."""
+
+from __future__ import annotations
+
+from types import ModuleType, SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock
+
+from skyvern.forge.sdk.api.llm.api_handler_factory import _enrich_llm_span
+from skyvern.forge.sdk.core import skyvern_context
+from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
+
+
+def _call_enrich(span: MagicMock) -> None:
+    _enrich_llm_span(
+        span,
+        model="gpt-5",
+        prompt_name="workflow-copilot",
+        prompt_tokens=10,
+        completion_tokens=20,
+        reasoning_tokens=0,
+        cached_tokens=0,
+        latency_ms=100,
+        llm_cost=0.001,
+    )
+
+
+def _set_attribute_keys(span: MagicMock) -> list[str]:
+    return [call.args[0] for call in span.set_attribute.call_args_list if call.args]
+
+
+class TestEnrichLlmSpan:
+    def test_stamps_attribute_when_context_has_session_id(self) -> None:
+        span = MagicMock()
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="chat_xyz")):
+            _call_enrich(span)
+        span.set_attribute.assert_any_call("copilot.session_id", "chat_xyz")
+
+    def test_no_attribute_when_context_has_no_session_id(self) -> None:
+        span = MagicMock()
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id=None)):
+            _call_enrich(span)
+        assert "copilot.session_id" not in _set_attribute_keys(span)
+
+    def test_no_attribute_when_no_ambient_context(self) -> None:
+        span = MagicMock()
+        skyvern_context.reset()
+        _call_enrich(span)
+        assert "copilot.session_id" not in _set_attribute_keys(span)
+
+
+class _FakeAgentSpanData:
+    def __init__(self, name: str = "workflow-copilot") -> None:
+        self.name = name
+
+
+class _FakeGenerationSpanData:
+    pass
+
+
+class _FakeFunctionSpanData:
+    def __init__(self, name: str = "some_tool") -> None:
+        self.name = name
+
+
+def _install_patch(monkeypatch: Any) -> Any:
+    # Wire ModuleType stubs for the full logfire chain — sys.modules entries alone aren't enough.
+    import sys
+
+    from skyvern.forge.sdk.copilot import tracing_setup
+
+    def _fake_original(span_data: Any, msg_template: str) -> dict[str, Any]:
+        attrs: dict[str, Any] = {}
+        if isinstance(span_data, _FakeAgentSpanData):
+            attrs["name"] = span_data.name
+        if isinstance(span_data, _FakeFunctionSpanData):
+            attrs["name"] = span_data.name
+        return attrs
+
+    class _FakeWrapper:
+        @staticmethod
+        def create_span(*args: Any, **kwargs: Any) -> Any:
+            return None
+
+    logfire_mod = ModuleType("logfire")
+    internal_mod = ModuleType("logfire._internal")
+    integrations_mod = ModuleType("logfire._internal.integrations")
+    oai_mod = ModuleType("logfire._internal.integrations.openai_agents")
+    oai_mod.attributes_from_span_data = _fake_original  # type: ignore[attr-defined]
+    oai_mod.LogfireTraceProviderWrapper = _FakeWrapper  # type: ignore[attr-defined]
+    logfire_mod._internal = internal_mod  # type: ignore[attr-defined]
+    internal_mod.integrations = integrations_mod  # type: ignore[attr-defined]
+    integrations_mod.openai_agents = oai_mod  # type: ignore[attr-defined]
+
+    monkeypatch.setitem(sys.modules, "logfire", logfire_mod)
+    monkeypatch.setitem(sys.modules, "logfire._internal", internal_mod)
+    monkeypatch.setitem(sys.modules, "logfire._internal.integrations", integrations_mod)
+    monkeypatch.setitem(sys.modules, "logfire._internal.integrations.openai_agents", oai_mod)
+    monkeypatch.setitem(
+        sys.modules,
+        "agents",
+        SimpleNamespace(
+            AgentSpanData=_FakeAgentSpanData,
+            GenerationSpanData=_FakeGenerationSpanData,
+            FunctionSpanData=_FakeFunctionSpanData,
+        ),
+    )
+
+    tracing_setup._patch_agent_span_attributes()
+    return oai_mod.attributes_from_span_data
+
+
+class TestPatchedSpanAttributes:
+    def test_stamps_on_agent_span_when_context_has_session_id(self, monkeypatch: Any) -> None:
+        patched = _install_patch(monkeypatch)
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="chat_xyz")):
+            attrs = patched(_FakeAgentSpanData(), "Agent run: {name!r}")
+        assert attrs["copilot.session_id"] == "chat_xyz"
+
+    def test_stamps_on_generation_span_when_context_has_session_id(self, monkeypatch: Any) -> None:
+        patched = _install_patch(monkeypatch)
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="chat_xyz")):
+            attrs = patched(_FakeGenerationSpanData(), "Generation")
+        assert attrs["copilot.session_id"] == "chat_xyz"
+
+    def test_stamps_on_function_span_when_context_has_session_id(self, monkeypatch: Any) -> None:
+        patched = _install_patch(monkeypatch)
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="chat_xyz")):
+            attrs = patched(_FakeFunctionSpanData(), "Function call")
+        assert attrs["copilot.session_id"] == "chat_xyz"
+
+    def test_no_attribute_when_context_has_no_session_id(self, monkeypatch: Any) -> None:
+        patched = _install_patch(monkeypatch)
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id=None)):
+            attrs_agent = patched(_FakeAgentSpanData(), "Agent run: {name!r}")
+            attrs_gen = patched(_FakeGenerationSpanData(), "Generation")
+            attrs_fn = patched(_FakeFunctionSpanData(), "Function call")
+        assert "copilot.session_id" not in attrs_agent
+        assert "copilot.session_id" not in attrs_gen
+        assert "copilot.session_id" not in attrs_fn
+
+    def test_no_attribute_when_no_ambient_context(self, monkeypatch: Any) -> None:
+        patched = _install_patch(monkeypatch)
+        skyvern_context.reset()
+        attrs = patched(_FakeAgentSpanData(), "Agent run: {name!r}")
+        assert "copilot.session_id" not in attrs

tests/unit/test_copilot_task_block_rejection.py

@@ -303,6 +303,7 @@ async def test_update_workflow_preserves_legacy_task_block_under_unchanged_label
         patch("skyvern.forge.sdk.copilot.tools._process_workflow_yaml", return_value=fake_workflow),
         patch("skyvern.forge.sdk.copilot.tools.app") as mock_app,
     ):
+        mock_app.WORKFLOW_SERVICE.get_workflow = AsyncMock(return_value=None)
         mock_app.WORKFLOW_SERVICE.update_workflow_definition = AsyncMock()
         result = await _update_workflow({"workflow_yaml": submitted}, ctx)
 
@@ -343,6 +344,7 @@ async def test_update_workflow_allows_all_allowed_block_types() -> None:
         patch("skyvern.forge.sdk.copilot.tools._process_workflow_yaml", return_value=fake_workflow),
         patch("skyvern.forge.sdk.copilot.tools.app") as mock_app,
     ):
+        mock_app.WORKFLOW_SERVICE.get_workflow = AsyncMock(return_value=None)
         mock_app.WORKFLOW_SERVICE.update_workflow_definition = AsyncMock()
         result = await _update_workflow({"workflow_yaml": submitted}, ctx)
 

tests/unit/test_workflow_copilot_session_context.py

@@ -0,0 +1,49 @@
+"""Tests for the bind_copilot_session_id context manager."""
+
+from __future__ import annotations
+
+import pytest
+
+from skyvern.forge.sdk.core import skyvern_context
+from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
+from skyvern.forge.sdk.routes.workflow_copilot import bind_copilot_session_id
+
+
+class TestBindCopilotSessionId:
+    def test_sets_id_during_scope_when_ambient_context_present(self) -> None:
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id=None)):
+            with bind_copilot_session_id("chat_xyz"):
+                ctx = skyvern_context.current()
+                assert ctx is not None
+                assert ctx.copilot_session_id == "chat_xyz"
+
+    def test_restores_prior_value_on_normal_exit(self) -> None:
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="outer")):
+            with bind_copilot_session_id("inner"):
+                assert skyvern_context.current().copilot_session_id == "inner"  # type: ignore[union-attr]
+            assert skyvern_context.current().copilot_session_id == "outer"  # type: ignore[union-attr]
+
+    def test_restores_prior_value_when_body_raises(self) -> None:
+        class _Boom(RuntimeError):
+            pass
+
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="outer")):
+            with pytest.raises(_Boom):
+                with bind_copilot_session_id("inner"):
+                    raise _Boom("body raised")
+            assert skyvern_context.current().copilot_session_id == "outer"  # type: ignore[union-attr]
+
+    def test_noop_when_chat_id_is_none(self) -> None:
+        with skyvern_context.scoped(SkyvernContext(copilot_session_id="outer")):
+            with bind_copilot_session_id(None):
+                # No overwrite — the outer value must stick.
+                assert skyvern_context.current().copilot_session_id == "outer"  # type: ignore[union-attr]
+            assert skyvern_context.current().copilot_session_id == "outer"  # type: ignore[union-attr]
+
+    def test_noop_when_no_ambient_context(self) -> None:
+        skyvern_context.reset()
+        # Helper must not raise when there is no context to mutate — the
+        # copilot route should still function, just without the tag.
+        with bind_copilot_session_id("chat_xyz"):
+            assert skyvern_context.current() is None
+        assert skyvern_context.current() is None