Add control tools

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/server/api/mcp/v1/schema.resolvers.go

@@ -1066,3 +1066,99 @@ func (r *Resolver) UpdateAuditTool(ctx context.Context, req *mcp.CallToolRequest
 		Audit: types.NewAudit(audit),
 	}, nil
 }
+
+func (r *Resolver) ListControlsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlsInput) (*mcp.CallToolResult, types.ListControlsOutput, error) {
+	r.MustBeAuthorized(ctx, input.OrganizationID, authz.ActionListControls)
+
+	prb := r.ProboService(ctx, input.OrganizationID)
+
+	pageOrderBy := page.OrderBy[coredata.ControlOrderField]{
+		Field:     coredata.ControlOrderFieldCreatedAt,
+		Direction: page.OrderDirectionDesc,
+	}
+	if input.OrderBy != nil {
+		pageOrderBy = page.OrderBy[coredata.ControlOrderField]{
+			Field:     input.OrderBy.Field,
+			Direction: input.OrderBy.Direction,
+		}
+	}
+
+	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)
+
+	var controlFilter = coredata.NewControlFilter(nil)
+	if input.Filter != nil {
+		controlFilter = coredata.NewControlFilter(input.Filter.Query)
+	}
+
+	page, err := prb.Controls.ListForOrganizationID(ctx, input.OrganizationID, cursor, controlFilter)
+	if err != nil {
+		panic(fmt.Errorf("cannot list organization controls: %w", err))
+	}
+
+	return nil, types.NewListControlsOutput(page), nil
+}
+
+func (r *Resolver) GetControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetControlInput) (*mcp.CallToolResult, types.GetControlOutput, error) {
+	r.MustBeAuthorized(ctx, input.ID, authz.ActionGet)
+
+	prb := r.ProboService(ctx, input.ID)
+
+	control, err := prb.Controls.Get(ctx, input.ID)
+	if err != nil {
+		return nil, types.GetControlOutput{}, fmt.Errorf("failed to get control: %w", err)
+	}
+
+	return nil, types.GetControlOutput{
+		Control: types.NewControl(control),
+	}, nil
+}
+
+func (r *Resolver) AddControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddControlInput) (*mcp.CallToolResult, types.AddControlOutput, error) {
+	r.MustBeAuthorized(ctx, input.FrameworkID, authz.ActionCreateControl)
+
+	svc := r.ProboService(ctx, input.FrameworkID)
+
+	control, err := svc.Controls.Create(
+		ctx,
+		probo.CreateControlRequest{
+			FrameworkID:            input.FrameworkID,
+			Name:                   input.Name,
+			Description:            input.Description,
+			SectionTitle:           input.SectionTitle,
+			Status:                 input.Status,
+			ExclusionJustification: input.ExclusionJustification,
+		},
+	)
+	if err != nil {
+		return nil, types.AddControlOutput{}, fmt.Errorf("failed to create control: %w", err)
+	}
+
+	return nil, types.AddControlOutput{
+		Control: types.NewControl(control),
+	}, nil
+}
+
+func (r *Resolver) UpdateControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateControlInput) (*mcp.CallToolResult, types.UpdateControlOutput, error) {
+	r.MustBeAuthorized(ctx, input.ID, authz.ActionUpdateControl)
+
+	svc := r.ProboService(ctx, input.ID)
+
+	control, err := svc.Controls.Update(
+		ctx,
+		probo.UpdateControlRequest{
+			ID:                     input.ID,
+			Name:                   input.Name,
+			Description:            UnwrapOmittable(input.Description),
+			SectionTitle:           input.SectionTitle,
+			Status:                 input.Status,
+			ExclusionJustification: input.ExclusionJustification,
+		},
+	)
+	if err != nil {
+		return nil, types.UpdateControlOutput{}, fmt.Errorf("failed to update control: %w", err)
+	}
+
+	return nil, types.UpdateControlOutput{
+		Control: types.NewControl(control),
+	}, nil
+}

pkg/server/api/mcp/v1/server/server.go

@@ -53,6 +53,10 @@ type ResolverInterface interface {
 	GetAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetAuditInput) (*mcp.CallToolResult, types.GetAuditOutput, error)
 	AddAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddAuditInput) (*mcp.CallToolResult, types.AddAuditOutput, error)
 	UpdateAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateAuditInput) (*mcp.CallToolResult, types.UpdateAuditOutput, error)
+	ListControlsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlsInput) (*mcp.CallToolResult, types.ListControlsOutput, error)
+	GetControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetControlInput) (*mcp.CallToolResult, types.GetControlOutput, error)
+	AddControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddControlInput) (*mcp.CallToolResult, types.AddControlOutput, error)
+	UpdateControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateControlInput) (*mcp.CallToolResult, types.UpdateControlOutput, error)
 }
 
 // New creates a new MCP server instance with all handlers registered.
@@ -502,4 +506,44 @@ func registerToolHandlers(server *mcp.Server, resolver ResolverInterface) {
 		},
 		resolver.UpdateAuditTool,
 	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "listControls",
+			Description:  "List all controls for the organization or framework",
+			InputSchema:  types.ListControlsToolInputSchema,
+			OutputSchema: types.ListControlsToolOutputSchema,
+		},
+		resolver.ListControlsTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "getControl",
+			Description:  "Get a control by ID",
+			InputSchema:  types.GetControlToolInputSchema,
+			OutputSchema: types.GetControlToolOutputSchema,
+		},
+		resolver.GetControlTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "addControl",
+			Description:  "Add a new control to a framework",
+			InputSchema:  types.AddControlToolInputSchema,
+			OutputSchema: types.AddControlToolOutputSchema,
+		},
+		resolver.AddControlTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "updateControl",
+			Description:  "Update an existing control",
+			InputSchema:  types.UpdateControlToolInputSchema,
+			OutputSchema: types.UpdateControlToolOutputSchema,
+		},
+		resolver.UpdateControlTool,
+	)
 }

pkg/server/api/mcp/v1/specification.yaml

@@ -2446,6 +2446,214 @@ components:
         audit:
           $ref: "#/components/schemas/Audit"
 
+    ControlStatus:
+      type: string
+      enum:
+        - INCLUDED
+        - EXCLUDED
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/coredata.ControlStatus
+
+    ControlOrderField:
+      type: string
+      enum:
+        - CREATED_AT
+        - SECTION_TITLE
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/coredata.ControlOrderField
+
+    ControlOrderBy:
+      type: object
+      required:
+        - field
+        - direction
+      properties:
+        field:
+          $ref: "#/components/schemas/ControlOrderField"
+          description: Control order field
+        direction:
+          $ref: "#/components/schemas/OrderDirection"
+          description: Control order direction
+
+    Control:
+      type: object
+      required:
+        - id
+        - organization_id
+        - framework_id
+        - section_title
+        - name
+        - status
+        - created_at
+        - updated_at
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Control ID
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        framework_id:
+          $ref: "#/components/schemas/GID"
+          description: Framework ID
+        section_title:
+          type: string
+          description: Section title
+        name:
+          type: string
+          description: Control name
+        description:
+          type:
+            - string
+            - "null"
+          description: Control description
+        status:
+          $ref: "#/components/schemas/ControlStatus"
+          description: Control status
+        exclusion_justification:
+          type:
+            - string
+            - "null"
+          description: Exclusion justification
+        created_at:
+          type: string
+          format: date-time
+          description: Creation timestamp
+        updated_at:
+          type: string
+          format: date-time
+          description: Update timestamp
+
+    ListControlsInput:
+      type: object
+      required:
+        - organization_id
+      properties:
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        framework_id:
+          $ref: "#/components/schemas/GID"
+          description: Framework ID
+        order_by:
+          $ref: "#/components/schemas/ControlOrderBy"
+          description: Control order by
+        size:
+          type: integer
+          description: Page size
+        cursor:
+          $ref: "#/components/schemas/CursorKey"
+          description: Page cursor
+        filter:
+          type: object
+          properties:
+            query:
+              type: string
+              description: Search query
+
+    ListControlsOutput:
+      type: object
+      required:
+        - controls
+      properties:
+        next_cursor:
+          $ref: "#/components/schemas/CursorKey"
+          description: Next cursor
+        controls:
+          type: array
+          items:
+            $ref: "#/components/schemas/Control"
+
+    GetControlInput:
+      type: object
+      required:
+        - id
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Control ID
+
+    GetControlOutput:
+      type: object
+      required:
+        - control
+      properties:
+        control:
+          $ref: "#/components/schemas/Control"
+
+    AddControlInput:
+      type: object
+      required:
+        - organization_id
+        - framework_id
+        - section_title
+        - name
+      properties:
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        framework_id:
+          $ref: "#/components/schemas/GID"
+          description: Framework ID
+        section_title:
+          type: string
+          description: Section title
+        name:
+          type: string
+          description: Control name
+        description:
+          type: string
+          description: Control description
+        status:
+          $ref: "#/components/schemas/ControlStatus"
+          description: Control status
+        exclusion_justification:
+          type: string
+          description: Exclusion justification
+
+    AddControlOutput:
+      type: object
+      required:
+        - control
+      properties:
+        control:
+          $ref: "#/components/schemas/Control"
+
+    UpdateControlInput:
+      type: object
+      required:
+        - id
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Control ID
+        section_title:
+          type: string
+          description: Section title
+        name:
+          type: string
+          description: Control name
+        description:
+          type: ["string", "null"]
+          description: Control description
+          go.probo.inc/mcpgen/omittable: true
+        status:
+          anyOf:
+            - $ref: "#/components/schemas/ControlStatus"
+              description: Control status
+            - type: "null"
+              description: No status
+          description: Control status
+        exclusion_justification:
+          type: string
+          description: Exclusion justification
+
+    UpdateControlOutput:
+      type: object
+      required:
+        - control
+      properties:
+        control:
+          $ref: "#/components/schemas/Control"
+
 tools:
   - name: listOrganizations
     description: List all organizations the user has access to
@@ -2748,3 +2956,31 @@ tools:
       $ref: "#/components/schemas/UpdateAuditInput"
     outputSchema:
       $ref: "#/components/schemas/UpdateAuditOutput"
+  - name: listControls
+    description: List all controls for the organization or framework
+    readonly: true
+    inputSchema:
+      $ref: "#/components/schemas/ListControlsInput"
+    outputSchema:
+      $ref: "#/components/schemas/ListControlsOutput"
+  - name: getControl
+    description: Get a control by ID
+    readonly: true
+    inputSchema:
+      $ref: "#/components/schemas/GetControlInput"
+    outputSchema:
+      $ref: "#/components/schemas/GetControlOutput"
+  - name: addControl
+    description: Add a new control to a framework
+    readonly: false
+    inputSchema:
+      $ref: "#/components/schemas/AddControlInput"
+    outputSchema:
+      $ref: "#/components/schemas/AddControlOutput"
+  - name: updateControl
+    description: Update an existing control
+    readonly: false
+    inputSchema:
+      $ref: "#/components/schemas/UpdateControlInput"
+    outputSchema:
+      $ref: "#/components/schemas/UpdateControlOutput"

pkg/server/api/mcp/v1/types/control.go

@@ -0,0 +1,52 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+
+package types
+
+import (
+	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/page"
+)
+
+func NewControl(c *coredata.Control) *Control {
+	return &Control{
+		ID:                     c.ID,
+		OrganizationID:         c.OrganizationID,
+		SectionTitle:           c.SectionTitle,
+		FrameworkID:            c.FrameworkID,
+		Name:                   c.Name,
+		Description:            c.Description,
+		Status:                 c.Status,
+		ExclusionJustification: c.ExclusionJustification,
+		CreatedAt:              c.CreatedAt,
+		UpdatedAt:              c.UpdatedAt,
+	}
+}
+
+func NewListControlsOutput(controlPage *page.Page[*coredata.Control, coredata.ControlOrderField]) ListControlsOutput {
+	controls := make([]*Control, 0, len(controlPage.Data))
+	for _, c := range controlPage.Data {
+		controls = append(controls, NewControl(c))
+	}
+
+	var nextCursor *page.CursorKey
+	if len(controlPage.Data) > 0 {
+		cursorKey := controlPage.Data[len(controlPage.Data)-1].CursorKey(controlPage.Cursor.OrderBy.Field)
+		nextCursor = &cursorKey
+	}
+
+	return ListControlsOutput{
+		NextCursor: nextCursor,
+		Controls:   controls,
+	}
+}

pkg/server/api/mcp/v1/types/types.go

@@ -18,6 +18,8 @@ var (
 	AddAuditToolOutputSchema                   = mcp.MustUnmarshalSchema(`{"type":"object","required":["audit"],"properties":{"audit":{"type":"object","required":["id","organization_id","framework_id","state","trust_center_visibility","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"description":"Audit name"},"organization_id":{"type":"string","format":"string"},"state":{"type":"string","enum":["NOT_STARTED","IN_PROGRESS","COMPLETED","REJECTED","OUTDATED"]},"trust_center_visibility":{"type":"string","enum":["NONE","PRIVATE","PUBLIC"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"},"valid_from":{"description":"Valid from date","format":"date-time"},"valid_until":{"description":"Valid until date","format":"date-time"}}}}}`)
 	AddContinualImprovementToolInputSchema     = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id","reference_id","owner_id"],"properties":{"description":{"type":"string","description":"Description"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"priority":{"description":"Priority","anyOf":[{"type":"string","enum":["LOW","MEDIUM","HIGH"]},{"type":"null","description":"No priority"}]},"reference_id":{"type":"string","description":"Reference ID"},"source":{"type":"string","description":"Source"},"status":{"description":"Status","anyOf":[{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},{"type":"null","description":"No status"}]},"target_date":{"type":"string","description":"Target date","format":"date-time"}}}`)
 	AddContinualImprovementToolOutputSchema    = mcp.MustUnmarshalSchema(`{"type":"object","required":["continual_improvement"],"properties":{"continual_improvement":{"type":"object","required":["id","organization_id","reference_id","owner_id","status","priority","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Description"},"id":{"type":"string","format":"string"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"priority":{"type":"string","enum":["LOW","MEDIUM","HIGH"]},"reference_id":{"type":"string","description":"Reference ID"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"source":{"description":"Source"},"source_id":{"description":"Source ID"},"status":{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},"target_date":{"description":"Target date","format":"date-time"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
+	AddControlToolInputSchema                  = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id","framework_id","section_title","name"],"properties":{"description":{"type":"string","description":"Control description"},"exclusion_justification":{"type":"string","description":"Exclusion justification"},"framework_id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"organization_id":{"type":"string","format":"string"},"section_title":{"type":"string","description":"Section title"},"status":{"type":"string","enum":["INCLUDED","EXCLUDED"]}}}`)
+	AddControlToolOutputSchema                 = mcp.MustUnmarshalSchema(`{"type":"object","required":["control"],"properties":{"control":{"type":"object","required":["id","organization_id","framework_id","section_title","name","status","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Control description"},"exclusion_justification":{"description":"Exclusion justification"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"organization_id":{"type":"string","format":"string"},"section_title":{"type":"string","description":"Section title"},"status":{"type":"string","enum":["INCLUDED","EXCLUDED"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	AddDatumToolInputSchema                    = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id","name","data_classification","owner_id"],"properties":{"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"name":{"type":"string","description":"Datum name"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"vendor_ids":{"type":"array","description":"Vendor IDs","items":{"type":"string","format":"string"}}}}`)
 	AddDatumToolOutputSchema                   = mcp.MustUnmarshalSchema(`{"type":"object","required":["datum"],"properties":{"datum":{"type":"object","required":["id","organization_id","name","data_classification","owner_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Datum name"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	AddFrameworkToolInputSchema                = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id","name"],"properties":{"description":{"type":"string","description":"Framework description"},"name":{"type":"string","description":"Framework name"},"organization_id":{"type":"string","format":"string"}}}`)
@@ -40,6 +42,8 @@ var (
 	GetAuditToolOutputSchema                   = mcp.MustUnmarshalSchema(`{"type":"object","required":["audit"],"properties":{"audit":{"type":"object","required":["id","organization_id","framework_id","state","trust_center_visibility","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"description":"Audit name"},"organization_id":{"type":"string","format":"string"},"state":{"type":"string","enum":["NOT_STARTED","IN_PROGRESS","COMPLETED","REJECTED","OUTDATED"]},"trust_center_visibility":{"type":"string","enum":["NONE","PRIVATE","PUBLIC"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"},"valid_from":{"description":"Valid from date","format":"date-time"},"valid_until":{"description":"Valid until date","format":"date-time"}}}}}`)
 	GetContinualImprovementToolInputSchema     = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"id":{"type":"string","format":"string"}}}`)
 	GetContinualImprovementToolOutputSchema    = mcp.MustUnmarshalSchema(`{"type":"object","required":["continual_improvement"],"properties":{"continual_improvement":{"type":"object","required":["id","organization_id","reference_id","owner_id","status","priority","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Description"},"id":{"type":"string","format":"string"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"priority":{"type":"string","enum":["LOW","MEDIUM","HIGH"]},"reference_id":{"type":"string","description":"Reference ID"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"source":{"description":"Source"},"source_id":{"description":"Source ID"},"status":{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},"target_date":{"description":"Target date","format":"date-time"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
+	GetControlToolInputSchema                  = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"id":{"type":"string","format":"string"}}}`)
+	GetControlToolOutputSchema                 = mcp.MustUnmarshalSchema(`{"type":"object","required":["control"],"properties":{"control":{"type":"object","required":["id","organization_id","framework_id","section_title","name","status","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Control description"},"exclusion_justification":{"description":"Exclusion justification"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"organization_id":{"type":"string","format":"string"},"section_title":{"type":"string","description":"Section title"},"status":{"type":"string","enum":["INCLUDED","EXCLUDED"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	GetDatumToolInputSchema                    = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"id":{"type":"string","format":"string"}}}`)
 	GetDatumToolOutputSchema                   = mcp.MustUnmarshalSchema(`{"type":"object","required":["datum"],"properties":{"datum":{"type":"object","required":["id","organization_id","name","data_classification","owner_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Datum name"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	GetFrameworkToolInputSchema                = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"id":{"type":"string","format":"string"}}}`)
@@ -60,6 +64,8 @@ var (
 	ListAuditsToolOutputSchema                 = mcp.MustUnmarshalSchema(`{"type":"object","required":["audits"],"properties":{"audits":{"type":"array","items":{"type":"object","required":["id","organization_id","framework_id","state","trust_center_visibility","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"description":"Audit name"},"organization_id":{"type":"string","format":"string"},"state":{"type":"string","enum":["NOT_STARTED","IN_PROGRESS","COMPLETED","REJECTED","OUTDATED"]},"trust_center_visibility":{"type":"string","enum":["NONE","PRIVATE","PUBLIC"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"},"valid_from":{"description":"Valid from date","format":"date-time"},"valid_until":{"description":"Valid until date","format":"date-time"}}}},"next_cursor":{"type":"string","format":"string"}}}`)
 	ListContinualImprovementsToolInputSchema   = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id"],"properties":{"cursor":{"type":"string","format":"string"},"filter":{"type":"object","properties":{"snapshot_id":{"type":"string","format":"string"}}},"order_by":{"type":"object","required":["field","direction"],"properties":{"direction":{"type":"string","enum":["ASC","DESC"]},"field":{"type":"string","enum":["CREATED_AT","REFERENCE_ID","TARGET_DATE","STATUS","PRIORITY"]}}},"organization_id":{"type":"string","format":"string"},"size":{"type":"integer","description":"Page size"}}}`)
 	ListContinualImprovementsToolOutputSchema  = mcp.MustUnmarshalSchema(`{"type":"object","required":["continual_improvements"],"properties":{"continual_improvements":{"type":"array","items":{"type":"object","required":["id","organization_id","reference_id","owner_id","status","priority","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Description"},"id":{"type":"string","format":"string"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"priority":{"type":"string","enum":["LOW","MEDIUM","HIGH"]},"reference_id":{"type":"string","description":"Reference ID"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"source":{"description":"Source"},"source_id":{"description":"Source ID"},"status":{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},"target_date":{"description":"Target date","format":"date-time"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}},"next_cursor":{"type":"string","format":"string"}}}`)
+	ListControlsToolInputSchema                = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id"],"properties":{"cursor":{"type":"string","format":"string"},"filter":{"type":"object","properties":{"query":{"type":"string","description":"Search query"}}},"framework_id":{"type":"string","format":"string"},"order_by":{"type":"object","required":["field","direction"],"properties":{"direction":{"type":"string","enum":["ASC","DESC"]},"field":{"type":"string","enum":["CREATED_AT","SECTION_TITLE"]}}},"organization_id":{"type":"string","format":"string"},"size":{"type":"integer","description":"Page size"}}}`)
+	ListControlsToolOutputSchema               = mcp.MustUnmarshalSchema(`{"type":"object","required":["controls"],"properties":{"controls":{"type":"array","items":{"type":"object","required":["id","organization_id","framework_id","section_title","name","status","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Control description"},"exclusion_justification":{"description":"Exclusion justification"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"organization_id":{"type":"string","format":"string"},"section_title":{"type":"string","description":"Section title"},"status":{"type":"string","enum":["INCLUDED","EXCLUDED"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}},"next_cursor":{"type":"string","format":"string"}}}`)
 	ListDataToolInputSchema                    = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id"],"properties":{"cursor":{"type":"string","format":"string"},"filter":{"type":"object","properties":{"snapshot_id":{"type":"string","format":"string"}}},"order_by":{"type":"object","required":["field","direction"],"properties":{"direction":{"type":"string","enum":["ASC","DESC"]},"field":{"type":"string","enum":["CREATED_AT","NAME","DATA_CLASSIFICATION"]}}},"organization_id":{"type":"string","format":"string"},"size":{"type":"integer","description":"Page size"}}}`)
 	ListDataToolOutputSchema                   = mcp.MustUnmarshalSchema(`{"type":"object","required":["data"],"properties":{"data":{"type":"array","items":{"type":"object","required":["id","organization_id","name","data_classification","owner_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Datum name"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}},"next_cursor":{"type":"string","format":"string"}}}`)
 	ListFrameworksToolInputSchema              = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id"],"properties":{"cursor":{"type":"string","format":"string"},"order_by":{"type":"object","required":["field","direction"],"properties":{"direction":{"type":"string","enum":["ASC","DESC"]},"field":{"type":"string","enum":["CREATED_AT"]}}},"organization_id":{"type":"string","format":"string"},"size":{"type":"integer","description":"Page size"}}}`)
@@ -84,6 +90,8 @@ var (
 	UpdateAuditToolOutputSchema                = mcp.MustUnmarshalSchema(`{"type":"object","required":["audit"],"properties":{"audit":{"type":"object","required":["id","organization_id","framework_id","state","trust_center_visibility","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"description":"Audit name"},"organization_id":{"type":"string","format":"string"},"state":{"type":"string","enum":["NOT_STARTED","IN_PROGRESS","COMPLETED","REJECTED","OUTDATED"]},"trust_center_visibility":{"type":"string","enum":["NONE","PRIVATE","PUBLIC"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"},"valid_from":{"description":"Valid from date","format":"date-time"},"valid_until":{"description":"Valid until date","format":"date-time"}}}}}`)
 	UpdateContinualImprovementToolInputSchema  = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"description":{"description":"Description"},"id":{"type":"string","format":"string"},"owner_id":{"description":"Owner ID","anyOf":[{"type":"string","format":"string"},{"type":"null"}]},"priority":{"description":"Priority","anyOf":[{"type":"string","enum":["LOW","MEDIUM","HIGH"]},{"type":"null","description":"No priority"}]},"reference_id":{"type":"string","description":"Reference ID"},"source":{"description":"Source"},"status":{"description":"Status","anyOf":[{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},{"type":"null","description":"No status"}]},"target_date":{"description":"Target date","format":"date-time"}}}`)
 	UpdateContinualImprovementToolOutputSchema = mcp.MustUnmarshalSchema(`{"type":"object","required":["continual_improvement"],"properties":{"continual_improvement":{"type":"object","required":["id","organization_id","reference_id","owner_id","status","priority","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Description"},"id":{"type":"string","format":"string"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"priority":{"type":"string","enum":["LOW","MEDIUM","HIGH"]},"reference_id":{"type":"string","description":"Reference ID"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"source":{"description":"Source"},"source_id":{"description":"Source ID"},"status":{"type":"string","enum":["OPEN","IN_PROGRESS","CLOSED"]},"target_date":{"description":"Target date","format":"date-time"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
+	UpdateControlToolInputSchema               = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"description":{"description":"Control description"},"exclusion_justification":{"type":"string","description":"Exclusion justification"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"section_title":{"type":"string","description":"Section title"},"status":{"description":"Control status","anyOf":[{"type":"string","enum":["INCLUDED","EXCLUDED"]},{"type":"null","description":"No status"}]}}}`)
+	UpdateControlToolOutputSchema              = mcp.MustUnmarshalSchema(`{"type":"object","required":["control"],"properties":{"control":{"type":"object","required":["id","organization_id","framework_id","section_title","name","status","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Control description"},"exclusion_justification":{"description":"Exclusion justification"},"framework_id":{"type":"string","format":"string"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Control name"},"organization_id":{"type":"string","format":"string"},"section_title":{"type":"string","description":"Section title"},"status":{"type":"string","enum":["INCLUDED","EXCLUDED"]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	UpdateDatumToolInputSchema                 = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Datum name"},"owner_id":{"description":"Owner ID","anyOf":[{"type":"string","format":"string"},{"type":"null"}]},"vendor_ids":{"type":"array","description":"Vendor IDs","items":{"type":"string","format":"string"}}}}`)
 	UpdateDatumToolOutputSchema                = mcp.MustUnmarshalSchema(`{"type":"object","required":["datum"],"properties":{"datum":{"type":"object","required":["id","organization_id","name","data_classification","owner_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"data_classification":{"type":"string","enum":["PUBLIC","INTERNAL","CONFIDENTIAL","SECRET"]},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Datum name"},"organization_id":{"type":"string","format":"string"},"owner_id":{"type":"string","format":"string"},"snapshot_id":{"description":"Snapshot ID","anyOf":[{"type":"string","format":"string"},{"type":"null","description":"No snapshot"}]},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
 	UpdateFrameworkToolInputSchema             = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"description":{"description":"Framework description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Framework name"}}}`)
@@ -169,6 +177,29 @@ type AddContinualImprovementOutput struct {
 	ContinualImprovement *ContinualImprovement `json:"continual_improvement"`
 }
 
+// AddControlInput represents the schema
+type AddControlInput struct {
+	// Control description
+	Description *string `json:"description,omitempty"`
+	// Exclusion justification
+	ExclusionJustification *string `json:"exclusion_justification,omitempty"`
+	// Framework ID
+	FrameworkID gid.GID `json:"framework_id"`
+	// Control name
+	Name string `json:"name"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Section title
+	SectionTitle string `json:"section_title"`
+	// Control status
+	Status *coredata.ControlStatus `json:"status,omitempty"`
+}
+
+// AddControlOutput represents the schema
+type AddControlOutput struct {
+	Control *Control `json:"control"`
+}
+
 // AddDatumInput represents the schema
 type AddDatumInput struct {
 	// Data classification
@@ -457,6 +488,38 @@ type ContinualImprovementOrderBy struct {
 	Field coredata.ContinualImprovementOrderField `json:"field"`
 }
 
+// Control represents the schema
+type Control struct {
+	// Creation timestamp
+	CreatedAt time.Time `json:"created_at"`
+	// Control description
+	Description *string `json:"description,omitempty"`
+	// Exclusion justification
+	ExclusionJustification *string `json:"exclusion_justification,omitempty"`
+	// Framework ID
+	FrameworkID gid.GID `json:"framework_id"`
+	// Control ID
+	ID gid.GID `json:"id"`
+	// Control name
+	Name string `json:"name"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Section title
+	SectionTitle string `json:"section_title"`
+	// Control status
+	Status coredata.ControlStatus `json:"status"`
+	// Update timestamp
+	UpdatedAt time.Time `json:"updated_at"`
+}
+
+// ControlOrderBy represents the schema
+type ControlOrderBy struct {
+	// Control order direction
+	Direction page.OrderDirection `json:"direction"`
+	// Control order field
+	Field coredata.ControlOrderField `json:"field"`
+}
+
 // Datum represents the schema
 type Datum struct {
 	// Creation timestamp
@@ -542,6 +605,17 @@ type GetContinualImprovementOutput struct {
 	ContinualImprovement *ContinualImprovement `json:"continual_improvement"`
 }
 
+// GetControlInput represents the schema
+type GetControlInput struct {
+	// Control ID
+	ID gid.GID `json:"id"`
+}
+
+// GetControlOutput represents the schema
+type GetControlOutput struct {
+	Control *Control `json:"control"`
+}
+
 // GetDatumInput represents the schema
 type GetDatumInput struct {
 	// Datum ID
@@ -678,6 +752,28 @@ type ListContinualImprovementsOutput struct {
 	NextCursor *page.CursorKey `json:"next_cursor,omitempty"`
 }
 
+// ListControlsInput represents the schema
+type ListControlsInput struct {
+	// Page cursor
+	Cursor *page.CursorKey          `json:"cursor,omitempty"`
+	Filter *ListControlsInputFilter `json:"filter,omitempty"`
+	// Framework ID
+	FrameworkID *gid.GID `json:"framework_id,omitempty"`
+	// Control order by
+	OrderBy *ControlOrderBy `json:"order_by,omitempty"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Page size
+	Size *int `json:"size,omitempty"`
+}
+
+// ListControlsOutput represents the schema
+type ListControlsOutput struct {
+	Controls []*Control `json:"controls"`
+	// Next cursor
+	NextCursor *page.CursorKey `json:"next_cursor,omitempty"`
+}
+
 // ListDataInput represents the schema
 type ListDataInput struct {
 	// Page cursor
@@ -1120,6 +1216,27 @@ type UpdateContinualImprovementOutput struct {
 	ContinualImprovement *ContinualImprovement `json:"continual_improvement"`
 }
 
+// UpdateControlInput represents the schema
+type UpdateControlInput struct {
+	// Control description
+	Description mcp.Omittable[*string] `json:"description,omitempty"`
+	// Exclusion justification
+	ExclusionJustification *string `json:"exclusion_justification,omitempty"`
+	// Control ID
+	ID gid.GID `json:"id"`
+	// Control name
+	Name *string `json:"name,omitempty"`
+	// Section title
+	SectionTitle *string `json:"section_title,omitempty"`
+	// Control status
+	Status *coredata.ControlStatus `json:"status,omitempty"`
+}
+
+// UpdateControlOutput represents the schema
+type UpdateControlOutput struct {
+	Control *Control `json:"control"`
+}
+
 // UpdateDatumInput represents the schema
 type UpdateDatumInput struct {
 	// Data classification
@@ -1315,6 +1432,12 @@ type ListContinualImprovementsInputFilter struct {
 	SnapshotID *gid.GID `json:"snapshot_id,omitempty"`
 }
 
+// ListControlsInputFilter represents the schema
+type ListControlsInputFilter struct {
+	// Search query
+	Query *string `json:"query,omitempty"`
+}
+
 // ListDataInputFilter represents the schema
 type ListDataInputFilter struct {
 	// Snapshot ID