feat: separate agent vs local commands

Copilot · BukeLy · Copilot · commit 217506283b0e · 2026-01-09T09:15:08.000Z
Co-authored-by: BukeLy &lt;19304666+BukeLy@users.noreply.github.com&gt;
diff --git a/agent-sdk-client/config.py b/agent-sdk-client/config.py
@@ -28,26 +28,34 @@ def extract_command(text: Optional[str]) -> Optional[str]:
     return command
 
 
-def load_command_whitelist(config_path: Path = DEFAULT_CONFIG_PATH) -> list[str]:
-    """Load command whitelist from TOML config file."""
+def _load_config(config_path: Path = DEFAULT_CONFIG_PATH) -> tuple[list[str], dict[str, str]]:
+    """Load agent/local commands from TOML config file."""
     if not config_path.exists():
-        return []
+        return [], {}
 
     try:
         with config_path.open('rb') as f:
             data = tomllib.load(f)
-        whitelist = data.get('white_list_commands', {}).get('whitelist', [])
-        if not isinstance(whitelist, list):
-            logger.warning("Command whitelist is not a list; ignoring configuration")
-            return []
-
-        commands = [cmd for cmd in whitelist if isinstance(cmd, str)]
-        if len(commands) != len(whitelist):
-            logger.warning("Ignoring non-string entries in command whitelist")
-        return commands
+        agent_commands = data.get('agent_commands', {}).get('commands', [])
+        if not isinstance(agent_commands, list):
+            logger.warning("Agent commands config is not a list; ignoring configuration")
+            agent_commands = []
+        agent_commands = [cmd for cmd in agent_commands if isinstance(cmd, str)]
+
+        local_commands_raw = data.get('local_commands', {})
+        if not isinstance(local_commands_raw, dict):
+            logger.warning("Local commands config is not a table; ignoring configuration")
+            local_commands_raw = {}
+        local_commands = {
+            f"/{name.lstrip('/')}" if not name.startswith('/') else name: str(value)
+            for name, value in local_commands_raw.items()
+            if isinstance(name, str) and isinstance(value, str)
+        }
+
+        return agent_commands, local_commands
     except (OSError, tomllib.TOMLDecodeError) as exc:  # pragma: no cover - defensive logging
-        logger.warning("Failed to load command whitelist: %s", exc)
-    return []
+        logger.warning("Failed to load command configuration: %s", exc)
+    return [], {}
 
 
 @dataclass
@@ -58,22 +66,40 @@ class Config:
     agent_server_url: str
     auth_token: str
     queue_url: str
-    command_whitelist: list[str]
+    agent_commands: list[str]
+    local_commands: dict[str, str]
 
     @classmethod
     def from_env(cls, config_path: Optional[Path] = None) -> 'Config':
         """Load configuration from environment variables."""
+        agent_cmds, local_cmds = _load_config(config_path or DEFAULT_CONFIG_PATH)
         return cls(
             telegram_token=os.getenv('TELEGRAM_BOT_TOKEN', ''),
             agent_server_url=os.getenv('AGENT_SERVER_URL', ''),
             auth_token=os.getenv('SDK_CLIENT_AUTH_TOKEN', 'default-token'),
             queue_url=os.getenv('QUEUE_URL', ''),
-            command_whitelist=load_command_whitelist(config_path or DEFAULT_CONFIG_PATH),
+            agent_commands=agent_cmds,
+            local_commands=local_cmds,
         )
 
-    def is_command_allowed(self, text: Optional[str]) -> bool:
-        """Check whether text should be forwarded to Agent backend."""
-        command = extract_command(text)
-        if command is None:
-            return True
-        return command in self.command_whitelist
+    def get_command(self, text: Optional[str]) -> Optional[str]:
+        return extract_command(text)
+
+    def is_agent_command(self, cmd: Optional[str]) -> bool:
+        return bool(cmd) and cmd in self.agent_commands
+
+    def is_local_command(self, cmd: Optional[str]) -> bool:
+        return bool(cmd) and cmd in self.local_commands
+
+    def local_response(self, cmd: str) -> str:
+        return self.local_commands.get(cmd, "Unsupported command.")
+
+    def unknown_command_message(self) -> str:
+        parts = []
+        if self.agent_commands:
+            parts.append("Agent commands:\n" + "\n".join(self.agent_commands))
+        if self.local_commands:
+            parts.append("Local commands:\n" + "\n".join(self.local_commands.keys()))
+        if not parts:
+            return "Unsupported command."
+        return "Unsupported command.\n\n" + "\n\n".join(parts)
diff --git a/agent-sdk-client/config.toml b/agent-sdk-client/config.toml
@@ -1,6 +1,10 @@
-[white_list_commands]
-# Only commands in this whitelist will be forwarded to the Agent backend
-whitelist = [
+[agent_commands]
+# Commands forwarded to the Agent backend
+commands = [
     "/custom-skill",
     "/hello-world",
 ]
+
+[local_commands]
+# Local-only commands handled by the client
+help = "Hello World"
diff --git a/agent-sdk-client/consumer.py b/agent-sdk-client/consumer.py
@@ -55,31 +55,43 @@ async def process_message(message_data: dict) -> None:
         logger.warning("Received update with no message or edited_message")
         return
 
-    if not config.is_command_allowed(message.text):
-        # Defensive guard: producer should already block non-whitelisted commands.
-        logger.info(
-            "Skipping non-whitelisted command (consumer fallback)",
-            extra={
-                'chat_id': message.chat_id,
-                'message_id': message.message_id,
-            },
-        )
-        allowed = config.command_whitelist
-        if allowed:
-            allowed_list = "\n".join(allowed)
-            text = f"Unsupported command. Allowed commands:\n{allowed_list}"
-        else:
-            text = "Unsupported command."
-        try:
-            await bot.send_message(
-                chat_id=message.chat_id,
-                text=text,
-                message_thread_id=message.message_thread_id,
-                reply_to_message_id=message.message_id,
+    cmd = config.get_command(message.text)
+    if cmd:
+        if config.is_local_command(cmd):
+            logger.info(
+                "Handling local command in consumer (fallback path)",
+                extra={'chat_id': message.chat_id, 'message_id': message.message_id},
             )
-        except Exception:
-            logger.warning("Failed to send local command response", exc_info=True)
-        return
+            try:
+                await bot.send_message(
+                    chat_id=message.chat_id,
+                    text=config.local_response(cmd),
+                    message_thread_id=message.message_thread_id,
+                    reply_to_message_id=message.message_id,
+                )
+            except Exception:
+                logger.warning("Failed to send local command response", exc_info=True)
+            return
+
+        if not config.is_agent_command(cmd):
+            # Defensive guard: producer should already block non-agent commands.
+            logger.info(
+                "Skipping non-agent command (consumer fallback)",
+                extra={
+                    'chat_id': message.chat_id,
+                    'message_id': message.message_id,
+                },
+            )
+            try:
+                await bot.send_message(
+                    chat_id=message.chat_id,
+                    text=config.unknown_command_message(),
+                    message_thread_id=message.message_thread_id,
+                    reply_to_message_id=message.message_id,
+                )
+            except Exception:
+                logger.warning("Failed to send local command response", exc_info=True)
+            return
 
     # Send typing indicator
     await bot.send_chat_action(
diff --git a/agent-sdk-client/handler.py b/agent-sdk-client/handler.py
@@ -120,17 +120,12 @@ def _send_to_sqs_safe(sqs, queue_url: str, message_body: dict) -> bool:
         return False
 
 
-def _handle_local_command(bot: Bot, message, config: Config) -> bool:
-    """Handle non-whitelisted commands locally to give user feedback."""
-    if config.is_command_allowed(message.text):
-        return False
-
-    allowed = config.command_whitelist
-    if allowed:
-        allowed_list = "\n".join(allowed)
-        text = f"Unsupported command. Allowed commands:\n{allowed_list}"
+def _handle_local_command(bot: Bot, message, config: Config, cmd: str) -> bool:
+    """Handle local commands or unknown commands."""
+    if config.is_local_command(cmd):
+        text = config.local_response(cmd)
     else:
-        text = "Unsupported command."
+        text = config.unknown_command_message()
 
     try:
         bot.send_message(
@@ -179,7 +174,13 @@ def lambda_handler(event: dict, context: Any) -> dict:
         logger.debug('Ignoring webhook without text message')
         return {'statusCode': 200}
 
-    if _handle_local_command(bot, message, config):
+    cmd = config.get_command(message.text)
+    if cmd and config.is_local_command(cmd):
+        _handle_local_command(bot, message, config, cmd)
+        return {'statusCode': 200}
+
+    if cmd and not config.is_agent_command(cmd):
+        _handle_local_command(bot, message, config, cmd)
         return {'statusCode': 200}
 
     # Write to SQS for async processing
diff --git a/tests/test_command_config.py b/tests/test_command_config.py
@@ -0,0 +1,117 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+CLIENT_CONFIG_PATH = Path(__file__).resolve().parent.parent / "agent-sdk-client" / "config.py"
+spec = importlib.util.spec_from_file_location("agent_sdk_client_config", CLIENT_CONFIG_PATH)
+config_module = importlib.util.module_from_spec(spec)
+assert spec.loader is not None
+spec.loader.exec_module(config_module)
+Config = config_module.Config
+extract_command = config_module.extract_command
+
+
+def load_config_from_text(text: str, tmp_path: Path) -> Config:
+    config_path = tmp_path / "config.toml"
+    config_path.write_text(text)
+    return Config.from_env(config_path=config_path)
+
+
+def test_load_agent_and_local_commands(tmp_path):
+    cfg = load_config_from_text(
+        """[agent_commands]
+commands = ["/a", "/b"]
+
+[local_commands]
+help = "Hello"
+""",
+        tmp_path,
+    )
+    assert cfg.agent_commands == ["/a", "/b"]
+    assert cfg.local_commands == {"/help": "Hello"}
+
+
+@pytest.mark.parametrize(
+    "text,cmd",
+    [
+        ("hello world", None),
+        ("/allowed", "/allowed"),
+        ("/allowed extra args", "/allowed"),
+        ("/allowed@bot", "/allowed"),
+        ("/@bot", None),
+        ("/", None),
+        (None, None),
+    ],
+)
+def test_extract_command(text, cmd):
+    assert extract_command(text) == cmd
+
+
+def test_command_classification(tmp_path):
+    cfg = load_config_from_text(
+        """[agent_commands]
+commands = ["/agent"]
+
+[local_commands]
+help = "Hello World"
+""",
+        tmp_path,
+    )
+    assert cfg.is_agent_command("/agent")
+    assert not cfg.is_agent_command("/other")
+    assert cfg.is_local_command("/help")
+    assert not cfg.is_local_command("/agent")
+
+
+def test_unknown_command_message_lists_known():
+    cfg = Config(
+        telegram_token="",
+        agent_server_url="",
+        auth_token="",
+        queue_url="",
+        agent_commands=["/agent1"],
+        local_commands={"/help": "hi"},
+    )
+    msg = cfg.unknown_command_message()
+    assert "Agent commands" in msg and "/agent1" in msg
+    assert "Local commands" in msg and "/help" in msg
+
+
+def test_invalid_agent_commands_type(tmp_path, caplog):
+    with caplog.at_level("WARNING"):
+        cfg = load_config_from_text(
+            """[agent_commands]
+commands = "not-a-list"
+""",
+            tmp_path,
+        )
+    assert cfg.agent_commands == []
+    assert any("Agent commands config is not a list" in rec.message for rec in caplog.records)
+
+
+def test_invalid_local_commands_type(tmp_path, caplog):
+    cfg = load_config_from_text(
+        """[local_commands]
+value = 1
+""",
+        tmp_path,
+    )
+    assert cfg.local_commands == {}
+
+
+def test_missing_config_file(tmp_path):
+    missing = tmp_path / "missing.toml"
+    cfg = Config.from_env(config_path=missing)
+    assert cfg.agent_commands == []
+    assert cfg.local_commands == {}
+
+
+def test_malformed_toml_returns_empty(tmp_path, caplog):
+    path = tmp_path / "bad.toml"
+    path.write_text("not = [ [")
+    with caplog.at_level("WARNING"):
+        cfg = Config.from_env(config_path=path)
+    assert cfg.agent_commands == []
+    assert cfg.local_commands == {}
+    assert any("Failed to load command configuration" in rec.message for rec in caplog.records)
diff --git a/tests/test_command_whitelist.py b/tests/test_command_whitelist.py
