Implements Livebook Agent authentication (#2403)

config/config.exs

@@ -25,6 +25,7 @@ config :mime, :types, %{
 
 config :livebook,
   teams_url: "https://teams.livebook.dev",
+  agent_name: nil,
   app_service_name: nil,
   app_service_url: nil,
   authentication_mode: :token,

config/test.exs

@@ -22,9 +22,10 @@ if File.exists?(data_path) do
   File.rm_rf!(data_path)
 end
 
-config :livebook, :data_path, data_path
+config :livebook,
-
+  data_path: data_path,
-config :livebook, :feature_flags, deployment_groups: true
+  feature_flags: [deployment_groups: true],
+  agent_name: "chonky-cat"
 
 # Use longnames when running tests in CI, so that no host resolution is required,
 # see https://github.com/livebook-dev/livebook/pull/173#issuecomment-819468549

lib/livebook.ex

@@ -229,6 +229,10 @@ defmodule Livebook do
     if dns_cluster_query = Livebook.Config.dns_cluster_query!("LIVEBOOK_CLUSTER") do
       config :livebook, :dns_cluster_query, dns_cluster_query
     end
+
+    if agent_name = Livebook.Config.agent_name!("LIVEBOOK_AGENT_NAME") do
+      config :livebook, :agent_name, agent_name
+    end
   end
 
   @doc """

lib/livebook/application.ex

@@ -261,8 +261,14 @@ defmodule Livebook.Application do
     cond do
       teams_key && auth ->
         case String.split(auth, ":") do
-          ["offline", name, public_key] -> create_offline_hub(teams_key, name, public_key)
+          ["offline", name, public_key] ->
-          _ -> Livebook.Config.abort!("Invalid LIVEBOOK_TEAMS_AUTH configuration.")
+            create_offline_hub(teams_key, name, public_key)
+
+          ["online", name, org_id, org_key_id, agent_key] ->
+            create_online_hub(teams_key, name, org_id, org_key_id, agent_key)
+
+          _ ->
+            Livebook.Config.abort!("Invalid LIVEBOOK_TEAMS_AUTH configuration.")
         end
 
       teams_key || auth ->
@@ -321,9 +327,9 @@ defmodule Livebook.Application do
       id: "team-#{name}",
       hub_name: name,
       hub_emoji: "⭐️",
-      user_id: 0,
+      user_id: nil,
-      org_id: 0,
+      org_id: nil,
-      org_key_id: 0,
+      org_key_id: nil,
       session_token: "",
       teams_key: teams_key,
       org_public_key: public_key,
@@ -334,6 +340,21 @@ defmodule Livebook.Application do
     })
   end
 
+  defp create_online_hub(teams_key, name, org_id, org_key_id, agent_key) do
+    Livebook.Hubs.save_hub(%Livebook.Hubs.Team{
+      id: "team-#{name}",
+      hub_name: name,
+      hub_emoji: "💡",
+      user_id: nil,
+      org_id: org_id,
+      org_key_id: org_key_id,
+      session_token: agent_key,
+      teams_key: teams_key,
+      org_public_key: nil,
+      offline: nil
+    })
+  end
+
   defp config_env_var?("LIVEBOOK_" <> _), do: true
   defp config_env_var?("RELEASE_" <> _), do: true
   defp config_env_var?("MIX_ENV"), do: true

lib/livebook/config.ex

@@ -220,6 +220,14 @@ defmodule Livebook.Config do
     Application.fetch_env!(:livebook, :teams_url)
   end
 
+  @doc """
+  Returns the configured name for the Livebook Agent session.
+  """
+  @spec agent_name() :: String.t()
+  def agent_name() do
+    Application.fetch_env!(:livebook, :agent_name)
+  end
+
   @doc """
   Returns if aws_credentials is enabled.
   """
@@ -573,6 +581,13 @@ defmodule Livebook.Config do
     System.get_env(env)
   end
 
+  @doc """
+  Parses agent name from env.
+  """
+  def agent_name!(env) do
+    System.get_env(env)
+  end
+
   @doc """
   Parses and validates default runtime from env.
   """

lib/livebook/hubs/personal.ex

@@ -230,7 +230,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Personal do
     :ok = Broadcasts.secret_deleted(secret)
   end
 
-  def connection_error(_personal), do: raise("not implemented")
+  def connection_status(_personal), do: raise("not implemented")
 
   def notebook_stamp(_hub, _notebook_source, metadata) when metadata == %{} do
     :skip

lib/livebook/hubs/provider.ex

@@ -80,10 +80,10 @@ defprotocol Livebook.Hubs.Provider do
   def delete_secret(hub, secret)
 
   @doc """
-  Gets the connection error from hub.
+  Gets the connection status from hub.
   """
-  @spec connection_error(t()) :: String.t() | nil
+  @spec connection_status(t()) :: String.t() | nil
-  def connection_error(hub)
+  def connection_status(hub)
 
   @doc """
   Generates a notebook stamp.

lib/livebook/hubs/team.ex

@@ -30,7 +30,7 @@ defmodule Livebook.Hubs.Team do
           offline: Offline.t() | nil
         }
 
-  @enforce_keys [:user_id, :org_id, :org_key_id, :session_token, :org_public_key, :teams_key]
+  @enforce_keys [:org_id, :org_key_id, :session_token, :teams_key]
 
   embedded_schema do
     field :org_id, :integer
@@ -140,12 +140,15 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
 
   def disconnect(team), do: TeamClient.stop(team.id)
 
-  def connection_error(team) do
+  def connection_status(team) do
     cond do
       team.offline ->
         "You are running an offline Hub for deployment. You cannot modify its settings."
 
-      reason = TeamClient.get_connection_error(team.id) ->
+      team.user_id == nil ->
+        "You are running a Livebook Agent instance with online Hub for deployment. You are in read-only mode."
+
+      reason = TeamClient.get_connection_status(team.id) ->
         "Cannot connect to Hub: #{reason}.\nWill attempt to reconnect automatically..."
 
       true ->

lib/livebook/hubs/team_client.ex

@@ -14,8 +14,9 @@ defmodule Livebook.Hubs.TeamClient do
 
   defstruct [
     :hub,
-    :connection_error,
+    :connection_status,
     :derived_key,
+    :deployment_group_id,
     connected?: false,
     secrets: [],
     file_systems: [],
@@ -61,11 +62,11 @@ defmodule Livebook.Hubs.TeamClient do
   end
 
   @doc """
-  Returns the latest error from connection.
+  Returns the latest status from connection.
   """
-  @spec get_connection_error(String.t()) :: String.t() | nil
+  @spec get_connection_status(String.t()) :: String.t() | nil
-  def get_connection_error(id) do
+  def get_connection_status(id) do
-    GenServer.call(registry_name(id), :get_connection_error)
+    GenServer.call(registry_name(id), :get_connection_status)
   catch
     :exit, _ -> "connection refused"
   end
@@ -102,13 +103,24 @@ defmodule Livebook.Hubs.TeamClient do
   def init(%Hubs.Team{offline: nil} = team) do
     derived_key = Teams.derive_key(team.teams_key)
 
-    headers = [
+    headers =
-      {"x-lb-version", Livebook.Config.app_version()},
+      if team.user_id do
-      {"x-user", to_string(team.user_id)},
+        [
-      {"x-org", to_string(team.org_id)},
+          {"x-lb-version", Livebook.Config.app_version()},
-      {"x-org-key", to_string(team.org_key_id)},
+          {"x-user", to_string(team.user_id)},
-      {"x-session-token", team.session_token}
+          {"x-org", to_string(team.org_id)},
-    ]
+          {"x-org-key", to_string(team.org_key_id)},
+          {"x-session-token", team.session_token}
+        ]
+      else
+        [
+          {"x-lb-version", Livebook.Config.app_version()},
+          {"x-org", to_string(team.org_id)},
+          {"x-org-key", to_string(team.org_key_id)},
+          {"x-agent-name", Livebook.Config.agent_name()},
+          {"x-agent-key", team.session_token}
+        ]
+      end
 
     {:ok, _pid} = Teams.Connection.start_link(self(), headers)
     {:ok, %__MODULE__{hub: team, derived_key: derived_key}}
@@ -127,18 +139,28 @@ defmodule Livebook.Hubs.TeamClient do
   end
 
   @impl true
-  def handle_call(:get_connection_error, _caller, state) do
+  def handle_call(:get_connection_status, _caller, state) do
-    {:reply, state.connection_error, state}
+    {:reply, state.connection_status, state}
   end
 
   def handle_call(:connected?, _caller, state) do
     {:reply, state.connected?, state}
   end
 
-  def handle_call(:get_secrets, _caller, state) do
+  def handle_call(:get_secrets, _caller, %{deployment_group_id: nil} = state) do
     {:reply, state.secrets, state}
   end
 
+  def handle_call(:get_secrets, _caller, state) do
+    case find_deployment_group(state) do
+      nil ->
+        {:reply, state.secrets, state}
+
+      %{secrets: agent_secrets} ->
+        {:reply, Enum.uniq_by(agent_secrets ++ state.secrets, & &1.name), state}
+    end
+  end
+
   def handle_call(:get_file_systems, _caller, state) do
     {:reply, state.file_systems, state}
   end
@@ -151,13 +173,13 @@ defmodule Livebook.Hubs.TeamClient do
   def handle_info(:connected, state) do
     Hubs.Broadcasts.hub_connected(state.hub.id)
 
-    {:noreply, %{state | connected?: true, connection_error: nil}}
+    {:noreply, %{state | connected?: true, connection_status: nil}}
   end
 
   def handle_info({:connection_error, reason}, state) do
     Hubs.Broadcasts.hub_connection_failed(state.hub.id, reason)
 
-    {:noreply, %{state | connected?: false, connection_error: reason}}
+    {:noreply, %{state | connected?: false, connection_status: reason}}
   end
 
   def handle_info({:server_error, reason}, state) do
@@ -360,11 +382,19 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
-  defp handle_event(:user_connected, user_connected, state) do
+  defp handle_event(:user_connected, connected, state) do
     state
-    |> dispatch_secrets(user_connected)
+    |> dispatch_secrets(connected)
-    |> dispatch_file_systems(user_connected)
+    |> dispatch_file_systems(connected)
-    |> dispatch_deployment_groups(user_connected)
+    |> dispatch_deployment_groups(connected)
+  end
+
+  defp handle_event(:agent_connected, agent_connected, state) do
+    %{state | deployment_group_id: to_string(agent_connected.deployment_group_id)}
+    |> update_hub(agent_connected)
+    |> dispatch_secrets(agent_connected)
+    |> dispatch_file_systems(agent_connected)
+    |> dispatch_deployment_groups(agent_connected)
   end
 
   defp dispatch_secrets(state, %{secrets: secrets}) do
@@ -418,6 +448,15 @@ defmodule Livebook.Hubs.TeamClient do
     )
   end
 
+  defp update_hub(state, %{public_key: org_public_key}) do
+    hub = %{state.hub | org_public_key: org_public_key}
+
+    # TODO: Fix this before merging
+    # ^hub = Hubs.save_hub(hub)
+
+    %{state | hub: hub}
+  end
+
   defp diff(old_list, new_list, fun, deleted_fun \\ nil, updated_fun \\ nil) do
     deleted_fun = unless deleted_fun, do: fun, else: deleted_fun
     updated_fun = unless updated_fun, do: fun, else: updated_fun
@@ -435,4 +474,10 @@ defmodule Livebook.Hubs.TeamClient do
         reduce: state,
         do: (acc -> handle_event(topic, event, acc))
   end
+
+  defp find_deployment_group(%{deployment_group_id: nil}),
+    do: nil
+
+  defp find_deployment_group(%{deployment_group_id: id, deployment_groups: groups}),
+    do: Enum.find(groups, &(&1.id == id))
 end

lib/livebook/teams/requests.ex

@@ -227,7 +227,12 @@ defmodule Livebook.Teams.Requests do
   end
 
   defp auth_headers(team) do
-    token = "#{team.user_id}:#{team.org_id}:#{team.org_key_id}:#{team.session_token}"
+    token =
+      if team.user_id do
+        "#{team.user_id}:#{team.org_id}:#{team.org_key_id}:#{team.session_token}"
+      else
+        "#{team.session_token}:#{Livebook.Config.agent_name()}:#{team.org_id}:#{team.org_key_id}"
+      end
 
     [
       {"x-lb-version", Livebook.Config.app_version()},
@@ -282,6 +287,10 @@ defmodule Livebook.Teams.Requests do
           do: {:error, Jason.decode!(body)},
           else: {:transport_error, body}
 
+      {:ok, 401, _headers, _body} ->
+        {:transport_error,
+         "You are not authorized to perform this action, make sure you have the access or you are not in a Livebook Agent instance"}
+
       _otherwise ->
         {:transport_error,
          "Something went wrong, try again later or please file a bug if it persists"}

lib/livebook_web/live/hub/edit/team_component.ex

@@ -65,11 +65,8 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
   def render(assigns) do
     ~H"""
     <div>
-      <LayoutHelpers.topbar
+      <LayoutHelpers.topbar :if={Provider.connection_status(@hub)} variant={:warning}>
-        :if={not @hub_metadata.connected? && Provider.connection_error(@hub)}
+        <%= Provider.connection_status(@hub) %>
-        variant={:warning}
-      >
-        <%= Provider.connection_error(@hub) %>
       </LayoutHelpers.topbar>
 
       <div class="p-4 md:px-12 md:py-7 max-w-screen-md mx-auto">

lib/livebook_web/live/hub/teams/deployment_group_live.ex

@@ -64,11 +64,8 @@ defmodule LivebookWeb.Hub.Teams.DeploymentGroupLive do
       saved_hubs={@saved_hubs}
     >
       <div>
-        <LayoutHelpers.topbar
+        <LayoutHelpers.topbar :if={Provider.connection_status(@hub)} variant={:warning}>
-          :if={not @hub_metadata.connected? && Provider.connection_error(@hub)}
+          <%= Provider.connection_status(@hub) %>
-          variant={:warning}
-        >
-          <%= Provider.connection_error(@hub) %>
         </LayoutHelpers.topbar>
 
         <div class="p-4 md:px-12 md:py-7 max-w-screen-md mx-auto">

lib/livebook_web/live/layout_helpers.ex

@@ -237,22 +237,22 @@ defmodule LivebookWeb.LayoutHelpers do
   defp sidebar_link_border_color(to, current) when to == current, do: "border-white"
   defp sidebar_link_border_color(_to, _current), do: "border-transparent"
 
-  defp hub_connection_link_opts(hub, to, current) do
+  defp hub_connection_link_opts(%{provider: hub}, to, current) do
     text_color = sidebar_link_text_color(to, current)
     border_color = sidebar_link_border_color(to, current)
 
     class =
       "h-7 flex items-center hover:text-white #{text_color} border-l-4 #{border_color} hover:border-white"
 
-    if hub.connected? do
+    if message = Provider.connection_status(hub) do
-      [id: "hub-#{hub.id}", navigate: to, class: class]
-    else
       [
         id: "hub-#{hub.id}",
         navigate: to,
-        "data-tooltip": Provider.connection_error(hub.provider),
+        "data-tooltip": message,
         class: "tooltip right " <> class
       ]
+    else
+      [id: "hub-#{hub.id}", navigate: to, class: class]
     end
   end
 

proto/lib/livebook_proto/agent_connected.pb.ex

@@ -0,0 +1,15 @@
+defmodule LivebookProto.AgentConnected do
+  use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.12.0"
+
+  field :id, 1, type: :int32
+  field :name, 2, type: :string
+  field :public_key, 3, type: :string, json_name: "publicKey"
+  field :deployment_group_id, 4, type: :int32, json_name: "deploymentGroupId"
+  field :secrets, 5, repeated: true, type: LivebookProto.Secret
+  field :file_systems, 6, repeated: true, type: LivebookProto.FileSystem, json_name: "fileSystems"
+
+  field :deployment_groups, 7,
+    repeated: true,
+    type: LivebookProto.DeploymentGroup,
+    json_name: "deploymentGroups"
+end

proto/lib/livebook_proto/deployment_group_secret.pb.ex

@@ -3,5 +3,5 @@ defmodule LivebookProto.DeploymentGroupSecret do
 
   field :name, 1, type: :string
   field :value, 2, type: :string
-  field :deployment_group_id, 3, type: :string
+  field :deployment_group_id, 3, type: :string, json_name: "deploymentGroupId"
 end

proto/lib/livebook_proto/event.pb.ex

@@ -1,7 +1,7 @@
 defmodule LivebookProto.Event do
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.12.0"
 
-  oneof(:type, 0)
+  oneof :type, 0
 
   field :secret_created, 1,
     type: LivebookProto.SecretCreated,
@@ -52,4 +52,9 @@ defmodule LivebookProto.Event do
     type: LivebookProto.DeploymentGroupDeleted,
     json_name: "deploymentGroupDeleted",
     oneof: 0
+
+  field :agent_connected, 11,
+    type: LivebookProto.AgentConnected,
+    json_name: "agentConnected",
+    oneof: 0
 end

proto/messages.proto

@@ -83,6 +83,17 @@ message UserConnected {
   string name = 1;
   repeated Secret secrets = 2;
   repeated FileSystem file_systems = 3;
+  repeated DeploymentGroup deployment_groups = 4;
+}
+
+message AgentConnected {
+  int32 id = 1;
+  string name = 2;
+  string public_key = 3;
+  int32 deployment_group_id = 4;
+  repeated Secret secrets = 5;
+  repeated FileSystem file_systems = 6;
+  repeated DeploymentGroup deployment_groups = 7;
 }
 
 message Event {
@@ -97,5 +108,6 @@ message Event {
     DeploymentGroupCreated deployment_group_created = 8;
     DeploymentGroupUpdated deployment_group_updated = 9;
     DeploymentGroupDeleted deployment_group_deleted = 10;
+    AgentConnected agent_connected = 11;
   }
 }

test/livebook/hubs/provider_test.exs

@@ -55,8 +55,8 @@ defmodule Livebook.Hubs.ProviderTest do
       refute secret in Provider.get_secrets(hub)
     end
 
-    test "connection_error/1", %{hub: hub} do
+    test "connection_status/1", %{hub: hub} do
-      assert_raise RuntimeError, "not implemented", fn -> Provider.connection_error(hub) end
+      assert_raise RuntimeError, "not implemented", fn -> Provider.connection_status(hub) end
     end
   end
 end

test/livebook_teams/hubs/team_client_test.exs

@@ -11,132 +11,68 @@ defmodule Livebook.Hubs.TeamClientTest do
     :ok
   end
 
-  describe "start_link/1" do
+  test "rejects the web socket connection with invalid credentials", %{user: user, token: token} do
-    test "successfully authenticates the web socket connection", %{user: user, node: node} do
+    team =
-      team = build_team_hub(user, node)
+      build(:team,
-      id = team.id
+        user_id: user.id,
+        org_id: 123_456,
+        org_key_id: 123_456,
+        session_token: token
+      )
 
-      refute TeamClient.connected?(team.id)
+    id = team.id
 
-      TeamClient.start_link(team)
+    TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
-      assert TeamClient.connected?(team.id)
-    end
 
-    test "rejects the web socket connection with invalid credentials", %{user: user, token: token} do
+    assert_receive {:hub_server_error, ^id, error}
-      team =
-        build(:team,
-          user_id: user.id,
-          org_id: 123_456,
-          org_key_id: 123_456,
-          session_token: token
-        )
 
-      id = team.id
+    assert error ==
+             "#{team.hub_name}: Your session is out-of-date. Please re-join the organization."
 
-      TeamClient.start_link(team)
+    refute Livebook.Hubs.hub_exists?(team.id)
-
-      assert_receive {:hub_server_error, ^id, error}
-
-      assert error ==
-               "#{team.hub_name}: Your session is out-of-date. Please re-join the organization."
-
-      refute Livebook.Hubs.hub_exists?(team.id)
-    end
   end
 
   describe "handle events" do
-    test "receives the secret_created event", %{user: user, node: node} do
+    setup %{user: user, node: node} do
       team = create_team_hub(user, node)
       id = team.id
 
       assert_receive {:hub_connected, ^id}
 
-      secret = build(:secret, name: "SECRET_CREATED_FOO", value: "BAR")
+      {:ok, team: team}
-      assert Livebook.Hubs.create_secret(team, secret) == :ok
-
-      name = secret.name
-      value = secret.value
-
-      # receives `{:event, :secret_created, secret_created}` event
-      # with the value decrypted
-      assert_receive {:secret_created, %{name: ^name, value: ^value}}
     end
 
-    test "receives the secret_updated event", %{user: user, node: node} do
+    test "receives the secret events", %{team: team} do
-      team = create_team_hub(user, node)
+      secret = build(:secret, name: "SECRET", value: "BAR")
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      secret = build(:secret, name: "SECRET_UPDATED_FOO", value: "BAR")
       assert Livebook.Hubs.create_secret(team, secret) == :ok
 
       name = secret.name
       value = secret.value
 
       # receives `{:secret_created, secret_created}` event
+      # with the value decrypted
       assert_receive {:secret_created, %{name: ^name, value: ^value}}
 
       # updates the secret
-      update_secret = Map.replace!(secret, :value, "BAZ")
+      updated_secret = Map.replace!(secret, :value, "BAZ")
-      assert Livebook.Hubs.update_secret(team, update_secret) == :ok
+      assert Livebook.Hubs.update_secret(team, updated_secret) == :ok
 
-      new_value = update_secret.value
+      new_value = updated_secret.value
 
       # receives `{:secret_updated, secret_updated}` event
       # with the value decrypted
       assert_receive {:secret_updated, %{name: ^name, value: ^new_value}}
-    end
-
-    test "receives the secret_deleted event", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      secret = build(:secret, name: "SECRET_DELETED_FOO", value: "BAR")
-      assert Livebook.Hubs.create_secret(team, secret) == :ok
-
-      name = secret.name
-      value = secret.value
-
-      # receives `{:secret_created, secret_created}` event
-      assert_receive {:secret_created, %{name: ^name, value: ^value}}
 
       # deletes the secret
-      assert Livebook.Hubs.delete_secret(team, secret) == :ok
+      assert Livebook.Hubs.delete_secret(team, updated_secret) == :ok
 
       # receives `{:secret_deleted, secret_deleted}` event
-      assert_receive {:secret_deleted, %{name: ^name, value: ^value}}
+      assert_receive {:secret_deleted, %{name: ^name, value: ^new_value}}
     end
 
-    test "receives the file_system_created event", %{user: user, node: node} do
+    test "receives the file system events", %{team: team} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      file_system = build(:fs_s3, bucket_url: "https://file_system_created.s3.amazonaws.com")
-      assert Livebook.Hubs.create_file_system(team, file_system) == :ok
-
-      bucket_url = file_system.bucket_url
-
-      # receives `{:event, :file_system_created, file_system_created}` event
-      assert_receive {:file_system_created, %{bucket_url: ^bucket_url}}
-    end
-
-    test "receives the file_system_updated event", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
       file_system =
-        build(:fs_s3,
+        build(:fs_s3, bucket_url: "https://file_system.s3.amazonaws.com", region: "us-east-1")
-          bucket_url: "https://file_system_updated.s3.amazonaws.com",
-          region: "us-east-1"
-        )
 
       assert Livebook.Hubs.create_file_system(team, file_system) == :ok
 
@@ -148,70 +84,31 @@ defmodule Livebook.Hubs.TeamClientTest do
                       %{external_id: id, bucket_url: ^bucket_url, region: ^region}}
 
       # updates the file system
-      update_file_system = %{file_system | region: "eu-central-1", external_id: id}
+      updated_file_system = %{file_system | region: "eu-central-1", external_id: id}
-      assert Livebook.Hubs.update_file_system(team, update_file_system) == :ok
+      assert Livebook.Hubs.update_file_system(team, updated_file_system) == :ok
 
-      new_region = update_file_system.region
+      new_region = updated_file_system.region
 
       # receives `{:file_system_updated, file_system_updated}` event
       assert_receive {:file_system_updated,
                       %{external_id: ^id, bucket_url: ^bucket_url, region: ^new_region}}
-    end
-
-    test "receives the file_system_deleted event", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      file_system = build(:fs_s3, bucket_url: "https://file_system_deleted.s3.amazonaws.com")
-      assert Livebook.Hubs.create_file_system(team, file_system) == :ok
-
-      bucket_url = file_system.bucket_url
-
-      # receives `{:file_system_created, file_system_created}` event
-      assert_receive {:file_system_created, %{external_id: id, bucket_url: ^bucket_url}}
 
       # deletes the file system
-      delete_file_system = %{file_system | external_id: id}
+      assert Livebook.Hubs.delete_file_system(team, updated_file_system) == :ok
-      assert Livebook.Hubs.delete_file_system(team, delete_file_system) == :ok
 
       # receives `{:file_system_deleted, file_system_deleted}` event
       assert_receive {:file_system_deleted, %{external_id: ^id, bucket_url: ^bucket_url}}
     end
 
-    test "receives the deployment_group_created event", %{user: user, node: node} do
+    test "receives the deployment group events", %{team: team} do
-      team = create_team_hub(user, node)
+      deployment_group = build(:deployment_group, name: "DEPLOYMENT_GROUP", mode: "online")
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      deployment_group =
-        build(:deployment_group, name: "DEPLOYMENT_GROUP_CREATED_FOO", mode: "online")
-
-      assert {:ok, _id} =
-               Livebook.Teams.create_deployment_group(team, deployment_group)
-
-      %{name: name, mode: mode} = deployment_group
-      # receives `{:event, :deployment_group_created, deployment_group_created}` event
-      assert_receive {:deployment_group_created, %{name: ^name, mode: ^mode}}
-    end
-
-    test "receives the deployment_group_updated event", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      deployment_group =
-        build(:deployment_group, name: "DEPLOYMENT_GROUP_UPDATED_FOO", mode: "offline")
 
       assert {:ok, id} =
                Livebook.Teams.create_deployment_group(team, deployment_group)
 
       %{name: name, mode: mode} = deployment_group
 
-      # receives `{:deployment_group_created, deployment_group_created}` event
+      # receives `{:event, :deployment_group_created, deployment_group_created}` event
       assert_receive {:deployment_group_created, %{name: ^name, mode: ^mode}}
 
       # updates the deployment group
@@ -227,170 +124,74 @@ defmodule Livebook.Hubs.TeamClientTest do
 
       # receives `{:deployment_group_updated, deployment_group_updated}` event
       assert_receive {:deployment_group_updated, %{name: ^name, mode: ^new_mode}}
-    end
-
-    test "receives the deployment_group_deleted event", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      id = team.id
-
-      assert_receive {:hub_connected, ^id}
-
-      deployment_group =
-        build(:deployment_group, name: "DEPLOYMENT_GROUP_DELETED_FOO", mode: "online")
-
-      assert {:ok, id} =
-               Livebook.Teams.create_deployment_group(team, deployment_group)
-
-      name = deployment_group.name
-      mode = deployment_group.mode
-
-      # receives `{:deployment_group_created, deployment_group_created}` event
-      assert_receive {:deployment_group_created, %{name: ^name, mode: ^mode}}
 
       # deletes the deployment group
-      assert Livebook.Teams.delete_deployment_group(team, %{
+      assert Livebook.Teams.delete_deployment_group(team, update_deployment_group) == :ok
-               deployment_group
-               | id: id
-             }) == :ok
 
       # receives `{:deployment_group_deleted, deployment_group_deleted}` event
-      assert_receive {:deployment_group_deleted, %{name: ^name, mode: ^mode}}
+      assert_receive {:deployment_group_deleted, %{name: ^name, mode: ^new_mode}}
     end
   end
 
-  describe "user connected event" do
+  describe "handle user_connected event" do
-    test "fills the secrets list", %{user: user, node: node} do
+    setup %{user: user, node: node} do
       team = build_team_hub(user, node)
-      id = team.id
 
+      user_connected =
+        LivebookProto.UserConnected.new!(
+          name: team.hub_name,
+          secrets: [],
+          file_systems: [],
+          deployment_groups: []
+        )
+
+      {:ok, team: team, user_connected: user_connected}
+    end
+
+    test "dispatches the secrets list", %{team: team, user_connected: user_connected} do
       secret =
         build(:secret,
-          name: "SECRET_CREATED",
+          name: "CHONKY_CAT",
           value: "an encrypted value",
-          hub_id: id
+          hub_id: team.id
         )
 
       secret_key = Livebook.Teams.derive_key(team.teams_key)
       secret_value = Livebook.Teams.encrypt(secret.value, secret_key)
       livebook_proto_secret = LivebookProto.Secret.new!(name: secret.name, value: secret_value)
 
-      user_connected =
+      # creates the secret
-        LivebookProto.UserConnected.new!(
+      user_connected = %{user_connected | secrets: [livebook_proto_secret]}
-          name: team.hub_name,
+      pid = connect_to_teams(team)
-          secrets: [livebook_proto_secret],
-          file_systems: []
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
       refute_receive {:secret_created, ^secret}
-
       send(pid, {:event, :user_connected, user_connected})
       assert_receive {:secret_created, ^secret}
       assert secret in TeamClient.get_secrets(team.id)
-    end
-
-    test "replaces the secret with updated value", %{user: user, node: node} do
-      team = build_team_hub(user, node)
-      id = team.id
-
-      secret =
-        build(:secret,
-          name: "SECRET_UPDATED",
-          value: "an encrypted value",
-          hub_id: id
-        )
-
-      secret_key = Livebook.Teams.derive_key(team.teams_key)
-      secret_value = Livebook.Teams.encrypt(secret.value, secret_key)
-      livebook_proto_secret = LivebookProto.Secret.new!(name: secret.name, value: secret_value)
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [livebook_proto_secret],
-          file_systems: []
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
-      refute_receive {:secret_created, ^secret}
-
-      send(pid, {:event, :user_connected, user_connected})
-      assert_receive {:secret_created, ^secret}
 
+      # updates the secret
       updated_secret = %{secret | value: "an updated value"}
       secret_value = Livebook.Teams.encrypt(updated_secret.value, secret_key)
-
+      updated_livebook_proto_secret = %{livebook_proto_secret | value: secret_value}
-      updated_livebook_proto_secret =
+      user_connected = %{user_connected | secrets: [updated_livebook_proto_secret]}
-        LivebookProto.Secret.new!(name: updated_secret.name, value: secret_value)
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [updated_livebook_proto_secret],
-          file_systems: []
-        )
-
       send(pid, {:event, :user_connected, user_connected})
       assert_receive {:secret_updated, ^updated_secret}
-
       refute secret in TeamClient.get_secrets(team.id)
       assert updated_secret in TeamClient.get_secrets(team.id)
+
+      # deletes the secret
+      user_connected = %{user_connected | secrets: []}
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:secret_deleted, ^updated_secret}
+      refute updated_secret in TeamClient.get_secrets(team.id)
     end
 
-    test "deletes the secret", %{user: user, node: node} do
+    test "dispatches the file systems list", %{team: team, user_connected: user_connected} do
-      team = build_team_hub(user, node)
-      id = team.id
-
-      secret =
-        build(:secret,
-          name: "SECRET_UPDATED",
-          value: "an encrypted value",
-          hub_id: id
-        )
-
-      secret_key = Livebook.Teams.derive_key(team.teams_key)
-      secret_value = Livebook.Teams.encrypt(secret.value, secret_key)
-      livebook_proto_secret = LivebookProto.Secret.new!(name: secret.name, value: secret_value)
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [livebook_proto_secret],
-          file_systems: []
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
-      refute_receive {:secret_created, ^secret}
-
-      send(pid, {:event, :user_connected, user_connected})
-      assert_receive {:secret_created, ^secret}
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [],
-          file_systems: []
-        )
-
-      send(pid, {:event, :user_connected, user_connected})
-      assert_receive {:secret_deleted, ^secret}
-
-      refute secret in TeamClient.get_secrets(team.id)
-    end
-
-    test "fills the file systems list", %{user: user, node: node} do
-      team = build_team_hub(user, node)
-      id = team.id
-
       bucket_url = "https://mybucket.s3.amazonaws.com"
       hash = :crypto.hash(:sha256, bucket_url)
-      fs_id = "#{id}-s3-#{Base.url_encode64(hash, padding: false)}"
+      fs_id = "#{team.id}-s3-#{Base.url_encode64(hash, padding: false)}"
 
       file_system =
-        build(:fs_s3, id: fs_id, bucket_url: bucket_url, external_id: "123456", hub_id: id)
+        build(:fs_s3, id: fs_id, bucket_url: bucket_url, external_id: "123456", hub_id: team.id)
 
       type = Livebook.FileSystems.type(file_system)
       %{name: name} = Livebook.FileSystem.external_metadata(file_system)
@@ -408,66 +209,18 @@ defmodule Livebook.Hubs.TeamClientTest do
           value: value
         )
 
-      user_connected =
+      # creates the file system
-        LivebookProto.UserConnected.new!(
+      user_connected = %{user_connected | file_systems: [livebook_proto_file_system]}
-          name: team.hub_name,
+      pid = connect_to_teams(team)
-          secrets: [],
-          file_systems: [livebook_proto_file_system]
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
       refute_receive {:file_system_created, ^file_system}
-
       send(pid, {:event, :user_connected, user_connected})
       assert_receive {:file_system_created, ^file_system}
       assert file_system in TeamClient.get_file_systems(team.id)
-    end
-
-    test "replaces the file system with updated value", %{user: user, node: node} do
-      team = build_team_hub(user, node)
-      id = team.id
-
-      bucket_url = "https://update_fs_994641.s3.amazonaws.com"
-      hash = :crypto.hash(:sha256, bucket_url)
-      fs_id = "#{id}-s3-#{Base.url_encode64(hash, padding: false)}"
-
-      file_system =
-        build(:fs_s3, id: fs_id, bucket_url: bucket_url, external_id: "994641", hub_id: id)
-
-      type = Livebook.FileSystems.type(file_system)
-      %{name: name} = Livebook.FileSystem.external_metadata(file_system)
-      attrs = Livebook.FileSystem.dump(file_system)
-      credentials = Jason.encode!(attrs)
-
-      secret_key = Livebook.Teams.derive_key(team.teams_key)
-      value = Livebook.Teams.encrypt(credentials, secret_key)
-
-      livebook_proto_file_system =
-        LivebookProto.FileSystem.new!(
-          id: file_system.external_id,
-          name: name,
-          type: to_string(type),
-          value: value
-        )
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [],
-          file_systems: [livebook_proto_file_system]
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
-      refute_receive {:file_system_created, ^file_system}
-
-      send(pid, {:event, :user_connected, user_connected})
-      assert_receive {:file_system_created, ^file_system}
 
+      # updates the file system
       updated_file_system = %{
         file_system
-        | id: "#{id}-s3-ATC52Lo7d-bS21OLjPQ8KFBPJN8ku4hCn2nic2jTGeI",
+        | id: "#{team.id}-s3-ATC52Lo7d-bS21OLjPQ8KFBPJN8ku4hCn2nic2jTGeI",
           bucket_url: "https://updated_name.s3.amazonaws.com"
       }
 
@@ -475,38 +228,140 @@ defmodule Livebook.Hubs.TeamClientTest do
       updated_credentials = Jason.encode!(updated_attrs)
       updated_value = Livebook.Teams.encrypt(updated_credentials, secret_key)
 
-      updated_livebook_proto_file_system =
+      updated_livebook_proto_file_system = %{
-        LivebookProto.FileSystem.new!(
+        livebook_proto_file_system
-          id: updated_file_system.external_id,
+        | name: updated_file_system.bucket_url,
-          name: updated_file_system.bucket_url,
-          type: to_string(type),
           value: updated_value
-        )
+      }
-
-      user_connected =
-        LivebookProto.UserConnected.new!(
-          name: team.hub_name,
-          secrets: [],
-          file_systems: [updated_livebook_proto_file_system]
-        )
 
+      user_connected = %{user_connected | file_systems: [updated_livebook_proto_file_system]}
       send(pid, {:event, :user_connected, user_connected})
       assert_receive {:file_system_updated, ^updated_file_system}
-
       refute file_system in TeamClient.get_file_systems(team.id)
       assert updated_file_system in TeamClient.get_file_systems(team.id)
+
+      # deletes the file system
+      user_connected = %{user_connected | file_systems: []}
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:file_system_deleted, ^updated_file_system}
+      refute updated_file_system in TeamClient.get_file_systems(team.id)
     end
 
-    test "deletes the file system", %{user: user, node: node} do
+    test "dispatches the deployment groups list", %{team: team, user_connected: user_connected} do
-      team = build_team_hub(user, node)
+      deployment_group =
-      id = team.id
+        build(:deployment_group,
+          id: "1",
+          name: "sleepy-cat-#{System.unique_integer([:positive])}",
+          mode: "offline",
+          hub_id: team.id,
+          secrets: []
+        )
 
-      bucket_url = "https://delete_fs_45465641.s3.amazonaws.com"
+      livebook_proto_deployment_group =
+        LivebookProto.DeploymentGroup.new!(
+          id: to_string(deployment_group.id),
+          name: deployment_group.name,
+          mode: to_string(deployment_group.mode),
+          secrets: []
+        )
+
+      # creates the deployment group
+      user_connected = %{user_connected | deployment_groups: [livebook_proto_deployment_group]}
+      pid = connect_to_teams(team)
+      refute_receive {:deployment_group_created, ^deployment_group}
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:deployment_group_created, ^deployment_group}
+      assert deployment_group in TeamClient.get_deployment_groups(team.id)
+
+      # updates the deployment group
+      updated_deployment_group = %{deployment_group | mode: "online"}
+
+      updated_livebook_proto_deployment_group = %{
+        livebook_proto_deployment_group
+        | mode: updated_deployment_group.mode
+      }
+
+      user_connected = %{
+        user_connected
+        | deployment_groups: [updated_livebook_proto_deployment_group]
+      }
+
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:deployment_group_updated, ^updated_deployment_group}
+      refute deployment_group in TeamClient.get_deployment_groups(team.id)
+      assert updated_deployment_group in TeamClient.get_deployment_groups(team.id)
+
+      # deletes the deployment group
+      user_connected = %{user_connected | deployment_groups: []}
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:deployment_group_deleted, ^updated_deployment_group}
+      refute updated_deployment_group in TeamClient.get_deployment_groups(team.id)
+    end
+  end
+
+  describe "handle agent_connected event" do
+    setup %{node: node} do
+      {agent_key, org, deployment_group, team} = build_agent_team_hub(node)
+      org_key_pair = erpc_call(node, :create_org_key_pair, [[org: org]])
+
+      agent_connected =
+        LivebookProto.AgentConnected.new!(
+          id: agent_key.id,
+          name: Livebook.Config.agent_name(),
+          public_key: org_key_pair.public_key,
+          deployment_group_id: deployment_group.id,
+          secrets: [],
+          file_systems: [],
+          deployment_groups: []
+        )
+
+      {:ok, team: team, deployment_group: deployment_group, agent_connected: agent_connected}
+    end
+
+    test "dispatches the secrets list", %{team: team, agent_connected: agent_connected} do
+      secret =
+        build(:secret,
+          name: "AGENT_SECRET",
+          value: "an encrypted value",
+          hub_id: team.id
+        )
+
+      secret_key = Livebook.Teams.derive_key(team.teams_key)
+      secret_value = Livebook.Teams.encrypt(secret.value, secret_key)
+      livebook_proto_secret = LivebookProto.Secret.new!(name: secret.name, value: secret_value)
+
+      # creates the secret
+      agent_connected = %{agent_connected | secrets: [livebook_proto_secret]}
+      pid = connect_to_teams(team)
+      refute_receive {:secret_created, ^secret}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:secret_created, ^secret}
+      assert secret in TeamClient.get_secrets(team.id)
+
+      # updates the secret
+      updated_secret = %{secret | value: "an updated value"}
+      secret_value = Livebook.Teams.encrypt(updated_secret.value, secret_key)
+      updated_livebook_proto_secret = %{livebook_proto_secret | value: secret_value}
+      agent_connected = %{agent_connected | secrets: [updated_livebook_proto_secret]}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:secret_updated, ^updated_secret}
+      refute secret in TeamClient.get_secrets(team.id)
+      assert updated_secret in TeamClient.get_secrets(team.id)
+
+      # deletes the secret
+      agent_connected = %{agent_connected | secrets: []}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:secret_deleted, ^updated_secret}
+      refute updated_secret in TeamClient.get_secrets(team.id)
+    end
+
+    test "dispatches the file systems list", %{team: team, agent_connected: agent_connected} do
+      bucket_url = "https://mybucket.s3.amazonaws.com"
       hash = :crypto.hash(:sha256, bucket_url)
-      fs_id = "#{id}-s3-#{Base.url_encode64(hash, padding: false)}"
+      fs_id = "#{team.id}-s3-#{Base.url_encode64(hash, padding: false)}"
 
       file_system =
-        build(:fs_s3, id: fs_id, bucket_url: bucket_url, external_id: "45465641", hub_id: id)
+        build(:fs_s3, id: fs_id, bucket_url: bucket_url, external_id: "123456", hub_id: team.id)
 
       type = Livebook.FileSystems.type(file_system)
       %{name: name} = Livebook.FileSystem.external_metadata(file_system)
@@ -524,31 +379,161 @@ defmodule Livebook.Hubs.TeamClientTest do
           value: value
         )
 
-      user_connected =
+      # creates the file system
-        LivebookProto.UserConnected.new!(
+      agent_connected = %{agent_connected | file_systems: [livebook_proto_file_system]}
-          name: team.hub_name,
+      pid = connect_to_teams(team)
-          secrets: [],
-          file_systems: [livebook_proto_file_system]
-        )
-
-      {:ok, pid} = TeamClient.start_link(team)
-      assert_receive {:hub_connected, ^id}
       refute_receive {:file_system_created, ^file_system}
-
+      send(pid, {:event, :agent_connected, agent_connected})
-      send(pid, {:event, :user_connected, user_connected})
       assert_receive {:file_system_created, ^file_system}
+      assert file_system in TeamClient.get_file_systems(team.id)
 
-      user_connected =
+      # updates the file system
-        LivebookProto.UserConnected.new!(
+      updated_file_system = %{
-          name: team.hub_name,
+        file_system
-          secrets: [],
+        | id: "#{team.id}-s3-ATC52Lo7d-bS21OLjPQ8KFBPJN8ku4hCn2nic2jTGeI",
-          file_systems: []
+          bucket_url: "https://updated_name.s3.amazonaws.com"
+      }
+
+      updated_attrs = Livebook.FileSystem.dump(updated_file_system)
+      updated_credentials = Jason.encode!(updated_attrs)
+      updated_value = Livebook.Teams.encrypt(updated_credentials, secret_key)
+
+      updated_livebook_proto_file_system = %{
+        livebook_proto_file_system
+        | name: updated_file_system.bucket_url,
+          value: updated_value
+      }
+
+      agent_connected = %{agent_connected | file_systems: [updated_livebook_proto_file_system]}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:file_system_updated, ^updated_file_system}
+      refute file_system in TeamClient.get_file_systems(team.id)
+      assert updated_file_system in TeamClient.get_file_systems(team.id)
+
+      # deletes the file system
+      agent_connected = %{agent_connected | file_systems: []}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:file_system_deleted, ^updated_file_system}
+      refute updated_file_system in TeamClient.get_file_systems(team.id)
+    end
+
+    test "dispatches the deployment groups list",
+         %{team: team, deployment_group: teams_deployment_group, agent_connected: agent_connected} do
+      deployment_group =
+        build(:deployment_group,
+          id: to_string(teams_deployment_group.id),
+          name: teams_deployment_group.name,
+          mode: to_string(teams_deployment_group.mode),
+          hub_id: team.id,
+          secrets: []
         )
 
-      send(pid, {:event, :user_connected, user_connected})
+      livebook_proto_deployment_group =
-      assert_receive {:file_system_deleted, ^file_system}
+        LivebookProto.DeploymentGroup.new!(
+          id: to_string(deployment_group.id),
+          name: deployment_group.name,
+          mode: to_string(deployment_group.mode),
+          secrets: []
+        )
 
-      refute file_system in TeamClient.get_file_systems(team.id)
+      # creates the deployment group
+      agent_connected = %{agent_connected | deployment_groups: [livebook_proto_deployment_group]}
+      pid = connect_to_teams(team)
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:deployment_group_created, ^deployment_group}
+      assert deployment_group in TeamClient.get_deployment_groups(team.id)
+
+      # updates the deployment group
+      updated_deployment_group = %{deployment_group | mode: "offline"}
+
+      updated_livebook_proto_deployment_group = %{
+        livebook_proto_deployment_group
+        | mode: updated_deployment_group.mode
+      }
+
+      agent_connected = %{
+        agent_connected
+        | deployment_groups: [updated_livebook_proto_deployment_group]
+      }
+
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:deployment_group_updated, ^updated_deployment_group}
+      refute deployment_group in TeamClient.get_deployment_groups(team.id)
+      assert updated_deployment_group in TeamClient.get_deployment_groups(team.id)
+
+      # deletes the deployment group
+      agent_connected = %{agent_connected | deployment_groups: []}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:deployment_group_deleted, ^updated_deployment_group}
+      refute updated_deployment_group in TeamClient.get_deployment_groups(team.id)
+    end
+
+    test "dispatches the secrets list and override with deployment group secret",
+         %{team: team, deployment_group: teams_deployment_group, agent_connected: agent_connected} do
+      secret =
+        build(:secret,
+          name: "ORG_SECRET",
+          value: "an encrypted value",
+          hub_id: team.id
+        )
+
+      secret_key = Livebook.Teams.derive_key(team.teams_key)
+      secret_value = Livebook.Teams.encrypt(secret.value, secret_key)
+      livebook_proto_secret = LivebookProto.Secret.new!(name: secret.name, value: secret_value)
+
+      # creates the secret
+      agent_connected = %{agent_connected | secrets: [livebook_proto_secret]}
+      pid = connect_to_teams(team)
+      refute_receive {:secret_created, ^secret}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:secret_created, ^secret}
+      assert secret in TeamClient.get_secrets(team.id)
+
+      # overrides the secret with deployment group secret
+      override_secret = %{
+        secret
+        | value: "an updated value",
+          deployment_group_id: teams_deployment_group.id
+      }
+
+      secret_value = Livebook.Teams.encrypt(override_secret.value, secret_key)
+
+      livebook_proto_deployment_group_secret =
+        LivebookProto.DeploymentGroupSecret.new!(
+          name: override_secret.name,
+          value: secret_value,
+          deployment_group_id: override_secret.deployment_group_id
+        )
+
+      deployment_group =
+        build(:deployment_group,
+          id: to_string(teams_deployment_group.id),
+          name: teams_deployment_group.name,
+          mode: to_string(teams_deployment_group.mode),
+          hub_id: team.id,
+          secrets: [override_secret]
+        )
+
+      livebook_proto_deployment_group =
+        LivebookProto.DeploymentGroup.new!(
+          id: to_string(deployment_group.id),
+          name: deployment_group.name,
+          mode: to_string(deployment_group.mode),
+          secrets: [livebook_proto_deployment_group_secret]
+        )
+
+      agent_connected = %{agent_connected | deployment_groups: [livebook_proto_deployment_group]}
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:deployment_group_created, ^deployment_group}
+      refute secret in TeamClient.get_secrets(team.id)
+      assert override_secret in TeamClient.get_secrets(team.id)
     end
   end
+
+  defp connect_to_teams(%{id: id} = team) do
+    {:ok, pid} = TeamClient.start_link(team)
+    assert_receive {:hub_connected, ^id}
+
+    pid
+  end
 end

test/livebook_teams/web/hub/deployment_group_live_test.exs

@@ -20,7 +20,7 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
   test "creates a deployment group", %{conn: conn, hub: hub} do
     deployment_group =
       build(:deployment_group,
-        name: "TEAM_ADD_DEPLOYMENT_GROUP",
+        name: "TEAMS_ADD_DEPLOYMENT_GROUP",
         mode: "offline",
         hub_id: hub.id
       )
@@ -49,11 +49,11 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     |> render_submit(attrs)
 
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{id: id, name: "TEAM_ADD_DEPLOYMENT_GROUP"} =
+                    %DeploymentGroup{id: id, name: "TEAMS_ADD_DEPLOYMENT_GROUP"} =
                       deployment_group}
 
     assert_patch(view, "/hub/#{hub.id}/deployment-groups/edit/#{id}")
-    assert render(view) =~ "Deployment group TEAM_ADD_DEPLOYMENT_GROUP added successfully"
+    assert render(view) =~ "Deployment group TEAMS_ADD_DEPLOYMENT_GROUP added successfully"
     assert deployment_group in Livebook.Teams.get_deployment_groups(hub)
 
     # Guarantee it shows the error from API
@@ -69,13 +69,13 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
   test "updates an existing deployment group", %{conn: conn, hub: hub} do
     insert_deployment_group(
-      name: "TEAM_EDIT_DEPLOYMENT_GROUP",
+      name: "TEAMS_EDIT_DEPLOYMENT_GROUP",
       mode: "online",
       hub_id: hub.id
     )
 
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAM_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
+                    %DeploymentGroup{name: "TEAMS_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
 
     attrs = %{
       deployment_group: %{
@@ -110,19 +110,24 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
     assert_receive {:deployment_group_updated, ^updated_deployment_group}
     assert_patch(view, "/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
-    assert render(view) =~ "Deployment group TEAM_EDIT_DEPLOYMENT_GROUP updated successfully"
+    assert render(view) =~ "Deployment group TEAMS_EDIT_DEPLOYMENT_GROUP updated successfully"
     assert updated_deployment_group in Livebook.Teams.get_deployment_groups(hub)
   end
 
   test "creates a secret", %{conn: conn, hub: hub} do
     insert_deployment_group(
-      name: "TEAM_EDIT_DEPLOYMENT_GROUP",
+      name: "TEAMS_EDIT_DEPLOYMENT_GROUP",
       mode: "online",
       hub_id: hub.id
     )
 
+    hub_id = hub.id
+
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAM_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
+                    %DeploymentGroup{name: "TEAMS_EDIT_DEPLOYMENT_GROUP", hub_id: ^hub_id} =
+                      deployment_group}
+
+    id = deployment_group.id
 
     {:ok, view, _html} =
       live(conn, ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
@@ -130,8 +135,8 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     secret =
       build(:secret,
         name: "DEPLOYMENT_GROUP_SECRET",
-        hub_id: hub.id,
+        hub_id: hub_id,
-        deployment_group_id: deployment_group.id
+        deployment_group_id: id
       )
 
     attrs = %{
@@ -151,7 +156,7 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
     assert_patch(
       view,
-      ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}/secrets/new"
+      ~p"/hub/#{hub_id}/deployment-groups/edit/#{id}/secrets/new"
     )
 
     assert render(view) =~ "Add secret"
@@ -169,9 +174,10 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     |> render_submit(attrs)
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{secrets: [^secret]} = deployment_group}
+                    %Livebook.Teams.DeploymentGroup{id: ^id, secrets: [^secret]} =
+                      deployment_group}
 
-    assert_patch(view, "/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
+    assert_patch(view, "/hub/#{hub_id}/deployment-groups/edit/#{id}")
     assert render(view) =~ "Secret DEPLOYMENT_GROUP_SECRET added successfully"
     assert render(element(view, "#deployment-group-secrets-list")) =~ secret.name
     assert secret in deployment_group.secrets
@@ -179,7 +185,7 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     # Guarantee it shows the error from API
 
     {:ok, view, _html} =
-      live(conn, ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}/secrets/new")
+      live(conn, ~p"/hub/#{hub_id}/deployment-groups/edit/#{id}/secrets/new")
 
     view
     |> element("#secrets-form")
@@ -190,26 +196,31 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
   test "updates an existing secret", %{conn: conn, hub: hub} do
     insert_deployment_group(
-      name: "TEAM_EDIT_DEPLOYMENT_GROUP",
+      name: "TEAMS_EDIT_DEPLOYMENT_GROUP",
       mode: "online",
       hub_id: hub.id
     )
 
+    hub_id = hub.id
+
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAM_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
+                    %DeploymentGroup{name: "TEAMS_EDIT_DEPLOYMENT_GROUP", hub_id: ^hub_id} =
+                      deployment_group}
+
+    id = deployment_group.id
 
     secret =
       insert_secret(
         name: "DEPLOYMENT_GROUP_EDIT_SECRET",
-        hub_id: hub.id,
+        hub_id: hub_id,
-        deployment_group_id: deployment_group.id
+        deployment_group_id: id
       )
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{secrets: [^secret]} = deployment_group}
+                    %Livebook.Teams.DeploymentGroup{id: ^id, secrets: [^secret]}}
 
     {:ok, view, _html} =
-      live(conn, ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
+      live(conn, ~p"/hub/#{hub_id}/deployment-groups/edit/#{id}")
 
     attrs = %{
       secret: %{
@@ -227,7 +238,7 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
     assert_patch(
       view,
-      ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}/secrets/edit/#{secret.name}"
+      ~p"/hub/#{hub_id}/deployment-groups/edit/#{id}/secrets/edit/#{secret.name}"
     )
 
     assert render(view) =~ "Edit secret"
@@ -247,9 +258,10 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     updated_secret = %{secret | value: new_value}
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{secrets: [^updated_secret]} = deployment_group}
+                    %Livebook.Teams.DeploymentGroup{id: ^id, secrets: [^updated_secret]} =
+                      deployment_group}
 
-    assert_patch(view, "/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
+    assert_patch(view, "/hub/#{hub_id}/deployment-groups/edit/#{id}")
     assert render(view) =~ "Secret DEPLOYMENT_GROUP_EDIT_SECRET updated successfully"
     assert render(element(view, "#deployment-group-secrets-list")) =~ secret.name
     assert updated_secret in deployment_group.secrets
@@ -257,26 +269,31 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
   test "deletes an existing secret", %{conn: conn, hub: hub} do
     insert_deployment_group(
-      name: "TEAM_EDIT_DEPLOYMENT_GROUP",
+      name: "TEAMS_EDIT_DEPLOYMENT_GROUP",
       mode: "online",
       hub_id: hub.id
     )
 
+    hub_id = hub.id
+
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAM_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
+                    %DeploymentGroup{name: "TEAMS_EDIT_DEPLOYMENT_GROUP", hub_id: ^hub_id} =
+                      deployment_group}
+
+    id = deployment_group.id
 
     secret =
       insert_secret(
         name: "DEPLOYMENT_GROUP_DELETE_SECRET",
-        hub_id: hub.id,
+        hub_id: hub_id,
-        deployment_group_id: deployment_group.id
+        deployment_group_id: id
       )
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{secrets: [^secret]} = deployment_group}
+                    %Livebook.Teams.DeploymentGroup{id: ^id, secrets: [^secret]}}
 
     {:ok, view, _html} =
-      live(conn, ~p"/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
+      live(conn, ~p"/hub/#{hub_id}/deployment-groups/edit/#{id}")
 
     refute view
            |> element("#secrets-form button[disabled]")
@@ -289,9 +306,9 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     render_confirm(view)
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{secrets: []} = deployment_group}
+                    %Livebook.Teams.DeploymentGroup{id: ^id, secrets: []} = deployment_group}
 
-    assert_patch(view, "/hub/#{hub.id}/deployment-groups/edit/#{deployment_group.id}")
+    assert_patch(view, "/hub/#{hub_id}/deployment-groups/edit/#{id}")
     assert render(view) =~ "Secret DEPLOYMENT_GROUP_DELETE_SECRET deleted successfully"
     refute render(element(view, "#deployment-group-secrets-list")) =~ secret.name
     refute secret in deployment_group.secrets

test/livebook_teams/web/hub/edit_live_test.exs

@@ -7,18 +7,18 @@ defmodule LivebookWeb.Integration.Hub.EditLiveTest do
   alias Livebook.Hubs
   alias Livebook.Teams.DeploymentGroup
 
-  setup %{user: user, node: node} do
+  describe "user" do
-    Livebook.Hubs.Broadcasts.subscribe([:crud, :connection, :secrets, :file_systems])
+    setup %{user: user, node: node} do
-    Livebook.Teams.Broadcasts.subscribe([:deployment_groups])
+      Livebook.Hubs.Broadcasts.subscribe([:crud, :connection, :secrets, :file_systems])
-    hub = create_team_hub(user, node)
+      Livebook.Teams.Broadcasts.subscribe([:deployment_groups])
-    id = hub.id
+      hub = create_team_hub(user, node)
+      id = hub.id
 
-    assert_receive {:hub_connected, ^id}
+      assert_receive {:hub_connected, ^id}
 
-    {:ok, hub: hub}
+      {:ok, hub: hub}
-  end
+    end
 
-  describe "team" do
     test "updates the hub", %{conn: conn, hub: hub} do
       {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
 
@@ -447,6 +447,99 @@ defmodule LivebookWeb.Integration.Hub.EditLiveTest do
     end
   end
 
+  describe "agent" do
+    setup %{node: node} do
+      Livebook.Hubs.Broadcasts.subscribe([:crud, :connection])
+      {agent_key, org, deployment_group, hub} = create_agent_team_hub(node)
+      id = hub.id
+
+      assert_receive {:hub_changed, ^id}
+      assert_receive {:hub_connected, ^id}
+
+      {:ok, hub: hub, agent_key: agent_key, org: org, deployment_group: deployment_group}
+    end
+
+    test "shows an error when creating a secret", %{conn: conn, hub: hub} do
+      {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
+      secret = build(:secret, name: "TEAM_ADD_SECRET", hub_id: hub.id)
+
+      attrs = %{
+        secret: %{
+          name: secret.name,
+          value: secret.value,
+          hub_id: secret.hub_id
+        }
+      }
+
+      refute render(view) =~ secret.name
+
+      view
+      |> element("#add-secret")
+      |> render_click(%{})
+
+      assert_patch(view, ~p"/hub/#{hub.id}/secrets/new")
+      assert render(view) =~ "Add secret"
+
+      view
+      |> element("#secrets-form")
+      |> render_change(attrs)
+
+      refute view
+             |> element("#secrets-form button[disabled]")
+             |> has_element?()
+
+      view
+      |> element("#secrets-form")
+      |> render_submit(attrs)
+
+      refute_receive {:secret_created, ^secret}
+
+      assert render(view) =~
+               "You are not authorized to perform this action, make sure you have the access or you are not in a Livebook Agent instance"
+
+      refute secret in Livebook.Hubs.get_secrets(hub)
+    end
+
+    test "shows an error when creating a file system", %{conn: conn, hub: hub} do
+      {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
+
+      bypass = Bypass.open()
+      file_system = build_bypass_file_system(bypass, hub.id)
+      id = file_system.id
+      attrs = %{file_system: Livebook.FileSystem.dump(file_system)}
+
+      expect_s3_listing(bypass)
+
+      refute render(view) =~ file_system.bucket_url
+
+      view
+      |> element("#add-file-system")
+      |> render_click(%{})
+
+      assert_patch(view, ~p"/hub/#{hub.id}/file-systems/new")
+      assert render(view) =~ "Add file storage"
+
+      view
+      |> element("#file-systems-form")
+      |> render_change(attrs)
+
+      refute view
+             |> element("#file-systems-form button[disabled]")
+             |> has_element?()
+
+      view
+      |> element("#file-systems-form")
+      |> render_submit(attrs)
+
+      refute_receive {:file_system_created, %{id: ^id}}
+
+      assert render(view) =~
+               "You are not authorized to perform this action, make sure you have the access or you are not in a Livebook Agent instance"
+
+      refute file_system in Livebook.Hubs.get_file_systems(hub)
+    end
+  end
+
   defp expect_s3_listing(bypass) do
     Bypass.expect_once(bypass, "GET", "/mybucket", fn conn ->
       conn

test/livebook_teams/web/session_live_test.exs

@@ -403,11 +403,12 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
         hub_id: team_id
       )
 
-      insert_deployment_group(
+      deployment_group =
-        name: "DEPLOYMENT_GROUP_TOBIAS",
+        insert_deployment_group(
-        mode: "online",
+          name: "DEPLOYMENT_GROUP_TOBIAS",
-        hub_id: team_id
+          mode: "online",
-      )
+          hub_id: team_id
+        )
 
       Session.subscribe(session.id)
       Session.set_notebook_hub(session.pid, team_id)
@@ -426,11 +427,13 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       assert render(view) =~ "Deployment Group"
       assert has_element?(view, "#select_deployment_group_form")
 
+      id = deployment_group.id
+
       view
-      |> form("#select_deployment_group_form", %{id: "2"})
+      |> form("#select_deployment_group_form", %{deployment_group_id: id})
       |> render_change()
 
-      assert_receive {:operation, {:set_notebook_deployment_group, _client, "2"}}
+      assert_receive {:operation, {:set_notebook_deployment_group, _client, ^id}}
     end
 
     @tag :tmp_dir

test/support/factory.ex

@@ -59,7 +59,10 @@ defmodule Livebook.Factory do
   def build(:deployment_group) do
     %Livebook.Teams.DeploymentGroup{
       name: "FOO",
-      mode: "offline"
+      mode: "offline",
+      clustering: "",
+      zta_key: "",
+      zta_provider: :""
     }
   end
 
@@ -107,8 +110,9 @@ defmodule Livebook.Factory do
   def insert_deployment_group(attrs \\ %{}) do
     deployment_group = build(:deployment_group, attrs)
     hub = Livebook.Hubs.fetch_hub!(deployment_group.hub_id)
-    {:ok, _id} = Livebook.Teams.create_deployment_group(hub, deployment_group)
+    {:ok, id} = Livebook.Teams.create_deployment_group(hub, deployment_group)
-    deployment_group
+
+    %{deployment_group | id: to_string(id)}
   end
 
   def insert_env_var(factory_name, attrs \\ %{}) do

test/support/hub_helpers.ex

@@ -13,9 +13,9 @@ defmodule Livebook.HubHelpers do
     teams_key: @offline_hub_key,
     org_public_key: @offline_hub_org_public_key,
     hub_name: @offline_hub_org_name,
-    user_id: 0,
+    user_id: nil,
-    org_id: 0,
+    org_id: nil,
-    org_key_id: 0,
+    org_key_id: nil,
     session_token: "",
     offline: %Livebook.Hubs.Team.Offline{
       secrets: []
@@ -27,6 +27,14 @@ defmodule Livebook.HubHelpers do
     Livebook.Hubs.save_hub(hub)
   end
 
+  def create_agent_team_hub(node) do
+    {agent_key, org, deployment_group, hub} = build_agent_team_hub(node)
+    erpc_call(node, :create_org_key_pair, [[org: org]])
+    ^hub = Livebook.Hubs.save_hub(hub)
+
+    {agent_key, org, deployment_group, hub}
+  end
+
   def build_team_headers(user, node) do
     hub = build_team_hub(user, node)
 
@@ -62,6 +70,40 @@ defmodule Livebook.HubHelpers do
     )
   end
 
+  def build_agent_team_hub(node) do
+    teams_org = build(:org)
+    teams_key = teams_org.teams_key
+    key_hash = Livebook.Teams.Org.key_hash(teams_org)
+
+    org = erpc_call(node, :create_org, [])
+    org_key = erpc_call(node, :create_org_key, [[org: org, key_hash: key_hash]])
+
+    deployment_group =
+      erpc_call(node, :create_deployment_group, [
+        [
+          name: "sleepy-cat-#{Ecto.UUID.generate()}",
+          mode: :online,
+          org: org
+        ]
+      ])
+
+    agent_key = erpc_call(node, :create_agent_key, [[deployment_group: deployment_group]])
+
+    team =
+      build(:team,
+        id: "team-#{org.name}",
+        hub_name: org.name,
+        user_id: nil,
+        org_id: org.id,
+        org_key_id: org_key.id,
+        org_public_key: nil,
+        session_token: agent_key.key,
+        teams_key: teams_key
+      )
+
+    {agent_key, org, deployment_group, team}
+  end
+
   def build_offline_team_hub(user, node) do
     teams_org = build(:org, teams_key: @offline_hub_key, name: @offline_hub_org_name)
     key_hash = Livebook.Teams.Org.key_hash(teams_org)
@@ -184,8 +226,8 @@ defmodule Livebook.HubHelpers do
     send(pid, {:event, :file_system_deleted, file_system_deleted})
   end
 
-  def create_teams_file_system(hub, node) do
+  def create_teams_file_system(hub, node, org_key \\ nil) do
-    org_key = erpc_call(node, :get_org_key!, [hub.org_key_id])
+    org_key = if org_key, do: org_key, else: erpc_call(node, :get_org_key!, [hub.org_key_id])
     erpc_call(node, :create_file_system, [[org_key: org_key]])
   end
 
@@ -208,6 +250,10 @@ defmodule Livebook.HubHelpers do
     :ok = Livebook.Hubs.create_file_system(hub, file_system)
   end
 
+  def erpc_call(node, fun, args) do
+    :erpc.call(node, TeamsRPC, fun, args)
+  end
+
   defp hub_pid(hub) do
     if pid = GenServer.whereis({:via, Registry, {Livebook.HubsRegistry, hub.id}}) do
       {:ok, pid}
@@ -215,8 +261,4 @@ defmodule Livebook.HubHelpers do
   end
 
   defp hub_element_id(id), do: "#hubs #hub-#{id}"
-
-  defp erpc_call(node, fun, args) do
-    :erpc.call(node, TeamsRPC, fun, args)
-  end
 end