diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/FlightDeckStreamsApp.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/FlightDeckStreamsApp.java index 717a300..744a351 100644 --- a/processor-apps/processing/src/main/java/io/flightdeck/streams/FlightDeckStreamsApp.java +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/FlightDeckStreamsApp.java @@ -6,6 +6,7 @@ import io.flightdeck.streams.processors.EnrichInputMessageProcessor; import io.flightdeck.streams.processors.ExtractToolUseItemsProcessor; import io.flightdeck.streams.processors.MemoirSessionEndProcessor; +import io.flightdeck.streams.processors.SessionConcurrencyGuardProcessor; import io.flightdeck.streams.processors.SessionEndProcessor; import io.flightdeck.streams.processors.TransformToolUseDoneProcessor; import io.flightdeck.streams.model.ThinkResponse; @@ -160,6 +161,10 @@ static Topology buildTopology(boolean memoirEnabled, boolean userSettingEnabled) ); // ── Register each processor fragment ────────────────────────────────── + // Admission control runs first: it consumes message-input and produces + // admitted-message-input, dropping new user messages whose session has a + // turn still in progress (so the LLM never sees an unanswered tool_use). + SessionConcurrencyGuardProcessor.register(builder, thinkStream, replyToTable); EnrichInputMessageProcessor.register(builder, memoirTable, thinkTable, userSettingsTable); ExtractToolUseItemsProcessor.register(builder, thinkStream); EndTurnProcessor.register(builder, thinkStream, replyToTable); @@ -194,6 +199,7 @@ private static void ensureTopicsExist(Properties streamsProps) { List requiredTopics = new java.util.ArrayList<>(List.of( Topics.MESSAGE_INPUT, + Topics.ADMITTED_MESSAGE_INPUT, Topics.ENRICHED_MESSAGE_INPUT, Topics.THINK_REQUEST_RESPONSE, Topics.TOOL_USE, diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/config/Topics.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/config/Topics.java index bfd914c..bc8df24 100644 --- a/processor-apps/processing/src/main/java/io/flightdeck/streams/config/Topics.java +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/config/Topics.java @@ -28,6 +28,15 @@ private static String requireEnv(String key) { /** Per-session accumulated conversation history (KTable backing store) */ public static final String SESSION_CONTEXT = PREFIX + "session-context"; + /** + * Admission-controlled user/tool messages: {@code message-input} records that + * have passed the {@code SessionConcurrencyGuardProcessor}. A new user message + * for a session whose previous turn is still in progress is rejected here + * (never reaching enrichment), so the LLM never receives a {@code tool_use} + * block that is not immediately followed by its {@code tool_result}. + */ + public static final String ADMITTED_MESSAGE_INPUT = PREFIX + "admitted-message-input"; + /** Merged: historical context + latest user message, ready for the LLM */ public static final String ENRICHED_MESSAGE_INPUT = PREFIX + "enriched-message-input"; diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/model/SessionStatus.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/model/SessionStatus.java new file mode 100644 index 0000000..61f7a8c --- /dev/null +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/model/SessionStatus.java @@ -0,0 +1,32 @@ +package io.flightdeck.streams.model; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * Per-session concurrency state owned by the + * {@code SessionConcurrencyGuardProcessor}'s {@code session-status-store}. + * + *

A session is {@code busy} from the moment a user message is admitted (a turn + * starts) until a clean end-turn response is observed + * ({@code endTurn == true} with no outstanding tool calls). While busy, a new + * incoming user message is rejected rather than appended to the conversation — + * appending it mid-turn would place a {@code tool_use} block immediately before a + * {@code user} message instead of its {@code tool_result}, which the LLM API + * rejects with a 400 and which kills the turn. + * + *

+ */ +@JsonIgnoreProperties(ignoreUnknown = true) +public record SessionStatus( + @JsonProperty("busy") boolean busy, + @JsonProperty("turn_started_ms") long turnStartedMs, + @JsonProperty("last_activity_ms") long lastActivityMs, + @JsonProperty("timestamp") String timestamp +) {} diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/model/UserResponse.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/model/UserResponse.java index 43f2bdb..bddc426 100644 --- a/processor-apps/processing/src/main/java/io/flightdeck/streams/model/UserResponse.java +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/model/UserResponse.java @@ -32,5 +32,6 @@ public record UserResponse( @JsonProperty("cost") Double cost, @JsonProperty("source_agent") String sourceAgent, // which agent produced this @JsonProperty("reply_to") Map replyTo, // multi-agent reply route (nullable) + @JsonProperty("status") String status, // null for normal LLM output; "rejected" when the agent was busy @JsonProperty("timestamp") String timestamp ) {} diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EndTurnProcessor.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EndTurnProcessor.java index a882133..ca41f80 100644 --- a/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EndTurnProcessor.java +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EndTurnProcessor.java @@ -158,6 +158,7 @@ static UserResponse toUserResponse(String sessionId, ThinkResponse response, Str response.totalSessionCost(), sourceAgent, parseReply(replyJson), + null, // status — normal end-turn output (not a busy rejection) Instant.now().toString() ); } diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EnrichInputMessageProcessor.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EnrichInputMessageProcessor.java index 9fd9be8..f070c19 100644 --- a/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EnrichInputMessageProcessor.java +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/EnrichInputMessageProcessor.java @@ -53,9 +53,9 @@ public static void register(StreamsBuilder builder, KTable thinkTable, KTable userSettingsTable) { - // ── Left side: incoming user messages ──────────────────────────────── + // ── Left side: incoming user messages, post-admission-control ──────── KStream inputStream = builder.stream( - Topics.MESSAGE_INPUT, + Topics.ADMITTED_MESSAGE_INPUT, Consumed.with(Serdes.String(), JsonSerde.of(MessageInput.class)) ); diff --git a/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessor.java b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessor.java new file mode 100644 index 0000000..6f3aa15 --- /dev/null +++ b/processor-apps/processing/src/main/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessor.java @@ -0,0 +1,344 @@ +package io.flightdeck.streams.processors; + +import io.flightdeck.streams.config.Topics; +import io.flightdeck.streams.model.MessageInput; +import io.flightdeck.streams.model.SessionStatus; +import io.flightdeck.streams.model.ThinkResponse; +import io.flightdeck.streams.model.UserResponse; +import io.flightdeck.streams.serdes.JsonSerde; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.StreamsBuilder; +import org.apache.kafka.streams.kstream.Consumed; +import org.apache.kafka.streams.kstream.Joined; +import org.apache.kafka.streams.kstream.KStream; +import org.apache.kafka.streams.kstream.KTable; +import org.apache.kafka.streams.kstream.Produced; +import org.apache.kafka.streams.processor.PunctuationType; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.StoreBuilder; +import org.apache.kafka.streams.state.Stores; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; + +/** + *

Session Concurrency Guard Processor

+ * + *

Admission control placed before {@link EnrichInputMessageProcessor}. + * It rejects a new user message for a session whose previous turn is still in + * progress, and lets the system's own tool-result continuation through. + * + *

Why

+ * History for the next LLM call is rebuilt as + * {@code previousMessages + [lastInputMessage] + lastInputResponse} + * (see {@link EnrichInputMessageProcessor}). When the previous turn issued tools, + * {@code lastInputResponse} ends with an assistant message carrying + * {@code tool_use} blocks. The matching {@code tool_result} blocks arrive as a + * {@code message-input} record with {@code role:"tool"}. If a {@code role:"user"} + * message is appended instead, the LLM sees a {@code tool_use} that is + * not immediately followed by its {@code tool_result} and returns a 400 — killing + * the turn. This guard prevents that by dropping the premature user message. + * + *

Topology

+ * 
+ *   message-input ─(re-key by session_id)─┐
+ *                                          ├─►[guard]─┬─► admitted-message-input ─►[enrich]─►…
+ *   think-request-response (KStream) ──────┘   │      │
+ *                      owns: session-status-store    └─(⋈ reply-to)─► message-output  (rejection)
+ *
+ * + *

Admitted messages are re-published to {@code admitted-message-input}, which + * {@link EnrichInputMessageProcessor} consumes in place of {@code message-input}. + * The intermediate topic keeps the guard a self-contained fragment and gives the + * admitted stream a natural session_id partitioning for downstream joins. + * Dedupe/admission state is solely the session-keyed {@code session-status-store}. + * + *

Rejections are left-joined with the per-session reply-to route before + * publication, mirroring {@link EndTurnProcessor}: a multi-agent caller whose + * message is rejected fails fast through its HTTP callback (and the one-shot + * route is tombstoned on delivery by the OutputConsumer) instead of hanging + * until its own timeout — which would also leave a stale route that the + * in-flight turn's reply could be mis-delivered to. + * + *

Busy lifecycle (per session)

+ *
    + *
  • A {@code role:"user"} message when not busy → admitted; session + * marked busy. When busy → rejected (a {@code status:"rejected"} + * {@link UserResponse} is emitted to {@code message-output}).
    • + *
  • A {@code role:"tool"} continuation → always admitted; activity refreshed. + * Never rejected.
    • + *
  • A {@code think-request-response} that is a clean end-turn + * ({@code endTurn} with no tool calls) → busy cleared (turn complete). + * Otherwise busy is kept and activity refreshed (tools still pending).
    • + *
  • A wall-clock punctuator force-clears any session idle longer than + * {@link #SESSION_BUSY_TTL_MS} — the safety net that unwedges a session + * whose turn crashed or whose end-turn signal was lost, so a session can + * never be locked permanently.
    • + *
+ */ +public class SessionConcurrencyGuardProcessor { + + private static final Logger log = LoggerFactory.getLogger(SessionConcurrencyGuardProcessor.class); + + /** RocksDB store: session_id → {@link SessionStatus}. */ + public static final String STATUS_STORE = "session-status-store"; + + /** Role that marks the system's own tool-result continuation — never rejected. */ + static final String ROLE_TOOL = "tool"; + + /** Content delivered to the caller when a message is rejected for being mid-turn. */ + static final String REJECTION_CONTENT = + "The agent is still processing your previous request. Please wait for it to finish before sending another message."; + + /** Status flag stamped onto a rejection {@link UserResponse}. */ + static final String STATUS_REJECTED = "rejected"; + + /** + * Idle timeout after which a busy session is force-cleared. Must exceed the + * longest legitimate quiet stretch within a turn — the async tool timeout + * ({@code ASYNC_TOOL_TIMEOUT_MS}, default 300s) — plus margin. + */ + static final long SESSION_BUSY_TTL_MS = + envLong("SESSION_BUSY_TTL_MS", 600_000L); + + /** How often the wall-clock punctuator scans for idle/wedged sessions. */ + static final long PUNCTUATE_INTERVAL_MS = + envLong("SESSION_GUARD_PUNCTUATE_INTERVAL_MS", 15_000L); + + /** + * Registers the guard: consumes {@code message-input} and produces the + * admitted subset (keyed by session_id) to {@code admitted-message-input}, + * plus a rejection to {@code message-output} for anything dropped. + */ + public static void register(StreamsBuilder builder, + KStream thinkStream, + KTable replyToTable) { + + StoreBuilder> storeBuilder = + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(STATUS_STORE), + Serdes.String(), + JsonSerde.of(SessionStatus.class) + ); + builder.addStateStore(storeBuilder); + + // Incoming messages, re-keyed by session_id (mirrors EnrichInputMessageProcessor). + KStream messageEvents = builder + .stream(Topics.MESSAGE_INPUT, + Consumed.with(Serdes.String(), JsonSerde.of(MessageInput.class))) + .selectKey((key, msg) -> msg.sessionId()) + .mapValues(GuardEvent::message); + + // Think responses drive the busy lifecycle (clear on clean end-turn). + KStream thinkEvents = thinkStream + .filter((sessionId, response) -> response != null) + .mapValues(GuardEvent::think); + + KStream processed = messageEvents + .merge(thinkEvents) + .process( + (ProcessorSupplier) GuardProcessor::new, + STATUS_STORE + ); + + // Admitted messages continue to enrichment via the intermediate topic. + processed + .filter((sessionId, result) -> result != null && result.admitted() != null) + .mapValues(GuardResult::admitted) + .to(Topics.ADMITTED_MESSAGE_INPUT, + Produced.with(Serdes.String(), JsonSerde.of(MessageInput.class))); + + // Rejections are delivered back through the normal output channel, with + // the session's reply-to route joined in (as EndTurnProcessor does) so a + // multi-agent caller gets the rejection via its HTTP callback. + processed + .filter((sessionId, result) -> result != null && result.rejection() != null) + .mapValues(GuardResult::rejection) + .leftJoin(replyToTable, SessionConcurrencyGuardProcessor::withReplyRoute, + Joined.with(Serdes.String(), JsonSerde.of(UserResponse.class), Serdes.String())) + .to(Topics.MESSAGE_OUTPUT, + Produced.with(Serdes.String(), JsonSerde.of(UserResponse.class))); + } + + /** + * Embeds the session's reply-to route (raw JSON from the reply-to table) + * into a rejection. Null/absent/invalid route → the rejection passes through + * unchanged and is delivered over the default (WebSocket) channel. + */ + static UserResponse withReplyRoute(UserResponse rejection, String replyJson) { + java.util.Map reply = EndTurnProcessor.parseReply(replyJson); + if (reply == null) return rejection; + return new UserResponse( + rejection.sessionId(), + rejection.userId(), + rejection.content(), + rejection.inputTokens(), + rejection.outputTokens(), + rejection.cost(), + rejection.sourceAgent(), + reply, + rejection.status(), + rejection.timestamp()); + } + + // ───────────────────────────────────────────────────────────────────────── + // Internal merged-event wrapper (never hits a topic — no serde needed) + // ───────────────────────────────────────────────────────────────────────── + + enum Kind { MESSAGE, THINK } + + record GuardEvent(Kind kind, MessageInput message, ThinkResponse think) { + static GuardEvent message(MessageInput m) { return new GuardEvent(Kind.MESSAGE, m, null); } + static GuardEvent think(ThinkResponse t) { return new GuardEvent(Kind.THINK, null, t); } + } + + /** Output wrapper: exactly one of {@code admitted} / {@code rejection} is set. */ + record GuardResult(MessageInput admitted, UserResponse rejection) { + static GuardResult admit(MessageInput m) { return new GuardResult(m, null); } + static GuardResult reject(UserResponse u) { return new GuardResult(null, u); } + } + + // ───────────────────────────────────────────────────────────────────────── + // Processor + // ───────────────────────────────────────────────────────────────────────── + + static class GuardProcessor implements Processor { + + private ProcessorContext context; + private KeyValueStore store; + + @Override + public void init(ProcessorContext ctx) { + this.context = ctx; + this.store = ctx.getStateStore(STATUS_STORE); + ctx.schedule(Duration.ofMillis(PUNCTUATE_INTERVAL_MS), + PunctuationType.WALL_CLOCK_TIME, this::sweep); + } + + @Override + public void process(Record record) { + String key = record.key(); + GuardEvent event = record.value(); + if (key == null || event == null) { + log.warn("Null key or event — skipping"); + return; + } + if (event.kind() == Kind.MESSAGE) { + handleMessage(key, event.message(), record.timestamp()); + } else { + handleThink(key, event.think()); + } + } + + // ── Incoming message: admit or reject ────────────────────────────────── + private void handleMessage(String sessionId, MessageInput msg, long ts) { + if (msg == null) { + log.warn("[{}] Null message — skipping", sessionId); + return; + } + long now = System.currentTimeMillis(); + SessionStatus current = store.get(sessionId); + + // The system's own tool-result continuation is always admitted — it is + // precisely the record that completes the pending tool_use blocks. + if (ROLE_TOOL.equals(msg.role())) { + if (current != null && current.busy()) { + store.put(sessionId, new SessionStatus( + true, current.turnStartedMs(), now, Instant.now().toString())); + } + context.forward(new Record<>(sessionId, GuardResult.admit(msg), ts)); + return; + } + + // A new user (or scheduler-triggered) message. + if (current != null && current.busy()) { + log.warn("[{}] Rejecting message — turn in progress since {}ms ago (role={})", + sessionId, now - current.turnStartedMs(), msg.role()); + context.forward(new Record<>(sessionId, GuardResult.reject(rejection(msg)), ts)); + return; + } + + // Admit and open a new turn. + store.put(sessionId, new SessionStatus(true, now, now, Instant.now().toString())); + log.info("[{}] Admitted message — turn started (role={})", sessionId, msg.role()); + context.forward(new Record<>(sessionId, GuardResult.admit(msg), ts)); + } + + // ── Think response: clear busy on clean end-turn, else keep busy ─────── + private void handleThink(String sessionId, ThinkResponse response) { + if (response == null) return; + long now = System.currentTimeMillis(); + + boolean noToolCalls = response.toolUses() == null || response.toolUses().isEmpty(); + boolean cleanEndTurn = response.endTurn() && noToolCalls; + + if (cleanEndTurn) { + if (store.get(sessionId) != null) { + store.delete(sessionId); + log.info("[{}] Turn complete — session marked idle", sessionId); + } + return; + } + + // Tools still pending (or not yet end-of-turn): keep the session busy + // and refresh activity so a long, legitimately-active turn is not swept. + SessionStatus current = store.get(sessionId); + long turnStarted = (current != null) ? current.turnStartedMs() : now; + store.put(sessionId, new SessionStatus(true, turnStarted, now, Instant.now().toString())); + } + + // ── Wall-clock sweep: force-clear wedged / idle sessions ─────────────── + private void sweep(long now) { + List toClear = new ArrayList<>(); + try (KeyValueIterator it = store.all()) { + while (it.hasNext()) { + KeyValue kv = it.next(); + SessionStatus s = kv.value; + if (s == null) continue; + if (now - s.lastActivityMs() > SESSION_BUSY_TTL_MS) { + toClear.add(kv.key); + } + } + } + for (String key : toClear) { + store.delete(key); + log.warn("[{}] Session idle > {}ms — force-clearing busy state (turn likely crashed)", + key, SESSION_BUSY_TTL_MS); + } + } + + private static UserResponse rejection(MessageInput msg) { + return new UserResponse( + msg.sessionId(), + msg.userId(), + REJECTION_CONTENT, + 0, 0, 0.0, + null, // source_agent — system-generated, not from an agent + null, // reply_to — joined in downstream (withReplyRoute) + STATUS_REJECTED, + Instant.now().toString()); + } + } + + private static long envLong(String key, long defaultValue) { + String v = System.getenv(key); + if (v == null || v.isBlank()) return defaultValue; + try { + return Long.parseLong(v.trim()); + } catch (NumberFormatException e) { + log.warn("Invalid {}={} — using default {}", key, v, defaultValue); + return defaultValue; + } + } +} diff --git a/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/EnrichInputMessageProcessorTest.java b/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/EnrichInputMessageProcessorTest.java index d35c4ea..dc7e4c1 100644 --- a/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/EnrichInputMessageProcessorTest.java +++ b/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/EnrichInputMessageProcessorTest.java @@ -70,8 +70,11 @@ void setUp() { driver = new TopologyTestDriver(builder.build(), props); + // Enrichment consumes the post-admission-control stream, so the test + // feeds admitted-message-input directly (bypassing the guard, which is + // exercised separately in SessionConcurrencyGuardProcessorTest). messageInput = driver.createInputTopic( - Topics.MESSAGE_INPUT, + Topics.ADMITTED_MESSAGE_INPUT, Serdes.String().serializer(), JsonSerde.of(MessageInput.class).serializer()); diff --git a/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessorTest.java b/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessorTest.java new file mode 100644 index 0000000..f2d58cf --- /dev/null +++ b/processor-apps/processing/src/test/java/io/flightdeck/streams/processors/SessionConcurrencyGuardProcessorTest.java @@ -0,0 +1,269 @@ +package io.flightdeck.streams.processors; + +import io.flightdeck.streams.config.Topics; +import io.flightdeck.streams.model.MessageInput; +import io.flightdeck.streams.model.ThinkResponse; +import io.flightdeck.streams.model.ToolUseItem; +import io.flightdeck.streams.model.UserResponse; +import io.flightdeck.streams.serdes.JsonSerde; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.*; +import org.apache.kafka.streams.kstream.Consumed; +import org.apache.kafka.streams.kstream.KStream; +import org.apache.kafka.streams.kstream.KTable; +import org.junit.jupiter.api.*; + +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.Properties; + +import static org.assertj.core.api.Assertions.*; + +/** + * Tests for {@link SessionConcurrencyGuardProcessor}. + * + * Invariants under test: + * - First user message is admitted and opens a turn. + * - A second user message while the turn is in progress is rejected (and a + * status:"rejected" UserResponse is emitted to message-output). + * - The system's own role:"tool" continuation is always admitted, even mid-turn. + * - A clean end-turn think response clears busy so the next user message is admitted. + * - A think response that still carries tool calls keeps the session busy. + * - A session whose turn never completes is force-cleared after the idle TTL. + * - Sessions are isolated from one another. + */ +class SessionConcurrencyGuardProcessorTest { + + private TopologyTestDriver driver; + private TestInputTopic messageInput; + private TestInputTopic thinkInput; + private TestInputTopic replyToInput; + private TestOutputTopic admittedOutput; + private TestOutputTopic rejectionOutput; + + @BeforeEach + void setUp() { + StreamsBuilder builder = new StreamsBuilder(); + + KStream thinkStream = builder.stream( + Topics.THINK_REQUEST_RESPONSE, + Consumed.with(Serdes.String(), JsonSerde.of(ThinkResponse.class))); + + KTable replyToTable = builder.table( + Topics.REPLY_TO, + Consumed.with(Serdes.String(), Serdes.String())); + + SessionConcurrencyGuardProcessor.register(builder, thinkStream, replyToTable); + + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, "test-session-guard"); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "dummy:9092"); + + // Anchor the mock wall clock so deadlines (System.currentTimeMillis()) + // line up with advanceWallClockTime(). + driver = new TopologyTestDriver(builder.build(), props, Instant.now()); + + messageInput = driver.createInputTopic( + Topics.MESSAGE_INPUT, + Serdes.String().serializer(), + JsonSerde.of(MessageInput.class).serializer()); + + thinkInput = driver.createInputTopic( + Topics.THINK_REQUEST_RESPONSE, + Serdes.String().serializer(), + JsonSerde.of(ThinkResponse.class).serializer()); + + replyToInput = driver.createInputTopic( + Topics.REPLY_TO, + Serdes.String().serializer(), + Serdes.String().serializer()); + + admittedOutput = driver.createOutputTopic( + Topics.ADMITTED_MESSAGE_INPUT, + Serdes.String().deserializer(), + JsonSerde.of(MessageInput.class).deserializer()); + + rejectionOutput = driver.createOutputTopic( + Topics.MESSAGE_OUTPUT, + Serdes.String().deserializer(), + JsonSerde.of(UserResponse.class).deserializer()); + } + + @AfterEach + void tearDown() { driver.close(); } + + // ── Admit / reject ───────────────────────────────────────────────────────── + + @Test + @DisplayName("First user message is admitted; no rejection emitted") + void firstMessage_admitted() { + messageInput.pipeInput("s1", userMsg("s1", "u", "hello")); + + assertThat(admittedOutput.isEmpty()).isFalse(); + assertThat(admittedOutput.readRecord().value().content()).isEqualTo("hello"); + assertThat(rejectionOutput.isEmpty()).isTrue(); + } + + @Test + @DisplayName("Second user message while a turn is in progress is rejected") + void secondMessage_whileBusy_rejected() { + messageInput.pipeInput("s2", userMsg("s2", "u", "first")); + admittedOutput.readRecord(); // first admitted + + messageInput.pipeInput("s2", userMsg("s2", "u", "second")); + + // The second message must NOT be admitted... + assertThat(admittedOutput.isEmpty()).isTrue(); + // ...and a rejection notice must be emitted instead. + assertThat(rejectionOutput.isEmpty()).isFalse(); + UserResponse rejection = rejectionOutput.readRecord().value(); + assertThat(rejection.sessionId()).isEqualTo("s2"); + assertThat(rejection.status()).isEqualTo("rejected"); + assertThat(rejection.cost()).isEqualTo(0.0); + assertThat(rejection.content()).isNotBlank(); + // No reply-to route stored → delivered over the default (WebSocket) channel. + assertThat(rejection.replyTo()).isNull(); + } + + @Test + @DisplayName("Rejection for a multi-agent caller carries the session's reply-to route") + void rejection_carriesReplyToRoute() { + replyToInput.pipeInput("s2r", "{\"callbackService\":\"caller-agent\",\"token\":\"tok-1\"}"); + + messageInput.pipeInput("s2r", userMsg("s2r", "u", "first")); + admittedOutput.readRecord(); + + messageInput.pipeInput("s2r", userMsg("s2r", "u", "second")); + + UserResponse rejection = rejectionOutput.readRecord().value(); + assertThat(rejection.status()).isEqualTo("rejected"); + assertThat(rejection.replyTo()) + .containsEntry("callbackService", "caller-agent") + .containsEntry("token", "tok-1"); + } + + // ── Tool continuation is never rejected ────────────────────────────────────── + + @Test + @DisplayName("role:tool continuation is admitted even while the turn is in progress") + void toolContinuation_admittedWhileBusy() { + messageInput.pipeInput("s3", userMsg("s3", "u", "do something")); + admittedOutput.readRecord(); // user message admitted, session now busy + + // The pipeline's own tool-result continuation re-enters message-input. + messageInput.pipeInput("s3", toolMsg("s3", "u")); + + assertThat(admittedOutput.isEmpty()).isFalse(); + assertThat(admittedOutput.readRecord().value().role()).isEqualTo("tool"); + // A continuation is not a rejection. + assertThat(rejectionOutput.isEmpty()).isTrue(); + } + + // ── Busy lifecycle via think responses ─────────────────────────────────────── + + @Test + @DisplayName("Clean end-turn clears busy so the next user message is admitted") + void cleanEndTurn_clearsBusy() { + messageInput.pipeInput("s4", userMsg("s4", "u", "first")); + admittedOutput.readRecord(); + + // LLM finishes the turn with no outstanding tool calls. + thinkInput.pipeInput("s4", endTurn("s4", "u")); + + // A follow-up user message is now admitted, not rejected. + messageInput.pipeInput("s4", userMsg("s4", "u", "second")); + assertThat(rejectionOutput.isEmpty()).isTrue(); + assertThat(admittedOutput.readRecord().value().content()).isEqualTo("second"); + } + + @Test + @DisplayName("Think response that still carries tool calls keeps the session busy") + void toolCallsPending_keepsBusy() { + messageInput.pipeInput("s5", userMsg("s5", "u", "first")); + admittedOutput.readRecord(); + + // LLM responded with a tool_use — turn is NOT complete. + thinkInput.pipeInput("s5", withTools("s5", "u", + new ToolUseItem("t1", "lookup", "lookup", Map.of(), "s5", 0, TS))); + + // A new user message must still be rejected. + messageInput.pipeInput("s5", userMsg("s5", "u", "second")); + assertThat(admittedOutput.isEmpty()).isTrue(); + assertThat(rejectionOutput.readRecord().value().status()).isEqualTo("rejected"); + } + + @Test + @DisplayName("A think response that opens a turn before any message keeps the session busy") + void thinkWithoutPriorAdmit_marksBusy() { + // Defensive: a think response (e.g. after a restart) with tools pending, + // with no prior admit recorded, still locks the session. + thinkInput.pipeInput("s5b", withTools("s5b", "u", + new ToolUseItem("t1", "lookup", "lookup", Map.of(), "s5b", 0, TS))); + + messageInput.pipeInput("s5b", userMsg("s5b", "u", "barge in")); + assertThat(admittedOutput.isEmpty()).isTrue(); + assertThat(rejectionOutput.readRecord().value().status()).isEqualTo("rejected"); + } + + // ── TTL safety net ─────────────────────────────────────────────────────────── + + @Test + @DisplayName("A wedged turn is force-cleared after the idle TTL, unblocking the session") + void wedgedTurn_forceClearedAfterTtl() { + messageInput.pipeInput("s6", userMsg("s6", "u", "first")); + admittedOutput.readRecord(); + + // No end-turn ever arrives (think consumer crashed). Advance past the TTL. + driver.advanceWallClockTime(Duration.ofMinutes(11)); + + // The session is no longer busy: a new message is admitted. + messageInput.pipeInput("s6", userMsg("s6", "u", "second")); + assertThat(rejectionOutput.isEmpty()).isTrue(); + assertThat(admittedOutput.readRecord().value().content()).isEqualTo("second"); + } + + // ── Session isolation ────────────────────────────────────────────────────── + + @Test + @DisplayName("A busy session does not block a different session") + void sessionIsolation() { + messageInput.pipeInput("A", userMsg("A", "u1", "a-first")); + admittedOutput.readRecord(); + + // Different session — must be admitted despite A being busy. + messageInput.pipeInput("B", userMsg("B", "u2", "b-first")); + assertThat(rejectionOutput.isEmpty()).isTrue(); + assertThat(admittedOutput.readRecord().value().sessionId()).isEqualTo("B"); + } + + @Test + @DisplayName("Admitted record is keyed by session_id") + void admittedKey_isSessionId() { + messageInput.pipeInput("orig-key", userMsg("key-sess", "u", "hi")); + assertThat(admittedOutput.readRecord().key()).isEqualTo("key-sess"); + } + + // ── Helpers ─────────────────────────────────────────────────────────────── + + private static final String TS = "2026-03-10T12:00:00Z"; + + private static MessageInput userMsg(String sessionId, String userId, String content) { + return new MessageInput(sessionId, userId, "user", content, TS, Map.of()); + } + + private static MessageInput toolMsg(String sessionId, String userId) { + return new MessageInput(sessionId, userId, "tool", "tool result", TS, Map.of()); + } + + private static ThinkResponse endTurn(String sessionId, String userId) { + return new ThinkResponse(sessionId, userId, 0.0, 0.0, 0.0, 0, 0, + List.of(), null, List.of(), List.of(), true, false, 0, 0, 0.0, TS); + } + + private static ThinkResponse withTools(String sessionId, String userId, ToolUseItem... items) { + return new ThinkResponse(sessionId, userId, 0.0, 0.0, 0.0, 0, 0, + List.of(), null, List.of(), List.of(items), false, false, 0, 0, 0.0, TS); + } +}