Add DataSource parameter for artifacts endpoints (#18)

* Add DataSource parameter for artifacts endpoints
* PR feedback for artifacts endpoints

Author: sciapanCA

src/tests/test_artifact_relationships.py

@@ -372,6 +372,117 @@ async def test_explicit_profile_maps_correctly(self, mock_get_api_key):
 
         call_args = mock_client.post.call_args
         assert call_args[1]["json"]["profile"] == "InheritanceOnly"
+        # No data_source supplied => omitted from the body.
+        assert "dataSource" not in call_args[1]["json"]
+
+    @pytest.mark.asyncio
+    @patch("tools.artifact_relationships.get_api_key_from_context")
+    async def test_forwards_data_source(self, mock_get_api_key):
+        mock_get_api_key.return_value = "test_key"
+
+        ctx = MagicMock(spec=Context)
+        ctx.debug = AsyncMock()
+        ctx.error = AsyncMock()
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "sourceIdentifier": "id",
+            "profile": "CallsOnly",
+            "found": True,
+            "relationships": [],
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mock_response
+
+        mock_context = MagicMock()
+        mock_context.client = mock_client
+        mock_context.base_url = "https://app.codealive.ai"
+        ctx.request_context.lifespan_context = mock_context
+
+        await get_artifact_relationships(
+            ctx=ctx,
+            identifier="id",
+            data_source="backend",
+        )
+
+        assert mock_client.post.call_args[1]["json"]["dataSource"] == "backend"
+
+    @pytest.mark.asyncio
+    @patch("tools.artifact_relationships.get_api_key_from_context")
+    async def test_whitespace_data_source_omitted(self, mock_get_api_key):
+        """A whitespace-only data_source normalizes to None: not sent to the backend
+        and not echoed in the not-found hint (preserves the 409-on-ambiguity fallback)."""
+        mock_get_api_key.return_value = "test_key"
+
+        ctx = MagicMock(spec=Context)
+        ctx.debug = AsyncMock()
+        ctx.error = AsyncMock()
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "sourceIdentifier": "id",
+            "profile": "CallsOnly",
+            "found": False,
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mock_response
+
+        mock_context = MagicMock()
+        mock_context.client = mock_client
+        mock_context.base_url = "https://app.codealive.ai"
+        ctx.request_context.lifespan_context = mock_context
+
+        result = await get_artifact_relationships(
+            ctx=ctx,
+            identifier="id",
+            data_source="   ",
+        )
+
+        assert "dataSource" not in mock_client.post.call_args[1]["json"]
+        # The confusing `... in data source "   "` hint must not appear.
+        assert '"   "' not in result["hint"]
+
+    @pytest.mark.asyncio
+    @patch("tools.artifact_relationships.get_api_key_from_context")
+    async def test_ambiguous_409_surfaces_candidate_data_sources(self, mock_get_api_key):
+        import httpx
+
+        mock_get_api_key.return_value = "test_key"
+
+        ctx = MagicMock(spec=Context)
+        ctx.debug = AsyncMock()
+        ctx.error = AsyncMock()
+
+        mock_response = MagicMock()
+        mock_response.status_code = 409
+        mock_response.text = (
+            '{"detail": "Identifier matches 2 data sources: '
+            "Name='backend' Id='ds-main', Name='backend-legacy' Id='ds-master'\"}"
+        )
+        mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
+            "Conflict", request=MagicMock(), response=mock_response
+        )
+
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mock_response
+
+        mock_context = MagicMock()
+        mock_context.client = mock_client
+        mock_context.base_url = "https://app.codealive.ai"
+        ctx.request_context.lifespan_context = mock_context
+
+        with pytest.raises(ToolError) as exc:
+            await get_artifact_relationships(ctx=ctx, identifier="org/repo::path::Symbol")
+
+        message = str(exc.value)
+        assert "409" in message
+        # The candidate data sources from the backend 409 must be surfaced, plus the data_source retry hint.
+        assert "backend" in message and "backend-legacy" in message
+        assert "data_source" in message
 
     @pytest.mark.asyncio
     async def test_empty_identifier_raises_tool_error(self):
@@ -446,3 +557,21 @@ async def test_not_found_response_renders_correctly(self, mock_get_api_key):
 
         assert data["found"] is False
         assert "relationships" not in data
+
+    def test_not_found_hint_with_data_source_suggests_retry_or_omit(self):
+        payload = _build_relationships_dict(
+            {"sourceIdentifier": "org/repo::path::S", "profile": "CallsOnly", "found": False},
+            data_source="backend",
+        )
+        hint = payload["hint"]
+        assert "backend" in hint
+        assert "data_source" in hint
+        assert "omit" in hint.lower()
+
+    def test_not_found_hint_without_data_source_is_generic(self):
+        payload = _build_relationships_dict(
+            {"sourceIdentifier": "org/repo::path::S", "profile": "CallsOnly", "found": False},
+        )
+        hint = payload["hint"]
+        assert "data_source" not in hint
+        assert "fresh identifier" in hint

src/tests/test_fetch_artifacts.py

@@ -365,6 +365,39 @@ def test_hint_absent_when_no_artifacts_have_content(self):
         assert "<hint>" not in result
 
 
+class TestBuildArtifactsXmlDataSourceMissHint:
+    """When a data_source was supplied but nothing was found, hint to retry or drop it."""
+
+    def test_hint_when_data_source_scoped_returns_nothing(self):
+        data = {"artifacts": [
+            {"identifier": "repo::a.ts::F", "content": None, "contentByteSize": None},
+        ]}
+        result = _build_artifacts_xml(data, data_source="backend")
+        assert "<hint>" in result
+        assert "backend" in result
+        # Guides toward the two recovery moves.
+        assert "data_source" in result
+        assert "omit" in result.lower()
+
+    def test_hint_when_empty_artifacts_and_data_source(self):
+        result = _build_artifacts_xml({"artifacts": []}, data_source="ds-main")
+        assert "ds-main" in result and "<hint>" in result
+
+    def test_no_miss_hint_when_data_source_resolved_content(self):
+        data = {"artifacts": [
+            {"identifier": "repo::a.ts::F", "content": "code", "contentByteSize": 4},
+        ]}
+        result = _build_artifacts_xml(data, data_source="backend")
+        assert "omit data_source" not in result
+
+    def test_no_miss_hint_without_data_source(self):
+        data = {"artifacts": [
+            {"identifier": "repo::a.ts::F", "content": None, "contentByteSize": None},
+        ]}
+        result = _build_artifacts_xml(data)
+        assert "<hint>" not in result
+
+
 @pytest.mark.asyncio
 @patch('tools.fetch_artifacts.get_api_key_from_context')
 async def test_fetch_artifacts_returns_xml(mock_get_api_key):
@@ -476,6 +509,81 @@ async def test_fetch_artifacts_posts_correct_body(mock_get_api_key):
     body = call_args.kwargs["json"]
     assert body["identifiers"] == ["id1", "id2"]
     assert "names" not in body
+    # No data_source supplied => the field is omitted (preserves the 409-on-ambiguity fallback).
+    assert "dataSource" not in body
+
+
+@pytest.mark.asyncio
+@patch('tools.fetch_artifacts.get_api_key_from_context')
+async def test_fetch_artifacts_forwards_data_source(mock_get_api_key):
+    """data_source (Name or Id) is forwarded as the DataSource body field when provided."""
+    mock_get_api_key.return_value = "test_key"
+
+    ctx = MagicMock(spec=Context)
+    ctx.info = AsyncMock()
+    ctx.warning = AsyncMock()
+    ctx.error = AsyncMock()
+
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"artifacts": []}
+    mock_response.raise_for_status = MagicMock()
+
+    mock_client = AsyncMock()
+    mock_client.post.return_value = mock_response
+
+    mock_codealive_context = MagicMock()
+    mock_codealive_context.client = mock_client
+    mock_codealive_context.base_url = "https://app.codealive.ai"
+
+    ctx.request_context.lifespan_context = mock_codealive_context
+    ctx.request_context.headers = {"authorization": "Bearer test_key"}
+
+    await fetch_artifacts(
+        ctx=ctx,
+        identifiers=["id1"],
+        data_source="backend",
+    )
+
+    body = mock_client.post.call_args.kwargs["json"]
+    assert body["dataSource"] == "backend"
+
+
+@pytest.mark.asyncio
+@patch('tools.fetch_artifacts.get_api_key_from_context')
+async def test_fetch_artifacts_whitespace_data_source_omitted(mock_get_api_key):
+    """A whitespace-only data_source normalizes to None: not sent to the backend
+    and not echoed in the not-found hint (preserves the 409-on-ambiguity fallback)."""
+    mock_get_api_key.return_value = "test_key"
+
+    ctx = MagicMock(spec=Context)
+    ctx.info = AsyncMock()
+    ctx.warning = AsyncMock()
+    ctx.error = AsyncMock()
+
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"artifacts": []}
+    mock_response.raise_for_status = MagicMock()
+
+    mock_client = AsyncMock()
+    mock_client.post.return_value = mock_response
+
+    mock_codealive_context = MagicMock()
+    mock_codealive_context.client = mock_client
+    mock_codealive_context.base_url = "https://app.codealive.ai"
+
+    ctx.request_context.lifespan_context = mock_codealive_context
+    ctx.request_context.headers = {"authorization": "Bearer test_key"}
+
+    result = await fetch_artifacts(
+        ctx=ctx,
+        identifiers=["id1"],
+        data_source="   ",
+    )
+
+    body = mock_client.post.call_args.kwargs["json"]
+    assert "dataSource" not in body
+    # The confusing `... data source "   "` hint must not appear.
+    assert '"   "' not in result
 
 
 @pytest.mark.asyncio

src/tests/test_response_transformer.py

@@ -294,10 +294,34 @@ def test_data_preservation(self):
         assert first["identifier"] == "CodeAlive-AI/codealive-mcp::src/tools/search.py::codebase_search"
         assert first["contentByteSize"] == 8500
         assert first["description"] == "Main search function"
+        # Data-source identity must be surfaced (not stripped) so the agent can feed it back
+        # as `data_source` to disambiguate a branch-blind identifier.
+        assert first["dataSource"] == {"id": "685b21230e3822f4efa9d073", "name": "codealive-mcp"}
 
         assert second["path"] == "README.md"
         assert second["kind"] == "Chunk"
         assert second["description"] == "Search documentation section"
+        assert second["dataSource"] == {"id": "685b21230e3822f4efa9d073", "name": "codealive-mcp"}
+
+    def test_grep_transform_surfaces_data_source(self):
+        response = {
+            "results": [
+                {
+                    "kind": "File",
+                    "identifier": "owner/repo::src/auth.py",
+                    "location": {"path": "src/auth.py"},
+                    "matchCount": 1,
+                    "matches": [
+                        {"lineNumber": 3, "startColumn": 0, "endColumn": 4, "lineText": "auth"}
+                    ],
+                    "dataSource": {"type": "repository", "id": "ds-main", "name": "backend"},
+                }
+            ]
+        }
+
+        result = transform_grep_response(response)
+
+        assert result["results"][0]["dataSource"] == {"id": "ds-main", "name": "backend"}
 
     def test_grep_transform_preserves_match_previews(self):
         response = {