From 1bfc647132d7b0bdcd55316a2cacee64c1c6b95e Mon Sep 17 00:00:00 2001
From: Lifei Zhou <lifei@squareup.com>
Date: Thu, 29 Aug 2024 09:03:12 +1000
Subject: [PATCH] Enable cli options for plugin (#22)

* added the entry point for plugin with cli group option
---
 pyproject.toml         |  6 ++++--
 src/goose/cli/main.py  | 25 +++++++++++++++++--------
 tests/cli/test_main.py | 42 +++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 18bff707f..124304d17 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "goose-ai"
 description = "a programming agent that runs on your machine"
-version = "0.8.0"
+version = "0.8.1"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
@@ -30,9 +30,11 @@ default = "goose.profile:default_profile"
 [project.entry-points."goose.command"]
 file = "goose.command.file:FileCommand"
 
-[project.entry-points."goose.cli"]
+[project.entry-points."goose.cli.group"]
 goose = "goose.cli.main:goose_cli"
 
+[project.entry-points."goose.cli.group_option"]
+
 [project.scripts]
 goose = "goose.cli.main:cli"
 
diff --git a/src/goose/cli/main.py b/src/goose/cli/main.py
index 30bf3d9b1..61c481ea4 100644
--- a/src/goose/cli/main.py
+++ b/src/goose/cli/main.py
@@ -12,12 +12,10 @@
 from goose.utils import load_plugins
 from goose.utils.session_file import list_sorted_session_files
 
-
 @click.group()
 def goose_cli() -> None:
     pass
 
-
 @goose_cli.command()
 def version() -> None:
     """Lists the version of goose and any plugins"""
@@ -81,7 +79,7 @@ def session_start(profile: str, plan: Optional[str] = None) -> None:
 @session.command(name="resume")
 @click.argument("name", required=False)
 @click.option("--profile")
-def session_resume(name: str, profile: str) -> None:
+def session_resume(name: Optional[str], profile: str) -> None:
     """Resume an existing goose session"""
     if name is None:
         session_files = get_session_files()
@@ -97,6 +95,7 @@ def session_resume(name: str, profile: str) -> None:
 
 @session.command(name="list")
 def session_list() -> None:
+    """List goose sessions"""
     session_files = get_session_files().items()
     for session_name, session_file in session_files:
         print(f"{datetime.fromtimestamp(session_file.stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S')}    {session_name}")
@@ -105,6 +104,7 @@ def session_list() -> None:
 @session.command(name="clear")
 @click.option("--keep", default=3, help="Keep this many entries, default 3")
 def session_clear(keep: int) -> None:
+    """Delete old goose sessions, keeping the most recent sessions up to the specified number"""
     for i, (_, session_file) in enumerate(get_session_files().items()):
         if i >= keep:
             session_file.unlink()
@@ -113,13 +113,22 @@ def session_clear(keep: int) -> None:
 def get_session_files() -> Dict[str, Path]:
     return list_sorted_session_files(SESSIONS_PATH)
 
+@click.group(
+        invoke_without_command=True,
+        name="goose",
+        help="AI-powered tool to assist in solving programming and operational tasks",)
+@click.pass_context
+def cli(_: click.Context, **kwargs: Dict) -> None:
+    pass
 
-# merging goose cli with additional cli plugins.
-def cli() -> None:
-    clis = load_plugins("goose.cli")
-    cli_list = list(clis.values()) or []
-    click.CommandCollection(sources=cli_list)()
+all_cli_group_options = load_plugins("goose.cli.group_option")
+for option in all_cli_group_options.values():
+    cli = option()(cli)
 
+all_cli_groups = load_plugins("goose.cli.group")
+for group in all_cli_groups.values():
+    for command in group.commands.values():
+        cli.add_command(command)
 
 if __name__ == "__main__":
     cli()
diff --git a/tests/cli/test_main.py b/tests/cli/test_main.py
index 38d4c6c74..253aa3a3c 100644
--- a/tests/cli/test_main.py
+++ b/tests/cli/test_main.py
@@ -1,11 +1,13 @@
 from datetime import datetime
+import importlib
 from time import time
 from unittest.mock import MagicMock, patch
 
+import click
 import pytest
 from click.testing import CliRunner
 from exchange import Message
-from goose.cli.main import goose_cli
+from goose.cli.main import cli, goose_cli
 
 
 @pytest.fixture
@@ -78,3 +80,41 @@ def test_session_clear_command(mock_session_files_path, create_session_file):
     session_files = list(mock_session_files_path.glob("*.jsonl"))
     assert len(session_files) == 1
     assert session_files[0].stem == "second"
+
+
+def test_combined_group_option():
+    with patch("goose.utils.load_plugins") as mock_load_plugin:
+        group_option_name = "--describe-commands"
+        def option_callback(ctx, *_):
+            click.echo("Option callback")
+            ctx.exit()
+        mock_group_options = {
+            'option1': lambda: click.option(
+                group_option_name,
+                is_flag=True,
+                callback=option_callback,
+            ),
+        }
+        def side_effect_func(param):
+            if param == "goose.cli.group_option":
+                return mock_group_options
+            elif param == "goose.cli.group":
+                return {  }
+        mock_load_plugin.side_effect = side_effect_func
+
+        # reload cli after mocking
+        importlib.reload(importlib.import_module('goose.cli.main'))
+        import goose.cli.main
+        cli = goose.cli.main.cli
+
+        runner = CliRunner()
+        result = runner.invoke(cli, [group_option_name])
+        assert result.exit_code == 0
+
+def test_combined_group_commands(mock_session):
+    mock_session_class, mock_session_instance = mock_session
+    runner = CliRunner()
+    runner.invoke(cli, ["session", "resume", "session1", "--profile", "default"])
+    mock_session_class.assert_called_once_with(name="session1", profile="default")
+    mock_session_instance.run.assert_called_once()
+