From bed026ac91c6b7b18f4e66be8a8a6f968f0f1690 Mon Sep 17 00:00:00 2001 From: zbl94 Date: Mon, 13 Jul 2026 23:58:13 +0000 Subject: [PATCH] Add Gemini Enterprise Skill Registry integration for local harnesses Materialize agentskills.io skills from the Gemini Enterprise Skill Registry (Vertex AI v1beta1) into on-disk folders before the harness starts, so the built-in harnesses can use registry-hosted skills on the local `ax exec` / `ax serve` path. - internal/skills/geminienterprise: harness-agnostic package that reads config.SkillsConfig, drives the registry client (ListSkills / GetSkill / GetSkillRevision / skills:retrieve), safe-unzips payloads to //, and reports what it wrote. First-wins on duplicate ids with a warning; fail-safe (a registry error never blocks harness startup). - config: top-level `skills.registries[]` (harness-agnostic -- each actor runs a single harness that consumes the materialized folder). Per-registry selection (skills / query / all), required target_dir, and a validated "exactly one selection mode" rule. Also wires the interactions harness's system_instruction from ax.yaml. - cliutil: materializes skills once, up front, at controller construction; for the interactions harness (no SKILLS_DIR concept) it appends a discovery pointer to the system instruction. Scope: local path only; the substrate/pod path does not yet read ax.yaml. Verified end-to-end against a live registry (by-id and by-query). --- cmd/ax/internal/cliutil/cliutil.go | 27 +- internal/config/config.go | 104 +++++ internal/config/config_test.go | 177 ++++++++ .../harness/antigravityinteractions/skills.go | 61 +++ .../antigravityinteractions/skills_test.go | 68 +++ internal/skills/geminienterprise/client.go | 400 +++++++++++++++++ .../skills/geminienterprise/client_test.go | 414 ++++++++++++++++++ .../skills/geminienterprise/materialize.go | 173 ++++++++ .../geminienterprise/materialize_test.go | 100 +++++ internal/skills/geminienterprise/unzip.go | 202 +++++++++ 10 files changed, 1723 insertions(+), 3 deletions(-) create mode 100644 internal/harness/antigravityinteractions/skills.go create mode 100644 internal/harness/antigravityinteractions/skills_test.go create mode 100644 internal/skills/geminienterprise/client.go create mode 100644 internal/skills/geminienterprise/client_test.go create mode 100644 internal/skills/geminienterprise/materialize.go create mode 100644 internal/skills/geminienterprise/materialize_test.go create mode 100644 internal/skills/geminienterprise/unzip.go diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go index f8d0ebb1..04b1b121 100644 --- a/cmd/ax/internal/cliutil/cliutil.go +++ b/cmd/ax/internal/cliutil/cliutil.go @@ -26,6 +26,7 @@ import ( "github.com/google/ax/internal/harness/antigravity" "github.com/google/ax/internal/harness/antigravityinteractions" "github.com/google/ax/internal/harness/substrate" + "github.com/google/ax/internal/skills/geminienterprise" ) // Controller is the active controller type for this build. @@ -64,6 +65,22 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont var defaultHarnessID string var err error + // Materialize registry skills once, up front (skills config is top-level and + // harness-agnostic; each actor runs a single harness that consumes the + // materialized folder). Unconditional when configured. Fail-safe: a registry + // error degrades capability but never blocks harness creation. The + // interactions harness (no SKILLS_DIR concept) is told where the materialized + // skills are via a pointer appended to its system instruction. Only the local + // path materializes; substrate/pod does not yet read ax.yaml. + // + // TODO(joycel): wire the Antigravity SDK harness too. It discovers skills via + // SKILLS_DIR, so its SKILLS_DIR needs to be pointed at the materialized + // target_dir; currently only the interactions harness is fully wired. + var skillsPointer string + if !substrateMode { + skillsPointer = antigravityinteractions.SkillsSystemInstruction(geminienterprise.Materialize(ctx, cfg.Skills)) + } + // Built-in Antigravity harness. var antigravityHarness harness.Harness if !substrateMode { @@ -98,7 +115,8 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont // Built-in Antigravity Interactions harness. var antigravityInteractionsHarness harness.Harness if !substrateMode { - agent := cfg.Harnesses.AntigravityInteractions.Agent + aiCfg := cfg.Harnesses.AntigravityInteractions + agent := aiCfg.Agent if agent == "" { agent = antigravityinteractions.DefaultAgent } @@ -108,9 +126,12 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont if sErr != nil { return nil, fmt.Errorf("antigravity-interactions harness: %w", sErr) } + // skillsPointer was built once, up front, from the top-level skills + // config (see above). Append it to any configured system instruction. antigravityInteractionsHarness, err = antigravityinteractions.New(antigravityinteractions.AntigravityInteractionsConfig{ - Agent: agent, - StateDir: stateDir, + Agent: agent, + SystemInstruction: antigravityinteractions.JoinSystemInstruction(aiCfg.SystemInstruction, skillsPointer), + StateDir: stateDir, }) } else { antigravityInteractionsHarness, err = substrate.New(config.AntigravityInteractionsHarnessID, "", "", config.AntigravityInteractionsTemplate, 80) diff --git a/internal/config/config.go b/internal/config/config.go index 21ec968b..58ec5008 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -48,6 +48,11 @@ type Config struct { Server ServerConfig `yaml:"server"` EventLog EventLogConfig `yaml:"eventlog"` Harnesses HarnessesConfig `yaml:"harnesses,omitempty"` + // Skills sources skills from the Gemini Enterprise Skill Registry into on-disk folders + // before the harness starts. It is harness-agnostic: each actor runs exactly + // one harness, which consumes the materialized folder(s). Optional; disabled + // when no registry is enabled. + Skills SkillsConfig `yaml:"skills,omitempty"` Telemetry TelemetryConfig `yaml:"telemetry,omitempty"` } @@ -105,6 +110,101 @@ type AntigravityHarnessConfig struct { type AntigravityInteractionsHarnessConfig struct { Default bool `yaml:"default,omitempty"` // Default harness or not Agent string `yaml:"agent,omitempty"` // Interactions API agent (default: antigravityinteractions.DefaultAgent) + // SystemInstruction is a free-form system prompt sent on every turn. + SystemInstruction string `yaml:"system_instruction,omitempty"` +} + +// SkillsConfig configures optional skill sources (top-level, harness-agnostic). +// Today the only source type is the Gemini Enterprise Skill Registry; it may source from more +// than one registry (e.g. a shared org-wide registry plus a team-specific one), +// each with its own project/location, selection, and target directory. +type SkillsConfig struct { + Registries []SkillsRegistryConfig `yaml:"registries,omitempty"` +} + +// Validate checks the (top-level) skills config. +func (s SkillsConfig) Validate() error { + for i := range s.Registries { + if err := s.Registries[i].validate(i); err != nil { + return err + } + } + return nil +} + +// validate enforces that, when this registry source is enabled, exactly one +// selection mode is set (skills, query, or all) — a config-level "oneof". idx is +// the registry's index within the registries list, for error context. +func (rc SkillsRegistryConfig) validate(idx int) error { + if !rc.Enabled { + return nil + } + if rc.TargetDir == "" { + return fmt.Errorf("skills.registries[%d] requires target_dir (skills materialize to //)", idx) + } + modes := 0 + if len(rc.Skills) > 0 { + modes++ + } + if rc.Query != nil { + modes++ + } + if rc.All { + modes++ + } + switch { + case modes == 0: + return fmt.Errorf("skills.registries[%d] requires exactly one selection mode (set one of skills, query, or all)", idx) + case modes > 1: + return fmt.Errorf("skills.registries[%d] sets multiple selection modes; set exactly one of skills, query, or all", idx) + } + if rc.Query != nil && rc.Query.Text == "" { + return fmt.Errorf("skills.registries[%d].query requires a non-empty text", idx) + } + return nil +} + +// SkillsRegistryConfig sources agentskills.io skills from the Gemini Skill +// Registry. When Enabled, exactly one selection mode should be set (Skills, +// Query, or All); if none is set, all skills are materialized. +type SkillsRegistryConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + // Project owns the skills (projects/{Project}/locations/{Location}/skills). + // Empty falls back to the GOOGLE_CLOUD_PROJECT environment variable. + Project string `yaml:"project,omitempty"` + // Location is the registry region, e.g. "us-central1". Empty falls back to + // GOOGLE_CLOUD_LOCATION, then a built-in default. + Location string `yaml:"location,omitempty"` + + // --- selection (choose one) --- + + // Skills is an explicit allowlist of skills, each optionally pinned to a + // revision. Takes precedence over Query and All. + Skills []SkillRefConfig `yaml:"skills,omitempty"` + // Query is a semantic search selection; its top matches are materialized. + // TopK lives inside it because it only has meaning for a query. + Query *SkillsQueryConfig `yaml:"query,omitempty"` + // All materializes every skill in the project/location (used when neither + // Skills nor Query is set; can also be set explicitly). + All bool `yaml:"all,omitempty"` + + // TargetDir is the base directory skills are materialized into: each skill + // is written to //. Required when Enabled. + TargetDir string `yaml:"target_dir,omitempty"` +} + +// SkillsQueryConfig selects skills by semantic search. +type SkillsQueryConfig struct { + // Text is the semantic search string (required for query selection). + Text string `yaml:"text"` + // TopK bounds the number of matches (<=0 uses the server default). + TopK int `yaml:"top_k,omitempty"` +} + +// SkillRefConfig identifies a skill to materialize, optionally pinned. +type SkillRefConfig struct { + ID string `yaml:"id"` + Revision string `yaml:"revision,omitempty"` } // SubstrateHarnessConfig registers a custom harness deployed on substrate @@ -216,6 +316,10 @@ func (c *Config) Validate() error { return fmt.Errorf("multiple harnesses marked as default") } + if err := c.Skills.Validate(); err != nil { + return err + } + return nil } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 149b7a64..baebb037 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -168,3 +168,180 @@ func TestLoadFromBytes_Invalid(t *testing.T) { t.Fatal("LoadFromBytes(invalid): got nil error, want error") } } + +func TestParse_SkillsByID(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: my-proj + location: us-central1 + target_dir: /tmp/ax-skills + skills: + - id: emoji + - id: lowercase + revision: rev-3 +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + if len(cfg.Skills.Registries) != 1 { + t.Fatalf("registries = %d, want 1", len(cfg.Skills.Registries)) + } + reg := cfg.Skills.Registries[0] + if !reg.Enabled || reg.Project != "my-proj" || reg.Location != "us-central1" || reg.TargetDir != "/tmp/ax-skills" { + t.Errorf("registry = %+v, want enabled my-proj/us-central1 /tmp/ax-skills", reg) + } + if len(reg.Skills) != 2 { + t.Fatalf("skills = %d, want 2", len(reg.Skills)) + } + if reg.Skills[0].ID != "emoji" || reg.Skills[0].Revision != "" { + t.Errorf("skills[0] = %+v, want {emoji }", reg.Skills[0]) + } + if reg.Skills[1].ID != "lowercase" || reg.Skills[1].Revision != "rev-3" { + t.Errorf("skills[1] = %+v, want {lowercase rev-3}", reg.Skills[1]) + } + if reg.Query != nil { + t.Errorf("Query = %+v, want nil in by-id mode", reg.Query) + } +} + +func TestParse_SkillsByQuery(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: my-proj + target_dir: /tmp/ax-skills + query: + text: "find gcp skills" + top_k: 5 +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + regs := cfg.Skills.Registries + if len(regs) != 1 || regs[0].Query == nil { + t.Fatalf("registries = %+v, want one with a query block", regs) + } + if regs[0].Query.Text != "find gcp skills" || regs[0].Query.TopK != 5 { + t.Errorf("Query = %+v, want {find gcp skills 5}", *regs[0].Query) + } +} + +func TestParse_MultipleRegistries(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: org-proj + target_dir: /tmp/org + all: true + - enabled: true + project: team-proj + target_dir: /tmp/team + skills: + - id: teamskill +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + regs := cfg.Skills.Registries + if len(regs) != 2 { + t.Fatalf("registries = %d, want 2", len(regs)) + } + if regs[0].Project != "org-proj" || regs[0].TargetDir != "/tmp/org" || !regs[0].All { + t.Errorf("registries[0] = %+v, want org-proj /tmp/org all", regs[0]) + } + if regs[1].Project != "team-proj" || regs[1].TargetDir != "/tmp/team" || len(regs[1].Skills) != 1 { + t.Errorf("registries[1] = %+v, want team-proj /tmp/team [teamskill]", regs[1]) + } + // Each registry has exactly one selection mode + target_dir, so validation passes. + if err := cfg.Skills.Validate(); err != nil { + t.Errorf("Skills.Validate() = %v, want nil", err) + } +} + +func TestValidate_SkillsSelectionOneof(t *testing.T) { + // withRegistry returns a valid config carrying a single top-level registry. + // It fills TargetDir (a required field) unless the caller already set one, so + // selection-mode assertions aren't masked by the target_dir check. + withRegistry := func(rc SkillsRegistryConfig) *Config { + if rc.Enabled && rc.TargetDir == "" { + rc.TargetDir = "/tmp/skills" + } + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{rc} + return c + } + + t.Run("disabled skips validation", func(t *testing.T) { + if err := withRegistry(SkillsRegistryConfig{Enabled: false}).Validate(); err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } + }) + + t.Run("enabled without target_dir is an error", func(t *testing.T) { + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{ + {Enabled: true, Project: "p", All: true}, // valid mode, but no target_dir + } + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "target_dir") { + t.Fatalf("Validate() = %v, want target_dir error", err) + } + }) + + t.Run("zero selection modes is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{Enabled: true, Project: "p"}).Validate() + if err == nil || !strings.Contains(err.Error(), "exactly one selection mode") { + t.Fatalf("Validate() = %v, want exactly-one error", err) + } + }) + + t.Run("multiple selection modes is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Skills: []SkillRefConfig{{ID: "emoji"}}, + All: true, + }).Validate() + if err == nil || !strings.Contains(err.Error(), "multiple selection modes") { + t.Fatalf("Validate() = %v, want multiple-modes error", err) + } + }) + + t.Run("exactly one is valid", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Skills: []SkillRefConfig{{ID: "emoji"}}, + }).Validate() + if err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } + }) + + t.Run("query with empty text is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Query: &SkillsQueryConfig{Text: ""}, + }).Validate() + if err == nil || !strings.Contains(err.Error(), "non-empty text") { + t.Fatalf("Validate() = %v, want empty-text error", err) + } + }) + + t.Run("second registry invalid is caught", func(t *testing.T) { + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{ + {Enabled: true, Project: "p", TargetDir: "/tmp/a", All: true}, + {Enabled: true, Project: "p", TargetDir: "/tmp/b"}, // no selection mode + } + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "registries[1]") { + t.Fatalf("Validate() = %v, want error citing registries[1]", err) + } + }) +} diff --git a/internal/harness/antigravityinteractions/skills.go b/internal/harness/antigravityinteractions/skills.go new file mode 100644 index 00000000..5e0d57f8 --- /dev/null +++ b/internal/harness/antigravityinteractions/skills.go @@ -0,0 +1,61 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "fmt" + "strings" + + "github.com/google/ax/internal/skills/geminienterprise" +) + +// SkillsSystemInstruction builds a system-instruction pointer telling the agent +// where its materialized skills live and lists them. +// +// This is discovery logic specific to the Antigravity Interactions harness: it +// has no SKILLS_DIR concept, so it must be told where to find skills via its +// system instruction (its built-in file tools then read that directory). +// Harnesses that auto-discover a skills directory (e.g. the Antigravity SDK +// harness via SKILLS_DIR) do not use this. +func SkillsSystemInstruction(res geminienterprise.Result) string { + if res.Empty() { + return "" + } + var b strings.Builder + for _, w := range res.Written { + if b.Len() > 0 { + b.WriteString("\n") + } + fmt.Fprintf(&b, "Agent skills are available under %s. Available skills:", w.Dir) + for _, s := range w.Skills { + fmt.Fprintf(&b, " %s", s.SkillID) + } + b.WriteString(". Read a skill's SKILL.md before using it.") + } + return b.String() +} + +// JoinSystemInstruction combines a user-configured system instruction with an +// optional skills pointer (from SkillsSystemInstruction), dropping empties. +func JoinSystemInstruction(base, pointer string) string { + parts := make([]string, 0, 2) + if strings.TrimSpace(base) != "" { + parts = append(parts, base) + } + if strings.TrimSpace(pointer) != "" { + parts = append(parts, pointer) + } + return strings.Join(parts, "\n\n") +} diff --git a/internal/harness/antigravityinteractions/skills_test.go b/internal/harness/antigravityinteractions/skills_test.go new file mode 100644 index 00000000..151e993b --- /dev/null +++ b/internal/harness/antigravityinteractions/skills_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "strings" + "testing" + + "github.com/google/ax/internal/skills/geminienterprise" +) + +func TestSkillsSystemInstruction(t *testing.T) { + t.Run("empty result yields empty pointer", func(t *testing.T) { + if got := SkillsSystemInstruction(geminienterprise.Result{}); got != "" { + t.Errorf("empty result pointer = %q, want empty", got) + } + }) + + t.Run("mentions dir and skill ids", func(t *testing.T) { + res := geminienterprise.Result{Written: []geminienterprise.Written{{ + Dir: "/workspace", + Skills: []geminienterprise.MaterializedSkill{{SkillID: "emoji"}, {SkillID: "lowercase"}}, + }}} + got := SkillsSystemInstruction(res) + if !strings.Contains(got, "/workspace") || !strings.Contains(got, "emoji") || !strings.Contains(got, "lowercase") { + t.Errorf("pointer = %q, want it to mention dir and skill ids", got) + } + }) + + t.Run("multiple registries produce multiple lines", func(t *testing.T) { + res := geminienterprise.Result{Written: []geminienterprise.Written{ + {Dir: "/a", Skills: []geminienterprise.MaterializedSkill{{SkillID: "s1"}}}, + {Dir: "/b", Skills: []geminienterprise.MaterializedSkill{{SkillID: "s2"}}}, + }} + got := SkillsSystemInstruction(res) + if !strings.Contains(got, "/a") || !strings.Contains(got, "/b") || !strings.Contains(got, "\n") { + t.Errorf("pointer = %q, want both dirs on separate lines", got) + } + }) +} + +func TestJoinSystemInstruction(t *testing.T) { + cases := []struct { + base, ptr, want string + }{ + {"", "", ""}, + {"base only", "", "base only"}, + {"", "ptr only", "ptr only"}, + {"base", "ptr", "base\n\nptr"}, + } + for _, c := range cases { + if got := JoinSystemInstruction(c.base, c.ptr); got != c.want { + t.Errorf("join(%q,%q) = %q, want %q", c.base, c.ptr, got, c.want) + } + } +} diff --git a/internal/skills/geminienterprise/client.go b/internal/skills/geminienterprise/client.go new file mode 100644 index 00000000..93e4f726 --- /dev/null +++ b/internal/skills/geminienterprise/client.go @@ -0,0 +1,400 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// This file is the internal Gemini Enterprise Skill Registry transport client: it talks to the +// Vertex AI v1beta1 REST API (ListSkills / GetSkill / GetSkillRevision / +// skills:retrieve), fetches skill payloads, and safe-unzips them to disk. The +// public, config-driven orchestration lives in materialize.go. + +package geminienterprise + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +// cloudPlatformScope is the OAuth2 scope required to call Vertex AI (read path +// needs roles/aiplatform.viewer). +const cloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform" + +// apiVersion is the Vertex AI API version the Skill Registry is exposed under. +const apiVersion = "v1beta1" + +// defaultHTTPTimeout bounds a single registry HTTP call. +const defaultHTTPTimeout = 60 * time.Second + +// client fetches a selected set of skills from the Gemini Enterprise Skill Registry and +// writes them into a target directory as agentskills.io skill folders. +type client struct { + baseURL string // .../v1beta1/projects/{project}/locations/{location}/skills + http *http.Client + ts oauth2.TokenSource + caps unzipCaps +} + +// clientOptions configures a client. newClient fills defaults for optional fields. +type clientOptions struct { + // --- Required --- + + // Project is the Google Cloud project that owns the skills. + Project string + // Location is the registry region, e.g. "us-central1". Selects the Vertex + // host (https://{Location}-aiplatform.googleapis.com) unless Endpoint is set. + Location string + + // --- Optional (newClient fills defaults) --- + + // Endpoint overrides the API host (scheme+host, no trailing slash), e.g. a + // sandbox host or an httptest.Server URL in unit tests. + Endpoint string + // HTTPClient is the client used for all registry calls. Defaults to + // http.Client{Timeout: 60s}. + HTTPClient *http.Client + // TokenSource provides the OAuth2 bearer token. If nil, ADC is used + // (roles/aiplatform.viewer). Set to a static source in tests. + TokenSource oauth2.TokenSource + // Caps bounds a defensive local unzip of untrusted payloads. + Caps unzipCaps +} + +// newClient constructs a client from options. It validates required fields and +// fills defaults; it performs no network I/O. +func newClient(opts clientOptions) (*client, error) { + if opts.Project == "" { + return nil, errors.New("registry: Project is required") + } + if opts.Location == "" { + return nil, errors.New("registry: Location is required") + } + + host := opts.Endpoint + if host == "" { + host = fmt.Sprintf("https://%s-aiplatform.googleapis.com", opts.Location) + } + host = strings.TrimSuffix(host, "/") + baseURL := fmt.Sprintf("%s/%s/projects/%s/locations/%s/skills", + host, apiVersion, opts.Project, opts.Location) + + httpClient := opts.HTTPClient + if httpClient == nil { + httpClient = &http.Client{Timeout: defaultHTTPTimeout} + } + + ts := opts.TokenSource + if ts == nil { + creds, err := google.FindDefaultCredentials(context.Background(), cloudPlatformScope) + if err != nil { + return nil, fmt.Errorf("registry: finding application default credentials: %w", err) + } + ts = creds.TokenSource + } + + return &client{ + baseURL: baseURL, + http: httpClient, + ts: ts, + caps: opts.Caps.withDefaults(), + }, nil +} + +// selection describes WHICH skills to fetch. Exactly one of the three modes +// should be set (all / by-id / by-query). +type selection struct { + // All fetches every skill returned by ListSkills for the project/location. + All bool + // SkillRefs is an explicit allowlist, each optionally pinned to a revision. + SkillRefs []skillRef + // Query is a semantic search string; the top matches are fetched. TopK bounds + // the result count (<=0 means the server default). + Query string + TopK int +} + +// skillRef identifies a single skill to fetch, optionally pinned. +type skillRef struct { + SkillID string + Revision string // empty => latest default revision +} + +// fetchResult reports the outcome of a fetch call. fetch is fail-safe: a skill +// that cannot be fetched or unzipped is recorded in skipped. +type fetchResult struct { + materialized []MaterializedSkill + skipped []skippedSkill +} + +// skippedSkill records a skill that was intentionally not written, and why. +type skippedSkill struct { + skillID string + reason error +} + +// fetch resolves the selection, fetches each skill's payload, and safe-unzips it +// into targetDir as //... targetDir is created if missing. +// +// claimed enforces first-wins across the whole operation: a skill id whose +// (targetDir, id) was already materialized (by this registry's own selection or +// an earlier registry sharing the dir) is skipped with a warning rather than +// overwriting the winner. regIdx is the registry's index, for log context. +// +// It is fail-safe: individual skill failures are collected in fetchResult.skipped; +// fetch returns a non-nil error only for whole-operation failures. +func (c *client) fetch(ctx context.Context, sel selection, targetDir string, claimed *claimSet, regIdx int) (*fetchResult, error) { + refs, err := c.resolveSelection(ctx, sel) + if err != nil { + return nil, fmt.Errorf("registry: resolving selection: %w", err) + } + if err := os.MkdirAll(targetDir, 0o755); err != nil { + return nil, fmt.Errorf("registry: creating target dir %q: %w", targetDir, err) + } + + res := &fetchResult{} + for _, ref := range refs { + // First-wins: skip (with a warning) any id already materialized into this + // dir, rather than overwriting it. + if won, byRegistry := claimed.claim(targetDir, ref.SkillID, regIdx); !won { + log.Printf("skills: registries[%d] skill %q already materialized into %s by registries[%d]; keeping the first, skipping this one", + regIdx, ref.SkillID, targetDir, byRegistry) + continue + } + mat, err := c.fetchAndWrite(ctx, ref, targetDir) + if err != nil { + res.skipped = append(res.skipped, skippedSkill{skillID: ref.SkillID, reason: err}) + continue + } + res.materialized = append(res.materialized, *mat) + } + return res, nil +} + +// resolveSelection turns a selection into a concrete list of skillRefs to fetch. +func (c *client) resolveSelection(ctx context.Context, sel selection) ([]skillRef, error) { + switch { + case len(sel.SkillRefs) > 0: + return sel.SkillRefs, nil + case sel.Query != "": + ids, err := c.retrieveSkillIDs(ctx, sel.Query, sel.TopK) + if err != nil { + return nil, err + } + return toRefs(ids), nil + case sel.All: + ids, err := c.listSkillIDs(ctx) + if err != nil { + return nil, err + } + return toRefs(ids), nil + default: + return nil, errors.New("empty selection: set All, SkillRefs, or Query") + } +} + +// fetchAndWrite fetches one skill's payload (latest or pinned) and unzips it. +func (c *client) fetchAndWrite(ctx context.Context, ref skillRef, targetDir string) (*MaterializedSkill, error) { + if ref.SkillID == "" { + return nil, errors.New("skill ref has empty SkillID") + } + payloadB64, revision, err := c.fetchPayload(ctx, ref) + if err != nil { + return nil, err + } + zipped, err := base64.StdEncoding.DecodeString(payloadB64) + if err != nil { + return nil, fmt.Errorf("decoding payload: %w", err) + } + skillDir := filepath.Join(targetDir, ref.SkillID) + // Replace any prior materialization of this skill so stale files don't linger. + if err := os.RemoveAll(skillDir); err != nil { + return nil, fmt.Errorf("clearing %q: %w", skillDir, err) + } + if err := safeUnzip(zipped, skillDir, c.caps); err != nil { + // Leave no partial dir behind on failure. + _ = os.RemoveAll(skillDir) + return nil, fmt.Errorf("unzipping: %w", err) + } + return &MaterializedSkill{SkillID: ref.SkillID, Revision: revision, Dir: skillDir}, nil +} + +// fetchPayload returns the base64 zippedFilesystem and the concrete revision id +// for a skill ref (GetSkill for latest, GetSkillRevision when pinned). +func (c *client) fetchPayload(ctx context.Context, ref skillRef) (payloadB64, revision string, err error) { + if ref.Revision != "" { + // GetSkillRevision: payload is nested under "skill". + var resp skillRevisionResponse + if err := c.getJSON(ctx, c.baseURL+"/"+ref.SkillID+"/revisions/"+ref.Revision, &resp); err != nil { + return "", "", err + } + if resp.Skill.ZippedFilesystem == "" { + return "", "", errors.New("GetSkillRevision: empty zippedFilesystem") + } + return resp.Skill.ZippedFilesystem, ref.Revision, nil + } + // GetSkill: payload at top level. + var resp skillResponse + if err := c.getJSON(ctx, c.baseURL+"/"+ref.SkillID, &resp); err != nil { + return "", "", err + } + if resp.ZippedFilesystem == "" { + return "", "", errors.New("GetSkill: empty zippedFilesystem") + } + return resp.ZippedFilesystem, resp.currentRevision(), nil +} + +// listSkillIDs pages through ListSkills and returns all skill ids. +func (c *client) listSkillIDs(ctx context.Context) ([]string, error) { + var ids []string + pageToken := "" + for { + u := c.baseURL + if pageToken != "" { + u += "?pageToken=" + url.QueryEscape(pageToken) + } + var resp listSkillsResponse + if err := c.getJSON(ctx, u, &resp); err != nil { + return nil, err + } + for _, s := range resp.Skills { + if id := s.id(); id != "" { + ids = append(ids, id) + } + } + if resp.NextPageToken == "" { + return ids, nil + } + pageToken = resp.NextPageToken + } +} + +// retrieveSkillIDs runs semantic search (skills:retrieve) and returns skill ids. +func (c *client) retrieveSkillIDs(ctx context.Context, query string, topK int) ([]string, error) { + u := c.baseURL + ":retrieve?query=" + url.QueryEscape(query) + if topK > 0 { + u += fmt.Sprintf("&topK=%d", topK) + } + var resp retrieveSkillsResponse + if err := c.getJSON(ctx, u, &resp); err != nil { + return nil, err + } + var ids []string + for _, r := range resp.RetrievedSkills { + if id := lastSegment(r.SkillName); id != "" { + ids = append(ids, id) + } + } + return ids, nil +} + +// getJSON issues an authenticated GET and decodes the JSON body into out. +func (c *client) getJSON(ctx context.Context, rawURL string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) + if err != nil { + return err + } + tok, err := c.ts.Token() + if err != nil { + return fmt.Errorf("obtaining token: %w", err) + } + tok.SetAuthHeader(req) + req.Header.Set("Content-Type", "application/json") + + resp, err := c.http.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(body), 512)) + } + if err := json.Unmarshal(body, out); err != nil { + return fmt.Errorf("decoding response: %w", err) + } + return nil +} + +// --- Wire types (AIP-standard shapes; see the design doc appendix). --- + +type listSkillsResponse struct { + Skills []skillResponse `json:"skills"` + NextPageToken string `json:"nextPageToken"` +} + +// retrieveSkillsResponse is a skills:retrieve result. Each hit carries the skill +// resource name in a flat "skillName" field (not a nested skill object). +type retrieveSkillsResponse struct { + RetrievedSkills []struct { + SkillName string `json:"skillName"` // projects/.../skills/{id} + } `json:"retrievedSkills"` +} + +// skillResponse is a Skill resource. GetSkill returns the payload at top level. +type skillResponse struct { + Name string `json:"name"` // projects/.../skills/{id} + DisplayName string `json:"displayName"` + Description string `json:"description"` + State string `json:"state"` + DefaultRevision string `json:"defaultRevision"` + ZippedFilesystem string `json:"zippedFilesystem"` +} + +// skillRevisionResponse is a GetSkillRevision result; the payload is nested +// under "skill". +type skillRevisionResponse struct { + Name string `json:"name"` // projects/.../skills/{id}/revisions/{rev} + Skill skillResponse `json:"skill"` +} + +// id extracts the trailing skill id from a resource name. +func (s skillResponse) id() string { return lastSegment(s.Name) } + +// currentRevision reports the skill's default revision id, if the server provided +// one (best-effort; may be empty for "latest"). +func (s skillResponse) currentRevision() string { return lastSegment(s.DefaultRevision) } + +func lastSegment(resourceName string) string { + if resourceName == "" { + return "" + } + parts := strings.Split(resourceName, "/") + return parts[len(parts)-1] +} + +func toRefs(ids []string) []skillRef { + refs := make([]skillRef, 0, len(ids)) + for _, id := range ids { + refs = append(refs, skillRef{SkillID: id}) + } + return refs +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "...(truncated)" +} diff --git a/internal/skills/geminienterprise/client_test.go b/internal/skills/geminienterprise/client_test.go new file mode 100644 index 00000000..39285b5c --- /dev/null +++ b/internal/skills/geminienterprise/client_test.go @@ -0,0 +1,414 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "archive/zip" + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "golang.org/x/oauth2" +) + +// zipFixtureB64 builds an in-memory zip from name->content and returns the base64 +// string the registry would return in `zippedFilesystem`. +func zipFixtureB64(t *testing.T, files map[string]string) string { + t.Helper() + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + for name, content := range files { + w, err := zw.Create(name) + if err != nil { + t.Fatalf("zip create %q: %v", name, err) + } + if _, err := w.Write([]byte(content)); err != nil { + t.Fatalf("zip write %q: %v", name, err) + } + } + if err := zw.Close(); err != nil { + t.Fatalf("zip close: %v", err) + } + return base64.StdEncoding.EncodeToString(buf.Bytes()) +} + +// rawZipFixture returns the raw zip bytes (for direct safeUnzip tests). +func rawZipFixture(t *testing.T, files map[string]string) []byte { + t.Helper() + b, err := base64.StdEncoding.DecodeString(zipFixtureB64(t, files)) + if err != nil { + t.Fatalf("decode fixture: %v", err) + } + return b +} + +// newFakeClient returns a client pointed at an httptest.Server whose handler is +// provided by the caller, with a no-op token source. +func newFakeClient(t *testing.T, handler http.HandlerFunc) *client { + t.Helper() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + c, err := newClient(clientOptions{ + Project: "test-project", + Location: "us-central1", + Endpoint: srv.URL, + TokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "fake"}), + }) + if err != nil { + t.Fatalf("newClient: %v", err) + } + return c +} + +func TestFetchByID_Latest(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{ + "SKILL.md": "---\nname: emoji\n---\n# do stuff", + "scripts/x.sh": "echo hi", + }) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + // GetSkill: .../skills/emoji -> zippedFilesystem at top level. + if !strings.HasSuffix(r.URL.Path, "/skills/emoji") { + http.Error(w, "unexpected path "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillResponse{ + Name: "projects/p/locations/us-central1/skills/emoji", + DefaultRevision: "projects/p/locations/us-central1/skills/emoji/revisions/rev-7", + ZippedFilesystem: payload, + }) + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "emoji"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.skipped) != 0 { + t.Fatalf("unexpected skipped: %+v", res.skipped) + } + if len(res.materialized) != 1 { + t.Fatalf("materialized = %d, want 1", len(res.materialized)) + } + got := res.materialized[0] + if got.SkillID != "emoji" || got.Revision != "rev-7" { + t.Errorf("got %+v, want SkillID=emoji Revision=rev-7", got) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "---\nname: emoji\n---\n# do stuff") + assertFile(t, filepath.Join(dir, "emoji", "scripts", "x.sh"), "echo hi") +} + +func TestFetchByID_PinnedRevision(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{"SKILL.md": "pinned"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + // GetSkillRevision: .../skills/emoji/revisions/rev-3 -> nested under "skill". + if !strings.HasSuffix(r.URL.Path, "/skills/emoji/revisions/rev-3") { + http.Error(w, "unexpected path "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillRevisionResponse{ + Name: "projects/p/locations/us-central1/skills/emoji/revisions/rev-3", + Skill: skillResponse{ZippedFilesystem: payload}, + }) + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "emoji", Revision: "rev-3"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 1 || res.materialized[0].Revision != "rev-3" { + t.Fatalf("got %+v, want one skill pinned at rev-3", res.materialized) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "pinned") +} + +func TestFetchAll_Paginated(t *testing.T) { + pa := zipFixtureB64(t, map[string]string{"SKILL.md": "a"}) + pb := zipFixtureB64(t, map[string]string{"SKILL.md": "b"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/skills") && r.URL.Query().Get("pageToken") == "": + writeJSON(t, w, listSkillsResponse{ + Skills: []skillResponse{{Name: ".../skills/a"}}, + NextPageToken: "TOK", + }) + case strings.HasSuffix(r.URL.Path, "/skills") && r.URL.Query().Get("pageToken") == "TOK": + writeJSON(t, w, listSkillsResponse{Skills: []skillResponse{{Name: ".../skills/b"}}}) + case strings.HasSuffix(r.URL.Path, "/skills/a"): + writeJSON(t, w, skillResponse{Name: ".../skills/a", ZippedFilesystem: pa}) + case strings.HasSuffix(r.URL.Path, "/skills/b"): + writeJSON(t, w, skillResponse{Name: ".../skills/b", ZippedFilesystem: pb}) + default: + http.Error(w, "unexpected "+r.URL.String(), http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), selection{All: true}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 2 { + t.Fatalf("materialized = %d, want 2 (paged)", len(res.materialized)) + } + assertFile(t, filepath.Join(dir, "a", "SKILL.md"), "a") + assertFile(t, filepath.Join(dir, "b", "SKILL.md"), "b") +} + +func TestFetchByQuery(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{"SKILL.md": "found"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.Contains(r.URL.Path, "/skills:retrieve"): + if got := r.URL.Query().Get("query"); got != "emoji stuff" { + t.Errorf("query = %q, want %q", got, "emoji stuff") + } + writeJSON(t, w, retrieveSkillsResponse{ + RetrievedSkills: []struct { + SkillName string `json:"skillName"` + }{{SkillName: ".../skills/emoji"}}, + }) + case strings.HasSuffix(r.URL.Path, "/skills/emoji"): + writeJSON(t, w, skillResponse{Name: ".../skills/emoji", ZippedFilesystem: payload}) + default: + http.Error(w, "unexpected "+r.URL.String(), http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{Query: "emoji stuff", TopK: 3}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 1 { + t.Fatalf("materialized = %d, want 1", len(res.materialized)) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "found") +} + +func TestFetchFailSafe_OneBadSkillSkipped(t *testing.T) { + good := zipFixtureB64(t, map[string]string{"SKILL.md": "ok"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/skills/good"): + writeJSON(t, w, skillResponse{Name: ".../skills/good", ZippedFilesystem: good}) + case strings.HasSuffix(r.URL.Path, "/skills/bad"): + http.Error(w, "boom", http.StatusInternalServerError) + default: + http.Error(w, "unexpected", http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "good"}, {SkillID: "bad"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch returned whole-op error, want fail-safe: %v", err) + } + if len(res.materialized) != 1 || res.materialized[0].SkillID != "good" { + t.Errorf("materialized = %+v, want only 'good'", res.materialized) + } + if len(res.skipped) != 1 || res.skipped[0].skillID != "bad" { + t.Errorf("skipped = %+v, want only 'bad'", res.skipped) + } + // The bad skill left no directory behind. + if _, err := os.Stat(filepath.Join(dir, "bad")); !os.IsNotExist(err) { + t.Errorf("expected no dir for 'bad', stat err = %v", err) + } +} + +func TestFetch_FirstWinsAcrossSharedClaimSet(t *testing.T) { + // Two fetches (simulating two registries) into the SAME dir with a SHARED + // claim set. Both offer skill id "dup"; only the first should be written. + payloadA := zipFixtureB64(t, map[string]string{"SKILL.md": "from-A"}) + payloadB := zipFixtureB64(t, map[string]string{"SKILL.md": "from-B"}) + + makeClient := func(payload string) *client { + return newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/skills/dup") { + http.Error(w, "unexpected "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillResponse{Name: ".../skills/dup", ZippedFilesystem: payload}) + }) + } + + dir := t.TempDir() + claimed := newClaimSet() + sel := selection{SkillRefs: []skillRef{{SkillID: "dup"}}} + + // Registry 0 wins. + res0, err := makeClient(payloadA).fetch(context.Background(), sel, dir, claimed, 0) + if err != nil { + t.Fatalf("fetch #0: %v", err) + } + if len(res0.materialized) != 1 { + t.Fatalf("registry 0 materialized %d, want 1", len(res0.materialized)) + } + + // Registry 1 offers the same id into the same dir -> skipped (not written). + res1, err := makeClient(payloadB).fetch(context.Background(), sel, dir, claimed, 1) + if err != nil { + t.Fatalf("fetch #1: %v", err) + } + if len(res1.materialized) != 0 { + t.Errorf("registry 1 materialized %d, want 0 (first-wins)", len(res1.materialized)) + } + // The content on disk must be the FIRST writer's. + assertFile(t, filepath.Join(dir, "dup", "SKILL.md"), "from-A") +} + +func TestNewClientValidation(t *testing.T) { + if _, err := newClient(clientOptions{Location: "us-central1"}); err == nil { + t.Error("expected error for missing Project") + } + if _, err := newClient(clientOptions{Project: "p"}); err == nil { + t.Error("expected error for missing Location") + } +} + +// --- safeUnzip unit tests --- + +func TestSafeUnzip_Basic(t *testing.T) { + z := rawZipFixture(t, map[string]string{ + "SKILL.md": "hi", + "scripts/tool.sh": "run", + "references/a.txt": "ref", + }) + dir := t.TempDir() + if err := safeUnzip(z, filepath.Join(dir, "s"), unzipCaps{}.withDefaults()); err != nil { + t.Fatalf("safeUnzip: %v", err) + } + assertFile(t, filepath.Join(dir, "s", "SKILL.md"), "hi") + assertFile(t, filepath.Join(dir, "s", "scripts", "tool.sh"), "run") +} + +func TestSafeUnzip_PreservesExecutableMode(t *testing.T) { + // Build a zip with an executable script (0755) and a normal file (0644), + // setting per-entry modes via CreateHeader. + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + writeEntry := func(name, content string, mode os.FileMode) { + hdr := &zip.FileHeader{Name: name, Method: zip.Deflate} + hdr.SetMode(mode) + w, err := zw.CreateHeader(hdr) + if err != nil { + t.Fatalf("CreateHeader %q: %v", name, err) + } + if _, err := w.Write([]byte(content)); err != nil { + t.Fatalf("write %q: %v", name, err) + } + } + writeEntry("scripts/tool.sh", "#!/bin/sh\necho hi\n", 0o755) + writeEntry("SKILL.md", "hi", 0o644) + if err := zw.Close(); err != nil { + t.Fatalf("zip close: %v", err) + } + + dir := t.TempDir() + if err := safeUnzip(buf.Bytes(), filepath.Join(dir, "s"), unzipCaps{}.withDefaults()); err != nil { + t.Fatalf("safeUnzip: %v", err) + } + + scriptInfo, err := os.Stat(filepath.Join(dir, "s", "scripts", "tool.sh")) + if err != nil { + t.Fatal(err) + } + if scriptInfo.Mode().Perm()&0o100 == 0 { + t.Errorf("script mode = %v, want owner-execute bit set (0755 preserved)", scriptInfo.Mode().Perm()) + } + mdInfo, err := os.Stat(filepath.Join(dir, "s", "SKILL.md")) + if err != nil { + t.Fatal(err) + } + if mdInfo.Mode().Perm()&0o111 != 0 { + t.Errorf("SKILL.md mode = %v, want no execute bits (0644)", mdInfo.Mode().Perm()) + } +} + +func TestSafeUnzip_ZipSlipRejected(t *testing.T) { + for _, bad := range []string{"../evil.txt", "a/../../evil.txt", "/etc/evil"} { + z := rawZipFixture(t, map[string]string{bad: "x"}) + dir := t.TempDir() + err := safeUnzip(z, filepath.Join(dir, "s"), unzipCaps{}.withDefaults()) + if err == nil { + t.Errorf("entry %q: expected rejection, got nil", bad) + continue + } + if _, statErr := os.Stat(filepath.Join(dir, "evil.txt")); !os.IsNotExist(statErr) { + t.Errorf("entry %q: file escaped destination", bad) + } + } +} + +func TestSafeUnzip_TooManyFiles(t *testing.T) { + files := map[string]string{} + for i := 0; i < 5; i++ { + files[fmt.Sprintf("f%d.txt", i)] = "x" + } + z := rawZipFixture(t, files) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxFiles: 3}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "exceeds cap") { + t.Fatalf("expected MaxFiles cap error, got %v", err) + } +} + +func TestSafeUnzip_TooLarge(t *testing.T) { + z := rawZipFixture(t, map[string]string{"big.txt": strings.Repeat("A", 1000)}) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxTotalUnzippedBytes: 100}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "size") { + t.Fatalf("expected size cap error, got %v", err) + } +} + +func TestSafeUnzip_TooDeep(t *testing.T) { + z := rawZipFixture(t, map[string]string{"a/b/c/d/e/f/g/h/i/j/deep.txt": "x"}) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxDepth: 3}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "depth") { + t.Fatalf("expected depth cap error, got %v", err) + } +} + +// --- helpers --- + +func writeJSON(t *testing.T, w http.ResponseWriter, v any) { + t.Helper() + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(v); err != nil { + t.Fatalf("encode response: %v", err) + } +} + +func assertFile(t *testing.T, path, want string) { + t.Helper() + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("reading %q: %v", path, err) + } + if string(got) != want { + t.Errorf("%q = %q, want %q", path, got, want) + } +} diff --git a/internal/skills/geminienterprise/materialize.go b/internal/skills/geminienterprise/materialize.go new file mode 100644 index 00000000..9c2f9421 --- /dev/null +++ b/internal/skills/geminienterprise/materialize.go @@ -0,0 +1,173 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package geminienterprise turns a harness's skills configuration into on-disk +// skill folders. It sources agentskills.io skills from the Gemini Enterprise +// Skill Registry (a managed, versioned catalog exposed over the Vertex AI +// v1beta1 REST API) and writes each skill to //. +// +// It is harness-agnostic: it only writes files and reports what it wrote (see +// Result). It knows nothing about specific harnesses, SKILLS_DIR, or discovery +// pointers — callers decide how a given harness is told where its skills are. +// +// Scope: read-only. This package never creates, updates, or deletes registry +// skills (authoring is out of scope). +package geminienterprise + +import ( + "context" + "log" + "os" + "strings" + + "github.com/google/ax/internal/config" +) + +// Environment fallbacks for project/location (the registry target_dir is a +// required config field, so it has no env fallback). +const ( + envCloudProject = "GOOGLE_CLOUD_PROJECT" + envCloudLocation = "GOOGLE_CLOUD_LOCATION" + + defaultRegistryLocation = "us-central1" +) + +// Result reports what Materialize wrote: one Written entry per registry that +// produced skills, grouping the skills with the directory they landed in. +type Result struct { + Written []Written +} + +// Written groups the skills materialized from one registry with the directory +// they were written into (each skill at //). +type Written struct { + Dir string + Skills []MaterializedSkill +} + +// Empty reports whether nothing was materialized. +func (r Result) Empty() bool { return len(r.Written) == 0 } + +// MaterializedSkill records one skill written to disk. +type MaterializedSkill struct { + SkillID string + Revision string + Dir string // path of the written skill folder (//) +} + +// Materialize materializes every enabled registry in sc into its configured +// target_dir (each skill at //) and reports what it wrote. +// +// target_dir is a required, validated config field (see config.SkillsConfig), so +// this does not fall back to any env var or the working directory. It is +// fail-safe: disabled registries are skipped, and any registry error is logged +// and swallowed so a skill problem never blocks harness creation. When the same +// skill id would be written into the same dir more than once (within one +// registry's selection, or across registries sharing a dir), the FIRST writer +// wins and later duplicates are skipped with a warning. Substrate/pod +// materialization is a separate, later path; this wires the local flow only. +func Materialize(ctx context.Context, sc config.SkillsConfig) Result { + var res Result + // claimed tracks (target_dir, skill-id) pairs already written across all + // registries so the FIRST writer of an id into a dir wins; later duplicates + // (within one registry's selection, or across registries sharing a dir) are + // skipped with a warning instead of silently overwriting. + claimed := newClaimSet() + // TODO: fetches are sequential (both here and per-skill in client.fetch); + // consider bounded-concurrent download/unzip to speed up materialization. + // TODO: bound overall materialization with a deadline (only per-call HTTP + // timeouts exist today); e.g. a default ~120s that users can override. + for i := range sc.Registries { + rc := sc.Registries[i] + if !rc.Enabled { + continue + } + project := firstNonEmpty(rc.Project, os.Getenv(envCloudProject)) + if project == "" { + log.Printf("skills: registries[%d] enabled but no project (config or %s); skipping", i, envCloudProject) + continue + } + location := firstNonEmpty(rc.Location, os.Getenv(envCloudLocation), defaultRegistryLocation) + + c, err := newClient(clientOptions{Project: project, Location: location}) + if err != nil { + log.Printf("skills: registries[%d] client init failed: %v; skipping", i, err) + continue + } + out, err := c.fetch(ctx, selectionFromConfig(rc), rc.TargetDir, claimed, i) + if err != nil { + log.Printf("skills: registries[%d] fetch failed: %v; continuing", i, err) + continue + } + for _, s := range out.skipped { + log.Printf("skills: registries[%d] skipped %q: %v", i, s.skillID, s.reason) + } + if len(out.materialized) == 0 { + continue + } + log.Printf("skills: registries[%d] materialized %d skill(s) into %s (skipped %d)", + i, len(out.materialized), rc.TargetDir, len(out.skipped)) + res.Written = append(res.Written, Written{Dir: rc.TargetDir, Skills: out.materialized}) + } + return res +} + +// claimSet tracks which (dir, skill-id) pairs have already been materialized, so +// the first writer of a given id into a given dir wins. +type claimSet struct { + seen map[string]int // key -> registry index that first claimed it +} + +func newClaimSet() *claimSet { return &claimSet{seen: map[string]int{}} } + +func claimKey(dir, id string) string { return dir + "\x00" + id } + +// claim records (dir, id) as owned by registry regIdx and returns true if this +// is the first claim. If already claimed, it returns false and the index of the +// registry that won. +func (c *claimSet) claim(dir, id string, regIdx int) (won bool, byRegistry int) { + key := claimKey(dir, id) + if prev, ok := c.seen[key]; ok { + return false, prev + } + c.seen[key] = regIdx + return true, regIdx +} + +// selectionFromConfig maps a SkillsRegistryConfig to a selection: +// - explicit Skills list => by-id (with optional revision pin) +// - else Query => by-query +// - else => all +func selectionFromConfig(rc config.SkillsRegistryConfig) selection { + if len(rc.Skills) > 0 { + refs := make([]skillRef, 0, len(rc.Skills)) + for _, s := range rc.Skills { + refs = append(refs, skillRef{SkillID: s.ID, Revision: s.Revision}) + } + return selection{SkillRefs: refs} + } + if rc.Query != nil && strings.TrimSpace(rc.Query.Text) != "" { + return selection{Query: rc.Query.Text, TopK: rc.Query.TopK} + } + return selection{All: true} +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} diff --git a/internal/skills/geminienterprise/materialize_test.go b/internal/skills/geminienterprise/materialize_test.go new file mode 100644 index 00000000..2779a377 --- /dev/null +++ b/internal/skills/geminienterprise/materialize_test.go @@ -0,0 +1,100 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "context" + "testing" + + "github.com/google/ax/internal/config" +) + +func TestSelectionFromConfig(t *testing.T) { + t.Run("explicit skills take precedence", func(t *testing.T) { + rc := config.SkillsRegistryConfig{ + Skills: []config.SkillRefConfig{{ID: "emoji"}, {ID: "lowercase", Revision: "rev-3"}}, + Query: &config.SkillsQueryConfig{Text: "ignored"}, + } + sel := selectionFromConfig(rc) + if sel.All || sel.Query != "" { + t.Fatalf("got %+v, want by-id", sel) + } + if len(sel.SkillRefs) != 2 || + sel.SkillRefs[0].SkillID != "emoji" || sel.SkillRefs[0].Revision != "" || + sel.SkillRefs[1].SkillID != "lowercase" || sel.SkillRefs[1].Revision != "rev-3" { + t.Fatalf("refs = %+v", sel.SkillRefs) + } + }) + + t.Run("query when no explicit skills", func(t *testing.T) { + sel := selectionFromConfig(config.SkillsRegistryConfig{ + Query: &config.SkillsQueryConfig{Text: "find gcp", TopK: 5}, + }) + if sel.Query != "find gcp" || sel.TopK != 5 { + t.Fatalf("got %+v, want query=find gcp topK=5", sel) + } + }) + + t.Run("empty query text falls through to all", func(t *testing.T) { + sel := selectionFromConfig(config.SkillsRegistryConfig{Query: &config.SkillsQueryConfig{Text: " "}}) + if !sel.All { + t.Fatalf("got %+v, want All (blank query text)", sel) + } + }) + + t.Run("all when nothing set", func(t *testing.T) { + if sel := selectionFromConfig(config.SkillsRegistryConfig{}); !sel.All { + t.Fatalf("got %+v, want All", sel) + } + }) +} + +func TestClaimSet_FirstWins(t *testing.T) { + cs := newClaimSet() + if won, by := cs.claim("/dir", "s1", 0); !won || by != 0 { + t.Fatalf("first claim = (%v,%d), want (true,0)", won, by) + } + // Same (dir,id) again -> loses, reports the first registry (0). + if won, by := cs.claim("/dir", "s1", 2); won || by != 0 { + t.Fatalf("dup claim = (%v,%d), want (false,0)", won, by) + } + // Same id, different dir -> independent, wins. + if won, _ := cs.claim("/other", "s1", 1); !won { + t.Fatal("same id in different dir should win") + } + // Different id, same dir -> wins. + if won, _ := cs.claim("/dir", "s2", 1); !won { + t.Fatal("different id in same dir should win") + } +} + +func TestMaterialize_DisabledIsEmpty(t *testing.T) { + // Disabled config => no-op, empty result, no error/panic. + if res := Materialize(context.Background(), config.SkillsConfig{}); !res.Empty() { + t.Errorf("disabled Materialize = %+v, want empty", res) + } +} + +func TestMaterialize_EnabledNoProjectIsEmpty(t *testing.T) { + // Enabled with a target_dir but no project (config empty, GOOGLE_CLOUD_PROJECT + // unset) => fail-safe empty result (no panic, no materialization). + t.Setenv(envCloudProject, "") + sc := config.SkillsConfig{Registries: []config.SkillsRegistryConfig{ + {Enabled: true, TargetDir: t.TempDir()}, + }} + if res := Materialize(context.Background(), sc); !res.Empty() { + t.Errorf("no-project Materialize = %+v, want empty", res) + } +} diff --git a/internal/skills/geminienterprise/unzip.go b/internal/skills/geminienterprise/unzip.go new file mode 100644 index 00000000..df1046b9 --- /dev/null +++ b/internal/skills/geminienterprise/unzip.go @@ -0,0 +1,202 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "archive/zip" + "bytes" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// unzipCaps bounds a defensive local unzip of an untrusted skill payload. Even +// though the registry validates payloads server-side, the local unzip must +// independently guard against Zip-Slip (entries with "..", absolute paths, or +// symlinks) and resource exhaustion. +type unzipCaps struct { + // MaxFiles caps the number of entries per skill archive. + MaxFiles int + // MaxTotalUnzippedBytes caps the total unzipped size per skill archive. + MaxTotalUnzippedBytes int64 + // MaxDepth caps directory nesting depth within a skill archive. + MaxDepth int +} + +// The registry's documented payload caps, used as defaults. +const ( + defaultMaxFiles = 10_000 + defaultMaxTotalBytes = 500 << 20 // 500 MiB + defaultMaxDepth = 8 +) + +// withDefaults fills any zero-valued cap with the registry's documented limit. +func (c unzipCaps) withDefaults() unzipCaps { + if c.MaxFiles <= 0 { + c.MaxFiles = defaultMaxFiles + } + if c.MaxTotalUnzippedBytes <= 0 { + c.MaxTotalUnzippedBytes = defaultMaxTotalBytes + } + if c.MaxDepth <= 0 { + c.MaxDepth = defaultMaxDepth + } + return c +} + +// safeUnzip extracts a zip archive into destDir, defensively rejecting unsafe +// entries and enforcing caps. destDir must not exist yet or must be safe to +// write into; callers clear it beforehand. +// +// Guards: +// - Zip-Slip: every entry must resolve to a path inside destDir. +// - Absolute paths and paths containing ".." are rejected. +// - Symlinks (and any non-regular, non-dir mode) are rejected. +// - MaxFiles / MaxDepth / MaxTotalUnzippedBytes are enforced; the running +// total is checked while copying so a lying uncompressed-size can't be used +// to blow past the cap. +func safeUnzip(archive []byte, destDir string, caps unzipCaps) error { + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + return fmt.Errorf("opening zip: %w", err) + } + if len(zr.File) > caps.MaxFiles { + return fmt.Errorf("archive has %d entries, exceeds cap %d", len(zr.File), caps.MaxFiles) + } + + // Resolve destDir to an absolute, clean base for containment checks. + absDest, err := filepath.Abs(destDir) + if err != nil { + return fmt.Errorf("resolving dest: %w", err) + } + if err := os.MkdirAll(absDest, 0o755); err != nil { + return fmt.Errorf("creating dest: %w", err) + } + + var total int64 + for _, f := range zr.File { + if err := validateEntryName(f.Name, caps.MaxDepth); err != nil { + return err + } + // Reject symlinks and any special modes; only dirs and regular files. + mode := f.Mode() + if mode&os.ModeSymlink != 0 { + return fmt.Errorf("entry %q is a symlink (rejected)", f.Name) + } + if !mode.IsDir() && !mode.IsRegular() { + return fmt.Errorf("entry %q has unsupported mode %v (rejected)", f.Name, mode) + } + + target := filepath.Join(absDest, filepath.FromSlash(f.Name)) + // Containment: target must be within absDest. + if !withinBase(absDest, target) { + return fmt.Errorf("entry %q escapes destination (zip-slip)", f.Name) + } + + if f.FileInfo().IsDir() { + if err := os.MkdirAll(target, 0o755); err != nil { + return fmt.Errorf("creating dir %q: %w", f.Name, err) + } + continue + } + + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return fmt.Errorf("creating parent of %q: %w", f.Name, err) + } + written, err := writeCappedFile(f, target, caps.MaxTotalUnzippedBytes-total) + if err != nil { + return err + } + total += written + if total > caps.MaxTotalUnzippedBytes { + return fmt.Errorf("archive exceeds unzipped size cap %d bytes", caps.MaxTotalUnzippedBytes) + } + } + return nil +} + +// validateEntryName rejects absolute paths, "..", and over-deep nesting. +func validateEntryName(name string, maxDepth int) error { + if name == "" { + return fmt.Errorf("empty entry name") + } + // Reject Windows-style and POSIX absolute paths. + if filepath.IsAbs(name) || strings.HasPrefix(name, "/") || strings.HasPrefix(name, `\`) { + return fmt.Errorf("entry %q is an absolute path (rejected)", name) + } + slashed := strings.ReplaceAll(name, `\`, "/") + for _, seg := range strings.Split(slashed, "/") { + if seg == ".." { + return fmt.Errorf("entry %q contains '..' (rejected)", name) + } + } + depth := 0 + for _, seg := range strings.Split(strings.Trim(slashed, "/"), "/") { + if seg != "" && seg != "." { + depth++ + } + } + if depth > maxDepth { + return fmt.Errorf("entry %q nesting depth %d exceeds cap %d", name, depth, maxDepth) + } + return nil +} + +// withinBase reports whether target is inside base (after cleaning). +func withinBase(base, target string) bool { + rel, err := filepath.Rel(base, target) + if err != nil { + return false + } + return rel != ".." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) +} + +// writeCappedFile copies one zip entry to disk, refusing to write more than +// remaining bytes (so a mismatched declared size can't overrun the cap). +func writeCappedFile(f *zip.File, target string, remaining int64) (int64, error) { + if remaining < 0 { + return 0, fmt.Errorf("unzipped size cap exceeded before %q", f.Name) + } + rc, err := f.Open() + if err != nil { + return 0, fmt.Errorf("opening entry %q: %w", f.Name, err) + } + defer rc.Close() + + // Preserve the archive entry's permission bits so executable skill scripts + // (e.g. under scripts/) stay executable after extraction. Chmod after create + // so the mode sticks regardless of the process umask. + perm := f.Mode().Perm() + out, err := os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm) + if err != nil { + return 0, fmt.Errorf("creating file %q: %w", f.Name, err) + } + defer out.Close() + if err := out.Chmod(perm); err != nil { + return 0, fmt.Errorf("setting mode on %q: %w", f.Name, err) + } + + // Limit the copy to remaining+1 so we can detect overrun deterministically. + n, err := io.Copy(out, io.LimitReader(rc, remaining+1)) + if err != nil { + return n, fmt.Errorf("writing %q: %w", f.Name, err) + } + if n > remaining { + return n, fmt.Errorf("entry %q exceeds remaining unzipped size budget", f.Name) + } + return n, nil +}