[a2a] Remove stress tests from regular a2a integration tests. Update a2a client according to PR comments

Author: EugeneTheDev

a2a/a2a-client/src/commonMain/kotlin/ai/koog/a2a/client/A2AClient.kt

@@ -15,17 +15,18 @@ import ai.koog.a2a.transport.ClientTransport
 import ai.koog.a2a.transport.Request
 import ai.koog.a2a.transport.Response
 import kotlinx.coroutines.flow.Flow
-import kotlin.concurrent.Volatile
+import kotlin.concurrent.atomics.AtomicReference
+import kotlin.concurrent.atomics.ExperimentalAtomicApi
 
 /**
  * A2A client responsible for sending requests to A2A server.
  */
+@OptIn(ExperimentalAtomicApi::class)
 public open class A2AClient(
     private val transport: ClientTransport,
     private val agentCardResolver: AgentCardResolver,
 ) {
-    @Volatile
-    protected var agentCard: AgentCard? = null
+    protected var agentCard: AtomicReference<AgentCard?> = AtomicReference(null)
 
     /**
      * Performs initialization logic.
@@ -41,7 +42,7 @@ public open class A2AClient(
      */
     public open suspend fun getAgentCard(): AgentCard {
         return agentCardResolver.resolve().also {
-            agentCard = it
+            agentCard.exchange(it)
         }
     }
 
@@ -51,7 +52,7 @@ public open class A2AClient(
      * @throws [IllegalStateException] if it's not initialized
      */
     public open fun cachedAgentCard(): AgentCard {
-        return checkNotNull(agentCard) { "Agent card is not initialized." }
+        return checkNotNull(agentCard.load()) { "Agent card is not initialized." }
     }
 
     /**
@@ -64,12 +65,12 @@ public open class A2AClient(
         request: Request<Nothing?>,
         ctx: ClientCallContext = ClientCallContext.Default
     ): Response<AgentCard> {
-        check(getAgentCard().supportsAuthenticatedExtendedCard == true) {
+        check(cachedAgentCard().supportsAuthenticatedExtendedCard == true) {
             "Agent card reports that authenticated extended agent card is not supported."
         }
 
         return transport.getAuthenticatedExtendedAgentCard(request, ctx).also {
-            agentCard = it.data
+            agentCard.exchange(it.data)
         }
     }
 

a2a/a2a-core/src/commonMain/kotlin/ai/koog/a2a/model/AgentCard.kt

@@ -110,7 +110,9 @@ public data class AgentCard(
 @JvmInline
 @Serializable
 public value class TransportProtocol(public val value: String) {
-    @Suppress("MissingKDocForPublicAPI")
+    /**
+     * List of known transport protocols.
+     */
     public companion object {
         /**
          * JSON-RPC protocol.

a2a/a2a-server/src/jvmTest/kotlin/ai/koog/a2a/server/jsonrpc/A2AServerJsonRpcIntegrationTest.kt

@@ -26,7 +26,6 @@ import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withContext
 import org.junit.jupiter.api.AfterAll
 import org.junit.jupiter.api.BeforeAll
-import org.junit.jupiter.api.RepeatedTest
 import org.junit.jupiter.api.TestInstance
 import org.junit.jupiter.api.parallel.Execution
 import org.junit.jupiter.api.parallel.ExecutionMode
@@ -248,70 +247,4 @@ class A2AServerJsonRpcIntegrationTest : BaseA2AServerJsonRpcTest() {
             }
         }
     }
-
-    @RepeatedTest(300, name = "{currentRepetition}/{totalRepetitions}")
-    fun `stress test cancel task`() = runTest(timeout = testTimeout) {
-        val createTaskRequest = Request(
-            data = MessageSendParams(
-                message = Message(
-                    messageId = Uuid.Companion.random().toString(),
-                    role = Role.User,
-                    parts = listOf(
-                        TextPart("do cancelable task"),
-                    ),
-                    contextId = "test-context"
-                ),
-            ),
-        )
-
-        val taskId = (client.sendMessage(createTaskRequest).data as Task).id
-
-        val cancelTaskRequest = Request(
-            data = TaskIdParams(
-                id = taskId,
-            )
-        )
-
-        val response = client.cancelTask(cancelTaskRequest)
-
-        response.data should {
-            it.id shouldBe taskId
-            it.contextId shouldBe "test-context"
-            it.status should {
-                it.state shouldBe TaskState.Canceled
-                it.message shouldNotBeNull {
-                    role shouldBe Role.Agent
-                    parts shouldBe listOf(TextPart("Task canceled"))
-                }
-            }
-        }
-    }
-
-    @RepeatedTest(300, name = "{currentRepetition}/{totalRepetitions}")
-    fun `stress test send message`() = runTest(timeout = testTimeout) {
-        val request = Request(
-            data = MessageSendParams(
-                message = Message(
-                    messageId = Uuid.Companion.random().toString(),
-                    role = Role.User,
-                    parts = listOf(
-                        TextPart("hello world"),
-                    ),
-                    contextId = "test-context"
-                ),
-            )
-        )
-
-        val response = client.sendMessage(request)
-
-        response should {
-            it.id shouldBe request.id
-
-            it.data.shouldBeInstanceOf<Message> {
-                it.role shouldBe Role.Agent
-                it.parts shouldBe listOf(TextPart("Hello World"))
-                it.contextId shouldBe "test-context"
-            }
-        }
-    }
 }