Repository navigation
Add aitools mcp serve, the stdio proxy for managed MCP servers - #7550
Draft
simonfaltum wants to merge 32 commits into
Draft
simonfaltum wants to merge 32 commits into
simonfaltum wants to merge 32 commits into
Conversation
simonfaltum
force-pushed
the
aitools-mcp/06b-serve
branch
from
November 12, 2026 12:37
694ba3c to
c3af024
Compare
Contributor
|
Commit: Regression test reportSelected 12 changed tests (limit=20, 0 not selected) 12 tests: regression=12 unreleased=0 coverage=0 failing=0
TestAccept/cmd/aitools/mcp/add-agent-not-installed: branch ✅ | main (8e70c8b) ❌ | latest (v1.26.0) ❌main (8e70c8b): latest (v1.26.0): TestAccept/cmd/aitools/mcp/add-default-config: branch ✅ | main (8e70c8b) ❌ | latest (v1.26.0) ❌main (8e70c8b): latest (v1.26.0): TestAccept/cmd/aitools/mcp/add-text: branch ✅ | main (8e70c8b) ❌ | latest (v1.26.0) ❌main (8e70c8b): latest (v1.26.0): TestAccept/cmd/aitools/mcp/add-two-agents: branch ✅ | main (8e70c8b) ❌ | latest (v1.26.0) ❌main (8e70c8b): latest (v1.26.0): TestAccept/cmd/aitools/mcp/add-two-profiles: branch ✅ | main (8e70c8b) ❌ | latest (v1.26.0) ❌main (8e70c8b): latest (v1.26.0): 7 more failure details omitted, see the full report in the run. |
Collaborator
✅ Integration test reportCommit: 10450c9
Top 9 slowest tests (at least 2 minutes):
|
simonfaltum
force-pushed
the
aitools-mcp/06a-bridge
branch
from
November 12, 2026 14:36
4ecc28f to
2f43e8a
Compare
simonfaltum
force-pushed
the
aitools-mcp/06b-serve
branch
from
November 12, 2026 14:37
eec9d1a to
5b1d279
Compare
simonfaltum
force-pushed
the
aitools-mcp/06a-bridge
branch
from
November 12, 2026 17:08
2f43e8a to
3c43783
Compare
simonfaltum
force-pushed
the
aitools-mcp/06b-serve
branch
from
November 12, 2026 17:10
5b1d279 to
9187003
Compare
simonfaltum
force-pushed
the
aitools-mcp/06a-bridge
branch
2 times, most recently
from
November 12, 2026 22:00
4688f64 to
ed4190f
Compare
simonfaltum
force-pushed
the
aitools-mcp/06b-serve
branch
from
November 13, 2026 00:15
9187003 to
48fac6f
Compare
`databricks aitools mcp serve --profile <p> --service <full name>` is the
command each installed MCP entry runs. It binds to the profile, refusing to
start when the profile is gone, re-pointed at another host, or cannot mint
a token (exit 2 with the message on stderr), then bridges newline-delimited
JSON-RPC on stdin/stdout to the service endpoint.
- Binding runs before stdin is read; the bound config is set on the command
context so the exit-time telemetry upload authenticates with it.
- The proxy sets the User-Agent extras aitools-mcp/proxy and
aitools-mcp-client/<agent>, plus the workspace routing header, and appends
the bound workspace and profile to the initialize instructions.
- A token refresh failure mid-session answers that request with an error
naming the profile, then exits 2.
- The root command skips the stale-version warning for serve, and the exit
status now comes from root.ExitStatus through main.
- Fake /ai-gateway/mcp-services/{name} endpoint in testserver, with
acceptance tests for stdout purity and binding errors.
Co-authored-by: Isaac <no-reply@databricks.com>
Stdout carries the MCP protocol for serve, so an explicit --log-file stdout would mix log lines into it. serve now fails at flag validation, as it does for --output, and the binding-errors scenario covers both flags. Co-authored-by: Isaac <no-reply@databricks.com>
macOS wc pads its count with spaces, so the golden differed from Linux. Strip the padding the way acceptance/ssh/bulk-transfer does. Co-authored-by: Isaac <no-reply@databricks.com>
Card 03 moved MCP records into the `mcp` section of ~/.databricks/aitools/skills/.state.json. serve-binding-errors wrote a v1 file at the old path, so the bind found no record and the proxy started. The fixture now writes the new file and pins its record to the test's DATABRICKS_CONFIG_FILE, as `add` does; otherwise card 03 reports a config-file mismatch before the host change. Golden changes, all from card 03: - The PROFILE_NOT_FOUND text names the config path as resolved (absolute) and has one space before "or remove", not three. - The PROFILE_HOST_CHANGED text and the no-record notice have no trailing period. Co-authored-by: Isaac <no-reply@databricks.com>
The proxy bridge now takes an http.RoundTripper. A nil one means http.DefaultTransport, whose dialer waits 30 s. serve passes a transport that keeps the profile's proxy and TLS settings, as cmd/apps/logs.go does, and dials with a 10 s connect timeout. Co-authored-by: Isaac <no-reply@databricks.com>
Card 03 and 04 changed the binding messages. A profile with no state record now says the server needs the profile and how to set it up. A host change names the server in the remove hint, not the profile. Co-authored-by: Isaac <no-reply@databricks.com>
ExitStatus honors an error with an ExitStatus() int method, found with errors.AsType, so serve no longer wraps its own status error. The method is not named ExitCode because os/exec's ExitError has that method, and a child process's code would become the CLI's exit status. mcp.BindError renames ExitCode to ExitStatus to match. The refresh failure is the AUTH_FAILED BindError, so it exits 2 as the binding failures do. The failed-execution log lines and the telemetry exit code now come from the same status, so the log no longer records exit_code 1 for an exit 2. Co-authored-by: Isaac <no-reply@databricks.com>
A command annotated with root.StdoutCarriesProtocol gets its log destination checked before the log file is opened. A log file that is stdout, by name or by identity (/dev/stdout), is refused, so the protocol on stdout is never shared with a log line. The stale-version warning is skipped for such commands, because an agent shows the server's stderr on every start. aitools mcp serve carries the annotation. Its own check for --log-file stdout and the match by command name are removed. Co-authored-by: Isaac <no-reply@databricks.com>
…ELETE A refresh failure stops the proxy only when it is a real token failure. A request whose own context ended (the client cancelled it, or shutdown began) fails that request alone, and the session DELETE at shutdown is best effort, so its auth failure is only logged. An unknown --agent no longer stops serve: it only feeds the User-Agent, and an entry written by a later release may name an agent this binary does not know. Unknown names are left out of the User-Agent and logged at debug. The first SIGINT or SIGTERM cancels the context and releases the signal handler before the context reports done, so a second signal during shutdown ends the process. The proxy ignores the profile's HTTPTransport: BoundConfig never sets it. The transport keeps the connect timeout and InsecureSkipVerify. Co-authored-by: Isaac <no-reply@databricks.com>
serve-binding-errors records the config file as the binary spells it (pwd -W on Git Bash), so the binding check matches on Windows. The unknown --agent case moves to serve-passthrough, which shows that such an agent is served and logged at debug. The out.* files are goldens; the stdout and stderr captures now go to LOG.* files. The last binding case checks a log file on stdout at debug level. The testserver's MCP endpoint names the JSON-RPC version once and documents its helpers. The proxy test's runOptions is documented. Co-authored-by: Isaac <no-reply@databricks.com>
… error A log file is refused for a protocol command when it is the process's stdout. The check now compares the opened file with os.Stdout: os.Stat of /dev/stdout does not match stdout on macOS when stdout is a pipe. The check runs after the file is opened and before the first log line. The log file is closed when Execute returns, so a test or an embedding caller does not keep it open (Windows cannot remove an open file). ExitStatus honors only the named type ExitStatusError, which serve returns for binding and auth failures. mcp.BindError no longer carries an exit status, so `mcp add` exits 1 in both output modes, as its SPEC says. An error with an ExitStatus method of its own no longer sets the status. Co-authored-by: Isaac <no-reply@databricks.com>
The check ran after initializeContext had installed the stdout logger, so a refused command still wrote its telemetry debug line to stdout. It now runs inside initializeContext, after the log file opens and before the default logger is set, and it closes the refused file. The log file is also recorded on the context here, so Execute closes it after the last log line. Co-authored-by: Isaac <no-reply@databricks.com>
…the service name - A token that cannot be refreshed mid-session is printed once: the bridge logs it on the request that needed it, and serve exits 2 without printing it again. The session DELETE after that failure is sent without a credential, so it does not log the refresh error a second time. - A state file that cannot be read is a binding failure (exit 2). - --service must be a full three-part name, with the message naming the flag. - An unknown --agent is reported on stderr once, not at debug level. - The bound profile is no longer set as the command's config: serve logs no telemetry yet, so the setting had no effect. Card 10a adds it with the telemetry it serves. - The signal context is released by the same stop that cancels the refresh guard, so a second signal during a refresh-failure shutdown ends the process. - serveTransport sets InsecureSkipVerify on the clone's TLS config instead of replacing it. - The fake MCP endpoint answers a client response (an id with no method) with 202 and no body. - The second-signal test asserts the child was ended by the signal, not by a panic or exit. Co-authored-by: Isaac <no-reply@databricks.com>
writesToStdout read the process's os.Stdout and os.Stderr. Under `go test -json` the testing package sets os.Stderr to os.Stdout, so the stdout cases were taken for stderr and the check passed the log file through. Take the two streams as parameters, and give the tests pipes and a separate stderr of their own. Production behavior is unchanged. Co-authored-by: Isaac <no-reply@databricks.com>
On Windows os.SameFile reports any two pipes or character devices as the same file, because their file ids are zero. With stdout piped, a null device or a named-pipe log file was refused as stdout. Compare by identity only for regular files there; elsewhere nothing changes. stdout itself is matched by pointer, so the stdout keyword is still refused on every OS. Co-authored-by: Isaac <no-reply@databricks.com>
When stdout and stderr are one file, as on a terminal, a log file opened as /dev/stderr or /dev/tty was refused as stdout, while --log-file stderr was allowed. The protocol already shares that file with the diagnostics, so the check cannot help there: a destination that is the same file as stderr is not refused. Other destinations are compared as before. Co-authored-by: Isaac <no-reply@databricks.com>
Execute closed the log file through the context of the command that ran. On a panic the command's context was never returned, so the file stayed open, and the default logger kept writing to it after Execute returned, or closed. Execute now keeps a slot in its own context that initializeContext fills, closes the file in a defer (which also runs on a panic), and restores the default logger that was in place before the command. Co-authored-by: Isaac <no-reply@databricks.com>
A caller whose credential is known to be unusable has no reason to send the session DELETE at shutdown: it would fail and log the failure again. Authenticate can now return ErrSkipRequest for the DELETE. The bridge then sends and logs nothing for it. Any other request that returns the sentinel is an error like any other. Co-authored-by: Isaac <no-reply@databricks.com>
After a token refresh failure, every later request asked for a token again and logged the login hint again, a request cancelled before its failure was logged left nothing on stderr, and a stdin error after the failure exited 1 and hid the auth status. Now the first failing request carries the hint, with the restart step; later requests fail at once with a short error that does not repeat it; the session DELETE is skipped (proxy.ErrSkipRequest); and serve prints one final line and exits 2, whatever the bridge returned. Co-authored-by: Isaac <no-reply@databricks.com>
… have The fake MCP endpoint answered any service name, so a wrong service name or path in the proxy's command still passed. A name with no service now gets 404, as the gateway answers a name that has none, and the tests and the serve acceptance case seed the service they use. Co-authored-by: Isaac <no-reply@databricks.com>
… logger The stale-version skip for a command whose stdout carries a protocol had no test. Move the decision into staleVersionSkipped, which takes the build info, and test it. In initializeContext the refused log file no longer needs a close: Execute closes the file it opened. Co-authored-by: Isaac <no-reply@databricks.com>
…its dialer The flag names of serve are the constants that add writes into agent configs, so an entry binds what add wrote; a test parses that argv with serve's flags. The dialer helper existed only for a test of a constant, so it is inlined and the test dropped. The binding status doc no longer cites a list command that is not in this tree. Co-authored-by: Isaac <no-reply@databricks.com>
The serve cases ran the command without printing it, so the flags under test did not appear in the goldens. Wrap each invocation in trace. Co-authored-by: Isaac <no-reply@databricks.com>
slog.SetDefault redirects the standard log package to the new handler, so after Execute closed the log file, log.Printf still wrote to it. Execute now saves and restores log.Writer() and log.Flags() with the default logger. Co-authored-by: Isaac <no-reply@databricks.com>
The stdout check ran for every protocol command, including the default stderr log, and a failed Stat returned an error that stopped the command, so `2>&-` made `serve` fail. A stat failure now means the stream is not stdout, and writesToStdout returns a bool. Co-authored-by: Isaac <no-reply@databricks.com>
Unity Catalog names are case-insensitive, and the fake's listing already matches them that way. The MCP endpoint looked the name up exactly, so a service addressed in another case got 404. Co-authored-by: Isaac <no-reply@databricks.com>
jq -s also accepts pretty-printed or several-per-line JSON, so the passthrough test did not check that each stdout line is one message. The filter now reads each line with fromjson. The binding-errors steps send stderr to the script output, so each traced command, its stderr and its exit code appear in order in the golden. Co-authored-by: Isaac <no-reply@databricks.com>
No behavior change. The refresh-failure tests share one fake endpoint and one Authenticate, and the exit-with-status error is built by one helper. Stale-version skipping is one predicate, the refresh-failure messages are constants in the tests, and the two no-body endpoint tests are one table. Comments are trimmed to their reasons, test comments start with the test name, and the silent_err.go import is back to its original form. Co-authored-by: Isaac <no-reply@databricks.com>
…ndows console case serve checks its binding only at start. A profile re-pointed to another host during the session keeps receiving the endpoint of the bound host, so that is documented on the command. The refresh hint's restart is needed because a new spawn re-runs the binding checks. The stdout check treats a Windows console as a file that is not stdout, so `--log-file stdout` is refused there even in a terminal. Co-authored-by: Isaac <no-reply@databricks.com>
depguard forbids the standard log package outside a few generators. The import is needed here because slog.SetDefault redirects that package, and libs/log has no way to restore it. The exception is annotated where it is used. Co-authored-by: Isaac <no-reply@databricks.com>
Serve stopped on any token refresh failure. A network error, a timeout or a 5xx from the token endpoint is not a credential the user must replace, so serve now fails only the request that needed the token, with the login hint, and keeps serving. The next request asks for a token again. A failure stops serve only when the token is rejected: the token endpoint answers 400 or 401 with invalid_grant or invalid_client, or the U2M client reports an invalid refresh token. Anything else is transient, including errors that cannot be classified. Co-authored-by: Isaac <no-reply@databricks.com>
A failure that is not a rejection, such as a network blip, says nothing about the credential, so the request's error no longer tells the user to log in. It says the next request tries again. Co-authored-by: Isaac <no-reply@databricks.com>
simonfaltum
force-pushed
the
aitools-mcp/06a-bridge
branch
from
November 13, 2026 15:05
6634449 to
5c42df0
Compare
simonfaltum
force-pushed
the
aitools-mcp/06b-serve
branch
from
November 13, 2026 15:05
7064bbc to
10450c9
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request is part of a stack that adds the hidden
databricks aitools mcpcommand group. It builds on #7540 and addsaitools mcp serve, the stdio command that each installed MCP entry runs.Changes
aitools mcp serve --profile <p> --service <full name> [--agent <name>], hidden. It runs the binding checks, then bridges stdin and stdout to the service endpoint with the profile's credentials.--outputand a log file on stdout are refused.mcp addstill exits 1 on a binding failure.libs/aitools/mcp/proxy:Authenticatemay returnErrSkipRequestto skip the session DELETE.cmd/root:ExitStatusError, whichExitStatushonors.Executecloses the log file and restores the default logger and standardlogoutput on return, also after a panic.libs/testserver: the fake MCP endpoint answers 404 for unknown services and matches names without regard to case.Why
Each entry that
mcp addwrites runsaitools mcp serve. It binds to one profile before it reads stdin, so a profile that points elsewhere at start fails with a message instead of sending a token to another host.Decisions to review
invalid_grantorinvalid_client, or the CLI's U2M client reporting an invalid refresh token. Everything else fails only its request, including network errors, timeouts, 5xx responses, and anything the check cannot classify. The reason is a laptop that wakes before the network is back: stopping there leaves the agent with a dead server until the entry restarts. The cost is that an unclassified error keeps failing requests instead of stopping serve.databricks auth login --profile P --host B, keeps sending its token to the host serve was bound to. The command's doc comment says so.--log-file stdoutworks on a Unix terminal. Windows refuses it even in a terminal, because a console is not a regular file there.--agentis accepted and left out of the User-Agent, since a later release may write agents this binary does not know.Tests
go test -raceoncmd/aitools,cmd/root,libs/aitoolsandlibs/testserver, also through the repo's gotestsum runner.cmd/root: the stdout refusal and check, including a closed stderr; logger restore after a panic and a normal return.cmd/aitools: a rejected token stops the proxy with exit 2, the hint, and no later token request. A transient failure fails only its request, and the next request succeeds while serve runs to stdin close (exit 0).TestCredentialRejectedchecks each error shape.serve-passthroughchecks each stdout line is one JSON-RPC message;serve-binding-errorschecks exit 2 and the stderr text.Not tested: a real token refresh end to end, because the fake cannot fail one. The typed U2M case has no exported constructor, so the tests do not reach it. Windows: only
go vetand a test build ran, so the console behavior is from reading the code.This pull request and its description were written by Isaac.