Add agent as subject type for agentic OBO tokens

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/handlers/claims/AgentAccessTokenClaimProvider.java

@@ -1,13 +1,32 @@
+/*
+ * Copyright (c) 2025, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ *
+ * WSO2 Inc. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 package org.wso2.carbon.identity.oauth2.token.handlers.claims;
 
 import org.apache.commons.lang.StringUtils;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.apache.oltu.oauth2.common.message.types.GrantType;
 import org.wso2.carbon.identity.core.util.IdentityUtil;
 import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
 import org.wso2.carbon.identity.oauth2.authz.OAuthAuthzReqMessageContext;
 import org.wso2.carbon.identity.oauth2.token.OAuthTokenReqMessageContext;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -16,6 +35,7 @@
  */
 public class AgentAccessTokenClaimProvider implements JWTAccessTokenClaimProvider {
 
+    private static final Log LOG = LogFactory.getLog(AgentAccessTokenClaimProvider.class);
     private static final String ACT = "act";
     private static final String SUB = "sub";
     private static final String AGENT = "AGENT";
@@ -30,16 +50,22 @@ public Map<String, Object> getAdditionalClaims(OAuthAuthzReqMessageContext conte
     @Override
     public Map<String, Object> getAdditionalClaims(OAuthTokenReqMessageContext context) throws IdentityOAuth2Exception {
 
+        LOG.debug("Processing additional claims for agent access token");
         String agentIdentityUserStoreName = IdentityUtil.getAgentIdentityUserstoreName();
         if (StringUtils.isNotEmpty(agentIdentityUserStoreName) && agentIdentityUserStoreName
                 .equalsIgnoreCase(context.getAuthorizedUser().getUserStoreDomain())) {
             Map<String, Object> agentMap = new HashMap<>();
             agentMap.put(AUT, AGENT);
+            LOG.debug("Added AGENT as AUT claim based on authorized agent user store domain");
             return agentMap;
         } else if (GrantType.AUTHORIZATION_CODE.toString().equals(context.getOauth2AccessTokenReqDTO().getGrantType())
             && context.getRequestedActor() != null) {
             Map<String, Object> agentMap = new HashMap<>();
-            agentMap.put(ACT, Collections.singletonMap(SUB, context.getRequestedActor()));
+            Map<String, Object> actMap = new HashMap<>();
+            actMap.put(SUB, context.getRequestedActor());
+            actMap.put(AUT, AGENT);
+            agentMap.put(ACT, actMap);
+            LOG.debug("Added agent claims in OBO tokens for authorization code grant as requested actor");
             return agentMap;
         }
         return null;

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/openidconnect/OIDCAgentClaimProviderImpl.java

@@ -1,18 +1,21 @@
 /*
- * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2025, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
+ * WSO2 Inc. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
  * You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
  */
+
 package org.wso2.carbon.identity.openidconnect;
 
 import org.apache.commons.lang.StringUtils;

components/org.wso2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/oauth2/token/handlers/claims/AgentAccessTokenClaimProviderTest.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright (c) 2025, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ *
+ * WSO2 Inc. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.oauth2.token.handlers.claims;
+
+import org.mockito.MockedStatic;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
+import org.wso2.carbon.identity.oauth2.authz.OAuthAuthzReqMessageContext;
+import org.wso2.carbon.identity.oauth2.token.OAuthTokenReqMessageContext;
+
+import java.util.Map;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+public class AgentAccessTokenClaimProviderTest {
+
+    private OAuthTokenReqMessageContext tokReqMsgCtx;
+    private static final String AGENT_USERSTORE = "AGENT_STORE";
+    private static final String AGENT = "AGENT";
+    private static final String AUT = "aut";
+    private static final String ACT = "act";
+    private static final String SUB = "sub";
+    private static final String REQUESTED_ACTOR = "requested-actor";
+
+    @BeforeMethod
+    public void setUp() {
+
+        tokReqMsgCtx = mock(OAuthTokenReqMessageContext.class);
+    }
+
+    @Test
+    public void testGetAdditionalClaims_AgentUserStore() throws Exception {
+
+        try (MockedStatic<IdentityUtil> identityUtil = mockStatic(IdentityUtil.class)) {
+            identityUtil.when(IdentityUtil::getAgentIdentityUserstoreName).thenReturn(AGENT_USERSTORE);
+            when(tokReqMsgCtx.getAuthorizedUser())
+                    .thenReturn(mock(org.wso2.carbon.identity.application.authentication.framework.model
+                            .AuthenticatedUser.class));
+            when(tokReqMsgCtx.getAuthorizedUser().getUserStoreDomain()).thenReturn(AGENT_USERSTORE);
+
+            AgentAccessTokenClaimProvider provider = new AgentAccessTokenClaimProvider();
+            Map<String, Object> claims = provider.getAdditionalClaims(tokReqMsgCtx);
+
+            assertNotNull(claims);
+            assertEquals(claims.get(AUT), AGENT);
+        }
+    }
+
+    @Test
+    public void testGetAdditionalClaims_AuthorizationCodeGrantWithRequestedActor() throws Exception {
+
+        when(tokReqMsgCtx.getOauth2AccessTokenReqDTO()).thenReturn(mock(org.wso2.carbon.identity.oauth2.dto
+                .OAuth2AccessTokenReqDTO.class));
+        when(tokReqMsgCtx.getOauth2AccessTokenReqDTO().getGrantType()).thenReturn("authorization_code");
+        when(tokReqMsgCtx.getRequestedActor()).thenReturn(REQUESTED_ACTOR);
+
+        try (MockedStatic<IdentityUtil> identityUtil = mockStatic(IdentityUtil.class)) {
+            identityUtil.when(IdentityUtil::getAgentIdentityUserstoreName).thenReturn("AGENT_USER_STORE");
+            when(tokReqMsgCtx.getAuthorizedUser()).thenReturn(mock(org.wso2.carbon.identity.application.authentication
+                    .framework.model.AuthenticatedUser.class));
+            when(tokReqMsgCtx.getAuthorizedUser().getUserStoreDomain()).thenReturn("PRIMARY_USER_STORE");
+
+            AgentAccessTokenClaimProvider provider = new AgentAccessTokenClaimProvider();
+            Map<String, Object> claims = provider.getAdditionalClaims(tokReqMsgCtx);
+
+            assertNotNull(claims);
+            assertTrue(claims.containsKey(ACT));
+            Map<String, Object> actMap = (Map<String, Object>) claims.get(ACT);
+            assertEquals(actMap.get(SUB), REQUESTED_ACTOR);
+            assertEquals(actMap.get(AUT), AGENT);
+        }
+    }
+
+    @Test
+    public void testGetAdditionalClaims_NoAgentClaims() throws Exception {
+
+        try (MockedStatic<IdentityUtil> identityUtil = mockStatic(IdentityUtil.class)) {
+            identityUtil.when(IdentityUtil::getAgentIdentityUserstoreName).thenReturn("AGENT_USER_STORE");
+            when(tokReqMsgCtx.getAuthorizedUser()).thenReturn(mock(org.wso2.carbon.identity.application.authentication
+                    .framework.model.AuthenticatedUser.class));
+            when(tokReqMsgCtx.getAuthorizedUser().getUserStoreDomain()).thenReturn("PRIMARY_USER_STORE");
+            when(tokReqMsgCtx.getOauth2AccessTokenReqDTO()).thenReturn(mock(org.wso2.carbon.identity.oauth2
+                    .dto.OAuth2AccessTokenReqDTO.class));
+            when(tokReqMsgCtx.getOauth2AccessTokenReqDTO().getGrantType()).thenReturn("password");
+            when(tokReqMsgCtx.getRequestedActor()).thenReturn(null);
+
+            AgentAccessTokenClaimProvider provider = new AgentAccessTokenClaimProvider();
+            Map<String, Object> claims = provider.getAdditionalClaims(tokReqMsgCtx);
+
+            assertNull(claims);
+        }
+    }
+
+    @Test
+    public void testGetAdditionalClaims_AuthzReqContext() throws Exception {
+
+        AgentAccessTokenClaimProvider provider = new AgentAccessTokenClaimProvider();
+        OAuthAuthzReqMessageContext authzCtx = mock(OAuthAuthzReqMessageContext.class);
+        assertNull(provider.getAdditionalClaims(authzCtx));
+    }
+}

components/org.wso2.carbon.identity.oauth/src/test/resources/testng.xml

@@ -132,6 +132,7 @@
             <class name="org.wso2.carbon.identity.oauth2.device.codegenerator.GenerateKeysTest"/>
             <class name="org.wso2.carbon.identity.oauth2.responsemode.provider.ResponseModeProviderTest"/>
             <class name="org.wso2.carbon.identity.oauth2.token.handlers.claims.ImpersonatedAccessTokenClaimProviderTest"/>
+            <class name="org.wso2.carbon.identity.oauth2.token.handlers.claims.AgentAccessTokenClaimProviderTest"/>
             <class name="org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheTest"/>
             <class name="org.wso2.carbon.identity.oauth2.rar.AuthorizationDetailsServiceTest" />
             <class name="org.wso2.carbon.identity.oauth2.rar.token.AccessTokenResponseRARHandlerTest"/>