Add sign-in challenges and the authenticate tool to aitools mcp serve

acceptance/cmd/aitools/mcp/serve-authenticate-tool/.databrickscfg

@@ -0,0 +1,3 @@
+[logfood]
+host = $DATABRICKS_HOST
+token = test-token-logfood

acceptance/cmd/aitools/mcp/serve-authenticate-tool/out.test.toml

@@ -0,0 +1 @@
+Cloud = false

acceptance/cmd/aitools/mcp/serve-authenticate-tool/output.txt

@@ -0,0 +1,43 @@
+
+=== Before sign-in, tools/list answers with one tool, authenticate; it opened no sign-in
+
+>>> jq -c select(.id==2) | [.result.tools[].name] serve.stdout.txt
+["authenticate"]
+
+>>> cat stderr-before-authenticate.txt
+No install record for system.ai.github on profile "logfood"; using its current host [DATABRICKS_URL]
+
+=== authenticate waits for the sign-in, then answers and asks the agent to re-list
+
+>>> jq -c select(.id==3) | .result | {isError, text: .content[0].text} serve.stdout.txt
+{"isError":false,"text":"Signed in to system.ai.github on profile logfood. Tools are available now; retry the user's request in this turn."}
+
+>>> jq -c select(.method=="notifications/tools/list_changed") serve.stdout.txt
+{"jsonrpc":"2.0","method":"notifications/tools/list_changed"}
+
+>>> cat serve.stderr.txt
+No install record for system.ai.github on profile "logfood"; using its current host [DATABRICKS_URL]
+Sign in with the same account you use for Databricks profile logfood.
+Open this URL to sign in to system.ai.github:
+[DATABRICKS_URL]/mcp-service-login?name=system.ai.github&o=[WORKSPACE_ID]
+Waiting for sign-in...
+Signed in to system.ai.github on profile logfood.
+Public repositories only until a repository owner installs https://github.com/apps/databricks-github-connector.
+
+=== After the sign-in, tools/list passes through and tools/call reaches the service
+
+>>> jq -c select(.id==4) | [.result.tools[].name] serve.stdout.txt
+["echo"]
+
+>>> jq -c select(.id==5) | .result.content[0].text serve.stdout.txt
+"hello"
+
+=== Stdout holds only JSON-RPC lines
+
+>>> jq -c .jsonrpc serve.stdout.txt
+"2.0"
+"2.0"
+"2.0"
+"2.0"
+"2.0"
+"2.0"

acceptance/cmd/aitools/mcp/serve-authenticate-tool/requests-after.ndjson

@@ -0,0 +1,2 @@
+{"jsonrpc":"2.0","id":4,"method":"tools/list"}
+{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"echo","arguments":{"text":"hello"}}}

acceptance/cmd/aitools/mcp/serve-authenticate-tool/requests-authenticate.ndjson

@@ -0,0 +1 @@
+{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"authenticate","arguments":{}}}

acceptance/cmd/aitools/mcp/serve-authenticate-tool/requests-list.ndjson

@@ -0,0 +1,3 @@
+{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test-client","version":"0.0.1"}}}
+{"jsonrpc":"2.0","method":"notifications/initialized"}
+{"jsonrpc":"2.0","id":2,"method":"tools/list"}

acceptance/cmd/aitools/mcp/serve-authenticate-tool/script

@@ -0,0 +1,43 @@
+# Isolate HOME so the proxy finds no state file and no agent config, as on a fresh machine.
+sethome home
+envsubst < .databrickscfg > out && mv out .databrickscfg
+export DATABRICKS_CONFIG_FILE=.databrickscfg
+# No browser: the proxy prints the sign-in URL and does not open anything.
+export BROWSER=none
+# The sign-in is polled every 2 s, so the waits below need more attempts than the default.
+export RETRY_MAX_ATTEMPTS=60
+
+# The fake gets the GitHub service with its OAuth connection; the caller has no credential for it yet.
+MSYS_NO_PATHCONV=1 $CLI api post "/api/2.1/unity-catalog/mcp-services?parent=schemas/system.ai&mcp_service_id=github" --profile logfood --json '{"comment":"Repositories, issues, pull requests","config":{"source_connection":{"name":"connections/github","options":{"oauth_provider":"github"}}}}' &> LOG.create-github
+
+# Stdin stays open until the last request is answered: the proxy shuts down at EOF and cancels a call
+# still in flight after 2 s. Each step waits for the previous answer, so the order is deterministic.
+# The last steps wait for ids 4 and 5 too, so the proxy has answered them before stdin closes.
+{
+  cat requests-list.ndjson
+  retry --until '"id":2' cat serve.stdout.txt > /dev/null
+  cat serve.stderr.txt > stderr-before-authenticate.txt
+  cat requests-authenticate.ndjson
+  retry --until "Waiting for sign-in" cat serve.stderr.txt > /dev/null
+  MSYS_NO_PATHCONV=1 $CLI api post /api/2.1/unity-catalog/mcp-services/system.ai.github/user-credentials --profile logfood --json '{}' &> LOG.sign-in-credential
+  retry --until '"id":3' cat serve.stdout.txt > /dev/null
+  cat requests-after.ndjson
+  retry --until '"id":4' cat serve.stdout.txt > /dev/null
+  retry --until '"id":5' cat serve.stdout.txt > /dev/null
+} | $CLI aitools mcp serve --profile logfood --service system.ai.github --agent claude-code > serve.stdout.txt 2> serve.stderr.txt
+
+title "Before sign-in, tools/list answers with one tool, authenticate; it opened no sign-in"
+trace jq -c 'select(.id==2) | [.result.tools[].name]' serve.stdout.txt
+trace cat stderr-before-authenticate.txt
+
+title "authenticate waits for the sign-in, then answers and asks the agent to re-list"
+trace jq -c 'select(.id==3) | .result | {isError, text: .content[0].text}' serve.stdout.txt
+trace jq -c 'select(.method=="notifications/tools/list_changed")' serve.stdout.txt
+trace cat serve.stderr.txt
+
+title "After the sign-in, tools/list passes through and tools/call reaches the service"
+trace jq -c 'select(.id==4) | [.result.tools[].name]' serve.stdout.txt
+trace jq -c 'select(.id==5) | .result.content[0].text' serve.stdout.txt
+
+title "Stdout holds only JSON-RPC lines"
+trace jq -c '.jsonrpc' serve.stdout.txt

acceptance/cmd/aitools/mcp/serve-authenticate-tool/test.toml

@@ -0,0 +1,3 @@
+# The script leaves the captured streams of the proxy and the snapshot of its stderr before authenticate in this
+# directory. The output is the trace lines, so the captures are ignored.
+Ignore = ["stderr-before-authenticate.txt", "serve.stdout.txt", "serve.stderr.txt"]

cmd/aitools/mcp_add.go

@@ -18,6 +18,7 @@ import (
 	"github.com/databricks/cli/libs/aitools/mcp/agentconfig"
 	"github.com/databricks/cli/libs/aitools/mcpconfig"
 	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/browser"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/telemetry/protos"
@@ -879,7 +880,7 @@ func signInInstalled(ctx context.Context, errOut io.Writer, w *databricks.Worksp
 			continue
 		}
 		url := mcp.SignInURL(w.Config.Host, plans[i].service.FullName, req.workspaceID)
-		state, _, err := signIn(ctx, errOut, w, req.profile, req.pinned, url, plans[i].service)
+		state, _, _, err := signIn(ctx, errOut, w, req.pinned, url, plans[i].service, signInOptions{Profile: req.profile, Open: browser.OpenOutputToStderr})
 		server.SignIn = &addSignInOutput{State: state}
 		if err != nil {
 			server.SignIn.ErrorCategory = errorCategory(err)

cmd/aitools/mcp_login.go

@@ -20,7 +20,10 @@ import (
 )
 
 // mcpLoginCommand is the command that signs a profile in to a service. The hints name it.
-const mcpLoginCommand = "databricks aitools mcp login"
+const (
+	mcpLoginCommand  = "databricks aitools mcp login"
+	mcpLogoutCommand = "databricks aitools mcp logout"
+)
 
 // Sign-in states of `mcp login` and `mcp logout`, reported in JSON.
 const (
@@ -180,7 +183,7 @@ func mcpLogin(ctx context.Context, errOut io.Writer, w *databricks.WorkspaceClie
 		return describeRequestError(out.Profile, pinnedFile, fmt.Errorf("cannot read the workspace ID of %s: %w", w.Config.Host, err))
 	}
 	url := mcp.SignInURL(out.Host, svc.FullName, workspaceID)
-	state, urlShown, err := signIn(ctx, errOut, w, out.Profile, pinnedFile, url, svc)
+	state, _, urlShown, err := signIn(ctx, errOut, w, pinnedFile, url, svc, signInOptions{Profile: out.Profile, Open: browser.OpenOutputToStderr})
 	out.State = state
 	// The URL is named when it was printed. A sign-in that failed before its page is printed names none.
 	if urlShown {
@@ -218,6 +221,11 @@ func mcpLoginHint(service, profileName, pinnedFile string) string {
 	return mcp.ConfigFileCommand(pinnedFile, mcpLoginCommand+" "+shellquote.BashArg(service)+" --profile "+shellquote.BashArg(profileName))
 }
 
+// mcpLogoutHint returns the command that removes profileName's sign-in to service. pinnedFile is as in mcpLoginHint.
+func mcpLogoutHint(service, profileName, pinnedFile string) string {
+	return mcp.ConfigFileCommand(pinnedFile, mcpLogoutCommand+" "+shellquote.BashArg(service)+" --profile "+shellquote.BashArg(profileName))
+}
+
 // lastCheckText names the failure of the last poll of a sign-in that timed out, for the timeout message. It is empty
 // when the last poll read the credential.
 func lastCheckText(timeout *mcp.SignInTimeoutError) string {
@@ -233,27 +241,38 @@ var (
 	signInWaitTimeout  = mcp.SignInWaitTimeout
 )
 
-// signIn runs the provider sign-in of svc for the profile and reports progress on errOut. It returns the sign-in state
-// name, whether url was printed, and the failure, which is categorized for JSON. `add --login` shares it.
-func signIn(ctx context.Context, errOut io.Writer, w *databricks.WorkspaceClient, profileName, pinnedFile, url string, svc mcp.Service) (state string, urlShown bool, err error) {
-	alreadySignedIn, urlShown, err := mcp.SignIn(ctx, w, url, svc.FullName, errOut, mcp.SignInOptions{
+// signInOptions is what a caller of signIn sets: the profile named in the notice before the page, the page opener, and
+// the poll callback. signIn adds the polling settings.
+type signInOptions struct {
+	Profile string
+	Open    func(ctx context.Context, url string) error
+	OnPoll  func()
+}
+
+// signIn runs the provider sign-in of svc for the profile named in opts and reports progress on errOut. It returns the
+// sign-in state name, whether the caller was already signed in, whether url was printed, and the failure, which is
+// categorized for JSON. `add --login` shares this.
+func signIn(ctx context.Context, errOut io.Writer, w *databricks.WorkspaceClient, pinnedFile, url string, svc mcp.Service, opts signInOptions) (state string, alreadySignedIn, urlShown bool, err error) {
+	profileName := opts.Profile
+	alreadySignedIn, urlShown, err = mcp.SignIn(ctx, w, url, svc.FullName, errOut, mcp.SignInOptions{
 		Profile: profileName,
 		Poll:    signInPollInterval,
 		Timeout: signInWaitTimeout,
-		Open:    browser.OpenOutputToStderr,
+		Open:    opts.Open,
+		OnPoll:  opts.OnPoll,
 	})
 	timeout, timedOut := errors.AsType[*mcp.SignInTimeoutError](err)
 	switch {
 	case timedOut:
-		return signInStateTimeout, urlShown, failure(categorySignInTimeout, err,
+		return signInStateTimeout, false, urlShown, failure(categorySignInTimeout, err,
 			"sign-in to %s did not complete in %g minutes%s. Finish at %s then retry: %s",
 			svc.FullName, signInWaitTimeout.Minutes(), lastCheckText(timeout), url, mcpLoginHint(svc.FullName, profileName, pinnedFile))
 	case errors.Is(err, mcp.ErrSignInFailed):
-		return signInStateFailed, urlShown, failure(protos.AitoolsErrorCategoryUncategorized, err,
+		return signInStateFailed, false, urlShown, failure(protos.AitoolsErrorCategoryUncategorized, err,
 			"sign-in to %s failed on profile %s. Run: %s",
 			svc.FullName, profileName, mcpLoginHint(svc.FullName, profileName, pinnedFile))
 	case err != nil:
-		return "", urlShown, describeRequestError(profileName, pinnedFile, err)
+		return "", false, urlShown, describeRequestError(profileName, pinnedFile, err)
 	}
 
 	if alreadySignedIn {
@@ -264,5 +283,5 @@ func signIn(ctx context.Context, errOut io.Writer, w *databricks.WorkspaceClient
 	if svc.Known != nil && svc.Known.Note != "" {
 		fmt.Fprintln(errOut, svc.Known.Note)
 	}
-	return signInStateActive, urlShown, nil
+	return signInStateActive, alreadySignedIn, urlShown, nil
 }

cmd/aitools/mcp_login_test.go

@@ -2,6 +2,7 @@ package aitools
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -344,7 +345,8 @@ func TestSignInTimeoutNamesTheLastFailedCheck(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx, _ := cmdio.NewTestContextWithStderr(t.Context())
-	state, _, err := signIn(ctx, io.Discard, w, "logfood", "", "https://example.test/login", mcp.Service{FullName: "system.ai.github", SignInRequired: true})
+	state, _, _, err := signIn(ctx, io.Discard, w, "", "https://example.test/login", mcp.Service{FullName: "system.ai.github", SignInRequired: true},
+		signInOptions{Profile: "logfood", Open: func(context.Context, string) error { return nil }})
 	require.Error(t, err)
 	assert.Equal(t, signInStateTimeout, state)
 	_, timedOut := errors.AsType[*mcp.SignInTimeoutError](err)