Let Tower SDK use the local session when running in --local mode

crates/tower-cmd/src/run.rs

@@ -51,6 +51,10 @@ pub fn run_cmd() -> Command {
 pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMatches)>) {
     let res = get_run_parameters(args, cmd);
 
+    // We always expect there to be an environment due to the fact that there is a
+    // default value.
+    let env = args.get_one::<String>("environment").unwrap();
+
     match res {
         Ok((local, path, params, app_name)) => {
             debug!(
@@ -64,13 +68,9 @@ pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMa
                 if app_name.is_some() {
                     output::die("Running apps by name locally is not supported yet.");
                 } else {
-                    do_run_local(config, path, params).await;
+                    do_run_local(config, path, env, params).await;
                 }
             } else {
-                // We always expect there to be an environmnt due to the fact that there is a
-                // default value.
-                let env = args.get_one::<String>("environment").unwrap();
-
                 do_run_remote(config, path, env, params, app_name).await;
             }
         }
@@ -83,10 +83,7 @@ pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMa
 /// do_run_local is the entrypoint for running an app locally. It will load the Towerfile, build
 /// the package, and launch the app. The relevant package is cleaned up after execution is
 /// complete.
-async fn do_run_local(config: Config, path: PathBuf, mut params: HashMap<String, String>) {
-    // There is always an implicit `local` environment when running in a local context.
-    let env = "local".to_string();
-
+async fn do_run_local(config: Config, path: PathBuf, env: &str, mut params: HashMap<String, String>) {
     let mut spinner = output::spinner("Setting up runtime environment...");
 
     // Load all the secrets and catalogs from the server
@@ -104,6 +101,18 @@ async fn do_run_local(config: Config, path: PathBuf, mut params: HashMap<String,
 
     spinner.success();
 
+    // We prepare all the other misc environment variables that we need to inject
+    let mut env_vars = HashMap::new();
+    env_vars.extend(catalogs);
+    env_vars.insert("TOWER_URL".to_string(), config.tower_url.to_string());
+
+    // There should always be a session, if there isn't one then I'm not sure how we got here?
+    let session = config.session.unwrap_or_else(|| {
+        output::die("No session found. Please log in to Tower first.");
+    });
+
+    env_vars.insert("TOWER_JWT".to_string(), session.token.jwt.to_string());
+
     // Load the Towerfile
     let towerfile_path = path.join("Towerfile");
     let towerfile = load_towerfile(&towerfile_path);
@@ -131,7 +140,7 @@ async fn do_run_local(config: Config, path: PathBuf, mut params: HashMap<String,
 
     let mut launcher: AppLauncher<LocalApp> = AppLauncher::default();
     if let Err(err) = launcher
-        .launch(Context::new(), sender, package, env, secrets, params, catalogs)
+        .launch(Context::new(), sender, package, env.to_string(), secrets, params, env_vars)
         .await
     {
         output::runtime_error(err);

src/tower/_client.py

@@ -5,6 +5,7 @@
 
 from ._context import TowerContext
 from .exceptions import (
+    AppNotFoundError,
     NotFoundException,
     UnauthorizedException,
     UnknownException,
@@ -24,6 +25,7 @@
     RunAppResponse,
 )
 from .tower_api_client.models.error_model import ErrorModel
+from .tower_api_client.errors import UnexpectedStatus
 
 # WAIT_TIMEOUT is the amount of time to wait between requests when polling the
 # Tower API.
@@ -84,18 +86,24 @@ def run_app(
         parameters=run_params,
     )
 
-    output: Optional[Union[ErrorModel, RunAppResponse]] = run_app_api.sync(
-        slug=slug, client=client, body=input_body
-    )
+    try:
+        output: Optional[Union[ErrorModel, RunAppResponse]] = run_app_api.sync(
+            slug=slug, client=client, body=input_body
+        )
 
-    if output is None:
-        raise RuntimeError("Error running app")
-    else:
-        if isinstance(output, ErrorModel):
-            raise RuntimeError(f"Error running app: {output.title}")
+        if output is None:
+            raise RuntimeError("Error running app")
         else:
-            return output.run
-
+            if isinstance(output, ErrorModel):
+                raise RuntimeError(f"Error running app: {output.title}")
+            else:
+                return output.run
+    except UnexpectedStatus as e:
+        # Raise an AppNotFoundError here if the app was, indeed, not found.
+        if e.status_code == 404:
+            raise AppNotFoundError(slug)
+        else:
+            raise UnknownException(f"Unexpected status code {e.status_code} when running app {slug}")
 
 def wait_for_run(
     run: Run,
@@ -303,13 +311,23 @@ def _env_client(ctx: TowerContext, timeout: Optional[float] = None) -> Authentic
         else:
             tower_url += "/v1"
 
+    if ctx.jwt is not None:
+        token = ctx.jwt
+        auth_header_name = "Authorization"
+        prefix = "Bearer"
+    else:
+        token = ctx.api_key
+        auth_header_name = "X-API-Key"
+        prefix = ""
+
     return AuthenticatedClient(
         verify_ssl=False,
         base_url=tower_url,
-        token=ctx.api_key,
-        auth_header_name="X-API-Key",
-        prefix="",
+        token=token,
+        auth_header_name=auth_header_name,
+        prefix=prefix,
         timeout=timeout,
+        raise_on_unexpected_status=True,
     )
 
 

src/tower/_context.py

@@ -1,10 +1,11 @@
 import os
 
 class TowerContext:
-    def __init__(self, tower_url: str, environment: str, api_key: str = None, hugging_face_provider: str = None, hugging_face_api_key: str = None):
+    def __init__(self, tower_url: str, environment: str, api_key: str = None, hugging_face_provider: str = None, hugging_face_api_key: str = None, jwt: str = None):
         self.tower_url = tower_url
         self.environment = environment
         self.api_key = api_key
+        self.jwt = jwt
         self.hugging_face_provider = hugging_face_provider
         self.hugging_face_api_key = hugging_face_api_key
 
@@ -18,9 +19,13 @@ def is_local(self) -> bool:
 
     @classmethod
     def build(cls):
-        tower_url = os.getenv("TOWER_URL")
-        tower_environment = os.getenv("TOWER_ENVIRONMENT")
+        tower_url = os.getenv("TOWER_URL", "https://api.tower.dev")
+        tower_environment = os.getenv("TOWER_ENVIRONMENT", "default")
         tower_api_key = os.getenv("TOWER_API_KEY")
+        tower_jwt = os.getenv("TOWER_JWT")
+
+        # NOTE: These are experimental, used only for our experimental Hugging
+        # Face integration for LLMs.
         hugging_face_provider = os.getenv("TOWER_HUGGING_FACE_PROVIDER")
         hugging_face_api_key = os.getenv("TOWER_HUGGING_FACE_API_KEY")
 
@@ -30,5 +35,6 @@ def build(cls):
             api_key = tower_api_key,
             hugging_face_provider = hugging_face_provider,
             hugging_face_api_key = hugging_face_api_key,
+            jwt = tower_jwt,
         )
 

src/tower/exceptions.py

@@ -27,3 +27,7 @@ def __init__(self, time: float):
 class RunFailedError(RuntimeError):
     def __init__(self, app_name: str, number: int, state: str):
         super().__init__(f"Run {app_name}#{number} failed with status '{state}'")
+
+class AppNotFoundError(RuntimeError):
+    def __init__(self, app_name: str):
+        super().__init__(f"App '{app_name}' not found in the Tower.")

tests/tower/test_client.py

@@ -324,3 +324,67 @@ def test_raising_an_error_during_partial_failure(
     # Now actually wait for the runs
     with pytest.raises(RunFailedError) as excinfo:
         tower.wait_for_runs(runs, raise_on_failure=True)
+
+
+def test_raising_an_error_for_a_not_found_app(
+    httpx_mock, 
+    mock_api_config, 
+    mock_run_response_factory, 
+    create_run_object
+):
+    tower = mock_api_config
+
+    # Mock a 404 response with error model
+    error_response = {
+        "$schema": "https://api.tower.dev/v1/schemas/ErrorModel.json",
+        "status": 404,
+        "title": "Not Found",
+        "detail": "The requested app 'non-existent-app' was not found",
+        "instance": "https://api.example.com/v1/apps/non-existent-app",
+        "type": "about:blank"
+    }
+
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.example.com/v1/apps/non-existent-app/runs",
+        json=error_response,
+        status_code=404,
+    )
+
+    # Attempt to run a non-existent app and verify it raises the correct error
+    with pytest.raises(tower.exceptions.AppNotFoundError) as excinfo:
+        tower.run_app("non-existent-app")
+
+    assert "not found" in str(excinfo.value).lower()
+
+
+def test_raising_an_unexpected_error_based_on_status_code(
+    httpx_mock, 
+    mock_api_config, 
+    mock_run_response_factory, 
+    create_run_object
+):
+    tower = mock_api_config
+
+    # Mock a 404 response with error model
+    error_response = {
+        "$schema": "https://api.tower.dev/v1/schemas/ErrorModel.json",
+        "status": 404,
+        "title": "Not Found",
+        "detail": "The requested app 'non-existent-app' was not found",
+        "instance": "https://api.example.com/v1/apps/non-existent-app",
+        "type": "about:blank"
+    }
+
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.example.com/v1/apps/non-existent-app/runs",
+        json=error_response,
+        status_code=400,
+    )
+
+    # Attempt to run a non-existent app and verify it raises the correct error
+    with pytest.raises(tower.exceptions.UnknownException) as excinfo:
+        tower.run_app("non-existent-app")
+
+    assert "unexpected status code" in str(excinfo.value).lower()