kiln_ai.datamodel

 7class BasePrompt(BaseModel):
 8    """
 9    A prompt for a task. This is the basic data storage format which can be used throughout a project.
10
11    The "Prompt" model name is reserved for the custom prompts parented by a task.
12    """
13
14    name: FilenameString = Field(description="The name of the prompt.")
15    description: str | None = Field(
16        default=None,
17        description="A more detailed description of the prompt.",
18    )
19    generator_id: str | None = Field(
20        default=None,
21        description="The id of the generator that created this prompt.",
22    )
23    prompt: str = Field(
24        description="The prompt for the task.",
25        min_length=1,
26    )
27    chain_of_thought_instructions: str | None = Field(
28        default=None,
29        description="Instructions for the model 'thinking' about the requirement prior to answering. Used for chain of thought style prompting. COT will not be used unless this is provided.",
30    )

11class DataGuide(KilnParentedModel):
12    """Persistent input data guide for synthetic data generation, stored as a
13    child of a Task.
14
15    The guide describes what realistic *inputs* to this task look like — input
16    shape, format, distribution, and the kinds of values inputs contain. It is
17    consumed at the topic and input generation stages of synthetic data
18    generation, never at the output stage. Output behavior is the job of the
19    task's system prompt, not this guide.
20
21    The guide is a single markdown body. Two canonical shapes are supported
22    depending on origin:
23
24    - **Manual flow** (user-authored examples): leads with `# Reference Inputs`
25      (user-owned `## Example N` blocks), followed by `# Semantics`, `# Style`,
26      `# Presentation Defaults`.
27    - **Kiln Pro / Copilot flow** (analyze pipeline): only `# Semantics`,
28      `# Style`, `# Presentation Defaults` — the analyze prompt derives rules
29      from input documents rather than quoting them, matching Mike's
30      GENERATE_CORPUS_GUIDELINES vocabulary.
31
32    The metaprompter treats the whole body as one editable artifact and returns
33    a refined version on each refine pass; refine auto-detects which shape it
34    is working on by checking for a `# Reference Inputs` heading.
35    """
36
37    guide: str = Field(
38        default="",
39        description="Markdown body of the input data guide. Manual-flow guides start with `# Reference Inputs`; Kiln Pro / Copilot-flow guides have only `# Semantics`, `# Style`, `# Presentation Defaults`.",
40    )
41
42    source: DataGuideSource = Field(
43        default="manual",
44        description="Which flow created this guide. Refine + verify pipelines branch on this to choose the right metaprompter (manual = user-curated examples + edits all sections; kiln_pro = LLM-derived, refine is feedback-only surgical edits).",
45    )

195class DataSource(BaseModel):
196    """
197    Represents the origin of data, either human, synthetic, file import, or tool call, with associated properties.
198
199    Properties vary based on the source type - for synthetic/tool_call sources this includes
200    model information, for human sources this includes creator information, for file imports
201    this includes file information.
202    """
203
204    type: DataSourceType = Field(description="The type of data source.")
205    properties: Dict[str, str | int | float] = Field(
206        default={},
207        description="Properties describing the data source. For synthetic things like model. For human: the human's name. For file_import: file information.",
208    )
209    run_config_id: Optional[str] = Field(
210        default=None,
211        description="The ID of the saved TaskRunConfig used to produce this data, if any. Only present when the run was initiated from a saved TaskRunConfig (e.g. via the saved-config dropdown or a tool call); runs configured ad-hoc from the run page leave this unset. Not validated against the file system: a TaskRunConfig may be deleted without invalidating historical runs that reference it.",
212    )
213    run_config: Optional[RunConfigProperties] = Field(
214        default=None,
215        description="The run config used to generate the data, if generated by a running a model in Kiln (only true for type=synthetic).",
216    )
217
218    @model_validator(mode="after")
219    def normalize_empty_run_config_id(self) -> Self:
220        # Some callers (e.g. tool wrappers reading from external properties)
221        # default missing IDs to "". Coerce to None so this field is either a
222        # real ID or unset — never an ambiguous empty string on disk.
223        if self.run_config_id == "":
224            self.run_config_id = None
225        return self
226
227    _data_source_properties = [
228        DataSourceProperty(
229            name="created_by",
230            type=str,
231            required_for=[DataSourceType.human],
232            not_allowed_for=[
233                DataSourceType.synthetic,
234                DataSourceType.file_import,
235                DataSourceType.tool_call,
236            ],
237        ),
238        DataSourceProperty(
239            name="model_name",
240            type=str,
241            required_for=[DataSourceType.synthetic],
242            not_allowed_for=[
243                DataSourceType.human,
244                DataSourceType.file_import,
245                DataSourceType.tool_call,
246            ],
247        ),
248        DataSourceProperty(
249            name="model_provider",
250            type=str,
251            required_for=[DataSourceType.synthetic],
252            not_allowed_for=[
253                DataSourceType.human,
254                DataSourceType.file_import,
255                DataSourceType.tool_call,
256            ],
257        ),
258        DataSourceProperty(
259            name="adapter_name",
260            type=str,
261            required_for=[DataSourceType.synthetic],
262            not_allowed_for=[
263                DataSourceType.human,
264                DataSourceType.file_import,
265                DataSourceType.tool_call,
266            ],
267        ),
268        DataSourceProperty(
269            # Legacy field -- allow loading from old runs, but we shouldn't be setting it.
270            name="prompt_builder_name",
271            type=str,
272            not_allowed_for=[
273                DataSourceType.human,
274                DataSourceType.file_import,
275                DataSourceType.tool_call,
276            ],
277        ),
278        DataSourceProperty(
279            # The PromptId of the prompt. Can be a saved prompt, fine-tune, generator name, etc. See PromptId type for more details.
280            name="prompt_id",
281            type=str,
282            not_allowed_for=[
283                DataSourceType.human,
284                DataSourceType.file_import,
285                DataSourceType.tool_call,
286            ],
287        ),
288        DataSourceProperty(
289            name="file_name",
290            type=str,
291            required_for=[DataSourceType.file_import],
292            not_allowed_for=[
293                DataSourceType.human,
294                DataSourceType.synthetic,
295                DataSourceType.tool_call,
296            ],
297        ),
298    ]
299
300    @model_validator(mode="after")
301    def validate_type(self) -> "DataSource":
302        if self.type not in DataSourceType:
303            raise ValueError(f"Invalid data source type: {self.type}")
304        return self
305
306    @model_validator(mode="after")
307    def validate_properties(self) -> "DataSource":
308        for prop in self._data_source_properties:
309            # Check the property type is correct
310            if prop.name in self.properties:
311                if not isinstance(self.properties[prop.name], prop.type):
312                    raise ValueError(
313                        f"'{prop.name}' must be of type {prop.type.__name__} for {self.type} data source"
314                    )
315            # Check the property is required for the data source type
316            if self.type in prop.required_for:
317                if prop.name not in self.properties:
318                    raise ValueError(
319                        f"'{prop.name}' is required for {self.type} data source"
320                    )
321            # Check the property is not allowed for the data source type
322            elif self.type in prop.not_allowed_for and prop.name in self.properties:
323                raise ValueError(
324                    f"'{prop.name}' is not allowed for {self.type} data source"
325                )
326        return self
327
328    @model_validator(mode="after")
329    def validate_no_empty_properties(self) -> Self:
330        for prop, value in self.properties.items():
331            if isinstance(value, str) and value == "":
332                raise ValueError(
333                    f"Property '{prop}' must be a non-empty string for {self.type} data source"
334                )
335        return self

181class DataSourceProperty(BaseModel):
182    """
183    Defines a property that can be associated with a data source.
184
185    Includes validation rules for when properties are required or not allowed
186    based on the data source type.
187    """
188
189    name: str
190    type: Type[Union[str, int, float]]
191    required_for: List[DataSourceType] = []
192    not_allowed_for: List[DataSourceType] = []

167class DataSourceType(str, Enum):
168    """
169    The source type of a piece of data.
170
171    Human: a human created the data
172    Synthetic: a model created the data
173    """
174
175    human = "human"
176    synthetic = "synthetic"
177    file_import = "file_import"
178    tool_call = "tool_call"

 83class DatasetSplit(KilnParentedModel):
 84    """
 85    A collection of task runs, with optional splits (train, test, validation).
 86
 87    Used to freeze a dataset into train/test/validation splits for repeatable fine-tuning or other tasks.
 88
 89    Maintains a list of IDs for each split, to avoid data duplication.
 90    """
 91
 92    name: FilenameString = Field(description="The name of the dataset split.")
 93    description: str | None = Field(
 94        default=None,
 95        description="A description of the dataset for you and your team. Not used in training.",
 96    )
 97    splits: list[DatasetSplitDefinition] = Field(
 98        default_factory=list,
 99        description="The splits in the dataset.",
100    )
101    split_contents: dict[str, list[str]] = Field(
102        description="The contents of each split in the dataset. The key is the split name, and the value is a list of task run IDs.",
103    )
104    filter: DatasetFilterId | None = Field(
105        default=None,
106        description="The filter used to build the dataset.",
107    )
108
109    @model_validator(mode="after")
110    def validate_split_percentages(self) -> "DatasetSplit":
111        total = sum(split.percentage for split in self.splits)
112        if not math.isclose(total, 1.0, rel_tol=1e-9):
113            raise ValueError(f"The sum of split percentages must be 1.0 (got {total})")
114        return self
115
116    @classmethod
117    def from_task(
118        cls,
119        name: str,
120        task: "Task",
121        splits: list[DatasetSplitDefinition],
122        filter_id: DatasetFilterId = "all",
123        description: str | None = None,
124    ):
125        """
126        Build a dataset split from a task.
127        """
128        filter = dataset_filter_from_id(filter_id)
129        split_contents = cls.build_split_contents(task, splits, filter)
130        return cls(
131            parent=task,
132            name=name,
133            description=description,
134            splits=splits,
135            split_contents=split_contents,
136            filter=filter_id,
137        )
138
139    @classmethod
140    def build_split_contents(
141        cls,
142        task: "Task",
143        splits: list[DatasetSplitDefinition],
144        filter: DatasetFilter,
145    ) -> dict[str, list[str]]:
146        valid_ids = []
147        for task_run in task.runs():
148            if filter(task_run):
149                valid_ids.append(task_run.id)
150
151        # Shuffle and split by split percentage
152        random.shuffle(valid_ids)
153        split_contents = {}
154        start_idx = 0
155        remaining_items = len(valid_ids)
156
157        # Handle all splits except the last one
158        for split in splits[:-1]:
159            split_size = round(len(valid_ids) * split.percentage)
160            split_contents[split.name] = valid_ids[start_idx : start_idx + split_size]
161            start_idx += split_size
162            remaining_items -= split_size
163
164        # Last split gets all remaining items (for rounding)
165        if splits:
166            split_contents[splits[-1].name] = valid_ids[start_idx:]
167
168        return split_contents
169
170    def parent_task(self) -> "Task | None":
171        # inline import to avoid circular import
172        from kiln_ai.datamodel import Task
173
174        if not isinstance(self.parent, Task):
175            return None
176        return self.parent
177
178    def missing_count(self) -> int:
179        """
180        Returns:
181            int: the number of task runs that have an ID persisted in this dataset split, but no longer exist in the dataset
182        """
183        parent = self.parent_task()
184        if parent is None:
185            raise ValueError("DatasetSplit has no parent task")
186
187        runs = parent.runs(readonly=True)
188        all_ids = set(run.id for run in runs)
189        all_ids_in_splits = set()
190        for ids in self.split_contents.values():
191            all_ids_in_splits.update(ids)
192        missing = all_ids_in_splits - all_ids
193        return len(missing)
194
195    def _get_runs(self) -> list[TaskRun]:
196        """
197        Get all task runs referenced in this dataset split.
198
199        Returns:
200            list[TaskRun]: list of task runs in this dataset split
201        """
202        parent = self.parent_task()
203        if parent is None:
204            return []
205
206        runs = []
207        all_run_ids = set()
208        for run_ids in self.split_contents.values():
209            all_run_ids.update(run_ids)
210
211        # Find all runs by their IDs
212        for task_run in parent.runs(readonly=True):
213            if task_run.id in all_run_ids:
214                runs.append(task_run)
215
216        return runs
217
218    @staticmethod
219    def compute_tool_info(runs: list[TaskRun]) -> DatasetToolInfo:
220        """
221        Compute tool info from a list of task runs.
222
223        Args:
224            runs: list of task runs to analyze
225
226        Returns:
227            DatasetToolInfo: information about tools used across the task runs
228        """
229
230        has_tool_mismatch = False
231        tools: set[str] | None = None
232
233        for run in runs:
234            # Extract tools from run config, treating missing source/run_config/tools_config as empty tools
235            run_tools: set[str] = set()
236            source = run.output.source if run.output else None
237            if source is not None and isinstance(
238                source.run_config, KilnAgentRunConfigProperties
239            ):
240                tools_config = source.run_config.tools_config
241                if tools_config is not None:
242                    run_tools = set(tools_config.tools)
243
244            # First run establishes the expected tool set (including empty)
245            if tools is None:
246                tools = run_tools
247            elif run_tools != tools:
248                # Mismatch found
249                has_tool_mismatch = True
250                tools = None
251                break
252
253        # If no valid runs were processed, return empty tools
254        if tools is None:
255            if not has_tool_mismatch:
256                tools = set()
257
258        return DatasetToolInfo(
259            has_tool_mismatch=has_tool_mismatch,
260            tools=None if tools is None else sorted(tools),
261        )
262
263    def tool_info(self) -> DatasetToolInfo:
264        """
265        Helper method to compute tool info for the dataset split. Iterate through all runs in the dataset split and check the tools used in each run config.
266
267        Returns:
268            DatasetToolInfo: information about tools used across task runs in this dataset split
269        """
270        runs = self._get_runs()
271        tool_info = self.compute_tool_info(runs)
272        return tool_info

39class DatasetSplitDefinition(BaseModel):
40    """
41    A definition of a split in a dataset.
42
43    Example: name="train", description="The training set", percentage=0.8 (80% of the dataset)
44    """
45
46    name: FilenameString = Field(
47        description="The name of the dataset split definition."
48    )
49    description: str | None = Field(
50        default=None,
51        description="A description of the dataset for you and your team. Not used in training.",
52    )
53    percentage: float = Field(
54        ge=0.0,
55        le=1.0,
56        description="The percentage of the dataset that this split represents (between 0 and 1).",
57    )

 64class ExternalToolServer(KilnParentedModel):
 65    """
 66    Configuration for communicating with a external MCP (Model Context Protocol) Server for LLM tool calls. External tool servers can be remote or local.
 67
 68    This model stores the necessary configuration to connect to and authenticate with
 69    external MCP servers that provide tools for LLM interactions.
 70    """
 71
 72    name: FilenameString = Field(description="The name of the external tool.")
 73    type: ToolServerType = Field(
 74        description="The type of external tool server. Remote tools are hosted on a remote server",
 75    )
 76    description: str | None = Field(
 77        default=None,
 78        description="A description of the external tool for you and your team. Will not be used in prompts/training/validation.",
 79    )
 80
 81    properties: (
 82        LocalServerProperties | RemoteServerProperties | KilnTaskServerProperties
 83    ) = Field(
 84        description="Configuration properties specific to the tool type.",
 85    )
 86
 87    # Private variable to store unsaved secrets
 88    _unsaved_secrets: dict[str, str] = PrivateAttr(default_factory=dict)
 89
 90    def model_post_init(self, __context: Any) -> None:
 91        # Process secrets after initialization (pydantic v2 hook)
 92        self._process_secrets_from_properties()
 93
 94    def _process_secrets_from_properties(self) -> None:
 95        """
 96        Extract secrets from properties and move them to _unsaved_secrets.
 97        This removes secrets from the properties dict so they aren't saved to file.
 98        Clears existing _unsaved_secrets first to handle property updates correctly.
 99        """
100        # Clear existing unsaved secrets since we're reprocessing
101        self._unsaved_secrets.clear()
102
103        secret_keys = self.get_secret_keys()
104
105        if not secret_keys:
106            return
107
108        # Extract secret values from properties based on server type
109        match self.type:
110            case ToolServerType.remote_mcp:
111                headers = self.properties.get("headers", {})
112                for key_name in secret_keys:
113                    if key_name in headers:
114                        self._unsaved_secrets[key_name] = headers[key_name]
115                        # Remove from headers immediately so they are not saved to file
116                        del headers[key_name]
117
118            case ToolServerType.local_mcp:
119                env_vars = self.properties.get("env_vars", {})
120                for key_name in secret_keys:
121                    if key_name in env_vars:
122                        self._unsaved_secrets[key_name] = env_vars[key_name]
123                        # Remove from env_vars immediately so they are not saved to file
124                        del env_vars[key_name]
125
126            case ToolServerType.kiln_task:
127                pass
128
129            case _:
130                raise_exhaustive_enum_error(self.type)
131
132    def __setattr__(self, name: str, value: Any) -> None:
133        """
134        Override __setattr__ to process secrets whenever properties are updated.
135        """
136        super().__setattr__(name, value)
137
138        # Process secrets whenever properties are updated
139        if name == "properties":
140            self._process_secrets_from_properties()
141
142    # Validation Helpers
143
144    @classmethod
145    def check_server_url(cls, server_url: str) -> None:
146        """Validate Server URL"""
147        if not isinstance(server_url, str):
148            raise ValueError("Server URL must be a string")
149
150        # Check for leading whitespace in URL
151        if server_url != server_url.lstrip():
152            raise ValueError("Server URL must not have leading whitespace")
153
154        parsed_url = urlparse(server_url)
155        if not parsed_url.netloc:
156            raise ValueError("Server URL is not a valid URL")
157        if parsed_url.scheme not in ["http", "https"]:
158            raise ValueError("Server URL must start with http:// or https://")
159
160    @classmethod
161    def check_headers(cls, headers: dict) -> None:
162        """Validate Headers"""
163        if not isinstance(headers, dict):
164            raise ValueError("headers must be a dictionary")
165
166        for key, value in headers.items():
167            if not key:
168                raise ValueError("Header name is required")
169            if not value:
170                raise ValueError("Header value is required")
171
172            # Reject invalid header names and CR/LF in names/values
173            token_re = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$")
174            if not token_re.match(key):
175                raise ValueError(f'Invalid header name: "{key}"')
176            if re.search(r"\r|\n", key) or re.search(r"\r|\n", value):
177                raise ValueError(
178                    "Header names/values must not contain invalid characters"
179                )
180
181    @classmethod
182    def check_secret_keys(
183        cls, secret_keys: list, key_type: str, tool_type: str
184    ) -> None:
185        """Validate Secret Keys (generic method for both header and env var keys)"""
186        if not isinstance(secret_keys, list):
187            raise ValueError(
188                f"{key_type} must be a list for external tools of type '{tool_type}'"
189            )
190        if not all(isinstance(k, str) for k in secret_keys):
191            raise ValueError(f"{key_type} must contain only strings")
192        if not all(key for key in secret_keys):
193            raise ValueError("Secret key is required")
194
195    @classmethod
196    def check_env_vars(cls, env_vars: dict) -> None:
197        """Validate Environment Variables"""
198        if not isinstance(env_vars, dict):
199            raise ValueError("environment variables must be a dictionary")
200
201        # Validate env_vars keys are in the correct format for Environment Variables
202        # According to POSIX specification, environment variable names must:
203        # - Start with a letter (a-z, A-Z) or underscore (_)
204        # - Contain only ASCII letters, digits, and underscores
205        for key, _ in env_vars.items():
206            if not key or not (
207                key[0].isascii() and (key[0].isalpha() or key[0] == "_")
208            ):
209                raise ValueError(
210                    f"Invalid environment variable key: {key}. Must start with a letter or underscore."
211                )
212
213            if not all(c.isascii() and (c.isalnum() or c == "_") for c in key):
214                raise ValueError(
215                    f"Invalid environment variable key: {key}. Can only contain letters, digits, and underscores."
216                )
217
218    @classmethod
219    def type_from_data(cls, data: dict) -> ToolServerType:
220        """Get the tool server type from the data for the the validators"""
221        raw_type = data.get("type")
222        if raw_type is None:
223            raise ValueError("type is required")
224        try:
225            return ToolServerType(raw_type)
226        except ValueError:
227            valid_types = ", ".join(type.value for type in ToolServerType)
228            raise ValueError(f"type must be one of: {valid_types}")
229
230    @model_validator(mode="before")
231    def upgrade_old_properties(cls, data: dict) -> dict:
232        """
233        Upgrade properties for backwards compatibility.
234        """
235        properties = data.get("properties")
236        if properties is not None and "is_archived" not in properties:
237            # Add is_archived field with default value back to data
238            properties["is_archived"] = False
239            data["properties"] = properties
240        return data
241
242    @model_validator(mode="before")
243    def validate_required_fields(cls, data: dict) -> dict:
244        """Validate that each tool type has the required configuration."""
245        server_type = ExternalToolServer.type_from_data(data)
246        properties = data.get("properties", {})
247
248        match server_type:
249            case ToolServerType.remote_mcp:
250                server_url = properties.get("server_url", None)
251                if server_url is None:
252                    raise ValueError(
253                        "Server URL is required to connect to a remote MCP server"
254                    )
255                ExternalToolServer.check_server_url(server_url)
256
257            case ToolServerType.local_mcp:
258                command = properties.get("command", None)
259                if command is None:
260                    raise ValueError("command is required to start a local MCP server")
261                if not isinstance(command, str):
262                    raise ValueError(
263                        "command must be a string to start a local MCP server"
264                    )
265                # Reject empty/whitespace-only command strings
266                if command.strip() == "":
267                    raise ValueError("command must be a non-empty string")
268
269                args = properties.get("args", None)
270                if args is not None:
271                    if not isinstance(args, list):
272                        raise ValueError(
273                            "arguments must be a list to start a local MCP server"
274                        )
275
276            case ToolServerType.kiln_task:
277                tool_name_validator(properties.get("name", ""))
278                err_msg_prefix = "Kiln task server properties:"
279                validate_return_dict_prop(
280                    properties, "description", str, err_msg_prefix
281                )
282                description = properties.get("description", "")
283                if len(description) > 128:
284                    raise ValueError("description must be 128 characters or less")
285                validate_return_dict_prop(
286                    properties, "is_archived", bool, err_msg_prefix
287                )
288                validate_return_dict_prop(properties, "task_id", str, err_msg_prefix)
289                validate_return_dict_prop(
290                    properties, "run_config_id", str, err_msg_prefix
291                )
292
293            case _:
294                # Type checking will catch missing cases
295                raise_exhaustive_enum_error(server_type)
296        return data
297
298    @model_validator(mode="before")
299    def validate_headers_and_env_vars(cls, data: dict) -> dict:
300        """
301        Validate secrets, these needs to be validated before model initlization because secrets will be processed and stripped
302        """
303        type = ExternalToolServer.type_from_data(data)
304
305        properties = data.get("properties", {})
306        if properties is None:
307            raise ValueError("properties is required")
308
309        match type:
310            case ToolServerType.remote_mcp:
311                # Validate headers
312                headers = properties.get("headers", None)
313                if headers is not None:
314                    ExternalToolServer.check_headers(headers)
315
316                # Secret header keys are optional, validate if they are set
317                secret_header_keys = properties.get("secret_header_keys", None)
318                if secret_header_keys is not None:
319                    ExternalToolServer.check_secret_keys(
320                        secret_header_keys, "secret_header_keys", "remote_mcp"
321                    )
322
323            case ToolServerType.local_mcp:
324                # Validate secret environment variable keys
325                env_vars = properties.get("env_vars", {})
326                if env_vars is not None:
327                    ExternalToolServer.check_env_vars(env_vars)
328
329                # Secret env var keys are optional, but if they are set, they must be a list of strings
330                secret_env_var_keys = properties.get("secret_env_var_keys", None)
331                if secret_env_var_keys is not None:
332                    ExternalToolServer.check_secret_keys(
333                        secret_env_var_keys, "secret_env_var_keys", "local_mcp"
334                    )
335
336            case ToolServerType.kiln_task:
337                pass
338
339            case _:
340                raise_exhaustive_enum_error(type)
341
342        return data
343
344    def get_secret_keys(self) -> list[str]:
345        """
346        Get the list of secret key names based on server type.
347
348        Returns:
349            List of secret key names (header names for remote, env var names for local)
350        """
351        match self.type:
352            case ToolServerType.remote_mcp:
353                return self.properties.get("secret_header_keys", [])
354            case ToolServerType.local_mcp:
355                return self.properties.get("secret_env_var_keys", [])
356            case ToolServerType.kiln_task:
357                return []
358            case _:
359                raise_exhaustive_enum_error(self.type)
360
361    def retrieve_secrets(self) -> tuple[dict[str, str], list[str]]:
362        """
363        Retrieve secrets from configuration system or in-memory storage.
364        Automatically determines which secret keys to retrieve based on the server type.
365        Config secrets take precedence over unsaved secrets.
366
367        Returns:
368            Tuple of (secrets_dict, missing_secrets_list) where:
369            - secrets_dict: Dictionary mapping key names to their secret values
370            - missing_secrets_list: List of secret key names that are missing values
371        """
372        secrets = {}
373        missing_secrets = []
374        secret_keys = self.get_secret_keys()
375
376        if secret_keys and len(secret_keys) > 0:
377            config = Config.shared()
378            mcp_secrets = config.get_value(MCP_SECRETS_KEY)
379
380            for key_name in secret_keys:
381                secret_value = None
382
383                # First check config secrets (persistent storage), key is mcp_server_id::key_name
384                secret_key = self._config_secret_key(key_name)
385                secret_value = mcp_secrets.get(secret_key) if mcp_secrets else None
386
387                # Fall back to unsaved secrets (in-memory storage)
388                if (
389                    not secret_value
390                    and hasattr(self, "_unsaved_secrets")
391                    and key_name in self._unsaved_secrets
392                ):
393                    secret_value = self._unsaved_secrets[key_name]
394
395                if secret_value:
396                    secrets[key_name] = secret_value
397                else:
398                    missing_secrets.append(key_name)
399
400        return secrets, missing_secrets
401
402    def _save_secrets(self) -> None:
403        """
404        Save unsaved secrets to the configuration system.
405        """
406        secret_keys = self.get_secret_keys()
407
408        # No secrets to save
409        if not secret_keys:
410            return
411
412        if self.id is None:
413            raise ValueError("Server ID cannot be None when saving secrets")
414
415        # Check if secrets are already saved
416        if not hasattr(self, "_unsaved_secrets") or not self._unsaved_secrets:
417            return
418
419        config = Config.shared()
420        mcp_secrets: dict[str, str] = config.get_value(MCP_SECRETS_KEY) or {}
421
422        # Store secrets with the pattern: mcp_server_id::key_name
423        for key_name, secret_value in self._unsaved_secrets.items():
424            secret_key = self._config_secret_key(key_name)
425            mcp_secrets[secret_key] = secret_value
426
427        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
428
429        # Clear unsaved secrets after saving
430        self._unsaved_secrets.clear()
431
432    def delete_secrets(self) -> None:
433        """
434        Delete all secrets for this tool server from the configuration system.
435        """
436        secret_keys = self.get_secret_keys()
437
438        config = Config.shared()
439        mcp_secrets = config.get_value(MCP_SECRETS_KEY) or dict[str, str]()
440
441        # Remove secrets with the pattern: mcp_server_id::key_name
442        for key_name in secret_keys:
443            secret_key = self._config_secret_key(key_name)
444            if secret_key in mcp_secrets:
445                del mcp_secrets[secret_key]
446
447        # Always call update_settings to maintain consistency with the old behavior
448        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
449
450    def save_to_file(self) -> None:
451        """
452        Override save_to_file to automatically save any unsaved secrets before saving to file.
453
454        This ensures that secrets are always saved when the object is saved,
455        preventing the issue where secrets could be lost if save_to_file is called
456        without explicitly saving secrets first.
457        """
458        # Save any unsaved secrets first
459        if hasattr(self, "_unsaved_secrets") and self._unsaved_secrets:
460            self._save_secrets()
461
462        # Call the parent save_to_file method
463        super().save_to_file()
464
465    #  Internal helpers
466
467    def _config_secret_key(self, key_name: str) -> str:
468        """
469        Generate the secret key pattern for storing/retrieving secrets.
470
471        Args:
472            key_name: The name of the secret key
473
474        Returns:
475            The formatted secret key: "{server_id}::{key_name}"
476        """
477        return f"{self.id}::{key_name}"