From 5e818555449cf5c07c30255395cf42cc421ae18b Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Fri, 15 Nov 2024 14:48:54 +0100
Subject: [PATCH 01/33] [rfc-update-tasks-a6-a7] tmp

---
 tasks/agent6.py             | 15 +++++++++++++++
 tasks/libs/common/agent6.py | 27 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 tasks/agent6.py
 create mode 100644 tasks/libs/common/agent6.py

diff --git a/tasks/agent6.py b/tasks/agent6.py
new file mode 100644
index 0000000000000..fbb9bf2dc68b1
--- /dev/null
+++ b/tasks/agent6.py
@@ -0,0 +1,15 @@
+from invoke import task
+
+from tasks.libs.common.agent6 import prepare
+
+
+@task
+def init(ctx):
+    prepare(ctx)
+
+
+@task
+def run(ctx, command):
+    """Run a command in the agent 6 environment."""
+    prepare(ctx)
+    ctx.run(f"cd {ctx.agent6_worktree} && {command}")
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
new file mode 100644
index 0000000000000..66d02ddf546ae
--- /dev/null
+++ b/tasks/libs/common/agent6.py
@@ -0,0 +1,27 @@
+# Agent 6 / 7 compatibility utilities, used to execute tasks from Agent 7 (main) on Agent 6 (6.53.x)
+
+from pathlib import Path
+
+AGENT6_BRANCH = "6.53.x"
+AGENT6_WORKTREE = Path.cwd().parent / "datadog-agent6"
+
+
+def prepare(ctx):
+    """Will prepare the environment for agent6 commands.
+
+    To be used before each agent6 command.
+    Will:
+    1. Add the agent6 worktree if not present.
+    2. Fetch the latest changes from the agent6 worktree.
+    """
+
+    if not AGENT6_WORKTREE.is_dir():
+        ctx.run(f"git worktree add '{AGENT6_WORKTREE}' origin/{AGENT6_BRANCH}", warn=True)
+
+    ctx.run(f"git -C '{AGENT6_WORKTREE}' fetch origin {AGENT6_BRANCH}", warn=True)
+
+
+def is_agent6(ctx):
+    """Will return True if the current environment is an agent6 environment."""
+
+    return Path.cwd() == AGENT6_WORKTREE

From 88d022ccb4484fd34a6ce9c858d567e8a2ce981a Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Fri, 15 Nov 2024 15:07:15 +0100
Subject: [PATCH 02/33] [rfc-update-tasks-a6-a7] tmp

---
 tasks/__init__.py           |  2 ++
 tasks/libs/common/agent6.py | 13 ++++++++++---
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/tasks/__init__.py b/tasks/__init__.py
index 12753e9a67e68..80ed9058300ef 100644
--- a/tasks/__init__.py
+++ b/tasks/__init__.py
@@ -9,6 +9,7 @@
 
 from tasks import (
     agent,
+    agent6,
     ami,
     bench,
     buildimages,
@@ -210,6 +211,7 @@
 ns.add_collection(collector)
 ns.add_collection(invoke_unit_tests)
 ns.add_collection(debug)
+ns.add_collection(agent6)
 ns.configure(
     {
         "run": {
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index 66d02ddf546ae..e5b6f9a92fa74 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -1,4 +1,10 @@
-# Agent 6 / 7 compatibility utilities, used to execute tasks from Agent 7 (main) on Agent 6 (6.53.x)
+"""Agent 6 / 7 compatibility utilities, used to execute tasks from Agent 7 (main) on Agent 6 (6.53.x).
+
+Common environment variables that can be used:
+- AGENT6_NO_PULL: If set to any value, the agent6 worktree will not be pulled before running the command.
+"""
+
+import os
 
 from pathlib import Path
 
@@ -18,10 +24,11 @@ def prepare(ctx):
     if not AGENT6_WORKTREE.is_dir():
         ctx.run(f"git worktree add '{AGENT6_WORKTREE}' origin/{AGENT6_BRANCH}", warn=True)
 
-    ctx.run(f"git -C '{AGENT6_WORKTREE}' fetch origin {AGENT6_BRANCH}", warn=True)
+    if not os.environ.get("AGENT6_NO_PULL"):
+        ctx.run(f"git -C '{AGENT6_WORKTREE}' fetch origin {AGENT6_BRANCH}", warn=True)
 
 
-def is_agent6(ctx):
+def is_agent6():
     """Will return True if the current environment is an agent6 environment."""
 
     return Path.cwd() == AGENT6_WORKTREE

From b36cef896c34f116708b0ad4fec36a7be917d8f5 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 10:44:15 +0100
Subject: [PATCH 03/33] [rfc-update-tasks-a6-a7] Implemented context + agent 6
 switch for modules commands [skip ci]

---
 tasks/agent6.py                | 10 +++---
 tasks/libs/common/agent6.py    | 57 ++++++++++++++++++++++++++++++----
 tasks/libs/common/gomodules.py |  3 ++
 tasks/modules.py               | 41 +++++++++++++-----------
 4 files changed, 83 insertions(+), 28 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index fbb9bf2dc68b1..82ce98730a5ef 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -1,15 +1,17 @@
 from invoke import task
 
-from tasks.libs.common.agent6 import prepare
+from tasks.libs.common.agent6 import _agent6_context, prepare
 
 
 @task
 def init(ctx):
+    """Will prepare the agent 6 context (git clone / pull of the agent 6 branch)."""
+
     prepare(ctx)
 
 
 @task
 def run(ctx, command):
-    """Run a command in the agent 6 environment."""
-    prepare(ctx)
-    ctx.run(f"cd {ctx.agent6_worktree} && {command}")
+    """Runs a command in the agent 6 environment."""
+    with _agent6_context(ctx):
+        ctx.run(command)
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index e5b6f9a92fa74..89f0b00d956d2 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -5,11 +5,12 @@
 """
 
 import os
-
+from contextlib import contextmanager
 from pathlib import Path
 
 AGENT6_BRANCH = "6.53.x"
-AGENT6_WORKTREE = Path.cwd().parent / "datadog-agent6"
+AGENT6_WORKING_DIRECTORY = Path.cwd().parent / "datadog-agent6"
+AGENT7_WORKING_DIRECTORY = Path.cwd()
 
 
 def prepare(ctx):
@@ -21,14 +22,58 @@ def prepare(ctx):
     2. Fetch the latest changes from the agent6 worktree.
     """
 
-    if not AGENT6_WORKTREE.is_dir():
-        ctx.run(f"git worktree add '{AGENT6_WORKTREE}' origin/{AGENT6_BRANCH}", warn=True)
+    if not AGENT6_WORKING_DIRECTORY.is_dir():
+        ctx.run(f"git worktree add '{AGENT6_WORKING_DIRECTORY}' origin/{AGENT6_BRANCH}", warn=True)
 
     if not os.environ.get("AGENT6_NO_PULL"):
-        ctx.run(f"git -C '{AGENT6_WORKTREE}' fetch origin {AGENT6_BRANCH}", warn=True)
+        ctx.run(f"git -C '{AGENT6_WORKING_DIRECTORY}' fetch origin {AGENT6_BRANCH}", warn=True)
 
 
 def is_agent6():
     """Will return True if the current environment is an agent6 environment."""
 
-    return Path.cwd() == AGENT6_WORKTREE
+    return Path.cwd() == AGENT6_WORKING_DIRECTORY
+
+
+@contextmanager
+def _agent6_context(ctx):
+    """To run code from the agent6 environment.
+
+    Prefer using agent_context(ctx, version).
+    """
+
+    prepare(ctx)
+
+    current_dir = Path.cwd()
+
+    try:
+        # Enter
+        os.chdir(AGENT6_WORKING_DIRECTORY)
+
+        # TODO: Load modules etc.
+        yield
+    finally:
+        # Exit
+        os.chdir(current_dir)
+
+
+@contextmanager
+def agent_context(ctx, version: str | int | None):
+    """Runs code from the agent6 environment if the version is 6.
+
+    Usage:
+        > with agent_context(ctx, version):
+        >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target version
+    """
+
+    if version == 6 or isinstance(version, str) and version.startswith("6"):
+        with _agent6_context(ctx):
+            yield
+    else:
+        yield
+
+
+def agent_working_directory():
+    """Returns the working directory for the current context (agent 6 / 7)."""
+
+    return AGENT6_WORKING_DIRECTORY if is_agent6() else AGENT7_WORKING_DIRECTORY
diff --git a/tasks/libs/common/gomodules.py b/tasks/libs/common/gomodules.py
index ffbdd9f03d10d..d4bf14ca354a9 100644
--- a/tasks/libs/common/gomodules.py
+++ b/tasks/libs/common/gomodules.py
@@ -14,6 +14,7 @@
 import yaml
 
 import tasks
+from tasks.libs.common.agent6 import agent_working_directory
 
 
 class ConfigDumper(yaml.SafeDumper):
@@ -313,6 +314,8 @@ def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
         base_dir: Root directory of the agent repository ('.' by default).
     """
 
+    base_dir = base_dir or agent_working_directory()
+
     return Configuration.from_file(base_dir).modules
 
 
diff --git a/tasks/modules.py b/tasks/modules.py
index 9c1d8dbf24d71..53e88f1f26437 100644
--- a/tasks/modules.py
+++ b/tasks/modules.py
@@ -12,6 +12,7 @@
 import yaml
 from invoke import Context, Exit, task
 
+from tasks.libs.common.agent6 import agent_context
 from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.gomodules import (
     ConfigDumper,
@@ -260,41 +261,45 @@ def get_module_by_path(path: Path) -> GoModule | None:
 
 
 @task
-def show(_, path: str, remove_defaults: bool = False, base_dir: str = '.'):
+def show(ctx, path: str, remove_defaults: bool = False, base_dir: str = '.', version: int = None):
     """Show the module information for the given path.
 
     Args:
         remove_defaults: If True, will remove default values from the output.
+        version: If 6, will show agent 6 modules.
     """
 
-    config = Configuration.from_file(Path(base_dir))
-    if path in config.ignored_modules:
-        print(f'Module {path} is ignored')
-        return
+    with agent_context(ctx, version):
+        config = Configuration.from_file(Path(base_dir))
+        if path in config.ignored_modules:
+            print(f'Module {path} is ignored')
+            return
 
-    module = config.modules.get(path)
+        module = config.modules.get(path)
 
-    assert module, f'Module {path} not found'
+        assert module, f'Module {path} not found'
 
-    yaml.dump(
-        {path: module.to_dict(remove_defaults=remove_defaults, remove_path=True)}, sys.stdout, Dumper=ConfigDumper
-    )
+        yaml.dump(
+            {path: module.to_dict(remove_defaults=remove_defaults, remove_path=True)}, sys.stdout, Dumper=ConfigDumper
+        )
 
 
 @task
-def show_all(_, base_dir: str = '.', ignored=False):
+def show_all(ctx, base_dir: str = '.', ignored=False, version: int = None):
     """Show the list of modules.
 
     Args:
         ignored: If True, will list ignored modules.
+        version: If 6, will show agent 6 modules.
     """
 
-    config = Configuration.from_file(Path(base_dir))
+    with agent_context(ctx, version):
+        config = Configuration.from_file(Path(base_dir))
 
-    if ignored:
-        names = config.ignored_modules
-    else:
-        names = list(config.modules.keys())
+        if ignored:
+            names = config.ignored_modules
+        else:
+            names = list(config.modules.keys())
 
-    print('\n'.join(sorted(names)))
-    print(len(names), 'modules')
+        print('\n'.join(sorted(names)))
+        print(len(names), 'modules')

From 7cc0f50ec892ee7b9c67d55a8577d9eecd1963dc Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 13:59:37 +0100
Subject: [PATCH 04/33] [rfc-update-tasks-a6-a7] tmp

---
 tasks/agent6.py                               |  9 ++++++++
 tasks/github_tasks.py                         | 14 ++++++++---
 .../libs/ciproviders/github_actions_tools.py  |  6 +++--
 tasks/libs/ciproviders/gitlab_api.py          |  5 ++--
 tasks/libs/common/agent6.py                   | 23 +++++++++++++++----
 tasks/libs/common/constants.py                |  4 +++-
 tasks/libs/common/git.py                      | 15 +++++++++---
 tasks/libs/common/utils.py                    | 11 +++++----
 tasks/libs/notify/pipeline_status.py          |  6 ++---
 tasks/libs/pipeline/tools.py                  |  6 +++--
 tasks/libs/releasing/notes.py                 |  5 ++--
 tasks/linter.py                               |  6 ++---
 tasks/pipeline.py                             | 10 ++++----
 tasks/release.py                              |  6 ++---
 14 files changed, 88 insertions(+), 38 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index 82ce98730a5ef..a2182e56d4a25 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -15,3 +15,12 @@ def run(ctx, command):
     """Runs a command in the agent 6 environment."""
     with _agent6_context(ctx):
         ctx.run(command)
+
+
+@task
+def t(ctx, version=None):
+    from tasks.libs.common.agent6 import agent_context
+    from tasks.libs.common.git import get_default_branch
+
+    with agent_context(ctx, version):
+        print(get_default_branch())
diff --git a/tasks/github_tasks.py b/tasks/github_tasks.py
index b5723f232cd93..abe46fd6ff366 100644
--- a/tasks/github_tasks.py
+++ b/tasks/github_tasks.py
@@ -19,8 +19,9 @@
     trigger_macos_workflow,
 )
 from tasks.libs.common.color import color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH, DEFAULT_INTEGRATIONS_CORE_BRANCH
+from tasks.libs.common.constants import DEFAULT_INTEGRATIONS_CORE_BRANCH
 from tasks.libs.common.datadog_api import create_gauge, send_event, send_metrics
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.junit_upload_core import repack_macos_junit_tar
 from tasks.libs.common.utils import get_git_pretty_ref
 from tasks.libs.owners.linter import codeowner_has_orphans, directory_has_packages_without_owner
@@ -36,7 +37,7 @@ def concurrency_key():
     current_ref = get_git_pretty_ref()
 
     # We want workflows to run to completion on the default branch and release branches
-    if re.search(rf'^({DEFAULT_BRANCH}|\d+\.\d+\.x)$', current_ref):
+    if re.search(rf'^({get_default_branch()}|\d+\.\d+\.x)$', current_ref):
         return None
 
     return current_ref
@@ -68,7 +69,7 @@ def _trigger_macos_workflow(release, destination=None, retry_download=0, retry_i
 def trigger_macos(
     _,
     workflow_type="build",
-    datadog_agent_ref=DEFAULT_BRANCH,
+    datadog_agent_ref=None,
     release_version="nightly-a7",
     major_version="7",
     destination=".",
@@ -79,6 +80,13 @@ def trigger_macos(
     test_washer=False,
     integrations_core_ref=DEFAULT_INTEGRATIONS_CORE_BRANCH,
 ):
+    """
+    Args:
+        datadog_agent_ref: If None, will be the default branch.
+    """
+
+    datadog_agent_ref = datadog_agent_ref or get_default_branch()
+
     if workflow_type == "build":
         conclusion = _trigger_macos_workflow(
             # Provide the release version to be able to fetch the associated
diff --git a/tasks/libs/ciproviders/github_actions_tools.py b/tasks/libs/ciproviders/github_actions_tools.py
index 8bef945f46460..9399302271b5d 100644
--- a/tasks/libs/ciproviders/github_actions_tools.py
+++ b/tasks/libs/ciproviders/github_actions_tools.py
@@ -11,13 +11,13 @@
 
 from tasks.libs.ciproviders.github_api import GithubAPI
 from tasks.libs.common.color import color_message
-from tasks.libs.common.utils import DEFAULT_BRANCH
+from tasks.libs.common.git import get_default_branch
 
 
 def trigger_macos_workflow(
     workflow_name="macos.yaml",
     github_action_ref="master",
-    datadog_agent_ref=DEFAULT_BRANCH,
+    datadog_agent_ref=None,
     release_version=None,
     major_version=None,
     gitlab_pipeline_id=None,
@@ -31,6 +31,8 @@ def trigger_macos_workflow(
     """
     Trigger a workflow to build a MacOS Agent.
     """
+
+    datadog_agent_ref = datadog_agent_ref or get_default_branch()
     inputs = {}
 
     if datadog_agent_ref is not None:
diff --git a/tasks/libs/ciproviders/gitlab_api.py b/tasks/libs/ciproviders/gitlab_api.py
index 212e2ce38f142..0cf4c429b59b0 100644
--- a/tasks/libs/ciproviders/gitlab_api.py
+++ b/tasks/libs/ciproviders/gitlab_api.py
@@ -23,8 +23,7 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH
-from tasks.libs.common.git import get_common_ancestor, get_current_branch
+from tasks.libs.common.git import get_common_ancestor, get_current_branch, get_default_branch
 from tasks.libs.common.utils import retry_function
 
 BASE_URL = "https://gitlab.ddbuild.io"
@@ -1214,7 +1213,7 @@ def compute_gitlab_ci_config_diff(ctx, before: str, after: str):
     after_name = after or "local files"
 
     # The before commit is the LCA commit between before and after
-    before = before or DEFAULT_BRANCH
+    before = before or get_default_branch()
     before = get_common_ancestor(ctx, before, after or "HEAD")
 
     print(f'Getting after changes config ({color_message(after_name, Color.BOLD)})')
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index 89f0b00d956d2..4259391e46641 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -42,19 +42,21 @@ def _agent6_context(ctx):
     Prefer using agent_context(ctx, version).
     """
 
-    prepare(ctx)
+    # Do not stack two agent 6 contexts
+    if is_agent6():
+        yield
+        return
 
-    current_dir = Path.cwd()
+    prepare(ctx)
 
     try:
         # Enter
         os.chdir(AGENT6_WORKING_DIRECTORY)
 
-        # TODO: Load modules etc.
         yield
     finally:
         # Exit
-        os.chdir(current_dir)
+        os.chdir(AGENT7_WORKING_DIRECTORY)
 
 
 @contextmanager
@@ -66,10 +68,15 @@ def agent_context(ctx, version: str | int | None):
         >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target version
     """
 
-    if version == 6 or isinstance(version, str) and version.startswith("6"):
+    switch_agent6 = version == 6 or isinstance(version, str) and version.startswith("6")
+
+    if switch_agent6:
         with _agent6_context(ctx):
             yield
     else:
+        # NOTE: This ensures that we don't push agent 7 context from agent 6 context (context might be switched within inner functions)
+        assert not is_agent6(), 'Agent 7 context cannot be used within an agent 6 context'
+
         yield
 
 
@@ -77,3 +84,9 @@ def agent_working_directory():
     """Returns the working directory for the current context (agent 6 / 7)."""
 
     return AGENT6_WORKING_DIRECTORY if is_agent6() else AGENT7_WORKING_DIRECTORY
+
+
+def get_default_branch():
+    from tasks.libs.common.constants import _DEFAULT_BRANCH
+
+    return AGENT6_BRANCH if is_agent6() else _DEFAULT_BRANCH
diff --git a/tasks/libs/common/constants.py b/tasks/libs/common/constants.py
index ff9dc7a285a42..b197c1e19df78 100644
--- a/tasks/libs/common/constants.py
+++ b/tasks/libs/common/constants.py
@@ -1,4 +1,3 @@
-DEFAULT_BRANCH = "main"
 DEFAULT_INTEGRATIONS_CORE_BRANCH = "master"
 GITHUB_ORG = "DataDog"
 REPO_NAME = "datadog-agent"
@@ -9,6 +8,9 @@
 ALLOWED_REPO_ALL_BRANCHES = ALLOWED_REPO_NON_NIGHTLY_BRANCHES.union(ALLOWED_REPO_NIGHTLY_BRANCHES)
 AGENT_VERSION_CACHE_NAME = "agent-version.cache"
 
+# WARNING: Do not use this constant, prefer using `utils.get_default_branch()` instead. This is used to change default branch within agent 6 context
+_DEFAULT_BRANCH = "main"
+
 # Metric Origin Constants:
 # https://github.com/DataDog/dd-source/blob/a060ce7a403c2215c44ebfbcc588e42cd9985aeb/domains/metrics/shared/libs/proto/origin/origin.proto#L144
 ORIGIN_PRODUCT = 17
diff --git a/tasks/libs/common/git.py b/tasks/libs/common/git.py
index 61bac36427b21..2124efb8da813 100644
--- a/tasks/libs/common/git.py
+++ b/tasks/libs/common/git.py
@@ -8,7 +8,6 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH
 from tasks.libs.common.user_interactions import yes_no_question
 
 if TYPE_CHECKING:
@@ -99,7 +98,17 @@ def get_current_branch(ctx) -> str:
     return ctx.run("git rev-parse --abbrev-ref HEAD", hide=True).stdout.strip()
 
 
-def get_common_ancestor(ctx, branch, base=DEFAULT_BRANCH) -> str:
+def get_default_branch():
+    """Returns the default git branch given the current context (agent 6 / 7)."""
+
+    from tasks.libs.common import agent6
+
+    return agent6.get_default_branch()
+
+
+def get_common_ancestor(ctx, branch, base=None) -> str:
+    base = base or get_default_branch()
+
     return ctx.run(f"git merge-base {branch} {base}", hide=True).stdout.strip()
 
 
@@ -139,7 +148,7 @@ def check_base_branch(branch, release_version):
     Checks if the given branch is either the default branch or the release branch associated
     with the given release version.
     """
-    return branch == DEFAULT_BRANCH or branch == release_version.branch()
+    return branch == get_default_branch() or branch == release_version.branch()
 
 
 def try_git_command(ctx, git_command):
diff --git a/tasks/libs/common/utils.py b/tasks/libs/common/utils.py
index 10882ddfeed7d..4df825e84460a 100644
--- a/tasks/libs/common/utils.py
+++ b/tasks/libs/common/utils.py
@@ -24,8 +24,8 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import ALLOWED_REPO_ALL_BRANCHES, DEFAULT_BRANCH, REPO_PATH
-from tasks.libs.common.git import get_commit_sha
+from tasks.libs.common.constants import ALLOWED_REPO_ALL_BRANCHES, REPO_PATH
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.owners.parsing import search_owners
 from tasks.libs.releasing.version import get_version
 from tasks.libs.types.arch import Arch
@@ -373,6 +373,9 @@ def get_version_ldflags(ctx, major_version='7', install_path=None):
     Compute the version from the git tags, and set the appropriate compiler
     flags
     """
+
+    from tasks.libs.common.git import get_commit_sha
+
     payload_v = get_payload_version()
     commit = get_commit_sha(ctx, short=True)
 
@@ -494,8 +497,8 @@ def environ(env):
 
 
 def is_pr_context(branch, pr_id, test_name):
-    if branch == DEFAULT_BRANCH:
-        print(f"Running on {DEFAULT_BRANCH}, skipping check for {test_name}.")
+    if branch == get_default_branch():
+        print(f"Running on {get_default_branch()}, skipping check for {test_name}.")
         return False
     if not pr_id:
         print(f"PR not found, skipping check for {test_name}.")
diff --git a/tasks/libs/notify/pipeline_status.py b/tasks/libs/notify/pipeline_status.py
index 25915ccc36d82..569973decfb0b 100644
--- a/tasks/libs/notify/pipeline_status.py
+++ b/tasks/libs/notify/pipeline_status.py
@@ -2,7 +2,7 @@
 import re
 
 from tasks.libs.ciproviders.gitlab_api import get_commit, get_pipeline
-from tasks.libs.common.constants import DEFAULT_BRANCH
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.notify.utils import DEPLOY_PIPELINES_CHANNEL, PIPELINES_CHANNEL, PROJECT_NAME
 from tasks.libs.pipeline.data import get_failed_jobs
 from tasks.libs.pipeline.notifications import (
@@ -40,7 +40,7 @@ def send_message(ctx, notification_type, dry_run):
     # For deploy pipelines not on the main branch, send notifications in a
     # dedicated channel.
     slack_channel = PIPELINES_CHANNEL
-    if notification_type == "deploy" and pipeline.ref != DEFAULT_BRANCH:
+    if notification_type == "deploy" and pipeline.ref != get_default_branch():
         slack_channel = DEPLOY_PIPELINES_CHANNEL
 
     header = ""
@@ -64,7 +64,7 @@ def send_message(ctx, notification_type, dry_run):
     else:
         send_slack_message(slack_channel, str(message))
 
-    if should_send_message_to_author(pipeline.ref, DEFAULT_BRANCH):
+    if should_send_message_to_author(pipeline.ref, get_default_branch()):
         author_email = commit.author_email
         if dry_run:
             print(f"Would send to {author_email}:\n{str(message)}")
diff --git a/tasks/libs/pipeline/tools.py b/tasks/libs/pipeline/tools.py
index fc38ce67e883a..7f9d08e8eca6e 100644
--- a/tasks/libs/pipeline/tools.py
+++ b/tasks/libs/pipeline/tools.py
@@ -13,8 +13,8 @@
 
 from tasks.libs.ciproviders.gitlab_api import refresh_pipeline
 from tasks.libs.common.color import Color, color_message
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.user_interactions import yes_no_question
-from tasks.libs.common.utils import DEFAULT_BRANCH
 
 PIPELINE_FINISH_TIMEOUT_SEC = 3600 * 5
 
@@ -115,7 +115,7 @@ def gracefully_cancel_pipeline(repo: Project, pipeline: ProjectPipeline, force_c
 
 def trigger_agent_pipeline(
     repo: Project,
-    ref=DEFAULT_BRANCH,
+    ref=None,
     release_version_6="nightly",
     release_version_7="nightly-a7",
     branch="nightly",
@@ -134,6 +134,8 @@ def trigger_agent_pipeline(
     - run a pipeline with all end-to-end tests,
     - run a deploy pipeline (includes all builds & kitchen tests + uploads artifacts to staging repositories);
     """
+
+    ref = ref or get_default_branch()
     args = {}
 
     if deploy:
diff --git a/tasks/libs/releasing/notes.py b/tasks/libs/releasing/notes.py
index 10edb3dab145d..d831f372b6197 100644
--- a/tasks/libs/releasing/notes.py
+++ b/tasks/libs/releasing/notes.py
@@ -2,7 +2,8 @@
 
 from invoke import Failure
 
-from tasks.libs.common.constants import DEFAULT_BRANCH, GITHUB_REPO_NAME
+from tasks.libs.common.constants import GITHUB_REPO_NAME
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.releasing.version import current_version
 
 
@@ -42,7 +43,7 @@ def _add_dca_prelude(ctx, agent7_version, agent6_version=""):
             f"""prelude:
     |
     Released on: {date.today()}
-    Pinned to datadog-agent v{agent7_version}: `CHANGELOG <https://github.com/{GITHUB_REPO_NAME}/blob/{DEFAULT_BRANCH}/CHANGELOG.rst#{agent7_version.replace('.', '')}{agent6_version}>`_."""
+    Pinned to datadog-agent v{agent7_version}: `CHANGELOG <https://github.com/{GITHUB_REPO_NAME}/blob/{get_default_branch()}/CHANGELOG.rst#{agent7_version.replace('.', '')}{agent6_version}>`_."""
         )
 
     ctx.run(f"git add {new_releasenote}")
diff --git a/tasks/linter.py b/tasks/linter.py
index 082c3065cc666..f549ae2005075 100644
--- a/tasks/linter.py
+++ b/tasks/linter.py
@@ -32,8 +32,8 @@
 )
 from tasks.libs.common.check_tools_version import check_tools_version
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH, GITHUB_REPO_NAME
-from tasks.libs.common.git import get_file_modifications, get_staged_files
+from tasks.libs.common.constants import GITHUB_REPO_NAME
+from tasks.libs.common.git import get_default_branch, get_file_modifications, get_staged_files
 from tasks.libs.common.utils import gitlab_section, is_pr_context, running_in_ci
 from tasks.libs.owners.parsing import read_owners
 from tasks.libs.types.copyright import CopyrightLinter, LintFailure
@@ -52,7 +52,7 @@ def python(ctx):
 
     print(
         f"""Remember to set up pre-commit to lint your files before committing:
-    https://github.com/DataDog/datadog-agent/blob/{DEFAULT_BRANCH}/docs/dev/agent_dev_env.md#pre-commit-hooks"""
+    https://github.com/DataDog/datadog-agent/blob/{get_default_branch()}/docs/dev/agent_dev_env.md#pre-commit-hooks"""
     )
 
     if running_in_ci():
diff --git a/tasks/pipeline.py b/tasks/pipeline.py
index 31b728798c334..3f4396d2869e8 100644
--- a/tasks/pipeline.py
+++ b/tasks/pipeline.py
@@ -19,8 +19,7 @@
     refresh_pipeline,
 )
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH
-from tasks.libs.common.git import get_commit_sha, get_current_branch
+from tasks.libs.common.git import get_commit_sha, get_current_branch, get_default_branch
 from tasks.libs.common.utils import (
     get_all_allowed_repo_branches,
     is_allowed_repo_branch,
@@ -94,7 +93,7 @@ def check_deploy_pipeline(repo: Project, git_ref: str, release_version_6, releas
 
 
 @task
-def clean_running_pipelines(ctx, git_ref=DEFAULT_BRANCH, here=False, use_latest_sha=False, sha=None):
+def clean_running_pipelines(ctx, git_ref=None, here=False, use_latest_sha=False, sha=None):
     """
     Fetch running pipelines on a target ref (+ optionally a git sha), and ask the user if they
     should be cancelled.
@@ -104,6 +103,8 @@ def clean_running_pipelines(ctx, git_ref=DEFAULT_BRANCH, here=False, use_latest_
 
     if here:
         git_ref = get_current_branch(ctx)
+    else:
+        git_ref = git_ref or get_default_branch()
 
     print(f"Fetching running pipelines on {git_ref}")
 
@@ -130,11 +131,12 @@ def workflow_rules(gitlab_file=".gitlab-ci.yml"):
 
 
 @task
-def trigger(_, git_ref=DEFAULT_BRANCH, release_version_6="dev", release_version_7="dev-a7", repo_branch="dev"):
+def trigger(_, git_ref=None, release_version_6="dev", release_version_7="dev-a7", repo_branch="dev"):
     """
     OBSOLETE: Trigger a deploy pipeline on the given git ref. Use pipeline.run with the --deploy option instead.
     """
 
+    git_ref = git_ref or get_default_branch()
     use_release_entries = ""
     major_versions = []
 
diff --git a/tasks/release.py b/tasks/release.py
index a82deec90161f..0cfea137b13de 100644
--- a/tasks/release.py
+++ b/tasks/release.py
@@ -20,7 +20,6 @@
 from tasks.libs.ciproviders.gitlab_api import get_gitlab_repo
 from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.constants import (
-    DEFAULT_BRANCH,
     GITHUB_REPO_NAME,
 )
 from tasks.libs.common.git import (
@@ -28,6 +27,7 @@
     check_clean_branch_state,
     clone,
     get_current_branch,
+    get_default_branch,
     get_last_commit,
     get_last_release_tag,
     try_git_command,
@@ -371,7 +371,7 @@ def create_rc(ctx, major_versions="6,7", patch_version=False, upstream="origin",
     if not check_base_branch(current_branch, new_highest_version):
         raise Exit(
             color_message(
-                f"The branch you are on is neither {DEFAULT_BRANCH} or the correct release branch ({new_highest_version.branch()}). Aborting.",
+                f"The branch you are on is neither {get_default_branch()} or the correct release branch ({new_highest_version.branch()}). Aborting.",
                 "red",
             ),
             code=1,
@@ -472,7 +472,7 @@ def build_rc(ctx, major_versions="6,7", patch_version=False, k8s_deployments=Fal
     if not check_base_branch(current_branch, new_version):
         raise Exit(
             color_message(
-                f"The branch you are on is neither {DEFAULT_BRANCH} or the correct release branch ({new_version.branch()}). Aborting.",
+                f"The branch you are on is neither {get_default_branch()} or the correct release branch ({new_version.branch()}). Aborting.",
                 "red",
             ),
             code=1,

From 1f8ed02ce8199a304d994256c135d9c8f4b0a8f0 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 14:25:30 +0100
Subject: [PATCH 05/33] [rfc-update-tasks-a6-a7] Fixed refactoring [skip ci]

---
 tasks/libs/common/constants.py | 2 +-
 tasks/libs/common/utils.py     | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tasks/libs/common/constants.py b/tasks/libs/common/constants.py
index b197c1e19df78..656d4d4a97245 100644
--- a/tasks/libs/common/constants.py
+++ b/tasks/libs/common/constants.py
@@ -8,7 +8,7 @@
 ALLOWED_REPO_ALL_BRANCHES = ALLOWED_REPO_NON_NIGHTLY_BRANCHES.union(ALLOWED_REPO_NIGHTLY_BRANCHES)
 AGENT_VERSION_CACHE_NAME = "agent-version.cache"
 
-# WARNING: Do not use this constant, prefer using `utils.get_default_branch()` instead. This is used to change default branch within agent 6 context
+# WARNING: Do not use this constant, prefer using `get_default_branch()` instead. This is used to change default branch within agent 6 context
 _DEFAULT_BRANCH = "main"
 
 # Metric Origin Constants:
diff --git a/tasks/libs/common/utils.py b/tasks/libs/common/utils.py
index 4df825e84460a..0a569f396c52f 100644
--- a/tasks/libs/common/utils.py
+++ b/tasks/libs/common/utils.py
@@ -25,7 +25,7 @@
 
 from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.constants import ALLOWED_REPO_ALL_BRANCHES, REPO_PATH
-from tasks.libs.common.git import get_default_branch
+from tasks.libs.common.git import get_commit_sha, get_default_branch
 from tasks.libs.owners.parsing import search_owners
 from tasks.libs.releasing.version import get_version
 from tasks.libs.types.arch import Arch
@@ -374,8 +374,6 @@ def get_version_ldflags(ctx, major_version='7', install_path=None):
     flags
     """
 
-    from tasks.libs.common.git import get_commit_sha
-
     payload_v = get_payload_version()
     commit = get_commit_sha(ctx, short=True)
 

From e50350e816a4699c62fe490def9728980c279c12 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 14:40:46 +0100
Subject: [PATCH 06/33] [rfc-update-tasks-a6-a7] modules: Updated cache

---
 tasks/agent6.py                | 10 +--------
 tasks/libs/common/gomodules.py | 13 +++++++++++-
 tasks/modules.py               | 39 ++++++++++++++++------------------
 3 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index a2182e56d4a25..7a82543d7360a 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -13,14 +13,6 @@ def init(ctx):
 @task
 def run(ctx, command):
     """Runs a command in the agent 6 environment."""
+
     with _agent6_context(ctx):
         ctx.run(command)
-
-
-@task
-def t(ctx, version=None):
-    from tasks.libs.common.agent6 import agent_context
-    from tasks.libs.common.git import get_default_branch
-
-    with agent_context(ctx, version):
-        print(get_default_branch())
diff --git a/tasks/libs/common/gomodules.py b/tasks/libs/common/gomodules.py
index d4bf14ca354a9..24ae2cfca0e9d 100644
--- a/tasks/libs/common/gomodules.py
+++ b/tasks/libs/common/gomodules.py
@@ -307,6 +307,17 @@ def dependency_path(self, agent_version):
 
 
 @lru_cache
+def _get_default_modules_abs(base_dir_abs: Path) -> dict[str, GoModule]:
+    """Get default modules from absolute path.
+
+    Used to lru cache the result.
+    """
+
+    assert base_dir_abs.is_absolute()
+
+    return Configuration.from_file(base_dir_abs).modules
+
+
 def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
     """Load the default modules from the modules.yml file.
 
@@ -316,7 +327,7 @@ def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
 
     base_dir = base_dir or agent_working_directory()
 
-    return Configuration.from_file(base_dir).modules
+    return _get_default_modules_abs(base_dir.resolve())
 
 
 def validate_module(
diff --git a/tasks/modules.py b/tasks/modules.py
index 53e88f1f26437..d063eb50ae2d2 100644
--- a/tasks/modules.py
+++ b/tasks/modules.py
@@ -12,7 +12,6 @@
 import yaml
 from invoke import Context, Exit, task
 
-from tasks.libs.common.agent6 import agent_context
 from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.gomodules import (
     ConfigDumper,
@@ -261,7 +260,7 @@ def get_module_by_path(path: Path) -> GoModule | None:
 
 
 @task
-def show(ctx, path: str, remove_defaults: bool = False, base_dir: str = '.', version: int = None):
+def show(_, path: str, remove_defaults: bool = False, base_dir: str = '.'):
     """Show the module information for the given path.
 
     Args:
@@ -269,23 +268,22 @@ def show(ctx, path: str, remove_defaults: bool = False, base_dir: str = '.', ver
         version: If 6, will show agent 6 modules.
     """
 
-    with agent_context(ctx, version):
-        config = Configuration.from_file(Path(base_dir))
-        if path in config.ignored_modules:
-            print(f'Module {path} is ignored')
-            return
+    config = Configuration.from_file(Path(base_dir))
+    if path in config.ignored_modules:
+        print(f'Module {path} is ignored')
+        return
 
-        module = config.modules.get(path)
+    module = config.modules.get(path)
 
-        assert module, f'Module {path} not found'
+    assert module, f'Module {path} not found'
 
-        yaml.dump(
-            {path: module.to_dict(remove_defaults=remove_defaults, remove_path=True)}, sys.stdout, Dumper=ConfigDumper
-        )
+    yaml.dump(
+        {path: module.to_dict(remove_defaults=remove_defaults, remove_path=True)}, sys.stdout, Dumper=ConfigDumper
+    )
 
 
 @task
-def show_all(ctx, base_dir: str = '.', ignored=False, version: int = None):
+def show_all(_, base_dir: str = '.', ignored=False):
     """Show the list of modules.
 
     Args:
@@ -293,13 +291,12 @@ def show_all(ctx, base_dir: str = '.', ignored=False, version: int = None):
         version: If 6, will show agent 6 modules.
     """
 
-    with agent_context(ctx, version):
-        config = Configuration.from_file(Path(base_dir))
+    config = Configuration.from_file(Path(base_dir))
 
-        if ignored:
-            names = config.ignored_modules
-        else:
-            names = list(config.modules.keys())
+    if ignored:
+        names = config.ignored_modules
+    else:
+        names = list(config.modules.keys())
 
-        print('\n'.join(sorted(names)))
-        print(len(names), 'modules')
+    print('\n'.join(sorted(names)))
+    print(len(names), 'modules')

From 728a0c02ed5bbc1b689c6b32e021d71e8d57c128 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 14:53:47 +0100
Subject: [PATCH 07/33] [rfc-update-tasks-a6-a7] Added agent6.invoke

---
 tasks/agent6.py             | 14 +++++++++++++-
 tasks/libs/common/agent6.py | 20 ++++++++++++++++----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index 7a82543d7360a..d6256968d45de 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -1,6 +1,6 @@
 from invoke import task
 
-from tasks.libs.common.agent6 import _agent6_context, prepare
+from tasks.libs.common.agent6 import _agent6_context, enter_agent6_context, prepare
 
 
 @task
@@ -16,3 +16,15 @@ def run(ctx, command):
 
     with _agent6_context(ctx):
         ctx.run(command)
+
+
+@task
+def invoke(ctx):
+    """Enters the agent 6 environment in order to invoke tasks in this context.
+
+    Usage:
+        > inv agent6.invoke modules.show-all  # Will show agent 6 modules
+    """
+
+    # The tasks run after this one will be using the agent 6 environment
+    enter_agent6_context(ctx)
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index 4259391e46641..05fc9da3dd2fa 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -35,6 +35,20 @@ def is_agent6():
     return Path.cwd() == AGENT6_WORKING_DIRECTORY
 
 
+def enter_agent6_context(ctx):
+    """Enters the agent 6 environment."""
+
+    prepare(ctx)
+
+    os.chdir(AGENT6_WORKING_DIRECTORY)
+
+
+def exit_agent6_context():
+    """Exits the agent 6 environment."""
+
+    os.chdir(AGENT7_WORKING_DIRECTORY)
+
+
 @contextmanager
 def _agent6_context(ctx):
     """To run code from the agent6 environment.
@@ -47,16 +61,14 @@ def _agent6_context(ctx):
         yield
         return
 
-    prepare(ctx)
-
     try:
         # Enter
-        os.chdir(AGENT6_WORKING_DIRECTORY)
+        enter_agent6_context(ctx)
 
         yield
     finally:
         # Exit
-        os.chdir(AGENT7_WORKING_DIRECTORY)
+        exit_agent6_context()
 
 
 @contextmanager

From aa21a8e4a90aad202aca071283e8d0e4e1fdee63 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 15:12:25 +0100
Subject: [PATCH 08/33] [rfc-update-tasks-a6-a7] Updated all "main"

---
 tasks/git.py                         |  4 ++--
 tasks/libs/ciproviders/github_api.py |  3 ++-
 tasks/libs/common/git.py             |  6 ++++--
 tasks/libs/releasing/json.py         | 16 ++++++++++++++--
 tasks/release.py                     | 12 ++++++------
 tasks/setup.py                       |  5 +++--
 tasks/system_probe.py                |  6 ++++--
 7 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/tasks/git.py b/tasks/git.py
index 658c6615cd621..2574ff19c7091 100644
--- a/tasks/git.py
+++ b/tasks/git.py
@@ -4,14 +4,14 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import color_message
-from tasks.libs.common.git import get_current_branch
+from tasks.libs.common.git import get_current_branch, get_default_branch
 
 
 @task
 def check_protected_branch(ctx):
     local_branch = get_current_branch(ctx)
 
-    if local_branch == 'main':
+    if local_branch == get_default_branch():
         print(color_message("You're about to commit or push to the main, are you sure this is what you want?", "red"))
         raise Exit(code=1)
 
diff --git a/tasks/libs/ciproviders/github_api.py b/tasks/libs/ciproviders/github_api.py
index 2d61567f5016b..f8ed33e3a3252 100644
--- a/tasks/libs/ciproviders/github_api.py
+++ b/tasks/libs/ciproviders/github_api.py
@@ -13,6 +13,7 @@
 
 from tasks.libs.common.color import color_message
 from tasks.libs.common.constants import GITHUB_REPO_NAME
+from tasks.libs.common.git import get_default_branch
 
 try:
     import semver
@@ -494,7 +495,7 @@ def create_release_pr(title, base_branch, target_branch, version, changelog_pr=F
     ]
 
     if changelog_pr:
-        labels.append("backport/main")
+        labels.append(f"backport/{get_default_branch()}")
 
     updated_pr = github.update_pr(
         pull_number=pr.number,
diff --git a/tasks/libs/common/git.py b/tasks/libs/common/git.py
index 2124efb8da813..ad4a7eb2c0b37 100644
--- a/tasks/libs/common/git.py
+++ b/tasks/libs/common/git.py
@@ -88,7 +88,9 @@ def get_file_modifications(
     return modifications
 
 
-def get_modified_files(ctx, base_branch="main") -> list[str]:
+def get_modified_files(ctx, base_branch=None) -> list[str]:
+    base_branch = base_branch or get_default_branch()
+
     return get_file_modifications(
         ctx, base_branch=base_branch, added=True, modified=True, only_names=True, no_renames=True
     )
@@ -140,7 +142,7 @@ def get_main_parent_commit(ctx) -> str:
     """
     Get the commit sha your current branch originated from
     """
-    return ctx.run("git merge-base HEAD origin/main", hide=True).stdout.strip()
+    return ctx.run(f"git merge-base HEAD origin/{get_default_branch()}", hide=True).stdout.strip()
 
 
 def check_base_branch(branch, release_version):
diff --git a/tasks/libs/releasing/json.py b/tasks/libs/releasing/json.py
index e0ea089e43fa1..ed281acb907b4 100644
--- a/tasks/libs/releasing/json.py
+++ b/tasks/libs/releasing/json.py
@@ -6,7 +6,9 @@
 
 from invoke.exceptions import Exit
 
+from tasks.libs.common.agent6 import is_agent6
 from tasks.libs.common.constants import TAG_FOUND_TEMPLATE
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.releasing.documentation import _stringify_config, nightly_entry_for, release_entry_for
 from tasks.libs.releasing.version import (
     VERSION_RE,
@@ -41,6 +43,12 @@
     "datadog-agent-macos-build": "master",
     "datadog-agent": "main",
 }
+DEFAULT_BRANCHES_AGENT6 = {
+    "omnibus-software": "6.53.x",
+    "omnibus-ruby": "6.53.x",
+    "datadog-agent-macos-build": "6.53.x",
+    "datadog-agent": "6.53.x",
+}
 
 
 def load_release_json():
@@ -335,8 +343,12 @@ def generate_repo_data(warning_mode, next_version, release_branch):
     data = {}
     for repo in repos:
         branch = release_branch
-        if branch == "main":
-            branch = next_version.branch() if repo == "integrations-core" else DEFAULT_BRANCHES.get(repo, "main")
+        if branch == get_default_branch():
+            branch = (
+                next_version.branch()
+                if repo == "integrations-core"
+                else (DEFAULT_BRANCHES_AGENT6 if is_agent6() else DEFAULT_BRANCHES).get(repo, get_default_branch())
+            )
         data[repo] = {
             'branch': branch,
             'previous_tag': previous_tags.get(repo, ""),
diff --git a/tasks/release.py b/tasks/release.py
index 0cfea137b13de..d5278e9c4cf05 100644
--- a/tasks/release.py
+++ b/tasks/release.py
@@ -663,7 +663,7 @@ def create_release_branches(ctx, base_directory="~/dd", major_versions="6,7", up
 
         create_release_pr(
             f"[release] Update current milestone to {next}",
-            "main",
+            get_default_branch(),
             milestone_branch,
             next,
         )
@@ -683,8 +683,8 @@ def create_release_branches(ctx, base_directory="~/dd", major_versions="6,7", up
 
             with open(file, "w") as gl:
                 for line in file_content:
-                    if re.search(r"compare_to: main", line):
-                        gl.write(line.replace("main", f"{release_branch}"))
+                    if re.search(rf"compare_to: {get_default_branch()}", line):
+                        gl.write(line.replace(get_default_branch(), f"{release_branch}"))
                     else:
                         gl.write(line)
 
@@ -754,7 +754,7 @@ def cleanup(ctx):
     current_milestone = _update_last_stable(ctx, version)
 
     # create pull request to update last stable version
-    main_branch = "main"
+    main_branch = get_default_branch()
     cleanup_branch = f"release/{version}-cleanup"
     ctx.run(f"git checkout -b {cleanup_branch}")
     ctx.run("git add release.json")
@@ -787,7 +787,7 @@ def cleanup(ctx):
 @task
 def check_omnibus_branches(ctx):
     base_branch = _get_release_json_value('base_branch')
-    if base_branch == 'main':
+    if base_branch == get_default_branch():
         omnibus_ruby_branch = 'datadog-5.5.0'
         omnibus_software_branch = 'master'
     else:
@@ -913,7 +913,7 @@ def get_active_release_branch(_):
     if release_branch:
         print(f"{release_branch.name}")
     else:
-        print("main")
+        print(get_default_branch())
 
 
 @task
diff --git a/tasks/setup.py b/tasks/setup.py
index 0426543a16010..5e1eeee98d1f1 100644
--- a/tasks/setup.py
+++ b/tasks/setup.py
@@ -17,6 +17,7 @@
 
 from tasks import vscode
 from tasks.libs.common.color import Color, color_message
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.status import Status
 from tasks.libs.common.utils import running_in_pyapp
 
@@ -95,7 +96,7 @@ def check_git_repo(ctx) -> SetupResult:
     ctx.run("git fetch", hide=True)
 
     print(color_message("Checking main branch...", Color.BLUE))
-    output = ctx.run('git rev-list "^HEAD" origin/main --count', hide=True)
+    output = ctx.run(f'git rev-list "^HEAD" origin/{get_default_branch()} --count', hide=True)
     count = output.stdout.strip()
 
     message = ""
@@ -103,7 +104,7 @@ def check_git_repo(ctx) -> SetupResult:
 
     if count != "0":
         status = Status.WARN
-        message = f"Your branch is {count} commit(s) behind main. Please update your branch."
+        message = f"Your branch is {count} commit(s) behind {get_default_branch()}. Please update your branch."
 
     return SetupResult("Check git repository", status, message)
 
diff --git a/tasks/system_probe.py b/tasks/system_probe.py
index 2225b234951f7..a1754b6a9bc49 100644
--- a/tasks/system_probe.py
+++ b/tasks/system_probe.py
@@ -25,7 +25,7 @@
 from tasks.flavor import AgentFlavor
 from tasks.libs.build.ninja import NinjaWriter
 from tasks.libs.common.color import color_message
-from tasks.libs.common.git import get_commit_sha
+from tasks.libs.common.git import get_commit_sha, get_default_branch
 from tasks.libs.common.utils import (
     REPO_PATH,
     bin_name,
@@ -1778,12 +1778,14 @@ def generate_minimized_btfs(ctx, source_dir, output_dir, bpf_programs):
 
 
 @task
-def process_btfhub_archive(ctx, branch="main"):
+def process_btfhub_archive(ctx, branch=None):
     """
     process btfhub-archive repo to only select BTF tarball files of a single architecture
     :param ctx: invoke context
     :param branch: branch of DataDog/btfhub-archive to clone
     """
+
+    branch = branch or get_default_branch()
     output_dir = os.getcwd()
     with tempfile.TemporaryDirectory() as temp_dir:
         with ctx.cd(temp_dir):

From d90e347ff0589c89d53b8088c24563f94c3b2248 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 15:37:02 +0100
Subject: [PATCH 09/33] [rfc-update-tasks-a6-a7] Fixes, cleaning

---
 tasks/agent6.py   | 9 ++++++---
 tasks/pipeline.py | 2 +-
 tasks/release.py  | 5 +++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index d6256968d45de..a1bcd28a8ca2a 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -1,6 +1,6 @@
 from invoke import task
 
-from tasks.libs.common.agent6 import _agent6_context, enter_agent6_context, prepare
+from tasks.libs.common.agent6 import agent_context, enter_agent6_context, prepare
 
 
 @task
@@ -14,7 +14,7 @@ def init(ctx):
 def run(ctx, command):
     """Runs a command in the agent 6 environment."""
 
-    with _agent6_context(ctx):
+    with agent_context(ctx, 6):
         ctx.run(command)
 
 
@@ -22,9 +22,12 @@ def run(ctx, command):
 def invoke(ctx):
     """Enters the agent 6 environment in order to invoke tasks in this context.
 
+    Note:
+        This task should be avoided when a --version, --major-version or --agent-version argument is available in the task.
+
     Usage:
         > inv agent6.invoke modules.show-all  # Will show agent 6 modules
     """
 
-    # The tasks run after this one will be using the agent 6 environment
+    # The tasks running after this one will be using the agent 6 environment
     enter_agent6_context(ctx)
diff --git a/tasks/pipeline.py b/tasks/pipeline.py
index 3f4396d2869e8..1221bf6227658 100644
--- a/tasks/pipeline.py
+++ b/tasks/pipeline.py
@@ -833,7 +833,7 @@ def test_merge_queue(ctx):
     timestamp = int(datetime.now(timezone.utc).timestamp())
     test_main = f"mq/test_{timestamp}"
     current_branch = get_current_branch(ctx)
-    ctx.run("git checkout main", hide=True)
+    ctx.run(f"git checkout {get_default_branch()}", hide=True)
     ctx.run("git pull", hide=True)
     ctx.run(f"git checkout -b {test_main}", hide=True)
     ctx.run(f"git push origin {test_main}", hide=True)
diff --git a/tasks/release.py b/tasks/release.py
index d5278e9c4cf05..c1cea2eb9f8d2 100644
--- a/tasks/release.py
+++ b/tasks/release.py
@@ -47,6 +47,7 @@
     release_manager,
 )
 from tasks.libs.releasing.json import (
+    DEFAULT_BRANCHES,
     UNFREEZE_REPO_AGENT,
     UNFREEZE_REPOS,
     _get_release_json_value,
@@ -788,8 +789,8 @@ def cleanup(ctx):
 def check_omnibus_branches(ctx):
     base_branch = _get_release_json_value('base_branch')
     if base_branch == get_default_branch():
-        omnibus_ruby_branch = 'datadog-5.5.0'
-        omnibus_software_branch = 'master'
+        omnibus_ruby_branch = DEFAULT_BRANCHES['omnibus-ruby']
+        omnibus_software_branch = DEFAULT_BRANCHES['omnibus-software']
     else:
         omnibus_ruby_branch = base_branch
         omnibus_software_branch = base_branch

From f195804d174c346322ef14935c1dfa2e5c5bd2ff Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 15:51:20 +0100
Subject: [PATCH 10/33] [rfc-update-tasks-a6-a7] Added tests

---
 tasks/unit_tests/libs/common/agent6_tests.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 tasks/unit_tests/libs/common/agent6_tests.py

diff --git a/tasks/unit_tests/libs/common/agent6_tests.py b/tasks/unit_tests/libs/common/agent6_tests.py
new file mode 100644
index 0000000000000..ce249481837e8
--- /dev/null
+++ b/tasks/unit_tests/libs/common/agent6_tests.py
@@ -0,0 +1,19 @@
+import unittest
+
+from invoke import Context
+
+from tasks.libs.common.agent6 import agent_context, is_agent6
+
+
+def get_ctx():
+    return Context()
+
+
+class TestAgent6(unittest.TestCase):
+    def test_context_is_agent6_true(self):
+        with agent_context(get_ctx(), 6):
+            self.assertTrue(is_agent6())
+
+    def test_context_is_agent6_false(self):
+        with agent_context(get_ctx(), 7):
+            self.assertFalse(is_agent6())

From 63761562cb298ddcf0f01d7750cedc948db79a72 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 18 Nov 2024 16:49:41 +0100
Subject: [PATCH 11/33] [rfc-update-tasks-a6-a7] unit-tests: Added agent 6
 tests

---
 tasks/unit_tests/libs/common/agent6_tests.py | 49 +++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/tasks/unit_tests/libs/common/agent6_tests.py b/tasks/unit_tests/libs/common/agent6_tests.py
index ce249481837e8..df4b414648430 100644
--- a/tasks/unit_tests/libs/common/agent6_tests.py
+++ b/tasks/unit_tests/libs/common/agent6_tests.py
@@ -1,8 +1,11 @@
+import os
 import unittest
 
 from invoke import Context
 
-from tasks.libs.common.agent6 import agent_context, is_agent6
+from tasks.libs.common.agent6 import agent_context, is_agent6, prepare
+from tasks.libs.common.git import get_default_branch
+from tasks.libs.common.gomodules import get_default_modules
 
 
 def get_ctx():
@@ -10,6 +13,11 @@ def get_ctx():
 
 
 class TestAgent6(unittest.TestCase):
+    def setUp(self):
+        # Pull only once
+        prepare(get_ctx())
+        os.environ['AGENT6_NO_PULL'] = '1'
+
     def test_context_is_agent6_true(self):
         with agent_context(get_ctx(), 6):
             self.assertTrue(is_agent6())
@@ -17,3 +25,42 @@ def test_context_is_agent6_true(self):
     def test_context_is_agent6_false(self):
         with agent_context(get_ctx(), 7):
             self.assertFalse(is_agent6())
+
+    def test_context_nested(self):
+        with agent_context(get_ctx(), 6):
+            with agent_context(get_ctx(), 6):
+                self.assertTrue(is_agent6())
+            self.assertTrue(is_agent6())
+
+    def test_context_pwd(self):
+        ctx = get_ctx()
+
+        with agent_context(ctx, 7):
+            pwd7 = ctx.run('pwd').stdout
+
+        with agent_context(ctx, 6):
+            pwd6 = ctx.run('pwd').stdout
+
+        self.assertNotEqual(pwd6, pwd7)
+
+    def test_context_modules(self):
+        ctx = get_ctx()
+
+        with agent_context(ctx, 7):
+            modules7 = get_default_modules()
+
+        with agent_context(ctx, 6):
+            modules6 = get_default_modules()
+
+        self.assertNotEqual(set(modules6.keys()), set(modules7.keys()))
+
+    def test_context_branch(self):
+        ctx = get_ctx()
+
+        with agent_context(ctx, 7):
+            branch7 = get_default_branch()
+
+        with agent_context(ctx, 6):
+            branch6 = get_default_branch()
+
+        self.assertNotEqual(branch6, branch7)

From 1905845e8edf635d678964750f9c86dba3b04e29 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Tue, 19 Nov 2024 11:12:04 +0100
Subject: [PATCH 12/33] [rfc-update-tasks-a6-a7] Applied suggestions, moved
 functions, removed agent 6 context ctx mgnr

---
 tasks/libs/common/agent6.py    | 47 +++++++++-------------------------
 tasks/libs/common/constants.py |  3 ---
 tasks/libs/common/git.py       |  4 +--
 tasks/libs/common/gomodules.py |  2 +-
 tasks/libs/common/utils.py     |  8 ++++++
 5 files changed, 23 insertions(+), 41 deletions(-)

diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index 05fc9da3dd2fa..0a6e901ba58eb 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -49,28 +49,6 @@ def exit_agent6_context():
     os.chdir(AGENT7_WORKING_DIRECTORY)
 
 
-@contextmanager
-def _agent6_context(ctx):
-    """To run code from the agent6 environment.
-
-    Prefer using agent_context(ctx, version).
-    """
-
-    # Do not stack two agent 6 contexts
-    if is_agent6():
-        yield
-        return
-
-    try:
-        # Enter
-        enter_agent6_context(ctx)
-
-        yield
-    finally:
-        # Exit
-        exit_agent6_context()
-
-
 @contextmanager
 def agent_context(ctx, version: str | int | None):
     """Runs code from the agent6 environment if the version is 6.
@@ -83,22 +61,21 @@ def agent_context(ctx, version: str | int | None):
     switch_agent6 = version == 6 or isinstance(version, str) and version.startswith("6")
 
     if switch_agent6:
-        with _agent6_context(ctx):
+        # Do not stack two agent 6 contexts
+        if is_agent6():
             yield
+            return
+
+        try:
+            # Enter
+            enter_agent6_context(ctx)
+
+            yield
+        finally:
+            # Exit
+            exit_agent6_context()
     else:
         # NOTE: This ensures that we don't push agent 7 context from agent 6 context (context might be switched within inner functions)
         assert not is_agent6(), 'Agent 7 context cannot be used within an agent 6 context'
 
         yield
-
-
-def agent_working_directory():
-    """Returns the working directory for the current context (agent 6 / 7)."""
-
-    return AGENT6_WORKING_DIRECTORY if is_agent6() else AGENT7_WORKING_DIRECTORY
-
-
-def get_default_branch():
-    from tasks.libs.common.constants import _DEFAULT_BRANCH
-
-    return AGENT6_BRANCH if is_agent6() else _DEFAULT_BRANCH
diff --git a/tasks/libs/common/constants.py b/tasks/libs/common/constants.py
index 656d4d4a97245..0b100a3432263 100644
--- a/tasks/libs/common/constants.py
+++ b/tasks/libs/common/constants.py
@@ -8,9 +8,6 @@
 ALLOWED_REPO_ALL_BRANCHES = ALLOWED_REPO_NON_NIGHTLY_BRANCHES.union(ALLOWED_REPO_NIGHTLY_BRANCHES)
 AGENT_VERSION_CACHE_NAME = "agent-version.cache"
 
-# WARNING: Do not use this constant, prefer using `get_default_branch()` instead. This is used to change default branch within agent 6 context
-_DEFAULT_BRANCH = "main"
-
 # Metric Origin Constants:
 # https://github.com/DataDog/dd-source/blob/a060ce7a403c2215c44ebfbcc588e42cd9985aeb/domains/metrics/shared/libs/proto/origin/origin.proto#L144
 ORIGIN_PRODUCT = 17
diff --git a/tasks/libs/common/git.py b/tasks/libs/common/git.py
index ad4a7eb2c0b37..605c14b151d36 100644
--- a/tasks/libs/common/git.py
+++ b/tasks/libs/common/git.py
@@ -103,9 +103,9 @@ def get_current_branch(ctx) -> str:
 def get_default_branch():
     """Returns the default git branch given the current context (agent 6 / 7)."""
 
-    from tasks.libs.common import agent6
+    from tasks.libs.common.agent6 import AGENT6_BRANCH, is_agent6
 
-    return agent6.get_default_branch()
+    return AGENT6_BRANCH if is_agent6() else 'main'
 
 
 def get_common_ancestor(ctx, branch, base=None) -> str:
diff --git a/tasks/libs/common/gomodules.py b/tasks/libs/common/gomodules.py
index 24ae2cfca0e9d..dd381a5e9db68 100644
--- a/tasks/libs/common/gomodules.py
+++ b/tasks/libs/common/gomodules.py
@@ -14,7 +14,7 @@
 import yaml
 
 import tasks
-from tasks.libs.common.agent6 import agent_working_directory
+from tasks.libs.common.utils import agent_working_directory
 
 
 class ConfigDumper(yaml.SafeDumper):
diff --git a/tasks/libs/common/utils.py b/tasks/libs/common/utils.py
index 0a569f396c52f..168abaf0d5239 100644
--- a/tasks/libs/common/utils.py
+++ b/tasks/libs/common/utils.py
@@ -750,3 +750,11 @@ def get_metric_origin(origin_product, origin_sub_product, origin_product_detail,
     if origin_field:
         return {"origin": metric_origin}
     return metric_origin
+
+
+def agent_working_directory():
+    """Returns the working directory for the current context (agent 6 / 7)."""
+
+    from tasks.libs.common.agent6 import AGENT6_WORKING_DIRECTORY, AGENT7_WORKING_DIRECTORY, is_agent6
+
+    return AGENT6_WORKING_DIRECTORY if is_agent6() else AGENT7_WORKING_DIRECTORY

From 5c37de722446128669c9761f533ef70c2ce0e363 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Tue, 19 Nov 2024 11:14:39 +0100
Subject: [PATCH 13/33] [rfc-update-tasks-a6-a7] Removed system probe branch

---
 tasks/system_probe.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tasks/system_probe.py b/tasks/system_probe.py
index a1754b6a9bc49..2225b234951f7 100644
--- a/tasks/system_probe.py
+++ b/tasks/system_probe.py
@@ -25,7 +25,7 @@
 from tasks.flavor import AgentFlavor
 from tasks.libs.build.ninja import NinjaWriter
 from tasks.libs.common.color import color_message
-from tasks.libs.common.git import get_commit_sha, get_default_branch
+from tasks.libs.common.git import get_commit_sha
 from tasks.libs.common.utils import (
     REPO_PATH,
     bin_name,
@@ -1778,14 +1778,12 @@ def generate_minimized_btfs(ctx, source_dir, output_dir, bpf_programs):
 
 
 @task
-def process_btfhub_archive(ctx, branch=None):
+def process_btfhub_archive(ctx, branch="main"):
     """
     process btfhub-archive repo to only select BTF tarball files of a single architecture
     :param ctx: invoke context
     :param branch: branch of DataDog/btfhub-archive to clone
     """
-
-    branch = branch or get_default_branch()
     output_dir = os.getcwd()
     with tempfile.TemporaryDirectory() as temp_dir:
         with ctx.cd(temp_dir):

From 76a6bd192f30a51477b1721049777cdb14905cf1 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Wed, 20 Nov 2024 10:46:27 +0100
Subject: [PATCH 14/33] [rfc-update-tasks-a6-a7] agent6: Add remove env

---
 tasks/agent6.py                              | 11 +++++++++--
 tasks/libs/common/agent6.py                  | 13 ++++++++++---
 tasks/unit_tests/libs/common/agent6_tests.py |  4 ++--
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/tasks/agent6.py b/tasks/agent6.py
index a1bcd28a8ca2a..67ac69b838fa6 100644
--- a/tasks/agent6.py
+++ b/tasks/agent6.py
@@ -1,13 +1,20 @@
 from invoke import task
 
-from tasks.libs.common.agent6 import agent_context, enter_agent6_context, prepare
+from tasks.libs.common.agent6 import agent_context, enter_agent6_context, init_env, remove_env
 
 
 @task
 def init(ctx):
     """Will prepare the agent 6 context (git clone / pull of the agent 6 branch)."""
 
-    prepare(ctx)
+    init_env(ctx)
+
+
+@task
+def remove(ctx):
+    """Will remove all the context git submodules."""
+
+    remove_env(ctx)
 
 
 @task
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
index 0a6e901ba58eb..e21fb09cd9c51 100644
--- a/tasks/libs/common/agent6.py
+++ b/tasks/libs/common/agent6.py
@@ -13,7 +13,7 @@
 AGENT7_WORKING_DIRECTORY = Path.cwd()
 
 
-def prepare(ctx):
+def init_env(ctx):
     """Will prepare the environment for agent6 commands.
 
     To be used before each agent6 command.
@@ -26,7 +26,14 @@ def prepare(ctx):
         ctx.run(f"git worktree add '{AGENT6_WORKING_DIRECTORY}' origin/{AGENT6_BRANCH}", warn=True)
 
     if not os.environ.get("AGENT6_NO_PULL"):
-        ctx.run(f"git -C '{AGENT6_WORKING_DIRECTORY}' fetch origin {AGENT6_BRANCH}", warn=True)
+        ctx.run(f"git -C '{AGENT6_WORKING_DIRECTORY}' fetch origin {AGENT6_BRANCH}", warn=True, hide=True)
+
+
+def remove_env(ctx):
+    """Will remove the environment for agent6 commands."""
+
+    if AGENT6_WORKING_DIRECTORY.is_dir():
+        ctx.run(f"git worktree remove '{AGENT6_WORKING_DIRECTORY}'", warn=True)
 
 
 def is_agent6():
@@ -38,7 +45,7 @@ def is_agent6():
 def enter_agent6_context(ctx):
     """Enters the agent 6 environment."""
 
-    prepare(ctx)
+    init_env(ctx)
 
     os.chdir(AGENT6_WORKING_DIRECTORY)
 
diff --git a/tasks/unit_tests/libs/common/agent6_tests.py b/tasks/unit_tests/libs/common/agent6_tests.py
index df4b414648430..dfffcdaa127c7 100644
--- a/tasks/unit_tests/libs/common/agent6_tests.py
+++ b/tasks/unit_tests/libs/common/agent6_tests.py
@@ -3,7 +3,7 @@
 
 from invoke import Context
 
-from tasks.libs.common.agent6 import agent_context, is_agent6, prepare
+from tasks.libs.common.agent6 import agent_context, init_env, is_agent6
 from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.gomodules import get_default_modules
 
@@ -15,7 +15,7 @@ def get_ctx():
 class TestAgent6(unittest.TestCase):
     def setUp(self):
         # Pull only once
-        prepare(get_ctx())
+        init_env(get_ctx())
         os.environ['AGENT6_NO_PULL'] = '1'
 
     def test_context_is_agent6_true(self):

From 6c12470d9ac3efe9e547e81a8fd66b6089282c13 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 10:27:27 +0100
Subject: [PATCH 15/33] [rfc-update-tasks-a6-a7] agent 6 -> worktree

---
 tasks/__init__.py                            |  4 +-
 tasks/agent6.py                              | 40 ---------
 tasks/libs/common/agent6.py                  | 88 -------------------
 tasks/libs/common/git.py                     | 11 ++-
 tasks/libs/common/utils.py                   |  4 +-
 tasks/libs/common/worktree.py                | 92 ++++++++++++++++++++
 tasks/libs/releasing/json.py                 |  4 +-
 tasks/unit_tests/libs/common/agent6_tests.py | 45 +++++-----
 tasks/worktree.py                            | 44 ++++++++++
 9 files changed, 174 insertions(+), 158 deletions(-)
 delete mode 100644 tasks/agent6.py
 delete mode 100644 tasks/libs/common/agent6.py
 create mode 100644 tasks/libs/common/worktree.py
 create mode 100644 tasks/worktree.py

diff --git a/tasks/__init__.py b/tasks/__init__.py
index 80ed9058300ef..93ce84f1c37b2 100644
--- a/tasks/__init__.py
+++ b/tasks/__init__.py
@@ -9,7 +9,6 @@
 
 from tasks import (
     agent,
-    agent6,
     ami,
     bench,
     buildimages,
@@ -63,6 +62,7 @@
     trace_agent,
     vim,
     vscode,
+    worktree,
 )
 from tasks.build_tags import audit_tag_impact, print_default_build_tags
 from tasks.components import lint_components, lint_fxutil_oneshot_test
@@ -211,7 +211,7 @@
 ns.add_collection(collector)
 ns.add_collection(invoke_unit_tests)
 ns.add_collection(debug)
-ns.add_collection(agent6)
+ns.add_collection(worktree)
 ns.configure(
     {
         "run": {
diff --git a/tasks/agent6.py b/tasks/agent6.py
deleted file mode 100644
index 67ac69b838fa6..0000000000000
--- a/tasks/agent6.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from invoke import task
-
-from tasks.libs.common.agent6 import agent_context, enter_agent6_context, init_env, remove_env
-
-
-@task
-def init(ctx):
-    """Will prepare the agent 6 context (git clone / pull of the agent 6 branch)."""
-
-    init_env(ctx)
-
-
-@task
-def remove(ctx):
-    """Will remove all the context git submodules."""
-
-    remove_env(ctx)
-
-
-@task
-def run(ctx, command):
-    """Runs a command in the agent 6 environment."""
-
-    with agent_context(ctx, 6):
-        ctx.run(command)
-
-
-@task
-def invoke(ctx):
-    """Enters the agent 6 environment in order to invoke tasks in this context.
-
-    Note:
-        This task should be avoided when a --version, --major-version or --agent-version argument is available in the task.
-
-    Usage:
-        > inv agent6.invoke modules.show-all  # Will show agent 6 modules
-    """
-
-    # The tasks running after this one will be using the agent 6 environment
-    enter_agent6_context(ctx)
diff --git a/tasks/libs/common/agent6.py b/tasks/libs/common/agent6.py
deleted file mode 100644
index e21fb09cd9c51..0000000000000
--- a/tasks/libs/common/agent6.py
+++ /dev/null
@@ -1,88 +0,0 @@
-"""Agent 6 / 7 compatibility utilities, used to execute tasks from Agent 7 (main) on Agent 6 (6.53.x).
-
-Common environment variables that can be used:
-- AGENT6_NO_PULL: If set to any value, the agent6 worktree will not be pulled before running the command.
-"""
-
-import os
-from contextlib import contextmanager
-from pathlib import Path
-
-AGENT6_BRANCH = "6.53.x"
-AGENT6_WORKING_DIRECTORY = Path.cwd().parent / "datadog-agent6"
-AGENT7_WORKING_DIRECTORY = Path.cwd()
-
-
-def init_env(ctx):
-    """Will prepare the environment for agent6 commands.
-
-    To be used before each agent6 command.
-    Will:
-    1. Add the agent6 worktree if not present.
-    2. Fetch the latest changes from the agent6 worktree.
-    """
-
-    if not AGENT6_WORKING_DIRECTORY.is_dir():
-        ctx.run(f"git worktree add '{AGENT6_WORKING_DIRECTORY}' origin/{AGENT6_BRANCH}", warn=True)
-
-    if not os.environ.get("AGENT6_NO_PULL"):
-        ctx.run(f"git -C '{AGENT6_WORKING_DIRECTORY}' fetch origin {AGENT6_BRANCH}", warn=True, hide=True)
-
-
-def remove_env(ctx):
-    """Will remove the environment for agent6 commands."""
-
-    if AGENT6_WORKING_DIRECTORY.is_dir():
-        ctx.run(f"git worktree remove '{AGENT6_WORKING_DIRECTORY}'", warn=True)
-
-
-def is_agent6():
-    """Will return True if the current environment is an agent6 environment."""
-
-    return Path.cwd() == AGENT6_WORKING_DIRECTORY
-
-
-def enter_agent6_context(ctx):
-    """Enters the agent 6 environment."""
-
-    init_env(ctx)
-
-    os.chdir(AGENT6_WORKING_DIRECTORY)
-
-
-def exit_agent6_context():
-    """Exits the agent 6 environment."""
-
-    os.chdir(AGENT7_WORKING_DIRECTORY)
-
-
-@contextmanager
-def agent_context(ctx, version: str | int | None):
-    """Runs code from the agent6 environment if the version is 6.
-
-    Usage:
-        > with agent_context(ctx, version):
-        >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target version
-    """
-
-    switch_agent6 = version == 6 or isinstance(version, str) and version.startswith("6")
-
-    if switch_agent6:
-        # Do not stack two agent 6 contexts
-        if is_agent6():
-            yield
-            return
-
-        try:
-            # Enter
-            enter_agent6_context(ctx)
-
-            yield
-        finally:
-            # Exit
-            exit_agent6_context()
-    else:
-        # NOTE: This ensures that we don't push agent 7 context from agent 6 context (context might be switched within inner functions)
-        assert not is_agent6(), 'Agent 7 context cannot be used within an agent 6 context'
-
-        yield
diff --git a/tasks/libs/common/git.py b/tasks/libs/common/git.py
index 605c14b151d36..8f3d6d6065454 100644
--- a/tasks/libs/common/git.py
+++ b/tasks/libs/common/git.py
@@ -5,6 +5,7 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
+from invoke import Context
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
@@ -100,12 +101,18 @@ def get_current_branch(ctx) -> str:
     return ctx.run("git rev-parse --abbrev-ref HEAD", hide=True).stdout.strip()
 
 
+def is_agent6(ctx) -> bool:
+    return get_current_branch(ctx).startswith("6.")
+
+
 def get_default_branch():
     """Returns the default git branch given the current context (agent 6 / 7)."""
 
-    from tasks.libs.common.agent6 import AGENT6_BRANCH, is_agent6
+    # We create a context to avoid passing context in each function
+    # This context is used to get the current branch so there is no side effect
+    ctx = Context()
 
-    return AGENT6_BRANCH if is_agent6() else 'main'
+    return '6.53.x' if is_agent6(ctx) else 'main'
 
 
 def get_common_ancestor(ctx, branch, base=None) -> str:
diff --git a/tasks/libs/common/utils.py b/tasks/libs/common/utils.py
index 168abaf0d5239..44758ae35b543 100644
--- a/tasks/libs/common/utils.py
+++ b/tasks/libs/common/utils.py
@@ -755,6 +755,6 @@ def get_metric_origin(origin_product, origin_sub_product, origin_product_detail,
 def agent_working_directory():
     """Returns the working directory for the current context (agent 6 / 7)."""
 
-    from tasks.libs.common.agent6 import AGENT6_WORKING_DIRECTORY, AGENT7_WORKING_DIRECTORY, is_agent6
+    from tasks.libs.common.worktree import LOCAL_DIRECTORY, WORKTREE_DIRECTORY, is_worktree
 
-    return AGENT6_WORKING_DIRECTORY if is_agent6() else AGENT7_WORKING_DIRECTORY
+    return WORKTREE_DIRECTORY if is_worktree() else LOCAL_DIRECTORY
diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
new file mode 100644
index 0000000000000..11d758d64954f
--- /dev/null
+++ b/tasks/libs/common/worktree.py
@@ -0,0 +1,92 @@
+"""Worktree utilities, used to execute tasks from this local repository (main) to a worktree with a different HEAD (e.g. 6.53.x).
+
+Common environment variables that can be used:
+- WORKTREE_NO_PULL: If set to any value, the worktree will not be pulled before running the command.
+"""
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+
+WORKTREE_DIRECTORY = Path.cwd().parent / "datadog-agent-worktree"
+LOCAL_DIRECTORY = Path.cwd().resolve()
+
+
+def init_env(ctx, branch: str | None = None):
+    """Will prepare the environment for commands applying to a worktree.
+
+    To be used before each worktree section.
+    Will:
+    1. Add the agent worktree if not present.
+    2. Fetch the latest changes from the agent worktree.
+    """
+
+    if not WORKTREE_DIRECTORY.is_dir():
+        ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}")
+
+    if branch:
+        worktree_branch = ctx.run(
+            f"git -C '{WORKTREE_DIRECTORY}' rev-parse --abbrev-ref HEAD", hide=True
+        ).stdout.strip()
+        if worktree_branch != branch:
+            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' switch '{branch}'", hide=True)
+
+        if not os.environ.get("AGENT_WORKTREE_NO_PULL"):
+            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' pull", hide=True)
+
+
+def remove_env(ctx):
+    """Will remove the environment for commands applying to a worktree."""
+
+    if WORKTREE_DIRECTORY.is_dir():
+        ctx.run(f"git worktree remove '{WORKTREE_DIRECTORY}'", warn=True)
+
+
+def is_worktree():
+    """Will return True if the current environment is a worktree environment."""
+
+    return Path.cwd() != LOCAL_DIRECTORY
+
+
+def enter_env(ctx, branch: str):
+    """Enters the worktree environment."""
+
+    init_env(ctx, branch)
+
+    os.chdir(WORKTREE_DIRECTORY)
+
+
+def exit_env():
+    """Exits the worktree environment."""
+
+    os.chdir(LOCAL_DIRECTORY)
+
+
+@contextmanager
+def agent_context(ctx, branch: str | None):
+    """Applies code to the worktree environment if the version is not None.
+
+    Usage:
+        > with agent_context(ctx, branch):
+        >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target branch
+    """
+
+    if branch is not None:
+        # Do not stack two environments
+        if is_worktree():
+            yield
+            return
+
+        try:
+            # Enter
+            enter_env(ctx, branch)
+
+            yield
+        finally:
+            # Exit
+            exit_env()
+    else:
+        # NOTE: This ensures that we don't push local context from a worktree context (context might be switched within inner functions)
+        assert not is_worktree(), 'Local context cannot be used within a worktree context'
+
+        yield
diff --git a/tasks/libs/releasing/json.py b/tasks/libs/releasing/json.py
index ed281acb907b4..85270b23ce70e 100644
--- a/tasks/libs/releasing/json.py
+++ b/tasks/libs/releasing/json.py
@@ -6,9 +6,9 @@
 
 from invoke.exceptions import Exit
 
-from tasks.libs.common.agent6 import is_agent6
 from tasks.libs.common.constants import TAG_FOUND_TEMPLATE
 from tasks.libs.common.git import get_default_branch
+from tasks.libs.common.worktree import is_worktree
 from tasks.libs.releasing.documentation import _stringify_config, nightly_entry_for, release_entry_for
 from tasks.libs.releasing.version import (
     VERSION_RE,
@@ -347,7 +347,7 @@ def generate_repo_data(warning_mode, next_version, release_branch):
             branch = (
                 next_version.branch()
                 if repo == "integrations-core"
-                else (DEFAULT_BRANCHES_AGENT6 if is_agent6() else DEFAULT_BRANCHES).get(repo, get_default_branch())
+                else (DEFAULT_BRANCHES_AGENT6 if is_worktree() else DEFAULT_BRANCHES).get(repo, get_default_branch())
             )
         data[repo] = {
             'branch': branch,
diff --git a/tasks/unit_tests/libs/common/agent6_tests.py b/tasks/unit_tests/libs/common/agent6_tests.py
index dfffcdaa127c7..826aefb775fc7 100644
--- a/tasks/unit_tests/libs/common/agent6_tests.py
+++ b/tasks/unit_tests/libs/common/agent6_tests.py
@@ -3,53 +3,54 @@
 
 from invoke import Context
 
-from tasks.libs.common.agent6 import agent_context, init_env, is_agent6
 from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.gomodules import get_default_modules
+from tasks.libs.common.worktree import agent_context, init_env, is_worktree
 
 
 def get_ctx():
     return Context()
 
 
-class TestAgent6(unittest.TestCase):
+class TestWorktree(unittest.TestCase):
     def setUp(self):
+        print('SETUP')
         # Pull only once
-        init_env(get_ctx())
-        os.environ['AGENT6_NO_PULL'] = '1'
+        init_env(get_ctx(), '6.53.x')
+        os.environ['AGENT_WORKTREE_NO_PULL'] = '1'
 
-    def test_context_is_agent6_true(self):
-        with agent_context(get_ctx(), 6):
-            self.assertTrue(is_agent6())
+    def test_context_is_worktree_true(self):
+        with agent_context(get_ctx(), '6.53.x'):
+            self.assertTrue(is_worktree())
 
-    def test_context_is_agent6_false(self):
-        with agent_context(get_ctx(), 7):
-            self.assertFalse(is_agent6())
+    def test_context_is_worktree_false(self):
+        with agent_context(get_ctx(), None):
+            self.assertFalse(is_worktree())
 
     def test_context_nested(self):
-        with agent_context(get_ctx(), 6):
-            with agent_context(get_ctx(), 6):
-                self.assertTrue(is_agent6())
-            self.assertTrue(is_agent6())
+        with agent_context(get_ctx(), '6.53.x'):
+            with agent_context(get_ctx(), '6.53.x'):
+                self.assertTrue(is_worktree())
+            self.assertTrue(is_worktree())
 
     def test_context_pwd(self):
         ctx = get_ctx()
 
-        with agent_context(ctx, 7):
-            pwd7 = ctx.run('pwd').stdout
+        with agent_context(ctx, None):
+            pwdlocal = ctx.run('pwd').stdout
 
-        with agent_context(ctx, 6):
+        with agent_context(ctx, '6.53.x'):
             pwd6 = ctx.run('pwd').stdout
 
-        self.assertNotEqual(pwd6, pwd7)
+        self.assertNotEqual(pwd6, pwdlocal)
 
     def test_context_modules(self):
         ctx = get_ctx()
 
-        with agent_context(ctx, 7):
+        with agent_context(ctx, 'main'):
             modules7 = get_default_modules()
 
-        with agent_context(ctx, 6):
+        with agent_context(ctx, '6.53.x'):
             modules6 = get_default_modules()
 
         self.assertNotEqual(set(modules6.keys()), set(modules7.keys()))
@@ -57,10 +58,10 @@ def test_context_modules(self):
     def test_context_branch(self):
         ctx = get_ctx()
 
-        with agent_context(ctx, 7):
+        with agent_context(ctx, 'main'):
             branch7 = get_default_branch()
 
-        with agent_context(ctx, 6):
+        with agent_context(ctx, '6.53.x'):
             branch6 = get_default_branch()
 
         self.assertNotEqual(branch6, branch7)
diff --git a/tasks/worktree.py b/tasks/worktree.py
new file mode 100644
index 0000000000000..3924a787e2e56
--- /dev/null
+++ b/tasks/worktree.py
@@ -0,0 +1,44 @@
+from invoke import task
+
+from tasks.libs.common.worktree import agent_context, enter_env, init_env, remove_env
+
+
+@task
+def init(ctx, branch: str | None = None):
+    """Will prepare the worktree context (git clone / pull of the agent branch)."""
+
+    init_env(ctx, branch)
+
+
+@task
+def remove(ctx):
+    """Will the git worktree context."""
+
+    remove_env(ctx)
+
+
+@task
+def run(ctx, branch: str, command: str):
+    """Runs a command in the target worktree environment.
+
+    Usage:
+        $ inv worktree.run 6.53.x "head CHANGELOG.rst"  # Displays the changelog of the target branch
+    """
+
+    with agent_context(ctx, branch):
+        ctx.run(command)
+
+
+@task
+def invoke(ctx, branch: str):
+    """Enters the worktree environment in order to invoke tasks in this context.
+
+    Note:
+        This task should be avoided when a --version, --major-version or --agent-version argument is available in the task.
+
+    Usage:
+        > inv worktree.invoke 6.53.x modules.show-all  # Will show agent 6 modules
+    """
+
+    # The tasks running after this one will be using the agent 6 environment
+    enter_env(ctx, branch)

From 3e68cb5b390cb500973563ebeb565914c5bcd020 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 10:46:42 +0100
Subject: [PATCH 16/33] [rfc-update-tasks-a6-a7] Fixed modules (removed lru
 cache), renamed agent 6 -> worktree tests

---
 tasks/libs/common/gomodules.py                    | 15 +--------------
 .../common/{agent6_tests.py => worktree_tests.py} |  0
 2 files changed, 1 insertion(+), 14 deletions(-)
 rename tasks/unit_tests/libs/common/{agent6_tests.py => worktree_tests.py} (100%)

diff --git a/tasks/libs/common/gomodules.py b/tasks/libs/common/gomodules.py
index dd381a5e9db68..357ab03dc515e 100644
--- a/tasks/libs/common/gomodules.py
+++ b/tasks/libs/common/gomodules.py
@@ -7,7 +7,6 @@
 import sys
 from collections.abc import Callable
 from dataclasses import dataclass
-from functools import lru_cache
 from pathlib import Path
 from typing import ClassVar
 
@@ -306,18 +305,6 @@ def dependency_path(self, agent_version):
 AGENT_MODULE_PATH_PREFIX = "github.com/DataDog/datadog-agent/"
 
 
-@lru_cache
-def _get_default_modules_abs(base_dir_abs: Path) -> dict[str, GoModule]:
-    """Get default modules from absolute path.
-
-    Used to lru cache the result.
-    """
-
-    assert base_dir_abs.is_absolute()
-
-    return Configuration.from_file(base_dir_abs).modules
-
-
 def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
     """Load the default modules from the modules.yml file.
 
@@ -327,7 +314,7 @@ def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
 
     base_dir = base_dir or agent_working_directory()
 
-    return _get_default_modules_abs(base_dir.resolve())
+    return Configuration.from_file(base_dir).modules
 
 
 def validate_module(
diff --git a/tasks/unit_tests/libs/common/agent6_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
similarity index 100%
rename from tasks/unit_tests/libs/common/agent6_tests.py
rename to tasks/unit_tests/libs/common/worktree_tests.py

From 96daa32edbeefb1916318fbdd606ec8bac386592 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 11:53:06 +0100
Subject: [PATCH 17/33] [rfc-update-tasks-a6-a7] Added no switch option

---
 tasks/libs/common/worktree.py                  | 15 +++++++++++----
 tasks/unit_tests/libs/common/worktree_tests.py | 15 +++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 11d758d64954f..4a7f806ae5646 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -48,10 +48,13 @@ def is_worktree():
     return Path.cwd() != LOCAL_DIRECTORY
 
 
-def enter_env(ctx, branch: str):
+def enter_env(ctx, branch: str, no_switch=False):
     """Enters the worktree environment."""
 
-    init_env(ctx, branch)
+    if not no_switch:
+        init_env(ctx, branch)
+    else:
+        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_switch is set to True"
 
     os.chdir(WORKTREE_DIRECTORY)
 
@@ -63,9 +66,13 @@ def exit_env():
 
 
 @contextmanager
-def agent_context(ctx, branch: str | None):
+def agent_context(ctx, branch: str | None, no_switch=False):
     """Applies code to the worktree environment if the version is not None.
 
+    Args:
+        branch: The branch to switch to.
+        no_switch: If True, the branch will not be switched (no pull will be performed too).
+
     Usage:
         > with agent_context(ctx, branch):
         >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target branch
@@ -79,7 +86,7 @@ def agent_context(ctx, branch: str | None):
 
         try:
             # Enter
-            enter_env(ctx, branch)
+            enter_env(ctx, branch, no_switch=no_switch)
 
             yield
         finally:
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index 826aefb775fc7..bf0eb9d70a9eb 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -65,3 +65,18 @@ def test_context_branch(self):
             branch6 = get_default_branch()
 
         self.assertNotEqual(branch6, branch7)
+
+    def test_context_no_switch(self):
+        ctx = get_ctx()
+
+        with agent_context(ctx, '6.53.x'):
+            branch6 = get_default_branch()
+
+        with agent_context(ctx, 'main'):
+            branch7 = get_default_branch()
+
+        with agent_context(ctx, 'main', no_switch=True):
+            branch_noswitch = get_default_branch()
+
+        self.assertNotEqual(branch6, branch7)
+        self.assertEqual(branch7, branch_noswitch)

From 210e96c34ad74d08e8b4b3f9d0b29f0e37ece719 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 12:08:30 +0100
Subject: [PATCH 18/33] [rfc-update-tasks-a6-a7] switch -> checkout

---
 tasks/libs/common/worktree.py                  | 10 +++++-----
 tasks/unit_tests/libs/common/worktree_tests.py |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 4a7f806ae5646..765de58b5ff1a 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -48,10 +48,10 @@ def is_worktree():
     return Path.cwd() != LOCAL_DIRECTORY
 
 
-def enter_env(ctx, branch: str, no_switch=False):
+def enter_env(ctx, branch: str, no_checkout=False):
     """Enters the worktree environment."""
 
-    if not no_switch:
+    if not no_checkout:
         init_env(ctx, branch)
     else:
         assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_switch is set to True"
@@ -66,12 +66,12 @@ def exit_env():
 
 
 @contextmanager
-def agent_context(ctx, branch: str | None, no_switch=False):
+def agent_context(ctx, branch: str | None, no_checkout=False):
     """Applies code to the worktree environment if the version is not None.
 
     Args:
         branch: The branch to switch to.
-        no_switch: If True, the branch will not be switched (no pull will be performed too).
+        no_checkout: If True, the branch will not be switched (no pull will be performed too).
 
     Usage:
         > with agent_context(ctx, branch):
@@ -86,7 +86,7 @@ def agent_context(ctx, branch: str | None, no_switch=False):
 
         try:
             # Enter
-            enter_env(ctx, branch, no_switch=no_switch)
+            enter_env(ctx, branch, no_checkout=no_checkout)
 
             yield
         finally:
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index bf0eb9d70a9eb..c2d65ed42bb2e 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -75,7 +75,7 @@ def test_context_no_switch(self):
         with agent_context(ctx, 'main'):
             branch7 = get_default_branch()
 
-        with agent_context(ctx, 'main', no_switch=True):
+        with agent_context(ctx, 'main', no_checkout=True):
             branch_noswitch = get_default_branch()
 
         self.assertNotEqual(branch6, branch7)

From 7e612583559f7926e7aafe0b15e7754eff8b7843 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 13:18:02 +0100
Subject: [PATCH 19/33] [rfc-update-tasks-a6-a7] switch -> checkout

---
 tasks/libs/common/worktree.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 765de58b5ff1a..ccf287fca1209 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -29,7 +29,7 @@ def init_env(ctx, branch: str | None = None):
             f"git -C '{WORKTREE_DIRECTORY}' rev-parse --abbrev-ref HEAD", hide=True
         ).stdout.strip()
         if worktree_branch != branch:
-            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' switch '{branch}'", hide=True)
+            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' checkout '{branch}'", hide=True)
 
         if not os.environ.get("AGENT_WORKTREE_NO_PULL"):
             ctx.run(f"git -C '{WORKTREE_DIRECTORY}' pull", hide=True)

From c0b62e11e318e25d7ed062bcb5b5925044002cc1 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 18:12:02 +0100
Subject: [PATCH 20/33] [rfc-update-tasks-a6-a7] remove: Fixed

---
 tasks/libs/common/worktree.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index ccf287fca1209..d3bb674213419 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -38,8 +38,7 @@ def init_env(ctx, branch: str | None = None):
 def remove_env(ctx):
     """Will remove the environment for commands applying to a worktree."""
 
-    if WORKTREE_DIRECTORY.is_dir():
-        ctx.run(f"git worktree remove '{WORKTREE_DIRECTORY}'", warn=True)
+    ctx.run(f"git worktree remove -f '{WORKTREE_DIRECTORY}'", warn=True)
 
 
 def is_worktree():

From bd9dd21cf2677f318ab000b5a6b63af9c7981495 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 18:12:47 +0100
Subject: [PATCH 21/33] [rfc-update-tasks-a6-a7] Added notes

---
 tasks/libs/common/worktree.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index d3bb674213419..54e654e282ecb 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -22,7 +22,12 @@ def init_env(ctx, branch: str | None = None):
     """
 
     if not WORKTREE_DIRECTORY.is_dir():
-        ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}")
+        try:
+            ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}")
+        except Exception as e:
+            e.add_note('You might want to reset the worktree directory with `inv worktree.remove`')
+
+            raise
 
     if branch:
         worktree_branch = ctx.run(

From 2466d2720508073c77af1ce48015363264b28748 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Thu, 21 Nov 2024 18:18:31 +0100
Subject: [PATCH 22/33] [rfc-update-tasks-a6-a7] Fix error note

---
 tasks/libs/common/worktree.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 54e654e282ecb..da54a2e219505 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -8,6 +8,10 @@
 from contextlib import contextmanager
 from pathlib import Path
 
+from invoke.exceptions import Exit
+
+from tasks.libs.common.color import Color, color_message
+
 WORKTREE_DIRECTORY = Path.cwd().parent / "datadog-agent-worktree"
 LOCAL_DIRECTORY = Path.cwd().resolve()
 
@@ -22,12 +26,11 @@ def init_env(ctx, branch: str | None = None):
     """
 
     if not WORKTREE_DIRECTORY.is_dir():
-        try:
-            ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}")
-        except Exception as e:
-            e.add_note('You might want to reset the worktree directory with `inv worktree.remove`')
-
-            raise
+        if not ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}", warn=True):
+            raise Exit(
+                f'{color_message("Error", Color.RED)}: Cannot initialize worktree environment. You might want to reset the worktree directory with `inv worktree.remove`',
+                code=1,
+            )
 
     if branch:
         worktree_branch = ctx.run(

From a1290b5d872b757fb282a0273ca0bdce32f20bf1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9lian=20Raimbault?=
 <161456554+CelianR@users.noreply.github.com>
Date: Mon, 25 Nov 2024 03:51:28 -0500
Subject: [PATCH 23/33] Apply suggestions from code review [skip ci]

Co-authored-by: sabrina lu <sabrina.lu@datadoghq.com>
---
 tasks/libs/common/worktree.py                  | 4 ++--
 tasks/unit_tests/libs/common/worktree_tests.py | 2 +-
 tasks/worktree.py                              | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index da54a2e219505..23cbda6a4ec70 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -61,7 +61,7 @@ def enter_env(ctx, branch: str, no_checkout=False):
     if not no_checkout:
         init_env(ctx, branch)
     else:
-        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_switch is set to True"
+        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_checkout is set to True"
 
     os.chdir(WORKTREE_DIRECTORY)
 
@@ -74,7 +74,7 @@ def exit_env():
 
 @contextmanager
 def agent_context(ctx, branch: str | None, no_checkout=False):
-    """Applies code to the worktree environment if the version is not None.
+    """Applies code to the worktree environment if the branch is not None.
 
     Args:
         branch: The branch to switch to.
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index c2d65ed42bb2e..d5a4baf7a2320 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -66,7 +66,7 @@ def test_context_branch(self):
 
         self.assertNotEqual(branch6, branch7)
 
-    def test_context_no_switch(self):
+    def test_context_no_checkout(self):
         ctx = get_ctx()
 
         with agent_context(ctx, '6.53.x'):
diff --git a/tasks/worktree.py b/tasks/worktree.py
index 3924a787e2e56..5adfd15b48d6c 100644
--- a/tasks/worktree.py
+++ b/tasks/worktree.py
@@ -12,7 +12,7 @@ def init(ctx, branch: str | None = None):
 
 @task
 def remove(ctx):
-    """Will the git worktree context."""
+    """Will remove the git worktree context."""
 
     remove_env(ctx)
 

From 1869b01f7fffcacd9b2e79b48cedd941940aa8ec Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 10:08:15 +0100
Subject: [PATCH 24/33] [rfc-update-tasks-a6-a7] Applied Sabrina's suggestions

---
 tasks/git.py                                   |  6 +++++-
 tasks/libs/common/worktree.py                  | 12 ++++++------
 tasks/modules.py                               |  4 ++--
 tasks/pipeline.py                              | 14 +++++++-------
 tasks/release.py                               |  7 +++++--
 tasks/unit_tests/libs/common/worktree_tests.py |  6 +++---
 6 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/tasks/git.py b/tasks/git.py
index 2574ff19c7091..8125035b04320 100644
--- a/tasks/git.py
+++ b/tasks/git.py
@@ -12,7 +12,11 @@ def check_protected_branch(ctx):
     local_branch = get_current_branch(ctx)
 
     if local_branch == get_default_branch():
-        print(color_message("You're about to commit or push to the main, are you sure this is what you want?", "red"))
+        print(
+            color_message(
+                f"You're about to commit or push to {get_default_branch()}, are you sure this is what you want?", "red"
+            )
+        )
         raise Exit(code=1)
 
     if re.fullmatch(r'^[0-9]+\.[0-9]+\.x$', local_branch):
diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 23cbda6a4ec70..4635076fefabf 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -52,13 +52,13 @@ def remove_env(ctx):
 def is_worktree():
     """Will return True if the current environment is a worktree environment."""
 
-    return Path.cwd() != LOCAL_DIRECTORY
+    return Path.cwd() == WORKTREE_DIRECTORY
 
 
-def enter_env(ctx, branch: str, no_checkout=False):
+def enter_env(ctx, branch: str, skip_checkout=False):
     """Enters the worktree environment."""
 
-    if not no_checkout:
+    if not skip_checkout:
         init_env(ctx, branch)
     else:
         assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_checkout is set to True"
@@ -73,12 +73,12 @@ def exit_env():
 
 
 @contextmanager
-def agent_context(ctx, branch: str | None, no_checkout=False):
+def agent_context(ctx, branch: str | None, skip_checkout=False):
     """Applies code to the worktree environment if the branch is not None.
 
     Args:
         branch: The branch to switch to.
-        no_checkout: If True, the branch will not be switched (no pull will be performed too).
+        skip_checkout: If True, the branch will not be switched (no pull will be performed too).
 
     Usage:
         > with agent_context(ctx, branch):
@@ -93,7 +93,7 @@ def agent_context(ctx, branch: str | None, no_checkout=False):
 
         try:
             # Enter
-            enter_env(ctx, branch, no_checkout=no_checkout)
+            enter_env(ctx, branch, skip_checkout=skip_checkout)
 
             yield
         finally:
diff --git a/tasks/modules.py b/tasks/modules.py
index d063eb50ae2d2..9172538d4a92b 100644
--- a/tasks/modules.py
+++ b/tasks/modules.py
@@ -265,7 +265,7 @@ def show(_, path: str, remove_defaults: bool = False, base_dir: str = '.'):
 
     Args:
         remove_defaults: If True, will remove default values from the output.
-        version: If 6, will show agent 6 modules.
+        base_dir: Where to load modules from.
     """
 
     config = Configuration.from_file(Path(base_dir))
@@ -287,8 +287,8 @@ def show_all(_, base_dir: str = '.', ignored=False):
     """Show the list of modules.
 
     Args:
+        base_dir: Where to load modules from.
         ignored: If True, will list ignored modules.
-        version: If 6, will show agent 6 modules.
     """
 
     config = Configuration.from_file(Path(base_dir))
diff --git a/tasks/pipeline.py b/tasks/pipeline.py
index 1221bf6227658..8d63109fee536 100644
--- a/tasks/pipeline.py
+++ b/tasks/pipeline.py
@@ -831,16 +831,16 @@ def test_merge_queue(ctx):
     # Create a new main and push it
     print("Creating a new main branch")
     timestamp = int(datetime.now(timezone.utc).timestamp())
-    test_main = f"mq/test_{timestamp}"
+    test_default = f"mq/test_{timestamp}"
     current_branch = get_current_branch(ctx)
     ctx.run(f"git checkout {get_default_branch()}", hide=True)
     ctx.run("git pull", hide=True)
-    ctx.run(f"git checkout -b {test_main}", hide=True)
-    ctx.run(f"git push origin {test_main}", hide=True)
+    ctx.run(f"git checkout -b {test_default}", hide=True)
+    ctx.run(f"git push origin {test_default}", hide=True)
     # Create a PR towards this new branch and adds it to the merge queue
     print("Creating a PR and adding it to the merge queue")
     gh = GithubAPI()
-    pr = gh.create_pr(f"Test MQ for {current_branch}", "", test_main, current_branch)
+    pr = gh.create_pr(f"Test MQ for {current_branch}", "", test_default, current_branch)
     pr.create_issue_comment("/merge")
     # Search for the generated pipeline
     print(f"PR {pr.html_url} is waiting for MQ pipeline generation")
@@ -850,7 +850,7 @@ def test_merge_queue(ctx):
         time.sleep(30)
         pipelines = agent.pipelines.list(per_page=100)
         try:
-            pipeline = next(p for p in pipelines if p.ref.startswith(f"mq-working-branch-{test_main}"))
+            pipeline = next(p for p in pipelines if p.ref.startswith(f"mq-working-branch-{test_default}"))
             print(f"Pipeline found: {pipeline.web_url}")
             break
         except StopIteration as e:
@@ -868,8 +868,8 @@ def test_merge_queue(ctx):
         pipeline.cancel()
     pr.edit(state="closed")
     ctx.run(f"git checkout {current_branch}", hide=True)
-    ctx.run(f"git branch -D {test_main}", hide=True)
-    ctx.run(f"git push origin :{test_main}", hide=True)
+    ctx.run(f"git branch -D {test_default}", hide=True)
+    ctx.run(f"git push origin :{test_default}", hide=True)
     if not success:
         raise Exit(message="Merge queue test failed", code=1)
 
diff --git a/tasks/release.py b/tasks/release.py
index c1cea2eb9f8d2..136400ec66ab5 100644
--- a/tasks/release.py
+++ b/tasks/release.py
@@ -30,6 +30,7 @@
     get_default_branch,
     get_last_commit,
     get_last_release_tag,
+    is_agent6,
     try_git_command,
 )
 from tasks.libs.common.gomodules import get_default_modules
@@ -48,6 +49,7 @@
 )
 from tasks.libs.releasing.json import (
     DEFAULT_BRANCHES,
+    DEFAULT_BRANCHES_AGENT6,
     UNFREEZE_REPO_AGENT,
     UNFREEZE_REPOS,
     _get_release_json_value,
@@ -789,8 +791,9 @@ def cleanup(ctx):
 def check_omnibus_branches(ctx):
     base_branch = _get_release_json_value('base_branch')
     if base_branch == get_default_branch():
-        omnibus_ruby_branch = DEFAULT_BRANCHES['omnibus-ruby']
-        omnibus_software_branch = DEFAULT_BRANCHES['omnibus-software']
+        default_branches = DEFAULT_BRANCHES_AGENT6 if is_agent6() else DEFAULT_BRANCHES
+        omnibus_ruby_branch = default_branches['omnibus-ruby']
+        omnibus_software_branch = default_branches['omnibus-software']
     else:
         omnibus_ruby_branch = base_branch
         omnibus_software_branch = base_branch
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index d5a4baf7a2320..caec486932432 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -75,8 +75,8 @@ def test_context_no_checkout(self):
         with agent_context(ctx, 'main'):
             branch7 = get_default_branch()
 
-        with agent_context(ctx, 'main', no_checkout=True):
-            branch_noswitch = get_default_branch()
+        with agent_context(ctx, 'main', skip_checkout=True):
+            branch_no_checkout = get_default_branch()
 
         self.assertNotEqual(branch6, branch7)
-        self.assertEqual(branch7, branch_noswitch)
+        self.assertEqual(branch7, branch_no_checkout)

From c127193eaa3c79c487616b88aa0f88429ad9a55a Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 10:53:42 +0100
Subject: [PATCH 25/33] [rfc-update-tasks-a6-a7] skip_checkout: Verified that
 the branch is already checked out

---
 tasks/libs/common/worktree.py                  |  6 ++++++
 tasks/unit_tests/libs/common/worktree_tests.py | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 4635076fefabf..0d4d6ea269432 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -11,6 +11,7 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
+from tasks.libs.common.git import get_current_branch
 
 WORKTREE_DIRECTORY = Path.cwd().parent / "datadog-agent-worktree"
 LOCAL_DIRECTORY = Path.cwd().resolve()
@@ -64,6 +65,11 @@ def enter_env(ctx, branch: str, skip_checkout=False):
         assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_checkout is set to True"
 
     os.chdir(WORKTREE_DIRECTORY)
+    if skip_checkout:
+        current_branch = get_current_branch(ctx)
+        assert (
+            current_branch == branch
+        ), f"skip_checkout is True but the current branch ({current_branch}) is not {branch}"
 
 
 def exit_env():
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index caec486932432..90bfd83239fef 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -80,3 +80,16 @@ def test_context_no_checkout(self):
 
         self.assertNotEqual(branch6, branch7)
         self.assertEqual(branch7, branch_no_checkout)
+
+    def test_context_no_checkout_error(self):
+        ctx = get_ctx()
+
+        with agent_context(ctx, '6.53.x'):
+            pass
+
+        def switch_context():
+            # The current branch is not main
+            with agent_context(ctx, 'main', skip_checkout=True):
+                pass
+
+        self.assertRaises(AssertionError, switch_context)

From 0ffb4d43e8650376c28a9a7b61537c66e26b9b13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9lian=20Raimbault?=
 <161456554+CelianR@users.noreply.github.com>
Date: Mon, 25 Nov 2024 08:01:21 -0500
Subject: [PATCH 26/33] Update tasks/libs/common/worktree.py

Co-authored-by: pducolin <45568537+pducolin@users.noreply.github.com>
---
 tasks/libs/common/worktree.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 0d4d6ea269432..e3bdcf01db590 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -62,7 +62,7 @@ def enter_env(ctx, branch: str, skip_checkout=False):
     if not skip_checkout:
         init_env(ctx, branch)
     else:
-        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_checkout is set to True"
+        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and skip_checkout is set to True"
 
     os.chdir(WORKTREE_DIRECTORY)
     if skip_checkout:

From 1e57702a3816c10cae888fcb62fcaca69332e774 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 14:35:30 +0100
Subject: [PATCH 27/33] [rfc-update-tasks-a6-a7] default_modules: Refactored
 since it is not cached anymore

---
 tasks/gotest.py                                | 11 +++++++----
 tasks/modules.py                               |  3 ++-
 tasks/unit_tests/libs/common/worktree_tests.py |  1 -
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/tasks/gotest.py b/tasks/gotest.py
index ceaee2f020bc2..3ef5ebc9de905 100644
--- a/tasks/gotest.py
+++ b/tasks/gotest.py
@@ -514,12 +514,13 @@ def get_modified_packages(ctx, build_tags=None, lint=False) -> list[GoModule]:
             modules_to_test[best_module_path] = GoModule(best_module_path, test_targets=[relative_target])
 
     # Clean up duplicated paths to reduce Go test cmd length
+    default_modules = get_default_modules()
     for module in modules_to_test:
         modules_to_test[module].test_targets = clean_nested_paths(modules_to_test[module].test_targets)
         if (
             len(modules_to_test[module].test_targets) >= WINDOWS_MAX_PACKAGES_NUMBER
         ):  # With more packages we can reach the limit of the command line length on Windows
-            modules_to_test[module].test_targets = get_default_modules()[module].test_targets
+            modules_to_test[module].test_targets = default_modules[module].test_targets
 
     print("Running tests for the following modules:")
     for module in modules_to_test:
@@ -752,16 +753,17 @@ def format_packages(ctx: Context, impacted_packages: set[str], build_tags: list[
     packages = [f'{package.replace("github.com/DataDog/datadog-agent/", "./")}' for package in impacted_packages]
     modules_to_test = {}
 
+    default_modules = get_default_modules()
     for package in packages:
         module_path = get_go_module(package)
 
         # Check if the module is in the target list of the modules we want to test
-        if module_path not in get_default_modules() or not get_default_modules()[module_path].should_test():
+        if module_path not in default_modules or not default_modules[module_path].should_test():
             continue
 
         # Check if the package is in the target list of the module we want to test
         targeted = False
-        for target in get_default_modules()[module_path].test_targets:
+        for target in default_modules[module_path].test_targets:
             if normpath(os.path.join(module_path, target)) in package:
                 targeted = True
                 break
@@ -784,12 +786,13 @@ def format_packages(ctx: Context, impacted_packages: set[str], build_tags: list[
             modules_to_test[module_path] = GoModule(module_path, test_targets=[relative_target])
 
     # Clean up duplicated paths to reduce Go test cmd length
+    default_modules = get_default_modules()
     for module in modules_to_test:
         modules_to_test[module].test_targets = clean_nested_paths(modules_to_test[module].test_targets)
         if (
             len(modules_to_test[module].test_targets) >= WINDOWS_MAX_PACKAGES_NUMBER
         ):  # With more packages we can reach the limit of the command line length on Windows
-            modules_to_test[module].test_targets = get_default_modules()[module].test_targets
+            modules_to_test[module].test_targets = default_modules[module].test_targets
 
     module_to_remove = []
     # Clean up to avoid running tests on package with no Go files matching build tags
diff --git a/tasks/modules.py b/tasks/modules.py
index 9172538d4a92b..8e71e4cf89ced 100644
--- a/tasks/modules.py
+++ b/tasks/modules.py
@@ -218,6 +218,7 @@ def validate_used_by_otel(ctx: Context):
     missing_used_by_otel_label: dict[str, list[str]] = defaultdict(list)
 
     # for every module labeled as "used_by_otel"
+    default_modules = get_default_modules()
     for otel_mod in otel_mods:
         gomod_path = f"{otel_mod}/go.mod"
         # get the go.mod data
@@ -237,7 +238,7 @@ def validate_used_by_otel(ctx: Context):
             # we need the relative path of module (without github.com/DataDog/datadog-agent/ prefix)
             rel_path = require['Path'].removeprefix("github.com/DataDog/datadog-agent/")
             # check if indirect module is labeled as "used_by_otel"
-            if rel_path not in get_default_modules() or not get_default_modules()[rel_path].used_by_otel:
+            if rel_path not in default_modules or not default_modules[rel_path].used_by_otel:
                 missing_used_by_otel_label[rel_path].append(otel_mod)
     if missing_used_by_otel_label:
         message = f"{color_message('ERROR', Color.RED)}: some indirect local dependencies of modules labeled \"used_by_otel\" are not correctly labeled in get_default_modules()\n"
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index 90bfd83239fef..ab8b6b3728317 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -14,7 +14,6 @@ def get_ctx():
 
 class TestWorktree(unittest.TestCase):
     def setUp(self):
-        print('SETUP')
         # Pull only once
         init_env(get_ctx(), '6.53.x')
         os.environ['AGENT_WORKTREE_NO_PULL'] = '1'

From 7daf382d2d4256e208c51334dee7876bfa859601 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 15:14:00 +0100
Subject: [PATCH 28/33] [rfc-update-tasks-a6-a7] Fixed missing arg

---
 tasks/release.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tasks/release.py b/tasks/release.py
index 136400ec66ab5..5e0b11483a938 100644
--- a/tasks/release.py
+++ b/tasks/release.py
@@ -791,7 +791,7 @@ def cleanup(ctx):
 def check_omnibus_branches(ctx):
     base_branch = _get_release_json_value('base_branch')
     if base_branch == get_default_branch():
-        default_branches = DEFAULT_BRANCHES_AGENT6 if is_agent6() else DEFAULT_BRANCHES
+        default_branches = DEFAULT_BRANCHES_AGENT6 if is_agent6(ctx) else DEFAULT_BRANCHES
         omnibus_ruby_branch = default_branches['omnibus-ruby']
         omnibus_software_branch = default_branches['omnibus-software']
     else:

From 14ad3d7021b8cff1cd8f12eeeaa897f7e695553b Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 15:22:13 +0100
Subject: [PATCH 29/33] [rfc-update-tasks-a6-a7] Added various git commands to
 worktree and skip_checkout option

---
 tasks/worktree.py | 45 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/tasks/worktree.py b/tasks/worktree.py
index 5adfd15b48d6c..6a6f08042fb9e 100644
--- a/tasks/worktree.py
+++ b/tasks/worktree.py
@@ -1,6 +1,7 @@
 from invoke import task
+from invoke.exceptions import Exit
 
-from tasks.libs.common.worktree import agent_context, enter_env, init_env, remove_env
+from tasks.libs.common.worktree import WORKTREE_DIRECTORY, agent_context, enter_env, init_env, remove_env
 
 
 @task
@@ -18,19 +19,53 @@ def remove(ctx):
 
 
 @task
-def run(ctx, branch: str, command: str):
+def status(ctx):
+    """Displays the status of the worktree environment."""
+
+    if not WORKTREE_DIRECTORY.is_dir():
+        raise Exit('No worktree environment found.')
+
+    ctx.run(f"git -C '{WORKTREE_DIRECTORY}' status", pty=True)
+
+
+@task
+def switch(ctx, ref):
+    """Changes the worktree environment to the specified ref.
+
+    Note:
+        This won't pull.
+    """
+
+    if not WORKTREE_DIRECTORY.is_dir():
+        raise Exit('No worktree environment found.')
+
+    ctx.run(f"git -C '{WORKTREE_DIRECTORY}' checkout '{ref}'", pty=True)
+
+
+@task
+def pull(ctx):
+    """Pulls the worktree environment."""
+
+    if not WORKTREE_DIRECTORY.is_dir():
+        raise Exit('No worktree environment found.')
+
+    ctx.run(f"git -C '{WORKTREE_DIRECTORY}' pull", pty=True)
+
+
+@task
+def run(ctx, branch: str, command: str, skip_checkout: bool = False):
     """Runs a command in the target worktree environment.
 
     Usage:
         $ inv worktree.run 6.53.x "head CHANGELOG.rst"  # Displays the changelog of the target branch
     """
 
-    with agent_context(ctx, branch):
+    with agent_context(ctx, branch, skip_checkout=skip_checkout):
         ctx.run(command)
 
 
 @task
-def invoke(ctx, branch: str):
+def invoke(ctx, branch: str, skip_checkout: bool = False):
     """Enters the worktree environment in order to invoke tasks in this context.
 
     Note:
@@ -41,4 +76,4 @@ def invoke(ctx, branch: str):
     """
 
     # The tasks running after this one will be using the agent 6 environment
-    enter_env(ctx, branch)
+    enter_env(ctx, branch, skip_checkout=skip_checkout)

From 3d2e68ba8211545ca64f0e7a9ddaf6406a00c275 Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 15:27:39 +0100
Subject: [PATCH 30/33] [rfc-update-tasks-a6-a7] Updated error message

---
 tasks/libs/common/worktree.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index e3bdcf01db590..23641992ca43f 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -69,7 +69,7 @@ def enter_env(ctx, branch: str, skip_checkout=False):
         current_branch = get_current_branch(ctx)
         assert (
             current_branch == branch
-        ), f"skip_checkout is True but the current branch ({current_branch}) is not {branch}"
+        ), f"skip_checkout is True but the current branch ({current_branch}) is not {branch}. You should switch branch before using this command, this can be safely done with `inv worktree.switch {branch}`."
 
 
 def exit_env():

From 3a2147ed63196c04af4ea1c16d78fc250b471e2a Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Mon, 25 Nov 2024 15:52:33 +0100
Subject: [PATCH 31/33] [rfc-update-tasks-a6-a7] agent_context: If branch is
 None, then it wont switch branch but will enter context

---
 tasks/libs/common/worktree.py                 | 35 +++++++++----------
 .../unit_tests/libs/common/worktree_tests.py  | 13 ++++---
 2 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 23641992ca43f..4f18a39448981 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -56,16 +56,19 @@ def is_worktree():
     return Path.cwd() == WORKTREE_DIRECTORY
 
 
-def enter_env(ctx, branch: str, skip_checkout=False):
+def enter_env(ctx, branch: str | None, skip_checkout=False):
     """Enters the worktree environment."""
 
+    if not branch:
+        assert skip_checkout, 'skip_checkout must be set to True if branch is None'
+
     if not skip_checkout:
         init_env(ctx, branch)
     else:
         assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and skip_checkout is set to True"
 
     os.chdir(WORKTREE_DIRECTORY)
-    if skip_checkout:
+    if skip_checkout and branch:
         current_branch = get_current_branch(ctx)
         assert (
             current_branch == branch
@@ -83,7 +86,7 @@ def agent_context(ctx, branch: str | None, skip_checkout=False):
     """Applies code to the worktree environment if the branch is not None.
 
     Args:
-        branch: The branch to switch to.
+        branch: The branch to switch to. If None, will enter the worktree environment without switching branch (ensures that skip_checkout is True).
         skip_checkout: If True, the branch will not be switched (no pull will be performed too).
 
     Usage:
@@ -91,22 +94,16 @@ def agent_context(ctx, branch: str | None, skip_checkout=False):
         >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target branch
     """
 
-    if branch is not None:
-        # Do not stack two environments
-        if is_worktree():
-            yield
-            return
-
-        try:
-            # Enter
-            enter_env(ctx, branch, skip_checkout=skip_checkout)
+    # Do not stack two environments
+    if is_worktree():
+        yield
+        return
 
-            yield
-        finally:
-            # Exit
-            exit_env()
-    else:
-        # NOTE: This ensures that we don't push local context from a worktree context (context might be switched within inner functions)
-        assert not is_worktree(), 'Local context cannot be used within a worktree context'
+    try:
+        # Enter
+        enter_env(ctx, branch, skip_checkout=skip_checkout)
 
         yield
+    finally:
+        # Exit
+        exit_env()
diff --git a/tasks/unit_tests/libs/common/worktree_tests.py b/tasks/unit_tests/libs/common/worktree_tests.py
index ab8b6b3728317..3998357c1accc 100644
--- a/tasks/unit_tests/libs/common/worktree_tests.py
+++ b/tasks/unit_tests/libs/common/worktree_tests.py
@@ -23,8 +23,7 @@ def test_context_is_worktree_true(self):
             self.assertTrue(is_worktree())
 
     def test_context_is_worktree_false(self):
-        with agent_context(get_ctx(), None):
-            self.assertFalse(is_worktree())
+        self.assertFalse(is_worktree())
 
     def test_context_nested(self):
         with agent_context(get_ctx(), '6.53.x'):
@@ -35,13 +34,17 @@ def test_context_nested(self):
     def test_context_pwd(self):
         ctx = get_ctx()
 
-        with agent_context(ctx, None):
-            pwdlocal = ctx.run('pwd').stdout
+        with agent_context(ctx, None, skip_checkout=True):
+            pwdnone = ctx.run('pwd').stdout
 
         with agent_context(ctx, '6.53.x'):
             pwd6 = ctx.run('pwd').stdout
 
-        self.assertNotEqual(pwd6, pwdlocal)
+        with agent_context(ctx, 'main'):
+            pwdmain = ctx.run('pwd').stdout
+
+        self.assertEqual(pwd6, pwdnone)
+        self.assertEqual(pwd6, pwdmain)
 
     def test_context_modules(self):
         ctx = get_ctx()

From 74155c391145c3919aa90a1129509fac21cacf9c Mon Sep 17 00:00:00 2001
From: Celian Raimbault <celian.raimbault@datadoghq.com>
Date: Tue, 26 Nov 2024 10:08:25 +0100
Subject: [PATCH 32/33] [rfc-update-tasks-a6-a7] invoke: Added warning

---
 tasks/worktree.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/tasks/worktree.py b/tasks/worktree.py
index 6a6f08042fb9e..57a9c90580293 100644
--- a/tasks/worktree.py
+++ b/tasks/worktree.py
@@ -1,6 +1,8 @@
 from invoke import task
 from invoke.exceptions import Exit
 
+from tasks.libs.common.color import Color, color_message
+from tasks.libs.common.user_interactions import yes_no_question
 from tasks.libs.common.worktree import WORKTREE_DIRECTORY, agent_context, enter_env, init_env, remove_env
 
 
@@ -65,15 +67,22 @@ def run(ctx, branch: str, command: str, skip_checkout: bool = False):
 
 
 @task
-def invoke(ctx, branch: str, skip_checkout: bool = False):
+def invoke(ctx, branch: str, skip_checkout: bool = False, yes: bool = False):
     """Enters the worktree environment in order to invoke tasks in this context.
 
     Note:
-        This task should be avoided when a --version, --major-version or --agent-version argument is available in the task.
+        This task should be avoided when a --branch or --release-branch argument is available in the task.
 
     Usage:
-        > inv worktree.invoke 6.53.x modules.show-all  # Will show agent 6 modules
+        > inv worktree.invoke 6.53.x --yes modules.show-all  # Will show agent 6 modules
     """
 
-    # The tasks running after this one will be using the agent 6 environment
-    enter_env(ctx, branch, skip_checkout=skip_checkout)
+    if yes or yes_no_question(
+        'Warning: This task should be avoided, use --branch or --release-branch argument if available in the task. Want to proceed?',
+        color=Color.ORANGE,
+        default=False,
+    ):
+        # The tasks running after this one will be using the agent 6 environment
+        enter_env(ctx, branch, skip_checkout=skip_checkout)
+    else:
+        raise Exit(color_message('Aborted.', Color.RED))

From bc87c2d77acae88c6ebee5d9c48919455b539632 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9lian=20Raimbault?=
 <161456554+CelianR@users.noreply.github.com>
Date: Tue, 26 Nov 2024 04:11:43 -0500
Subject: [PATCH 33/33] Apply suggestions from code review

Co-authored-by: sabrina lu <sabrina.lu@datadoghq.com>
---
 tasks/libs/common/worktree.py | 4 ++--
 tasks/worktree.py             | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tasks/libs/common/worktree.py b/tasks/libs/common/worktree.py
index 4f18a39448981..7e7da54af73df 100644
--- a/tasks/libs/common/worktree.py
+++ b/tasks/libs/common/worktree.py
@@ -72,7 +72,7 @@ def enter_env(ctx, branch: str | None, skip_checkout=False):
         current_branch = get_current_branch(ctx)
         assert (
             current_branch == branch
-        ), f"skip_checkout is True but the current branch ({current_branch}) is not {branch}. You should switch branch before using this command, this can be safely done with `inv worktree.switch {branch}`."
+        ), f"skip_checkout is True but the current branch ({current_branch}) is not {branch}. You should check out the branch before using this command, this can be safely done with `inv worktree.checkout {branch}`."
 
 
 def exit_env():
@@ -87,7 +87,7 @@ def agent_context(ctx, branch: str | None, skip_checkout=False):
 
     Args:
         branch: The branch to switch to. If None, will enter the worktree environment without switching branch (ensures that skip_checkout is True).
-        skip_checkout: If True, the branch will not be switched (no pull will be performed too).
+        skip_checkout: If True, the branch will not be checked out (no pull will be performed too).
 
     Usage:
         > with agent_context(ctx, branch):
diff --git a/tasks/worktree.py b/tasks/worktree.py
index 57a9c90580293..bb00ed751b370 100644
--- a/tasks/worktree.py
+++ b/tasks/worktree.py
@@ -31,7 +31,7 @@ def status(ctx):
 
 
 @task
-def switch(ctx, ref):
+def checkout(ctx, ref):
     """Changes the worktree environment to the specified ref.
 
     Note: