diff --git a/pkg/hookdeck/auth.go b/pkg/hookdeck/auth.go index 9ff6f02d..a0f4e246 100644 --- a/pkg/hookdeck/auth.go +++ b/pkg/hookdeck/auth.go @@ -175,7 +175,13 @@ func (c *Client) withoutProjectScope() *Client { // ValidateAPIKey validates an API key and returns user/project information func (c *Client) ValidateAPIKey() (*ValidateAPIKeyResponse, error) { - res, err := c.clientForCLIAuthValidate().Get(context.Background(), APIPathPrefix+"/cli-auth/validate", "", nil) + return c.ValidateAPIKeyContext(context.Background()) +} + +// ValidateAPIKeyContext is ValidateAPIKey bounded by ctx, for callers that must +// not wait indefinitely on the API. +func (c *Client) ValidateAPIKeyContext(ctx context.Context) (*ValidateAPIKeyResponse, error) { + res, err := c.clientForCLIAuthValidate().Get(ctx, APIPathPrefix+"/cli-auth/validate", "", nil) if err != nil { return nil, err } diff --git a/pkg/hookdeck/projects.go b/pkg/hookdeck/projects.go index cab8aedf..6c2013b7 100644 --- a/pkg/hookdeck/projects.go +++ b/pkg/hookdeck/projects.go @@ -12,7 +12,13 @@ type Project struct { } func (c *Client) ListProjects() ([]Project, error) { - res, err := c.clientForCLIAuthValidate().Get(context.Background(), APIPathPrefix+"/projects", "", nil) + return c.ListProjectsContext(context.Background()) +} + +// ListProjectsContext is ListProjects bounded by ctx, for callers that must not +// wait indefinitely on the API. +func (c *Client) ListProjectsContext(ctx context.Context) ([]Project, error) { + res, err := c.clientForCLIAuthValidate().Get(ctx, APIPathPrefix+"/projects", "", nil) if err != nil { return []Project{}, err } diff --git a/pkg/mcpcore/project_display.go b/pkg/mcpcore/project_display.go index 6de8767a..33d62c3f 100644 --- a/pkg/mcpcore/project_display.go +++ b/pkg/mcpcore/project_display.go @@ -1,15 +1,31 @@ package mcpcore import ( + "context" + "time" + "github.com/hookdeck/hookdeck-cli/pkg/hookdeck" "github.com/hookdeck/hookdeck-cli/pkg/project" ) +// projectDisplayLookupTimeout bounds the whole display-name lookup. The lookup +// runs before every tool call until it succeeds, including calls that never +// reach the API (help, a write refused in read-only mode, a rejected argument), +// and the API client has no deadline of its own. Without this bound one slow +// API response held the tool call open for as long as the API took: the client +// saw initialize answered and then nothing. The name is cosmetic, so giving up +// costs only an id where a name would have been, and the next call tries again. +// A variable so tests can shorten it. +var projectDisplayLookupTimeout = 5 * time.Second + // FillProjectDisplayNameIfNeeded sets target.ProjectOrg and target.ProjectName // when target has an API key and project id but no cached org/name (typical // after loading the profile from disk). Fails silently on API errors. Stdio MCP // invokes tools sequentially, so this is safe without locking. // +// The lookup gives up after projectDisplayLookupTimeout, or sooner when ctx is +// done, so it can never hold up the tool call it runs ahead of. +// // lookup is the client account-level requests are made from, which is not always // target: a product API served from its own host does not answer them, so the // lookup has to go to the account API. @@ -19,16 +35,21 @@ import ( // projects only works for a user-associated key, so a project-scoped key from // `hookdeck ci` — a common way to configure an MCP server — would otherwise // leave the name empty and callers with nothing but an opaque id to show. -func FillProjectDisplayNameIfNeeded(lookup, target *hookdeck.Client) { +func FillProjectDisplayNameIfNeeded(ctx context.Context, lookup, target *hookdeck.Client) { if lookup == nil || target == nil || target.APIKey == "" || target.ProjectID == "" { return } if target.ProjectName != "" || target.ProjectOrg != "" { return } + if ctx == nil { + ctx = context.Background() + } + ctx, cancel := context.WithTimeout(ctx, projectDisplayLookupTimeout) + defer cancel() // The key's own project, available whatever kind of key it is. - if response, err := lookup.ValidateAPIKey(); err == nil && response.ProjectID == target.ProjectID { + if response, err := lookup.ValidateAPIKeyContext(ctx); err == nil && response.ProjectID == target.ProjectID { target.ProjectName = response.ProjectName target.ProjectOrg = response.OrganizationName if target.ProjectName != "" || target.ProjectOrg != "" { @@ -36,9 +57,14 @@ func FillProjectDisplayNameIfNeeded(lookup, target *hookdeck.Client) { } } + // Out of time: the listing would fail too. + if ctx.Err() != nil { + return + } + // The active project differs from the key's, so it has to be looked up. // Only a user-associated key can do this. - projects, err := lookup.ListProjects() + projects, err := lookup.ListProjectsContext(ctx) if err != nil { return } diff --git a/pkg/mcpcore/project_display_test.go b/pkg/mcpcore/project_display_test.go index 18cd1253..82ed5591 100644 --- a/pkg/mcpcore/project_display_test.go +++ b/pkg/mcpcore/project_display_test.go @@ -1,18 +1,111 @@ package mcpcore import ( + "context" "encoding/json" "net/http" "net/http/httptest" "net/url" "strings" + "sync/atomic" "testing" + "time" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/hookdeck/hookdeck-cli/pkg/hookdeck" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +// stalledAPI answers nothing until the request is abandoned, which is what a +// slow API looks like to a client with no deadline of its own. +func stalledAPI(t *testing.T) (*hookdeck.Client, *atomic.Int64) { + t.Helper() + var requests atomic.Int64 + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + select { + case <-r.Context().Done(): + case <-release: + } + })) + t.Cleanup(func() { + close(release) + server.Close() + }) + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + return &hookdeck.Client{BaseURL: baseURL, APIKey: "k", ProjectID: "tm_1"}, &requests +} + +func shortenProjectDisplayLookup(t *testing.T, d time.Duration) { + t.Helper() + previous := projectDisplayLookupTimeout + projectDisplayLookupTimeout = d + t.Cleanup(func() { projectDisplayLookupTimeout = previous }) +} + +// TestFillProjectDisplayName_GivesUpOnAStalledAPI: the lookup is cosmetic, so a +// slow API must cost at most the timeout, not however long the API takes. +func TestFillProjectDisplayName_GivesUpOnAStalledAPI(t *testing.T) { + shortenProjectDisplayLookup(t, 200*time.Millisecond) + client, requests := stalledAPI(t) + + start := time.Now() + FillProjectDisplayNameIfNeeded(context.Background(), client, client) + + assert.Less(t, time.Since(start), 3*time.Second, "the lookup must give up at its timeout") + assert.GreaterOrEqual(t, requests.Load(), int64(1), "the lookup should have been attempted") + assert.Empty(t, client.ProjectName, "a lookup that gave up leaves the name for the next call to fill") + assert.Empty(t, client.ProjectOrg) +} + +// TestToolCallIsNotHeldByAStalledDisplayNameLookup is the case the acceptance +// suite kept hitting: a tool that never needs the API (here, one standing in for +// a write refused in read-only mode) answered nothing, because the lookup ahead +// of it waited on the API with no deadline. +func TestToolCallIsNotHeldByAStalledDisplayNameLookup(t *testing.T) { + shortenProjectDisplayLookup(t, 200*time.Millisecond) + client, _ := stalledAPI(t) + s := &Server{client: client, accountClient: client} + + refusal := mcpsdk.ToolHandler(func(ctx context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { + return &mcpsdk.CallToolResult{IsError: true}, nil + }) + + done := make(chan *mcpsdk.CallToolResult, 1) + go func() { + result, _ := s.wrapWithTelemetry("gateway_sources_read", refusal)(context.Background(), newCallToolRequest(`{"action":"delete"}`)) + done <- result + }() + + select { + case result := <-done: + require.NotNil(t, result) + assert.True(t, result.IsError) + case <-time.After(5 * time.Second): + t.Fatal("the tool call is still waiting on the display-name lookup") + } +} + +// TestFillProjectDisplayName_StopsWhenTheCallIsCancelled: a client that cancels +// the tool call must not leave the lookup running to its own timeout. +func TestFillProjectDisplayName_StopsWhenTheCallIsCancelled(t *testing.T) { + shortenProjectDisplayLookup(t, time.Minute) + client, _ := stalledAPI(t) + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + start := time.Now() + FillProjectDisplayNameIfNeeded(ctx, client, client) + + assert.Less(t, time.Since(start), 3*time.Second, "cancelling the call must end the lookup") + assert.Empty(t, client.ProjectName) +} + func TestFillProjectDisplayNameIfNeeded_SetsNameFromAPI(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/2026-09-01/projects" { @@ -32,14 +125,14 @@ func TestFillProjectDisplayNameIfNeeded_SetsNameFromAPI(t *testing.T) { APIKey: "k", ProjectID: "proj_x", } - FillProjectDisplayNameIfNeeded(client, client) + FillProjectDisplayNameIfNeeded(context.Background(), client, client) require.Equal(t, "Acme", client.ProjectOrg) require.Equal(t, "production", client.ProjectName) } func TestFillProjectDisplayNameIfNeeded_NoOpWhenNameSet(t *testing.T) { client := &hookdeck.Client{ProjectID: "p", ProjectName: "already"} - FillProjectDisplayNameIfNeeded(client, client) + FillProjectDisplayNameIfNeeded(context.Background(), client, client) require.Equal(t, "already", client.ProjectName) } @@ -72,7 +165,7 @@ func TestFillProjectDisplayNameIfNeeded_LooksUpThroughTheAccountClient(t *testin account := &hookdeck.Client{BaseURL: accountURL, APIKey: "k", ProjectID: "proj_x"} product := &hookdeck.Client{BaseURL: productURL, APIKey: "k", ProjectID: "proj_x"} - FillProjectDisplayNameIfNeeded(account, product) + FillProjectDisplayNameIfNeeded(context.Background(), account, product) require.Equal(t, "Acme", product.ProjectOrg) require.Equal(t, "production", product.ProjectName) } @@ -100,7 +193,7 @@ func TestFillProjectDisplayName_ProjectScopedKey(t *testing.T) { require.NoError(t, err) client := &hookdeck.Client{BaseURL: baseURL, APIKey: "ci-key", ProjectID: "tm_1"} - FillProjectDisplayNameIfNeeded(client, client) + FillProjectDisplayNameIfNeeded(context.Background(), client, client) assert.Equal(t, "cli outpost testing", client.ProjectName) assert.Equal(t, "Automated Testing", client.ProjectOrg) @@ -126,7 +219,7 @@ func TestFillProjectDisplayName_FallsBackToListing(t *testing.T) { require.NoError(t, err) client := &hookdeck.Client{BaseURL: baseURL, APIKey: "user-key", ProjectID: "tm_1"} - FillProjectDisplayNameIfNeeded(client, client) + FillProjectDisplayNameIfNeeded(context.Background(), client, client) assert.Equal(t, "the active one", client.ProjectName, "the key's own project must not be used when it is not the active one") } diff --git a/pkg/mcpcore/server.go b/pkg/mcpcore/server.go index 98bd68a0..bde3be9c 100644 --- a/pkg/mcpcore/server.go +++ b/pkg/mcpcore/server.go @@ -270,7 +270,7 @@ func (s *Server) wrapWithTelemetry(toolName string, handler mcpsdk.ToolHandler) } }() - FillProjectDisplayNameIfNeeded(s.accountClient, s.client) + FillProjectDisplayNameIfNeeded(ctx, s.accountClient, s.client) return handler(ctx, req) } diff --git a/test/acceptance/helpers.go b/test/acceptance/helpers.go index b6ab8047..d5f6472b 100644 --- a/test/acceptance/helpers.go +++ b/test/acceptance/helpers.go @@ -717,10 +717,27 @@ func RunOutpostMCPSubprocess(t *testing.T, projectRoot, configPath string, args return RunMCPSubprocess(t, projectRoot, configPath, append([]string{"outpost", "mcp"}, args...), extraEnv, stdin, runDuration) } -// RunMCPSubprocess builds the CLI binary, runs the given MCP command with optional stdin, -// lets it run for runDuration, then kills the process. Returns stdout, stderr, and the -// error from Wait (often non-nil because the process was killed). configPath, when non-empty, -// is passed as HOOKDECK_CONFIG_FILE. extraEnv entries override the process environment. +// mcpShutdownGrace is how long a server gets to exit after its stdin is closed +// before it is killed. +const mcpShutdownGrace = 5 * time.Second + +// RunMCPSubprocess builds the CLI binary and drives the given MCP command over +// stdio the way an MCP client does. stdin holds newline-delimited JSON-RPC +// messages. An initialize request is sent on its own, and the messages after it +// only once the server has answered it. Stdout is read as it arrives, and as soon +// as every request in stdin has a response, stdin is closed (how a client ends a +// stdio session) and the process is left to exit. +// +// runDuration is an upper bound, not a fixed wait: if the responses have not all +// arrived by then, the process is killed and the caller sees whatever it did +// write. With no requests in stdin (none at all, or only notifications) there is +// nothing to wait for, so the process runs until it exits or runDuration passes. +// An empty stdin reads as EOF at once, for commands expected to fail before they +// serve. +// +// Returns stdout, stderr, and the error from Wait (non-nil when the process was +// killed or exited non-zero). configPath, when non-empty, is passed as +// HOOKDECK_CONFIG_FILE. extraEnv entries override the process environment. func RunMCPSubprocess(t *testing.T, projectRoot, configPath string, args []string, extraEnv map[string]string, stdin string, runDuration time.Duration) (stdout, stderr string, err error) { t.Helper() tmpBinary := filepath.Join(projectRoot, "hookdeck-mcp-test-"+generateTimestamp()) @@ -743,44 +760,213 @@ func RunMCPSubprocess(t *testing.T, projectRoot, configPath string, args []strin } cmd.Env = env - var stdoutBuf, stderrBuf bytes.Buffer - cmd.Stdout = &stdoutBuf + out := newMCPStdout() + var stderrBuf bytes.Buffer + cmd.Stdout = out cmd.Stderr = &stderrBuf - closeStdin := func() {} + messages := splitMCPMessages(stdin) + var stdinPipe io.WriteCloser if stdin != "" { - pr, pw, pipeErr := os.Pipe() + pipe, pipeErr := cmd.StdinPipe() if pipeErr != nil { return "", "", pipeErr } - cmd.Stdin = pr - if _, werr := io.WriteString(pw, stdin); werr != nil { - _ = pr.Close() - _ = pw.Close() - return "", "", werr - } - if !strings.HasSuffix(stdin, "\n") { - _, _ = pw.Write([]byte("\n")) - } - closeStdin = func() { - _ = pw.Close() - _ = pr.Close() - } + stdinPipe = pipe } if err := cmd.Start(); err != nil { - closeStdin() return "", "", err } + exited := make(chan struct{}) + var waitErr error + go func() { + waitErr = cmd.Wait() + close(exited) + }() + // Killing here, not only signalling, is what bounds a write to a server + // that has stopped reading. + expired := make(chan struct{}) timer := time.AfterFunc(runDuration, func() { - if cmd.Process != nil { - _ = cmd.Process.Kill() - } + close(expired) + _ = cmd.Process.Kill() }) defer timer.Stop() - waitErr := cmd.Wait() - closeStdin() - return stdoutBuf.String(), stderrBuf.String(), waitErr + + // awaitResponses reports whether every id has a response, waiting until + // they do, the process exits, or runDuration passes. + awaitResponses := func(ids []string) bool { + for { + if out.hasResponses(ids) { + return true + } + select { + case <-out.changed: + case <-exited: + return out.hasResponses(ids) + case <-expired: + return false + } + } + } + + requests := mcpRequestIDs(messages) + answered := false + if stdinPipe != nil { + head, tail, initID := splitAtMCPInitialize(messages) + writeMCPMessages(stdinPipe, head) + if initID == "" || awaitResponses([]string{initID}) { + writeMCPMessages(stdinPipe, tail) + answered = len(requests) > 0 && awaitResponses(requests) + } + } + + if answered { + _ = stdinPipe.Close() + select { + case <-exited: + case <-expired: + case <-time.After(mcpShutdownGrace): + } + } else { + select { + case <-exited: + case <-expired: + } + } + _ = cmd.Process.Kill() + <-exited + return out.String(), stderrBuf.String(), waitErr +} + +// mcpStdout collects a server's stdout and tracks which JSON-RPC responses have +// arrived, so the harness can stop waiting as soon as the one it needs is there. +type mcpStdout struct { + mu sync.Mutex + buf bytes.Buffer + partial []byte + responses map[string]bool + // changed is signalled after every write. + changed chan struct{} +} + +func newMCPStdout() *mcpStdout { + return &mcpStdout{responses: map[string]bool{}, changed: make(chan struct{}, 1)} +} + +func (o *mcpStdout) Write(p []byte) (int, error) { + o.mu.Lock() + o.buf.Write(p) + o.partial = append(o.partial, p...) + for { + i := bytes.IndexByte(o.partial, '\n') + if i < 0 { + break + } + if id, ok := mcpResponseID(o.partial[:i]); ok { + o.responses[id] = true + } + o.partial = o.partial[i+1:] + } + o.mu.Unlock() + select { + case o.changed <- struct{}{}: + default: + } + return len(p), nil +} + +func (o *mcpStdout) hasResponses(ids []string) bool { + o.mu.Lock() + defer o.mu.Unlock() + for _, id := range ids { + if !o.responses[id] { + return false + } + } + return true +} + +func (o *mcpStdout) String() string { + o.mu.Lock() + defer o.mu.Unlock() + return o.buf.String() +} + +// jsonRPCEnvelope is the part of a JSON-RPC message the harness routes on. +type jsonRPCEnvelope struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` +} + +// mcpMessageID returns a JSON-RPC message's id as a comparable key, and whether +// it has one. The raw JSON is the key, so 2 and "2" stay distinct. +func mcpMessageID(env jsonRPCEnvelope) (string, bool) { + id := strings.TrimSpace(string(env.ID)) + if id == "" || id == "null" { + return "", false + } + return id, true +} + +// mcpResponseID returns the id of a response line. Requests the server sends to +// the client carry an id too, so a method rules a line out. +func mcpResponseID(line []byte) (string, bool) { + var env jsonRPCEnvelope + if json.Unmarshal(line, &env) != nil || env.Method != "" { + return "", false + } + return mcpMessageID(env) +} + +// splitMCPMessages returns the non-blank lines of stdin, one message each. +func splitMCPMessages(stdin string) []string { + var messages []string + for _, line := range strings.Split(stdin, "\n") { + if strings.TrimSpace(line) != "" { + messages = append(messages, line) + } + } + return messages +} + +// mcpRequestIDs returns the ids of the requests in messages: everything that +// expects a response. Notifications and lines that are not JSON are skipped. +func mcpRequestIDs(messages []string) []string { + var ids []string + for _, message := range messages { + var env jsonRPCEnvelope + if json.Unmarshal([]byte(message), &env) != nil || env.Method == "" { + continue + } + if id, ok := mcpMessageID(env); ok { + ids = append(ids, id) + } + } + return ids +} + +// splitAtMCPInitialize splits messages after the initialize request, returning +// its id. Without one, everything is head and the id is empty. +func splitAtMCPInitialize(messages []string) (head, tail []string, initID string) { + for i, message := range messages { + var env jsonRPCEnvelope + if json.Unmarshal([]byte(message), &env) != nil || env.Method != "initialize" { + continue + } + if id, ok := mcpMessageID(env); ok { + return messages[:i+1], messages[i+1:], id + } + } + return messages, nil, "" +} + +// writeMCPMessages writes each message as one line. Errors are ignored: a server +// that has exited cannot read them, and its output says why. +func writeMCPMessages(w io.Writer, messages []string) { + for _, message := range messages { + _, _ = io.WriteString(w, message+"\n") + } } // MCPToolCallResult is the parsed JSON-RPC response for an MCP tools/call request.