feat(mcp): accept locationCode/languageCode on Labs research tools (#113)

Co-authored-by: isatimur <isatimur@users.noreply.github.com>

src/server/mcp/tools/dataforseo-research-tools.market.test.ts

@@ -60,7 +60,13 @@ function setProject(market: { locationCode: number; languageCode: string }) {
   });
 }
 
-async function runRankedKeywords(args: { market?: { country: "US" } }) {
+type MarketArgs = {
+  market?: { country: "US" };
+  locationCode?: number;
+  languageCode?: string;
+};
+
+async function runRankedKeywords(args: MarketArgs) {
   const rankedKeywords = vi.fn().mockResolvedValue({
     items: [],
     totalCount: 0,
@@ -76,6 +82,20 @@ async function runRankedKeywords(args: { market?: { country: "US" } }) {
   return rankedKeywords;
 }
 
+async function runSerpCompetitors(args: MarketArgs) {
+  const serpCompetitors = vi.fn().mockResolvedValue([]);
+  mocks.createDataforseoClient.mockReturnValue({
+    labs: { serpCompetitors },
+  });
+  const { findSerpCompetitorsTool } =
+    await import("./dataforseo-research-tools");
+  await findSerpCompetitorsTool.handler(
+    { projectId: "project_1", keywords: ["seo"], ...args },
+    toolExtra,
+  );
+  return serpCompetitors;
+}
+
 describe("market resolution for Labs tools", () => {
   beforeEach(() => {
     vi.resetModules();
@@ -94,6 +114,85 @@ describe("market resolution for Labs tools", () => {
     );
   });
 
+  it("exposes explicit location and language selectors on both tool schemas", async () => {
+    const { findSerpCompetitorsTool, getRankedKeywordsTool } =
+      await import("./dataforseo-research-tools");
+
+    expect(getRankedKeywordsTool.config.inputSchema.locationCode).toBeDefined();
+    expect(getRankedKeywordsTool.config.inputSchema.languageCode).toBeDefined();
+    expect(
+      findSerpCompetitorsTool.config.inputSchema.locationCode,
+    ).toBeDefined();
+    expect(
+      findSerpCompetitorsTool.config.inputSchema.languageCode,
+    ).toBeDefined();
+  });
+
+  it("passes an explicit non-US market to both Labs tools", async () => {
+    setProject({ locationCode: 2704, languageCode: "vi" });
+
+    const rankedKeywords = await runRankedKeywords({
+      locationCode: 2756,
+      languageCode: "de",
+    });
+    expect(rankedKeywords).toHaveBeenCalledWith(
+      expect.objectContaining({ locationCode: 2756, languageCode: "de" }),
+    );
+
+    const serpCompetitors = await runSerpCompetitors({
+      locationCode: 2756,
+      languageCode: "de",
+    });
+    expect(serpCompetitors).toHaveBeenCalledWith(
+      expect.objectContaining({ locationCode: 2756, languageCode: "de" }),
+    );
+  });
+
+  it("uses the selected location's default language when only locationCode is explicit", async () => {
+    setProject({ locationCode: 2704, languageCode: "vi" });
+    const rankedKeywords = await runRankedKeywords({ locationCode: 2276 });
+    expect(rankedKeywords).toHaveBeenCalledWith(
+      expect.objectContaining({ locationCode: 2276, languageCode: "de" }),
+    );
+  });
+
+  it("prefers an explicit locationCode over the legacy market object", async () => {
+    setProject({ locationCode: 2704, languageCode: "vi" });
+    const rankedKeywords = await runRankedKeywords({
+      locationCode: 2756,
+      languageCode: "de",
+      market: { country: "US" },
+    });
+    expect(rankedKeywords).toHaveBeenCalledWith(
+      expect.objectContaining({ locationCode: 2756, languageCode: "de" }),
+    );
+  });
+
+  it("accepts a non-default language the location serves", async () => {
+    setProject({ locationCode: 2704, languageCode: "vi" });
+    // Switzerland serves fr/de/it; de is the default.
+    const serpCompetitors = await runSerpCompetitors({
+      locationCode: 2756,
+      languageCode: "fr",
+    });
+    expect(serpCompetitors).toHaveBeenCalledWith(
+      expect.objectContaining({ locationCode: 2756, languageCode: "fr" }),
+    );
+  });
+
+  it("rejects a language the location does not serve", async () => {
+    setProject({ locationCode: 2704, languageCode: "vi" });
+    await expect(
+      runRankedKeywords({ locationCode: 2276, languageCode: "fr" }),
+    ).rejects.toThrow("is not available for this location");
+  });
+
+  it("rejects an explicit non-Labs country before making a paid call", async () => {
+    await expect(
+      runRankedKeywords({ locationCode: 2352, languageCode: "is" }),
+    ).rejects.toThrow("Domain analytics is not available for this country");
+  });
+
   it("follows the project's default market when the market object is omitted", async () => {
     setProject({ locationCode: 2704, languageCode: "vi" });
     const rankedKeywords = await runRankedKeywords({});

src/server/mcp/tools/dataforseo-research-tools.ts

@@ -18,7 +18,10 @@ import {
   type McpTableColumn,
 } from "@/server/mcp/table";
 import { resolveLabsMarket, resolveMarket } from "@/shared/keyword-locations";
-import { assertLanguageForLocation } from "@/server/lib/market";
+import {
+  assertLabsLocationCode,
+  assertLanguageForLocation,
+} from "@/server/lib/market";
 import {
   DEFAULT_LOCATION_CODE,
   languageCodeSchema,
@@ -52,7 +55,7 @@ const marketSchema = z
   })
   .optional()
   .describe(
-    "Optional market object. Omitted = the project's default market (United States unless the project overrides it).",
+    "Legacy US selector. Prefer locationCode/languageCode for any Labs market. Explicit locationCode takes precedence; otherwise omitted = the project's default market.",
   );
 
 const nearSchema = z
@@ -126,6 +129,16 @@ const getRankedKeywordsInputSchema = {
     "Domain (no protocol/www) or absolute page URL to list ranked keywords for.",
   ),
   market: marketSchema,
+  locationCode: locationCodeSchema
+    .optional()
+    .describe(
+      "Country-level DataForSEO Labs location code. Defaults to the project's market; takes precedence over the legacy market object.",
+    ),
+  languageCode: languageCodeSchema
+    .optional()
+    .describe(
+      "Language for locationCode. Defaults to that location's primary language when locationCode overrides the project market.",
+    ),
   resultTypes: z
     .array(rankedResultTypeSchema)
     .min(1)
@@ -256,6 +269,16 @@ const findSerpCompetitorsInputSchema = {
     .max(100)
     .describe("Keywords whose SERPs are compared (1-100)."),
   market: marketSchema,
+  locationCode: locationCodeSchema
+    .optional()
+    .describe(
+      "Country-level DataForSEO Labs location code. Defaults to the project's market; takes precedence over the legacy market object.",
+    ),
+  languageCode: languageCodeSchema
+    .optional()
+    .describe(
+      "Language for locationCode. Defaults to that location's primary language when locationCode overrides the project market.",
+    ),
   resultTypes: z
     .array(serpCompetitorResultTypeSchema)
     .min(1)
@@ -349,19 +372,37 @@ const QUESTIONS_ANSWERS_MIN_RADIUS = 200;
 const QUESTIONS_ANSWERS_MAX_RADIUS = 199999;
 
 /**
- * Resolves the market selector to a Labs location + language. An explicit
+ * Resolves a Labs location + language. Explicit location/language fields win,
- * country wins; omitted inherits the project's default via resolveLabsMarket,
+ * followed by the legacy explicit-US selector; omitted fields inherit the
- * which keeps these Labs-only tools off an Ads-served project market.
+ * project's default via resolveLabsMarket, which keeps these Labs-only tools
+ * off an Ads-served project market. Validate before starting a paid request.
  */
 function resolveMarketSelector(
-  market: Market | undefined,
+  selector: {
+    market?: Market;
+    locationCode?: number;
+    languageCode?: string;
+  },
   project: { locationCode: number; languageCode: string },
 ): { locationCode: number; languageCode: string } {
-  if (market?.country != null) {
+  let resolved: { locationCode: number; languageCode: string };
+  if (selector.locationCode != null || selector.languageCode != null) {
+    resolved = resolveLabsMarket(
+      {
+        locationCode: selector.locationCode,
+        languageCode: selector.languageCode,
+      },
+      project,
+    );
+  } else if (selector.market?.country != null) {
     // The Zod enum already restricts explicit values to United States variants.
-    return { locationCode: DEFAULT_LOCATION_CODE, languageCode: "en" };
+    resolved = { locationCode: DEFAULT_LOCATION_CODE, languageCode: "en" };
+  } else {
+    resolved = resolveLabsMarket({}, project);
   }
-  return resolveLabsMarket({}, project);
+  assertLabsLocationCode(resolved.locationCode);
+  assertLanguageForLocation(resolved.locationCode, resolved.languageCode);
+  return resolved;
 }
 
 function formatCoordinate(value: number): string {
@@ -592,7 +633,7 @@ export const getRankedKeywordsTool = {
   config: {
     title: "Get ranked keywords",
     description:
-      "Returns exact keyword, URL, rank, search volume, CPC, intent, and traffic rows for a domain or page. Use this for strategy evidence; use get_domain_overview for aggregate domain footprint. Charges credits.",
+      "Returns market-specific keyword, URL, rank, search volume, CPC, intent, and traffic rows for a domain or page. Accepts country-level DataForSEO Labs location/language codes. Use this for strategy evidence; use get_domain_overview for aggregate domain footprint. Charges credits.",
     inputSchema: getRankedKeywordsInputSchema,
     outputSchema: {
       keywords: z.array(looseObjectOutputSchema),
@@ -608,7 +649,7 @@ export const getRankedKeywordsTool = {
   handler: withMcpProjectAuth(async (args: GetRankedKeywordsArgs, context) => {
     const client = createDataforseoClient(context.billing);
     const targetIsPage = /^https?:\/\//.test(args.target);
-    const market = resolveMarketSelector(args.market, context.project);
+    const market = resolveMarketSelector(args, context.project);
     const keywords = await client.domain.rankedKeywords({
       target: args.target,
       locationCode: market.locationCode,
@@ -773,7 +814,7 @@ export const findSerpCompetitorsTool = {
   config: {
     title: "Find SERP competitors",
     description:
-      "Compares domains competing in Google results for a supplied keyword set. Useful for market and search-intelligence reports; not radius-based local SEO. Charges credits.",
+      "Compares domains competing in Google results for a supplied keyword set in a country-level DataForSEO Labs market. Accepts location/language codes; not radius-based local SEO. Charges credits.",
     inputSchema: findSerpCompetitorsInputSchema,
     outputSchema: {
       competitors: z.array(looseObjectOutputSchema),
@@ -788,7 +829,7 @@ export const findSerpCompetitorsTool = {
   handler: withMcpProjectAuth(
     async (args: FindSerpCompetitorsArgs, context) => {
       const client = createDataforseoClient(context.billing);
-      const market = resolveMarketSelector(args.market, context.project);
+      const market = resolveMarketSelector(args, context.project);
       const competitors = await client.labs.serpCompetitors({
         keywords: args.keywords,
         locationCode: market.locationCode,