use crate::error::SkillError;
use crate::registry::extract::{extract_skill_archive, write_files_to_dir};
use adk_core::{AdkError, ErrorCategory, ErrorComponent, Result};
use adk_gcp::{GcpErrorCodes, GcpErrorContext, GcpHttpClient, truncate_for_error};
use base64::Engine as _;
use google_cloud_auth::credentials::Credentials;
use reqwest::Method;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::Duration;
const SKILL_REGISTRY_API_VERSION: &str = "v1beta1";
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
const AUTH_HEADERS_TIMEOUT: Duration = Duration::from_secs(30);
const ENV_GOOGLE_CLOUD_PROJECT: &str = "GOOGLE_CLOUD_PROJECT";
const ENV_GOOGLE_CLOUD_LOCATION: &str = "GOOGLE_CLOUD_LOCATION";
const RETRIEVE_MAX_TOP_K: u32 = 100;
#[derive(Debug, Clone)]
pub struct SkillRegistryConfig {
pub project_id: String,
pub location: String,
pub endpoint: Option<String>,
}
impl SkillRegistryConfig {
pub fn new(project_id: impl Into<String>, location: impl Into<String>) -> Self {
Self { project_id: project_id.into(), location: location.into(), endpoint: None }
}
pub fn from_env() -> Result<Self> {
let read = |key: &str| {
std::env::var(key)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
};
let project_id = read(ENV_GOOGLE_CLOUD_PROJECT);
let location = read(ENV_GOOGLE_CLOUD_LOCATION);
match (project_id, location) {
(Some(project_id), Some(location)) => Ok(Self::new(project_id, location)),
(project_id, location) => {
let missing = [
(ENV_GOOGLE_CLOUD_PROJECT, project_id.is_none()),
(ENV_GOOGLE_CLOUD_LOCATION, location.is_none()),
]
.into_iter()
.filter_map(|(key, is_missing)| is_missing.then_some(key))
.collect::<Vec<_>>()
.join(", ");
Err(AdkError::new(
ErrorComponent::Tool,
ErrorCategory::InvalidInput,
"skill.registry.missing_env",
format!(
"missing or blank environment variable(s): {missing}. Set them, or construct the config with SkillRegistryConfig::new",
),
)
.with_provider("vertex_ai"))
}
}
}
pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
self.endpoint = Some(endpoint.into());
self
}
fn endpoint(&self) -> String {
let endpoint = self
.endpoint
.clone()
.unwrap_or_else(|| format!("https://{}-aiplatform.googleapis.com", self.location));
if endpoint.contains("://") { endpoint } else { format!("https://{endpoint}") }
}
fn parent_path(&self) -> String {
format!("projects/{}/locations/{}", self.project_id, self.location)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum SkillState {
#[default]
StateUnspecified,
Active,
Creating,
Failed,
Deleting,
#[serde(other)]
Unknown,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum SkillSource {
#[default]
User,
System,
#[serde(other)]
Unknown,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct Skill {
pub name: String,
pub display_name: String,
pub description: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub license: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub compatibility: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub zipped_filesystem: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub state: Option<SkillState>,
#[serde(skip_serializing_if = "BTreeMap::is_empty")]
pub labels: BTreeMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sha256: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub skill_source: Option<SkillSource>,
#[serde(skip_serializing_if = "Option::is_none")]
pub create_time: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub update_time: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct SkillRevision {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub create_time: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub skill: Option<Skill>,
#[serde(skip_serializing_if = "Option::is_none")]
pub state: Option<SkillState>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct ListSkillsResponse {
#[serde(skip_serializing_if = "Vec::is_empty")]
pub skills: Vec<Skill>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_page_token: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct ListSkillRevisionsResponse {
#[serde(skip_serializing_if = "Vec::is_empty")]
pub skill_revisions: Vec<SkillRevision>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_page_token: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct RetrievedSkill {
pub skill_name: String,
pub description: String,
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
struct RetrieveSkillsResponse {
retrieved_skills: Vec<RetrievedSkill>,
}
#[derive(Debug, Clone)]
pub struct SkillContent {
pub skill: Skill,
pub sha256: String,
pub files: BTreeMap<String, Vec<u8>>,
}
impl SkillContent {
pub const SKILL_MD: &'static str = "SKILL.md";
pub fn skill_md(&self) -> Option<&[u8]> {
self.files.get(Self::SKILL_MD).map(Vec::as_slice)
}
pub fn write_to_dir(&self, dir: &Path) -> std::result::Result<Vec<PathBuf>, SkillError> {
write_files_to_dir(&self.files, dir)
}
}
const GCP_ERROR_CODES: GcpErrorCodes = GcpErrorCodes {
invalid_input: "skill.registry.invalid_input",
unauthorized: "skill.registry.unauthorized",
forbidden: "skill.registry.forbidden",
not_found: "skill.registry.not_found",
rate_limited: "skill.registry.rate_limited",
timeout: "skill.registry.timeout",
unavailable: "skill.registry.unavailable",
credentials_unavailable: "skill.registry.credentials_unavailable",
invalid_response: "skill.registry.invalid_response",
invalid_request: "skill.registry.invalid_request",
upstream_error: "skill.registry.upstream_error",
operation_failed: "skill.registry.operation_failed",
};
pub struct SkillRegistryClient {
client: GcpHttpClient,
parent: String,
}
impl std::fmt::Debug for SkillRegistryClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SkillRegistryClient").field("parent", &self.parent).finish_non_exhaustive()
}
}
impl SkillRegistryClient {
pub fn new_with_adc(config: SkillRegistryConfig) -> Result<Self> {
Self::build(config, None)
}
pub fn with_credentials(config: SkillRegistryConfig, credentials: Credentials) -> Result<Self> {
Self::build(config, Some(credentials))
}
fn build(config: SkillRegistryConfig, credentials: Option<Credentials>) -> Result<Self> {
let errors = GcpErrorContext::new(ErrorComponent::Tool, GCP_ERROR_CODES, "skill registry");
let mut builder = GcpHttpClient::builder(errors, config.endpoint())
.api_version(SKILL_REGISTRY_API_VERSION)
.connect_timeout(HTTP_CONNECT_TIMEOUT)
.request_timeout(HTTP_REQUEST_TIMEOUT)
.auth_timeout(AUTH_HEADERS_TIMEOUT);
if let Some(credentials) = credentials {
builder = builder.credentials(credentials);
}
Ok(Self { client: builder.build()?, parent: config.parent_path() })
}
pub fn parent_resource_name(&self) -> &str {
&self.parent
}
pub(crate) fn error_context(&self) -> &GcpErrorContext {
self.client.errors()
}
fn skill_path(&self, skill: &str) -> String {
if skill.contains('/') {
skill.to_string()
} else {
format!("{}/skills/{skill}", self.parent)
}
}
pub async fn get_skill(&self, skill: &str) -> Result<Skill> {
self.get_json(&self.skill_path(skill), &[]).await
}
pub async fn list_skills(
&self,
page_size: Option<u32>,
page_token: Option<&str>,
) -> Result<ListSkillsResponse> {
let mut query = Vec::new();
if let Some(page_size) = page_size {
query.push(("pageSize", page_size.to_string()));
}
if let Some(page_token) = page_token {
query.push(("pageToken", page_token.to_string()));
}
self.get_json(&format!("{}/skills", self.parent), &query).await
}
pub async fn search_skills(
&self,
query: &str,
top_k: Option<u32>,
) -> Result<Vec<RetrievedSkill>> {
if let Some(top_k) = top_k
&& top_k > RETRIEVE_MAX_TOP_K
{
return Err(self.client.errors().invalid_input(format!(
"topK {top_k} exceeds the skills:retrieve maximum of {RETRIEVE_MAX_TOP_K}. Pass a value between 1 and {RETRIEVE_MAX_TOP_K}, or None for the server default of 10",
)));
}
let mut params = vec![("query", query.to_string())];
if let Some(top_k) = top_k {
params.push(("topK", top_k.to_string()));
}
let response: RetrieveSkillsResponse =
self.get_json(&format!("{}/skills:retrieve", self.parent), ¶ms).await?;
Ok(response.retrieved_skills)
}
pub async fn list_skill_revisions(
&self,
skill: &str,
page_size: Option<u32>,
page_token: Option<&str>,
filter: Option<&str>,
) -> Result<ListSkillRevisionsResponse> {
let mut query = Vec::new();
if let Some(page_size) = page_size {
query.push(("pageSize", page_size.to_string()));
}
if let Some(page_token) = page_token {
query.push(("pageToken", page_token.to_string()));
}
if let Some(filter) = filter {
query.push(("filter", filter.to_string()));
}
self.get_json(&format!("{}/revisions", self.skill_path(skill)), &query).await
}
pub async fn get_skill_revision(&self, skill: &str, revision: &str) -> Result<SkillRevision> {
let path = if revision.contains('/') {
revision.to_string()
} else {
format!("{}/revisions/{revision}", self.skill_path(skill))
};
self.get_json(&path, &[]).await
}
pub async fn fetch_skill_content(&self, skill: &str) -> Result<SkillContent> {
let skill = self.get_skill(skill).await?;
self.content_from_skill(skill)
}
pub async fn fetch_skill_revision_content(
&self,
skill: &str,
revision: &str,
) -> Result<SkillContent> {
let revision = self.get_skill_revision(skill, revision).await?;
let snapshot = revision.skill.ok_or_else(|| {
self.client.errors().invalid_response(format!(
"revision `{}` carried no embedded skill snapshot; fetch the latest payload with fetch_skill_content instead",
revision.name,
))
})?;
self.content_from_skill(snapshot)
}
fn content_from_skill(&self, mut skill: Skill) -> Result<SkillContent> {
let encoded = skill
.zipped_filesystem
.take()
.filter(|payload| !payload.trim().is_empty())
.ok_or_else(|| {
self.client.errors().invalid_response(format!(
"skill `{}` response carried no zippedFilesystem payload; the skill may still be CREATING",
skill.name,
))
})?;
let bytes =
base64::engine::general_purpose::STANDARD.decode(encoded.trim()).map_err(|error| {
SkillError::RegistryPayloadDecode {
message: truncate_for_error(&error.to_string()),
}
})?;
let computed = format!("{:x}", Sha256::digest(&bytes));
match skill.sha256.as_deref().map(str::trim).filter(|digest| !digest.is_empty()) {
Some(expected) if !expected.eq_ignore_ascii_case(&computed) => {
return Err(SkillError::RegistryChecksumMismatch {
expected: expected.to_string(),
actual: computed,
}
.into());
}
Some(_) => {}
None => {
tracing::warn!(
skill.name = %skill.name,
"skill registry response carried no sha256 digest; skipping verification"
);
}
}
let files = extract_skill_archive(&bytes)?;
tracing::debug!(
skill.name = %skill.name,
file.count = files.len(),
"fetched and extracted skill content"
);
Ok(SkillContent { skill, sha256: computed, files })
}
async fn get_json<R: for<'de> Deserialize<'de>>(
&self,
path: &str,
query: &[(&str, String)],
) -> Result<R> {
tracing::debug!(skill_registry.path = path, "sending skill registry request");
let mut request = self.client.request(Method::GET, path).await?;
if !query.is_empty() {
request = request.query(query);
}
let value = self.client.send_value(request).await?;
serde_json::from_value(value).map_err(|error| {
let error = truncate_for_error(&error.to_string());
self.client
.errors()
.invalid_response(format!("failed to parse skill registry response JSON: {error}"))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_config_resolves_paths_and_default_endpoint() {
let config = SkillRegistryConfig::new("p", "europe-west4");
assert_eq!(config.parent_path(), "projects/p/locations/europe-west4");
assert_eq!(config.endpoint(), "https://europe-west4-aiplatform.googleapis.com");
}
#[test]
fn test_skill_deserializes_leniently_with_unknown_state_and_fields() {
let skill: Skill = serde_json::from_value(json!({
"name": "projects/p/locations/l/skills/s",
"displayName": "s",
"description": "d",
"state": "SOME_FUTURE_STATE",
"skillSource": "SOME_FUTURE_SOURCE",
"undocumentedField": { "nested": true },
}))
.unwrap();
assert_eq!(skill.state, Some(SkillState::Unknown));
assert_eq!(skill.skill_source, Some(SkillSource::Unknown));
assert!(skill.zipped_filesystem.is_none());
let revision: SkillRevision = serde_json::from_value(json!({
"name": "projects/p/locations/l/skills/s/revisions/1",
"updateTime": "2026-01-01T00:00:00Z",
}))
.unwrap();
assert!(revision.skill.is_none());
}
#[tokio::test]
async fn test_endpoint_rejects_cleartext_and_decorated_origins() {
let credentials =
google_cloud_auth::credentials::api_key_credentials::Builder::new("k").build();
let config =
SkillRegistryConfig::new("p", "us-central1").with_endpoint("http://example.com");
let error = SkillRegistryClient::with_credentials(config, credentials.clone()).unwrap_err();
assert!(error.message.contains("HTTPS"), "unexpected error: {}", error.message);
let config =
SkillRegistryConfig::new("p", "us-central1").with_endpoint("https://example.com/path");
let error = SkillRegistryClient::with_credentials(config, credentials).unwrap_err();
assert!(error.message.contains("origin"), "unexpected error: {}", error.message);
}
}