From 638204ab3c09bb2508e53b7a18df2f76c9b95139 Mon Sep 17 00:00:00 2001
From: Devansh Jain <31609257+devanshrj@users.noreply.github.com>
Date: Thu, 5 Mar 2026 20:50:45 -0800
Subject: [PATCH] fix: restrict --system flag to known preset IDs (#1290)

---
 src/agent/promptAssets.ts | 47 ++++++++++++++++++++++++++++++++-------
 src/cli/args.ts           |  3 +--
 src/index.ts              | 25 +++++++++------------
 3 files changed, 51 insertions(+), 24 deletions(-)

diff --git a/src/agent/promptAssets.ts b/src/agent/promptAssets.ts
index e31a727..715c76a 100644
--- a/src/agent/promptAssets.ts
+++ b/src/agent/promptAssets.ts
@@ -114,32 +114,62 @@ export const SYSTEM_PROMPTS: SystemPromptOption[] = [
   },
 ];
 
+/**
+ * Validate a system prompt preset ID.
+ *
+ * Known preset IDs are always accepted. Subagent names are only accepted
+ * when `allowSubagentNames` is true (internal subagent launches).
+ *
+ * @throws Error with a descriptive message listing valid options
+ */
+export async function validateSystemPromptPreset(
+  id: string,
+  opts?: { allowSubagentNames?: boolean },
+): Promise<void> {
+  const validPresets = SYSTEM_PROMPTS.map((p) => p.id);
+  if (validPresets.includes(id)) return;
+
+  if (opts?.allowSubagentNames) {
+    const { getAllSubagentConfigs } = await import("./subagents");
+    const subagentConfigs = await getAllSubagentConfigs();
+    if (subagentConfigs[id]) return;
+
+    const allValid = [...validPresets, ...Object.keys(subagentConfigs)];
+    throw new Error(
+      `Invalid system prompt "${id}". Must be one of: ${allValid.join(", ")}.`,
+    );
+  }
+
+  throw new Error(
+    `Invalid system prompt "${id}". Must be one of: ${validPresets.join(", ")}.`,
+  );
+}
+
 /**
  * Resolve a system prompt ID to its content.
  *
  * Resolution order:
- * 1. If it matches an ID from SYSTEM_PROMPTS, use its content
- * 2. If it matches a subagent name, use that subagent's system prompt
- * 3. Otherwise, use the default system prompt
+ * 1. No input → default system prompt
+ * 2. Known preset ID → preset content
+ * 3. Subagent name → subagent's system prompt
+ * 4. Unknown → throws (callers should validate first via validateSystemPromptPreset)
  *
  * @param systemPromptPreset - The system prompt preset (e.g., "letta-claude") or subagent name (e.g., "explore")
  * @returns The resolved system prompt content
+ * @throws Error if the ID doesn't match any preset or subagent
  */
 export async function resolveSystemPrompt(
   systemPromptPreset: string | undefined,
 ): Promise<string> {
-  // No input - use default
   if (!systemPromptPreset) {
     return SYSTEM_PROMPT;
   }
 
-  // 1. Check if it matches a system prompt ID
   const matchedPrompt = SYSTEM_PROMPTS.find((p) => p.id === systemPromptPreset);
   if (matchedPrompt) {
     return matchedPrompt.content;
   }
 
-  // 2. Check if it matches a subagent name
   const { getAllSubagentConfigs } = await import("./subagents");
   const subagentConfigs = await getAllSubagentConfigs();
   const matchedSubagent = subagentConfigs[systemPromptPreset];
@@ -147,6 +177,7 @@ export async function resolveSystemPrompt(
     return matchedSubagent.systemPrompt;
   }
 
-  // 3. Fall back to default
-  return SYSTEM_PROMPT;
+  throw new Error(
+    `Unknown system prompt "${systemPromptPreset}" — does not match any preset or subagent`,
+  );
 }
diff --git a/src/cli/args.ts b/src/cli/args.ts
index b439239..ac4faae 100644
--- a/src/cli/args.ts
+++ b/src/cli/args.ts
@@ -106,8 +106,7 @@ export const CLI_FLAG_CATALOG = {
     mode: "both",
     help: {
       argLabel: "<id>",
-      description:
-        "System prompt ID or subagent name (applies to new or existing agent)",
+      description: "System prompt preset ID (applies to new or existing agent)",
     },
   },
   "system-custom": { parser: { type: "string" }, mode: "both" },
diff --git a/src/index.ts b/src/index.ts
index 52d53cf..6f03e56 100755
--- a/src/index.ts
+++ b/src/index.ts
@@ -576,22 +576,19 @@ async function main(): Promise<void> {
     process.exit(1);
   }
 
-  // Validate system prompt preset if provided (can be a system prompt ID or subagent name)
+  // Validate system prompt preset if provided.
+  // Known preset IDs are always accepted. Subagent names are only accepted
+  // for internal subagent launches (LETTA_CODE_AGENT_ROLE=subagent).
   if (systemPromptPreset) {
-    const { SYSTEM_PROMPTS } = await import("./agent/promptAssets");
-    const { getAllSubagentConfigs } = await import("./agent/subagents");
-
-    const validSystemPrompts = SYSTEM_PROMPTS.map((p) => p.id);
-    const subagentConfigs = await getAllSubagentConfigs();
-    const validSubagentNames = Object.keys(subagentConfigs);
-
-    const isValidSystemPrompt = validSystemPrompts.includes(systemPromptPreset);
-    const isValidSubagent = validSubagentNames.includes(systemPromptPreset);
-
-    if (!isValidSystemPrompt && !isValidSubagent) {
-      const allValid = [...validSystemPrompts, ...validSubagentNames];
+    const { validateSystemPromptPreset } = await import("./agent/promptAssets");
+    const allowSubagentNames = process.env.LETTA_CODE_AGENT_ROLE === "subagent";
+    try {
+      await validateSystemPromptPreset(systemPromptPreset, {
+        allowSubagentNames,
+      });
+    } catch (err) {
       console.error(
-        `Error: Invalid system prompt "${systemPromptPreset}". Must be one of: ${allValid.join(", ")}.`,
+        `Error: ${err instanceof Error ? err.message : String(err)}`,
       );
       process.exit(1);
     }