diff --git a/agent/handlers/agentapi/taskprotection/factory.go b/agent/handlers/agentapi/taskprotection/factory.go
index 5e94ad4b362..b9bbca48e2f 100644
--- a/agent/handlers/agentapi/taskprotection/factory.go
+++ b/agent/handlers/agentapi/taskprotection/factory.go
@@ -18,7 +18,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/version"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/httpclient"
 
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	ecsapi "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 
@@ -37,7 +37,7 @@ type TaskProtectionClientFactory struct {
 // Helper function for retrieving credential from credentials manager and create ecs client
 func (factory TaskProtectionClientFactory) NewTaskProtectionClient(
 	taskRoleCredential credentials.TaskIAMRoleCredentials,
-) api.ECSTaskProtectionSDK {
+) ecsapi.ECSTaskProtectionSDK {
 	taskCredential := taskRoleCredential.GetIAMRoleCredentials()
 	cfg := aws.NewConfig().
 		WithCredentials(awscreds.NewStaticCredentials(taskCredential.AccessKeyID,
diff --git a/agent/handlers/agentapi/taskprotection/factory_test.go b/agent/handlers/agentapi/taskprotection/factory_test.go
index 122622729d7..a2f60975690 100644
--- a/agent/handlers/agentapi/taskprotection/factory_test.go
+++ b/agent/handlers/agentapi/taskprotection/factory_test.go
@@ -18,7 +18,7 @@ package taskprotection
 import (
 	"testing"
 
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
@@ -52,7 +52,7 @@ func TestGetECSClientHappyCase(t *testing.T) {
 	defer ctrl.Finish()
 
 	ret := factory.NewTaskProtectionClient(testIAMRoleCredentials)
-	_, ok := ret.(api.ECSTaskProtectionSDK)
+	_, ok := ret.(ecs.ECSTaskProtectionSDK)
 
 	// Assert response
 	assert.True(t, ok)
diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
index cf66da79c09..203a6240140 100644
--- a/agent/handlers/task_server_setup_test.go
+++ b/agent/handlers/task_server_setup_test.go
@@ -37,8 +37,8 @@ import (
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
 	mock_stats "github.com/aws/amazon-ecs-agent/agent/stats/mock"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
+	mock_taskprotection "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-	mock_taskprotection "github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/generate_mocks.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/generate_mocks.go
new file mode 100644
index 00000000000..012f038b3bc
--- /dev/null
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/generate_mocks.go
@@ -0,0 +1,17 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package ecs
+
+//go:generate mockgen -destination=mocks/api_mocks.go -copyright_file=../../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs ECSStandardSDK,ECSSubmitStateSDK,ECSClient,ECSTaskProtectionSDK
+//go:generate mockgen -destination=mocks/statechange/statechange_mocks.go -package=mock_statechange -copyright_file=../../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs ContainerMetadataGetter,TaskMetadataGetter
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/interface.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/interface.go
new file mode 100644
index 00000000000..218cdb9c8ee
--- /dev/null
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/interface.go
@@ -0,0 +1,91 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package ecs
+
+import (
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
+)
+
+// ECSClient is an interface over the ECSSDK interface which abstracts away some
+// details around constructing the request and reading the response down to the
+// parts the agent cares about.
+// For example, the ever-present 'Cluster' member is abstracted out so that it
+// may be configured once and used throughout transparently.
+type ECSClient interface {
+	// RegisterContainerInstance calculates the appropriate resources, creates
+	// the default cluster if necessary, and returns the registered
+	// ContainerInstanceARN if successful. Supplying a non-empty container
+	// instance ARN allows a container instance to update its registered
+	// resources.
+	RegisterContainerInstance(existingContainerInstanceArn string,
+		attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice,
+		outpostARN string) (string, string, error)
+	// SubmitTaskStateChange sends a state change and returns an error
+	// indicating if it was submitted
+	SubmitTaskStateChange(change TaskStateChange) error
+	// SubmitContainerStateChange sends a state change and returns an error
+	// indicating if it was submitted
+	SubmitContainerStateChange(change ContainerStateChange) error
+	// SubmitAttachmentStateChange sends an attachment state change and returns an error
+	// indicating if it was submitted
+	SubmitAttachmentStateChange(change AttachmentStateChange) error
+	// DiscoverPollEndpoint takes a ContainerInstanceARN and returns the
+	// endpoint at which this Agent should contact ACS
+	DiscoverPollEndpoint(containerInstanceArn string) (string, error)
+	// DiscoverTelemetryEndpoint takes a ContainerInstanceARN and returns the
+	// endpoint at which this Agent should contact Telemetry Service
+	DiscoverTelemetryEndpoint(containerInstanceArn string) (string, error)
+	// DiscoverServiceConnectEndpoint takes a ContainerInstanceARN and returns the
+	// endpoint at which this Agent should contact ServiceConnect
+	DiscoverServiceConnectEndpoint(containerInstanceArn string) (string, error)
+	// GetResourceTags retrieves the Tags associated with a certain resource
+	GetResourceTags(resourceArn string) ([]*ecs.Tag, error)
+	// UpdateContainerInstancesState updates the given container Instance ID with
+	// the given status. Only valid statuses are ACTIVE and DRAINING.
+	UpdateContainerInstancesState(instanceARN, status string) error
+	// GetHostResources retrieves a map that map the resource name to the corresponding resource
+	GetHostResources() (map[string]*ecs.Resource, error)
+}
+
+// ECSSDK is an interface that specifies the subset of the AWS Go SDK's ECS
+// client that the Agent uses.  This interface is meant to allow injecting a
+// mock for testing.
+type ECSStandardSDK interface {
+	CreateCluster(*ecs.CreateClusterInput) (*ecs.CreateClusterOutput, error)
+	RegisterContainerInstance(*ecs.RegisterContainerInstanceInput) (*ecs.RegisterContainerInstanceOutput, error)
+	DiscoverPollEndpoint(*ecs.DiscoverPollEndpointInput) (*ecs.DiscoverPollEndpointOutput, error)
+	ListTagsForResource(*ecs.ListTagsForResourceInput) (*ecs.ListTagsForResourceOutput, error)
+	UpdateContainerInstancesState(input *ecs.UpdateContainerInstancesStateInput) (*ecs.UpdateContainerInstancesStateOutput, error)
+}
+
+// ECSSubmitStateSDK is an interface with customized ecs client that
+// implements the SubmitTaskStateChange and SubmitContainerStateChange
+type ECSSubmitStateSDK interface {
+	SubmitContainerStateChange(*ecs.SubmitContainerStateChangeInput) (*ecs.SubmitContainerStateChangeOutput, error)
+	SubmitTaskStateChange(*ecs.SubmitTaskStateChangeInput) (*ecs.SubmitTaskStateChangeOutput, error)
+	SubmitAttachmentStateChanges(*ecs.SubmitAttachmentStateChangesInput) (*ecs.SubmitAttachmentStateChangesOutput, error)
+}
+
+// ECSTaskProtectionSDK is an interface with customized ecs client that
+// implements the UpdateTaskProtection and GetTaskProtection
+type ECSTaskProtectionSDK interface {
+	UpdateTaskProtection(input *ecs.UpdateTaskProtectionInput) (*ecs.UpdateTaskProtectionOutput, error)
+	UpdateTaskProtectionWithContext(ctx aws.Context, input *ecs.UpdateTaskProtectionInput,
+		opts ...request.Option) (*ecs.UpdateTaskProtectionOutput, error)
+	GetTaskProtection(input *ecs.GetTaskProtectionInput) (*ecs.GetTaskProtectionOutput, error)
+	GetTaskProtectionWithContext(ctx aws.Context, input *ecs.GetTaskProtectionInput,
+		opts ...request.Option) (*ecs.GetTaskProtectionOutput, error)
+}
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks/api_mocks.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks/api_mocks.go
new file mode 100644
index 00000000000..ec4f752661d
--- /dev/null
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks/api_mocks.go
@@ -0,0 +1,458 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//     http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+//
+
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs (interfaces: ECSStandardSDK,ECSSubmitStateSDK,ECSClient,ECSTaskProtectionSDK)
+
+// Package mock_ecs is a generated GoMock package.
+package mock_ecs
+
+import (
+	context "context"
+	reflect "reflect"
+
+	ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
+	ecs0 "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
+	request "github.com/aws/aws-sdk-go/aws/request"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockECSStandardSDK is a mock of ECSStandardSDK interface.
+type MockECSStandardSDK struct {
+	ctrl     *gomock.Controller
+	recorder *MockECSStandardSDKMockRecorder
+}
+
+// MockECSStandardSDKMockRecorder is the mock recorder for MockECSStandardSDK.
+type MockECSStandardSDKMockRecorder struct {
+	mock *MockECSStandardSDK
+}
+
+// NewMockECSStandardSDK creates a new mock instance.
+func NewMockECSStandardSDK(ctrl *gomock.Controller) *MockECSStandardSDK {
+	mock := &MockECSStandardSDK{ctrl: ctrl}
+	mock.recorder = &MockECSStandardSDKMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockECSStandardSDK) EXPECT() *MockECSStandardSDKMockRecorder {
+	return m.recorder
+}
+
+// CreateCluster mocks base method.
+func (m *MockECSStandardSDK) CreateCluster(arg0 *ecs0.CreateClusterInput) (*ecs0.CreateClusterOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CreateCluster", arg0)
+	ret0, _ := ret[0].(*ecs0.CreateClusterOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CreateCluster indicates an expected call of CreateCluster.
+func (mr *MockECSStandardSDKMockRecorder) CreateCluster(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateCluster", reflect.TypeOf((*MockECSStandardSDK)(nil).CreateCluster), arg0)
+}
+
+// DiscoverPollEndpoint mocks base method.
+func (m *MockECSStandardSDK) DiscoverPollEndpoint(arg0 *ecs0.DiscoverPollEndpointInput) (*ecs0.DiscoverPollEndpointOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DiscoverPollEndpoint", arg0)
+	ret0, _ := ret[0].(*ecs0.DiscoverPollEndpointOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DiscoverPollEndpoint indicates an expected call of DiscoverPollEndpoint.
+func (mr *MockECSStandardSDKMockRecorder) DiscoverPollEndpoint(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverPollEndpoint", reflect.TypeOf((*MockECSStandardSDK)(nil).DiscoverPollEndpoint), arg0)
+}
+
+// ListTagsForResource mocks base method.
+func (m *MockECSStandardSDK) ListTagsForResource(arg0 *ecs0.ListTagsForResourceInput) (*ecs0.ListTagsForResourceOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ListTagsForResource", arg0)
+	ret0, _ := ret[0].(*ecs0.ListTagsForResourceOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ListTagsForResource indicates an expected call of ListTagsForResource.
+func (mr *MockECSStandardSDKMockRecorder) ListTagsForResource(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTagsForResource", reflect.TypeOf((*MockECSStandardSDK)(nil).ListTagsForResource), arg0)
+}
+
+// RegisterContainerInstance mocks base method.
+func (m *MockECSStandardSDK) RegisterContainerInstance(arg0 *ecs0.RegisterContainerInstanceInput) (*ecs0.RegisterContainerInstanceOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RegisterContainerInstance", arg0)
+	ret0, _ := ret[0].(*ecs0.RegisterContainerInstanceOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// RegisterContainerInstance indicates an expected call of RegisterContainerInstance.
+func (mr *MockECSStandardSDKMockRecorder) RegisterContainerInstance(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterContainerInstance", reflect.TypeOf((*MockECSStandardSDK)(nil).RegisterContainerInstance), arg0)
+}
+
+// UpdateContainerInstancesState mocks base method.
+func (m *MockECSStandardSDK) UpdateContainerInstancesState(arg0 *ecs0.UpdateContainerInstancesStateInput) (*ecs0.UpdateContainerInstancesStateOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateContainerInstancesState", arg0)
+	ret0, _ := ret[0].(*ecs0.UpdateContainerInstancesStateOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateContainerInstancesState indicates an expected call of UpdateContainerInstancesState.
+func (mr *MockECSStandardSDKMockRecorder) UpdateContainerInstancesState(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateContainerInstancesState", reflect.TypeOf((*MockECSStandardSDK)(nil).UpdateContainerInstancesState), arg0)
+}
+
+// MockECSSubmitStateSDK is a mock of ECSSubmitStateSDK interface.
+type MockECSSubmitStateSDK struct {
+	ctrl     *gomock.Controller
+	recorder *MockECSSubmitStateSDKMockRecorder
+}
+
+// MockECSSubmitStateSDKMockRecorder is the mock recorder for MockECSSubmitStateSDK.
+type MockECSSubmitStateSDKMockRecorder struct {
+	mock *MockECSSubmitStateSDK
+}
+
+// NewMockECSSubmitStateSDK creates a new mock instance.
+func NewMockECSSubmitStateSDK(ctrl *gomock.Controller) *MockECSSubmitStateSDK {
+	mock := &MockECSSubmitStateSDK{ctrl: ctrl}
+	mock.recorder = &MockECSSubmitStateSDKMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockECSSubmitStateSDK) EXPECT() *MockECSSubmitStateSDKMockRecorder {
+	return m.recorder
+}
+
+// SubmitAttachmentStateChanges mocks base method.
+func (m *MockECSSubmitStateSDK) SubmitAttachmentStateChanges(arg0 *ecs0.SubmitAttachmentStateChangesInput) (*ecs0.SubmitAttachmentStateChangesOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitAttachmentStateChanges", arg0)
+	ret0, _ := ret[0].(*ecs0.SubmitAttachmentStateChangesOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SubmitAttachmentStateChanges indicates an expected call of SubmitAttachmentStateChanges.
+func (mr *MockECSSubmitStateSDKMockRecorder) SubmitAttachmentStateChanges(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitAttachmentStateChanges", reflect.TypeOf((*MockECSSubmitStateSDK)(nil).SubmitAttachmentStateChanges), arg0)
+}
+
+// SubmitContainerStateChange mocks base method.
+func (m *MockECSSubmitStateSDK) SubmitContainerStateChange(arg0 *ecs0.SubmitContainerStateChangeInput) (*ecs0.SubmitContainerStateChangeOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitContainerStateChange", arg0)
+	ret0, _ := ret[0].(*ecs0.SubmitContainerStateChangeOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SubmitContainerStateChange indicates an expected call of SubmitContainerStateChange.
+func (mr *MockECSSubmitStateSDKMockRecorder) SubmitContainerStateChange(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitContainerStateChange", reflect.TypeOf((*MockECSSubmitStateSDK)(nil).SubmitContainerStateChange), arg0)
+}
+
+// SubmitTaskStateChange mocks base method.
+func (m *MockECSSubmitStateSDK) SubmitTaskStateChange(arg0 *ecs0.SubmitTaskStateChangeInput) (*ecs0.SubmitTaskStateChangeOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitTaskStateChange", arg0)
+	ret0, _ := ret[0].(*ecs0.SubmitTaskStateChangeOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SubmitTaskStateChange indicates an expected call of SubmitTaskStateChange.
+func (mr *MockECSSubmitStateSDKMockRecorder) SubmitTaskStateChange(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitTaskStateChange", reflect.TypeOf((*MockECSSubmitStateSDK)(nil).SubmitTaskStateChange), arg0)
+}
+
+// MockECSClient is a mock of ECSClient interface.
+type MockECSClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockECSClientMockRecorder
+}
+
+// MockECSClientMockRecorder is the mock recorder for MockECSClient.
+type MockECSClientMockRecorder struct {
+	mock *MockECSClient
+}
+
+// NewMockECSClient creates a new mock instance.
+func NewMockECSClient(ctrl *gomock.Controller) *MockECSClient {
+	mock := &MockECSClient{ctrl: ctrl}
+	mock.recorder = &MockECSClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockECSClient) EXPECT() *MockECSClientMockRecorder {
+	return m.recorder
+}
+
+// DiscoverPollEndpoint mocks base method.
+func (m *MockECSClient) DiscoverPollEndpoint(arg0 string) (string, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DiscoverPollEndpoint", arg0)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DiscoverPollEndpoint indicates an expected call of DiscoverPollEndpoint.
+func (mr *MockECSClientMockRecorder) DiscoverPollEndpoint(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverPollEndpoint", reflect.TypeOf((*MockECSClient)(nil).DiscoverPollEndpoint), arg0)
+}
+
+// DiscoverServiceConnectEndpoint mocks base method.
+func (m *MockECSClient) DiscoverServiceConnectEndpoint(arg0 string) (string, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DiscoverServiceConnectEndpoint", arg0)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DiscoverServiceConnectEndpoint indicates an expected call of DiscoverServiceConnectEndpoint.
+func (mr *MockECSClientMockRecorder) DiscoverServiceConnectEndpoint(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverServiceConnectEndpoint", reflect.TypeOf((*MockECSClient)(nil).DiscoverServiceConnectEndpoint), arg0)
+}
+
+// DiscoverTelemetryEndpoint mocks base method.
+func (m *MockECSClient) DiscoverTelemetryEndpoint(arg0 string) (string, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DiscoverTelemetryEndpoint", arg0)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DiscoverTelemetryEndpoint indicates an expected call of DiscoverTelemetryEndpoint.
+func (mr *MockECSClientMockRecorder) DiscoverTelemetryEndpoint(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverTelemetryEndpoint", reflect.TypeOf((*MockECSClient)(nil).DiscoverTelemetryEndpoint), arg0)
+}
+
+// GetHostResources mocks base method.
+func (m *MockECSClient) GetHostResources() (map[string]*ecs0.Resource, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetHostResources")
+	ret0, _ := ret[0].(map[string]*ecs0.Resource)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetHostResources indicates an expected call of GetHostResources.
+func (mr *MockECSClientMockRecorder) GetHostResources() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHostResources", reflect.TypeOf((*MockECSClient)(nil).GetHostResources))
+}
+
+// GetResourceTags mocks base method.
+func (m *MockECSClient) GetResourceTags(arg0 string) ([]*ecs0.Tag, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetResourceTags", arg0)
+	ret0, _ := ret[0].([]*ecs0.Tag)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetResourceTags indicates an expected call of GetResourceTags.
+func (mr *MockECSClientMockRecorder) GetResourceTags(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetResourceTags", reflect.TypeOf((*MockECSClient)(nil).GetResourceTags), arg0)
+}
+
+// RegisterContainerInstance mocks base method.
+func (m *MockECSClient) RegisterContainerInstance(arg0 string, arg1 []*ecs0.Attribute, arg2 []*ecs0.Tag, arg3 string, arg4 []*ecs0.PlatformDevice, arg5 string) (string, string, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RegisterContainerInstance", arg0, arg1, arg2, arg3, arg4, arg5)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(string)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
+}
+
+// RegisterContainerInstance indicates an expected call of RegisterContainerInstance.
+func (mr *MockECSClientMockRecorder) RegisterContainerInstance(arg0, arg1, arg2, arg3, arg4, arg5 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterContainerInstance", reflect.TypeOf((*MockECSClient)(nil).RegisterContainerInstance), arg0, arg1, arg2, arg3, arg4, arg5)
+}
+
+// SubmitAttachmentStateChange mocks base method.
+func (m *MockECSClient) SubmitAttachmentStateChange(arg0 ecs.AttachmentStateChange) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitAttachmentStateChange", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SubmitAttachmentStateChange indicates an expected call of SubmitAttachmentStateChange.
+func (mr *MockECSClientMockRecorder) SubmitAttachmentStateChange(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitAttachmentStateChange", reflect.TypeOf((*MockECSClient)(nil).SubmitAttachmentStateChange), arg0)
+}
+
+// SubmitContainerStateChange mocks base method.
+func (m *MockECSClient) SubmitContainerStateChange(arg0 ecs.ContainerStateChange) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitContainerStateChange", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SubmitContainerStateChange indicates an expected call of SubmitContainerStateChange.
+func (mr *MockECSClientMockRecorder) SubmitContainerStateChange(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitContainerStateChange", reflect.TypeOf((*MockECSClient)(nil).SubmitContainerStateChange), arg0)
+}
+
+// SubmitTaskStateChange mocks base method.
+func (m *MockECSClient) SubmitTaskStateChange(arg0 ecs.TaskStateChange) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SubmitTaskStateChange", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SubmitTaskStateChange indicates an expected call of SubmitTaskStateChange.
+func (mr *MockECSClientMockRecorder) SubmitTaskStateChange(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubmitTaskStateChange", reflect.TypeOf((*MockECSClient)(nil).SubmitTaskStateChange), arg0)
+}
+
+// UpdateContainerInstancesState mocks base method.
+func (m *MockECSClient) UpdateContainerInstancesState(arg0, arg1 string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateContainerInstancesState", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateContainerInstancesState indicates an expected call of UpdateContainerInstancesState.
+func (mr *MockECSClientMockRecorder) UpdateContainerInstancesState(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateContainerInstancesState", reflect.TypeOf((*MockECSClient)(nil).UpdateContainerInstancesState), arg0, arg1)
+}
+
+// MockECSTaskProtectionSDK is a mock of ECSTaskProtectionSDK interface.
+type MockECSTaskProtectionSDK struct {
+	ctrl     *gomock.Controller
+	recorder *MockECSTaskProtectionSDKMockRecorder
+}
+
+// MockECSTaskProtectionSDKMockRecorder is the mock recorder for MockECSTaskProtectionSDK.
+type MockECSTaskProtectionSDKMockRecorder struct {
+	mock *MockECSTaskProtectionSDK
+}
+
+// NewMockECSTaskProtectionSDK creates a new mock instance.
+func NewMockECSTaskProtectionSDK(ctrl *gomock.Controller) *MockECSTaskProtectionSDK {
+	mock := &MockECSTaskProtectionSDK{ctrl: ctrl}
+	mock.recorder = &MockECSTaskProtectionSDKMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockECSTaskProtectionSDK) EXPECT() *MockECSTaskProtectionSDKMockRecorder {
+	return m.recorder
+}
+
+// GetTaskProtection mocks base method.
+func (m *MockECSTaskProtectionSDK) GetTaskProtection(arg0 *ecs0.GetTaskProtectionInput) (*ecs0.GetTaskProtectionOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetTaskProtection", arg0)
+	ret0, _ := ret[0].(*ecs0.GetTaskProtectionOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetTaskProtection indicates an expected call of GetTaskProtection.
+func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtection(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtection), arg0)
+}
+
+// GetTaskProtectionWithContext mocks base method.
+func (m *MockECSTaskProtectionSDK) GetTaskProtectionWithContext(arg0 context.Context, arg1 *ecs0.GetTaskProtectionInput, arg2 ...request.Option) (*ecs0.GetTaskProtectionOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetTaskProtectionWithContext", varargs...)
+	ret0, _ := ret[0].(*ecs0.GetTaskProtectionOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetTaskProtectionWithContext indicates an expected call of GetTaskProtectionWithContext.
+func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtectionWithContext), varargs...)
+}
+
+// UpdateTaskProtection mocks base method.
+func (m *MockECSTaskProtectionSDK) UpdateTaskProtection(arg0 *ecs0.UpdateTaskProtectionInput) (*ecs0.UpdateTaskProtectionOutput, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateTaskProtection", arg0)
+	ret0, _ := ret[0].(*ecs0.UpdateTaskProtectionOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateTaskProtection indicates an expected call of UpdateTaskProtection.
+func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtection(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtection), arg0)
+}
+
+// UpdateTaskProtectionWithContext mocks base method.
+func (m *MockECSTaskProtectionSDK) UpdateTaskProtectionWithContext(arg0 context.Context, arg1 *ecs0.UpdateTaskProtectionInput, arg2 ...request.Option) (*ecs0.UpdateTaskProtectionOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UpdateTaskProtectionWithContext", varargs...)
+	ret0, _ := ret[0].(*ecs0.UpdateTaskProtectionOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateTaskProtectionWithContext indicates an expected call of UpdateTaskProtectionWithContext.
+func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtectionWithContext), varargs...)
+}
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/statechange.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/statechange.go
new file mode 100644
index 00000000000..076926bd368
--- /dev/null
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/statechange.go
@@ -0,0 +1,154 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package ecs
+
+import (
+	"fmt"
+	"strconv"
+	"time"
+
+	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
+	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
+	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
+)
+
+// ContainerMetadataGetter retrieves specific information about a given container that ECS client is concerned with.
+type ContainerMetadataGetter interface {
+	GetContainerIsNil() bool
+	GetContainerSentStatusString() string
+	GetContainerRuntimeID() string
+	GetContainerIsEssential() bool
+}
+
+// TaskMetadataGetter retrieves specific information about a given task that ECS client is concerned with.
+type TaskMetadataGetter interface {
+	GetTaskIsNil() bool
+	GetTaskSentStatusString() string
+	GetTaskPullStartedAt() time.Time
+	GetTaskPullStoppedAt() time.Time
+	GetTaskExecutionStoppedAt() time.Time
+}
+
+// ContainerStateChange represents a state change that needs to be sent to the
+// SubmitContainerStateChange API.
+type ContainerStateChange struct {
+	// TaskArn is the unique identifier for the task.
+	TaskArn string
+	// RuntimeID is the dockerID of the container.
+	RuntimeID string
+	// ContainerName is the name of the container.
+	ContainerName string
+	// Status is the status to send.
+	Status apicontainerstatus.ContainerStatus
+	// ImageDigest is the sha-256 digest of the container image as pulled from the
+	// repository.
+	ImageDigest string
+	// Reason may contain details of why the container stopped.
+	Reason string
+	// ExitCode is the exit code of the container, if available.
+	ExitCode *int
+	// NetworkBindings contains the details of the host ports picked for the specified
+	// container ports.
+	NetworkBindings []*ecs.NetworkBinding
+	// MetadataGetter is used to retrieve other relevant information about the
+	// container.
+	MetadataGetter ContainerMetadataGetter
+}
+
+// TaskStateChange represents a state change that needs to be sent to the
+// SubmitTaskStateChange API.
+type TaskStateChange struct {
+	// Attachment is the ENI attachment object to send.
+	Attachment *ni.ENIAttachment
+	// TaskArn is the unique identifier for the task.
+	TaskARN string
+	// Status is the status to send.
+	Status apitaskstatus.TaskStatus
+	// Reason may contain details of why the task stopped.
+	Reason string
+	// Containers holds the events generated by containers owned by this task.
+	Containers []*ecs.ContainerStateChange
+	// ManagedAgents contain the name and status of Agents running inside the
+	// container.
+	ManagedAgents []*ecs.ManagedAgentStateChange
+	// PullStartedAt is the timestamp when the task start pulling.
+	PullStartedAt *time.Time
+	// PullStoppedAt is the timestamp when the task finished pulling.
+	PullStoppedAt *time.Time
+	// ExecutionStoppedAt is the timestamp when the essential container stopped.
+	ExecutionStoppedAt *time.Time
+	// MetadataGetter is used to retrieve other relevant information about the task.
+	MetadataGetter TaskMetadataGetter
+}
+
+// AttachmentStateChange represents a state change that needs to be sent to the
+// SubmitAttachmentStateChanges API.
+type AttachmentStateChange struct {
+	// Attachment is the ENI attachment object to send.
+	Attachment *ni.ENIAttachment
+}
+
+// String returns a human readable string representation of a ContainerStateChange.
+func (c *ContainerStateChange) String() string {
+	res := fmt.Sprintf("containerName=%s containerStatus=%s", c.ContainerName, c.Status.String())
+	if c.ExitCode != nil {
+		res += " containerExitCode=" + strconv.Itoa(*c.ExitCode)
+	}
+	if c.Reason != "" {
+		res += " containerReason=" + c.Reason
+	}
+	if len(c.NetworkBindings) != 0 {
+		res += fmt.Sprintf(" containerNetworkBindings=%v", c.NetworkBindings)
+	}
+	if c.MetadataGetter != nil && !c.MetadataGetter.GetContainerIsNil() {
+		res += fmt.Sprintf(" containerKnownSentStatus=%s containerRuntimeID=%s containerIsEssential=%v",
+			c.MetadataGetter.GetContainerSentStatusString(), c.MetadataGetter.GetContainerRuntimeID(),
+			c.MetadataGetter.GetContainerIsEssential())
+	}
+	return res
+}
+
+// String returns a human readable string representation of a TaskStateChange.
+func (change *TaskStateChange) String() string {
+	res := fmt.Sprintf("%s -> %s", change.TaskARN, change.Status.String())
+	if change.MetadataGetter != nil && !change.MetadataGetter.GetTaskIsNil() {
+		res += fmt.Sprintf(", Known Sent: %s, PullStartedAt: %s, PullStoppedAt: %s, ExecutionStoppedAt: %s",
+			change.MetadataGetter.GetTaskSentStatusString(),
+			change.MetadataGetter.GetTaskPullStartedAt(),
+			change.MetadataGetter.GetTaskPullStoppedAt(),
+			change.MetadataGetter.GetTaskExecutionStoppedAt())
+	}
+	if change.Attachment != nil {
+		res += ", " + change.Attachment.String()
+	}
+	for _, containerChange := range change.Containers {
+		res += ", container change: " + containerChange.String()
+	}
+	for _, managedAgentChange := range change.ManagedAgents {
+		res += ", managed agent: " + managedAgentChange.String()
+	}
+
+	return res
+}
+
+// String returns a human readable string representation of an AttachmentStateChange.
+func (change *AttachmentStateChange) String() string {
+	if change.Attachment != nil {
+		return fmt.Sprintf("%s -> %s, %s", change.Attachment.AttachmentARN, change.Attachment.Status.String(),
+			change.Attachment.String())
+	}
+
+	return ""
+}
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/generate_mocks.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/generate_mocks.go
index e21481704df..57e12a724fc 100644
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/generate_mocks.go
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/generate_mocks.go
@@ -11,6 +11,6 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
-//go:generate mockgen -destination=mocks/api_mocks.go -copyright_file=../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api ECSTaskProtectionSDK,ECSDiscoverEndpointSDK
+//go:generate mockgen -destination=mocks/api_mocks.go -copyright_file=../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api ECSDiscoverEndpointSDK
 
 package api
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/interface.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/interface.go
index a673293b79e..22324f0a367 100644
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/interface.go
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/interface.go
@@ -12,24 +12,6 @@
 // permissions and limitations under the License.
 package api
 
-import (
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/request"
-)
-
-// ECSTaskProtectionSDK is an interface with customized ecs client that
-// implements the UpdateTaskProtection and GetTaskProtection
-type ECSTaskProtectionSDK interface {
-	UpdateTaskProtection(input *ecs.UpdateTaskProtectionInput) (*ecs.UpdateTaskProtectionOutput, error)
-	UpdateTaskProtectionWithContext(ctx aws.Context, input *ecs.UpdateTaskProtectionInput,
-		opts ...request.Option) (*ecs.UpdateTaskProtectionOutput, error)
-	GetTaskProtection(input *ecs.GetTaskProtectionInput) (*ecs.GetTaskProtectionOutput, error)
-	GetTaskProtectionWithContext(ctx aws.Context, input *ecs.GetTaskProtectionInput,
-		opts ...request.Option) (*ecs.GetTaskProtectionOutput, error)
-}
-
 // ECSDiscoverEndpointSDK is an interface with customized ecs client that
 // implements the DiscoverPollEndpoint, DiscoverTelemetryEndpoint, and DiscoverServiceConnectEndpoint
 type ECSDiscoverEndpointSDK interface {
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks/api_mocks.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks/api_mocks.go
deleted file mode 100644
index 2c4f8b62894..00000000000
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks/api_mocks.go
+++ /dev/null
@@ -1,189 +0,0 @@
-// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//     http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-//
-
-// Code generated by MockGen. DO NOT EDIT.
-// Source: github.com/aws/amazon-ecs-agent/ecs-agent/api (interfaces: ECSTaskProtectionSDK,ECSDiscoverEndpointSDK)
-
-// Package mock_api is a generated GoMock package.
-package mock_api
-
-import (
-	context "context"
-	reflect "reflect"
-
-	ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-	request "github.com/aws/aws-sdk-go/aws/request"
-	gomock "github.com/golang/mock/gomock"
-)
-
-// MockECSTaskProtectionSDK is a mock of ECSTaskProtectionSDK interface.
-type MockECSTaskProtectionSDK struct {
-	ctrl     *gomock.Controller
-	recorder *MockECSTaskProtectionSDKMockRecorder
-}
-
-// MockECSTaskProtectionSDKMockRecorder is the mock recorder for MockECSTaskProtectionSDK.
-type MockECSTaskProtectionSDKMockRecorder struct {
-	mock *MockECSTaskProtectionSDK
-}
-
-// NewMockECSTaskProtectionSDK creates a new mock instance.
-func NewMockECSTaskProtectionSDK(ctrl *gomock.Controller) *MockECSTaskProtectionSDK {
-	mock := &MockECSTaskProtectionSDK{ctrl: ctrl}
-	mock.recorder = &MockECSTaskProtectionSDKMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use.
-func (m *MockECSTaskProtectionSDK) EXPECT() *MockECSTaskProtectionSDKMockRecorder {
-	return m.recorder
-}
-
-// GetTaskProtection mocks base method.
-func (m *MockECSTaskProtectionSDK) GetTaskProtection(arg0 *ecs.GetTaskProtectionInput) (*ecs.GetTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetTaskProtection", arg0)
-	ret0, _ := ret[0].(*ecs.GetTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetTaskProtection indicates an expected call of GetTaskProtection.
-func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtection(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtection), arg0)
-}
-
-// GetTaskProtectionWithContext mocks base method.
-func (m *MockECSTaskProtectionSDK) GetTaskProtectionWithContext(arg0 context.Context, arg1 *ecs.GetTaskProtectionInput, arg2 ...request.Option) (*ecs.GetTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	varargs := []interface{}{arg0, arg1}
-	for _, a := range arg2 {
-		varargs = append(varargs, a)
-	}
-	ret := m.ctrl.Call(m, "GetTaskProtectionWithContext", varargs...)
-	ret0, _ := ret[0].(*ecs.GetTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetTaskProtectionWithContext indicates an expected call of GetTaskProtectionWithContext.
-func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtectionWithContext), varargs...)
-}
-
-// UpdateTaskProtection mocks base method.
-func (m *MockECSTaskProtectionSDK) UpdateTaskProtection(arg0 *ecs.UpdateTaskProtectionInput) (*ecs.UpdateTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpdateTaskProtection", arg0)
-	ret0, _ := ret[0].(*ecs.UpdateTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// UpdateTaskProtection indicates an expected call of UpdateTaskProtection.
-func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtection(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtection), arg0)
-}
-
-// UpdateTaskProtectionWithContext mocks base method.
-func (m *MockECSTaskProtectionSDK) UpdateTaskProtectionWithContext(arg0 context.Context, arg1 *ecs.UpdateTaskProtectionInput, arg2 ...request.Option) (*ecs.UpdateTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	varargs := []interface{}{arg0, arg1}
-	for _, a := range arg2 {
-		varargs = append(varargs, a)
-	}
-	ret := m.ctrl.Call(m, "UpdateTaskProtectionWithContext", varargs...)
-	ret0, _ := ret[0].(*ecs.UpdateTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// UpdateTaskProtectionWithContext indicates an expected call of UpdateTaskProtectionWithContext.
-func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtectionWithContext), varargs...)
-}
-
-// MockECSDiscoverEndpointSDK is a mock of ECSDiscoverEndpointSDK interface.
-type MockECSDiscoverEndpointSDK struct {
-	ctrl     *gomock.Controller
-	recorder *MockECSDiscoverEndpointSDKMockRecorder
-}
-
-// MockECSDiscoverEndpointSDKMockRecorder is the mock recorder for MockECSDiscoverEndpointSDK.
-type MockECSDiscoverEndpointSDKMockRecorder struct {
-	mock *MockECSDiscoverEndpointSDK
-}
-
-// NewMockECSDiscoverEndpointSDK creates a new mock instance.
-func NewMockECSDiscoverEndpointSDK(ctrl *gomock.Controller) *MockECSDiscoverEndpointSDK {
-	mock := &MockECSDiscoverEndpointSDK{ctrl: ctrl}
-	mock.recorder = &MockECSDiscoverEndpointSDKMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use.
-func (m *MockECSDiscoverEndpointSDK) EXPECT() *MockECSDiscoverEndpointSDKMockRecorder {
-	return m.recorder
-}
-
-// DiscoverPollEndpoint mocks base method.
-func (m *MockECSDiscoverEndpointSDK) DiscoverPollEndpoint(arg0 string) (string, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DiscoverPollEndpoint", arg0)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// DiscoverPollEndpoint indicates an expected call of DiscoverPollEndpoint.
-func (mr *MockECSDiscoverEndpointSDKMockRecorder) DiscoverPollEndpoint(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverPollEndpoint", reflect.TypeOf((*MockECSDiscoverEndpointSDK)(nil).DiscoverPollEndpoint), arg0)
-}
-
-// DiscoverServiceConnectEndpoint mocks base method.
-func (m *MockECSDiscoverEndpointSDK) DiscoverServiceConnectEndpoint(arg0 string) (string, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DiscoverServiceConnectEndpoint", arg0)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// DiscoverServiceConnectEndpoint indicates an expected call of DiscoverServiceConnectEndpoint.
-func (mr *MockECSDiscoverEndpointSDKMockRecorder) DiscoverServiceConnectEndpoint(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverServiceConnectEndpoint", reflect.TypeOf((*MockECSDiscoverEndpointSDK)(nil).DiscoverServiceConnectEndpoint), arg0)
-}
-
-// DiscoverTelemetryEndpoint mocks base method.
-func (m *MockECSDiscoverEndpointSDK) DiscoverTelemetryEndpoint(arg0 string) (string, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DiscoverTelemetryEndpoint", arg0)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// DiscoverTelemetryEndpoint indicates an expected call of DiscoverTelemetryEndpoint.
-func (mr *MockECSDiscoverEndpointSDKMockRecorder) DiscoverTelemetryEndpoint(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverTelemetryEndpoint", reflect.TypeOf((*MockECSDiscoverEndpointSDK)(nil).DiscoverTelemetryEndpoint), arg0)
-}
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
index 22dc4c1fc2d..4dcc9b837cc 100644
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
@@ -21,7 +21,7 @@ package handlers
 import (
 	reflect "reflect"
 
-	api "github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	gomock "github.com/golang/mock/gomock"
 )
@@ -50,10 +50,10 @@ func (m *MockTaskProtectionClientFactoryInterface) EXPECT() *MockTaskProtectionC
 }
 
 // NewTaskProtectionClient mocks base method.
-func (m *MockTaskProtectionClientFactoryInterface) NewTaskProtectionClient(arg0 credentials.TaskIAMRoleCredentials) api.ECSTaskProtectionSDK {
+func (m *MockTaskProtectionClientFactoryInterface) NewTaskProtectionClient(arg0 credentials.TaskIAMRoleCredentials) ecs.ECSTaskProtectionSDK {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "NewTaskProtectionClient", arg0)
-	ret0, _ := ret[0].(api.ECSTaskProtectionSDK)
+	ret0, _ := ret[0].(ecs.ECSTaskProtectionSDK)
 	return ret0
 }
 
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
index 02c8c4c4dd8..6adb7ac308e 100644
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
@@ -1,10 +1,10 @@
 package handlers
 
 import (
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 )
 
 type TaskProtectionClientFactoryInterface interface {
-	NewTaskProtectionClient(taskRoleCredential credentials.TaskIAMRoleCredentials) api.ECSTaskProtectionSDK
+	NewTaskProtectionClient(taskRoleCredential credentials.TaskIAMRoleCredentials) ecs.ECSTaskProtectionSDK
 }
diff --git a/agent/vendor/modules.txt b/agent/vendor/modules.txt
index 538a5b5d6b4..2d536767a17 100644
--- a/agent/vendor/modules.txt
+++ b/agent/vendor/modules.txt
@@ -16,9 +16,10 @@ github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment
 github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource
 github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource/mocks
 github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status
+github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs
+github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks
 github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs
 github.com/aws/amazon-ecs-agent/ecs-agent/api/errors
-github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks
 github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status
 github.com/aws/amazon-ecs-agent/ecs-agent/async
 github.com/aws/amazon-ecs-agent/ecs-agent/async/mocks
diff --git a/ecs-agent/api/generate_mocks.go b/ecs-agent/api/generate_mocks.go
index e21481704df..57e12a724fc 100644
--- a/ecs-agent/api/generate_mocks.go
+++ b/ecs-agent/api/generate_mocks.go
@@ -11,6 +11,6 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
-//go:generate mockgen -destination=mocks/api_mocks.go -copyright_file=../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api ECSTaskProtectionSDK,ECSDiscoverEndpointSDK
+//go:generate mockgen -destination=mocks/api_mocks.go -copyright_file=../../scripts/copyright_file github.com/aws/amazon-ecs-agent/ecs-agent/api ECSDiscoverEndpointSDK
 
 package api
diff --git a/ecs-agent/api/interface.go b/ecs-agent/api/interface.go
index a673293b79e..22324f0a367 100644
--- a/ecs-agent/api/interface.go
+++ b/ecs-agent/api/interface.go
@@ -12,24 +12,6 @@
 // permissions and limitations under the License.
 package api
 
-import (
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/request"
-)
-
-// ECSTaskProtectionSDK is an interface with customized ecs client that
-// implements the UpdateTaskProtection and GetTaskProtection
-type ECSTaskProtectionSDK interface {
-	UpdateTaskProtection(input *ecs.UpdateTaskProtectionInput) (*ecs.UpdateTaskProtectionOutput, error)
-	UpdateTaskProtectionWithContext(ctx aws.Context, input *ecs.UpdateTaskProtectionInput,
-		opts ...request.Option) (*ecs.UpdateTaskProtectionOutput, error)
-	GetTaskProtection(input *ecs.GetTaskProtectionInput) (*ecs.GetTaskProtectionOutput, error)
-	GetTaskProtectionWithContext(ctx aws.Context, input *ecs.GetTaskProtectionInput,
-		opts ...request.Option) (*ecs.GetTaskProtectionOutput, error)
-}
-
 // ECSDiscoverEndpointSDK is an interface with customized ecs client that
 // implements the DiscoverPollEndpoint, DiscoverTelemetryEndpoint, and DiscoverServiceConnectEndpoint
 type ECSDiscoverEndpointSDK interface {
diff --git a/ecs-agent/api/mocks/api_mocks.go b/ecs-agent/api/mocks/api_mocks.go
index 2c4f8b62894..f2a57c4d1c5 100644
--- a/ecs-agent/api/mocks/api_mocks.go
+++ b/ecs-agent/api/mocks/api_mocks.go
@@ -13,113 +13,17 @@
 //
 
 // Code generated by MockGen. DO NOT EDIT.
-// Source: github.com/aws/amazon-ecs-agent/ecs-agent/api (interfaces: ECSTaskProtectionSDK,ECSDiscoverEndpointSDK)
+// Source: github.com/aws/amazon-ecs-agent/ecs-agent/api (interfaces: ECSDiscoverEndpointSDK)
 
 // Package mock_api is a generated GoMock package.
 package mock_api
 
 import (
-	context "context"
 	reflect "reflect"
 
-	ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-	request "github.com/aws/aws-sdk-go/aws/request"
 	gomock "github.com/golang/mock/gomock"
 )
 
-// MockECSTaskProtectionSDK is a mock of ECSTaskProtectionSDK interface.
-type MockECSTaskProtectionSDK struct {
-	ctrl     *gomock.Controller
-	recorder *MockECSTaskProtectionSDKMockRecorder
-}
-
-// MockECSTaskProtectionSDKMockRecorder is the mock recorder for MockECSTaskProtectionSDK.
-type MockECSTaskProtectionSDKMockRecorder struct {
-	mock *MockECSTaskProtectionSDK
-}
-
-// NewMockECSTaskProtectionSDK creates a new mock instance.
-func NewMockECSTaskProtectionSDK(ctrl *gomock.Controller) *MockECSTaskProtectionSDK {
-	mock := &MockECSTaskProtectionSDK{ctrl: ctrl}
-	mock.recorder = &MockECSTaskProtectionSDKMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use.
-func (m *MockECSTaskProtectionSDK) EXPECT() *MockECSTaskProtectionSDKMockRecorder {
-	return m.recorder
-}
-
-// GetTaskProtection mocks base method.
-func (m *MockECSTaskProtectionSDK) GetTaskProtection(arg0 *ecs.GetTaskProtectionInput) (*ecs.GetTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetTaskProtection", arg0)
-	ret0, _ := ret[0].(*ecs.GetTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetTaskProtection indicates an expected call of GetTaskProtection.
-func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtection(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtection), arg0)
-}
-
-// GetTaskProtectionWithContext mocks base method.
-func (m *MockECSTaskProtectionSDK) GetTaskProtectionWithContext(arg0 context.Context, arg1 *ecs.GetTaskProtectionInput, arg2 ...request.Option) (*ecs.GetTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	varargs := []interface{}{arg0, arg1}
-	for _, a := range arg2 {
-		varargs = append(varargs, a)
-	}
-	ret := m.ctrl.Call(m, "GetTaskProtectionWithContext", varargs...)
-	ret0, _ := ret[0].(*ecs.GetTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetTaskProtectionWithContext indicates an expected call of GetTaskProtectionWithContext.
-func (mr *MockECSTaskProtectionSDKMockRecorder) GetTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).GetTaskProtectionWithContext), varargs...)
-}
-
-// UpdateTaskProtection mocks base method.
-func (m *MockECSTaskProtectionSDK) UpdateTaskProtection(arg0 *ecs.UpdateTaskProtectionInput) (*ecs.UpdateTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpdateTaskProtection", arg0)
-	ret0, _ := ret[0].(*ecs.UpdateTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// UpdateTaskProtection indicates an expected call of UpdateTaskProtection.
-func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtection(arg0 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtection", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtection), arg0)
-}
-
-// UpdateTaskProtectionWithContext mocks base method.
-func (m *MockECSTaskProtectionSDK) UpdateTaskProtectionWithContext(arg0 context.Context, arg1 *ecs.UpdateTaskProtectionInput, arg2 ...request.Option) (*ecs.UpdateTaskProtectionOutput, error) {
-	m.ctrl.T.Helper()
-	varargs := []interface{}{arg0, arg1}
-	for _, a := range arg2 {
-		varargs = append(varargs, a)
-	}
-	ret := m.ctrl.Call(m, "UpdateTaskProtectionWithContext", varargs...)
-	ret0, _ := ret[0].(*ecs.UpdateTaskProtectionOutput)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// UpdateTaskProtectionWithContext indicates an expected call of UpdateTaskProtectionWithContext.
-func (mr *MockECSTaskProtectionSDKMockRecorder) UpdateTaskProtectionWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskProtectionWithContext", reflect.TypeOf((*MockECSTaskProtectionSDK)(nil).UpdateTaskProtectionWithContext), varargs...)
-}
-
 // MockECSDiscoverEndpointSDK is a mock of ECSDiscoverEndpointSDK interface.
 type MockECSDiscoverEndpointSDK struct {
 	ctrl     *gomock.Controller
diff --git a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
index 22dc4c1fc2d..4dcc9b837cc 100644
--- a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
+++ b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_mocks.go
@@ -21,7 +21,7 @@ package handlers
 import (
 	reflect "reflect"
 
-	api "github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	gomock "github.com/golang/mock/gomock"
 )
@@ -50,10 +50,10 @@ func (m *MockTaskProtectionClientFactoryInterface) EXPECT() *MockTaskProtectionC
 }
 
 // NewTaskProtectionClient mocks base method.
-func (m *MockTaskProtectionClientFactoryInterface) NewTaskProtectionClient(arg0 credentials.TaskIAMRoleCredentials) api.ECSTaskProtectionSDK {
+func (m *MockTaskProtectionClientFactoryInterface) NewTaskProtectionClient(arg0 credentials.TaskIAMRoleCredentials) ecs.ECSTaskProtectionSDK {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "NewTaskProtectionClient", arg0)
-	ret0, _ := ret[0].(api.ECSTaskProtectionSDK)
+	ret0, _ := ret[0].(ecs.ECSTaskProtectionSDK)
 	return ret0
 }
 
diff --git a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_test.go b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_test.go
index ca2696294b0..11e80aaffc4 100644
--- a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_test.go
+++ b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/handlers_test.go
@@ -24,8 +24,8 @@ import (
 	"testing"
 	"time"
 
+	mock_api "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
-	mock_api "github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/metrics"
diff --git a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
index 02c8c4c4dd8..6adb7ac308e 100644
--- a/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
+++ b/ecs-agent/tmds/handlers/taskprotection/v1/handlers/interface.go
@@ -1,10 +1,10 @@
 package handlers
 
 import (
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 )
 
 type TaskProtectionClientFactoryInterface interface {
-	NewTaskProtectionClient(taskRoleCredential credentials.TaskIAMRoleCredentials) api.ECSTaskProtectionSDK
+	NewTaskProtectionClient(taskRoleCredential credentials.TaskIAMRoleCredentials) ecs.ECSTaskProtectionSDK
 }