TCPMv2: Operating as a SNK, clear the End AMS flag

While operating as a sink, the End AMS flag is cleared
at the end of the Atomic Message Sequence.

BUG=b:145812132
BRANCH=none
TEST=make -j buildall
	manual tests:
	Using a kohaku and a Kensington dock, I verified that
	the END_AMS flag is cleard and PD communication does
	not stop.

Change-Id: Ia8b40ab42916dc5b2a690a6fb7f519015912a6fa
Signed-off-by: Sam Hurst <shurst@google.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/1954256
Reviewed-by: Jett Rink <jettrink@chromium.org>
diff --git a/common/usbc/usb_prl_sm.c b/common/usbc/usb_prl_sm.c
index 9a78b59..a295d61 100644
--- a/common/usbc/usb_prl_sm.c
+++ b/common/usbc/usb_prl_sm.c
@@ -495,7 +495,7 @@
 				(PRL_FLAGS_START_AMS | PRL_FLAGS_END_AMS))) {
 		if (tc_get_power_role(port) == PD_ROLE_SOURCE) {
 			/*
-			 * Start of AMS notification received from
+			 * Start of SRC AMS notification received from
 			 * Policy Engine
 			 */
 			if (PRL_TX_CHK_FLAG(port, PRL_FLAGS_START_AMS)) {
@@ -504,18 +504,22 @@
 				return;
 			}
 			/*
-			 * End of AMS notification received from
+			 * End of SRC AMS notification received from
 			 * Policy Engine
 			 */
 			else if (PRL_TX_CHK_FLAG(port, PRL_FLAGS_END_AMS)) {
-				PRL_TX_CLR_FLAG(port, PRL_FLAGS_END_AMS);
 				/* Set Rp = SinkTxOk */
 				tcpm_select_rp_value(port, SINK_TX_OK);
 				tcpm_set_cc(port, TYPEC_CC_RP);
 				prl_tx[port].retry_counter = 0;
+				/* PRL_FLAGS_END AMS is cleared here */
 				prl_tx[port].flags = 0;
 			}
 		} else {
+			/*
+			 * Start of SNK AMS notification received from
+			 * Policy Engine
+			 */
 			if (PRL_TX_CHK_FLAG(port, PRL_FLAGS_START_AMS)) {
 				PRL_TX_CLR_FLAG(port, PRL_FLAGS_START_AMS);
 				/*
@@ -525,6 +529,16 @@
 				set_state_prl_tx(port, PRL_TX_SNK_START_AMS);
 				return;
 			}
+			/*
+			 * End of SNK AMS notification received from
+			 * Policy Engine
+			 */
+			else if (PRL_TX_CHK_FLAG(port, PRL_FLAGS_END_AMS)) {
+				prl_tx[port].retry_counter = 0;
+				/* PRL_FLAGS_END AMS is cleared here */
+				prl_tx[port].flags = 0;
+			}
+
 		}
 	} else if (PRL_TX_CHK_FLAG(port, PRL_FLAGS_MSG_XMIT)) {
 		PRL_TX_CLR_FLAG(port, PRL_FLAGS_MSG_XMIT);