From 30526127ed5397e693d9b890335fe3e8635329cb Mon Sep 17 00:00:00 2001
From: Chris Penner <christopher.penner@gmail.com>
Date: Mon, 18 Nov 2024 15:32:31 -0800
Subject: [PATCH] Add indicator to prompt when on staging

---
 unison-cli/src/Unison/CommandLine/Main.hs | 11 ++++++++-
 unison-cli/src/Unison/Share/Codeserver.hs | 30 ++++++++++++++++-------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/unison-cli/src/Unison/CommandLine/Main.hs b/unison-cli/src/Unison/CommandLine/Main.hs
index 64e070be74..3b86508eb0 100644
--- a/unison-cli/src/Unison/CommandLine/Main.hs
+++ b/unison-cli/src/Unison/CommandLine/Main.hs
@@ -47,6 +47,8 @@ import Unison.Prelude
 import Unison.PrettyTerminal
 import Unison.Runtime.IOSource qualified as IOSource
 import Unison.Server.CodebaseServer qualified as Server
+import Unison.Share.Codeserver (isCustomCodeserver)
+import Unison.Share.Codeserver qualified as Codeserver
 import Unison.Symbol (Symbol)
 import Unison.Syntax.Parser qualified as Parser
 import Unison.Util.Pretty qualified as P
@@ -75,10 +77,17 @@ getUserInput codebase authHTTPClient pp currentProjectRoot numberedArgs =
       Line.handleInterrupt (pure Nothing) (Line.withInterrupt (Just <$> act)) >>= \case
         Nothing -> haskelineCtrlCHandling act
         Just a -> pure a
+
+    codeserverPrompt :: String
+    codeserverPrompt =
+      if isCustomCodeserver Codeserver.defaultCodeserver
+        then "🌐" <> Codeserver.codeserverRegName Codeserver.defaultCodeserver <> maybe "" (":" <>) (show <$> Codeserver.codeserverPort Codeserver.defaultCodeserver) <> "\n"
+        else ""
+
     go :: Line.InputT IO Input
     go = do
       let promptString = P.prettyProjectPath pp
-      let fullPrompt = P.toANSI 80 (promptString <> fromString prompt)
+      let fullPrompt = P.toANSI 80 (P.red (P.string codeserverPrompt) <> promptString <> fromString prompt)
       line <- Line.getInputLine fullPrompt
       case line of
         Nothing -> pure QuitI
diff --git a/unison-cli/src/Unison/Share/Codeserver.hs b/unison-cli/src/Unison/Share/Codeserver.hs
index a1617a4411..ea7aee4b73 100644
--- a/unison-cli/src/Unison/Share/Codeserver.hs
+++ b/unison-cli/src/Unison/Share/Codeserver.hs
@@ -1,4 +1,10 @@
-module Unison.Share.Codeserver where
+module Unison.Share.Codeserver
+  ( isCustomCodeserver,
+    defaultCodeserver,
+    resolveCodeserver,
+    CodeserverURI (..),
+  )
+where
 
 import Network.URI (parseURI)
 import System.IO.Unsafe (unsafePerformIO)
@@ -8,18 +14,24 @@ import Unison.Share.Types
 import Unison.Share.Types qualified as Share
 import UnliftIO.Environment (lookupEnv)
 
+shareProd :: CodeserverURI
+shareProd =
+  CodeserverURI
+    { codeserverScheme = Share.Https,
+      codeserverUserInfo = "",
+      codeserverRegName = "api.unison-lang.org",
+      codeserverPort = Nothing,
+      codeserverPath = []
+    }
+
+isCustomCodeserver :: CodeserverURI -> Bool
+isCustomCodeserver = (/=) shareProd
+
 -- | This is the URI where the share API is based.
 defaultCodeserver :: CodeserverURI
 defaultCodeserver = unsafePerformIO $ do
   lookupEnv "UNISON_SHARE_HOST" <&> \case
-    Nothing ->
-      CodeserverURI
-        { codeserverScheme = Share.Https,
-          codeserverUserInfo = "",
-          codeserverRegName = "api.unison-lang.org",
-          codeserverPort = Nothing,
-          codeserverPath = []
-        }
+    Nothing -> shareProd
     Just shareHost ->
       fromMaybe (error $ "Share Host is not a valid URI: " <> shareHost) $ do
         uri <- parseURI shareHost