[a2a] Add stress tests, fix race conditions, and refine testing approach

Author: EugeneTheDev

a2a/a2a-client/build.gradle.kts

@@ -37,6 +37,7 @@ kotlin {
                 implementation(kotlin("test-junit5"))
                 implementation(project(":a2a:a2a-test"))
                 implementation(project(":a2a:a2a-transport:a2a-transport-client-jsonrpc-http"))
+                implementation(project(":test-utils"))
 
                 implementation(libs.ktor.client.cio)
                 implementation(libs.ktor.client.logging)

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-client/src/jvmTest/kotlin/ai/koog/a2a/client/A2AClientJsonRpcIntegrationTest.kt

@@ -2,22 +2,23 @@ package ai.koog.a2a.client
 
 import ai.koog.a2a.test.BaseA2AProtocolTest
 import ai.koog.a2a.transport.client.jsonrpc.http.HttpJSONRPCClientTransport
+import ai.koog.test.utils.DockerAvailableCondition
 import io.ktor.client.HttpClient
 import io.ktor.client.plugins.logging.LogLevel
 import io.ktor.client.plugins.logging.Logging
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.AfterAll
 import org.junit.jupiter.api.BeforeAll
 import org.junit.jupiter.api.TestInstance
-import org.junit.jupiter.api.condition.EnabledOnOs
-import org.junit.jupiter.api.condition.OS
+import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.parallel.Execution
 import org.junit.jupiter.api.parallel.ExecutionMode
 import org.testcontainers.containers.GenericContainer
 import org.testcontainers.containers.wait.strategy.Wait
 import org.testcontainers.junit.jupiter.Container
 import org.testcontainers.junit.jupiter.Testcontainers
-import kotlin.time.Duration.Companion.minutes
+import kotlin.test.Test
+import kotlin.time.Duration.Companion.seconds
 
 /**
  * Integration test class for testing the JSON-RPC HTTP communication in the A2A client context.
@@ -26,7 +27,7 @@ import kotlin.time.Duration.Companion.minutes
  */
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 @Testcontainers
-@EnabledOnOs(OS.LINUX)
+@ExtendWith(DockerAvailableCondition::class)
 @Execution(ExecutionMode.SAME_THREAD, reason = "Working with the same instance of test server.")
 class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
     companion object {
@@ -37,7 +38,7 @@ class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
                 .waitingFor(Wait.forListeningPort())
     }
 
-    override val testTimeout = 1.minutes
+    override val testTimeout = 10.seconds
 
     private val httpClient = HttpClient {
         install(Logging) {
@@ -74,4 +75,36 @@ class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
     fun tearDown() = runTest {
         transport.close()
     }
+
+    @Test
+    override fun `test get agent card`() =
+        super.`test get agent card`()
+
+    @Test
+    override fun `test get authenticated extended agent card`() =
+        super.`test get authenticated extended agent card`()
+
+    @Test
+    override fun `test send message`() =
+        super.`test send message`()
+
+    @Test
+    override fun `test send message streaming`() =
+        super.`test send message streaming`()
+
+    @Test
+    override fun `test get task`() =
+        super.`test get task`()
+
+    @Test
+    override fun `test cancel task`() =
+        super.`test cancel task`()
+
+    @Test
+    override fun `test resubscribe task`() =
+        super.`test resubscribe task`()
+
+    @Test
+    override fun `test push notification configs`() =
+        super.`test push notification configs`()
 }

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/commonMain/kotlin/ai/koog/a2a/server/A2AServer.kt

@@ -39,16 +39,20 @@ import ai.koog.a2a.transport.Response
 import ai.koog.a2a.transport.ServerCallContext
 import ai.koog.a2a.utils.KeyedMutex
 import ai.koog.a2a.utils.withLock
+import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CompletableJob
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.channelFlow
-import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.last
+import kotlinx.coroutines.flow.lastOrNull
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.launch
 
 /**
@@ -321,6 +325,10 @@ public open class A2AServer(
     protected val idGenerator: IdGenerator = UuidIdGenerator,
     protected val coroutineScope: CoroutineScope = CoroutineScope(SupervisorJob()),
 ) : RequestHandler {
+    private companion object {
+        private val logger = KotlinLogging.logger {}
+    }
+
     /**
      * Mutex for locking specific tasks by their IDs.
      */
@@ -374,7 +382,7 @@ public open class A2AServer(
 
         val taskId = message.taskId ?: idGenerator.generateTaskId(message)
 
-        val session = tasksMutex.withLock(taskId) {
+        val (session, monitoringStarted) = tasksMutex.withLock(taskId) {
             // If there's a currently running session for the same task, wait for it to finish.
             sessionManager.getSession(taskId)?.join()
 
@@ -412,22 +420,36 @@ public open class A2AServer(
                 eventProcessor = eventProcessor,
             ) {
                 agentExecutor.execute(requestContext, eventProcessor)
-            }.also {
-                sessionManager.addSession(it)
+            }.let {
+                it to sessionManager.addSession(it)
             }
         }
 
+        // Signal that event collection is setup
+        val collectionStarted: CompletableJob = Job()
+
         // Subscribe to events stream and start emitting them.
         launch {
             session.events
+                .onStart {
+                    collectionStarted.complete()
+                }
                 .collect { event ->
                     send(Response(data = event, id = request.id))
                 }
         }
 
-        // Start the session to execute the agent and wait for it to finish.
-        // Using await here to propagate any exceptions thrown by the agent execution.
+        // Ensure event collection is setup to stream events in response.
+        collectionStarted.join()
+        // Ensure monitoring is ready to monitor the session.
+        monitoringStarted.join()
+
+        /*
+         Start the session to execute the agent and wait for it to finish.
+         Using await here to propagate any exceptions thrown by the agent execution.
+         */
         session.agentJob.await()
+        session.join()
     }
 
     override suspend fun onSendMessage(
@@ -440,10 +462,10 @@ public open class A2AServer(
 
         val event = if (messageConfiguration?.blocking == true) {
             // If blocking is requested, attempt to wait for the last event, until the current turn of the agent execution is finished.
-            eventStream.last()
+            eventStream.lastOrNull()
         } else {
-            eventStream.first()
-        }
+            eventStream.firstOrNull()
+        } ?: throw IllegalStateException("Can't get response from the agent: event stream is empty")
 
         return when (val eventData = event.data) {
             is Message -> Response(data = eventData, id = event.id)

a2a/a2a-server/src/commonMain/kotlin/ai/koog/a2a/server/session/Session.kt

@@ -14,15 +14,25 @@ import kotlinx.coroutines.flow.collect
  *
  * @property eventProcessor The session event processor
  * @property agentJob The execution process associated with this session's execution
- * @property events A stream of events generated during this session
  */
 public class Session(
     public val eventProcessor: SessionEventProcessor,
     public val agentJob: Deferred<Unit>
 ) {
-    public val contextId: String get() = eventProcessor.contextId
-    public val taskId: String get() = eventProcessor.taskId
-    public val events: Flow<Event> get() = eventProcessor.events
+    /**
+     * Context ID associated with this session.
+     */
+    public val contextId: String = eventProcessor.contextId
+
+    /**
+     * Task ID associated with this session.
+     */
+    public val taskId: String = eventProcessor.taskId
+
+    /**
+     * A stream of events associated with this session.
+     */
+    public val events: Flow<Event> = eventProcessor.events
 
     /**
      * Starts the [agentJob], if it hasn't already been started.
@@ -31,7 +41,7 @@ public class Session(
         agentJob.start()
     }
 
-    /*
+    /**
      * Suspends until the session, i.e., event stream and agent job, complete.
      * Waits for the event stream to finish first, to avoid triggering the agent job prematurely.
      * Assumes that by the time event stream is finished, agent job will already be completed or canceled.

a2a/a2a-server/src/commonMain/kotlin/ai/koog/a2a/server/session/SessionEventProcessor.kt

@@ -10,10 +10,9 @@ import ai.koog.a2a.server.exceptions.SessionNotActiveException
 import ai.koog.a2a.server.tasks.TaskStorage
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.filterIsInstance
-import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onSubscription
 import kotlinx.coroutines.flow.takeWhile
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -42,8 +41,6 @@ import kotlin.jvm.JvmInline
  * from the incoming request or a newly generated ID that must be used if creating a new task.
  * Note: This taskId might not correspond to an actually existing task initially - it serves as the
  * identifier that will be validated against all [TaskEvent] in this session.
- * @property isOpen Whether the session is open.
- * @property events A hot flow of events in this session that can be subscribed to.
  */
 @OptIn(ExperimentalAtomicApi::class)
 public class SessionEventProcessor(
@@ -63,6 +60,10 @@ public class SessionEventProcessor(
     }
 
     private val _isOpen: AtomicBoolean = AtomicBoolean(true)
+
+    /**
+     * Whether the session is open.
+     */
     public val isOpen: Boolean get() = _isOpen.load()
 
     /**
@@ -82,17 +83,15 @@ public class SessionEventProcessor(
     }
 
     private val _events = MutableSharedFlow<FlowEvent>()
-    public val events: Flow<Event>
-        get() = flow {
-            if (isOpen) {
-                _events
-                    .takeWhile { it !is FlowEvent.Close }
-                    .filterIsInstance<FlowEvent.Data>()
-                    .map { it.data }
-            } else {
-                emptyFlow()
-            }.collect(this)
-        }
+
+    /**
+     * A hot flow of events in this session that can be subscribed to.
+     */
+    public val events: Flow<Event> = _events
+        .onSubscription { if (!_isOpen.load()) emit(FlowEvent.Close) }
+        .takeWhile { it !is FlowEvent.Close }
+        .filterIsInstance<FlowEvent.Data>()
+        .map { it.data }
 
     /**
      * Sends a [Message] to the session event processor. Validates the message against the session context and updates