Skip to content

Commit c745ab9

Browse files
FelixTJDietrichclaude
FelixTJDietrich
and
claude
authored
refactor(server): migrate LLM proxy to per-job JWTs (#1662)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent bdcfdcf commit c745ab9

21 files changed

Lines changed: 503 additions & 222 deletions

‎.changeset/tidy-job-jwts.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"hephaestus": patch
3+
---
4+
5+
Agent sandboxes now use short-lived, proxy-scoped JWTs instead of reusable database-backed job secrets.

‎server/application/src/main/java/de/tum/cit/aet/hephaestus/agent/job/AgentJobExecutor.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import de.tum.cit.aet.hephaestus.agent.usage.LlmUsageRecorder;
3535
import de.tum.cit.aet.hephaestus.core.WorkspaceAgnostic;
3636
import de.tum.cit.aet.hephaestus.core.runtime.RuntimeRole;
37+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtIssuer;
3738
import de.tum.cit.aet.hephaestus.evidence.AutomatedReviewReadinessReport;
3839
import de.tum.cit.aet.hephaestus.observability.StructuredLogKeys;
3940
import io.micrometer.core.instrument.Counter;
@@ -148,6 +149,7 @@ public class AgentJobExecutor {
148149
private final WorkspaceAgentBindingRepository bindingRepository;
149150
private final JobTypeHandlerRegistry handlerRegistry;
150151
private final PracticePiAdapter practiceAgent;
152+
private final WorkerJwtIssuer workerJwtIssuer;
151153

152154
private final SandboxManager sandboxManager;
153155
private final AsyncTaskExecutor sandboxExecutor;
@@ -187,6 +189,7 @@ public AgentJobExecutor(
187189
WorkspaceAgentBindingRepository bindingRepository,
188190
JobTypeHandlerRegistry handlerRegistry,
189191
PracticePiAdapter practiceAgent,
192+
WorkerJwtIssuer workerJwtIssuer,
190193
SandboxManager sandboxManager,
191194
@Qualifier("sandboxExecutor") AsyncTaskExecutor sandboxExecutor,
192195
TransactionTemplate transactionTemplate,
@@ -203,6 +206,7 @@ public AgentJobExecutor(
203206
this.bindingRepository = bindingRepository;
204207
this.handlerRegistry = handlerRegistry;
205208
this.practiceAgent = practiceAgent;
209+
this.workerJwtIssuer = workerJwtIssuer;
206210
this.sandboxManager = sandboxManager;
207211
this.sandboxExecutor = sandboxExecutor;
208212
this.transactionTemplate = transactionTemplate;
@@ -715,15 +719,19 @@ private PreparedSandbox prepareSandboxSpec(UUID jobId, AgentJob job, ConfigSnaps
715719
return handler.prepareInputs(managedJob);
716720
});
717721

718-
// Every sandbox reaches the provider through the in-app LLM proxy with the job's own token;
719-
// there is no worker-side BYO-LLM override.
722+
// Sandboxes access providers through the LLM proxy with an attempt-scoped credential.
723+
String jobToken = workerJwtIssuer.issueForJob(
724+
jobId,
725+
job.getWorkspace().getId(),
726+
job.getRetryCount(),
727+
Duration.ofSeconds(snapshot.timeoutSeconds()).plusMinutes(5));
720728
PracticeAgentRequest adapterRequest = new PracticeAgentRequest(
721729
snapshot.apiProtocol(),
722730
snapshot.upstreamModelId(),
723731
snapshot.contextWindow(),
724732
snapshot.maxOutputTokens(),
725733
snapshot.supportsReasoning(),
726-
job.getJobToken(),
734+
jobToken,
727735
snapshot.allowInternet(),
728736
snapshot.timeoutSeconds());
729737

‎server/application/src/main/java/de/tum/cit/aet/hephaestus/agent/proxy/JobTokenAuthenticationFilter.java‎

Lines changed: 30 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
import de.tum.cit.aet.hephaestus.agent.job.AgentJobStatus;
77
import de.tum.cit.aet.hephaestus.agent.usage.LlmPriceSnapshot;
88
import de.tum.cit.aet.hephaestus.agent.usage.LlmUsageSourceType;
9+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.JobJwt;
10+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtInvalidException;
11+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtVerifier;
912
import jakarta.servlet.FilterChain;
1013
import jakarta.servlet.ServletException;
1114
import jakarta.servlet.http.HttpServletRequest;
@@ -14,7 +17,6 @@
1417
import java.math.BigDecimal;
1518
import java.net.InetAddress;
1619
import java.net.UnknownHostException;
17-
import java.security.MessageDigest;
1820
import java.util.Optional;
1921
import java.util.regex.Pattern;
2022
import org.jspecify.annotations.Nullable;
@@ -36,20 +38,21 @@ public class JobTokenAuthenticationFilter extends OncePerRequestFilter {
3638

3739
private static final Logger log = LoggerFactory.getLogger(JobTokenAuthenticationFilter.class);
3840

39-
/** Base64-URL characters (no padding). */
40-
private static final Pattern BASE64_URL_PATTERN = Pattern.compile("^[A-Za-z0-9_-]+$");
41-
4241
private static final String BEARER_PREFIX = "Bearer ";
42+
private static final String LLM_PROXY_SCOPE = "llm_proxy";
4343

4444
private final AgentJobRepository agentJobRepository;
45+
private final WorkerJwtVerifier jwtVerifier;
4546
private final MentorProxyCredentialRegistry mentorRegistry;
4647
private final ObjectMapper objectMapper;
4748

4849
JobTokenAuthenticationFilter(
4950
AgentJobRepository agentJobRepository,
51+
WorkerJwtVerifier jwtVerifier,
5052
MentorProxyCredentialRegistry mentorRegistry,
5153
ObjectMapper objectMapper) {
5254
this.agentJobRepository = agentJobRepository;
55+
this.jwtVerifier = jwtVerifier;
5356
this.mentorRegistry = mentorRegistry;
5457
this.objectMapper = objectMapper;
5558
}
@@ -69,12 +72,25 @@ protected void doFilterInternal(HttpServletRequest request, HttpServletResponse
6972
return;
7073
}
7174

72-
if (!BASE64_URL_PATTERN.matcher(token).matches()) {
73-
response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Invalid token format");
74-
return;
75+
Optional<ProxyRouting> routing = mentorRegistry.validate(token);
76+
if (routing.isEmpty()) {
77+
JobJwt jwt;
78+
try {
79+
if (!(jwtVerifier.verify(token) instanceof JobJwt jobJwt)) {
80+
response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Invalid or expired token");
81+
return;
82+
}
83+
jwt = jobJwt;
84+
} catch (WorkerJwtInvalidException e) {
85+
response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Invalid or expired token");
86+
return;
87+
}
88+
if (!jwt.scopes().contains(LLM_PROXY_SCOPE)) {
89+
response.sendError(HttpServletResponse.SC_FORBIDDEN, "Insufficient token scope");
90+
return;
91+
}
92+
routing = resolveJobRouting(jwt);
7593
}
76-
77-
Optional<ProxyRouting> routing = resolveJobRouting(token).or(() -> mentorRegistry.validate(token));
7894
if (routing.isEmpty()) {
7995
response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Invalid or expired token");
8096
return;
@@ -88,20 +104,15 @@ protected void doFilterInternal(HttpServletRequest request, HttpServletResponse
88104
}
89105
}
90106

91-
/**
92-
* Look up an {@code AgentJob} by token and translate its frozen {@link ConfigSnapshot} into routing.
93-
* The attempt's identity and spend-so-far are read here because the row and snapshot already loaded
94-
* carry both, so the budget gate costs no extra query.
95-
*/
96-
private Optional<ProxyRouting> resolveJobRouting(String token) {
97-
String hash = AgentJob.computeTokenHash(token);
98-
Optional<AgentJob> optionalJob = agentJobRepository.findByJobTokenHashAndStatus(hash, AgentJobStatus.RUNNING);
107+
private Optional<ProxyRouting> resolveJobRouting(JobJwt jwt) {
108+
Optional<AgentJob> optionalJob = agentJobRepository.findByIdWithWorkspace(jwt.jobId());
99109
if (optionalJob.isEmpty()) {
100110
return Optional.empty();
101111
}
102112
AgentJob job = optionalJob.get();
103-
if (!MessageDigest.isEqual(token.getBytes(), job.getJobToken().getBytes())) {
104-
log.warn("Token hash matched but constant-time comparison failed — possible collision");
113+
if (job.getStatus() != AgentJobStatus.RUNNING
114+
|| !job.getWorkspace().getId().equals(jwt.workspaceId())
115+
|| job.getRetryCount() != jwt.attempt()) {
105116
return Optional.empty();
106117
}
107118
if (job.getConfigSnapshot() == null) {

‎server/application/src/main/java/de/tum/cit/aet/hephaestus/agent/proxy/LlmProxySecurityConfig.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import de.tum.cit.aet.hephaestus.agent.job.AgentJobRepository;
44
import de.tum.cit.aet.hephaestus.core.runtime.RuntimeRole;
5+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtVerifier;
56
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
67
import org.springframework.context.annotation.Bean;
78
import org.springframework.context.annotation.Configuration;
@@ -14,7 +15,8 @@
1415

1516
/**
1617
* Separate security filter chain for the internal LLM proxy endpoints, ordered ahead of the main JWT
17-
* chain so these requests authenticate with proxy-scoped bearer tokens instead of JWTs.
18+
* chain so these requests authenticate with proxy-scoped bearer credentials instead of the
19+
* application user-authentication chain.
1820
*
1921
* <p>Gated on the worker/sandbox capability rather than the practice-job feature flag, because
2022
* disabling practice reviews must not break mentor turns.
@@ -28,6 +30,7 @@ class LlmProxySecurityConfig {
2830
SecurityFilterChain llmProxyFilterChain(
2931
HttpSecurity http,
3032
AgentJobRepository agentJobRepository,
33+
WorkerJwtVerifier jwtVerifier,
3134
MentorProxyCredentialRegistry mentorRegistry,
3235
ObjectMapper objectMapper)
3336
throws Exception {
@@ -36,7 +39,7 @@ SecurityFilterChain llmProxyFilterChain(
3639
.sessionManagement(session -> session.sessionCreationPolicy(SessionCreationPolicy.STATELESS))
3740
.authorizeHttpRequests(auth -> auth.anyRequest().authenticated())
3841
.addFilterBefore(
39-
new JobTokenAuthenticationFilter(agentJobRepository, mentorRegistry, objectMapper),
42+
new JobTokenAuthenticationFilter(agentJobRepository, jwtVerifier, mentorRegistry, objectMapper),
4043
UsernamePasswordAuthenticationFilter.class);
4144

4245
return http.build();

‎server/application/src/main/java/de/tum/cit/aet/hephaestus/core/runtime/hub/HubConfiguration.java‎

Lines changed: 0 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,15 @@
11
package de.tum.cit.aet.hephaestus.core.runtime.hub;
22

33
import de.tum.cit.aet.hephaestus.core.runtime.RuntimeRole;
4-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.JavaJwtWorkerJwtVerifier;
54
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtHandshakeInterceptor;
6-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtIssuer;
75
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtVerifier;
8-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerKeyRing;
9-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerTokenDenylistRepository;
10-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerTokenDenylistService;
11-
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerTokenProperties;
126
import de.tum.cit.aet.hephaestus.core.runtime.worker.protocol.FrameCodec;
137
import io.micrometer.core.instrument.MeterRegistry;
14-
import org.slf4j.Logger;
15-
import org.slf4j.LoggerFactory;
168
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
179
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
18-
import org.springframework.boot.context.properties.EnableConfigurationProperties;
1910
import org.springframework.context.ApplicationEventPublisher;
2011
import org.springframework.context.annotation.Bean;
2112
import org.springframework.context.annotation.Configuration;
22-
import org.springframework.core.env.Environment;
23-
import org.springframework.core.env.Profiles;
2413
import org.springframework.web.socket.config.annotation.EnableWebSocket;
2514
import tools.jackson.databind.ObjectMapper;
2615

@@ -31,12 +20,9 @@
3120
*/
3221
@Configuration(proxyBeanMethods = false)
3322
@ConditionalOnProperty(name = RuntimeRole.SERVER_PROPERTY, havingValue = "true", matchIfMissing = true)
34-
@EnableConfigurationProperties({WorkerTokenProperties.class})
3523
@EnableWebSocket
3624
public class HubConfiguration {
3725

38-
private static final Logger log = LoggerFactory.getLogger(HubConfiguration.class);
39-
4026
@Bean
4127
@ConditionalOnMissingBean(FrameCodec.class)
4228
FrameCodec frameCodec(ObjectMapper objectMapper) {
@@ -53,45 +39,6 @@ WorkerJobCancelDispatcher workerJobCancelDispatcher(WorkerSessionRegistry regist
5339
return new WorkerJobCancelDispatcher(registry);
5440
}
5541

56-
@Bean
57-
WorkerKeyRing workerKeyRing(WorkerTokenProperties properties, Environment environment) {
58-
WorkerKeyRing ring = WorkerKeyRing.fromConfig(properties);
59-
if (ring.active().ephemeral()) {
60-
// Ephemeral keys are fine for dev but unsafe in prod (worker reconnects across pod
61-
// restarts must verify already-issued JWTs), so only the prod profile warns.
62-
if (environment.acceptsProfiles(Profiles.of("prod"))) {
63-
log.warn(
64-
"Worker JWT is using an EPHEMERAL signing key (kid={}). Configure a stable key via "
65-
+ "hephaestus.worker.hub.token.keys[*].private-key for production.",
66-
ring.active().kid());
67-
} else {
68-
log.info(
69-
"Worker JWT using ephemeral signing key (kid={}) — fine for dev.",
70-
ring.active().kid());
71-
}
72-
}
73-
return ring;
74-
}
75-
76-
@Bean
77-
WorkerTokenDenylistService workerTokenDenylistService(WorkerTokenDenylistRepository repository) {
78-
return new WorkerTokenDenylistService(repository);
79-
}
80-
81-
@Bean
82-
WorkerJwtVerifier workerJwtVerifier(
83-
WorkerKeyRing keyRing,
84-
WorkerTokenProperties properties,
85-
WorkerTokenDenylistService denylist,
86-
MeterRegistry meterRegistry) {
87-
return new JavaJwtWorkerJwtVerifier(keyRing, properties, denylist, meterRegistry);
88-
}
89-
90-
@Bean
91-
WorkerJwtIssuer workerJwtIssuer(WorkerKeyRing keyRing, WorkerTokenProperties properties) {
92-
return new WorkerJwtIssuer(keyRing, properties);
93-
}
94-
9542
@Bean
9643
WorkerJwtHandshakeInterceptor workerJwtHandshakeInterceptor(WorkerJwtVerifier verifier) {
9744
return new WorkerJwtHandshakeInterceptor(verifier);

‎server/application/src/main/java/de/tum/cit/aet/hephaestus/core/runtime/hub/WorkerControlWebSocketHandler.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwt;
44
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerJwtHandshakeInterceptor;
5+
import de.tum.cit.aet.hephaestus.core.runtime.hub.auth.WorkerSessionJwt;
56
import de.tum.cit.aet.hephaestus.core.runtime.worker.protocol.CancelJob;
67
import de.tum.cit.aet.hephaestus.core.runtime.worker.protocol.CapacityReport;
78
import de.tum.cit.aet.hephaestus.core.runtime.worker.protocol.ForceReconnect;
@@ -67,14 +68,17 @@ public void afterConnectionEstablished(WebSocketSession rawTransport) {
6768
WebSocketSession transport = new ConcurrentWebSocketSessionDecorator(
6869
rawTransport, (int) HubProperties.SEND_TIME_LIMIT.toMillis(), HubProperties.SEND_BUFFER_SIZE_BYTES);
6970
String sessionId = UUID.randomUUID().toString();
70-
WorkerSession session =
71-
new WorkerSession(jwt.workerId(), sessionId, jwt.jti(), jwt.expiresAt(), transport, codec);
71+
if (!(jwt instanceof WorkerSessionJwt workerJwt)) {
72+
throw new IllegalStateException("job-scoped JWT cannot open a worker control session");
73+
}
74+
String workerId = workerJwt.workerId();
75+
WorkerSession session = new WorkerSession(workerId, sessionId, jwt.jti(), jwt.expiresAt(), transport, codec);
7276
rawTransport.getAttributes().put(ATTR_WORKER_SESSION, session);
7377
// Close half-open sessions that authenticate but never send WorkerHello.
7478
ScheduledFuture<?> helloDeadline = helloTimeoutScheduler.schedule(
7579
() -> {
7680
if (!rawTransport.isOpen()) return;
77-
log.warn("WorkerHello timeout for workerId={} sessionId={}; closing.", jwt.workerId(), sessionId);
81+
log.warn("WorkerHello timeout for workerId={} sessionId={}; closing.", workerId, sessionId);
7882
meterRegistry.counter("worker.hub.hello.timeout").increment();
7983
close(rawTransport, CloseStatus.SESSION_NOT_RELIABLE);
8084
},
@@ -83,7 +87,7 @@ public void afterConnectionEstablished(WebSocketSession rawTransport) {
8387
session.armHelloDeadline(helloDeadline);
8488
log.info(
8589
"WSS connection opened: workerId={}, sessionId={}, jwtExpiresAt={}",
86-
jwt.workerId(),
90+
workerId,
8791
sessionId,
8892
jwt.expiresAt());
8993
}

0 commit comments

Comments
 (0)