Require a timeout on every agent tool

internal/mcp/server.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"time"
 )
 
 const protocolVersion = "2025-06-18"
@@ -39,7 +40,72 @@ func TextResult(text string, isError bool) Result {
 	return Result{Content: []Content{{Type: "text", Text: text}}, IsError: isError}
 }
 
+const (
+	MaxTimeout = 10 * time.Minute
+	Grace      = 5 * time.Second
+)
+
+type timeoutKey struct{}
+
+func Timeout(ctx context.Context) time.Duration {
+	return ctx.Value(timeoutKey{}).(time.Duration)
+}
+
+func (t Tool) Schema() json.RawMessage {
+	s := map[string]any{}
+	err := json.Unmarshal(t.InputSchema, &s)
+	if err != nil {
+		panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
+	}
+	props, _ := s["properties"].(map[string]any)
+	if props == nil {
+		props = map[string]any{}
+		s["properties"] = props
+	}
+	if props["timeout_ms"] == nil {
+		props["timeout_ms"] = map[string]any{"type": "integer", "description": fmt.Sprintf("Give up on the call after this long (at most %d).", MaxTimeout.Milliseconds())}
+	}
+	required, _ := s["required"].([]any)
+	s["required"] = append(required, "timeout_ms")
+	out, err := json.Marshal(s)
+	if err != nil {
+		panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
+	}
+	return out
+}
+
 func (t Tool) Call(ctx context.Context, args json.RawMessage) Result {
+	p := struct {
+		TimeoutMs *int `json:"timeout_ms"`
+	}{}
+	err := json.Unmarshal(args, &p)
+	if err != nil {
+		return TextResult(fmt.Sprintf("malformed %s call: %s", t.Name, err), true)
+	}
+	if p.TimeoutMs == nil || *p.TimeoutMs <= 0 || *p.TimeoutMs > int(MaxTimeout.Milliseconds()) {
+		return TextResult(fmt.Sprintf("timeout_ms is required, between 1 and %d", MaxTimeout.Milliseconds()), true)
+	}
+	timeout := time.Duration(*p.TimeoutMs) * time.Millisecond
+
+	cctx, cancel := context.WithTimeout(context.WithValue(ctx, timeoutKey{}, timeout), timeout)
+	defer cancel()
+	done := make(chan Result, 1)
+	go func() {
+		done <- t.run(cctx, args)
+	}()
+	late := time.NewTimer(timeout + Grace)
+	defer late.Stop()
+	select {
+	case res := <-done:
+		return res
+	case <-ctx.Done():
+		return TextResult(ctx.Err().Error(), true)
+	case <-late.C:
+		return TextResult(fmt.Sprintf("%s gave no result within %s", t.Name, timeout), true)
+	}
+}
+
+func (t Tool) run(ctx context.Context, args json.RawMessage) Result {
 	if t.Rich != nil {
 		return t.Rich(ctx, args)
 	}
@@ -127,7 +193,7 @@ func (s *Server) dispatch(ctx context.Context, req request) (any, *rpcError) {
 	case "tools/list":
 		tools := []toolDesc{}
 		for _, t := range s.Tools {
-			tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
+			tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.Schema()})
 		}
 		return map[string]any{"tools": tools}, nil