Merge pull request #5840 from unisonweb/cp/share-http-clients

aryairani · web-flow · commit c9c45fc5241b · 2025-08-19T11:06:15.000-04:00
diff --git a/unison-cli/src/Unison/Auth/Tokens.hs b/unison-cli/src/Unison/Auth/Tokens.hs
@@ -8,6 +8,7 @@ import Data.Time.Clock (getCurrentTime)
 import Network.HTTP.Client qualified as HTTP
 import Network.HTTP.Client.TLS qualified as HTTP
 import Network.HTTP.Types qualified as Network
+import System.Environment (lookupEnv)
 import Unison.Auth.CredentialManager
 import Unison.Auth.Discovery (fetchDiscoveryDoc)
 import Unison.Auth.Types
@@ -20,21 +21,34 @@ import UnliftIO qualified
 -- The TokenProvider may automatically refresh access tokens if we have a refresh token.
 type TokenProvider = CodeserverId -> IO (Either CredentialFailure AccessToken)
 
+-- | If provided, this access token will be used on all
+-- requests which use the Authenticated HTTP Client; i.e. all codeserver interactions.
+--
+-- It's useful in scripted contexts or when running transcripts against a codeserver.
+accessTokenEnvVarKey :: String
+accessTokenEnvVarKey = "UNISON_SHARE_ACCESS_TOKEN"
+
 -- | Creates a 'TokenProvider' using the given 'CredentialManager'
 newTokenProvider :: CredentialManager -> TokenProvider
 newTokenProvider manager host = UnliftIO.try @_ @CredentialFailure $ do
-  creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCredentials manager host
-  let Tokens {accessToken = currentAccessToken} = tokens
-  expired <- isExpired creds
-  if expired
-    then do
-      discoveryDoc <- throwEitherM $ fetchDiscoveryDoc discoveryURI
-      fetchTime <- getCurrentTime
-      newTokens@(Tokens {accessToken = newAccessToken}) <- throwEitherM $ performTokenRefresh discoveryDoc tokens
-      userInfo <- throwEitherM $ getUserInfo discoveryDoc newAccessToken
-      saveCredentials manager host (codeserverCredentials discoveryURI newTokens fetchTime userInfo)
-      pure $ newAccessToken
-    else pure currentAccessToken
+  mayShareAccessToken <- fmap Text.pack <$> lookupEnv accessTokenEnvVarKey
+  case mayShareAccessToken of
+    Just accessToken -> do
+      -- If the access token is provided via environment variable, we don't need to refresh it.
+      pure accessToken
+    Nothing -> do
+      creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCredentials manager host
+      let Tokens {accessToken = currentAccessToken} = tokens
+      expired <- isExpired creds
+      if expired
+        then do
+          discoveryDoc <- throwEitherM $ fetchDiscoveryDoc discoveryURI
+          fetchTime <- getCurrentTime
+          newTokens@(Tokens {accessToken = newAccessToken}) <- throwEitherM $ performTokenRefresh discoveryDoc tokens
+          userInfo <- throwEitherM $ getUserInfo discoveryDoc newAccessToken
+          saveCredentials manager host (codeserverCredentials discoveryURI newTokens fetchTime userInfo)
+          pure newAccessToken
+        else pure currentAccessToken
 
 -- | Don't yet support automatically refreshing tokens.
 --
diff --git a/unison-cli/src/Unison/Codebase/Transcript/Runner.hs b/unison-cli/src/Unison/Codebase/Transcript/Runner.hs
@@ -24,7 +24,6 @@ import Data.Text.Encoding qualified as Text
 import Data.These (These (..))
 import Data.UUID.V4 qualified as UUID
 import Network.HTTP.Client qualified as HTTP
-import System.Environment (lookupEnv)
 import System.IO qualified as IO
 import Text.Megaparsec qualified as P
 import U.Codebase.Sqlite.DbId qualified as Db
@@ -79,13 +78,6 @@ import Prelude hiding (readFile, writeFile)
 terminalWidth :: Pretty.Width
 terminalWidth = 65
 
--- | If provided, this access token will be used on all
--- requests which use the Authenticated HTTP Client; i.e. all codeserver interactions.
---
--- It's useful in scripted contexts or when running transcripts against a codeserver.
-accessTokenEnvVarKey :: String
-accessTokenEnvVarKey = "UNISON_SHARE_ACCESS_TOKEN"
-
 type Runner =
   -- | The name of the transcript to run.
   String ->
@@ -104,14 +96,18 @@ withRunner ::
   (Runner -> m r) ->
   m r
 withRunner isTest verbosity ucmVersion action = do
+  credMan <- AuthN.newCredentialManager
+  authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient credMan
+
   -- If we're in a transcript test, configure the environment to use a non-existent fzf binary
   -- so that errors are consistent.
   -- This also prevents automated transcript tests from mistakenly opening fzf and waiting for user input.
   when isTest $ do
     liftIO $ setEnv Fuzzy.fzfPathEnvVar "NONE"
   withRuntimes \runtime sbRuntime ->
     action \transcriptName transcriptSrc codebase -> do
-      mcpServerConfig <- MCP.initServer codebase runtime sbRuntime Nothing ucmVersion
+      let workDir = Nothing
+      mcpServerConfig <- MCP.initServer codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient
       Server.startServer
         isTest
         Backend.BackendEnv {Backend.useNamesIndex = False}
@@ -121,19 +117,33 @@ withRunner isTest verbosity ucmVersion action = do
         (MCP.mcpServer mcpServerConfig)
         \case
           Nothing -> pure $ Left PortBindingFailure
-          Just baseUrl ->
-            either
-              (pure . Left . ParseError)
-              ( run isTest verbosity codebase runtime sbRuntime ucmVersion $
-                  tShow @Server.BaseUrl baseUrl
-              )
-              $ Transcript.parse transcriptName transcriptSrc
+          Just baseUrl -> do
+            let baseUrlText = tShow @Server.BaseUrl baseUrl
+            case (Transcript.parse transcriptName transcriptSrc) of
+              Left parseError -> pure $ Left (ParseError parseError)
+              Right stanzas ->
+                run
+                  isTest
+                  verbosity
+                  codebase
+                  runtime
+                  sbRuntime
+                  ucmVersion
+                  baseUrlText
+                  authenticatedHTTPClient
+                  credMan
+                  stanzas
   where
     withRuntimes :: (Runtime.Runtime Symbol -> Runtime.Runtime Symbol -> m a) -> m a
     withRuntimes action =
       RTI.withRuntime False RTI.Persistent ucmVersion \runtime ->
         RTI.withRuntime True RTI.Persistent ucmVersion \sbRuntime ->
           action runtime sbRuntime
+    initTranscriptAuthenticatedHTTPClient :: AuthN.CredentialManager -> m AuthN.AuthenticatedHttpClient
+    initTranscriptAuthenticatedHTTPClient credMan = liftIO $ do
+      let tokenProvider :: AuthN.TokenProvider
+          tokenProvider = AuthN.newTokenProvider credMan
+      AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
 
 isGeneratedBlock :: ProcessedBlock -> Bool
 isGeneratedBlock = generated . getCommonInfoTags
@@ -147,9 +157,11 @@ run ::
   Runtime.Runtime Symbol ->
   UCMVersion ->
   Text ->
+  AuthN.AuthenticatedHttpClient ->
+  AuthN.CredentialManager ->
   Transcript ->
   IO (Either Error Transcript)
-run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript = UnliftIO.try do
+run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL authenticatedHTTPClient credMan transcript = UnliftIO.try do
   let behaviors = extractBehaviors $ settings transcript
   let stanzas' = stanzas transcript
   httpManager <- HTTP.newManager HTTP.defaultManagerSettings
@@ -163,14 +175,6 @@ run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript =
         "Running the provided transcript file...",
         ""
       ]
-  mayShareAccessToken <- fmap Text.pack <$> lookupEnv accessTokenEnvVarKey
-  credMan <- AuthN.newCredentialManager
-  let tokenProvider :: AuthN.TokenProvider
-      tokenProvider =
-        maybe
-          (AuthN.newTokenProvider credMan)
-          (\accessToken _codeserverID -> pure $ Right accessToken)
-          mayShareAccessToken
   -- Queue of Stanzas and Just index, or Nothing if the stanza was programmatically generated
   -- e.g. a unison-file update by a command like 'edit'
   inputQueue <-
@@ -510,8 +514,6 @@ run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript =
               \issues."
           (_, _, _) -> pure ()
 
-  authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
-
   seedRef <- newIORef (0 :: Int)
 
   let env =
diff --git a/unison-cli/src/Unison/CommandLine/Main.hs b/unison-cli/src/Unison/CommandLine/Main.hs
@@ -21,10 +21,9 @@ import System.Console.Haskeline.History qualified as Line
 import System.FSNotify qualified as FSNotify
 import System.IO (hGetEcho, hPutStrLn, hSetEcho, stderr, stdin)
 import System.IO.Error (isDoesNotExistError)
-import Unison.Auth.CredentialManager (newCredentialManager)
+import Unison.Auth.CredentialManager qualified as AuthN
 import Unison.Auth.HTTPClient (AuthenticatedHttpClient)
 import Unison.Auth.HTTPClient qualified as AuthN
-import Unison.Auth.Tokens qualified as AuthN
 import Unison.Cli.Monad qualified as Cli
 import Unison.Cli.Pretty qualified as P
 import Unison.Cli.ProjectUtils qualified as ProjectUtils
@@ -143,10 +142,12 @@ main ::
   Codebase IO Symbol Ann ->
   Maybe Server.BaseUrl ->
   UCMVersion ->
+  AuthN.AuthenticatedHttpClient ->
+  AuthN.CredentialManager ->
   (PP.ProjectPathIds -> IO ()) ->
   ShouldWatchFiles ->
   IO ()
-main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion lspCheckForChanges shouldWatchFiles = do
+main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion authHTTPClient credentialManager lspCheckForChanges shouldWatchFiles = do
   -- we don't like FSNotify's debouncing (it seems to drop later events)
   -- so we will be doing our own instead
   let config = FSNotify.defaultConfig
@@ -175,9 +176,6 @@ main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl uc
       initialInputsRef <- newIORef $ Welcome.run welcome ++ initialInputs
       pageOutput <- newIORef True
 
-      credentialManager <- newCredentialManager
-      let tokenProvider = AuthN.newTokenProvider credentialManager
-      authHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
       initialEcho <- hGetEcho stdin
       let restoreEcho = (\currentEcho -> when (currentEcho /= initialEcho) $ hSetEcho stdin initialEcho)
       let getInput :: Cli.LoopState -> IO Input
diff --git a/unison-cli/src/Unison/MCP.hs b/unison-cli/src/Unison/MCP.hs
@@ -4,9 +4,7 @@ import Network.MCP.Server qualified as MCP
 import Network.MCP.Server.StdIO qualified as MCP
 import Network.MCP.Types
 import Text.RawString.QQ (r)
-import Unison.Auth.CredentialManager qualified as AuthN
 import Unison.Auth.HTTPClient qualified as AuthN
-import Unison.Auth.Tokens qualified as AuthN
 import Unison.Codebase (Codebase)
 import Unison.Codebase.Runtime (Runtime)
 import Unison.MCP.Prompts (prompts)
@@ -28,12 +26,8 @@ serverDescription =
         Before doing any work in unison please read the file://unison-guide resource for information on how to write unison.
     |]
 
-initServer :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> Maybe FilePath -> Text -> IO MCP.Server
-initServer codebase runtime sbRuntime workDir ucmVersion = do
-  credMan <- AuthN.newCredentialManager
-  let tokenProvider :: AuthN.TokenProvider
-      tokenProvider = AuthN.newTokenProvider credMan
-  authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
+initServer :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> Maybe FilePath -> Text -> AuthN.AuthenticatedHttpClient -> IO MCP.Server
+initServer codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient = do
   let env =
         Env
           { codebase,
@@ -49,8 +43,8 @@ initServer codebase runtime sbRuntime workDir ucmVersion = do
   runMCP env $ MCPWrapper.mkServer serverInfo serverDescription staticResources tools prompts
 
 -- | Run the MCP server until we hit EOF.
-runOnStdIO :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> IO ()
-runOnStdIO codebase runtime sbRuntime workDir ucmVersion = do
-  server <- initServer codebase runtime sbRuntime (pure workDir) ucmVersion
+runOnStdIO :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> AuthN.AuthenticatedHttpClient -> IO ()
+runOnStdIO codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient = do
+  server <- initServer codebase runtime sbRuntime (Just workDir) ucmVersion authenticatedHTTPClient
   -- Start the server with StdIO transport
   MCP.runServerWithSTDIO server
diff --git a/unison-cli/src/Unison/Main.hs b/unison-cli/src/Unison/Main.hs
@@ -60,10 +60,14 @@ import System.IO.Temp qualified as Temp
 import System.Path qualified as Path
 import Text.Megaparsec qualified as MP
 import U.Codebase.Sqlite.Queries qualified as Queries
+import Unison.Auth.CredentialManager qualified as AuthN
+import Unison.Auth.HTTPClient qualified as AuthN
+import Unison.Auth.Tokens qualified as AuthN
 import Unison.Cli.ProjectUtils qualified as ProjectUtils
 import Unison.Codebase (Codebase, CodebasePath)
 import Unison.Codebase qualified as Codebase
 import Unison.Codebase.Editor.Input qualified as Input
+import Unison.Codebase.Editor.UCMVersion (UCMVersion)
 import Unison.Codebase.Execute (execute)
 import Unison.Codebase.Init (CodebaseInitOptions (..), InitError (..), InitResult (..), SpecifiedCodebase (..))
 import Unison.Codebase.Init qualified as CodebaseInit
@@ -151,9 +155,12 @@ main version = do
         PrintVersion ->
           Text.putStrLn $ Text.pack progName <> " version: " <> Version.gitDescribeWithDate version
         MCPServer -> do
+          let ucmVersion = Version.gitDescribeWithDate version
+          credMan <- AuthN.newCredentialManager
+          authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
           getCodebaseOrExit mCodePathOption SC.DontLock (SC.MigrateAfterPrompt SC.Backup SC.Vacuum) \(_initRes, _, theCodebase) -> do
             withRuntimes RTI.Persistent \(runtime, sbRuntime) -> do
-              MCP.runOnStdIO theCodebase runtime sbRuntime currentDir (Version.gitDescribeWithDate version)
+              MCP.runOnStdIO theCodebase runtime sbRuntime currentDir ucmVersion authenticatedHTTPClient
         Init -> do
           exitError
             ( P.lines
@@ -184,6 +191,9 @@ main version = do
                       let fileEvent = Input.UnisonFileChanged (Text.pack file) contents
                       let noOpCheckForChanges _ = pure ()
                       let serverUrl = Nothing
+                      let ucmVersion = Version.gitDescribeWithDate version
+                      credMan <- liftIO $ AuthN.newCredentialManager
+                      authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
                       startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
                       launch
                         version
@@ -192,6 +202,8 @@ main version = do
                         sbrt
                         theCodebase
                         [Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI]
+                        authenticatedHTTPClient
+                        credMan
                         serverUrl
                         (PP.toIds startProjectPath)
                         initRes
@@ -207,6 +219,9 @@ main version = do
                   let fileEvent = Input.UnisonFileChanged (Text.pack "<standard input>") contents
                   let noOpCheckForChanges _ = pure ()
                   let serverUrl = Nothing
+                  let ucmVersion = Version.gitDescribeWithDate version
+                  credMan <- liftIO $ AuthN.newCredentialManager
+                  authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
                   startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
                   launch
                     version
@@ -215,6 +230,8 @@ main version = do
                     sbrt
                     theCodebase
                     [Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI]
+                    authenticatedHTTPClient
+                    credMan
                     serverUrl
                     (PP.toIds startProjectPath)
                     initRes
@@ -322,8 +339,10 @@ main version = do
               -- https://gitlab.haskell.org/ghc/ghc/-/merge_requests/1224
               void . Ki.fork scope $ LSP.spawnLsp lspFormattingConfig theCodebase runtime changeSignal
               let isTest = False
-              mcpServerConfig <-
-                MCP.initServer theCodebase runtime sbRuntime (pure currentDir) $ Version.gitDescribeWithDate version
+              let ucmVersion = Version.gitDescribeWithDate version
+              credMan <- liftIO $ AuthN.newCredentialManager
+              authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
+              mcpServerConfig <- MCP.initServer theCodebase runtime sbRuntime (Just currentDir) ucmVersion authenticatedHTTPClient
               Server.startServer
                 isTest
                 Backend.BackendEnv {Backend.useNamesIndex = False}
@@ -358,14 +377,15 @@ main version = do
                       takeMVar mvar
                     WithCLI -> do
                       PT.putPrettyLn $ P.string "Now starting the Unison Codebase Manager (UCM)..."
-
                       launch
                         version
                         currentDir
                         runtime
                         sbRuntime
                         theCodebase
                         []
+                        authenticatedHTTPClient
+                        credMan
                         mayBaseUrl
                         (PP.toIds startingProjectPath)
                         initRes
@@ -379,6 +399,9 @@ main version = do
       RTI.withRuntime False mode (Version.gitDescribeWithDate version) \runtime -> do
         RTI.withRuntime True mode (Version.gitDescribeWithDate version) \sbRuntime ->
           action (runtime, sbRuntime)
+    initTranscriptAuthenticatedHTTPClient :: UCMVersion -> AuthN.CredentialManager -> IO AuthN.AuthenticatedHttpClient
+    initTranscriptAuthenticatedHTTPClient ucmVersion credMan = do
+      AuthN.newAuthenticatedHTTPClient (AuthN.newTokenProvider credMan) ucmVersion
 
 isExitSuccess :: SomeException -> Bool
 isExitSuccess =
@@ -583,13 +606,15 @@ launch ::
   Rt.Runtime Symbol ->
   Codebase.Codebase IO Symbol Ann ->
   [Either Input.Event Input.Input] ->
+  AuthN.AuthenticatedHttpClient ->
+  AuthN.CredentialManager ->
   Maybe Server.BaseUrl ->
   PP.ProjectPathIds ->
   InitResult ->
   (PP.ProjectPathIds -> IO ()) ->
   CommandLine.ShouldWatchFiles ->
   IO ()
-launch version dir runtime sbRuntime codebase inputs serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
+launch version dir runtime sbRuntime codebase inputs authenticatedHTTPClient credMan serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
   showWelcomeHint <- Codebase.runTransaction codebase Queries.doProjectsExist
   let isNewCodebase = case initResult of
         CreatedCodebase -> NewlyCreatedCodebase
@@ -606,6 +631,8 @@ launch version dir runtime sbRuntime codebase inputs serverBaseUrl startingPath
         codebase
         serverBaseUrl
         ucmVersion
+        authenticatedHTTPClient
+        credMan
         lspCheckForChanges
         shouldWatchFiles
 
