From 00861af42c06eafe23c9b5b44f5b63a07adc6ea8 Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:39 -0500 Subject: [PATCH 1/5] feat(turns): lease, heartbeat, cancel request and phase columns on assistant_turn heartbeat_at, lease_owner, cancel_requested_at, phase and concurrency_mode, a (user_id, status) index, and a partial unique index so a turn gets one message.end even when two replicas reap it. Store helpers for a conditional update of a still-running turn. Co-Authored-By: Claude Opus 5.5 --- docs/schema-postgres.sql | 15 ++ .../qbits/assistant/model/AssistantTurn.java | 171 ++++++++++++++++++ .../assistant/persistence/AssistantStore.java | 67 +++++++ 3 files changed, 253 insertions(+) diff --git a/docs/schema-postgres.sql b/docs/schema-postgres.sql index e9e618f..46a3ab7 100644 --- a/docs/schema-postgres.sql +++ b/docs/schema-postgres.sql @@ -577,3 +577,18 @@ CREATE INDEX assistant_scheduled_run_owner ON assistant_scheduled_run(user_id, s -- X1/Y1: encoded audit filters can exceed the old 500-character record links. ALTER TABLE assistant_tool_call ALTER COLUMN audit_href TYPE TEXT; + + +-- Durable turns: apply this additive block before deploying the new runtime. +-- Like earlier additions, fresh installs run the CREATEs followed by ALTERs. +ALTER TABLE assistant_turn ADD COLUMN heartbeat_at TIMESTAMP; +ALTER TABLE assistant_turn ADD COLUMN lease_owner VARCHAR(100); +ALTER TABLE assistant_turn ADD COLUMN cancel_requested_at TIMESTAMP; +ALTER TABLE assistant_turn ADD COLUMN phase VARCHAR(64); +ALTER TABLE assistant_turn ADD COLUMN concurrency_mode VARCHAR(16); +CREATE INDEX assistant_turn_user_status ON assistant_turn (user_id, status); + +-- Fence concurrent orphan reapers at the database boundary. A turn ends once. +-- Check for pre-existing duplicate message.end rows before applying this index. +CREATE UNIQUE INDEX assistant_turn_event_terminal ON assistant_turn_event (turn_id) + WHERE event_type = 'message.end'; diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantTurn.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantTurn.java index 716127d..9eec8f2 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantTurn.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantTurn.java @@ -92,6 +92,21 @@ public class AssistantTurn extends QRecordEntity @QField(isEditable = false) private Instant endedAt; + @QField(isEditable = false) + private Instant heartbeatAt; + + @QField(isEditable = false, maxLength = 100) + private String leaseOwner; + + @QField(isEditable = false) + private Instant cancelRequestedAt; + + @QField(isEditable = false, maxLength = 64) + private String phase; + + @QField(isEditable = false, maxLength = 16) + private String concurrencyMode; + @QField(isEditable = false) private Integer toolCallCount; @@ -1011,4 +1026,160 @@ public class AssistantTurn extends QRecordEntity this.providerId = providerId; return (this); } + + + + /******************************************************************************* + ** Getter for heartbeatAt + *******************************************************************************/ + public Instant getHeartbeatAt() + { + return (this.heartbeatAt); + } + + + + /******************************************************************************* + ** Setter for heartbeatAt + *******************************************************************************/ + public void setHeartbeatAt(Instant heartbeatAt) + { + this.heartbeatAt = heartbeatAt; + } + + + + /******************************************************************************* + ** Fluent setter for heartbeatAt + *******************************************************************************/ + public AssistantTurn withHeartbeatAt(Instant heartbeatAt) + { + this.heartbeatAt = heartbeatAt; + return (this); + } + + + + /******************************************************************************* + ** Getter for leaseOwner + *******************************************************************************/ + public String getLeaseOwner() + { + return (this.leaseOwner); + } + + + + /******************************************************************************* + ** Setter for leaseOwner + *******************************************************************************/ + public void setLeaseOwner(String leaseOwner) + { + this.leaseOwner = leaseOwner; + } + + + + /******************************************************************************* + ** Fluent setter for leaseOwner + *******************************************************************************/ + public AssistantTurn withLeaseOwner(String leaseOwner) + { + this.leaseOwner = leaseOwner; + return (this); + } + + + + /******************************************************************************* + ** Getter for cancelRequestedAt + *******************************************************************************/ + public Instant getCancelRequestedAt() + { + return (this.cancelRequestedAt); + } + + + + /******************************************************************************* + ** Setter for cancelRequestedAt + *******************************************************************************/ + public void setCancelRequestedAt(Instant cancelRequestedAt) + { + this.cancelRequestedAt = cancelRequestedAt; + } + + + + /******************************************************************************* + ** Fluent setter for cancelRequestedAt + *******************************************************************************/ + public AssistantTurn withCancelRequestedAt(Instant cancelRequestedAt) + { + this.cancelRequestedAt = cancelRequestedAt; + return (this); + } + + + + /******************************************************************************* + ** Getter for phase + *******************************************************************************/ + public String getPhase() + { + return (this.phase); + } + + + + /******************************************************************************* + ** Setter for phase + *******************************************************************************/ + public void setPhase(String phase) + { + this.phase = phase; + } + + + + /******************************************************************************* + ** Fluent setter for phase + *******************************************************************************/ + public AssistantTurn withPhase(String phase) + { + this.phase = phase; + return (this); + } + + + + /******************************************************************************* + ** Getter for concurrencyMode + *******************************************************************************/ + public String getConcurrencyMode() + { + return (this.concurrencyMode); + } + + + + /******************************************************************************* + ** Setter for concurrencyMode + *******************************************************************************/ + public void setConcurrencyMode(String concurrencyMode) + { + this.concurrencyMode = concurrencyMode; + } + + + + /******************************************************************************* + ** Fluent setter for concurrencyMode + *******************************************************************************/ + public AssistantTurn withConcurrencyMode(String concurrencyMode) + { + this.concurrencyMode = concurrencyMode; + return (this); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/persistence/AssistantStore.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/persistence/AssistantStore.java index 712a31c..028c165 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/persistence/AssistantStore.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/persistence/AssistantStore.java @@ -31,6 +31,7 @@ import com.kingsrook.qbits.assistant.AssistantQBitConfig; import com.kingsrook.qbits.assistant.AssistantRuntime; import com.kingsrook.qbits.assistant.classification.Classifications; import com.kingsrook.qbits.assistant.engine.TurnClassification; +import com.kingsrook.qbits.assistant.engine.TurnEvent; import com.kingsrook.qbits.assistant.model.AssistantBudget; import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent; import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride; @@ -384,6 +385,27 @@ public class AssistantStore + /******************************************************************************* + ** Re-read before a lifecycle update. A completed row is never reopened by + ** a late local Stop or overwritten by a second local finalizer. Every read + ** and update remains under the caller's own session and owner filter. + *******************************************************************************/ + public boolean updateRunningTurn(String userId, QRecord changes) throws QException + { + synchronized(QContext.getQInstance()) + { + AssistantTurn current = getTurn(userId, changes.getValueInteger("id")); + if(current == null || !("RUNNING".equals(current.getStatus()) || "CANCELLING".equals(current.getStatus()))) + { + return (false); + } + updateFields(AssistantTurn.TABLE_NAME, changes); + return (true); + } + } + + + /******************************************************************************* ** One of the user's own pending actions, or null. *******************************************************************************/ @@ -411,6 +433,51 @@ public class AssistantStore ** Persist one streamed event. *******************************************************************************/ public void insertEvent(Integer turnId, String userId, int seq, String eventType, String dataJson) throws QException + { + if(TurnEvent.MESSAGE_END.equals(eventType) || TurnEvent.ERROR.equals(eventType)) + { + // Serialize memory-backend callers. The Postgres partial unique index + // arbitrates this same race between independent replicas. + synchronized(QContext.getQInstance()) + { + if(hasTerminalEvent(userId, turnId)) + { + return; + } + try + { + insertStreamEvent(turnId, userId, seq, eventType, dataJson); + } + catch(QException e) + { + if(!hasTerminalEvent(userId, turnId)) + { + throw (e); + } + } + } + return; + } + insertStreamEvent(turnId, userId, seq, eventType, dataJson); + } + + + + /******************************************************************************* + ** Whether a terminal replay frame has already been persisted by any owner. + *******************************************************************************/ + private boolean hasTerminalEvent(String userId, Integer turnId) throws QException + { + return (!query(AssistantTurnEvent.TABLE_NAME, owned(userId).withCriteria("turnId", QCriteriaOperator.EQUALS, turnId) + .withCriteria("eventType", QCriteriaOperator.EQUALS, TurnEvent.MESSAGE_END).withLimit(1), AssistantTurnEvent.class).isEmpty()); + } + + + + /******************************************************************************* + ** Persist one frame through the normal encrypted, user-owned store path. + *******************************************************************************/ + private void insertStreamEvent(Integer turnId, String userId, int seq, String eventType, String dataJson) throws QException { insert(new AssistantTurnEvent() .withTurnId(turnId) -- 2.54.0 From ffa7e25bf472933f44ea52587d03259b9f5f6ab5 Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:41 -0500 Subject: [PATCH 2/5] fix(turns): every turn is finalised: watchdog, lease, hard deadline, shutdown A per-turn watchdog (the user's own session) renews the lease, picks up a Stop requested on another replica, enforces maxTurnSeconds and force- finalises a wedged engine. One idempotent finaliser: accounting or event failures can no longer leave a turn RUNNING or registered. On shutdown the runtime stops taking turns, cancels, waits a grace and finalises the rest as interrupted. The in-memory concurrent counter is gone. Runs are published through the AssistantRunTracker stand-in SPI (qqq#1210). Co-Authored-By: Claude Opus 5.5 --- .../qbits/assistant/AssistantQBitConfig.java | 249 ++++++++++- .../qbits/assistant/AssistantRuntime.java | 151 ++++++- .../qbits/assistant/engine/RateLimiter.java | 14 +- .../qbits/assistant/engine/RunningTurn.java | 168 +++++++- .../qbits/assistant/engine/TurnEngine.java | 388 +++++++++++++----- .../qbits/assistant/engine/TurnEvent.java | 1 + .../assistant/jobs/AssistantRunTracker.java | 75 ++++ 7 files changed, 924 insertions(+), 122 deletions(-) create mode 100644 qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/jobs/AssistantRunTracker.java diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java index 21fa600..834fc2a 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java @@ -17,6 +17,7 @@ package com.kingsrook.qbits.assistant; +import javax.net.ssl.SSLContext; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; @@ -31,6 +32,7 @@ import java.util.regex.Pattern; import com.fasterxml.jackson.annotation.JsonIgnore; import com.kingsrook.qbits.assistant.classification.ClassificationSource; import com.kingsrook.qbits.assistant.classification.ConfiguredClassificationSource; +import com.kingsrook.qbits.assistant.jobs.AssistantRunTracker; import com.kingsrook.qbits.assistant.rag.VectorStore; import com.kingsrook.qbits.assistant.spi.AssistantModelProvider; import com.kingsrook.qbits.assistant.spi.AssistantProviderType; @@ -44,7 +46,6 @@ import com.kingsrook.qqq.backend.core.model.metadata.producers.MetaDataCustomize import com.kingsrook.qqq.backend.core.model.metadata.qbits.QBitConfig; import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData; import com.kingsrook.qqq.backend.core.utils.StringUtils; -import javax.net.ssl.SSLContext; /******************************************************************************* @@ -99,12 +100,20 @@ public class AssistantQBitConfig implements QBitConfig private Integer maxParallelToolCalls = 4; private Integer maxRowsToModel = 50; private Integer maxCellChars = 500; + private Integer turnHeartbeatSeconds = 5; + private Integer turnLeaseSeconds = 45; + private Integer turnForceFinalizeGraceSeconds = 20; + private Integer shutdownGraceSeconds = 8; + private String defaultConcurrencyMode = "SUPERSEDE"; + private Boolean registerShutdownHook = true; + private AssistantRunTracker runTracker = AssistantRunTracker.NOOP; + private Integer maxTurnSeconds = 300; private Integer maxUserMessageChars = 8000; private Integer historyMessageLimit = 40; private AssistantBudgets budgets = new AssistantBudgets(); private Integer maxTurnsPerUserPerMinute = 10; - private Integer maxConcurrentTurnsPerUser = 1; + private Integer maxConcurrentTurnsPerUser = 3; private Integer conversationRetentionDays = 90; private Integer toolResultRetentionDays = 30; private Integer approvalRetentionDays = 365; @@ -206,6 +215,23 @@ public class AssistantQBitConfig implements QBitConfig requirePositive(errors, "maxRowsToModel", maxRowsToModel); requirePositive(errors, "maxCellChars", maxCellChars); requirePositive(errors, "maxTurnSeconds", maxTurnSeconds); + requirePositive(errors, "turnHeartbeatSeconds", turnHeartbeatSeconds); + requirePositive(errors, "turnLeaseSeconds", turnLeaseSeconds); + requirePositive(errors, "turnForceFinalizeGraceSeconds", turnForceFinalizeGraceSeconds); + requirePositive(errors, "shutdownGraceSeconds", shutdownGraceSeconds); + if(turnHeartbeatSeconds != null && turnLeaseSeconds != null && turnLeaseSeconds < 2 * turnHeartbeatSeconds) + { + errors.add("turnLeaseSeconds must be at least twice turnHeartbeatSeconds."); + } + if(!"SUPERSEDE".equals(defaultConcurrencyMode) && !"CONCURRENT".equals(defaultConcurrencyMode)) + { + errors.add("defaultConcurrencyMode must be SUPERSEDE or CONCURRENT."); + } + if(turnHeartbeatSeconds != null && turnLeaseSeconds != null && turnHeartbeatSeconds >= turnLeaseSeconds) + { + errors.add("turnHeartbeatSeconds must be less than turnLeaseSeconds."); + } + requirePositive(errors, "maxUserMessageChars", maxUserMessageChars); requirePositive(errors, "historyMessageLimit", historyMessageLimit); requirePositive(errors, "maxTurnsPerUserPerMinute", maxTurnsPerUserPerMinute); @@ -2008,4 +2034,223 @@ public class AssistantQBitConfig implements QBitConfig this.scheduledTasksSchedulerName = scheduledTasksSchedulerName; return (this); } + + + + /******************************************************************************* + ** Getter for turnHeartbeatSeconds + *******************************************************************************/ + public Integer getTurnHeartbeatSeconds() + { + return (this.turnHeartbeatSeconds); + } + + + + /******************************************************************************* + ** Setter for turnHeartbeatSeconds + *******************************************************************************/ + public void setTurnHeartbeatSeconds(Integer turnHeartbeatSeconds) + { + this.turnHeartbeatSeconds = turnHeartbeatSeconds; + } + + + + /******************************************************************************* + ** Fluent setter for turnHeartbeatSeconds + *******************************************************************************/ + public AssistantQBitConfig withTurnHeartbeatSeconds(Integer turnHeartbeatSeconds) + { + this.turnHeartbeatSeconds = turnHeartbeatSeconds; + return (this); + } + + + + /******************************************************************************* + ** Getter for turnLeaseSeconds + *******************************************************************************/ + public Integer getTurnLeaseSeconds() + { + return (this.turnLeaseSeconds); + } + + + + /******************************************************************************* + ** Setter for turnLeaseSeconds + *******************************************************************************/ + public void setTurnLeaseSeconds(Integer turnLeaseSeconds) + { + this.turnLeaseSeconds = turnLeaseSeconds; + } + + + + /******************************************************************************* + ** Fluent setter for turnLeaseSeconds + *******************************************************************************/ + public AssistantQBitConfig withTurnLeaseSeconds(Integer turnLeaseSeconds) + { + this.turnLeaseSeconds = turnLeaseSeconds; + return (this); + } + + + + /******************************************************************************* + ** Getter for turnForceFinalizeGraceSeconds + *******************************************************************************/ + public Integer getTurnForceFinalizeGraceSeconds() + { + return (this.turnForceFinalizeGraceSeconds); + } + + + + /******************************************************************************* + ** Setter for turnForceFinalizeGraceSeconds + *******************************************************************************/ + public void setTurnForceFinalizeGraceSeconds(Integer turnForceFinalizeGraceSeconds) + { + this.turnForceFinalizeGraceSeconds = turnForceFinalizeGraceSeconds; + } + + + + /******************************************************************************* + ** Fluent setter for turnForceFinalizeGraceSeconds + *******************************************************************************/ + public AssistantQBitConfig withTurnForceFinalizeGraceSeconds(Integer turnForceFinalizeGraceSeconds) + { + this.turnForceFinalizeGraceSeconds = turnForceFinalizeGraceSeconds; + return (this); + } + + + + /******************************************************************************* + ** Getter for shutdownGraceSeconds + *******************************************************************************/ + public Integer getShutdownGraceSeconds() + { + return (this.shutdownGraceSeconds); + } + + + + /******************************************************************************* + ** Setter for shutdownGraceSeconds + *******************************************************************************/ + public void setShutdownGraceSeconds(Integer shutdownGraceSeconds) + { + this.shutdownGraceSeconds = shutdownGraceSeconds; + } + + + + /******************************************************************************* + ** Fluent setter for shutdownGraceSeconds + *******************************************************************************/ + public AssistantQBitConfig withShutdownGraceSeconds(Integer shutdownGraceSeconds) + { + this.shutdownGraceSeconds = shutdownGraceSeconds; + return (this); + } + + + + /******************************************************************************* + ** Getter for defaultConcurrencyMode + *******************************************************************************/ + public String getDefaultConcurrencyMode() + { + return (this.defaultConcurrencyMode); + } + + + + /******************************************************************************* + ** Setter for defaultConcurrencyMode + *******************************************************************************/ + public void setDefaultConcurrencyMode(String defaultConcurrencyMode) + { + this.defaultConcurrencyMode = defaultConcurrencyMode; + } + + + + /******************************************************************************* + ** Fluent setter for defaultConcurrencyMode + *******************************************************************************/ + public AssistantQBitConfig withDefaultConcurrencyMode(String defaultConcurrencyMode) + { + this.defaultConcurrencyMode = defaultConcurrencyMode; + return (this); + } + + + + /******************************************************************************* + ** Getter for registerShutdownHook + *******************************************************************************/ + public Boolean getRegisterShutdownHook() + { + return (this.registerShutdownHook); + } + + + + /******************************************************************************* + ** Setter for registerShutdownHook + *******************************************************************************/ + public void setRegisterShutdownHook(Boolean registerShutdownHook) + { + this.registerShutdownHook = registerShutdownHook; + } + + + + /******************************************************************************* + ** Fluent setter for registerShutdownHook + *******************************************************************************/ + public AssistantQBitConfig withRegisterShutdownHook(Boolean registerShutdownHook) + { + this.registerShutdownHook = registerShutdownHook; + return (this); + } + + + + /******************************************************************************* + ** Getter for runTracker + *******************************************************************************/ + @JsonIgnore + public AssistantRunTracker getRunTracker() + { + return (this.runTracker == null ? AssistantRunTracker.NOOP : this.runTracker); + } + + + + /******************************************************************************* + ** Setter for runTracker + *******************************************************************************/ + public void setRunTracker(AssistantRunTracker runTracker) + { + this.runTracker = runTracker; + } + + + + /******************************************************************************* + ** Fluent setter for runTracker + *******************************************************************************/ + public AssistantQBitConfig withRunTracker(AssistantRunTracker runTracker) + { + this.runTracker = runTracker; + return (this); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantRuntime.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantRuntime.java index d28fe05..522cca2 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantRuntime.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantRuntime.java @@ -17,8 +17,11 @@ package com.kingsrook.qbits.assistant; +import java.net.InetAddress; import java.time.Clock; +import java.time.Duration; import java.util.Map; +import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -50,6 +53,7 @@ import com.kingsrook.qqq.backend.core.actions.ai.AiToolKernelSettings; import com.kingsrook.qqq.backend.core.actions.ai.exposure.AiExposurePolicy; import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor; import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.logging.QLogger; import com.kingsrook.qqq.backend.core.model.metadata.QInstance; import com.kingsrook.qqq.backend.core.model.metadata.qbits.QBitMetaData; import com.kingsrook.qqq.backend.core.utils.StringUtils; @@ -62,6 +66,11 @@ import com.kingsrook.qqq.backend.core.utils.StringUtils; *******************************************************************************/ public class AssistantRuntime { + private static final QLogger LOG = QLogger.getLogger(AssistantRuntime.class); + private final String instanceId = instanceId(); + private volatile boolean shuttingDown; + private Thread shutdownHook; + private final ScheduledTurn scheduledTurn = new ScheduledTurn(); private final RagService ragService; @@ -103,12 +112,22 @@ public class AssistantRuntime ** *******************************************************************************/ public AssistantRuntime(QInstance qInstance, AssistantQBitConfig config, Clock clock) + { + this(qInstance, config, clock, new AssistantStore(config)); + } + + + + /******************************************************************************* + ** Explicit persistence collaborator for embedded runtimes and fault tests. + *******************************************************************************/ + public AssistantRuntime(QInstance qInstance, AssistantQBitConfig config, Clock clock, AssistantStore store) { AssistantChannelWall.validate(qInstance); this.qInstance = qInstance; this.config = config; this.clock = clock; - this.store = new AssistantStore(config); + this.store = store; this.policyService = new PolicyService(config, store, clock); this.providerService = new ProviderService(config, store, clock); this.budgetService = new BudgetService(store, providerService, clock); @@ -157,6 +176,21 @@ public class AssistantRuntime this.knowledgeTool = new KnowledgeTool(this, ragService, vectors); this.bulkLoadTools = new BulkLoadTools(qInstance, dataAccessService, config.getBulkLoadRunner()); this.contextBuilder = new ContextBuilder(config, kernel, toolbox).withKnowledgeTool(knowledgeTool).withBulkLoadTools(bulkLoadTools); + if(Boolean.TRUE.equals(config.getRegisterShutdownHook())) + { + shutdownHook = new Thread(() -> + { + try + { + shutdown(Duration.ofSeconds(config.getShutdownGraceSeconds())); + } + catch(Throwable ignored) + { + // Shutdown must never fail the JVM hook. + } + }, "assistant-shutdown-" + instanceId); + Runtime.getRuntime().addShutdownHook(shutdownHook); + } } @@ -312,8 +346,13 @@ public class AssistantRuntime /******************************************************************************* ** Track a turn running here. *******************************************************************************/ - public RunningTurn register(RunningTurn turn) + public synchronized RunningTurn register(RunningTurn turn) { + if(shuttingDown) + { + turn.forceFinalize("SHUTDOWN"); + return (turn); + } runningTurns.put(turn.getTurnId(), turn); return (turn); } @@ -345,7 +384,17 @@ public class AssistantRuntime *******************************************************************************/ public void cancelAll(String reason) { - runningTurns.values().forEach(t -> t.cancel(reason)); + runningTurns.values().forEach(t -> + { + try + { + t.cancel(reason); + } + catch(Exception e) + { + LOG.warn("Assistant cancellation failed; finalization will still run"); + } + }); } @@ -355,8 +404,7 @@ public class AssistantRuntime *******************************************************************************/ public void close() { - runningTurns.values().forEach(t -> t.cancel("SHUTDOWN")); - executor.shutdownNow(); + shutdown(Duration.ZERO); } @@ -504,4 +552,97 @@ public class AssistantRuntime return (bulkLoadTools); } + + + + /******************************************************************************* + ** Stop accepting work, cancel, wait a bounded grace, then force finalization. + *******************************************************************************/ + public void shutdown(Duration grace) + { + synchronized(this) + { + shuttingDown = true; + } + cancelAll("SHUTDOWN"); + long until = System.nanoTime() + Math.max(0, grace.toNanos()); + while(!runningTurns.isEmpty() && System.nanoTime() < until) + { + try + { + Thread.sleep(10); + } + catch(InterruptedException e) + { + Thread.currentThread().interrupt(); + break; + } + } + runningTurns.values().forEach(t -> + { + try + { + t.forceFinalize("SHUTDOWN"); + } + catch(Exception e) + { + LOG.warn("Assistant shutdown finalization failed"); + } + }); + executor.shutdownNow(); + if(shutdownHook != null && Thread.currentThread() != shutdownHook) + { + try + { + Runtime.getRuntime().removeShutdownHook(shutdownHook); + } + catch(IllegalStateException ignored) + { + // The JVM is already shutting down. + } + } + } + + + + /******************************************************************************* + ** Getter for shuttingDown + *******************************************************************************/ + public boolean isShuttingDown() + { + return (shuttingDown); + } + + + + /******************************************************************************* + ** Getter for instanceId + *******************************************************************************/ + public String getInstanceId() + { + return (instanceId); + } + + + + /******************************************************************************* + ** A stable, bounded hostname plus a per-runtime random suffix. + *******************************************************************************/ + private static String instanceId() + { + String host = System.getenv("HOSTNAME"); + if(!StringUtils.hasContent(host)) + { + try + { + host = InetAddress.getLocalHost().getHostName(); + } + catch(Exception ignored) + { + host = "unknown"; + } + } + return (host.substring(0, Math.min(87, host.length())) + ":" + UUID.randomUUID().toString().substring(0, 8)); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RateLimiter.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RateLimiter.java index 73b3c05..1e644da 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RateLimiter.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RateLimiter.java @@ -26,15 +26,12 @@ import java.util.Map; /******************************************************************************* - ** Per-user turn rate limits (spec 3.9): turns started per minute and turns - ** running at once. In memory, per replica - with several replicas the - ** effective limit is per replica (the concurrent limit is also checked - ** against the turn table by the API). + ** Per-user starts per minute, in memory per replica. The API implements + ** concurrency by superseding turns, never by refusing another question. *******************************************************************************/ public class RateLimiter { private final Map> starts = new HashMap<>(); - private final Map running = new HashMap<>(); private final Clock clock; @@ -61,17 +58,12 @@ public class RateLimiter queue.pollFirst(); } - if(running.getOrDefault(userId, 0) >= concurrent) - { - return ("You already have an assistant answer in progress. Wait for it to finish or stop it."); - } if(queue.size() >= perMinute) { return ("Too many assistant requests in the last minute. Try again shortly."); } queue.addLast(now); - running.merge(userId, 1, Integer::sum); return (null); } @@ -82,7 +74,7 @@ public class RateLimiter *******************************************************************************/ public synchronized void release(String userId) { - running.computeIfPresent(userId, (k, v) -> v <= 1 ? null : v - 1); + // Only starts per minute are counted; completion does not refund a start. } } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RunningTurn.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RunningTurn.java index 73afa0d..4d6ba75 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RunningTurn.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/RunningTurn.java @@ -17,7 +17,10 @@ package com.kingsrook.qbits.assistant.engine; +import java.time.Instant; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import com.kingsrook.qbits.assistant.spi.StreamHandle; @@ -30,6 +33,14 @@ public class RunningTurn private final String userId; private final TurnEventLog log = new TurnEventLog(); private final AtomicReference cancelReason = new AtomicReference<>(); + private final AtomicBoolean finalized = new AtomicBoolean(); + private Instant startedAt; + private volatile String phase = "starting"; + private volatile String toolName; + private volatile Thread engineThread; + private volatile Consumer forceFinalizer; + private volatile long cancelledNanos; + private volatile String trackingId; private volatile StreamHandle currentHandle; @@ -39,6 +50,7 @@ public class RunningTurn *******************************************************************************/ public RunningTurn(Integer turnId, String userId) { + this.startedAt = Instant.now(); this.turnId = turnId; this.userId = userId; } @@ -50,12 +62,24 @@ public class RunningTurn *******************************************************************************/ public void cancel(String reason) { - if(cancelReason.compareAndSet(null, reason)) + if(!finalized.get() && cancelReason.compareAndSet(null, reason)) { - StreamHandle handle = currentHandle; - if(handle != null) + cancelledNanos = System.nanoTime(); + try { - handle.cancel(); + StreamHandle handle = currentHandle; + if(handle != null) + { + handle.cancel(); + } + } + finally + { + Thread thread = engineThread; + if(thread != null && thread != Thread.currentThread()) + { + thread.interrupt(); + } } } } @@ -124,4 +148,140 @@ public class RunningTurn return (log); } + + + + /******************************************************************************* + ** Claim terminal processing exactly once. + *******************************************************************************/ + public boolean tryFinalize() + { + return (finalized.compareAndSet(false, true)); + } + + + + /******************************************************************************* + ** Whether terminal processing has been claimed. + *******************************************************************************/ + public boolean isFinalized() + { + return (finalized.get()); + } + + + + /******************************************************************************* + ** Getter for startedAt + *******************************************************************************/ + public Instant getStartedAt() + { + return (startedAt); + } + + + + /******************************************************************************* + ** Setter for startedAt + *******************************************************************************/ + public void setStartedAt(Instant startedAt) + { + this.startedAt = startedAt; + } + + + + /******************************************************************************* + ** Getter for phase + *******************************************************************************/ + public String getPhase() + { + return (phase); + } + + + + /******************************************************************************* + ** Getter for toolName + *******************************************************************************/ + public String getToolName() + { + return (toolName); + } + + + + /******************************************************************************* + ** Record the current phase and optional tool. + *******************************************************************************/ + public void setPhase(String phase, String toolName) + { + this.phase = phase; + this.toolName = toolName; + } + + + + /******************************************************************************* + ** Record the thread to interrupt when cancelling. + *******************************************************************************/ + public void setEngineThread(Thread engineThread) + { + this.engineThread = engineThread; + } + + + + /******************************************************************************* + ** Install the engine finalizer before registration. + *******************************************************************************/ + public void setForceFinalizer(Consumer forceFinalizer) + { + this.forceFinalizer = forceFinalizer; + } + + + + /******************************************************************************* + ** Force terminal processing using the original user session. + *******************************************************************************/ + public void forceFinalize(String reason) + { + Consumer finalizer = forceFinalizer; + if(finalizer != null) + { + finalizer.accept(reason); + } + } + + + + /******************************************************************************* + ** Whether cancellation has exceeded its bounded grace period. + *******************************************************************************/ + public boolean cancelGraceElapsed(int seconds) + { + return (isCancelled() && cancelledNanos != 0 && System.nanoTime() - cancelledNanos >= seconds * 1_000_000_000L); + } + + + + /******************************************************************************* + ** Getter for trackingId + *******************************************************************************/ + public String getTrackingId() + { + return (trackingId); + } + + + + /******************************************************************************* + ** Setter for trackingId + *******************************************************************************/ + public void setTrackingId(String trackingId) + { + this.trackingId = trackingId; + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java index fa92d3a..afc65a0 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java @@ -102,6 +102,10 @@ public class TurnEngine implements Runnable { private static final QLogger LOG = QLogger.getLogger(TurnEngine.class); + public static final String SUPERSEDED_TEXT = "Stopped because you asked something else."; + public static final String TIMEOUT_TEXT = "The assistant took too long to answer, so it was stopped. Ask again."; + public static final String RESTART_TEXT = "The server restarted while answering. Ask again."; + private static final int MAX_INPUT_JSON_CHARS = 4000; private static final AiRedactor INPUT_REDACTOR = new AiRedactor(List.of()); @@ -175,6 +179,10 @@ public class TurnEngine implements Runnable this.policy = policy; this.running = running; this.softBudgetNotice = softBudgetNotice; + this.startedAt = turn.getStartedAt() == null ? runtime.getClock().instant() : turn.getStartedAt(); + this.deadline = startedAt.plusSeconds(config.getMaxTurnSeconds()); + running.setStartedAt(startedAt); + running.setForceFinalizer(this::forceFinish); } @@ -198,8 +206,11 @@ public class TurnEngine implements Runnable @Override public void run() { - startedAt = runtime.getClock().instant(); - deadline = startedAt.plusSeconds(config.getMaxTurnSeconds()); + if(running.isFinalized()) + { + return; + } + running.setEngineThread(Thread.currentThread()); String status = "COMPLETE"; String stopReason = null; @@ -209,6 +220,7 @@ public class TurnEngine implements Runnable try { QContext.init(runtime.getQInstance(), session); + startWatchdog(); scheduledTask = runtime.getScheduledTurn().task(conversation.getId()); if(conversation.getScheduledTaskId() != null && scheduledTask == null) { @@ -218,9 +230,9 @@ public class TurnEngine implements Runnable ScheduledTurn.checkCost(runtime, conversation.getId(), turn.getId()); classification = new TurnClassification(config.getClassificationSource(), conversation, pageContext); store.saveClassification(userId, conversation.getId(), turn.getId(), classification); - startKillSwitchWatcher(); - emit(TurnEvent.MESSAGE_START, map("turnId", turn.getId(), "conversationId", conversation.getId(), "userMessageId", userMessage.getId(), "messageId", assistantMessage.getId())); + + emit(TurnEvent.MESSAGE_START, map("turnId", turn.getId(), "conversationId", conversation.getId(), "userMessageId", userMessage.getId(), "messageId", assistantMessage.getId(), "startedAt", startedAt.toString())); if(softBudgetNotice != null) { notice("warning", softBudgetNotice); @@ -232,6 +244,7 @@ public class TurnEngine implements Runnable ModelRouter.Route route; try { + phase("thinking", null); route = router.route(session, conversation, turn, userText, pageContext, policy, classification, running, deadline, text -> notice("warning", text)); } finally @@ -420,8 +433,21 @@ public class TurnEngine implements Runnable } finally { - finished = true; + // Clear interruption while persisting the terminal state. + boolean interrupted = Thread.interrupted(); + if(running.isCancelled()) + { + Stop stop = cancelledStop(); + status = stop.status(); + stopReason = stop.stopReason(); + errorCode = stop.errorCode(); + errorText = stop.message(); + } finish(status, stopReason, errorCode, errorText); + if(interrupted) + { + Thread.currentThread().interrupt(); + } QContext.clear(); } } @@ -429,35 +455,56 @@ public class TurnEngine implements Runnable /******************************************************************************* - ** Poll the kill switch while the turn runs (with the user's own session), - ** so a flip cancels the turn - and every other turn here - even while it - ** waits on the model. + ** Renew the lease, poll cancellation and the kill switch with the user's + ** own session, and enforce a hard deadline even if a provider is wedged. *******************************************************************************/ - private void startKillSwitchWatcher() + private void startWatchdog() { - long pollMillis = config.getKillSwitchPollMillis() == null ? 2000 : config.getKillSwitchPollMillis(); - Thread.ofVirtual().name("assistant-killswitch-" + turn.getId()).start(() -> + long heartbeatMillis = config.getTurnHeartbeatSeconds() * 1000L; + long pollMillis = Math.min(heartbeatMillis, config.getKillSwitchPollMillis() == null ? 2000 : config.getKillSwitchPollMillis()); + Thread.ofVirtual().name("assistant-watchdog-" + turn.getId()).start(() -> { + long lastHeartbeat = System.nanoTime(); try { - while(!finished && !running.isCancelled()) + QContext.init(runtime.getQInstance(), session); + while(!running.isFinalized()) { Thread.sleep(pollMillis); - if(finished) + if(running.isFinalized()) { - return; + break; } - QContext.init(runtime.getQInstance(), session); try { + if(System.nanoTime() - lastHeartbeat >= heartbeatMillis * 1_000_000L) + { + lastHeartbeat = System.nanoTime(); + heartbeat(); + } if(runtime.getPolicyService().isKillSwitchOn()) { runtime.cancelAll("KILL_SWITCH"); } } + catch(Exception e) + { + LOG.warn("Assistant watchdog tick failed", hipaaSafe(e), logPair("turnId", turn.getId())); + } finally { - QContext.clear(); + // Deadlines do not depend on successful database or policy reads. + Instant now = runtime.getClock().instant(); + if(!now.isBefore(deadline)) + { + safely(() -> running.cancel("TIMEOUT")); + } + if(!now.isBefore(deadline.plusSeconds(config.getTurnForceFinalizeGraceSeconds())) + || running.cancelGraceElapsed(config.getTurnForceFinalizeGraceSeconds())) + { + Stop stop = running.isCancelled() ? cancelledStop() : new Stop("ERROR", "TIMEOUT", TIMEOUT_TEXT, "timeout"); + finish(stop.status(), stop.stopReason(), stop.errorCode(), stop.message()); + } } } } @@ -465,11 +512,63 @@ public class TurnEngine implements Runnable { Thread.currentThread().interrupt(); } + finally + { + QContext.clear(); + } }); } + /******************************************************************************* + ** Renew only an active row, then observe persisted cancellation requests. + *******************************************************************************/ + private synchronized void heartbeat() throws QException + { + if(running.isFinalized()) + { + return; + } + AssistantTurn row = store.getTurn(userId, turn.getId()); + if(row == null || !("RUNNING".equals(row.getStatus()) || "CANCELLING".equals(row.getStatus()))) + { + safely(() -> running.cancel("USER")); + finish(row == null ? "ERROR" : row.getStatus(), row == null ? "interrupted" : row.getStopReason(), + row == null ? "INTERNAL" : row.getErrorCode(), null); + return; + } + store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", turn.getId()) + .withValue("heartbeatAt", runtime.getClock().instant()).withValue("phase", running.getPhase())); + row = store.getTurn(userId, turn.getId()); + if(row != null && (row.getCancelRequestedAt() != null || "CANCELLING".equals(row.getStatus()))) + { + running.cancel(row.getCancelReason() == null ? "USER" : row.getCancelReason()); + } + } + + + + /******************************************************************************* + ** Finalize from shutdown using the original authenticated user's context. + *******************************************************************************/ + private void forceFinish(String reason) + { + try + { + QContext.init(runtime.getQInstance(), session); + safely(() -> running.cancel(reason)); + Stop stop = "SHUTDOWN".equals(reason) ? new Stop("ERROR", "INTERRUPTED", RESTART_TEXT, "interrupted") : cancelledStop(); + finish(stop.status(), stop.stopReason(), stop.errorCode(), stop.message()); + } + finally + { + QContext.clear(); + } + } + + + /******************************************************************************* ** Why a turn stops. *******************************************************************************/ @@ -504,14 +603,15 @@ public class TurnEngine implements Runnable } if(!runtime.getClock().instant().isBefore(deadline)) { - return (new Stop("ERROR", "TIMEOUT", "The assistant ran out of time for this question.", "timeout")); + running.cancel("TIMEOUT"); + return (cancelledStop()); } try { AssistantTurn row = store.getTurn(userId, turn.getId()); if(row != null && "CANCELLING".equals(row.getStatus())) { - running.cancel("USER"); + running.cancel(row.getCancelReason() == null ? "USER" : row.getCancelReason()); return (cancelledStop()); } } @@ -529,6 +629,18 @@ public class TurnEngine implements Runnable *******************************************************************************/ private Stop cancelledStop() { + if("TIMEOUT".equals(running.getCancelReason())) + { + return (new Stop("ERROR", "TIMEOUT", TIMEOUT_TEXT, "timeout")); + } + if("SHUTDOWN".equals(running.getCancelReason())) + { + return (new Stop("ERROR", "INTERRUPTED", RESTART_TEXT, "interrupted")); + } + if("SUPERSEDED".equals(running.getCancelReason())) + { + return (new Stop("CANCELLED", null, SUPERSEDED_TEXT, ModelResponse.STOP_CANCELLED)); + } if("KILL_SWITCH".equals(running.getCancelReason())) { String reason = StringUtils.hasContent(policy.killSwitchReason()) ? " " + policy.killSwitchReason() : ""; @@ -618,6 +730,7 @@ public class TurnEngine implements Runnable @Override public void onTextDelta(String text) { + phase("writing", null); bufferText(text); } }; @@ -639,6 +752,7 @@ public class TurnEngine implements Runnable TurnDataRouting.beforeModelCall(runtime, classification, turn, active); fallback.beforeCall(active); modelCallCount++; + phase("thinking", null); StreamHandle handle = provider.stream(request.copy(), fallback.listener(listener)); running.setCurrentHandle(handle); try @@ -707,6 +821,7 @@ public class TurnEngine implements Runnable continue; } toolCallCount++; + phase("tool", use.getToolName()); emit(TurnEvent.TOOL_START, map("messageId", assistantMessage.getId(), "toolCallId", use.getToolUseId(), "name", use.getToolName(), "summary", ToolLabels.running(use.getToolName()))); if(!allowedTools.contains(use.getToolName())) @@ -996,93 +1111,100 @@ public class TurnEngine implements Runnable /******************************************************************************* ** Persist the message, turn, conversation and usage; emit the end events. *******************************************************************************/ - private void finish(String status, String stopReason, String errorCode, String errorText) + private synchronized void finish(String status, String stopReason, String errorCode, String errorText) { + if(!running.tryFinalize()) + { + return; + } try { - persistClassification(); - flushText(true); - Instant now = runtime.getClock().instant(); - - if(errorText != null) - { - parts.addNotice("error", errorText); - } - - String model = this.model != null ? this.model : (policy == null ? null : policy.mainModel()); + AssistantTurn persistedTurn = null; try { - store.updateFields(AssistantMessage.TABLE_NAME, new QRecord() - .withValue("id", assistantMessage.getId()) - .withValue("status", status) - .withValue("model", model) - .withValue("partsJson", parts.toJson())); - - AssistantConversation current = store.getConversation(userId, conversation.getId()); - if(current != null) - { - store.updateFields(AssistantConversation.TABLE_NAME, new QRecord() - .withValue("id", current.getId()) - .withValue("lastMessageAt", now) - .withValue("unreadCount", (current.getUnreadCount() == null ? 0 : current.getUnreadCount()) + 1) - .withValue("totalTokensIn", nz(current.getTotalTokensIn()) + nz(usage.getInputTokens()) + nz(usage.getCacheReadTokens()) + nz(usage.getCacheWriteTokens())) - .withValue("totalTokensOut", nz(current.getTotalTokensOut()) + nz(usage.getOutputTokens())) - .withValue("totalCostMicros", nz(current.getTotalCostMicros()) + costMicros)); - } - - store.addUsage(userId, LocalDate.now(runtime.getClock().withZone(ZoneOffset.UTC)), usage, costMicros, true); + persistedTurn = store.getTurn(userId, turn.getId()); } - catch(QException e) + catch(Exception e) { - LOG.warn("Could not persist the end of an assistant turn", hipaaSafe(e), logPair("turnId", turn.getId())); + LOG.warn("Could not read final assistant status", hipaaSafe(e), logPair("turnId", turn.getId())); + } + boolean alreadyTerminal = persistedTurn != null && !("RUNNING".equals(persistedTurn.getStatus()) || "CANCELLING".equals(persistedTurn.getStatus())); + if(alreadyTerminal) + { + status = persistedTurn.getStatus(); + stopReason = persistedTurn.getStopReason(); + errorCode = persistedTurn.getErrorCode(); + } + String finalStatus = status; + String finalStopReason = stopReason; + String finalErrorCode = errorCode; + String model = this.model != null ? this.model : (policy == null ? null : policy.mainModel()); + Instant now = runtime.getClock().instant(); + safely(() -> flushText(true)); + if(errorText != null && !alreadyTerminal) + { + parts.addNotice("CANCELLED".equals(status) ? "info" : "error", errorText); + } + running.setPhase("done", null); + if(!alreadyTerminal) + { + safely(() -> persistClassification()); + safely(() -> store.updateFields(AssistantMessage.TABLE_NAME, new QRecord() + .withValue("id", assistantMessage.getId()).withValue("status", finalStatus) + .withValue("model", model).withValue("partsJson", parts.toJson()))); } - ///////////////////////////////////////////////////////////////////// - // terminal events are persisted BEFORE the turn row turns final, // - // so a replay of a finished turn always ends with message.end // - ///////////////////////////////////////////////////////////////////// emit(TurnEvent.USAGE, map("turnId", turn.getId(), "inputTokens", usage.getInputTokens(), "outputTokens", usage.getOutputTokens(), "cacheReadTokens", usage.getCacheReadTokens(), "cacheWriteTokens", usage.getCacheWriteTokens(), "costMicros", costMicros)); if(errorCode != null) { emit(TurnEvent.ERROR, map("code", errorCode, "message", errorText)); } - Map end = map("messageId", assistantMessage.getId(), "status", status); - if(stopReason != null) + emit(TurnEvent.MESSAGE_END, map("messageId", assistantMessage.getId(), "status", status, "stopReason", stopReason)); + if(!alreadyTerminal) { - end.put("stopReason", stopReason); + safely(this::persistPendingEvents); + // Accounting failures cannot skip the terminal row or local cleanup. + safely(() -> + { + store.updateRunningTurn(userId, new QRecord() + .withValue("id", turn.getId()) + .withValue("status", finalStatus) + .withValue("phase", "done") + .withValue("stopReason", truncate(finalStopReason, 64)) + .withValue("errorCode", finalErrorCode) + .withValue("cancelReason", running.getCancelReason()) + .withValue("model", model) + .withValue("providerId", active == null ? null : active.providerId()) + .withValue("routeDecision", turn.getRouteDecision()) + .withValue("routeSource", turn.getRouteSource()) + .withValue("fallbackCount", turn.getFallbackCount()) + .withValue("endedAt", now) + .withValue("toolCallCount", toolCallCount) + .withValue("modelCallCount", modelCallCount) + .withValue("inputTokens", usage.getInputTokens()) + .withValue("outputTokens", usage.getOutputTokens()) + .withValue("cacheReadTokens", usage.getCacheReadTokens()) + .withValue("cacheWriteTokens", usage.getCacheWriteTokens()) + .withValue("costMicros", costMicros) + .withValue("providerRequestIds", requestIds.isEmpty() ? null : String.join(",", requestIds))); + }); + safely(() -> + { + AssistantConversation current = store.getConversation(userId, conversation.getId()); + if(current != null) + { + store.updateFields(AssistantConversation.TABLE_NAME, new QRecord() + .withValue("id", current.getId()) + .withValue("lastMessageAt", now) + .withValue("unreadCount", (current.getUnreadCount() == null ? 0 : current.getUnreadCount()) + 1) + .withValue("totalTokensIn", nz(current.getTotalTokensIn()) + nz(usage.getInputTokens()) + nz(usage.getCacheReadTokens()) + nz(usage.getCacheWriteTokens())) + .withValue("totalTokensOut", nz(current.getTotalTokensOut()) + nz(usage.getOutputTokens())) + .withValue("totalCostMicros", nz(current.getTotalCostMicros()) + costMicros)); + } + }); + safely(() -> store.addUsage(userId, LocalDate.now(runtime.getClock().withZone(ZoneOffset.UTC)), usage, costMicros, true)); } - emit(TurnEvent.MESSAGE_END, end); - persistPendingEvents(); - - try - { - store.updateFields(AssistantTurn.TABLE_NAME, new QRecord() - .withValue("id", turn.getId()) - .withValue("status", status) - .withValue("stopReason", truncate(stopReason, 64)) - .withValue("errorCode", errorCode) - .withValue("cancelReason", running.getCancelReason()) - .withValue("model", model) - .withValue("providerId", active == null ? null : active.providerId()) - .withValue("routeDecision", turn.getRouteDecision()) - .withValue("routeSource", turn.getRouteSource()) - .withValue("fallbackCount", turn.getFallbackCount()) - .withValue("endedAt", now) - .withValue("toolCallCount", toolCallCount) - .withValue("modelCallCount", modelCallCount) - .withValue("inputTokens", usage.getInputTokens()) - .withValue("outputTokens", usage.getOutputTokens()) - .withValue("cacheReadTokens", usage.getCacheReadTokens()) - .withValue("cacheWriteTokens", usage.getCacheWriteTokens()) - .withValue("costMicros", costMicros) - .withValue("providerRequestIds", requestIds.isEmpty() ? null : String.join(",", requestIds))); - } - catch(QException e) - { - LOG.warn("Could not persist the end of an assistant turn", hipaaSafe(e), logPair("turnId", turn.getId())); - } - LOG.info("Assistant turn finished", logPair("turnId", turn.getId()), logPair("conversationId", conversation.getId()), @@ -1091,6 +1213,7 @@ public class TurnEngine implements Runnable logPair("status", status), logPair("stopReason", stopReason), logPair("errorCode", errorCode), + logPair("cancelReason", running.getCancelReason()), logPair("modelCalls", modelCallCount), logPair("toolCalls", toolCallCount), logPair("inputTokens", usage.getInputTokens()), @@ -1100,29 +1223,90 @@ public class TurnEngine implements Runnable logPair("costMicros", costMicros), logPair("durationMs", Duration.between(startedAt, now).toMillis())); } - catch(QException | RuntimeException e) - { - LOG.warn("Assistant turn finish failed", hipaaSafe(e), logPair("turnId", turn.getId())); - } finally { + finished = true; running.getLog().close(); - if(scheduledTask != null) + safely(() -> { - runtime.getScheduledTurn().detach(conversation.getId(), session.getUuid()); + if(scheduledTask != null) + { + runtime.getScheduledTurn().detach(conversation.getId(), session.getUuid()); + } + }); + String finalStatus = status; + String finalText = errorText; + try + { + safely(() -> config.getRunTracker().finish(running.getTrackingId(), finalStatus, finalText)); + } + finally + { + runtime.unregister(turn.getId()); + runtime.getRateLimiter().release(userId); } - runtime.unregister(turn.getId()); - runtime.getRateLimiter().release(userId); } } + /******************************************************************************* + ** One independently fallible finalization operation. + *******************************************************************************/ + @FunctionalInterface + private interface FinishStep + { + /******************************************************************************* + ** Execute the step. + *******************************************************************************/ + void run() throws Exception; + } + + + + /******************************************************************************* + ** A failure in a collaborator must not strand the run. + *******************************************************************************/ + private void safely(FinishStep step) + { + try + { + step.run(); + } + catch(Exception e) + { + LOG.warn("Assistant lifecycle step failed", hipaaSafe(e), logPair("turnId", turn.getId())); + } + } + + + + /******************************************************************************* + ** Emit and track only phase changes, including the current tool label. + *******************************************************************************/ + private synchronized void phase(String phase, String toolName) + { + if(running.isFinalized() || (phase.equals(running.getPhase()) && Objects.equals(toolName, running.getToolName()))) + { + return; + } + running.setPhase(phase, toolName); + String label = toolName == null ? phase : ToolLabels.running(toolName); + emit(TurnEvent.TURN_PHASE, map("phase", phase, "toolName", toolName, "label", label)); + safely(() -> config.getRunTracker().progress(running.getTrackingId(), phase, label)); + } + + + /******************************************************************************* ** Buffer streamed text; emit it in small chunks (provider thread). *******************************************************************************/ - private void bufferText(String text) + private synchronized void bufferText(String text) { + if(running.isFinalized()) + { + return; + } synchronized(textBuffer) { textBuffer.append(text); @@ -1135,7 +1319,7 @@ public class TurnEngine implements Runnable /******************************************************************************* ** Emit buffered text as one text.delta (when big or old enough, or forced). *******************************************************************************/ - private void flushText(boolean force) + private synchronized void flushText(boolean force) { String chunk; synchronized(textBuffer) @@ -1170,8 +1354,12 @@ public class TurnEngine implements Runnable ** Append to the live log; queue for persistence (done on the engine ** thread, which has the QContext). *******************************************************************************/ - private void emit(String type, Map data) + private synchronized void emit(String type, Map data) { + if(finished) + { + return; + } toPersist.add(running.getLog().append(type, data)); } @@ -1189,7 +1377,7 @@ public class TurnEngine implements Runnable { store.insertEvent(turn.getId(), userId, event.getSeq(), event.getType(), event.toJson()); } - catch(QException e) + catch(Exception e) { LOG.warn("Could not persist assistant turn event", hipaaSafe(e), logPair("turnId", turn.getId()), logPair("seq", event.getSeq())); } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEvent.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEvent.java index 946e944..c655ec6 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEvent.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEvent.java @@ -29,6 +29,7 @@ import com.kingsrook.qqq.backend.core.utils.JsonUtils; *******************************************************************************/ public class TurnEvent { + public static final String TURN_PHASE = "turn.phase"; public static final String MESSAGE_START = "message.start"; public static final String TEXT_DELTA = "text.delta"; public static final String TOOL_START = "tool.start"; diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/jobs/AssistantRunTracker.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/jobs/AssistantRunTracker.java new file mode 100644 index 0000000..1d62be0 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/jobs/AssistantRunTracker.java @@ -0,0 +1,75 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.jobs; + + +/******************************************************************************* + ** STAND-IN(core=QRun-IO/qqq#1210, removeBy=RC.9): core job tracking + ** (jobs_job tracking rows) will implement this contract; hosts may wire their own. + *******************************************************************************/ +public interface AssistantRunTracker +{ + AssistantRunTracker NOOP = new AssistantRunTracker() + { + /******************************************************************************* + ** No tracking is configured. + *******************************************************************************/ + @Override + public String start(String userId, Integer turnId, Integer conversationId, String label) + { + return (null); + } + + /******************************************************************************* + ** No tracking is configured. + *******************************************************************************/ + @Override + public void progress(String trackingId, String phase, String message) + { + } + + /******************************************************************************* + ** No tracking is configured. + *******************************************************************************/ + @Override + public void finish(String trackingId, String status, String message) + { + } + }; + + /******************************************************************************* + ** Publish a run and return its tracking id. + *******************************************************************************/ + String start(String userId, Integer turnId, Integer conversationId, String label); + + /******************************************************************************* + ** Publish a changed phase. + *******************************************************************************/ + void progress(String trackingId, String phase, String message); + + /******************************************************************************* + ** Publish the terminal outcome. + *******************************************************************************/ + void finish(String trackingId, String status, String message); + + /******************************************************************************* + ** Wire the host's Stop action to the owner's cancellation path. + *******************************************************************************/ + default void onCancel(String trackingId, Runnable cancel) + { + } +} -- 2.54.0 From f60009f59513734ebe473d346d34a20b96dfa6fa Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:41 -0500 Subject: [PATCH 3/5] feat(api): never refuse a question; cross-pod Stop; GET /turns/active A new question supersedes the running one in the conversation (default) or runs concurrently; the soft per-user cap supersedes the oldest. Stop works on any replica and on dead runs, and is idempotent. Orphans are reaped once their lease lapses (45 s) instead of after maxTurnSeconds + 60 s. 503 SHUTTING_DOWN while a replica stops. turn.phase events and startedAt on message.start. UI API extension assistant 1.2.0. Co-Authored-By: Claude Opus 5.5 --- .../qbits/assistant/api/AssistantApi.java | 331 ++++++++++++++---- .../qbits/assistant/api/AssistantRoutes.java | 4 +- .../qbits/assistant/api/TurnStreams.java | 11 +- .../assistant/api/specs/AssistantSpecs.java | 91 ++++- .../resources/openapi/qqq-v1/assistant.yaml | 109 +++++- 5 files changed, 461 insertions(+), 85 deletions(-) diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantApi.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantApi.java index df9d926..fce8160 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantApi.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantApi.java @@ -17,11 +17,14 @@ package com.kingsrook.qbits.assistant.api; +import java.time.Duration; import java.time.Instant; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import com.fasterxml.jackson.core.type.TypeReference; import com.kingsrook.qbits.assistant.AssistantQBitConfig; @@ -51,11 +54,13 @@ import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome; import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult; import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor; import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.logging.QLogger; import com.kingsrook.qqq.backend.core.model.data.QRecord; import com.kingsrook.qqq.backend.core.model.session.QSession; import com.kingsrook.qqq.backend.core.utils.JsonUtils; import com.kingsrook.qqq.backend.core.utils.StringUtils; import com.kingsrook.qqq.middleware.javalin.specs.invoke.UiApiCredentials; +import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair; /******************************************************************************* @@ -65,6 +70,8 @@ import com.kingsrook.qqq.middleware.javalin.specs.invoke.UiApiCredentials; *******************************************************************************/ public class AssistantApi { + private static final QLogger LOG = QLogger.getLogger(AssistantApi.class); + public static final int TITLE_MAX_CHARS = 60; private final AssistantRuntime runtime; @@ -289,6 +296,19 @@ public class AssistantApi public RunningTurn startTurn(Integer conversationId, Map body) throws QException, AssistantApiException { requireAccess(); + if(runtime.isShuttingDown()) + { + throw (new AssistantApiException(503, "SHUTTING_DOWN", "The server is restarting. Try again in a moment.")); + } + String concurrency = config.getDefaultConcurrencyMode(); + if(body.containsKey("concurrency")) + { + if(!(body.get("concurrency") instanceof String mode) || !("supersede".equalsIgnoreCase(mode) || "concurrent".equalsIgnoreCase(mode))) + { + throw (new AssistantApiException(400, "concurrency must be supersede or concurrent.")); + } + concurrency = mode.toUpperCase(Locale.ROOT); + } AssistantConversation conversation = requireConversation(conversationId); if(!(body.get("text") instanceof String text) || !StringUtils.hasContent(text)) @@ -324,25 +344,36 @@ public class AssistantApi throw (new AssistantApiException(429, "BUDGET", budget.message())); } - long running = 0; - for(AssistantTurn other : store.listRunningTurns(userId)) - { - if(isRunning(reapIfOrphaned(other))) - { - running++; - } - } - if(running >= config.getMaxConcurrentTurnsPerUser()) - { - throw (new AssistantApiException(429, "RATE_LIMIT", "You already have an assistant answer in progress. Wait for it to finish or stop it.")); - } - String limited = runtime.getRateLimiter().tryAcquire(userId, config.getMaxTurnsPerUserPerMinute(), config.getMaxConcurrentTurnsPerUser()); if(limited != null) { throw (new AssistantApiException(429, "RATE_LIMIT", limited)); } + List remaining = new ArrayList<>(); + for(AssistantTurn other : store.listRunningTurns(userId)) + { + other = reapIfOrphaned(other); + if(!isRunning(other) || "SUPERSEDED".equals(other.getCancelReason())) + { + continue; + } + if("SUPERSEDE".equals(concurrency) && conversationId.equals(other.getConversationId())) + { + requestStop(other, "SUPERSEDED"); + } + else + { + remaining.add(other); + } + } + remaining.sort(Comparator.comparing(AssistantTurn::getStartedAt, Comparator.nullsFirst(Comparator.naturalOrder())) + .thenComparing(AssistantTurn::getId)); + while(remaining.size() >= config.getMaxConcurrentTurnsPerUser()) + { + requestStop(remaining.removeFirst(), "SUPERSEDED"); + } + try { if(!StringUtils.hasContent(conversation.getTitle())) @@ -361,13 +392,18 @@ public class AssistantApi .withSeq(seq) .withPartsJson(new MessageParts().addText(text).toJson())); + Instant startedAt = runtime.getClock().instant(); AssistantTurn turn = store.insert(new AssistantTurn() .withConversationId(conversationId) .withUserId(userId) .withUserMessageId(userMessage.getId()) .withStatus("RUNNING") .withModel(active.model()) - .withStartedAt(runtime.getClock().instant()) + .withStartedAt(startedAt) + .withHeartbeatAt(startedAt) + .withLeaseOwner(runtime.getInstanceId()) + .withPhase("starting") + .withConcurrencyMode(concurrency) .withToolCallCount(0) .withModelCallCount(0)); @@ -385,10 +421,37 @@ public class AssistantApi store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", turn.getId()).withValue("assistantMessageId", assistantMessage.getId())); turn.setAssistantMessageId(assistantMessage.getId()); - RunningTurn runningTurn = runtime.register(new RunningTurn(turn.getId(), userId)); + RunningTurn runningTurn = new RunningTurn(turn.getId(), userId); TurnEngine engine = new TurnEngine(runtime, session, userId, conversation, userMessage, assistantMessage, turn, text, pageContext, policy, runningTurn, budget.softExceeded() ? budget.message() : null).withCredentials(credentials); - runtime.getExecutor().submit(engine); + try + { + if(config.getRunTracker() != null) + { + /////////////////////////////////////////////////////////////////// + // never the conversation title: it is derived from the prompt // + /////////////////////////////////////////////////////////////////// + runningTurn.setTrackingId(config.getRunTracker().start(userId, turn.getId(), conversationId, "Assistant answer")); + config.getRunTracker().onCancel(runningTurn.getTrackingId(), () -> runningTurn.cancel("USER")); + } + } + catch(Exception e) + { + LOG.warn("Assistant run tracker start failed", logPair("turnId", turn.getId())); + } + runtime.register(runningTurn); + if(!runningTurn.isFinalized()) + { + try + { + runtime.getExecutor().submit(engine); + } + catch(RuntimeException e) + { + runningTurn.forceFinalize("SHUTDOWN"); + throw (e); + } + } return (runningTurn); } catch(QException | RuntimeException e) @@ -443,22 +506,94 @@ public class AssistantApi { requireAccess(); AssistantTurn turn = requireTurn(turnId); - if(runtime.getRunningTurn(turnId) == null) + if(isRunning(turn)) + { + turn = requestStop(turn, "USER"); + } + return (Map.of("status", turn.getStatus())); + } + + + + /******************************************************************************* + ** Durable Stop, shared by the API and superseding another question. + *******************************************************************************/ + private AssistantTurn requestStop(AssistantTurn turn, String reason) throws QException + { + synchronized(store) + { + turn = store.getTurn(userId, turn.getId()); + if(!isRunning(turn)) + { + return (turn); + } + RunningTurn live = runtime.getRunningTurn(turn.getId()); + if(live == null && orphaned(turn)) + { + return (reap(turn, reason)); + } + Instant now = runtime.getClock().instant(); + if(!store.updateRunningTurn(userId, new QRecord().withValue("id", turn.getId()) + .withValue("status", "CANCELLING").withValue("cancelReason", reason).withValue("cancelRequestedAt", now))) + { + return (store.getTurn(userId, turn.getId())); + } + turn.setStatus("CANCELLING"); + turn.setCancelReason(reason); + turn.setCancelRequestedAt(now); + if(live != null) + { + try + { + live.cancel(reason); + } + catch(Exception e) + { + LOG.warn("Assistant stream cancellation failed; watchdog will finalize"); + } + } + return (turn); + } + } + + + + /******************************************************************************* + ** GET /turns/active: the caller's runs, newest first, after orphan recovery. + *******************************************************************************/ + public Map activeTurns() throws QException, AssistantApiException + { + requireAccess(); + List turns = new ArrayList<>(); + for(AssistantTurn turn : store.listRunningTurns(userId)) { turn = reapIfOrphaned(turn); + if(isRunning(turn)) + { + turns.add(turn); + } } - if(!isRunning(turn)) + turns.sort(Comparator.comparing(AssistantTurn::getStartedAt, Comparator.nullsFirst(Comparator.naturalOrder())) + .thenComparing(AssistantTurn::getId).reversed()); + List> rows = new ArrayList<>(); + for(AssistantTurn turn : turns) { - throw (new AssistantApiException(409, "This answer has already finished.")); + RunningTurn live = runtime.getRunningTurn(turn.getId()); + Map row = new LinkedHashMap<>(); + row.put("turnId", turn.getId()); + row.put("conversationId", turn.getConversationId()); + row.put("assistantMessageId", turn.getAssistantMessageId()); + row.put("userMessageId", turn.getUserMessageId()); + row.put("status", turn.getStatus()); + row.put("phase", live == null ? turn.getPhase() : live.getPhase()); + row.put("startedAt", turn.getStartedAt() == null ? null : turn.getStartedAt().toString()); + row.put("heartbeatAt", turn.getHeartbeatAt() == null ? null : turn.getHeartbeatAt().toString()); + row.put("elapsedMs", turn.getStartedAt() == null ? 0 : Math.max(0, Duration.between(turn.getStartedAt(), runtime.getClock().instant()).toMillis())); + row.put("live", live != null); + row.put("concurrencyMode", turn.getConcurrencyMode()); + rows.add(row); } - - store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", turnId).withValue("status", "CANCELLING")); - RunningTurn live = runtime.getRunningTurn(turnId); - if(live != null && userId.equals(live.getUserId())) - { - live.cancel("USER"); - } - return (Map.of("status", "CANCELLING")); + return (Map.of("turns", rows)); } @@ -474,60 +609,104 @@ public class AssistantApi /******************************************************************************* - ** A turn whose replica died: it still says RUNNING / CANCELLING, no turn - ** of this runtime owns it, and it started longer ago than maxTurnSeconds - ** + 60 s (a live engine always ends itself by then). On first sight it is - ** finished as ERROR: the assistant message gets status ERROR and a - ** notice, and synthetic error (INTERNAL) + message.end events are - ** persisted, so replays end. Returns the (possibly updated) turn. + ** A remote run loses its lease, exceeds the absolute deadline, or has a + ** cancellation request that its owner did not acknowledge in time. + *******************************************************************************/ + private boolean orphaned(AssistantTurn turn) + { + Instant now = runtime.getClock().instant(); + Instant heartbeat = turn.getHeartbeatAt() == null ? turn.getStartedAt() : turn.getHeartbeatAt(); + return ((heartbeat != null && heartbeat.isBefore(now.minusSeconds(config.getTurnLeaseSeconds()))) + || (turn.getStartedAt() != null && turn.getStartedAt().isBefore(now.minusSeconds(maxStreamSeconds()))) + || ("CANCELLING".equals(turn.getStatus()) && turn.getCancelRequestedAt() != null + && turn.getCancelRequestedAt().isBefore(now.minusSeconds(config.getTurnLeaseSeconds())))); + } + + + + /******************************************************************************* + ** Recover an abandoned run and leave a terminal replay for refreshes. *******************************************************************************/ AssistantTurn reapIfOrphaned(AssistantTurn turn) throws QException { - if(!isRunning(turn) || runtime.getRunningTurn(turn.getId()) != null || turn.getStartedAt() == null) + synchronized(store) { - return (turn); - } - Instant stale = runtime.getClock().instant().minusSeconds(config.getMaxTurnSeconds() + 60L); - if(!turn.getStartedAt().isBefore(stale)) - { - return (turn); - } - - String text = "This answer was interrupted (the server restarted). Ask again."; - if(turn.getAssistantMessageId() != null) - { - AssistantMessage message = store.getMessage(userId, turn.getAssistantMessageId()); - if(message != null) + if(!isRunning(turn) || runtime.getRunningTurn(turn.getId()) != null || !orphaned(turn)) { - store.updateFields(AssistantMessage.TABLE_NAME, new QRecord() - .withValue("id", message.getId()) - .withValue("status", "ERROR") - .withValue("partsJson", MessageParts.fromJson(message.getPartsJson()).addNotice("error", text).toJson())); + return (turn); + } + String reason = "CANCELLING".equals(turn.getStatus()) && turn.getCancelRequestedAt() != null + && turn.getCancelRequestedAt().isBefore(runtime.getClock().instant().minusSeconds(config.getTurnLeaseSeconds())) + ? (turn.getCancelReason() == null ? "USER" : turn.getCancelReason()) : null; + return (reap(turn, reason)); + } + } + + + + /******************************************************************************* + ** Re-read before recovery; never overwrite an already terminal row or + ** append another terminal event. Each persistence step is independent. + *******************************************************************************/ + private AssistantTurn reap(AssistantTurn turn, String cancelReason) throws QException + { + AssistantTurn current = store.getTurn(userId, turn.getId()); + if(!isRunning(current)) + { + return (current); + } + if(!orphaned(current)) + { + return (cancelReason == null ? current : requestStop(current, cancelReason)); + } + String status = cancelReason == null ? "ERROR" : "CANCELLED"; + String text = cancelReason == null ? "This answer was interrupted (the server restarted). Ask again." + : ("SUPERSEDED".equals(cancelReason) ? TurnEngine.SUPERSEDED_TEXT : "Stopped by you."); + if(current.getAssistantMessageId() != null) + { + try + { + AssistantMessage message = store.getMessage(userId, current.getAssistantMessageId()); + if(message != null) + { + store.updateFields(AssistantMessage.TABLE_NAME, new QRecord().withValue("id", message.getId()).withValue("status", status) + .withValue("partsJson", MessageParts.fromJson(message.getPartsJson()).addNotice(cancelReason == null ? "error" : "info", text).toJson())); + } + } + catch(Exception e) + { + LOG.warn("Could not persist recovered assistant message"); } } - - int seq = store.lastEventSeq(userId, turn.getId()); - store.insertEvent(turn.getId(), userId, seq + 1, TurnEvent.ERROR, new TurnEvent(seq + 1, TurnEvent.ERROR, Map.of("code", "INTERNAL", "message", text)).toJson()); - Map end = new java.util.LinkedHashMap<>(); - if(turn.getAssistantMessageId() != null) + try { - end.put("messageId", turn.getAssistantMessageId()); + int seq = store.lastEventSeq(userId, current.getId()); + List last = store.listEvents(userId, current.getId(), Math.max(0, seq - 1), 1); + if(last.isEmpty() || !TurnEvent.MESSAGE_END.equals(last.getFirst().getEventType())) + { + if(cancelReason == null) + { + seq++; + store.insertEvent(current.getId(), userId, seq, TurnEvent.ERROR, + new TurnEvent(seq, TurnEvent.ERROR, Map.of("code", "INTERNAL", "message", text)).toJson()); + } + seq++; + Map end = new LinkedHashMap<>(); + end.put("messageId", current.getAssistantMessageId()); + end.put("status", status); + end.put("stopReason", cancelReason == null ? "interrupted" : "cancelled"); + store.insertEvent(current.getId(), userId, seq, TurnEvent.MESSAGE_END, new TurnEvent(seq, TurnEvent.MESSAGE_END, end).toJson()); + } + } + catch(Exception e) + { + LOG.warn("Could not persist recovered assistant events"); } - end.put("status", "ERROR"); - end.put("stopReason", "interrupted"); - store.insertEvent(turn.getId(), userId, seq + 2, TurnEvent.MESSAGE_END, new TurnEvent(seq + 2, TurnEvent.MESSAGE_END, end).toJson()); - Instant now = runtime.getClock().instant(); - store.updateFields(AssistantTurn.TABLE_NAME, new QRecord() - .withValue("id", turn.getId()) - .withValue("status", "ERROR") - .withValue("errorCode", "INTERNAL") - .withValue("stopReason", "interrupted") - .withValue("endedAt", now)); - turn.setStatus("ERROR"); - turn.setErrorCode("INTERNAL"); - turn.setEndedAt(now); - return (turn); + store.updateRunningTurn(userId, new QRecord().withValue("id", current.getId()).withValue("status", status) + .withValue("errorCode", cancelReason == null ? "INTERNAL" : null).withValue("cancelReason", cancelReason) + .withValue("stopReason", cancelReason == null ? "interrupted" : "cancelled").withValue("phase", "done").withValue("endedAt", now)); + return (store.getTurn(userId, current.getId())); } @@ -786,6 +965,16 @@ public class AssistantApi + /******************************************************************************* + ** Hard upper bound for remote stream polling and stale-run recovery. + *******************************************************************************/ + public long maxStreamSeconds() + { + return (config.getMaxTurnSeconds() + (long) config.getTurnForceFinalizeGraceSeconds() + config.getTurnLeaseSeconds()); + } + + + /******************************************************************************* ** The user may use the panel: enabled, signed in, assistant.hasAccess. *******************************************************************************/ diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantRoutes.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantRoutes.java index e5a1728..7edbe89 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantRoutes.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/AssistantRoutes.java @@ -17,12 +17,12 @@ package com.kingsrook.qbits.assistant.api; +import com.kingsrook.qbits.assistant.api.specs.AssistantSpecs; import com.kingsrook.qqq.backend.core.model.metadata.QInstance; import com.kingsrook.qqq.backend.core.model.metadata.code.QCodeReference; -import com.kingsrook.qbits.assistant.api.specs.AssistantSpecs; import com.kingsrook.qqq.backend.core.utils.CollectionUtils; -import com.kingsrook.qqq.middleware.javalin.specs.extensions.MiddlewareSpecExtensionsMetaData; import com.kingsrook.qqq.middleware.javalin.QJavalinMetaData; +import com.kingsrook.qqq.middleware.javalin.specs.extensions.MiddlewareSpecExtensionsMetaData; /******************************************************************************* diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/TurnStreams.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/TurnStreams.java index 9e9ccb3..b710b4a 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/TurnStreams.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/TurnStreams.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.io.OutputStream; import java.nio.charset.StandardCharsets; import java.util.List; +import java.util.Map; import com.kingsrook.qbits.assistant.engine.TurnEvent; import com.kingsrook.qbits.assistant.engine.TurnEventLog; import com.kingsrook.qqq.backend.core.logging.QLogger; @@ -101,7 +102,7 @@ public final class TurnStreams OutputStream out = startStream(context); int last = after; long lastWrite = System.currentTimeMillis(); - long giveUpAt = System.currentTimeMillis() + (api.maxTurnSeconds() + 60L) * 1000L; + long giveUpAt = System.currentTimeMillis() + api.maxStreamSeconds() * 1000L; AssistantApi.TurnEvents current = first; while(true) { @@ -113,10 +114,16 @@ public final class TurnStreams lastWrite = System.currentTimeMillis(); ended |= event.isTerminal(); } - if(ended || (current.finished() && !current.more()) || System.currentTimeMillis() > giveUpAt) + if(ended || (current.finished() && !current.more())) { return; } + if(System.currentTimeMillis() > giveUpAt) + { + write(out, new TurnEvent(last + 1, TurnEvent.ERROR, Map.of("code", "TIMEOUT", "message", "The event stream timed out. Reconnect to check this answer.")).toSseFrame()); + write(out, new TurnEvent(last + 2, TurnEvent.MESSAGE_END, Map.of("status", "ERROR", "stopReason", "stream_timeout")).toSseFrame()); + return; + } if(!current.more()) { Thread.sleep(REMOTE_POLL_MS); diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/specs/AssistantSpecs.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/specs/AssistantSpecs.java index b4aa2f6..44d863f 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/specs/AssistantSpecs.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/api/specs/AssistantSpecs.java @@ -61,7 +61,7 @@ import io.javalin.http.HttpStatus; public final class AssistantSpecs { public static final String EXTENSION_NAME = "assistant"; - public static final String EXTENSION_VERSION = "1.1.0"; + public static final String EXTENSION_VERSION = "1.2.0"; public static final String PREFIX = "/assistant"; @@ -94,6 +94,7 @@ public final class AssistantSpecs .withSpec(StartTurnSpec.class) .withSpec(TurnEventsSpec.class) .withSpec(CancelTurnSpec.class) + .withSpec(ActiveTurnsSpec.class) .withSpec(FeedbackSpec.class) .withSpec(GetActionSpec.class) .withSpec(ConfirmActionSpec.class) @@ -218,9 +219,14 @@ public final class AssistantSpecs "href", string("After a confirm: the written record's /app link.").withNullable(true))); rs.put("AssistantAccessResponseV1", object("Explain access: per table, the user's UI permissions, the assistant's data access, and each tool's verdict.", "tables", array(freeObject("table, label, dataAccess, dataAccessSource, ui {visible, read, insert, edit, delete}, operations [{tool, operation, allowed, reason}]."), null))); - rs.put("AssistantTurnEventV1", object("One server-sent event of a turn (event: , id: , data: this object). Types: message.start, text.delta, tool.start, tool.result, ui.navigate, notice, usage, message.end, error, ping.", + rs.put("AssistantActiveTurnsResponseV1", object("The caller's active runs, newest first; stale leases are reaped.", + "turns", array(object("An active assistant run.", "turnId", integer(null), "conversationId", integer(null), + "assistantMessageId", integer(null), "userMessageId", integer(null), "status", string("RUNNING or CANCELLING."), + "phase", string(null), "startedAt", string("ISO-8601."), "heartbeatAt", string("ISO-8601."), + "elapsedMs", integer(null), "live", bool("Running on this replica."), "concurrencyMode", string("SUPERSEDE or CONCURRENT.")), null))); + rs.put("AssistantTurnEventV1", object("One server-sent event of a turn (event: , id: , data: this object). Types: message.start, turn.phase, text.delta, tool.start, tool.result, ui.navigate, notice, usage, message.end, error, ping.", "seq", integer("Sequence number; resume with ?after=."), "type", string("The event type (also the SSE event name)."), - "data", freeObject("The event's fields, by type."))); + "data", freeObject("The event's fields, by type. message.start includes startedAt (ISO-8601). turn.phase includes phase (thinking, tool, writing), optional toolName and label."))); rs.put("AssistantErrorV1", object("An error.", "error", string("A message for the user."), "code", string("BUDGET, RATE_LIMIT, NOT_CONFIGURED, KILL_SWITCH, DISABLED, NO_PERMISSION...").withNullable(true))); return (rs); } @@ -562,6 +568,67 @@ public final class AssistantSpecs + /******************************************************************************* + ** GET /assistant/turns/active + *******************************************************************************/ + public static class ActiveTurnsSpec extends AssistantJsonSpec + { + /******************************************************************************* + ** HTTP method. + *******************************************************************************/ + @Override + protected HttpMethod method() + { + return (HttpMethod.GET); + } + + + + /******************************************************************************* + ** Route path. + *******************************************************************************/ + @Override + protected String path() + { + return (PREFIX + "/turns/active"); + } + + + + /******************************************************************************* + ** Operation summary. + *******************************************************************************/ + @Override + protected String summary() + { + return ("Find the user's active assistant runs"); + } + + + + /******************************************************************************* + ** Response schema. + *******************************************************************************/ + @Override + protected String getResponseSchemaName() + { + return ("AssistantActiveTurnsResponseV1"); + } + + + + /******************************************************************************* + ** Runs in the invoker's authenticated user session. + *******************************************************************************/ + @Override + protected Object handle(AssistantApi api, JsonInput input) throws Exception + { + return (api.activeTurns()); + } + } + + + /******************************************************************************* ** POST /assistant/conversations *******************************************************************************/ @@ -844,7 +911,7 @@ public final class AssistantSpecs @Override protected Schema defineRequestBodySchema() { - return (object("The user's text and the page context hint.", "text", string("The user's message."), "pageContext", freeObject("Where the user is in the app."))); + return (object("The user's text and the page context hint.", "text", string("The user's message."), "pageContext", freeObject("Where the user is in the app."), "concurrency", string("supersede or concurrent (case-insensitive); defaults to the host configuration."))); } @@ -1038,10 +1105,24 @@ public final class AssistantSpecs + /******************************************************************************* + ** Older replicas can still return 202 during a rolling upgrade. New + ** replicas return 200 for both active and already finished turns. + *******************************************************************************/ + @Override + public List defineAdditionalBasicResponses() + { + List responses = new java.util.ArrayList<>(super.defineAdditionalBasicResponses()); + responses.add(new BasicResponse(HttpStatus.ACCEPTED, "Legacy replica: cancellation accepted", "AssistantStatusResponseV1")); + return (responses); + } + + + @Override protected HttpStatus successStatus() { - return (HttpStatus.ACCEPTED); + return (HttpStatus.OK); } } diff --git a/qbit-assistant/src/main/resources/openapi/qqq-v1/assistant.yaml b/qbit-assistant/src/main/resources/openapi/qqq-v1/assistant.yaml index 0b2e636..0e63520 100644 --- a/qbit-assistant/src/main/resources/openapi/qqq-v1/assistant.yaml +++ b/qbit-assistant/src/main/resources/openapi/qqq-v1/assistant.yaml @@ -53,6 +53,43 @@ components: nullable: true type: "string" type: "object" + AssistantActiveTurnsResponseV1: + description: "The caller's active runs, newest first; stale leases are reaped." + properties: + turns: + items: + description: "An active assistant run." + properties: + turnId: + type: "integer" + conversationId: + type: "integer" + assistantMessageId: + type: "integer" + userMessageId: + type: "integer" + status: + description: "RUNNING or CANCELLING." + type: "string" + phase: + type: "string" + startedAt: + description: "ISO-8601." + type: "string" + heartbeatAt: + description: "ISO-8601." + type: "string" + elapsedMs: + type: "integer" + live: + description: "Running on this replica." + type: "boolean" + concurrencyMode: + description: "SUPERSEDE or CONCURRENT." + type: "string" + type: "object" + type: "array" + type: "object" AssistantConfigResponseV1: description: "Whether the assistant is available to this user, and how." properties: @@ -187,8 +224,8 @@ components: type: "object" AssistantTurnEventV1: description: "One server-sent event of a turn (event: , id: , data:\ - \ this object). Types: message.start, text.delta, tool.start, tool.result,\ - \ ui.navigate, notice, usage, message.end, error, ping." + \ this object). Types: message.start, turn.phase, text.delta, tool.start,\ + \ tool.result, ui.navigate, notice, usage, message.end, error, ping." properties: seq: description: "Sequence number; resume with ?after=." @@ -198,7 +235,9 @@ components: type: "string" data: additionalProperties: true - description: "The event's fields, by type." + description: "The event's fields, by type. message.start includes startedAt\ + \ (ISO-8601). turn.phase includes phase (thinking, tool, writing), optional\ + \ toolName and label." type: "object" type: "object" securitySchemes: @@ -212,7 +251,7 @@ info: description: "QQQ Assistant: conversations, streamed turns, feedback, pending actions\ \ and data access explanations" title: "QQQ Middleware API extension: assistant" - version: "1.1.0" + version: "1.2.0" openapi: "3.0.3" paths: /qqq/v1/assistant/config: @@ -605,6 +644,10 @@ paths: additionalProperties: true description: "Where the user is in the app." type: "object" + concurrency: + description: "supersede or concurrent (case-insensitive); defaults\ + \ to the host configuration." + type: "string" type: "object" required: true responses: @@ -754,6 +797,12 @@ paths: schema: $ref: "#/components/schemas/AssistantErrorV1" description: "Disabled, not configured, or the kill switch is on" + 200: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantStatusResponseV1" + description: "Successful Response" 409: content: application/json: @@ -765,7 +814,7 @@ paths: application/json: schema: $ref: "#/components/schemas/AssistantStatusResponseV1" - description: "Successful Response" + description: "Legacy replica: cancellation accepted" 429: content: application/json: @@ -776,6 +825,56 @@ paths: - sessionUuidCookie: - "N/A" summary: "Stop a running turn" + /qqq/v1/assistant/turns/active: + get: + description: "Find the user's active assistant runs" + responses: + 400: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "Invalid input" + 403: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "No access to the assistant" + 404: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "Not found (or not the user's own)" + 503: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "Disabled, not configured, or the kill switch is on" + 200: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantActiveTurnsResponseV1" + description: "Successful Response" + 409: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "Already finished or decided" + 429: + content: + application/json: + schema: + $ref: "#/components/schemas/AssistantErrorV1" + description: "Budget or rate limit" + security: + - sessionUuidCookie: + - "N/A" + summary: "Find the user's active assistant runs" /qqq/v1/assistant/messages/{id}/feedback: post: description: "Thumbs up or down on an assistant message" -- 2.54.0 From 50898a927df4f267f99879c13a42c6f9de64dadc Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:42 -0500 Subject: [PATCH 4/5] test(turns): completion, error after streaming, disconnect, refresh, supersede, concurrency, wedged provider, pod kill, rolling deploy, cross-pod Stop Co-Authored-By: Claude Opus 5.5 --- .../api/AssistantRouteProviderTest.java | 103 +- .../assistant/engine/DurableTurnsTest.java | 902 ++++++++++++++++++ .../engine/EngineComponentsTest.java | 3 +- .../assistant/engine/RobustnessTest.java | 2 +- .../testing/AssistantTestFixture.java | 1 + 5 files changed, 1001 insertions(+), 10 deletions(-) create mode 100644 qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/DurableTurnsTest.java diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/api/AssistantRouteProviderTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/api/AssistantRouteProviderTest.java index 4dd4316..e721617 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/api/AssistantRouteProviderTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/api/AssistantRouteProviderTest.java @@ -30,9 +30,13 @@ import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import com.fasterxml.jackson.core.type.TypeReference; import com.kingsrook.qbits.assistant.model.AssistantPolicy; import com.kingsrook.qbits.assistant.persistence.AssistantStore; +import com.kingsrook.qbits.assistant.spi.ModelPart; +import com.kingsrook.qbits.assistant.spi.ModelResponse; import com.kingsrook.qbits.assistant.testing.AssistantTestFixture; import com.kingsrook.qbits.assistant.testing.BaseAssistantTest; import com.kingsrook.qbits.assistant.testing.UiApiTestServer; @@ -314,7 +318,7 @@ class AssistantRouteProviderTest extends BaseAssistantTest assertThat(send("POST", "/messages/" + messages.get(0).get("id") + "/feedback", EDITOR, "{\"rating\":1}").statusCode()).isEqualTo(400); assertThat(send("POST", "/messages/" + msgId + "/feedback", OTHER_EDITOR, "{\"rating\":-1}").statusCode()).isEqualTo(404); - assertThat(send("POST", "/turns/" + turnId + "/cancel", EDITOR, "{}").statusCode()).isEqualTo(409); + assertThat(send("POST", "/turns/" + turnId + "/cancel", EDITOR, "{}").statusCode()).isEqualTo(200); } @@ -351,15 +355,10 @@ class AssistantRouteProviderTest extends BaseAssistantTest } assertThat(turnId).isNotNull(); - /////////////////////////////////////////////////////////////// - // a second turn while one runs: refused (concurrency = 1) // - /////////////////////////////////////////////////////////////// - assertThat(send("POST", "/conversations/" + id + "/turns", EDITOR, "{\"text\":\"again\"}").statusCode()).isEqualTo(429); - Thread.sleep(500); assertThat(send("POST", "/turns/" + turnId + "/cancel", OTHER_EDITOR, "{}").statusCode()).isEqualTo(404); HttpResponse cancel = send("POST", "/turns/" + turnId + "/cancel", EDITOR, "{}"); - assertThat(cancel.statusCode()).isEqualTo(202); + assertThat(cancel.statusCode()).isEqualTo(200); assertThat(json(cancel)).containsEntry("status", "CANCELLING"); List frames = frames(future.get().body()); @@ -371,6 +370,96 @@ class AssistantRouteProviderTest extends BaseAssistantTest + /******************************************************************************* + ** A disconnected POST keeps running; active discovery and replay can attach + ** again while the answer is blocked, then receive the terminal frame. + *******************************************************************************/ + @Test + @SuppressWarnings("unchecked") + void testDisconnectRefreshAndReattach() throws Exception + { + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + provider.thenToolUses("I am checking the count before answering. ", List.of(new Object[] { "count_records", Map.of("table", "blog_post") })) + .thenRespond(request -> + { + entered.countDown(); + try + { + if(!release.await(5, TimeUnit.SECONDS)) + { + throw (new IllegalStateException("test did not release provider")); + } + } + catch(InterruptedException e) + { + Thread.currentThread().interrupt(); + } + return (new ModelResponse().withParts(List.of(ModelPart.text("Finished after reconnect."))).withStopReason(ModelResponse.STOP_END_TURN)); + }); + Integer conversation = (Integer) json(send("POST", "/conversations", EDITOR, "{}")).get("id"); + HttpRequest request = HttpRequest.newBuilder(URI.create(base + "/conversations/" + conversation + "/turns")) + .header("Cookie", "sessionId=" + EDITOR).header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString("{\"text\":\"count posts\"}")).build(); + int lastSeq = 0; + try + { + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofInputStream()); + assertThat(response.statusCode()).isEqualTo(200); + try(BufferedReader reader = new BufferedReader(new InputStreamReader(response.body(), StandardCharsets.UTF_8))) + { + String event = null; + String line; + while((line = reader.readLine()) != null) + { + if(line.startsWith("id: ")) + { + lastSeq = Integer.parseInt(line.substring(4)); + } + if(line.startsWith("event: ")) + { + event = line.substring(7); + } + if(line.isEmpty() && "text.delta".equals(event)) + { + break; + } + } + } + assertThat(lastSeq).isPositive(); + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue(); + HttpResponse activeResponse = send("GET", "/turns/active", EDITOR, null); + assertThat(activeResponse.statusCode()).isEqualTo(200); + List> active = (List>) json(activeResponse).get("turns"); + assertThat(active).hasSize(1); + Map turn = active.getFirst(); + assertThat(turn).containsEntry("live", true).containsEntry("phase", "thinking").containsEntry("conversationId", conversation) + .containsKeys("startedAt", "heartbeatAt", "elapsedMs", "concurrencyMode"); + assertThat(json(send("GET", "/turns/active", OTHER_EDITOR, null)).get("turns")).isEqualTo(List.of()); + assertThat(send("GET", "/turns/active", NO_ACCESS, null).statusCode()).isEqualTo(403); + HttpRequest attach = HttpRequest.newBuilder(URI.create(base + "/turns/" + turn.get("turnId") + "/events?after=0")) + .header("Cookie", "sessionId=" + EDITOR).GET().build(); + HttpResponse attached = client.send(attach, HttpResponse.BodyHandlers.ofInputStream()); + release.countDown(); + List followed = frames(attached.body()); + assertThat(followed.getFirst().event()).isEqualTo("message.start"); + assertThat(followed.getFirst().data()).containsKey("startedAt"); + assertThat(followed.getLast().event()).isEqualTo("message.end"); + assertThat(followed.getLast().data()).containsEntry("status", "COMPLETE"); + HttpRequest replay = HttpRequest.newBuilder(URI.create(base + "/turns/" + turn.get("turnId") + "/events?after=" + lastSeq)) + .header("Cookie", "sessionId=" + EDITOR).GET().build(); + List rest = frames(client.send(replay, HttpResponse.BodyHandlers.ofInputStream()).body()); + assertThat(Integer.parseInt(rest.getFirst().id())).isEqualTo(lastSeq + 1); + assertThat(rest.getLast().event()).isEqualTo("message.end"); + } + finally + { + release.countDown(); + } + } + + + /******************************************************************************* ** *******************************************************************************/ diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/DurableTurnsTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/DurableTurnsTest.java new file mode 100644 index 0000000..14484eb --- /dev/null +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/DurableTurnsTest.java @@ -0,0 +1,902 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.engine; + + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; +import com.kingsrook.qbits.assistant.AssistantQBitConfig; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.api.AssistantApi; +import com.kingsrook.qbits.assistant.api.AssistantApiException; +import com.kingsrook.qbits.assistant.jobs.AssistantRunTracker; +import com.kingsrook.qbits.assistant.model.AssistantConversation; +import com.kingsrook.qbits.assistant.model.AssistantMessage; +import com.kingsrook.qbits.assistant.model.AssistantTurn; +import com.kingsrook.qbits.assistant.persistence.AssistantStore; +import com.kingsrook.qbits.assistant.provider.fake.ScriptedModelProvider; +import com.kingsrook.qbits.assistant.spi.ModelEventListener; +import com.kingsrook.qbits.assistant.spi.ModelPart; +import com.kingsrook.qbits.assistant.spi.ModelRequest; +import com.kingsrook.qbits.assistant.spi.ModelResponse; +import com.kingsrook.qbits.assistant.spi.ModelUsage; +import com.kingsrook.qbits.assistant.spi.StreamHandle; +import com.kingsrook.qbits.assistant.testing.AssistantTestFixture; +import com.kingsrook.qbits.assistant.testing.BaseAssistantTest; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import com.kingsrook.qqq.backend.core.model.data.QRecordEntity; +import com.kingsrook.qqq.backend.core.model.session.QSession; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.EDITOR; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + + +/******************************************************************************* + ** Durable run lifecycle, using latches instead of provider timing guesses. + *******************************************************************************/ +class DurableTurnsTest extends BaseAssistantTest +{ + private QSession session; + private AssistantApi api; + private AssistantStore store; + + + + /******************************************************************************* + ** Short, valid lease settings keep watchdog tests bounded. + *******************************************************************************/ + @Override + protected void customizeConfig(AssistantQBitConfig config) + { + config.withTurnHeartbeatSeconds(1).withTurnLeaseSeconds(3).withTurnForceFinalizeGraceSeconds(1) + .withMaxTurnSeconds(10).withKillSwitchPollMillis(50); + } + + + + /******************************************************************************* + ** Every operation uses the user's own fixture session. + *******************************************************************************/ + @BeforeEach + void atomicBackend() + { + config.getRuntime().close(); + config.setRuntime(new AssistantRuntime(qInstance, config, Clock.systemUTC(), new AtomicCallsStore(config))); + userContext(); + } + + + + /******************************************************************************* + ** Establish the caller's context, also after a test replaces the runtime. + *******************************************************************************/ + private void userContext() + { + session = AssistantTestFixture.ownerSession(EDITOR, config); + QContext.init(qInstance, session); + api = new AssistantApi(config.getRuntime(), session); + store = config.getRuntime().getStore(); + } + + + + /******************************************************************************* + ** Completion closes the lease, emits progress, and frees the next question. + *******************************************************************************/ + @Test + void completion() throws Exception + { + provider.thenText("Done").thenText("Again"); + Integer conversation = conversation(); + RunningTurn first = api.startTurn(conversation, Map.of("text", "hello")); + AssistantTurn row = completed(first); + assertThat(row.getStatus()).isEqualTo("COMPLETE"); + assertThat(row.getPhase()).isEqualTo("done"); + assertThat(row.getEndedAt()).isNotNull(); + assertThat(row.getHeartbeatAt()).isAfterOrEqualTo(row.getStartedAt()); + assertThat(row.getLeaseOwner()).isEqualTo(config.getRuntime().getInstanceId()).hasSizeLessThanOrEqualTo(100); + assertThat(row.getConcurrencyMode()).isEqualTo("SUPERSEDE"); + assertThat(store.listRunningTurns(EDITOR)).isEmpty(); + assertThat(store.updateRunningTurn(EDITOR, new QRecord().withValue("id", first.getTurnId()).withValue("status", "CANCELLING"))).isFalse(); + assertThat(first.getLog().after(0, 1).getFirst().getData()).containsEntry("startedAt", row.getStartedAt().toString()); + assertThat(first.getLog().after(0, 1)).filteredOn(e -> TurnEvent.TURN_PHASE.equals(e.getType())) + .extracting(e -> e.getData().get("phase")).containsExactly("thinking", "writing"); + assertThat(completed(api.startTurn(conversation, Map.of("text", "again"))).getStatus()).isEqualTo("COMPLETE"); + } + + + + /******************************************************************************* + ** A throwing post-processing collaborator cannot strand a streamed answer. + *******************************************************************************/ + @Test + void accountingFailureStillFinalizes() throws Exception + { + AssistantStore failing = new AssistantStore(config) + { + /******************************************************************************* + ** Inject a failure after streaming has finished. + *******************************************************************************/ + @Override + public void addUsage(String userId, LocalDate date, ModelUsage usage, long costMicros, boolean countTurn) + { + throw (new IllegalStateException("accounting unavailable")); + } + + /******************************************************************************* + ** Conversation totals fail independently of the turn update. + *******************************************************************************/ + @Override + public void updateFields(String table, QRecord record) throws QException + { + if(AssistantConversation.TABLE_NAME.equals(table) && record.getValue("totalTokensIn") != null) + { + throw (new IllegalStateException("totals unavailable")); + } + super.updateFields(table, record); + } + }; + config.getRuntime().close(); + AssistantRuntime runtime = new AssistantRuntime(qInstance, config, Clock.systemUTC()) + { + /******************************************************************************* + ** Substitute only the engine/API persistence collaborator. + *******************************************************************************/ + @Override + public AssistantStore getStore() + { + return (failing); + } + }; + config.setRuntime(runtime); + userContext(); + provider.thenText("A streamed answer").thenText("Next answer"); + Integer conversation = conversation(); + RunningTurn run = api.startTurn(conversation, Map.of("text", "hello")); + assertThat(completed(run).getStatus()).isEqualTo("COMPLETE"); + assertThat(api.turnEvents(run.getTurnId(), 0).persisted()).extracting(TurnEvent::getType).contains(TurnEvent.TEXT_DELTA, TurnEvent.MESSAGE_END); + assertThat(completed(api.startTurn(conversation, Map.of("text", "next"))).getStatus()).isEqualTo("COMPLETE"); + } + + + + /******************************************************************************* + ** Explicit concurrent mode permits overlap within and across conversations. + *******************************************************************************/ + @Test + void concurrentQuestions() throws Exception + { + ControlledProvider controlled = controlled(false); + Integer firstConversation = conversation(); + List runs = new ArrayList<>(); + runs.add(api.startTurn(firstConversation, Map.of("text", "one", "concurrency", "CoNcUrReNt"))); + controlled.awaitCalls(1); + runs.add(api.startTurn(firstConversation, Map.of("text", "two", "concurrency", "concurrent"))); + controlled.awaitCalls(2); + runs.add(api.startTurn(conversation(), Map.of("text", "three", "concurrency", "concurrent"))); + controlled.awaitCalls(3); + assertThat(active()).hasSize(3); + controlled.release(); + for(RunningTurn run : runs) + { + assertThat(completed(run).getStatus()).isEqualTo("COMPLETE"); + } + } + + + + /******************************************************************************* + ** Default supersede frees the slot immediately and records the explanation. + *******************************************************************************/ + @Test + void supersedeAndSoftCap() throws Exception + { + ControlledProvider controlled = controlled(false); + Integer conversation = conversation(); + RunningTurn first = api.startTurn(conversation, Map.of("text", "one")); + controlled.awaitCalls(1); + RunningTurn second = api.startTurn(conversation, Map.of("text", "two")); + controlled.awaitCalls(2); + AssistantTurn superseded = completed(first); + assertThat(superseded.getStatus()).isEqualTo("CANCELLED"); + assertThat(superseded.getCancelReason()).isEqualTo("SUPERSEDED"); + assertThat(store.getMessage(EDITOR, superseded.getAssistantMessageId()).getPartsJson()).contains(TurnEngine.SUPERSEDED_TEXT); + config.setMaxConcurrentTurnsPerUser(2); + RunningTurn third = api.startTurn(conversation(), Map.of("text", "three", "concurrency", "concurrent")); + controlled.awaitCalls(3); + RunningTurn fourth = api.startTurn(conversation(), Map.of("text", "four", "concurrency", "concurrent")); + controlled.awaitCalls(4); + assertThat(completed(second).getCancelReason()).isEqualTo("SUPERSEDED"); + controlled.release(); + assertThat(completed(third).getStatus()).isEqualTo("COMPLETE"); + assertThat(completed(fourth).getStatus()).isEqualTo("COMPLETE"); + } + + + + /******************************************************************************* + ** A provider ignoring both cancel and interruption is forcibly finalized. + *******************************************************************************/ + @Test + void wedgedProviderTimeout() throws Exception + { + ControlledProvider controlled = controlled(true); + controlled.emitText = false; + RecordingTracker tracker = new RecordingTracker(); + config.setRunTracker(tracker); + config.setMaxTurnSeconds(1); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wedged")); + try + { + controlled.awaitCalls(1); + AssistantTurn row = completed(run); + assertThat(row.getStatus()).isEqualTo("ERROR"); + assertThat(row.getErrorCode()).isEqualTo("TIMEOUT"); + assertThat(row.getCancelReason()).isEqualTo("TIMEOUT"); + assertThat(store.getMessage(EDITOR, row.getAssistantMessageId()).getPartsJson()).contains(TurnEngine.TIMEOUT_TEXT); + assertThat(controlled.handles.getFirst().cancelled).isTrue(); + assertThat(tracker.finishes.get()).isEqualTo(1); + int events = run.getLog().after(0, 1).size(); + controlled.release(); + await(() -> controlled.handles.getFirst().returned); + assertThat(run.getLog().after(0, 1)).hasSize(events); + assertThat(tracker.finishes.get()).isEqualTo(1); + } + finally + { + controlled.release(); + } + } + + + + /******************************************************************************* + ** Stop has its own bounded force path even when the provider ignores it. + *******************************************************************************/ + @Test + void stopWedgedProvider() throws Exception + { + ControlledProvider controlled = controlled(true); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + try + { + controlled.awaitCalls(1); + assertThat(api.cancelTurn(run.getTurnId())).containsEntry("status", "CANCELLING"); + AssistantTurn row = completed(run); + assertThat(row.getStatus()).isEqualTo("CANCELLED"); + assertThat(row.getCancelReason()).isEqualTo("USER"); + } + finally + { + controlled.release(); + } + } + + + + /******************************************************************************* + ** A failed heartbeat tick does not stop cancellation polling on later ticks. + *******************************************************************************/ + @Test + void heartbeatFailureDoesNotKillWatchdog() throws Exception + { + AtomicInteger heartbeats = new AtomicInteger(); + AssistantStore failingOnce = new AssistantStore(config) + { + /******************************************************************************* + ** Fail only the first renewal, then behave like the real store. + *******************************************************************************/ + @Override + public void updateFields(String table, QRecord record) throws QException + { + if(AssistantTurn.TABLE_NAME.equals(table) && record.getValue("heartbeatAt") != null && heartbeats.incrementAndGet() == 1) + { + throw (new QException("temporary heartbeat failure")); + } + super.updateFields(table, record); + } + }; + config.getRuntime().close(); + config.setRuntime(new AssistantRuntime(qInstance, config) + { + /******************************************************************************* + ** A transient persistence failure local to the owner. + *******************************************************************************/ + @Override + public AssistantStore getStore() + { + return (failingOnce); + } + }); + userContext(); + ControlledProvider controlled = controlled(false); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + controlled.awaitCalls(1); + AssistantRuntime other = new AssistantRuntime(qInstance, config); + try + { + new AssistantApi(other, session).cancelTurn(run.getTurnId()); + assertThat(completed(run).getStatus()).isEqualTo("CANCELLED"); + assertThat(heartbeats.get()).isGreaterThanOrEqualTo(2); + } + finally + { + controlled.release(); + other.close(); + } + } + + + + /******************************************************************************* + ** Stale leases, old cancel requests and absolute deadlines release remotely. + *******************************************************************************/ + @Test + void orphanRecoveryAndIdempotentStop() throws Exception + { + AssistantTurn stale = remote(Instant.now().minusSeconds(5)); + AssistantTurn fresh = remote(Instant.now()); + AssistantTurn deadline = remote(Instant.now()); + store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", deadline.getId()).withValue("startedAt", Instant.now().minusSeconds(30))); + assertThat(active()).extracting(r -> r.get("turnId")).containsExactly(fresh.getId()); + assertThat(store.getTurn(EDITOR, stale.getId()).getStatus()).isEqualTo("ERROR"); + assertThat(store.getTurn(EDITOR, deadline.getId()).getStatus()).isEqualTo("ERROR"); + assertThat(api.turnEvents(stale.getId(), 0).persisted()).extracting(TurnEvent::getType).containsExactly(TurnEvent.ERROR, TurnEvent.MESSAGE_END); + assertThat(api.turnEvents(stale.getId(), 0).persisted()).hasSize(2); + assertThat(api.cancelTurn(stale.getId())).containsEntry("status", "ERROR"); + assertThat(api.cancelTurn(fresh.getId())).containsEntry("status", "CANCELLING"); + assertThat(store.getTurn(EDITOR, fresh.getId()).getCancelRequestedAt()).isNotNull(); + store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", fresh.getId()).withValue("cancelRequestedAt", Instant.now().minusSeconds(5))); + assertThat(active()).isEmpty(); + assertThat(store.getTurn(EDITOR, fresh.getId()).getStatus()).isEqualTo("CANCELLED"); + AssistantTurn stopped = remote(Instant.now().minusSeconds(5)); + assertThat(api.cancelTurn(stopped.getId())).containsEntry("status", "CANCELLED"); + assertThat(api.cancelTurn(stopped.getId())).containsEntry("status", "CANCELLED"); + provider.thenText("Recovered"); + assertThat(completed(api.startTurn(conversation(), Map.of("text", "next"))).getStatus()).isEqualTo("COMPLETE"); + } + + + + /******************************************************************************* + ** Two replicas reaping the same row produce only one terminal frame. + *******************************************************************************/ + @Test + void simultaneousReapers() throws Exception + { + AssistantTurn stale = remote(Instant.now().minusSeconds(5)); + AssistantRuntime other = new AssistantRuntime(qInstance, config); + CountDownLatch start = new CountDownLatch(1); + try(var executor = java.util.concurrent.Executors.newVirtualThreadPerTaskExecutor()) + { + List> futures = new ArrayList<>(); + for(AssistantRuntime runtime : List.of(config.getRuntime(), other)) + { + futures.add(executor.submit(() -> + { + start.await(); + QContext.init(qInstance, session); + try + { + new AssistantApi(runtime, session).activeTurns(); + } + finally + { + QContext.clear(); + } + return (null); + })); + } + start.countDown(); + for(var future : futures) + { + future.get(5, java.util.concurrent.TimeUnit.SECONDS); + } + List events = api.turnEvents(stale.getId(), 0).persisted(); + assertThat(events).filteredOn(TurnEvent::isTerminal).hasSize(1); + assertThat(store.getTurn(EDITOR, stale.getId()).getStatus()).isEqualTo("ERROR"); + } + finally + { + other.close(); + } + } + + + + /******************************************************************************* + ** A terminal outcome written elsewhere wins over the local engine's finish. + *******************************************************************************/ + @Test + void remoteTerminalOutcomeIsPreserved() throws Exception + { + ControlledProvider controlled = controlled(false); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + controlled.awaitCalls(1); + store.insertEvent(run.getTurnId(), EDITOR, 50, TurnEvent.MESSAGE_END, + new TurnEvent(50, TurnEvent.MESSAGE_END, Map.of("status", "ERROR", "stopReason", "interrupted")).toJson()); + store.updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", run.getTurnId()).withValue("status", "ERROR") + .withValue("phase", "done").withValue("errorCode", "INTERNAL").withValue("stopReason", "interrupted").withValue("endedAt", Instant.now())); + controlled.release(); + assertThat(completed(run).getStatus()).isEqualTo("ERROR"); + assertThat(api.turnEvents(run.getTurnId(), 0).persisted()).filteredOn(TurnEvent::isTerminal).hasSize(1); + assertThat(run.getLog().after(0, 1)).last().extracting(e -> e.getData().get("status")).isEqualTo("ERROR"); + } + + + + /******************************************************************************* + ** Configuration rejects invalid lease timing and concurrency modes. + *******************************************************************************/ + @Test + void configurationValidation() + { + config.withTurnHeartbeatSeconds(5).withTurnLeaseSeconds(4).withDefaultConcurrencyMode("invalid"); + List errors = new ArrayList<>(); + config.validate(qInstance, errors); + assertThat(errors).contains("turnHeartbeatSeconds must be less than turnLeaseSeconds.", "defaultConcurrencyMode must be SUPERSEDE or CONCURRENT."); + config.withRunTracker(null); + assertThat(config.getRunTracker()).isSameAs(AssistantRunTracker.NOOP); + } + + + + /******************************************************************************* + ** A different runtime writes Stop; the owner sees it within one heartbeat. + *******************************************************************************/ + @Test + void crossPodStopAndHeartbeat() throws Exception + { + ControlledProvider controlled = controlled(false); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + controlled.awaitCalls(1); + AssistantRuntime remote = new AssistantRuntime(qInstance, config); + try + { + assertThat(remote.getInstanceId()).isNotEqualTo(config.getRuntime().getInstanceId()); + AssistantApi otherPod = new AssistantApi(remote, session); + assertThat(((List) otherPod.activeTurns().get("turns"))).hasSize(1); + assertThat(active().getFirst()).containsEntry("live", true).containsEntry("phase", "writing").containsKey("elapsedMs"); + assertThat(otherPod.cancelTurn(run.getTurnId())).containsEntry("status", "CANCELLING"); + AssistantTurn row = completed(run); + assertThat(row.getStatus()).isEqualTo("CANCELLED"); + assertThat(row.getCancelReason()).isEqualTo("USER"); + assertThat(row.getHeartbeatAt()).isAfter(row.getStartedAt()); + assertThat(otherPod.cancelTurn(run.getTurnId())).containsEntry("status", "CANCELLED"); + } + finally + { + controlled.release(); + remote.close(); + } + } + + + + /******************************************************************************* + ** Shutdown force-finalizes a wedged engine and refuses new inserts with 503. + *******************************************************************************/ + @Test + void rollingDeploy() throws Exception + { + ControlledProvider controlled = controlled(true); + RecordingTracker tracker = new RecordingTracker(); + config.setRunTracker(tracker); + Integer conversation = conversation(); + RunningTurn run = api.startTurn(conversation, Map.of("text", "wait")); + try + { + controlled.awaitCalls(1); + config.getRuntime().shutdown(Duration.ZERO); + QContext.init(qInstance, session); + AssistantTurn row = completed(run); + assertThat(row.getStatus()).isEqualTo("ERROR"); + assertThat(row.getErrorCode()).isEqualTo("INTERRUPTED"); + assertThat(store.getMessage(EDITOR, row.getAssistantMessageId()).getPartsJson()).contains(TurnEngine.RESTART_TEXT); + assertThat(tracker.finishes.get()).isEqualTo(1); + int messages = store.listMessages(EDITOR, conversation, null, 100).size(); + assertThatThrownBy(() -> api.startTurn(conversation, Map.of("text", "retry"))) + .isInstanceOfSatisfying(AssistantApiException.class, e -> + { + assertThat(e.getStatus()).isEqualTo(503); + assertThat(e.getCode()).isEqualTo("SHUTTING_DOWN"); + }); + assertThat(store.listMessages(EDITOR, conversation, null, 100)).hasSize(messages); + } + finally + { + controlled.release(); + } + } + + + + /******************************************************************************* + ** Graceful cancellation also uses the restart outcome, not a user cancel. + *******************************************************************************/ + @Test + void gracefulShutdown() throws Exception + { + ControlledProvider controlled = controlled(false); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + controlled.awaitCalls(1); + config.getRuntime().shutdown(Duration.ofSeconds(1)); + QContext.init(qInstance, session); + assertThat(completed(run).getErrorCode()).isEqualTo("INTERRUPTED"); + } + + + + /******************************************************************************* + ** Tracker progress is phase-based; its Stop callback cancels the local run. + *******************************************************************************/ + @Test + void trackerAndTrackerFailures() throws Exception + { + ControlledProvider controlled = controlled(false); + RecordingTracker tracker = new RecordingTracker(); + config.setRunTracker(tracker); + RunningTurn run = api.startTurn(conversation(), Map.of("text", "wait")); + controlled.awaitCalls(1); + await(() -> tracker.cancel != null); + tracker.cancel.run(); + assertThat(completed(run).getStatus()).isEqualTo("CANCELLED"); + assertThat(tracker.starts.get()).isEqualTo(1); + assertThat(tracker.phases).containsExactly("thinking", "writing"); + assertThat(tracker.finishes.get()).isEqualTo(1); + assertThat(tracker.label).isEqualTo("Assistant answer"); + tracker.fail = true; + config.withModelProvider(provider); + config.getRuntime().getProviderService().invalidate(); + provider.thenText("Still works"); + RunningTurn throwing = api.startTurn(conversation(), Map.of("text", "hello")); + assertThat(completed(throwing).getStatus()).isEqualTo("COMPLETE"); + assertThat(tracker.finishes.get()).isEqualTo(2); + } + + + + /******************************************************************************* + ** Build a remote persisted turn with an assistant message to recover. + *******************************************************************************/ + private AssistantTurn remote(Instant heartbeat) throws Exception + { + Integer conversation = conversation(); + AssistantMessage message = store.insert(new AssistantMessage().withConversationId(conversation).withUserId(EDITOR) + .withRole("ASSISTANT").withStatus("STREAMING").withSeq(1).withPartsJson("[]")); + return (store.insert(new AssistantTurn().withConversationId(conversation).withUserId(EDITOR).withAssistantMessageId(message.getId()) + .withStatus("RUNNING").withStartedAt(heartbeat).withHeartbeatAt(heartbeat).withLeaseOwner("dead-pod").withPhase("thinking"))); + } + + + + /******************************************************************************* + ** Create a user-owned conversation. + *******************************************************************************/ + private Integer conversation() throws Exception + { + return ((Integer) api.createConversation(Map.of()).get("id")); + } + + + + /******************************************************************************* + ** Read the active-run response. + *******************************************************************************/ + @SuppressWarnings("unchecked") + private List> active() throws Exception + { + return ((List>) api.activeTurns().get("turns")); + } + + + + /******************************************************************************* + ** Wait for both persisted completion and local cleanup, then check replay. + *******************************************************************************/ + private AssistantTurn completed(RunningTurn run) throws Exception + { + await(() -> config.getRuntime().getRunningTurn(run.getTurnId()) == null); + assertThat(run.getLog().isClosed()).isTrue(); + assertThat(run.getLog().after(0, 1)).filteredOn(TurnEvent::isTerminal).hasSize(1); + AssistantTurn row = store.getTurn(EDITOR, run.getTurnId()); + assertThat(row.getStatus()).isNotIn("RUNNING", "CANCELLING"); + assertThat(api.turnEvents(run.getTurnId(), 0).persisted()).last().extracting(TurnEvent::getType).isEqualTo(TurnEvent.MESSAGE_END); + return (row); + } + + + + /******************************************************************************* + ** Install a deterministic provider before the first call resolves it. + *******************************************************************************/ + private ControlledProvider controlled(boolean stubborn) + { + ControlledProvider result = new ControlledProvider(stubborn); + config.withModelProvider(result); + return (result); + } + + + + /******************************************************************************* + ** Short polling with a hard bound, without timing-sensitive fixed waits. + *******************************************************************************/ + private static void await(BooleanSupplier condition) throws InterruptedException + { + long until = System.nanoTime() + Duration.ofSeconds(6).toNanos(); + while(!condition.getAsBoolean() && System.nanoTime() < until) + { + Thread.sleep(5); + } + assertThat(condition.getAsBoolean()).as("condition reached within six seconds").isTrue(); + } + + + + /******************************************************************************* + ** MemoryRecordStore iterates its synchronized maps without holding their + ** lock. Serialize individual backend calls, as in RobustnessTest, so the + ** fixture models atomic database operations while turns still overlap. + *******************************************************************************/ + private static class AtomicCallsStore extends AssistantStore + { + private static final Object BACKEND_LOCK = new Object(); + + /******************************************************************************* + ** Preserve encryption and other configured persistence behavior. + *******************************************************************************/ + AtomicCallsStore(AssistantQBitConfig config) + { + super(config); + } + + /******************************************************************************* + ** Make a memory query atomic with respect to insert/update operations. + *******************************************************************************/ + @Override + public List query(String table, QQueryFilter filter, Class type) throws QException + { + synchronized(BACKEND_LOCK) + { + return (super.query(table, filter, type)); + } + } + + /******************************************************************************* + ** Make an insert atomic with respect to memory queries. + *******************************************************************************/ + @Override + public T insert(T entity) throws QException + { + synchronized(BACKEND_LOCK) + { + return (super.insert(entity)); + } + } + + /******************************************************************************* + ** Make an update atomic with respect to memory queries. + *******************************************************************************/ + @Override + public void updateFields(String table, QRecord record) throws QException + { + synchronized(BACKEND_LOCK) + { + super.updateFields(table, record); + } + } + } + + + + /******************************************************************************* + ** Streams partial text, then blocks until explicitly released. + *******************************************************************************/ + private static class ControlledProvider extends ScriptedModelProvider + { + private final boolean stubborn; + private final List handles = new CopyOnWriteArrayList<>(); + private boolean emitText = true; + + /******************************************************************************* + ** Configure whether cancel and interruption are deliberately ignored. + *******************************************************************************/ + ControlledProvider(boolean stubborn) + { + this.stubborn = stubborn; + } + + /******************************************************************************* + ** Start a controlled stream. + *******************************************************************************/ + @Override + public StreamHandle stream(ModelRequest request, ModelEventListener listener) + { + ControlledHandle handle = new ControlledHandle(stubborn); + if(emitText) + { + listener.onTextDelta("A partial answer that has already begun streaming to the user. "); + } + handles.add(handle); + return (handle); + } + + /******************************************************************************* + ** Await a known number of started calls. + *******************************************************************************/ + void awaitCalls(int count) throws InterruptedException + { + await(() -> handles.size() >= count); + } + + /******************************************************************************* + ** Always release test threads, including deliberately wedged ones. + *******************************************************************************/ + void release() + { + handles.forEach(h -> h.release.countDown()); + } + } + + + + /******************************************************************************* + ** A model that can deliberately violate the provider timeout contract. + *******************************************************************************/ + private static class ControlledHandle implements StreamHandle + { + private final boolean stubborn; + private final CountDownLatch release = new CountDownLatch(1); + private volatile boolean cancelled; + private volatile boolean returned; + + /******************************************************************************* + ** Configure cancellation behavior. + *******************************************************************************/ + ControlledHandle(boolean stubborn) + { + this.stubborn = stubborn; + } + + /******************************************************************************* + ** Abort unless this is the wedged-provider test. + *******************************************************************************/ + @Override + public void cancel() + { + cancelled = true; + if(!stubborn) + { + release.countDown(); + } + } + + /******************************************************************************* + ** Observe cancellation. + *******************************************************************************/ + @Override + public boolean isCancelled() + { + return (cancelled); + } + + /******************************************************************************* + ** Block deterministically; stubborn mode ignores timeouts and interrupts. + *******************************************************************************/ + @Override + public ModelResponse await(Duration timeout) throws InterruptedException + { + while(release.getCount() != 0) + { + try + { + release.await(); + } + catch(InterruptedException e) + { + if(!stubborn) + { + throw (e); + } + } + } + returned = true; + return (new ModelResponse().withParts(List.of(ModelPart.text("Done"))) + .withStopReason(cancelled ? ModelResponse.STOP_CANCELLED : ModelResponse.STOP_END_TURN)); + } + } + + + + /******************************************************************************* + ** Records the host job-tray contract and can throw from every callback. + *******************************************************************************/ + private static class RecordingTracker implements AssistantRunTracker + { + private final AtomicInteger starts = new AtomicInteger(); + private final AtomicInteger finishes = new AtomicInteger(); + private final List phases = new CopyOnWriteArrayList<>(); + private volatile Runnable cancel; + private String label; + private boolean fail; + + /******************************************************************************* + ** Record publication. + *******************************************************************************/ + @Override + public String start(String userId, Integer turnId, Integer conversationId, String label) + { + starts.incrementAndGet(); + this.label = label; + if(fail) + { + throw (new IllegalStateException("tracker start failed")); + } + return ("job-" + turnId); + } + + /******************************************************************************* + ** Record a phase transition. + *******************************************************************************/ + @Override + public void progress(String trackingId, String phase, String message) + { + phases.add(phase); + if(fail) + { + throw (new IllegalStateException("tracker progress failed")); + } + } + + /******************************************************************************* + ** Record terminal notification exactly once. + *******************************************************************************/ + @Override + public void finish(String trackingId, String status, String message) + { + finishes.incrementAndGet(); + if(fail) + { + throw (new IllegalStateException("tracker finish failed")); + } + } + + /******************************************************************************* + ** Retain the host's Stop callback. + *******************************************************************************/ + @Override + public void onCancel(String trackingId, Runnable cancel) + { + this.cancel = cancel; + } + } +} diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/EngineComponentsTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/EngineComponentsTest.java index 45a3028..ef995c4 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/EngineComponentsTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/EngineComponentsTest.java @@ -151,8 +151,7 @@ class EngineComponentsTest extends BaseAssistantTest RateLimiter limiter = new RateLimiter(clock); assertThat(limiter.tryAcquire("u", 2, 1)).isNull(); - assertThat(limiter.tryAcquire("u", 2, 1)).contains("in progress"); - limiter.release("u"); + assertThat(limiter.tryAcquire("u", 2, 1)).isNull(); limiter.release("u"); assertThat(limiter.tryAcquire("u", 2, 1)).contains("last minute"); diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/RobustnessTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/RobustnessTest.java index b6dc9e7..961422f 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/RobustnessTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/engine/RobustnessTest.java @@ -119,7 +119,7 @@ class RobustnessTest extends BaseAssistantTest // idempotent: a second look adds nothing // //////////////////////////////////////////////////////// assertThat(api.turnEvents(turn.getId(), 0).persisted()).hasSize(4); - assertThatThrownBy(() -> api.cancelTurn(turn.getId())).isInstanceOfSatisfying(AssistantApiException.class, e -> assertThat(e.getStatus()).isEqualTo(409)); + assertThat(api.cancelTurn(turn.getId())).containsEntry("status", "ERROR"); assertThat(api.turnEvents(young.getId(), 0).finished()).isFalse(); assertThatThrownBy(() -> api.deleteConversation(conversation.getId())).isInstanceOfSatisfying(AssistantApiException.class, e -> assertThat(e.getStatus()).isEqualTo(409)); diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/testing/AssistantTestFixture.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/testing/AssistantTestFixture.java index 1b36b2d..2f9cfe2 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/testing/AssistantTestFixture.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/testing/AssistantTestFixture.java @@ -99,6 +99,7 @@ public class AssistantTestFixture *******************************************************************************/ public static QInstance build(AssistantQBitConfig config) throws Exception { + config.setRegisterShutdownHook(false); MemoryRecordStore.fullReset(); QContext.clear(); -- 2.54.0 From c751fc055cc7e7ebfbec99328d25684b31ae49a4 Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:43 -0500 Subject: [PATCH 5/5] docs: durable turns upgrade notes and changelog Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 20 ++++++++++++++++++++ UPGRADE.md | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbb3d2f..9f769a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # Changelog +## Unreleased - Durable turns + +- Renew user-owned run leases with heartbeats, poll cross-replica Stop, enforce + hard deadlines, and finalize active runs during graceful shutdown. Recover + abandoned runs on replay, active discovery, or the next question. +- Default to superseding a prior answer in the same conversation. Accept explicit + concurrent questions; the per-user cap supersedes the oldest run instead of + refusing another question. Per-minute limits and budgets still apply. +- Add `GET /qqq/v1/assistant/turns/active`, `turn.phase` progress and `startedAt` + on `message.start`. Cancel is idempotent and returns HTTP 200. +- Add the optional `AssistantRunTracker` Jobs tray SPI (no-op by default), + pending QQQ core job tracking (#1210). +- Add the durable-turn schema upgrade; apply it before deploying this runtime. + A partial unique index fences duplicate terminal events from concurrent reapers. +- UI API extension `assistant` 1.2.0 (additive): `turns/active`, the optional + `concurrency` on a new turn, 503 `SHUTTING_DOWN`, `turn.phase`, + `message.start.startedAt`, cancelReason `SUPERSEDED`. The 429 response stays + documented for budgets and per-minute limits; the "answer already in progress" + refusal is no longer returned. Cancel keeps 202 documented for older replicas. + ## Unreleased - X1/Y1 - Add default-off `withRecordHistory(true)` host opt-in: remove only `audit*` on RC.6, or invoke core #1193 audit-read/write-deny support when available. Preserve user permissions and all other channel denials; kof22 does not opt in. diff --git a/UPGRADE.md b/UPGRADE.md index c68d1ef..e04d64b 100644 --- a/UPGRADE.md +++ b/UPGRADE.md @@ -1,5 +1,45 @@ # Upgrading qbit-assistant +## Durable turns (Unreleased) + +Apply the **Durable turns** block at the end of `docs/schema-postgres.sql` before +starting the new runtime. It adds nullable `heartbeat_at`, `lease_owner`, +`cancel_requested_at`, `phase`, `concurrency_mode` and the `(user_id, status)` +index. A partial unique index on terminal events also prevents two replicas from +inserting separate `message.end` frames for the same turn. Check for historical +duplicate terminal frames before applying it (resolve any duplicates through your +normal migration process). Fresh installs run the full schema, including ALTERs. +Existing rows with no heartbeat use `started_at` for lease expiry. During the +first upgrade from a runtime without heartbeats, drain old replicas' turns first: +a new replica cannot distinguish an old live turn from a dead owner after 45 s. + +Defaults: heartbeat 5 s, lease 45 s, hard duration 300 s, force-finalization grace +20 s, shutdown grace 8 s. Configure these through the corresponding +`AssistantQBitConfig` withers. Heartbeat must be shorter than the lease. The JVM +shutdown hook is enabled by default; tests disable `registerShutdownHook` and +managed hosts may call `runtime.shutdown(Duration)` themselves. Give the process +termination budget enough time for the configured shutdown grace and persistence. + +`POST /conversations/{id}/turns` accepts `concurrency: "supersede" | "concurrent"` +(case-insensitive). Default `SUPERSEDE` stops prior runs in that conversation; +`CONCURRENT` permits them. The default per-user cap is now 3; reaching it stops +the oldest run instead of returning the old in-progress 429. Per-minute and +budget 429 responses still apply. Superseded turns record `cancelReason=SUPERSEDED`. + +Cancel now returns HTTP **200**, including for terminal turns, with the current +`status`. The OpenAPI fragment retains the old 202 response for mixed-version +rollouts. Active runs are discoverable through `GET /qqq/v1/assistant/turns/active`; +reattach through the existing events route. Consumers should recognize +`turn.phase` and `message.start.startedAt`. Starts during shutdown return +503 `SHUTTING_DOWN`; clients may retry against a ready replica. + +Hosts may wire `AssistantRunTracker` for Jobs tray publication, progress, Stop, +and terminal notification. It is a no-op by default and all callbacks are +isolated from turn execution. Tracking ids are local to the owning runtime; +reaping an already dead owner does not invent another tracking job. Hosts should +associate their durable tracking row with the supplied turn id. The SPI is a +stand-in for QQQ #1210. + ## X1/Y1: channel wall and audit attribution Requires QQQ 4.1.0-RC.6. Every assistant table, process, report and widget is -- 2.54.0