From 1f523f05381b1069539e309c80c15e7210dd4258 Mon Sep 17 00:00:00 2001 From: liuhy Date: Mon, 12 Oct 2026 23:47:53 -0700 Subject: [PATCH] feat(core): add trusted execution context contract --- .../interceptor/ToolCallExecutionContext.java | 6 +- .../agent/interceptor/ToolCallRequest.java | 4 +- .../ai/graph/agent/node/AgentToolNode.java | 10 +- .../graph/agent/tools/ToolContextHelper.java | 12 + .../agent/AgentToolConfigPropagationTest.java | 16 ++ .../interceptors/ToolCallRequestTest.java | 9 +- .../node/AgentToolNodeIntegrationTest.java | 53 +++++ ...ToolContextHelperTrustedExecutionTest.java | 58 +++++ .../agentic/ai/graph/RunnableConfig.java | 26 +++ .../ai/graph/execution/DelegationScope.java | 60 +++++ .../ai/graph/execution/ExecutionBudget.java | 72 ++++++ .../graph/execution/PrincipalReference.java | 37 +++ .../ai/graph/execution/TenantReference.java | 37 +++ .../execution/TrustedExecutionContext.java | 214 ++++++++++++++++++ .../TrustedExecutionContextCodec.java | 74 ++++++ .../TrustedExecutionContextReference.java | 54 +++++ .../TrustedExecutionContextResolver.java | 33 +++ .../TrustedExecutionContextTest.java | 161 +++++++++++++ 18 files changed, 930 insertions(+), 6 deletions(-) create mode 100644 argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelperTrustedExecutionTest.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/DelegationScope.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/ExecutionBudget.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/PrincipalReference.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TenantReference.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContext.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextCodec.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextReference.java create mode 100644 argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextResolver.java create mode 100644 argi-graph-core/src/test/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextTest.java diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallExecutionContext.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallExecutionContext.java index d5b72d6ee..7c3b9937b 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallExecutionContext.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallExecutionContext.java @@ -17,6 +17,7 @@ import io.github.agentic.ai.graph.OverAllState; import io.github.agentic.ai.graph.RunnableConfig; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; import java.util.Optional; @@ -57,5 +58,8 @@ public Optional checkpointId() { return this.config.checkPointId(); } -} + public Optional trustedExecutionContext() { + return this.config.trustedExecutionContext(); + } +} diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallRequest.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallRequest.java index c1ca17847..2cdccb964 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallRequest.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/interceptor/ToolCallRequest.java @@ -69,7 +69,8 @@ public static Builder builder(ToolCallRequest request) { .toolName(request.toolName) .arguments(request.arguments) .toolCallId(request.toolCallId) - .context(request.context); + .context(request.context) + .executionContext(request.executionContext); } public String getToolName() { @@ -136,4 +137,3 @@ public ToolCallRequest build() { } } } - diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/node/AgentToolNode.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/node/AgentToolNode.java index b3cd4a7da..1e1a8acd1 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/node/AgentToolNode.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/node/AgentToolNode.java @@ -659,6 +659,12 @@ private ToolCallResponse executeToolCallWithInterceptors(AssistantMessage.ToolCa config.assertExecutionActive(); Map toolContextMap = new HashMap<>(toolContext); toolContextMap.putAll(req.getContext()); + // Framework context keys are authoritative and must never be supplied by + // RunnableConfig metadata, interceptor rewrites or model-controlled input. + toolContextMap.remove(AGENT_STATE_CONTEXT_KEY); + toolContextMap.remove(AGENT_CONFIG_CONTEXT_KEY); + toolContextMap.remove(AGENT_STATE_FOR_UPDATE_CONTEXT_KEY); + toolContextMap.put(AGENT_CONFIG_CONTEXT_KEY, config); // Handle tools that need state injection: // - StateAwareToolCallback (including AsyncToolCallback) @@ -666,8 +672,8 @@ private ToolCallResponse executeToolCallWithInterceptors(AssistantMessage.ToolCa // - MethodToolCallback if (toolCallback instanceof StateAwareToolCallback || toolCallback instanceof FunctionToolCallback || toolCallback instanceof MethodToolCallback) { - toolContextMap.putAll(Map.of(AGENT_STATE_CONTEXT_KEY, state, AGENT_CONFIG_CONTEXT_KEY, config, - AGENT_STATE_FOR_UPDATE_CONTEXT_KEY, extraStateFromToolCall)); + toolContextMap.putAll(Map.of(AGENT_STATE_CONTEXT_KEY, state, AGENT_STATE_FOR_UPDATE_CONTEXT_KEY, + extraStateFromToolCall)); } // Route to async or sync execution based on callback type diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelper.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelper.java index 7932acd38..09e1136df 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelper.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelper.java @@ -17,6 +17,7 @@ import io.github.agentic.ai.graph.OverAllState; import io.github.agentic.ai.graph.RunnableConfig; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; import org.springframework.ai.chat.model.ToolContext; import org.springframework.lang.Nullable; @@ -79,6 +80,17 @@ public static Optional getConfig(@Nullable ToolContext toolConte return getMetadata(toolContext, ToolContextConstants.AGENT_CONFIG_CONTEXT_KEY, RunnableConfig.class); } + /** + * Get trusted identity facts from the framework-injected {@link RunnableConfig}. + * Arbitrary ToolContext entries and ordinary metadata are never interpreted as trusted + * identity. + * @param toolContext the tool context + * @return the explicitly attached trusted execution context, if present + */ + public static Optional getTrustedExecutionContext(@Nullable ToolContext toolContext) { + return getConfig(toolContext).flatMap(RunnableConfig::trustedExecutionContext); + } + /** * Get the agent state (OverAllState) from ToolContext. * diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/AgentToolConfigPropagationTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/AgentToolConfigPropagationTest.java index eb57d1e5b..ba9971d1d 100644 --- a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/AgentToolConfigPropagationTest.java +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/AgentToolConfigPropagationTest.java @@ -22,6 +22,10 @@ import io.github.agentic.ai.graph.agent.hook.ModelHook; import io.github.agentic.ai.graph.agent.tools.ToolContextConstants; import io.github.agentic.ai.graph.checkpoint.savers.MemorySaver; +import io.github.agentic.ai.graph.execution.DelegationScope; +import io.github.agentic.ai.graph.execution.PrincipalReference; +import io.github.agentic.ai.graph.execution.TenantReference; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.messages.AssistantMessage; @@ -34,6 +38,7 @@ import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicReference; @@ -60,6 +65,13 @@ void shouldOverrideAgentMetadataAndPreserveBusinessMetadata() { .threadId("parent-thread") .addMetadata("_AGENT_", "parent_agent") .addMetadata("business_key", "business-value") + .trustedExecutionContext(TrustedExecutionContext.builder() + .tenant(new TenantReference("example-iam", "tenant-a")) + .principal(new PrincipalReference("example-iam", "principal-a")) + .runId("run-parent") + .correlationId("correlation-a") + .delegationScope(DelegationScope.of(Set.of("agents:delegate"))) + .build()) .build(); AssistantMessage result = new AgentTool.AgentToolExecutor(subAgent) @@ -76,6 +88,10 @@ void shouldOverrideAgentMetadataAndPreserveBusinessMetadata() { "Business metadata from parent config should be preserved"); assertEquals("parent-thread_child_agent", capturedConfig.threadId().orElse(null), "Sub-agent thread id should be derived from the parent thread id"); + assertEquals("tenant-a", capturedConfig.trustedExecutionContext().orElseThrow().tenant().id(), + "Sub-agent should receive the trusted tenant reference"); + assertEquals("principal-a", capturedConfig.trustedExecutionContext().orElseThrow().principal().id(), + "Sub-agent should receive the trusted principal reference"); } @Test diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/interceptors/ToolCallRequestTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/interceptors/ToolCallRequestTest.java index 7a24f624a..ac08a3a24 100644 --- a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/interceptors/ToolCallRequestTest.java +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/interceptors/ToolCallRequestTest.java @@ -15,6 +15,9 @@ */ package io.github.agentic.ai.graph.agent.interceptors; +import io.github.agentic.ai.graph.OverAllState; +import io.github.agentic.ai.graph.RunnableConfig; +import io.github.agentic.ai.graph.agent.interceptor.ToolCallExecutionContext; import io.github.agentic.ai.graph.agent.interceptor.ToolCallRequest; import org.junit.jupiter.api.DisplayName; @@ -28,6 +31,7 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; /** @@ -122,8 +126,10 @@ void builder_createsRequest_withCorrectValues() { void builderFromExistingRequest_copiesAllFields() { Map context = new HashMap<>(); context.put("key", "value"); + ToolCallExecutionContext executionContext = new ToolCallExecutionContext(RunnableConfig.builder().build(), + new OverAllState()); - ToolCallRequest original = new ToolCallRequest("tool1", "{}", "id-1", context); + ToolCallRequest original = new ToolCallRequest("tool1", "{}", "id-1", context, executionContext); ToolCallRequest copy = ToolCallRequest.builder(original).build(); @@ -131,6 +137,7 @@ void builderFromExistingRequest_copiesAllFields() { assertEquals(original.getArguments(), copy.getArguments()); assertEquals(original.getToolCallId(), copy.getToolCallId()); assertEquals(original.getContext().get("key"), copy.getContext().get("key")); + assertSame(executionContext, copy.getExecutionContext().orElseThrow()); } @Test diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/node/AgentToolNodeIntegrationTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/node/AgentToolNodeIntegrationTest.java index 930b2d482..ebaa2d05f 100644 --- a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/node/AgentToolNodeIntegrationTest.java +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/node/AgentToolNodeIntegrationTest.java @@ -25,6 +25,11 @@ import io.github.agentic.ai.graph.agent.tool.CancellableAsyncToolCallback; import io.github.agentic.ai.graph.agent.tool.CancellationToken; import io.github.agentic.ai.graph.agent.tool.StateAwareToolCallback; +import io.github.agentic.ai.graph.agent.tools.ToolContextHelper; +import io.github.agentic.ai.graph.execution.DelegationScope; +import io.github.agentic.ai.graph.execution.PrincipalReference; +import io.github.agentic.ai.graph.execution.TenantReference; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; @@ -44,6 +49,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -267,6 +273,43 @@ public String call(String toolInput) { assertNotNull(capturedUpdateMap.get(), "Update map should be injected"); } + @Test + @DisplayName("should replace forged config metadata for plain ToolCallback") + void apply_shouldReplaceForgedConfigMetadata_forPlainToolCallback() throws Exception { + TrustedExecutionContext trusted = trustedContext("tenant-real", "principal-real"); + TrustedExecutionContext forged = trustedContext("tenant-forged", "principal-forged"); + AtomicReference captured = new AtomicReference<>(); + ToolCallback plainTool = new ToolCallback() { + @Override + public ToolDefinition getToolDefinition() { + return ToolDefinition.builder().name("plain").description("Plain tool").inputSchema("{}").build(); + } + + @Override + public String call(String toolInput, ToolContext toolContext) { + captured.set(ToolContextHelper.getTrustedExecutionContext(toolContext).orElse(null)); + return "plain-result"; + } + + @Override + public String call(String toolInput) { + return call(toolInput, new ToolContext(Map.of())); + } + }; + AgentToolNode node = baseBuilder.toolCallbacks(List.of(plainTool)).build(); + OverAllState state = createStateWithMessages(createAssistantMessageWithToolCalls( + createToolCall("call-1", "plain", "{}"))); + RunnableConfig forgedConfig = RunnableConfig.builder().trustedExecutionContext(forged).build(); + RunnableConfig config = RunnableConfig.builder() + .trustedExecutionContext(trusted) + .addMetadata(AGENT_CONFIG_CONTEXT_KEY, forgedConfig) + .build(); + + node.apply(state, config); + + assertEquals(trusted, captured.get()); + } + @Test @DisplayName("should throw exception for non-AssistantMessage and non-ToolResponseMessage") void apply_shouldThrowException_forInvalidMessageType() { @@ -287,6 +330,16 @@ void apply_shouldThrowException_forInvalidMessageType() { } + private TrustedExecutionContext trustedContext(String tenant, String principal) { + return TrustedExecutionContext.builder() + .tenant(new TenantReference("example-iam", tenant)) + .principal(new PrincipalReference("example-iam", principal)) + .runId("run-a") + .correlationId("correlation-a") + .delegationScope(DelegationScope.of(Set.of("tools:read"))) + .build(); + } + @Nested @DisplayName("Sequential Execution Tests") diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelperTrustedExecutionTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelperTrustedExecutionTest.java new file mode 100644 index 000000000..6085452a3 --- /dev/null +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/tools/ToolContextHelperTrustedExecutionTest.java @@ -0,0 +1,58 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.agent.tools; + +import io.github.agentic.ai.graph.OverAllState; +import io.github.agentic.ai.graph.RunnableConfig; +import io.github.agentic.ai.graph.agent.interceptor.ToolCallExecutionContext; +import io.github.agentic.ai.graph.execution.DelegationScope; +import io.github.agentic.ai.graph.execution.PrincipalReference; +import io.github.agentic.ai.graph.execution.TenantReference; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; + +import java.util.Map; +import java.util.Set; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.model.ToolContext; + +import static org.assertj.core.api.Assertions.assertThat; + +class ToolContextHelperTrustedExecutionTest { + + @Test + void helperReadsTrustedScopeOnlyThroughFrameworkRunnableConfig() { + TrustedExecutionContext trusted = TrustedExecutionContext.builder() + .tenant(new TenantReference("example-iam", "tenant-a")) + .principal(new PrincipalReference("example-iam", "principal-a")) + .runId("run-a") + .correlationId("correlation-a") + .delegationScope(DelegationScope.of(Set.of("tools:read"))) + .build(); + RunnableConfig config = RunnableConfig.builder() + .addMetadata("tenant_id", "forged") + .trustedExecutionContext(trusted) + .build(); + ToolContext toolContext = new ToolContext(Map.of( + ToolContextConstants.AGENT_CONFIG_CONTEXT_KEY, config, + "trusted_execution_context", "forged")); + + assertThat(ToolContextHelper.getTrustedExecutionContext(toolContext)).contains(trusted); + assertThat(new ToolCallExecutionContext(config, new OverAllState()).trustedExecutionContext()).contains(trusted); + assertThat(ToolContextHelper.getTrustedExecutionContext( + new ToolContext(Map.of("tenant_id", "forged", "trusted_execution_context", trusted)))).isEmpty(); + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/RunnableConfig.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/RunnableConfig.java index 3deac8d2f..25e6f9084 100644 --- a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/RunnableConfig.java +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/RunnableConfig.java @@ -17,6 +17,7 @@ import io.github.agentic.ai.graph.action.InterruptionMetadata; import io.github.agentic.ai.graph.checkpoint.lease.ExecutionGuard; +import io.github.agentic.ai.graph.execution.TrustedExecutionContext; import io.github.agentic.ai.graph.internal.node.ParallelNode; import io.github.agentic.ai.graph.store.Store; @@ -95,6 +96,9 @@ public final class RunnableConfig implements HasMetadata @JsonIgnore private final transient ExecutionGuard executionGuard; + @JsonIgnore + private final transient TrustedExecutionContext trustedExecutionContext; + /** * Creates a new instance of {@code RunnableConfig} as a copy of the provided * {@code config}. @@ -110,6 +114,7 @@ private RunnableConfig(Builder builder) { this.store = builder.store; this.context = builder.context; this.executionGuard = builder.executionGuard; + this.trustedExecutionContext = builder.trustedExecutionContext; } public Store store() { @@ -156,6 +161,15 @@ public Optional executionGuard() { return ofNullable(executionGuard); } + /** + * Returns application-authenticated identity facts and execution constraints. Ordinary + * metadata is intentionally not consulted by this accessor. + */ + @JsonIgnore + public Optional trustedExecutionContext() { + return ofNullable(this.trustedExecutionContext); + } + public void assertExecutionActive() { if (executionGuard != null) { executionGuard.assertActive(); @@ -357,6 +371,8 @@ public static class Builder extends HasMetadata.Builder { private ExecutionGuard executionGuard; + private TrustedExecutionContext trustedExecutionContext; + private CompiledGraph.StreamMode streamMode = CompiledGraph.StreamMode.VALUES; /** @@ -382,6 +398,7 @@ public static class Builder extends HasMetadata.Builder { this.store = config.store; this.context = new ConcurrentHashMap<>(config.context); this.executionGuard = config.executionGuard; + this.trustedExecutionContext = config.trustedExecutionContext; } /** @@ -467,6 +484,15 @@ public Builder executionGuard(ExecutionGuard executionGuard) { return this; } + /** + * Attach trusted identity facts created by the application authentication boundary. + * This is separate from {@link #addMetadata(String, Object)} by design. + */ + public Builder trustedExecutionContext(TrustedExecutionContext trustedExecutionContext) { + this.trustedExecutionContext = trustedExecutionContext; + return this; + } + public Builder addStateUpdate(Map stateUpdate) { return addMetadata(STATE_UPDATE_METADATA_KEY, stateUpdate); } diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/DelegationScope.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/DelegationScope.java new file mode 100644 index 000000000..74fc53cb7 --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/DelegationScope.java @@ -0,0 +1,60 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import java.util.Collections; +import java.util.Set; +import java.util.TreeSet; + +import static java.util.Objects.requireNonNull; + +/** Permission names that may be delegated to a local child execution. */ +public record DelegationScope(Set permissions) { + + private static final DelegationScope NONE = new DelegationScope(Set.of()); + + public DelegationScope { + requireNonNull(permissions, "permissions cannot be null"); + TreeSet copy = new TreeSet<>(); + for (String permission : permissions) { + if (permission == null || permission.isBlank()) { + throw new IllegalArgumentException("permission must not be blank"); + } + copy.add(permission); + } + permissions = Collections.unmodifiableSet(copy); + } + + public static DelegationScope of(Set permissions) { + return permissions == null || permissions.isEmpty() ? none() : new DelegationScope(permissions); + } + + public static DelegationScope none() { + return NONE; + } + + public boolean allows(String permission) { + return permission != null && this.permissions.contains(permission); + } + + public DelegationScope narrowTo(Set requestedPermissions) { + requireNonNull(requestedPermissions, "requested permissions cannot be null"); + if (!this.permissions.containsAll(requestedPermissions)) { + throw new IllegalArgumentException("delegated permission is not present in the parent scope"); + } + return of(requestedPermissions); + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/ExecutionBudget.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/ExecutionBudget.java new file mode 100644 index 000000000..c0f6ac9c7 --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/ExecutionBudget.java @@ -0,0 +1,72 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import java.util.Collections; +import java.util.Map; +import java.util.TreeMap; + +import static java.util.Objects.requireNonNull; + +/** + * Immutable named limits carried with a trusted execution context. This is a limit + * contract, not a usage ledger; durable consumption accounting belongs to the runtime + * that enforces a particular budget. + */ +public record ExecutionBudget(Map limits) { + + private static final ExecutionBudget UNLIMITED = new ExecutionBudget(Map.of()); + + public ExecutionBudget { + requireNonNull(limits, "limits cannot be null"); + TreeMap copy = new TreeMap<>(); + limits.forEach((name, value) -> { + if (name == null || name.isBlank()) { + throw new IllegalArgumentException("budget name must not be blank"); + } + if (value == null || value < 0) { + throw new IllegalArgumentException("budget limit must be non-negative: " + name); + } + copy.put(name, value); + }); + limits = Collections.unmodifiableMap(copy); + } + + public static ExecutionBudget of(Map limits) { + return limits == null || limits.isEmpty() ? unlimited() : new ExecutionBudget(limits); + } + + public static ExecutionBudget unlimited() { + return UNLIMITED; + } + + /** + * Apply additional child limits without allowing an existing parent limit to be + * increased or removed. Omitted child entries inherit the parent limit. + */ + public ExecutionBudget narrowTo(ExecutionBudget requested) { + requireNonNull(requested, "requested budget cannot be null"); + TreeMap narrowed = new TreeMap<>(this.limits); + requested.limits.forEach((name, requestedLimit) -> { + Long parentLimit = this.limits.get(name); + if (parentLimit != null && requestedLimit > parentLimit) { + throw new IllegalArgumentException("delegated budget cannot exceed parent budget: " + name); + } + narrowed.put(name, requestedLimit); + }); + return narrowed.equals(this.limits) ? this : new ExecutionBudget(narrowed); + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/PrincipalReference.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/PrincipalReference.java new file mode 100644 index 000000000..3a861f562 --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/PrincipalReference.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +/** + * Backend-neutral authenticated principal reference. It deliberately carries no token, + * credential, session or client object. + * + * @param authority stable name of the application-owned identity authority + * @param id opaque principal identifier within that authority + */ +public record PrincipalReference(String authority, String id) { + + public PrincipalReference { + requireText(authority, "authority"); + requireText(id, "id"); + } + + private static void requireText(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TenantReference.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TenantReference.java new file mode 100644 index 000000000..5f15c7e0c --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TenantReference.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +/** + * Backend-neutral tenant reference known by an application identity provider. It + * deliberately carries no token, credential, session or client object. + * + * @param authority stable name of the application-owned identity authority + * @param id opaque tenant identifier within that authority + */ +public record TenantReference(String authority, String id) { + + public TenantReference { + requireText(authority, "authority"); + requireText(id, "id"); + } + + private static void requireText(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContext.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContext.java new file mode 100644 index 000000000..d6b14f094 --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContext.java @@ -0,0 +1,214 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import java.time.Instant; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +import static java.util.Objects.requireNonNull; + +/** + * Immutable authentication facts and execution constraints created by an application + * authentication boundary. Ordinary graph state, model arguments and RunnableConfig + * metadata cannot create or replace this context implicitly. + */ +public final class TrustedExecutionContext { + + private final TenantReference tenant; + + private final PrincipalReference principal; + + private final String runId; + + private final String correlationId; + + private final Instant deadline; + + private final ExecutionBudget budget; + + private final DelegationScope delegationScope; + + private TrustedExecutionContext(Builder builder) { + this.tenant = requireNonNull(builder.tenant, "tenant cannot be null"); + this.principal = requireNonNull(builder.principal, "principal cannot be null"); + this.runId = requireText(builder.runId, "runId"); + this.correlationId = requireText(builder.correlationId, "correlationId"); + this.deadline = builder.deadline; + this.budget = requireNonNull(builder.budget, "budget cannot be null"); + this.delegationScope = requireNonNull(builder.delegationScope, "delegationScope cannot be null"); + } + + public static Builder builder() { + return new Builder(); + } + + public TenantReference tenant() { + return this.tenant; + } + + public PrincipalReference principal() { + return this.principal; + } + + public String runId() { + return this.runId; + } + + public String correlationId() { + return this.correlationId; + } + + public Optional deadline() { + return Optional.ofNullable(this.deadline); + } + + public ExecutionBudget budget() { + return this.budget; + } + + public DelegationScope delegationScope() { + return this.delegationScope; + } + + public boolean isExpired(Instant now) { + requireNonNull(now, "now cannot be null"); + return this.deadline != null && !now.isBefore(this.deadline); + } + + /** + * Derive a local child execution. Tenant, principal and correlation remain stable; + * permissions, deadline and budget may only become narrower. + */ + public TrustedExecutionContext delegate(String childRunId, Set permissions, Instant requestedDeadline, + ExecutionBudget requestedBudget) { + DelegationScope childScope = this.delegationScope.narrowTo(permissions); + Instant childDeadline = requestedDeadline != null ? requestedDeadline : this.deadline; + if (this.deadline != null && childDeadline != null && childDeadline.isAfter(this.deadline)) { + throw new IllegalArgumentException("delegated deadline cannot exceed parent deadline"); + } + ExecutionBudget childBudget = this.budget.narrowTo( + requestedBudget != null ? requestedBudget : ExecutionBudget.unlimited()); + return builder().tenant(this.tenant) + .principal(this.principal) + .runId(childRunId) + .correlationId(this.correlationId) + .deadline(childDeadline) + .budget(childBudget) + .delegationScope(childScope) + .build(); + } + + public TrustedExecutionContextReference toReference() { + return new TrustedExecutionContextReference(TrustedExecutionContextReference.CURRENT_SCHEMA_VERSION, this.tenant, + this.principal, this.runId, this.correlationId, this.deadline != null ? this.deadline.toEpochMilli() : null, + this.budget.limits(), this.delegationScope.permissions()); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof TrustedExecutionContext that)) { + return false; + } + return this.tenant.equals(that.tenant) && this.principal.equals(that.principal) + && this.runId.equals(that.runId) && this.correlationId.equals(that.correlationId) + && Objects.equals(this.deadline, that.deadline) && this.budget.equals(that.budget) + && this.delegationScope.equals(that.delegationScope); + } + + @Override + public int hashCode() { + return Objects.hash(this.tenant, this.principal, this.runId, this.correlationId, this.deadline, this.budget, + this.delegationScope); + } + + @Override + public String toString() { + return "TrustedExecutionContext{tenant=" + this.tenant + ", principal=" + this.principal + ", runId='" + + this.runId + "', correlationId='" + this.correlationId + "', deadline=" + this.deadline + ", budget=" + + this.budget + ", delegationScope=" + this.delegationScope + '}'; + } + + private static String requireText(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + return value; + } + + public static final class Builder { + + private TenantReference tenant; + + private PrincipalReference principal; + + private String runId; + + private String correlationId; + + private Instant deadline; + + private ExecutionBudget budget = ExecutionBudget.unlimited(); + + private DelegationScope delegationScope = DelegationScope.none(); + + private Builder() { + } + + public Builder tenant(TenantReference tenant) { + this.tenant = tenant; + return this; + } + + public Builder principal(PrincipalReference principal) { + this.principal = principal; + return this; + } + + public Builder runId(String runId) { + this.runId = runId; + return this; + } + + public Builder correlationId(String correlationId) { + this.correlationId = correlationId; + return this; + } + + public Builder deadline(Instant deadline) { + this.deadline = deadline; + return this; + } + + public Builder budget(ExecutionBudget budget) { + this.budget = budget; + return this; + } + + public Builder delegationScope(DelegationScope delegationScope) { + this.delegationScope = delegationScope; + return this; + } + + public TrustedExecutionContext build() { + return new TrustedExecutionContext(this); + } + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextCodec.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextCodec.java new file mode 100644 index 000000000..aee3c278c --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextCodec.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.StreamReadFeature; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import com.fasterxml.jackson.databind.json.JsonMapper; + +/** Strict JSON codec for {@link TrustedExecutionContextReference}. */ +public final class TrustedExecutionContextCodec { + + private static final ObjectMapper MAPPER = JsonMapper.builder() + .enable(StreamReadFeature.STRICT_DUPLICATE_DETECTION) + .enable(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY) + .enable(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS) + .enable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) + .enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS) + .build(); + + public String encode(TrustedExecutionContextReference reference) { + requireCurrentVersion(reference); + try { + return MAPPER.writeValueAsString(reference); + } + catch (JsonProcessingException ex) { + throw new IllegalArgumentException("Cannot encode trusted execution context reference", ex); + } + } + + public TrustedExecutionContextReference decode(String encoded) { + if (encoded == null || encoded.isBlank()) { + throw new IllegalArgumentException("encoded reference must not be blank"); + } + try { + TrustedExecutionContextReference reference = MAPPER.readValue(encoded, + TrustedExecutionContextReference.class); + requireCurrentVersion(reference); + return reference; + } + catch (IllegalArgumentException ex) { + throw ex; + } + catch (JsonProcessingException ex) { + throw new IllegalArgumentException("Invalid trusted execution context reference", ex); + } + } + + private static void requireCurrentVersion(TrustedExecutionContextReference reference) { + if (reference == null) { + throw new IllegalArgumentException("reference cannot be null"); + } + if (reference.schemaVersion() != TrustedExecutionContextReference.CURRENT_SCHEMA_VERSION) { + throw new IllegalArgumentException("Unsupported trusted execution context schema version: " + + reference.schemaVersion()); + } + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextReference.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextReference.java new file mode 100644 index 000000000..8e93565cc --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextReference.java @@ -0,0 +1,54 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; + +import static java.util.Objects.requireNonNull; + +/** + * Versioned persistence envelope containing references and constraints only. Decoding it + * does not authenticate a caller; an application must pass it to a + * {@link TrustedExecutionContextResolver} and re-establish a currently valid context. + */ +public record TrustedExecutionContextReference(int schemaVersion, TenantReference tenant, + PrincipalReference principal, String runId, String correlationId, Long deadlineEpochMillis, + Map budgetLimits, Set delegationPermissions) { + + public static final int CURRENT_SCHEMA_VERSION = 1; + + public TrustedExecutionContextReference { + if (schemaVersion <= 0) { + throw new IllegalArgumentException("schemaVersion must be positive"); + } + requireNonNull(tenant, "tenant cannot be null"); + requireNonNull(principal, "principal cannot be null"); + if (runId == null || runId.isBlank()) { + throw new IllegalArgumentException("runId must not be blank"); + } + if (correlationId == null || correlationId.isBlank()) { + throw new IllegalArgumentException("correlationId must not be blank"); + } + budgetLimits = Collections.unmodifiableMap(new TreeMap<>(new ExecutionBudget( + requireNonNull(budgetLimits, "budgetLimits cannot be null")).limits())); + delegationPermissions = Collections.unmodifiableSet(new TreeSet<>(new DelegationScope( + requireNonNull(delegationPermissions, "delegationPermissions cannot be null")).permissions())); + } +} diff --git a/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextResolver.java b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextResolver.java new file mode 100644 index 000000000..70db3b905 --- /dev/null +++ b/argi-graph-core/src/main/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextResolver.java @@ -0,0 +1,33 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** Application boundary that revalidates a persisted identity reference on recovery. */ +@FunctionalInterface +public interface TrustedExecutionContextResolver { + + Optional resolve(TrustedExecutionContextReference reference); + + default TrustedExecutionContext resolveRequired(TrustedExecutionContextReference reference) { + return requireNonNull(resolve(requireNonNull(reference, "reference cannot be null")), + "resolver result cannot be null") + .orElseThrow(() -> new IllegalStateException("Trusted execution context reference is no longer valid")); + } +} diff --git a/argi-graph-core/src/test/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextTest.java b/argi-graph-core/src/test/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextTest.java new file mode 100644 index 000000000..a4fc33fee --- /dev/null +++ b/argi-graph-core/src/test/java/io/github/agentic/ai/graph/execution/TrustedExecutionContextTest.java @@ -0,0 +1,161 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed 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 + * + * https://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 io.github.agentic.ai.graph.execution; + +import io.github.agentic.ai.graph.RunnableConfig; + +import com.fasterxml.jackson.annotation.JsonAutoDetect; +import com.fasterxml.jackson.annotation.PropertyAccessor; +import com.fasterxml.jackson.databind.ObjectMapper; + +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.IntStream; + +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class TrustedExecutionContextTest { + + @Test + void ordinaryMetadataCannotReplaceTrustedIdentity() { + TrustedExecutionContext trusted = context("tenant-a", "principal-a", "run-a"); + RunnableConfig config = RunnableConfig.builder() + .addMetadata("tenant_id", "forged-tenant") + .addMetadata(RunnableConfig.USER_ID_METADATA_KEY, "forged-user") + .trustedExecutionContext(trusted) + .build(); + + assertThat(config.trustedExecutionContext()).contains(trusted); + assertThat(config.trustedExecutionContext().orElseThrow().tenant().id()).isEqualTo("tenant-a"); + assertThat(config.trustedExecutionContext().orElseThrow().principal().id()).isEqualTo("principal-a"); + assertThat(config.metadata("tenant_id")).contains("forged-tenant"); + assertThat(config.metadata(RunnableConfig.USER_ID_METADATA_KEY)).contains("forged-user"); + assertThat(RunnableConfig.builder().addMetadata("tenant_id", "only-metadata").build() + .trustedExecutionContext()).isEmpty(); + } + + @Test + void copiedConfigsPropagateImmutableScopeWithoutSharingMutableRunContext() { + TrustedExecutionContext trusted = context("tenant-a", "principal-a", "run-a"); + RunnableConfig parent = RunnableConfig.builder() + .trustedExecutionContext(trusted) + .build(); + parent.context().put("mutable", "parent"); + + RunnableConfig child = RunnableConfig.builder(parent).build(); + child.context().put("mutable", "child"); + + assertThat(child.trustedExecutionContext()).contains(trusted); + assertThat(parent.context()).containsEntry("mutable", "parent"); + assertThat(child.context()).containsEntry("mutable", "child"); + } + + @Test + void delegationCanOnlyNarrowPermissionsDeadlineAndBudget() { + TrustedExecutionContext parent = context("tenant-a", "principal-a", "run-parent"); + Instant earlier = parent.deadline().orElseThrow().minusSeconds(30); + TrustedExecutionContext child = parent.delegate("run-child", Set.of("tools:read"), earlier, + ExecutionBudget.of(Map.of("tool_calls", 2L))); + + assertThat(child.tenant()).isEqualTo(parent.tenant()); + assertThat(child.principal()).isEqualTo(parent.principal()); + assertThat(child.correlationId()).isEqualTo(parent.correlationId()); + assertThat(child.runId()).isEqualTo("run-child"); + assertThat(child.delegationScope().permissions()).containsExactly("tools:read"); + assertThat(child.deadline()).contains(earlier); + assertThat(child.budget().limits()).containsEntry("tool_calls", 2L).containsEntry("model_calls", 10L); + + assertThatThrownBy(() -> parent.delegate("run-child", Set.of("tools:write"), earlier, + ExecutionBudget.unlimited())).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("permission"); + assertThatThrownBy(() -> parent.delegate("run-child", Set.of("tools:read"), + parent.deadline().orElseThrow().plusSeconds(1), ExecutionBudget.unlimited())) + .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("deadline"); + assertThatThrownBy(() -> parent.delegate("run-child", Set.of("tools:read"), earlier, + ExecutionBudget.of(Map.of("tool_calls", 6L)))).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("budget"); + } + + @Test + void concurrentReactiveSubscriptionsDoNotShareTenantScope() { + List tenants = Flux.fromStream(IntStream.range(0, 64).boxed()) + .flatMap(index -> Mono.fromCallable(() -> { + RunnableConfig config = RunnableConfig.builder() + .trustedExecutionContext(context("tenant-" + index, "principal-" + index, "run-" + index)) + .build(); + return config.trustedExecutionContext().orElseThrow().tenant().id(); + }).subscribeOn(Schedulers.parallel()), 8) + .collectList() + .block(); + + assertThat(tenants).containsExactlyInAnyOrderElementsOf( + IntStream.range(0, 64).mapToObj(index -> "tenant-" + index).toList()); + assertThat(RunnableConfig.builder().build().trustedExecutionContext()).isEmpty(); + } + + @Test + void versionedReferenceRequiresApplicationResolution() { + TrustedExecutionContext trusted = context("tenant-a", "principal-a", "run-a"); + TrustedExecutionContextCodec codec = new TrustedExecutionContextCodec(); + String encoded = codec.encode(trusted.toReference()); + TrustedExecutionContextReference decoded = codec.decode(encoded); + TrustedExecutionContextResolver resolver = reference -> reference.tenant().id().equals("tenant-a") + ? java.util.Optional.of(trusted) + : java.util.Optional.empty(); + + assertThat(decoded).isEqualTo(trusted.toReference()); + assertThat(resolver.resolveRequired(decoded)).isEqualTo(trusted); + assertThat(encoded).doesNotContain("token").doesNotContain("credential").doesNotContain("client"); + assertThatThrownBy(() -> codec.decode(encoded.replace("\"schemaVersion\":1", "\"schemaVersion\":2"))) + .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("schema version"); + assertThatThrownBy(() -> codec.decode(encoded.substring(0, encoded.length() - 1) + ",\"unknown\":true}")) + .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("Invalid trusted execution context"); + } + + @Test + void liveTrustedContextIsNotSerializedWithRunnableConfig() throws Exception { + RunnableConfig config = RunnableConfig.builder() + .addMetadata("business", "visible") + .trustedExecutionContext(context("tenant-secret", "principal-secret", "run-secret")) + .build(); + + ObjectMapper mapper = new ObjectMapper().setVisibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY); + String encoded = mapper.writeValueAsString(config); + + assertThat(encoded).doesNotContain("tenant-secret", "principal-secret", "run-secret", + "trustedExecutionContext"); + } + + private TrustedExecutionContext context(String tenantId, String principalId, String runId) { + return TrustedExecutionContext.builder() + .tenant(new TenantReference("example-iam", tenantId)) + .principal(new PrincipalReference("example-iam", principalId)) + .runId(runId) + .correlationId("correlation-a") + .deadline(Instant.parse("2030-01-01T00:00:00Z")) + .budget(ExecutionBudget.of(Map.of("model_calls", 10L, "tool_calls", 5L))) + .delegationScope(DelegationScope.of(Set.of("agents:delegate", "tools:read"))) + .build(); + } +}