diff --git a/agent/handlers/v4/response.go b/agent/handlers/v4/response.go
index 5656a1871a4..238c4a66cfd 100644
--- a/agent/handlers/v4/response.go
+++ b/agent/handlers/v4/response.go
@@ -28,7 +28,7 @@ import (
 )
 
 // NewTaskResponse creates a new v4 response object for the task. It augments v2 task response
-// with additional network interface fields.
+// with additional fields for the v4 response.
 func NewTaskResponse(
 	taskARN string,
 	state dockerstate.TaskEngineState,
@@ -55,20 +55,15 @@ func NewTaskResponse(
 		if err != nil {
 			return nil, err
 		}
-		dockerContainer, ok := state.ContainerByID(container.ID)
-		if !ok {
-			return nil, errors.Errorf(
-				"v4 container response: unable to find container '%s'", container.ID)
-		}
-		var restartCount int
-		if dockerContainer.Container.RestartPolicyEnabled() {
-			restartCount = dockerContainer.Container.RestartTracker.GetRestartCount()
-		}
-		containers = append(containers, tmdsv4.ContainerResponse{
+		v4Response := &tmdsv4.ContainerResponse{
 			ContainerResponse: &v2Resp.Containers[i],
 			Networks:          networks,
-			RestartCount:      restartCount,
-		})
+		}
+		v4Response, err = augmentContainerResponse(container.ID, state, v4Response)
+		if err != nil {
+			return nil, err
+		}
+		containers = append(containers, *v4Response)
 	}
 
 	return &tmdsv4.TaskResponse{
@@ -79,8 +74,8 @@ func NewTaskResponse(
 	}, nil
 }
 
-// NewContainerResponse creates a new v4 container response based on container id.  It augments
-// v4 container response with additional network interface fields.
+// NewContainerResponse creates a new v4 container response based on container id. It augments
+// v2 container response with additional fields for v4.
 func NewContainerResponse(
 	containerID string,
 	state dockerstate.TaskEngineState,
@@ -97,20 +92,11 @@ func NewContainerResponse(
 	if err != nil {
 		return nil, err
 	}
-	dockerContainer, ok := state.ContainerByID(containerID)
-	if !ok {
-		return nil, errors.Errorf(
-			"v4 container response: unable to find container '%s'", containerID)
-	}
-	var restartCount int
-	if dockerContainer.Container.RestartPolicyEnabled() {
-		restartCount = dockerContainer.Container.RestartTracker.GetRestartCount()
-	}
-	return &tmdsv4.ContainerResponse{
+	v4Response := tmdsv4.ContainerResponse{
 		ContainerResponse: container,
 		Networks:          networks,
-		RestartCount:      restartCount,
-	}, nil
+	}
+	return augmentContainerResponse(containerID, state, &v4Response)
 }
 
 // toV4NetworkResponse converts v2 network response to v4. Additional fields are only
@@ -141,6 +127,24 @@ func toV4NetworkResponse(
 	return resp, nil
 }
 
+// augmentContainerResponse augments the container response with additional fields.
+func augmentContainerResponse(
+	containerID string,
+	state dockerstate.TaskEngineState,
+	v4Response *tmdsv4.ContainerResponse,
+) (*tmdsv4.ContainerResponse, error) {
+	dockerContainer, ok := state.ContainerByID(containerID)
+	if !ok {
+		return nil, errors.Errorf(
+			"v4 container response: unable to find container '%s'", containerID)
+	}
+	if dockerContainer.Container.RestartPolicyEnabled() {
+		restartCount := dockerContainer.Container.RestartTracker.GetRestartCount()
+		v4Response.RestartCount = &restartCount
+	}
+	return v4Response, nil
+}
+
 // newNetworkInterfaceProperties creates the NetworkInterfaceProperties object for a given
 // task.
 func newNetworkInterfaceProperties(task *apitask.Task) (tmdsv4.NetworkInterfaceProperties, error) {