feat(agents): persist agent capabilities

buf.lock

@@ -4,8 +4,8 @@ deps:
   - remote: buf.build
     owner: agynio
     repository: api
-    commit: 8352aa52d6d24a16a83fcaac76a44348
-    digest: shake256:f6f74b097b9349ff456e23e8763f9fc0f99bc8e45402eaeb573351544e717ac1f0ceeb8851c6a0f224af6ea33e665ecc9ecb85e2246219b45c3a4bde6d820857
+    commit: f3017f9d17f34204bb16b2d4965e8e58
+    digest: shake256:d7a844b0d81eb46525351d250d8b7df473f5dd907aa1b1e26273cf27fab84f0edde9844c07924e86f1f7583ed5204371aa1647f7d09165563973b0228206c086
   - remote: buf.build
     owner: opentelemetry
     repository: opentelemetry

internal/server/converter.go

@@ -35,6 +35,7 @@ func toProtoAgent(agent store.Agent) *agentsv1.Agent {
 		Configuration:  agent.Configuration,
 		Image:          agent.Image,
 		InitImage:      agent.InitImage,
+		Capabilities:   append([]string(nil), agent.Capabilities...),
 		Resources:      toProtoComputeResources(agent.Resources),
 	}
 	if agent.IdleTimeout != nil {

internal/server/server.go

@@ -98,6 +98,7 @@ func (s *Server) CreateAgent(ctx context.Context, req *agentsv1.CreateAgentReque
 		Image:         req.GetImage(),
 		InitImage:     req.GetInitImage(),
 		IdleTimeout:   &idleTimeout,
+		Capabilities:  append([]string(nil), req.GetCapabilities()...),
 		Resources:     resources,
 	})
 	if err != nil {
@@ -158,7 +159,11 @@ func (s *Server) UpdateAgent(ctx context.Context, req *agentsv1.UpdateAgentReque
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "id: %v", err)
 	}
-	if req.Name == nil && req.Nickname == nil && req.Role == nil && req.Model == nil && req.Description == nil && req.Configuration == nil && req.Image == nil && req.InitImage == nil && req.IdleTimeout == nil && req.Resources == nil {
+	// NOTE: proto3 repeated fields do not track presence on the wire. A nil
+	// slice indicates the caller did not set capabilities; when provided, the
+	// list replaces existing capabilities.
+	capabilitiesProvided := req.Capabilities != nil
+	if req.Name == nil && req.Nickname == nil && req.Role == nil && req.Model == nil && req.Description == nil && req.Configuration == nil && req.Image == nil && req.InitImage == nil && req.IdleTimeout == nil && req.Resources == nil && !capabilitiesProvided {
 		return nil, status.Error(codes.InvalidArgument, "at least one field must be provided")
 	}
 	if req.InitImage != nil && req.GetInitImage() == "" {
@@ -221,6 +226,10 @@ func (s *Server) UpdateAgent(ctx context.Context, req *agentsv1.UpdateAgentReque
 		}
 		update.IdleTimeout = &value
 	}
+	if capabilitiesProvided {
+		value := append([]string(nil), req.GetCapabilities()...)
+		update.Capabilities = &value
+	}
 	if req.Resources != nil {
 		resources := toStoreComputeResources(req.GetResources())
 		update.Resources = &resources

internal/store/store.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 
@@ -13,7 +14,7 @@ import (
 )
 
 const (
-	agentColumns                     = `id, organization_id, name, nickname, role, model, description, configuration, image, init_image, idle_timeout, resources_requests_cpu, resources_requests_memory, resources_limits_cpu, resources_limits_memory, created_at, updated_at`
+	agentColumns                     = `id, organization_id, name, nickname, role, model, description, configuration, image, init_image, idle_timeout, capabilities, resources_requests_cpu, resources_requests_memory, resources_limits_cpu, resources_limits_memory, created_at, updated_at`
 	volumeColumns                    = `id, organization_id, persistent, mount_path, size, description, ttl, created_at, updated_at`
 	volumeAttachmentColumns          = `id, volume_id, agent_id, mcp_id, hook_id, created_at, updated_at`
 	imagePullSecretAttachmentColumns = `id, image_pull_secret_id, agent_id, mcp_id, hook_id, created_at, updated_at`
@@ -48,9 +49,35 @@ func stringPtrFromPg(value pgtype.Text) *string {
 	return &text
 }
 
+func decodeCapabilities(value []byte) ([]string, error) {
+	if value == nil {
+		return nil, fmt.Errorf("capabilities is NULL")
+	}
+	var capabilities []string
+	if err := json.Unmarshal(value, &capabilities); err != nil {
+		return nil, fmt.Errorf("decode capabilities: %w", err)
+	}
+	if capabilities == nil {
+		return nil, fmt.Errorf("capabilities must be a JSON array")
+	}
+	return capabilities, nil
+}
+
+func encodeCapabilities(capabilities []string) ([]byte, error) {
+	if capabilities == nil {
+		capabilities = []string{}
+	}
+	data, err := json.Marshal(capabilities)
+	if err != nil {
+		return nil, fmt.Errorf("encode capabilities: %w", err)
+	}
+	return data, nil
+}
+
 func scanAgent(row pgx.Row) (Agent, error) {
 	var agent Agent
 	var idleTimeout pgtype.Text
+	var capabilities []byte
 	if err := row.Scan(
 		&agent.Meta.ID,
 		&agent.OrganizationID,
@@ -63,6 +90,7 @@ func scanAgent(row pgx.Row) (Agent, error) {
 		&agent.Image,
 		&agent.InitImage,
 		&idleTimeout,
+		&capabilities,
 		&agent.Resources.RequestsCPU,
 		&agent.Resources.RequestsMemory,
 		&agent.Resources.LimitsCPU,
@@ -73,6 +101,11 @@ func scanAgent(row pgx.Row) (Agent, error) {
 		return Agent{}, err
 	}
 	agent.IdleTimeout = stringPtrFromPg(idleTimeout)
+	decodedCapabilities, err := decodeCapabilities(capabilities)
+	if err != nil {
+		return Agent{}, err
+	}
+	agent.Capabilities = decodedCapabilities
 	return agent, nil
 }
 
@@ -251,9 +284,13 @@ func scanInitScript(row pgx.Row) (InitScript, error) {
 }
 
 func (s *Store) CreateAgent(ctx context.Context, organizationID uuid.UUID, input AgentInput) (Agent, error) {
+	capabilitiesJSON, err := encodeCapabilities(input.Capabilities)
+	if err != nil {
+		return Agent{}, err
+	}
 	row := s.pool.QueryRow(ctx,
-		fmt.Sprintf(`INSERT INTO agents (organization_id, name, nickname, role, model, description, configuration, image, init_image, idle_timeout, resources_requests_cpu, resources_requests_memory, resources_limits_cpu, resources_limits_memory)
-		 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
+		fmt.Sprintf(`INSERT INTO agents (organization_id, name, nickname, role, model, description, configuration, image, init_image, idle_timeout, capabilities, resources_requests_cpu, resources_requests_memory, resources_limits_cpu, resources_limits_memory)
+		 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
 		 RETURNING %s`, agentColumns),
 		organizationID,
 		input.Name,
@@ -265,6 +302,7 @@ func (s *Store) CreateAgent(ctx context.Context, organizationID uuid.UUID, input
 		input.Image,
 		input.InitImage,
 		input.IdleTimeout,
+		capabilitiesJSON,
 		input.Resources.RequestsCPU,
 		input.Resources.RequestsMemory,
 		input.Resources.LimitsCPU,
@@ -321,6 +359,13 @@ func (s *Store) UpdateAgent(ctx context.Context, id uuid.UUID, update AgentUpdat
 	if update.IdleTimeout != nil {
 		builder.add("idle_timeout", *update.IdleTimeout)
 	}
+	if update.Capabilities != nil {
+		capabilitiesJSON, err := encodeCapabilities(*update.Capabilities)
+		if err != nil {
+			return Agent{}, err
+		}
+		builder.add("capabilities", capabilitiesJSON)
+	}
 	if update.Resources != nil {
 		builder.add("resources_requests_cpu", update.Resources.RequestsCPU)
 		builder.add("resources_requests_memory", update.Resources.RequestsMemory)

internal/store/types.go

@@ -31,6 +31,7 @@ type Agent struct {
 	Image          string
 	InitImage      string
 	IdleTimeout    *string
+	Capabilities   []string
 	Resources      ComputeResources
 }
 
@@ -118,6 +119,7 @@ type AgentInput struct {
 	Image         string
 	InitImage     string
 	IdleTimeout   *string
+	Capabilities  []string
 	Resources     ComputeResources
 }
 
@@ -131,6 +133,7 @@ type AgentUpdate struct {
 	Image         *string
 	InitImage     *string
 	IdleTimeout   *string
+	Capabilities  *[]string
 	Resources     *ComputeResources
 }
 

migrations/0012_agent_capabilities.sql

@@ -0,0 +1,2 @@
+ALTER TABLE agents
+    ADD COLUMN capabilities JSONB NOT NULL DEFAULT '[]'::jsonb;