diff --git a/compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AbstractA2ARequestHandlerTest_v0_3.java b/compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AbstractA2ARequestHandlerTest_v0_3.java index 3862b4ac9..00a1717e6 100644 --- a/compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AbstractA2ARequestHandlerTest_v0_3.java +++ b/compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AbstractA2ARequestHandlerTest_v0_3.java @@ -139,9 +139,15 @@ public void cancel(RequestContext context, AgentEmitter agentEmitter) throws A2A // Create v1.0 DefaultRequestHandler org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler v10Handler = - DefaultRequestHandler.create( - agentExecutor, taskStore, queueManager, pushConfigStore, - mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler.builder() + .agentExecutor(agentExecutor) + .taskStore(taskStore) + .queueManager(queueManager) + .pushConfigStore(pushConfigStore) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); // Wrap in v0.3 conversion handler convert03To10Handler = new Convert_v0_3_To10RequestHandler(v10Handler); diff --git a/compat-0.3/transport/jsonrpc/src/test/java/org/a2aproject/sdk/compat03/transport/jsonrpc/handler/JSONRPCHandler_v0_3_Test.java b/compat-0.3/transport/jsonrpc/src/test/java/org/a2aproject/sdk/compat03/transport/jsonrpc/handler/JSONRPCHandler_v0_3_Test.java index 61c763873..55e150589 100644 --- a/compat-0.3/transport/jsonrpc/src/test/java/org/a2aproject/sdk/compat03/transport/jsonrpc/handler/JSONRPCHandler_v0_3_Test.java +++ b/compat-0.3/transport/jsonrpc/src/test/java/org/a2aproject/sdk/compat03/transport/jsonrpc/handler/JSONRPCHandler_v0_3_Test.java @@ -858,9 +858,14 @@ public void testDeletePushNotificationConfig() { public void testOnGetPushNotificationNoPushNotifierConfig() { // Create v1.0 request handler without push config store org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler v10Handler = - org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.create( - agentExecutor, taskStore, queueManager, null, mainEventBusProcessor, - internalExecutor, internalExecutor); + org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.builder() + .agentExecutor(agentExecutor) + .taskStore(taskStore) + .queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); // Wrap in v0.3 conversion handler Convert_v0_3_To10RequestHandler handlerWithoutPushConfig = new Convert_v0_3_To10RequestHandler(v10Handler); @@ -885,9 +890,14 @@ public void testOnGetPushNotificationNoPushNotifierConfig() { public void testOnSetPushNotificationNoPushNotifierConfig() { // Create v1.0 request handler without push config store org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler v10Handler = - org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.create( - agentExecutor, taskStore, queueManager, null, mainEventBusProcessor, - internalExecutor, internalExecutor); + org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.builder() + .agentExecutor(agentExecutor) + .taskStore(taskStore) + .queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); // Wrap in v0.3 conversion handler Convert_v0_3_To10RequestHandler handlerWithoutPushConfig = new Convert_v0_3_To10RequestHandler(v10Handler); @@ -952,9 +962,14 @@ public void testDeletePushNotificationConfigNotSupported() { public void testDeletePushNotificationConfigNoPushConfigStore() { // Create v1.0 request handler without push config store org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler v10Handler = - org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.create( - agentExecutor, taskStore, queueManager, null, mainEventBusProcessor, - internalExecutor, internalExecutor); + org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler.builder() + .agentExecutor(agentExecutor) + .taskStore(taskStore) + .queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); // Wrap in v0.3 conversion handler Convert_v0_3_To10RequestHandler handlerWithoutPushConfig = new Convert_v0_3_To10RequestHandler(v10Handler); diff --git a/docs/content/dev/configuration.md b/docs/content/dev/configuration.md index d38d04de0..2458b9c09 100644 --- a/docs/content/dev/configuration.md +++ b/docs/content/dev/configuration.md @@ -56,6 +56,15 @@ a2a.blocking.reconciliation.timeout.seconds=1 a2a.agent-card.cache.max-age=3600 ``` +### Request Context + +```properties +# Load referenced tasks from the TaskStore and enforce authorization checks (default: true) +a2a.request-context.populate-referred-tasks=true +``` + +When enabled, task IDs referenced in incoming messages are looked up in the `TaskStore` and made available to the `AgentExecutor` via `RequestContext.getRelatedTasks()`. This is useful for multi-task conversations where the agent needs access to state from related tasks. Enabled by default; set to `false` to avoid extra `TaskStore` lookups when not needed. + ### Tuning Guidelines - **Streaming Performance**: The executor handles streaming subscriptions. Too few threads can cause timeouts under concurrent load. diff --git a/extras/http-client-vertx/src/main/java/org/a2aproject/sdk/client/http/vertx/VertxA2AHttpClient.java b/extras/http-client-vertx/src/main/java/org/a2aproject/sdk/client/http/vertx/VertxA2AHttpClient.java index 7cf2c57ca..924f91047 100644 --- a/extras/http-client-vertx/src/main/java/org/a2aproject/sdk/client/http/vertx/VertxA2AHttpClient.java +++ b/extras/http-client-vertx/src/main/java/org/a2aproject/sdk/client/http/vertx/VertxA2AHttpClient.java @@ -455,10 +455,20 @@ private CompletableFuture executeAsyncSSE( }); }); } else { - // Non-SSE response (error body): deliver lines to messageConsumer so - // the SSEEventListener up the call stack can parse the JSON-RPC error. - response.pipe().to(new PlainBodyWriteStream(messageConsumer)) - .onSuccess(v -> { + // Non-SSE response (error body): collect the full body and deliver it + // to messageConsumer so the SSEEventListener can parse the JSON-RPC error. + // Using body() instead of pipe() avoids a race where small responses + // are already ended before pipe() can attach, causing + // "Response already ended" IllegalStateException. + response.resume(); + response.body() + .onSuccess(body -> { + if (body != null && body.length() > 0) { + String text = body.toString(StandardCharsets.UTF_8).trim(); + if (!text.isEmpty()) { + messageConsumer.accept(new ServerSentEvent(text)); + } + } if (futureCompleted.compareAndSet(false, true)) { completeRunnable.run(); future.complete(null); diff --git a/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java index 52a7fd0e4..2166b06b9 100644 --- a/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java +++ b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java @@ -27,6 +27,7 @@ import org.a2aproject.sdk.spec.ListTaskPushNotificationConfigsResult; import org.a2aproject.sdk.spec.ListTasksParams; import org.a2aproject.sdk.spec.MessageSendParams; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.spec.StreamingEventKind; import org.a2aproject.sdk.spec.Task; import org.a2aproject.sdk.spec.TaskIdParams; @@ -458,8 +459,9 @@ public void onDeleteTaskPushNotificationConfig(DeleteTaskPushNotificationConfigP } @Override - public void validateRequestedTask(@Nullable String requestedTaskId) throws A2AError { - delegate.validateRequestedTask(requestedTaskId); + public void authorizeTaskAccess(@Nullable String requestedTaskId, ServerCallContext context, + TaskOperation operation) throws A2AError { + delegate.authorizeTaskAccess(requestedTaskId, context, operation); } private boolean extractRequest() { diff --git a/extras/task-store-database-jpa/pom.xml b/extras/task-store-database-jpa/pom.xml index 021a212c5..c034178d1 100644 --- a/extras/task-store-database-jpa/pom.xml +++ b/extras/task-store-database-jpa/pom.xml @@ -72,6 +72,11 @@ rest-assured test + + org.mockito + mockito-core + test + jakarta.transaction jakarta.transaction-api diff --git a/extras/task-store-database-jpa/src/main/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStore.java b/extras/task-store-database-jpa/src/main/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStore.java index 8f7d3c696..4019f287d 100644 --- a/extras/task-store-database-jpa/src/main/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStore.java +++ b/extras/task-store-database-jpa/src/main/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStore.java @@ -24,7 +24,6 @@ import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; import org.a2aproject.sdk.server.ServerCallContext; import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; -import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.server.config.A2AConfigProvider; import org.a2aproject.sdk.server.tasks.TaskStateProvider; import org.a2aproject.sdk.server.tasks.TaskStore; @@ -292,7 +291,12 @@ public ListTasksResult list(ListTasksParams params, @Nullable ServerCallContext boolean hasMore; int totalSize; - if (authorizationProvider != null && context != null) { + if (authorizationProvider != null && context == null) { + LOGGER.warn("Authorization provider is configured but no ServerCallContext available — " + + "returning empty result (fail-closed)"); + return new ListTasksResult(List.of(), 0, 0, null); + } + if (authorizationProvider != null) { // Iterative fetch: accumulate pageSize authorized results across DB pages tasks = new ArrayList<>(pageSize); PageToken cursor = PageToken.fromString(params.pageToken()); @@ -314,7 +318,7 @@ public ListTasksResult list(ListTasksParams params, @Nullable ServerCallContext for (JpaTask jpaTask : batch) { processedCount++; Task task = deserializeTask(jpaTask); - if (authorizationProvider.checkRead(context, task.id(), TaskOperation.LIST_TASKS)) { + if (isReadAuthorized(authorizationProvider, context, task.id())) { tasks.add(task); if (tasks.size() == pageSize) { break; diff --git a/extras/task-store-database-jpa/src/test/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStoreAuthorizationTest.java b/extras/task-store-database-jpa/src/test/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStoreAuthorizationTest.java new file mode 100644 index 000000000..eb363f873 --- /dev/null +++ b/extras/task-store-database-jpa/src/test/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStoreAuthorizationTest.java @@ -0,0 +1,33 @@ +package org.a2aproject.sdk.extras.taskstore.database.jpa; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import jakarta.enterprise.inject.Instance; + +import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; +import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; +import org.a2aproject.sdk.spec.ListTasksParams; +import org.junit.jupiter.api.Test; + +class JpaDatabaseTaskStoreAuthorizationTest { + + @Test + @SuppressWarnings("unchecked") + void listFailsClosedWhenContextIsNull() { + TaskAuthorizationProvider authProvider = mock(TaskAuthorizationProvider.class); + Instance instance = mock(Instance.class); + when(instance.isResolvable()).thenReturn(true); + when(instance.get()).thenReturn(authProvider); + + JpaDatabaseTaskStore store = new JpaDatabaseTaskStore(instance); + + ListTasksResult result = store.list(new ListTasksParams(), null); + + assertNotNull(result); + assertEquals(0, result.tasks().size()); + assertEquals(0, result.totalSize()); + } +} diff --git a/reference/jsonrpc/src/main/java/org/a2aproject/sdk/server/apps/quarkus/A2AServerRoutes.java b/reference/jsonrpc/src/main/java/org/a2aproject/sdk/server/apps/quarkus/A2AServerRoutes.java index 953e65452..44e267cf7 100644 --- a/reference/jsonrpc/src/main/java/org/a2aproject/sdk/server/apps/quarkus/A2AServerRoutes.java +++ b/reference/jsonrpc/src/main/java/org/a2aproject/sdk/server/apps/quarkus/A2AServerRoutes.java @@ -29,6 +29,7 @@ import io.vertx.ext.web.RoutingContext; import io.vertx.ext.web.handler.BodyHandler; import org.a2aproject.sdk.common.A2AHeaders; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.grpc.utils.JSONRPCUtils; import org.a2aproject.sdk.jsonrpc.common.json.IdJsonMappingException; import org.a2aproject.sdk.jsonrpc.common.json.InvalidParamsJsonMappingException; @@ -484,9 +485,11 @@ private A2AResponse processNonStreamingRequest(NonStreamingJSONRPCRequest private Multi> processStreamingRequest( A2ARequest request, ServerCallContext context) throws A2AError { if (request instanceof SendStreamingMessageRequest req) { - jsonRpcHandler.validateRequestedTask(req.getParams().message().taskId()); + jsonRpcHandler.authorizeTaskAccess(req.getParams().message().taskId(), context, + TaskOperation.MESSAGE_SEND_STREAM); } else if (request instanceof SubscribeToTaskRequest req) { - jsonRpcHandler.validateRequestedTask(req.getParams().id()); + jsonRpcHandler.authorizeTaskAccess(req.getParams().id(), context, + TaskOperation.SUBSCRIBE_TO_TASK); } try { Flow.Publisher> publisher; diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/RequestContext.java b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/RequestContext.java index 05be4d3ee..129668328 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/RequestContext.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/RequestContext.java @@ -373,6 +373,15 @@ public Builder setServerCallContext(@Nullable ServerCallContext serverCallContex return task; } + /** + * Returns the server call context set on this builder. + * + * @return the server call context, or null if not set + */ + @Nullable ServerCallContext getServerCallContext() { + return serverCallContext; + } + /** * Builds the RequestContext with ID generation and validation. * diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilder.java b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilder.java index a340ee451..3aabe184b 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilder.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilder.java @@ -3,16 +3,28 @@ import java.util.ArrayList; import java.util.List; +import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.server.tasks.TaskStore; import org.a2aproject.sdk.spec.Task; +import org.a2aproject.sdk.spec.TaskNotFoundError; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class SimpleRequestContextBuilder extends RequestContext.Builder { + private static final Logger LOGGER = LoggerFactory.getLogger(SimpleRequestContextBuilder.class); + private final TaskStore taskStore; private final boolean shouldPopulateReferredTasks; + private final @Nullable TaskAuthorizationProvider authorizationProvider; - public SimpleRequestContextBuilder(TaskStore taskStore, boolean shouldPopulateReferredTasks) { + public SimpleRequestContextBuilder(TaskStore taskStore, boolean shouldPopulateReferredTasks, + @Nullable TaskAuthorizationProvider authorizationProvider) { this.taskStore = taskStore; this.shouldPopulateReferredTasks = shouldPopulateReferredTasks; + this.authorizationProvider = authorizationProvider; } @Override @@ -21,10 +33,19 @@ public RequestContext build() { if (taskStore != null && shouldPopulateReferredTasks && getParams() != null && getParams().message().referenceTaskIds() != null) { relatedTasks = new ArrayList<>(); + ServerCallContext callContext = getServerCallContext(); for (String taskId : getParams().message().referenceTaskIds()) { + if (authorizationProvider != null) { + if (callContext == null + || !authorizationProvider.checkRead(callContext, taskId, TaskOperation.MESSAGE_SEND)) { + throw new TaskNotFoundError(); + } + } Task task = taskStore.get(taskId); if (task != null) { relatedTasks.add(task); + } else { + LOGGER.warn("Referenced task '{}' not found in TaskStore", taskId); } } } diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/auth/TaskAuthorizationProvider.java b/server-common/src/main/java/org/a2aproject/sdk/server/auth/TaskAuthorizationProvider.java index 8387001c3..6390a9314 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/auth/TaskAuthorizationProvider.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/auth/TaskAuthorizationProvider.java @@ -2,6 +2,7 @@ import org.a2aproject.sdk.server.ServerCallContext; import org.a2aproject.sdk.spec.A2AError; +import org.jspecify.annotations.Nullable; /** * SPI for per-user task authorization. @@ -93,7 +94,7 @@ * so the first writer wins and the second is a harmless no-op. *
  • CDI injection requirement: When task authorization is required, always obtain * {@code RequestHandler} through CDI injection. Manual instantiation via - * {@code DefaultRequestHandler.create()} bypasses the + * {@code DefaultRequestHandler.builder().build()} bypasses the * {@code AuthorizationRequestHandlerDecorator}.
    • * * @@ -156,4 +157,28 @@ public interface TaskAuthorizationProvider { * @throws A2AError if recording fails */ void recordOwnership(ServerCallContext context, String taskId, TaskOperation operation) throws A2AError; + + /** + * Fail-closed read-access check that handles absent provider and missing call context. + *

    + * Returns {@code true} (allow) when no provider is configured. + * Returns {@code false} (deny) when a provider is configured but no call context is available. + * Otherwise delegates to {@link #checkRead}. + * + * @param provider the authorization provider, or {@code null} if authorization is disabled + * @param context the server call context, or {@code null} if unavailable + * @param taskId the task being accessed + * @param operation which RequestHandler method triggered the check + * @return {@code true} to allow, {@code false} to deny + */ + static boolean checkReadAccess(@Nullable TaskAuthorizationProvider provider, + @Nullable ServerCallContext context, String taskId, TaskOperation operation) { + if (provider == null) { + return true; + } + if (context == null) { + return false; + } + return provider.checkRead(context, taskId, operation); + } } diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecorator.java b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecorator.java index 366b55d50..39a7b5925 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecorator.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecorator.java @@ -261,7 +261,10 @@ public void onDeleteTaskPushNotificationConfig(DeleteTaskPushNotificationConfigP } @Override - public void validateRequestedTask(@Nullable String requestedTaskId) throws A2AError { - delegate.validateRequestedTask(requestedTaskId); + public void authorizeTaskAccess(@Nullable String requestedTaskId, ServerCallContext context, TaskOperation operation) throws A2AError { + if (requestedTaskId != null) { + enforceRead(context, requestedTaskId, operation); + } + delegate.authorizeTaskAccess(requestedTaskId, context, operation); } } diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java index addac8e69..522145bb7 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java @@ -23,6 +23,8 @@ import jakarta.annotation.PostConstruct; import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.inject.Any; +import jakarta.enterprise.inject.Instance; import jakarta.inject.Inject; import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; @@ -30,6 +32,8 @@ import org.a2aproject.sdk.server.agentexecution.AgentExecutor; import org.a2aproject.sdk.server.agentexecution.RequestContext; import org.a2aproject.sdk.server.agentexecution.SimpleRequestContextBuilder; +import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.server.config.A2AConfigProvider; import org.a2aproject.sdk.server.events.EnhancedRunnable; import org.a2aproject.sdk.server.events.EventConsumer; @@ -67,6 +71,7 @@ import org.a2aproject.sdk.spec.TaskQueryParams; import org.a2aproject.sdk.spec.TaskState; import org.a2aproject.sdk.spec.UnsupportedOperationError; +import org.a2aproject.sdk.util.Assert; import org.jspecify.annotations.NonNull; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -191,10 +196,20 @@ public class DefaultRequestHandler implements RequestHandler { private static final String A2A_BLOCKING_AGENT_TIMEOUT_SECONDS = "a2a.blocking.agent.timeout.seconds"; private static final String A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS = "a2a.blocking.consumption.timeout.seconds"; private static final String A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS = "a2a.blocking.reconciliation.timeout.seconds"; + private static final String A2A_REQUEST_CONTEXT_POPULATE_REFERRED_TASKS = "a2a.request-context.populate-referred-tasks"; @Inject A2AConfigProvider configProvider; + // Only used inside initConfig() (CDI lifecycle). In static create() paths this field + // remains null and is never accessed, hence the NullAway suppression. + @Inject + @Any + @SuppressWarnings("NullAway") + @Nullable Instance authorizationProviderInstance; + + private @Nullable TaskAuthorizationProvider authorizationProvider; + /** * Timeout in seconds to wait for agent execution to complete in blocking calls. * This allows slow agents (LLM-based, data processing, external APIs) sufficient time. @@ -281,7 +296,7 @@ public DefaultRequestHandler(AgentExecutor agentExecutor, TaskStore taskStore, // implementation if the parameter is null. Skip that for now, since otherwise I get CDI errors, and // I am unsure about the correct scope. // Also reworked to make a Supplier since otherwise the builder gets polluted with wrong tasks - this.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, false); + this.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, false, null); } @SuppressWarnings("NullAway.Init") @@ -293,24 +308,95 @@ void initConfig() { configProvider.getValue(A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS)); reconciliationTimeoutSeconds = Integer.parseInt( configProvider.getValue(A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS)); + if (authorizationProviderInstance != null && authorizationProviderInstance.isResolvable()) { + authorizationProvider = authorizationProviderInstance.get(); + } + boolean populateReferredTasks = Boolean.parseBoolean( + configProvider.getValue(A2A_REQUEST_CONTEXT_POPULATE_REFERRED_TASKS)); + this.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, populateReferredTasks, authorizationProvider); } - /** - * For testing - */ - public static DefaultRequestHandler create(AgentExecutor agentExecutor, TaskStore taskStore, - QueueManager queueManager, PushNotificationConfigStore pushConfigStore, - MainEventBusProcessor mainEventBusProcessor, - Executor executor, Executor eventConsumerExecutor) { - DefaultRequestHandler handler = - new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore, - mainEventBusProcessor, executor, eventConsumerExecutor); - handler.agentCompletionTimeoutSeconds = 5; - handler.consumptionCompletionTimeoutSeconds = 2; - handler.reconciliationTimeoutSeconds = 1; - - return handler; + public static Builder builder() { + return new Builder(); + } + + @SuppressWarnings("NullAway.Init") + public static class Builder { + private AgentExecutor agentExecutor; + private TaskStore taskStore; + private QueueManager queueManager; + private @Nullable PushNotificationConfigStore pushConfigStore; + private MainEventBusProcessor mainEventBusProcessor; + private Executor executor; + private Executor eventConsumerExecutor; + private @Nullable TaskAuthorizationProvider authorizationProvider; + private boolean populateReferredTasks; + + public Builder agentExecutor(AgentExecutor agentExecutor) { + this.agentExecutor = agentExecutor; + return this; + } + + public Builder taskStore(TaskStore taskStore) { + this.taskStore = taskStore; + return this; + } + + public Builder queueManager(QueueManager queueManager) { + this.queueManager = queueManager; + return this; + } + + public Builder pushConfigStore(@Nullable PushNotificationConfigStore pushConfigStore) { + this.pushConfigStore = pushConfigStore; + return this; + } + + public Builder mainEventBusProcessor(MainEventBusProcessor mainEventBusProcessor) { + this.mainEventBusProcessor = mainEventBusProcessor; + return this; + } + + public Builder executor(Executor executor) { + this.executor = executor; + return this; + } + + public Builder eventConsumerExecutor(Executor eventConsumerExecutor) { + this.eventConsumerExecutor = eventConsumerExecutor; + return this; + } + + public Builder authorizationProvider(@Nullable TaskAuthorizationProvider authorizationProvider) { + this.authorizationProvider = authorizationProvider; + return this; + } + + public Builder populateReferredTasks(boolean populateReferredTasks) { + this.populateReferredTasks = populateReferredTasks; + return this; + } + + @SuppressWarnings("NullAway") // pushConfigStore is intentionally @Nullable + public DefaultRequestHandler build() { + Assert.checkNotNullParam("agentExecutor", agentExecutor); + Assert.checkNotNullParam("taskStore", taskStore); + Assert.checkNotNullParam("queueManager", queueManager); + Assert.checkNotNullParam("mainEventBusProcessor", mainEventBusProcessor); + Assert.checkNotNullParam("executor", executor); + Assert.checkNotNullParam("eventConsumerExecutor", eventConsumerExecutor); + DefaultRequestHandler handler = + new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore, + mainEventBusProcessor, executor, eventConsumerExecutor); + handler.agentCompletionTimeoutSeconds = 5; + handler.consumptionCompletionTimeoutSeconds = 2; + handler.reconciliationTimeoutSeconds = 1; + handler.authorizationProvider = authorizationProvider; + handler.requestContextBuilder = + () -> new SimpleRequestContextBuilder(taskStore, populateReferredTasks, authorizationProvider); + return handler; + } } @Override @@ -1063,7 +1149,7 @@ private CompletableFuture cleanupProducer(@Nullable CompletableFuture authorizationProvider == null || context == null || - authorizationProvider.checkRead(context, task.id(), TaskOperation.LIST_TASKS)) + .filter(task -> isReadAuthorized(authorizationProvider, context, task.id())) .sorted(Comparator.comparing( (Task t) -> (t.status() != null && t.status().timestamp() != null) // Truncate to milliseconds for consistency with pageToken precision diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskStore.java b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskStore.java index 6f798cc7d..b1a5925bf 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskStore.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskStore.java @@ -2,6 +2,8 @@ import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.spec.ListTasksParams; import org.a2aproject.sdk.spec.Task; import org.jspecify.annotations.Nullable; @@ -215,4 +217,21 @@ public interface TaskStore { * @throws TaskStoreException for other listing failures not covered by specific subclasses */ ListTasksResult list(ListTasksParams params, @Nullable ServerCallContext context); + + /** + * Checks whether a task is authorized for reading during list operations. + *

    + * Delegates to {@link TaskAuthorizationProvider#checkReadAccess} with + * {@link TaskOperation#LIST_TASKS}. Implementations should call this method when + * filtering tasks in {@link #list} to ensure consistent authorization behavior. + * + * @param provider the authorization provider, or {@code null} if authorization is disabled + * @param context the server call context, or {@code null} if unavailable + * @param taskId the task being checked + * @return {@code true} to include the task, {@code false} to exclude it + */ + default boolean isReadAuthorized(@Nullable TaskAuthorizationProvider provider, + @Nullable ServerCallContext context, String taskId) { + return TaskAuthorizationProvider.checkReadAccess(provider, context, taskId, TaskOperation.LIST_TASKS); + } } diff --git a/server-common/src/main/resources/META-INF/a2a-defaults.properties b/server-common/src/main/resources/META-INF/a2a-defaults.properties index 9e0d13480..e1a71fe24 100644 --- a/server-common/src/main/resources/META-INF/a2a-defaults.properties +++ b/server-common/src/main/resources/META-INF/a2a-defaults.properties @@ -27,3 +27,8 @@ a2a.executor.keep-alive-seconds=60 # Queue capacity for pending tasks (must be bounded to enable pool growth) # When queue is full, new threads are created up to max-pool-size a2a.executor.queue-capacity=100 + +# SimpleRequestContextBuilder - Referenced task population +# When true, referenced task IDs in messages are resolved from the TaskStore +# and made available via RequestContext.getRelatedTasks() +a2a.request-context.populate-referred-tasks=true diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilderTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilderTest.java new file mode 100644 index 000000000..e0647537a --- /dev/null +++ b/server-common/src/test/java/org/a2aproject/sdk/server/agentexecution/SimpleRequestContextBuilderTest.java @@ -0,0 +1,192 @@ +package org.a2aproject.sdk.server.agentexecution; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import java.util.List; + +import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.auth.TaskAuthorizationProvider; +import org.a2aproject.sdk.server.auth.TaskOperation; +import org.a2aproject.sdk.server.tasks.InMemoryTaskStore; +import org.a2aproject.sdk.spec.Message; +import org.a2aproject.sdk.spec.MessageSendConfiguration; +import org.a2aproject.sdk.spec.MessageSendParams; +import org.a2aproject.sdk.spec.Task; +import org.a2aproject.sdk.spec.TaskNotFoundError; +import org.a2aproject.sdk.spec.TaskState; +import org.a2aproject.sdk.spec.TaskStatus; +import org.a2aproject.sdk.spec.TextPart; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class SimpleRequestContextBuilderTest { + + @Mock + private TaskAuthorizationProvider authorizationProvider; + + @Mock + private ServerCallContext callContext; + + private InMemoryTaskStore taskStore; + + @BeforeEach + void setUp() { + taskStore = new InMemoryTaskStore(); + } + + private static Task testTask(String id) { + return Task.builder() + .id(id) + .contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + } + + private static MessageSendConfiguration defaultConfiguration() { + return MessageSendConfiguration.builder() + .acceptedOutputModes(List.of()) + .returnImmediately(true) + .build(); + } + + private static MessageSendParams paramsWithReferenceTaskIds(List referenceTaskIds) { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart(""))) + .referenceTaskIds(referenceTaskIds) + .build(); + return MessageSendParams.builder() + .message(message) + .configuration(defaultConfiguration()) + .build(); + } + + private static MessageSendParams paramsWithoutReferenceTaskIds() { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart(""))) + .build(); + return MessageSendParams.builder() + .message(message) + .configuration(defaultConfiguration()) + .build(); + } + + @Test + void populateDisabled_ignoresReferenceTaskIds() { + taskStore.save(testTask("task-1"), false); + + RequestContext ctx = new SimpleRequestContextBuilder(taskStore, false, authorizationProvider) + .setParams(paramsWithReferenceTaskIds(List.of("task-1"))) + .setServerCallContext(callContext) + .build(); + + assertTrue(ctx.getRelatedTasks().isEmpty()); + verifyNoInteractions(authorizationProvider); + } + + @Test + void noReferenceTaskIds_emptyRelatedTasks() { + RequestContext ctx = new SimpleRequestContextBuilder(taskStore, true, null) + .setParams(paramsWithoutReferenceTaskIds()) + .build(); + + assertTrue(ctx.getRelatedTasks().isEmpty()); + } + + @Test + void noAuthorizationProvider_loadsAllReferencedTasks() { + taskStore.save(testTask("task-1"), false); + taskStore.save(testTask("task-2"), false); + + RequestContext ctx = new SimpleRequestContextBuilder(taskStore, true, null) + .setParams(paramsWithReferenceTaskIds(List.of("task-1", "task-2"))) + .build(); + + assertEquals(2, ctx.getRelatedTasks().size()); + assertTrue(ctx.getRelatedTasks().stream().anyMatch(t -> t.id().equals("task-1"))); + assertTrue(ctx.getRelatedTasks().stream().anyMatch(t -> t.id().equals("task-2"))); + } + + @Test + void withAuthorizationProvider_rejectsWhenAnyTaskUnauthorized() { + taskStore.save(testTask("task-1"), false); + taskStore.save(testTask("task-2"), false); + taskStore.save(testTask("task-3"), false); + + when(authorizationProvider.checkRead(eq(callContext), eq("task-1"), eq(TaskOperation.MESSAGE_SEND))) + .thenReturn(true); + when(authorizationProvider.checkRead(eq(callContext), eq("task-2"), eq(TaskOperation.MESSAGE_SEND))) + .thenReturn(false); + + SimpleRequestContextBuilder builder = new SimpleRequestContextBuilder(taskStore, true, authorizationProvider); + builder.setParams(paramsWithReferenceTaskIds(List.of("task-1", "task-2", "task-3"))); + builder.setServerCallContext(callContext); + + assertThrows(TaskNotFoundError.class, builder::build); + } + + @Test + void withAuthorizationProvider_allDenied_throwsTaskNotFoundError() { + taskStore.save(testTask("task-1"), false); + taskStore.save(testTask("task-2"), false); + + when(authorizationProvider.checkRead(eq(callContext), eq("task-1"), eq(TaskOperation.MESSAGE_SEND))) + .thenReturn(false); + + SimpleRequestContextBuilder builder = new SimpleRequestContextBuilder(taskStore, true, authorizationProvider); + builder.setParams(paramsWithReferenceTaskIds(List.of("task-1", "task-2"))); + builder.setServerCallContext(callContext); + + assertThrows(TaskNotFoundError.class, builder::build); + } + + @Test + void nonExistentTask_silentlySkipped() { + taskStore.save(testTask("task-1"), false); + + RequestContext ctx = new SimpleRequestContextBuilder(taskStore, true, null) + .setParams(paramsWithReferenceTaskIds(List.of("task-1", "nonexistent"))) + .build(); + + assertEquals(1, ctx.getRelatedTasks().size()); + assertEquals("task-1", ctx.getRelatedTasks().get(0).id()); + } + + @Test + void noServerCallContext_throwsWhenAuthProviderPresent() { + taskStore.save(testTask("task-1"), false); + taskStore.save(testTask("task-2"), false); + + SimpleRequestContextBuilder builder = new SimpleRequestContextBuilder(taskStore, true, authorizationProvider); + builder.setParams(paramsWithReferenceTaskIds(List.of("task-1", "task-2"))); + + assertThrows(TaskNotFoundError.class, builder::build); + verifyNoInteractions(authorizationProvider); + } + + @Test + void authorizationCheckUsesCorrectOperation() { + taskStore.save(testTask("task-1"), false); + + when(authorizationProvider.checkRead(eq(callContext), eq("task-1"), eq(TaskOperation.MESSAGE_SEND))) + .thenReturn(true); + + new SimpleRequestContextBuilder(taskStore, true, authorizationProvider) + .setParams(paramsWithReferenceTaskIds(List.of("task-1"))) + .setServerCallContext(callContext) + .build(); + + verify(authorizationProvider).checkRead(callContext, "task-1", TaskOperation.MESSAGE_SEND); + } +} diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/auth/SimpleTaskAuthorizationProvider.java b/server-common/src/test/java/org/a2aproject/sdk/server/auth/SimpleTaskAuthorizationProvider.java new file mode 100644 index 000000000..0680584d0 --- /dev/null +++ b/server-common/src/test/java/org/a2aproject/sdk/server/auth/SimpleTaskAuthorizationProvider.java @@ -0,0 +1,39 @@ +package org.a2aproject.sdk.server.auth; + +import java.util.concurrent.ConcurrentHashMap; + +import org.a2aproject.sdk.server.ServerCallContext; + +/** + * Owner-based {@link TaskAuthorizationProvider} for tests. + */ +public class SimpleTaskAuthorizationProvider implements TaskAuthorizationProvider { + private final ConcurrentHashMap owners = new ConcurrentHashMap<>(); + + @Override + public boolean checkRead(ServerCallContext context, String taskId, TaskOperation operation) { + String owner = owners.get(taskId); + return owner == null || owner.equals(context.getUser().getUsername()); + } + + @Override + public boolean checkWrite(ServerCallContext context, String taskId, TaskOperation operation) { + return checkRead(context, taskId, operation); + } + + @Override + public boolean checkCreate(ServerCallContext context, TaskOperation operation) { + return context.getUser().isAuthenticated(); + } + + // Only returns true after an explicit recordOwnership call — unrecorded tasks are allowed by checkRead/checkWrite. + @Override + public boolean isTaskRecorded(String taskId) { + return owners.containsKey(taskId); + } + + @Override + public void recordOwnership(ServerCallContext context, String taskId, TaskOperation operation) { + owners.putIfAbsent(taskId, context.getUser().getUsername()); + } +} diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AbstractA2ARequestHandlerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AbstractA2ARequestHandlerTest.java index e3d7add5a..6d9fac04b 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AbstractA2ARequestHandlerTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AbstractA2ARequestHandlerTest.java @@ -116,8 +116,15 @@ public void cancel(RequestContext context, AgentEmitter agentEmitter) throws A2A mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, pushSender, queueManager); EventQueueUtil.start(mainEventBusProcessor); - requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, pushConfigStore, mainEventBusProcessor, internalExecutor, internalExecutor); + requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor) + .taskStore(taskStore) + .queueManager(queueManager) + .pushConfigStore(pushConfigStore) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); } @AfterEach diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecoratorTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecoratorTest.java index 901a92b1b..0ebad3f4c 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecoratorTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/AuthorizationRequestHandlerDecoratorTest.java @@ -115,6 +115,13 @@ void onCancelTask_delegatesWithoutChecks() throws A2AError { assertEquals(expected, result); verify(delegate).onCancelTask(params, context); } + + @Test + void authorizeTaskAccess_delegatesWithoutChecks() throws A2AError { + decorator.authorizeTaskAccess("task-1", context, TaskOperation.SUBSCRIBE_TO_TASK); + + verify(delegate).authorizeTaskAccess(eq("task-1"), eq(context), eq(TaskOperation.SUBSCRIBE_TO_TASK)); + } } @Nested @@ -179,6 +186,43 @@ void onListTaskPushNotificationConfigs_denied() throws A2AError { verifyNoInteractions(delegate); } + @Test + void authorizeTaskAccess_allowed() throws A2AError { + when(authorizationProvider.checkRead(context, "task-1", TaskOperation.SUBSCRIBE_TO_TASK)).thenReturn(true); + + decorator.authorizeTaskAccess("task-1", context, TaskOperation.SUBSCRIBE_TO_TASK); + + verify(authorizationProvider).checkRead(context, "task-1", TaskOperation.SUBSCRIBE_TO_TASK); + verify(delegate).authorizeTaskAccess(eq("task-1"), eq(context), eq(TaskOperation.SUBSCRIBE_TO_TASK)); + } + + @Test + void authorizeTaskAccess_denied() throws A2AError { + when(authorizationProvider.checkRead(context, "task-1", TaskOperation.SUBSCRIBE_TO_TASK)).thenReturn(false); + + assertThrows(TaskNotFoundError.class, + () -> decorator.authorizeTaskAccess("task-1", context, TaskOperation.SUBSCRIBE_TO_TASK)); + verifyNoInteractions(delegate); + } + + @Test + void authorizeTaskAccess_nullTaskId_skipsAuthCheck() throws A2AError { + decorator.authorizeTaskAccess(null, context, TaskOperation.SUBSCRIBE_TO_TASK); + + verifyNoInteractions(authorizationProvider); + verify(delegate).authorizeTaskAccess(null, context, TaskOperation.SUBSCRIBE_TO_TASK); + } + + @Test + void authorizeTaskAccess_passesOperationToEnforceRead() throws A2AError { + when(authorizationProvider.checkRead(context, "task-1", TaskOperation.GET_TASK)).thenReturn(true); + + decorator.authorizeTaskAccess("task-1", context, TaskOperation.GET_TASK); + + verify(authorizationProvider).checkRead(context, "task-1", TaskOperation.GET_TASK); + verify(delegate).authorizeTaskAccess(eq("task-1"), eq(context), eq(TaskOperation.GET_TASK)); + } + @Test void spiException_propagates() throws A2AError { TaskQueryParams params = new TaskQueryParams("task-1", null, null); diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerReferenceTaskAuthorizationTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerReferenceTaskAuthorizationTest.java new file mode 100644 index 000000000..c56b57265 --- /dev/null +++ b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerReferenceTaskAuthorizationTest.java @@ -0,0 +1,263 @@ +package org.a2aproject.sdk.server.requesthandlers; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicReference; + +import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.agentexecution.AgentExecutor; +import org.a2aproject.sdk.server.agentexecution.RequestContext; +import org.a2aproject.sdk.server.auth.SimpleTaskAuthorizationProvider; +import org.a2aproject.sdk.server.auth.TaskOperation; +import org.a2aproject.sdk.server.auth.User; +import org.a2aproject.sdk.server.events.EventQueueUtil; +import org.a2aproject.sdk.server.events.InMemoryQueueManager; +import org.a2aproject.sdk.server.events.MainEventBus; +import org.a2aproject.sdk.server.events.MainEventBusProcessor; +import org.a2aproject.sdk.server.tasks.AgentEmitter; +import org.a2aproject.sdk.server.tasks.InMemoryPushNotificationConfigStore; +import org.a2aproject.sdk.server.tasks.InMemoryTaskStore; +import org.a2aproject.sdk.server.tasks.PushNotificationSender; +import org.a2aproject.sdk.spec.A2AError; +import org.a2aproject.sdk.spec.Message; +import org.a2aproject.sdk.spec.TaskNotFoundError; +import org.a2aproject.sdk.spec.MessageSendConfiguration; +import org.a2aproject.sdk.spec.MessageSendParams; +import org.a2aproject.sdk.spec.Task; +import org.a2aproject.sdk.spec.TaskState; +import org.a2aproject.sdk.spec.TaskStatus; +import org.a2aproject.sdk.spec.TextPart; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class DefaultRequestHandlerReferenceTaskAuthorizationTest { + + private static final MessageSendConfiguration DEFAULT_CONFIG = MessageSendConfiguration.builder() + .returnImmediately(true) + .acceptedOutputModes(List.of()) + .build(); + + private static final PushNotificationSender NOOP_PUSH_SENDER = (event, snapshot) -> {}; + + private InMemoryTaskStore taskStore; + private InMemoryQueueManager queueManager; + private MainEventBusProcessor mainEventBusProcessor; + private ExecutorService executor; + private RequestHandler requestHandler; + private AtomicReference> capturedRelatedTasks; + private SimpleTaskAuthorizationProvider authProvider; + + @BeforeEach + void setUp() { + capturedRelatedTasks = new AtomicReference<>(); + authProvider = new SimpleTaskAuthorizationProvider(); + + AgentExecutor agentExecutor = capturingAgentExecutor(capturedRelatedTasks); + + taskStore = new InMemoryTaskStore(); + MainEventBus mainEventBus = new MainEventBus(); + queueManager = new InMemoryQueueManager(taskStore, mainEventBus); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSH_SENDER, queueManager); + EventQueueUtil.start(mainEventBusProcessor); + + executor = Executors.newCachedThreadPool(); + requestHandler = DefaultRequestHandler.builder() + .agentExecutor(agentExecutor) + .taskStore(taskStore) + .queueManager(queueManager) + .pushConfigStore(new InMemoryPushNotificationConfigStore()) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(executor) + .eventConsumerExecutor(executor) + .authorizationProvider(authProvider) + .populateReferredTasks(true) + .build(); + } + + @AfterEach + void tearDown() { + if (mainEventBusProcessor != null) { + EventQueueUtil.stop(mainEventBusProcessor); + } + if (executor != null) { + executor.shutdownNow(); + } + } + + @Test + void ownerCanAccessOwnTasksViaReferenceTaskIds() throws Exception { + Task ownedTask = saveTaskOwnedBy("userA"); + + ServerCallContext contextA = contextForUser("userA"); + sendMessageWithReferences(List.of(ownedTask.id()), contextA); + + assertEquals(1, capturedRelatedTasks.get().size()); + assertEquals(ownedTask.id(), capturedRelatedTasks.get().get(0).id()); + } + + @Test + void nonOwnerCannotAccessTasksViaReferenceTaskIds() { + Task taskOwnedByA = saveTaskOwnedBy("userA"); + + ServerCallContext contextB = contextForUser("userB"); + assertThrows(TaskNotFoundError.class, + () -> sendMessageWithReferences(List.of(taskOwnedByA.id()), contextB)); + } + + @Test + void mixedOwnership_rejectsRequestWhenAnyTaskUnauthorized() { + Task taskA = saveTaskOwnedBy("userA"); + Task taskB = saveTaskOwnedBy("userB"); + + ServerCallContext contextA = contextForUser("userA"); + assertThrows(TaskNotFoundError.class, + () -> sendMessageWithReferences(List.of(taskA.id(), taskB.id()), contextA)); + } + + @Test + void noAuthorizationProvider_loadsAllReferencedTasks() throws Exception { + AtomicReference> localCapture = new AtomicReference<>(); + RequestHandler handlerNoAuth = DefaultRequestHandler.builder() + .agentExecutor(capturingAgentExecutor(localCapture)) + .taskStore(taskStore) + .queueManager(queueManager) + .pushConfigStore(new InMemoryPushNotificationConfigStore()) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(executor) + .eventConsumerExecutor(executor) + .populateReferredTasks(true) + .build(); + + Task task1 = Task.builder() + .id("t1").contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + Task task2 = Task.builder() + .id("t2").contextId("ctx-2") + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + taskStore.save(task1, false); + taskStore.save(task2, false); + + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("test"))) + .referenceTaskIds(List.of("t1", "t2")) + .build(); + MessageSendParams params = MessageSendParams.builder() + .message(message) + .configuration(DEFAULT_CONFIG) + .build(); + + handlerNoAuth.onMessageSend(params, contextForUser("anyone")); + + assertEquals(2, localCapture.get().size()); + } + + @Nested + class ValidateRequestedTaskAuthorizationTests { + + private RequestHandler decoratedHandler; + + @BeforeEach + void setUpDecorated() { + decoratedHandler = new AuthorizationRequestHandlerDecorator(requestHandler, authProvider); + } + + @Test + void ownerCanValidateOwnActiveTask() throws A2AError { + Task task = saveActiveTaskOwnedBy("userA"); + + decoratedHandler.authorizeTaskAccess(task.id(), contextForUser("userA"), TaskOperation.SUBSCRIBE_TO_TASK); + } + + @Test + void nonOwnerCannotValidateTask() { + Task task = saveActiveTaskOwnedBy("userA"); + + assertThrows(TaskNotFoundError.class, + () -> decoratedHandler.authorizeTaskAccess(task.id(), contextForUser("userB"), + TaskOperation.SUBSCRIBE_TO_TASK)); + } + + @Test + void nullTaskId_skipsAuthorizationCheck() throws A2AError { + decoratedHandler.authorizeTaskAccess(null, contextForUser("anyone"), TaskOperation.SUBSCRIBE_TO_TASK); + } + + private Task saveActiveTaskOwnedBy(String owner) { + String taskId = "active-task-" + owner + "-" + System.nanoTime(); + Task task = Task.builder() + .id(taskId).contextId("ctx-" + taskId) + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .build(); + taskStore.save(task, false); + authProvider.recordOwnership(contextForUser(owner), taskId, TaskOperation.MESSAGE_SEND); + return task; + } + } + + private Task saveTaskOwnedBy(String owner) { + String taskId = "task-" + owner + "-" + System.nanoTime(); + Task task = Task.builder() + .id(taskId).contextId("ctx-" + taskId) + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + taskStore.save(task, false); + authProvider.recordOwnership(contextForUser(owner), taskId, TaskOperation.MESSAGE_SEND); + return task; + } + + private void sendMessageWithReferences(List referenceTaskIds, ServerCallContext context) throws A2AError { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("test"))) + .referenceTaskIds(referenceTaskIds) + .build(); + MessageSendParams params = MessageSendParams.builder() + .message(message) + .configuration(DEFAULT_CONFIG) + .build(); + requestHandler.onMessageSend(params, context); + } + + private static AgentExecutor capturingAgentExecutor(AtomicReference> capture) { + return new AgentExecutor() { + @Override + public void execute(RequestContext context, AgentEmitter agentEmitter) throws A2AError { + capture.set(List.copyOf(context.getRelatedTasks())); + agentEmitter.complete(); + } + + @Override + public void cancel(RequestContext context, AgentEmitter agentEmitter) throws A2AError { + } + }; + } + + private static ServerCallContext contextForUser(String username) { + return new ServerCallContext(new TestUser(username), Map.of(), Set.of()); + } + + private record TestUser(String username) implements User { + @Override + public boolean isAuthenticated() { + return true; + } + + @Override + public String getUsername() { + return username; + } + } + +} diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java index 85ba534e3..12b7ebeed 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java @@ -127,8 +127,15 @@ public void cancel(RequestContext context, AgentEmitter agentEmitter) throws A2A EventQueueUtil.start(mainEventBusProcessor); // Create DefaultRequestHandler - requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, pushConfigStore, mainEventBusProcessor, internalExecutor, internalExecutor); + requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor) + .taskStore(taskStore) + .queueManager(queueManager) + .pushConfigStore(pushConfigStore) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor) + .eventConsumerExecutor(internalExecutor) + .build(); } @AfterEach diff --git a/transport/grpc/src/test/java/org/a2aproject/sdk/transport/grpc/handler/GrpcHandlerTest.java b/transport/grpc/src/test/java/org/a2aproject/sdk/transport/grpc/handler/GrpcHandlerTest.java index 8a0b13530..66c620164 100644 --- a/transport/grpc/src/test/java/org/a2aproject/sdk/transport/grpc/handler/GrpcHandlerTest.java +++ b/transport/grpc/src/test/java/org/a2aproject/sdk/transport/grpc/handler/GrpcHandlerTest.java @@ -277,7 +277,11 @@ public void testPushNotificationsNotSupportedError() throws Exception { @Test public void testOnGetPushNotificationNoPushNotifierConfig() throws Exception { // Create request handler without a push notifier - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); AgentCard card = AbstractA2ARequestHandlerTest.createAgentCard(false, true); GrpcHandler handler = new TestGrpcHandler(card, requestHandler, internalExecutor); StreamRecorder streamRecorder = getTaskPushNotificationConfigRequest(handler, @@ -288,7 +292,11 @@ public void testOnGetPushNotificationNoPushNotifierConfig() throws Exception { @Test public void testOnSetPushNotificationNoPushNotifierConfig() throws Exception { // Create request handler without a push notifier - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); AgentCard card = AbstractA2ARequestHandlerTest.createAgentCard(false, true); GrpcHandler handler = new TestGrpcHandler(card, requestHandler, internalExecutor); StreamRecorder streamRecorder = createTaskPushNotificationConfigRequest(handler, @@ -687,7 +695,11 @@ public void testListPushNotificationConfigNotSupported() throws Exception { @Test public void testListPushNotificationConfigNoPushConfigStore() { - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); GrpcHandler handler = new TestGrpcHandler(AbstractA2ARequestHandlerTest.CARD, requestHandler, internalExecutor); taskStore.save(AbstractA2ARequestHandlerTest.MINIMAL_TASK, false); agentExecutorExecute = (context, agentEmitter) -> { @@ -758,7 +770,11 @@ public void testDeletePushNotificationConfigNotSupported() throws Exception { @Test public void testDeletePushNotificationConfigNoPushConfigStore() { - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); GrpcHandler handler = new TestGrpcHandler(AbstractA2ARequestHandlerTest.CARD, requestHandler, internalExecutor); DeleteTaskPushNotificationConfigRequest request = DeleteTaskPushNotificationConfigRequest.newBuilder() .setId(AbstractA2ARequestHandlerTest.MINIMAL_TASK.id()) diff --git a/transport/jsonrpc/src/main/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandler.java b/transport/jsonrpc/src/main/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandler.java index ed7008bc4..e6ad92568 100644 --- a/transport/jsonrpc/src/main/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandler.java +++ b/transport/jsonrpc/src/main/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandler.java @@ -56,6 +56,7 @@ import org.a2aproject.sdk.spec.TaskPushNotificationConfig; import mutiny.zero.ZeroPublisher; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.jspecify.annotations.Nullable; /** @@ -381,7 +382,7 @@ public Flow.Publisher onSubscribeToTask( request.getId(), new InvalidRequestError("Streaming is not supported by the agent"))); } - requestHandler.validateRequestedTask(request.getParams().id()); + requestHandler.authorizeTaskAccess(request.getParams().id(), context, TaskOperation.SUBSCRIBE_TO_TASK); try { Flow.Publisher publisher = requestHandler.onSubscribeToTask(request.getParams(), context); @@ -743,7 +744,7 @@ public void onComplete() { }); } - public void validateRequestedTask(String requestedTaskId) { - requestHandler.validateRequestedTask(requestedTaskId); + public void authorizeTaskAccess(String requestedTaskId, ServerCallContext context, TaskOperation operation) { + requestHandler.authorizeTaskAccess(requestedTaskId, context, operation); } } diff --git a/transport/jsonrpc/src/test/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandlerTest.java b/transport/jsonrpc/src/test/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandlerTest.java index 9aa8fbccb..431c089ae 100644 --- a/transport/jsonrpc/src/test/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandlerTest.java +++ b/transport/jsonrpc/src/test/java/org/a2aproject/sdk/transport/jsonrpc/handler/JSONRPCHandlerTest.java @@ -1125,7 +1125,11 @@ public void testPushNotificationsNotSupportedError() { @Test public void testOnGetPushNotificationNoPushNotifierConfig() { // Create request handler without a push notifier - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); AgentCard card = createAgentCard(false, true); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1144,7 +1148,11 @@ public void testOnGetPushNotificationNoPushNotifierConfig() { @Test public void testOnSetPushNotificationNoPushNotifierConfig() { // Create request handler without a push notifier - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); AgentCard card = createAgentCard(false, true); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1235,7 +1243,11 @@ public void testDefaultRequestHandlerWithCustomComponents() { @Test public void testOnMessageSendErrorHandling() { - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); AgentCard card = createAgentCard(false, true); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1396,7 +1408,11 @@ public void testListPushNotificationConfigNotSupported() { @Test public void testListPushNotificationConfigNoPushConfigStore() { - DefaultRequestHandler requestHandler = DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler requestHandler = DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); taskStore.save(MINIMAL_TASK, false); agentExecutorExecute = (context, agentEmitter) -> { @@ -1490,7 +1506,11 @@ public void testDeletePushNotificationConfigNotSupported() { @Test public void testDeletePushNotificationConfigNoPushConfigStore() { DefaultRequestHandler requestHandler = - DefaultRequestHandler.create(executor, taskStore, queueManager, null, mainEventBusProcessor, internalExecutor, internalExecutor); + DefaultRequestHandler.builder() + .agentExecutor(executor).taskStore(taskStore).queueManager(queueManager) + .mainEventBusProcessor(mainEventBusProcessor) + .executor(internalExecutor).eventConsumerExecutor(internalExecutor) + .build(); JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); taskStore.save(MINIMAL_TASK, false); agentExecutorExecute = (context, agentEmitter) -> { diff --git a/transport/rest/src/main/java/org/a2aproject/sdk/transport/rest/handler/RestHandler.java b/transport/rest/src/main/java/org/a2aproject/sdk/transport/rest/handler/RestHandler.java index 8a1c426df..131d5b9b7 100644 --- a/transport/rest/src/main/java/org/a2aproject/sdk/transport/rest/handler/RestHandler.java +++ b/transport/rest/src/main/java/org/a2aproject/sdk/transport/rest/handler/RestHandler.java @@ -35,6 +35,7 @@ import org.a2aproject.sdk.server.ExtendedAgentCard; import org.a2aproject.sdk.server.PublicAgentCard; import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.server.extensions.A2AExtensions; import org.a2aproject.sdk.server.requesthandlers.RequestHandler; import org.a2aproject.sdk.server.util.async.Internal; @@ -301,7 +302,7 @@ public HTTPRestResponse sendStreamingMessage(ServerCallContext context, String t request.setTenant(tenant); MessageSendParams params = ProtoUtils.FromProto.messageSendParams(request); try { - requestHandler.validateRequestedTask(params.message().taskId()); + requestHandler.authorizeTaskAccess(params.message().taskId(), context, TaskOperation.MESSAGE_SEND_STREAM); } catch (A2AError e) { return createErrorResponse(e); } @@ -428,7 +429,7 @@ public HTTPRestResponse subscribeToTask(ServerCallContext context, String tenant } TaskIdParams params = TaskIdParams.builder().id(taskId).tenant(tenant).build(); try { - requestHandler.validateRequestedTask(params.id()); + requestHandler.authorizeTaskAccess(params.id(), context, TaskOperation.SUBSCRIBE_TO_TASK); } catch (A2AError e) { return createErrorResponse(e); }