[AIGTWY-5288] Reduce Desktop integration test flakiness

docs/claude-desktop-integration-testing.md

@@ -34,8 +34,9 @@ prevention, unrelated-profile preservation, deletion, and prior-selection restor
 It executes the exact generated helper and validates its credential against the live models API.
 Both the system.ai and subscription-relay journeys use the generated helper credentials and
 profile headers for real Messages inference. Each installs Cowork through the real Claude CLI,
-adds system.ai.dbsql through public ug commands, and uses the installed plugin's command to
-discover tools and execute a read-only query with a computed receipt. Both check repeat
+adds system.ai.web_search through public ug commands, and uses the installed plugin's command to
+discover tools and find Databricks' official website. Both require HTTP 200 and nonempty assistant
+text for inference, without matching generated wording. Both check repeat
 configure, preservation of unrelated settings, and profile/plugin cleanup on revert, within
 the existing Linux/Windows integration lanes and their model workspace.
 CUJ1 and CUJ2 compare complete Desktop profile/metadata JSON for system.ai models and non-relay

docs/claude-desktop.md

@@ -393,7 +393,7 @@ Slow metadata discovery does not delay unrelated MCP responses.
 Two automated Desktop journeys run in the existing Linux/Windows Claude integration jobs:
 system.ai models and subscription-relay MPS. Each asserts the generated profile, executes its
 credential helper, completes real Messages inference, installs Cowork through the Claude CLI,
-and discovers/calls DBSQL through the installed plugin using a read-only computed query.
+and discovers/calls web_search through the installed plugin to find Databricks' official website.
 Both check repeated configuration, preservation of unrelated settings, and revert cleanup.
 All calls use the journey's existing model workspace. CUJ1 and CUJ2 compare complete Desktop
 profile/metadata JSON for system.ai models and non-relay Anthropic MPS. CUJ3's existing configure

tests/README.md

@@ -301,7 +301,9 @@ they do not establish native Windows or Linux support.
 The installed-CLI integration suite covers two complete Desktop journeys on Linux and Windows:
 system.ai models and subscription-relay Anthropic MPS. Both verify profile configuration,
 credential-helper execution, live Messages inference, installed Cowork MCP calls, preservation,
-and revert. See [integration coverage](integration/README.md) for prerequisites and evidence.
+and revert. Inference requires HTTP 200 and nonempty assistant text; MCP calls use web_search
+to find Databricks' official website without SQL compute. See
+[integration coverage](integration/README.md) for prerequisites and evidence.
 CUJ1/CUJ2 compare complete Desktop profile/metadata JSON for system.ai models and non-relay
 Anthropic MPS. CUJ3's existing picker journey checks its incompatible scoped aliases preserve
 all Desktop profile/metadata bytes after its existing configure call.

tests/integration/README.md

@@ -47,29 +47,29 @@ path, preserves an unrelated profile and metadata, then runs the exact generated
 helper and validates its token against the paginated live models API. It requires system.ai
 Claude models without an MPS override and uses the profile's credentials and headers for a small
 Messages request on `system.ai.claude-haiku-4-5`, requiring HTTP 200, a nonempty response model ID,
-and the exact fresh marker. The service may return a provider model ID rather than its alias;
-this does not pin or prove its backing model revision. The relay journey selects
+and nonempty assistant text without matching its wording. The service may return a provider model
+ID rather than its alias; this does not pin or prove its backing model revision. The relay journey selects
 the real subscription provider through configure and completes
 Messages inference using the subscription token and separate Databricks header returned by its
 helper. Non-relay MPS Desktop configuration is asserted in CUJ2's existing MPS configuration
 journey; no dedicated Desktop MPS integration test is needed.
 
 Each of these two journeys supplies pre-existing Desktop account directories as inputs and
-requires real Claude CLI installation of the Cowork plugin. It registers system.ai.dbsql through
-`ug mcp add --names system.ai.dbsql`, then runs the command from the installed plugin with the
-MCP SDK. The check initializes the protocol, discovers execute_sql/poll_sql_result, executes
-`SELECT sha2('<fresh run id>', 256)`, polls pending statements, and requires the computed receipt
-in the successful result's single typed cell. This proves execution rather than an echoed query.
+requires real Claude CLI installation of the Cowork plugin. It registers system.ai.web_search through
+`ug mcp add --names system.ai.web_search`, then runs the command from the installed plugin with the
+MCP SDK. The check initializes the protocol, discovers web_search, and searches for Databricks'
+official website. A successful result must contain a URL on databricks.com or www.databricks.com;
+the response wording and search rankings are not asserted. No SQL compute is needed.
 Both journeys verify reconfiguration preserves the profile/plugin and revert removes UG's entries
 while retaining unrelated profiles and metadata. A separate local personal plugin is installed
 through the real Claude CLI before configuration; its registry entry and installed manifest must
 survive installation, reconfiguration, and revert unchanged. Revert also removes UG's marketplace
 registration and generated source directory.
 
 The existing Linux/Windows **Full journeys · Claude** jobs run both complete journeys against
-the same workspace as their model calls. The workspace must expose system.ai.dbsql and allow
-the test identity to run a constant read-only query on its default SQL compute. No dedicated
-Desktop MCP test, separate fixture-workspace step, or extra credentials are required.
+the same workspace as their model calls. The workspace must expose system.ai.web_search and allow
+the test identity to invoke it. No dedicated Desktop MCP test, separate fixture-workspace step,
+or extra credentials are required.
 
 These tests do not launch Desktop or establish GUI, skill execution, browser consent,
 token renewal, or injected authentication-failure/reconnect recovery. CI supplies the supported

tests/integration/test_ug_configure_claude_desktop.py

@@ -18,7 +18,6 @@
     installed_cowork_servers,
     prepare_cowork_home,
 )
-from utils.evidence import fresh_id
 from utils.harness import process_group_options, stop_process
 from utils.provider_catalog import fetch_anthropic_parent_catalog
 from utils.terminal import ConfigureTerminal, TerminalProcess
@@ -172,14 +171,13 @@ def _run_credential_helper(session, profile):
 
 
 def _assert_desktop_message(session, profile, model, headers, evidence_name):
-    marker = fresh_id()
     request = urllib.request.Request(
         profile["inferenceGatewayBaseUrl"] + "/v1/messages",
         data=json.dumps(
             {
                 "model": model,
                 "max_tokens": 32,
-                "messages": [{"role": "user", "content": f"Reply with exactly: {marker}"}],
+                "messages": [{"role": "user", "content": "Reply with a short greeting."}],
             }
         ).encode(),
         headers={
@@ -197,15 +195,15 @@ def _assert_desktop_message(session, profile, model, headers, evidence_name):
         details = session.redact(error.read().decode("utf-8", errors="replace"))
         session.record(evidence_name, {"status": error.code, "error": details})
         raise AssertionError(f"Desktop Messages failed ({error.code}): {details}") from None
-    session.record(evidence_name, {"marker": marker, "requested_model": model, "response": message})
+    session.record(
+        evidence_name,
+        {"status": response.status, "requested_model": model, "response": message},
+    )
     assert message["type"] == "message" and message["role"] == "assistant"
     # A model-service alias can return its provider's model ID.
     assert isinstance(message["model"], str) and message["model"].strip()
     assert message["stop_reason"] == "end_turn"
-    assert (
-        "".join(part["text"] for part in message["content"] if part["type"] == "text").strip()
-        == marker
-    )
+    assert "".join(part["text"] for part in message["content"] if part["type"] == "text").strip()
     return message
 
 
@@ -216,7 +214,7 @@ def test_ug_configure_claude_desktop_profile_lifecycle(live_session, workspace):
     changing unrelated profile data, and later removes only its own profile while restoring the
     prior selection. The exact generated helper returns a credential accepted by the live models
     API and completes a Messages request on a configured system.ai Claude model. The installed
-    Cowork plugin discovers DBSQL tools and executes a read-only query with a computed receipt.
+    Cowork plugin discovers web_search and finds Databricks' official website.
     Repeat configure preserves both configurations; an unrelated personal plugin survives
     configure and revert unchanged. This checks native Linux/Windows commands
     and protocols, not Desktop GUI, skill execution, or browser login; CI supplies a bearer.
@@ -271,10 +269,10 @@ def test_ug_configure_claude_desktop_profile_lifecycle(live_session, workspace):
     )
 
     _assert_desktop_selection(config_library, configured=True)
-    session.run("mcp", "add", "--names", "system.ai.dbsql", timeout=240)
+    session.run("mcp", "add", "--names", "system.ai.web_search", timeout=240)
     servers = installed_cowork_servers(session, root)
     _assert_unrelated_cowork_plugin(root, unrelated_plugin)
-    anyio.run(assert_desktop_mcp, session, servers["dbsql"], workspace, fresh_id())
+    anyio.run(assert_desktop_mcp, session, servers["web_search"], workspace)
 
     session.run(*configure_args, timeout=240)
     assert _desktop_profile(session, config_library, workspace) == profile
@@ -292,9 +290,9 @@ def test_ug_configure_claude_desktop_subscription_relay(
 
     Expected: the profile selects relay mode and its installed helper returns both the Claude
     subscription token and Databricks gateway header. Those exact credentials and profile
-    headers complete a real Anthropic Messages request with a fresh marker on the configured
-    subscription model. Its installed Cowork plugin discovers DBSQL tools and returns a computed
-    query receipt. Reconfigure preserves the profile/plugin; revert restores unrelated settings
+    headers complete a real Anthropic Messages request on the configured subscription model.
+    Its installed Cowork plugin discovers web_search and finds Databricks' official website.
+    Reconfigure preserves the profile/plugin; revert restores unrelated settings
     and unregisters only UG's plugin, preserving an unrelated personal plugin. No Desktop GUI,
     connector skill, or browser OAuth is exercised.
     """
@@ -374,10 +372,10 @@ def test_ug_configure_claude_desktop_subscription_relay(
     # This dated native relay ID must be returned unchanged; it is not a service alias.
     assert message["model"] == CLAUDE_RELAY_TEST_MODEL
     _assert_desktop_selection(config_library, configured=True)
-    session.run("mcp", "add", "--names", "system.ai.dbsql", timeout=240)
+    session.run("mcp", "add", "--names", "system.ai.web_search", timeout=240)
     servers = installed_cowork_servers(session, root)
     _assert_unrelated_cowork_plugin(root, unrelated_plugin)
-    anyio.run(assert_desktop_mcp, session, servers["dbsql"], workspace, fresh_id())
+    anyio.run(assert_desktop_mcp, session, servers["web_search"], workspace)
 
     with ConfigureTerminal(session, "claude", command, "reconfigure-desktop-relay") as configure:
         configure.select_agent("Claude Code")

tests/integration/utils/desktop.py

@@ -2,10 +2,11 @@
 
 import json
 import os
+import re
 import sys
 import tempfile
-from hashlib import sha256
 from pathlib import Path
+from urllib.parse import urlparse
 
 import anyio
 from mcp import ClientSession, StdioServerParameters
@@ -40,56 +41,38 @@ def installed_cowork_servers(session, root: Path) -> dict:
     return json.loads((plugin / manifest["mcpServers"]).read_text())["mcpServers"]
 
 
-async def assert_desktop_mcp(session, registration: dict, workspace: str, run_id: str) -> None:
+async def assert_desktop_mcp(session, registration: dict, workspace: str) -> None:
     assert registration["type"] == "stdio", registration
     arguments = registration["args"]
     assert arguments[arguments.index("--url") + 1] == (
-        workspace + "/ai-gateway/mcp-services/system.ai.dbsql"
+        workspace + "/ai-gateway/mcp-services/system.ai.web_search"
     ), registration
     parameters = StdioServerParameters(**registration)
     parameters.env = {**session.env, **(parameters.env or {})}
-    responses = []
     with tempfile.TemporaryFile(mode="w+", encoding="utf-8") as diagnostics:
         try:
-            with anyio.fail_after(180):
+            with anyio.fail_after(60):
                 async with stdio_client(parameters, errlog=diagnostics) as streams:
                     async with ClientSession(*streams) as client:
                         await client.initialize()
                         tools = {tool.name for tool in (await client.list_tools()).tools}
-                        assert {"execute_sql", "poll_sql_result"} <= tools, tools
-
-                        async def call(name, arguments):
-                            result = await client.call_tool(name, arguments)
-                            # MCP 1.x and 2.x expose different Python attribute names.
-                            response = result.model_dump(by_alias=True)
-                            responses.append(response)
-                            session.record(
-                                "desktop-mcp.json", {"run_id": run_id, "responses": responses}
-                            )
-                            assert not response["isError"], session.redact(json.dumps(response))
-                            blocks = [
-                                block["text"]
-                                for block in response["content"]
-                                if block["type"] == "text"
-                            ]
-                            assert len(blocks) == 1, response
-                            return json.loads(blocks[0])
-
-                        result = await call(
-                            "execute_sql",
-                            {"query": f"SELECT sha2('{run_id}', 256) AS desktop_mcp_receipt"},
+                        assert "web_search" in tools, tools
+                        result = await client.call_tool(
+                            "web_search", {"query": "Databricks official website"}
+                        )
+                        # MCP 1.x and 2.x expose different Python attribute names.
+                        response = result.model_dump(by_alias=True)
+                        session.record("desktop-mcp.json", response)
+                        assert not response["isError"], session.redact(json.dumps(response))
+                        text = "\n".join(
+                            block["text"]
+                            for block in response["content"]
+                            if block["type"] == "text"
                         )
-                        while result["status"]["state"] in {"PENDING", "RUNNING"}:
-                            await anyio.sleep(1)
-                            result = await call(
-                                "poll_sql_result", {"statement_id": result["statement_id"]}
-                            )
-                        assert result["status"]["state"] == "SUCCEEDED", result
-                        assert result["manifest"]["total_row_count"] == 1, result
-                        # A computed receipt proves execution rather than an echoed SQL request.
-                        assert result["result"]["data_typed_array"] == [
-                            {"values": [{"str": sha256(run_id.encode()).hexdigest()}]}
-                        ], result
+                        assert any(
+                            urlparse(url).hostname in {"databricks.com", "www.databricks.com"}
+                            for url in re.findall(r"https?://[^\s)<>]+", text)
+                        ), session.redact(text)
         finally:
             diagnostics.seek(0)
             stderr = session.redact(diagnostics.read())