Support branches in the existing MCP tools

CLAUDE.md

@@ -143,7 +143,7 @@ A command that identifies a service takes a **ref** — its ID, a read replica s
 - Help text shows `[name-or-id]` and says the service "can be given by ID or name". Shell completion offers names, falling back to IDs when what is typed matches an ID instead.
 - A name is accepted only as an argument. A configured default must be an ID, and one that resolves by name is refused. The `--service-id` flag completion only offers IDs.
 - Destructive commands accept a name, but their confirmation prompt takes only the ID.
-- MCP tools stay IDs-only — an intentional divergence, documented at `setServiceIDSchemaProperties`. A tool that goes through the resolve endpoint (the branch tools) refuses a match by name.
+- MCP tools stay IDs-only — an intentional divergence, documented at `setServiceIDSchemaProperties`. A tool that goes through the resolve endpoint (the branch tools, `db_query`, `db_schema`, `db_extensions_list`) refuses a match by name.
 
 ## Output
 

internal/cmd/integration_test.go

@@ -493,7 +493,7 @@ func TestServiceLifecycleIntegration(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to build API client: %v", err)
 		}
-		target, err := common.ResolveConnectionTargetByID(t.Context(), client, projectID, serviceID)
+		target, err := common.ResolveConnectionTargetByID(t.Context(), client, projectID, serviceID, false)
 		if err != nil {
 			t.Fatalf("failed to resolve connection target: %v", err)
 		}

internal/common/connection_target.go

@@ -188,6 +188,19 @@ func ResolveConnectionTarget(ctx context.Context, client api.ClientWithResponses
 	}
 }
 
+// ResolveConnectionTargetByID is ResolveConnectionTarget for a caller that
+// takes only IDs: a ref that matched by name is reported as not found.
+func ResolveConnectionTargetByID(ctx context.Context, client api.ClientWithResponsesInterface, projectID, id string, allowBranches bool) (*ConnectionTarget, error) {
+	target, err := ResolveConnectionTarget(ctx, client, projectID, id, allowBranches)
+	if err != nil {
+		return nil, err
+	}
+	if target.ID != id {
+		return nil, ExitWithCode(ExitServiceNotFound, fmt.Errorf("%s '%s' not found", connectionTargetKind(allowBranches), id))
+	}
+	return target, nil
+}
+
 // connectionTargetKind names what ResolveConnectionTarget looks for, in its
 // errors. Branches are left out unless they are allowed.
 func connectionTargetKind(allowBranches bool) string {
@@ -197,43 +210,6 @@ func connectionTargetKind(allowBranches bool) string {
 	return "service or read replica set"
 }
 
-// ResolveConnectionTargetByID fetches the service named by id — which may be a
-// primary service ID or a read replica set ID, both of which GetService
-// resolves — and works out whose credentials it uses. Callers that already
-// hold the service use NewConnectionTarget instead.
-func ResolveConnectionTargetByID(ctx context.Context, client api.ClientWithResponsesInterface, projectID, id string) (*ConnectionTarget, error) {
-	service, err := GetService(ctx, client, projectID, id)
-	if err != nil {
-		return nil, err
-	}
-	return NewConnectionTarget(ctx, client, projectID, *service)
-}
-
-// NewConnectionTarget returns the ConnectionTarget of an already-fetched
-// service. A standby read replica connects to its own endpoint but shares the
-// parent primary's credentials, so the parent is fetched here.
-func NewConnectionTarget(ctx context.Context, client api.ClientWithResponsesInterface, projectID string, service api.Service) (*ConnectionTarget, error) {
-	target := NewServiceConnectionTarget(service)
-	if !IsReadReplica(service) {
-		return target, nil
-	}
-	target.Kind = TargetReadReplica
-
-	// A replica with no parent recorded has nowhere else to look, so it stands
-	// in as its own credential owner.
-	parentID := util.DerefStr(service.ForkedFrom.ServiceID)
-	if parentID == "" {
-		return target, nil
-	}
-
-	parent, err := GetService(ctx, client, projectID, parentID)
-	if err != nil {
-		return nil, fmt.Errorf("failed to fetch parent service %q for read replica: %w", parentID, err)
-	}
-	target.Credentials = ServiceCredentials(*parent)
-	return target, nil
-}
-
 // environmentTag reads an optional environment tag. Anything but PROD reads
 // as DEV, since the tag is optional and treating unknowns as protected would
 // have read_only=prod block untagged services.

internal/common/connection_target_test.go

@@ -15,20 +15,48 @@ import (
 	"github.com/timescale/tiger-cli/internal/util"
 )
 
-// serviceTestClient builds an API client backed by an httptest server that
-// serves the getService endpoint from services keyed by service ID (404 when
-// absent).
-func serviceTestClient(t *testing.T, services map[string]api.Service) *api.ClientWithResponses {
+// resolveTestClient builds an API client backed by an httptest server that
+// answers the resolve endpoint from resolved, keyed by ref, and the getService
+// endpoint from services, keyed by service ID (404 when absent). Each resolved
+// value is a service, branch, or read replica set, encoded with the given
+// discriminator.
+func resolveTestClient(t *testing.T, resolved map[string]resolvedFixture, services map[string]api.Service) *api.ClientWithResponses {
 	t.Helper()
 
+	notFound := func(w http.ResponseWriter) {
+		w.WriteHeader(http.StatusNotFound)
+		_ = json.NewEncoder(w).Encode(map[string]string{"message": "not found"})
+	}
+
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 		parts := strings.Split(strings.Trim(r.URL.Path, "/"), "/")
-		id := parts[len(parts)-1]
-		svc, ok := services[id]
+		last := parts[len(parts)-1]
+
+		if last == "resolve" {
+			fixture, ok := resolved[r.URL.Query().Get("ref")]
+			if !ok {
+				notFound(w)
+				return
+			}
+			raw, err := json.Marshal(fixture.value)
+			if err != nil {
+				t.Errorf("failed to encode fixture: %v", err)
+				return
+			}
+			var body map[string]any
+			if err := json.Unmarshal(raw, &body); err != nil {
+				t.Errorf("failed to decode fixture: %v", err)
+				return
+			}
+			body["type"] = fixture.typ
+			_ = json.NewEncoder(w).Encode(body)
+			return
+		}
+
+		svc, ok := services[last]
 		if !ok {
-			w.WriteHeader(http.StatusNotFound)
-			_ = json.NewEncoder(w).Encode(map[string]string{"message": "service not found"})
+			notFound(w)
 			return
 		}
 		_ = json.NewEncoder(w).Encode(svc)
@@ -42,6 +70,12 @@ func serviceTestClient(t *testing.T, services map[string]api.Service) *api.Clien
 	return client
 }
 
+// resolvedFixture is one answer of the resolve endpoint in resolveTestClient.
+type resolvedFixture struct {
+	typ   api.ResolveType
+	value any
+}
+
 func primaryService() api.Service {
 	host := "svcprimary.example.com"
 	port := 5432
@@ -66,51 +100,87 @@ func replicaSet() api.ReadReplicaSet {
 	}
 }
 
-// replicaService returns a standby read replica of primaryService, shaped
-// the way GetService returns it when given a read replica set ID.
-func replicaService() api.Service {
-	host := "replica.example.com"
-	port := 5432
-	return api.Service{
-		ServiceID: "rep1234567",
-		ProjectID: "proj1",
-		Name:      "reporting-replica",
-		Endpoint:  &api.Endpoint{Host: &host, Port: &port},
-		ForkedFrom: &api.ForkSpec{
-			IsStandby: new(true),
-			ProjectID: new("proj1"),
-			ServiceID: new("svcprimary"),
-		},
-	}
-}
-
+// TestResolveConnectionTargetByID is helper-level because the MCP db tools
+// that call it can't reach the branch and replica cases without a live
+// connection past it.
 func TestResolveConnectionTargetByID(t *testing.T) {
 	primary := primaryService()
-	replica := replicaService()
-	client := serviceTestClient(t, map[string]api.Service{"svcprimary": primary, "rep1234567": replica})
+	replicaHost := "replica.example.com"
+	replica := api.ReadReplicaSet{
+		ID:        "rep1234567",
+		ServiceID: "svcprimary",
+		Name:      "reporting-replica",
+		Status:    api.ReadReplicaSetStatusActive,
+		Endpoint:  &api.Endpoint{Host: &replicaHost, Port: new(5432)},
+	}
+	branch := api.Branch{
+		ID:        "br12345678",
+		ProjectID: "proj1",
+		ServiceID: "svcprimary",
+		Name:      "my-branch",
+		Status:    api.DeployStatusREADY,
+		Endpoint:  &api.Endpoint{Host: new("br12345678.example.com"), Port: new(5432)},
+	}
+	client := resolveTestClient(t, map[string]resolvedFixture{
+		"svcprimary": {api.ResolveTypeSERVICE, primary},
+		"my-db":      {api.ResolveTypeSERVICE, primary},
+		"rep1234567": {api.ResolveTypeREADREPLICASET, replica},
+		"br12345678": {api.ResolveTypeBRANCH, branch},
+	}, map[string]api.Service{"svcprimary": primary})
+	ctx := context.Background()
 
 	// Primary ID → primary target.
-	target, err := ResolveConnectionTargetByID(context.Background(), client, "proj1", "svcprimary")
+	target, err := ResolveConnectionTargetByID(ctx, client, "proj1", "svcprimary", false)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 	if target.Kind != TargetService || target.ID != "svcprimary" || target.Credentials.ID != "svcprimary" {
 		t.Errorf("expected primary target for a primary ID, got %+v", target)
 	}
 
-	// Replica ID → replica target with parent credentials.
-	target, err = ResolveConnectionTargetByID(context.Background(), client, "proj1", "rep1234567")
+	// Replica set ID → replica target with the parent's credentials.
+	target, err = ResolveConnectionTargetByID(ctx, client, "proj1", "rep1234567", false)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 	if target.Kind != TargetReadReplica || target.ID != "rep1234567" || target.Credentials.ID != "svcprimary" {
 		t.Errorf("expected replica target with parent credentials, got %+v", target)
 	}
 
-	// Unknown ID → error.
-	if _, err := ResolveConnectionTargetByID(context.Background(), client, "proj1", "missing9999"); err == nil {
-		t.Error("expected an error for an unknown ID, got nil")
+	// Branch ID → branch target with its own credentials, when allowed.
+	target, err = ResolveConnectionTargetByID(ctx, client, "proj1", "br12345678", true)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if target.Kind != TargetBranch || target.ID != "br12345678" || target.Credentials.ID != "br12345678" {
+		t.Errorf("expected branch target with its own credentials, got %+v", target)
+	}
+
+	// Branch ID without branches allowed → not found, without naming branches.
+	_, err = ResolveConnectionTargetByID(ctx, client, "proj1", "br12345678", false)
+	if diff := cmp.Diff("service or read replica set 'br12345678' not found", errString(err)); diff != "" {
+		t.Errorf("gated branch error mismatch (-want +got):\n%s", diff)
+	}
+
+	// A name resolves, but isn't an ID.
+	_, err = ResolveConnectionTargetByID(ctx, client, "proj1", "my-db", true)
+	if diff := cmp.Diff("service, branch, or read replica set 'my-db' not found", errString(err)); diff != "" {
+		t.Errorf("name match error mismatch (-want +got):\n%s", diff)
+	}
+
+	// Unknown ID → not found.
+	_, err = ResolveConnectionTargetByID(ctx, client, "proj1", "missing999", true)
+	if diff := cmp.Diff("service, branch, or read replica set 'missing999' not found", errString(err)); diff != "" {
+		t.Errorf("unknown ID error mismatch (-want +got):\n%s", diff)
+	}
+}
+
+// errString returns err's message, or "" for nil.
+func errString(err error) string {
+	if err == nil {
+		return ""
 	}
+	return err.Error()
 }
 
 func TestNewReplicaConnectionTarget(t *testing.T) {
@@ -324,29 +394,13 @@ func TestReplicaPoolerWarning(t *testing.T) {
 }
 
 func TestResolveConnectionTargetByID_ReplicaParentFetchFails(t *testing.T) {
-	// The replica itself resolves, but its parent is absent from the server.
-	client := serviceTestClient(t, map[string]api.Service{"rep1234567": replicaService()})
+	// The replica set itself resolves, but its parent is absent from the server.
+	replica := api.ReadReplicaSet{ID: "rep1234567", ServiceID: "svcprimary", Name: "reporting-replica"}
+	client := resolveTestClient(t, map[string]resolvedFixture{
+		"rep1234567": {api.ResolveTypeREADREPLICASET, replica},
+	}, nil)
 
-	if _, err := ResolveConnectionTargetByID(context.Background(), client, "proj1", "rep1234567"); err == nil {
+	if _, err := ResolveConnectionTargetByID(context.Background(), client, "proj1", "rep1234567", false); err == nil {
 		t.Fatal("expected an error when the parent service can't be fetched, got nil")
 	}
 }
-
-func TestResolveConnectionTargetByID_ReplicaWithoutParentID(t *testing.T) {
-	// A standby with no parent recorded has nowhere else to look for
-	// credentials, so it stands in as its own credential owner.
-	replica := replicaService()
-	replica.ForkedFrom.ServiceID = nil
-	client := serviceTestClient(t, map[string]api.Service{"rep1234567": replica})
-
-	target, err := ResolveConnectionTargetByID(context.Background(), client, "proj1", "rep1234567")
-	if err != nil {
-		t.Fatalf("unexpected error: %v", err)
-	}
-	if target.Kind != TargetReadReplica {
-		t.Error("expected a replica target")
-	}
-	if target.Credentials.ID != "rep1234567" {
-		t.Errorf("expected the replica as its own credential owner, got %q", target.Credentials.ID)
-	}
-}

internal/mcp/db_extensions_list.go

@@ -17,10 +17,10 @@ type DBExtensionsListInput struct {
 	Role      string `json:"role,omitempty"`
 }
 
-func (DBExtensionsListInput) Schema() *jsonschema.Schema {
+func (DBExtensionsListInput) Schema(experimental bool) *jsonschema.Schema {
 	schema := util.Must(jsonschema.For[DBExtensionsListInput](nil))
 
-	schema.Properties["service_id"].Description = "Unique identifier of the service (10-character alphanumeric string). Use service_list to find service IDs. A read replica set ID is also accepted here — passing one lists the extensions on that read replica instead of the primary service."
+	schema.Properties["service_id"].Description = "Unique identifier of the service (10-character alphanumeric string). Use service_list to find service IDs. A read replica set ID is also accepted here — passing one lists the extensions on that read replica instead of the primary service." + branchIDHelp(experimental, "lists the extensions on that branch")
 	schema.Properties["service_id"].Examples = []any{"e6ue9697jf", "u8me885b93"}
 	schema.Properties["service_id"].Pattern = "^[a-z0-9]{10}$"
 
@@ -60,14 +60,14 @@ func (DBExtensionsListOutput) Schema() *jsonschema.Schema {
 	return schema
 }
 
-func newDBExtensionsListTool() *mcp.Tool {
+func newDBExtensionsListTool(experimental bool) *mcp.Tool {
 	return &mcp.Tool{
 		Name:  toolDBExtensionsList,
 		Title: "List Database Extensions",
 		Description: `List the PostgreSQL extensions installed in a service database.
 
 Connects to a PostgreSQL/TimescaleDB service in Tiger Cloud and returns each installed extension's name, version, schema, and description. The connection is opened in immutable read-only mode.`,
-		InputSchema:  DBExtensionsListInput{}.Schema(),
+		InputSchema:  DBExtensionsListInput{}.Schema(experimental),
 		OutputSchema: DBExtensionsListOutput{}.Schema(),
 		Annotations: &mcp.ToolAnnotations{
 			ReadOnlyHint:  true,
@@ -84,14 +84,16 @@ func (s *Server) handleDBExtensionsList(ctx context.Context, req *mcp.CallToolRe
 		return nil, DBExtensionsListOutput{}, err
 	}
 
-	target, err := common.ResolveConnectionTargetByID(ctx, client, projectID, input.ServiceID)
+	// service_id may name a service, one of its read replicas, or (behind the
+	// experimental gate) a branch.
+	target, err := common.ResolveConnectionTargetByID(ctx, client, projectID, input.ServiceID, s.app.Experimental)
 	if err != nil {
 		return nil, DBExtensionsListOutput{}, err
 	}
 
 	extensions, err := common.FetchExtensions(ctx, cfg, target, input.Role)
 	if err != nil {
-		return nil, DBExtensionsListOutput{}, handleDatabaseError(err)
+		return nil, DBExtensionsListOutput{}, handleDatabaseError(err, target)
 	}
 
 	return nil, DBExtensionsListOutput{Extensions: extensions}, nil