From 00861af42c06eafe23c9b5b44f5b63a07adc6ea8 Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 16:36:39 -0500 Subject: [PATCH 1/6] 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/6] 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/6] 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/6] 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/6] 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 From 9c3af04a9152d6ccba142312bf17e5610e216bfc Mon Sep 17 00:00:00 2001 From: James Maes Date: Wed, 14 Oct 2026 17:08:13 -0500 Subject: [PATCH 6/6] fix(defaults): empty rollbacks so a rollback never deletes host data A seeded row cannot be told apart from a host's own row with the same name ('default' profile, policy, starter skill names), and Liquibase also rolls back changesets it marked ran, so the name-based rollback lines could delete a host's data. Every seed changeset now rolls back to nothing; the test runs a real Liquibase rollback over host-owned rows and checks they survive. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 ++ .../assistant-defaults.changelog.sql | 23 ++++++++------- .../AssistantDefaultsChangelogTest.java | 29 ++++++++++++++++++- 3 files changed, 43 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dad22bd..9a64446 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased - Starter defaults +- Every seed changeset's rollback is empty: a seeded row cannot be told apart from a host's own row with the same name, and Liquibase also rolls back changesets it marked ran, so a name-based undo would delete host data. + - Fix the defaults changelog: every precondition is a portable `sql-check` (the formatted-SQL `table-exists` form failed to parse); the policy, budget and prices are seeded only into empty tables and the profile only when no default exists, so an existing host never changes on deploy. The test now runs the changelog through Liquibase 4.33, including an already configured host and a missing table. - The Setup card prices the default model the way calls are costed (registry, price rows, configured prices); "Add rows for every table" reads fresh data-access rows. diff --git a/qbit-assistant/src/main/resources/db/liquibase/qbit-assistant/assistant-defaults.changelog.sql b/qbit-assistant/src/main/resources/db/liquibase/qbit-assistant/assistant-defaults.changelog.sql index b0320f0..f72d4bb 100644 --- a/qbit-assistant/src/main/resources/db/liquibase/qbit-assistant/assistant-defaults.changelog.sql +++ b/qbit-assistant/src/main/resources/db/liquibase/qbit-assistant/assistant-defaults.changelog.sql @@ -9,6 +9,9 @@ -- Rules for this file: -- * insert-only: every row is inserted only when its natural key is absent, -- so an administrator's edits always survive and nothing is ever updated; +-- * rollback is empty: a seeded row cannot be told apart from a host's own +-- row with the same name, and Liquibase also rolls back changesets it marked +-- ran, so undoing would delete the host's data (remove seeded rows by hand); -- * no secrets: environment variable NAMES only, never values; -- * portable SQL (PostgreSQL and H2): INSERT ... SELECT ... WHERE NOT EXISTS; -- * every changeset has a precondition on its table, so a host that has not @@ -22,7 +25,7 @@ INSERT INTO assistant_policy (name, kill_switch, writes_require_confirmation, redact_pii, max_rows_per_call, create_date, modify_date) SELECT 'default', FALSE, TRUE, TRUE, 50, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_policy); --- rollback DELETE FROM assistant_policy WHERE name = 'default'; +-- rollback empty -- changeset qbit-assistant:defaults-2 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -37,7 +40,7 @@ Which tables can I look at? What changed in the last 7 days?', TRUE, TRUE, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_profile WHERE name = 'default' OR is_default = TRUE); --- rollback DELETE FROM assistant_profile WHERE name = 'default'; +-- rollback empty -- changeset qbit-assistant:defaults-3 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -53,7 +56,7 @@ SELECT 'find-and-count', 'Finding, filtering and counting records.', 5. If the result was truncated, say so and offer the link instead of listing more.', TRUE, 'find,show,list,which,how many,count,filter,search,look up', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_skill WHERE name = 'find-and-count'); --- rollback DELETE FROM assistant_skill WHERE name = 'find-and-count'; +-- rollback empty -- changeset qbit-assistant:defaults-4 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -68,7 +71,7 @@ SELECT 'explain-a-record', 'Summarizing a record and who changed it.', 4. Link to the record.', TRUE, 'explain,summarize,summary,tell me about,what is,status,history,changed,who changed,audit', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_skill WHERE name = 'explain-a-record'); --- rollback DELETE FROM assistant_skill WHERE name = 'explain-a-record'; +-- rollback empty -- changeset qbit-assistant:defaults-5 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -83,7 +86,7 @@ SELECT 'safe-changes', 'Proposing creates, updates and deletes safely.', 4. For many rows from a file, use the bulk load tools and always run bulk_load_dry_run and show its summary before bulk_load_run.', TRUE, 'change,update,set,edit,create,add,new,delete,remove,fix,rename,import,load', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_skill WHERE name = 'safe-changes'); --- rollback DELETE FROM assistant_skill WHERE name = 'safe-changes'; +-- rollback empty -- changeset qbit-assistant:defaults-6 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -97,7 +100,7 @@ SELECT 'access-questions', 'Explaining why data is missing or not available.', 3. Suggest the next step: check the filter, open the list screen, or ask an administrator for access.', TRUE, 'cannot see,can''t see,missing,not found,denied,access,permission,not allowed,hidden', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_skill WHERE name = 'access-questions'); --- rollback DELETE FROM assistant_skill WHERE name = 'access-questions'; +-- rollback empty -- changeset qbit-assistant:defaults-7 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -109,7 +112,7 @@ FROM assistant_profile p, assistant_skill s WHERE p.name = 'default' AND s.name IN ('find-and-count', 'explain-a-record', 'safe-changes', 'access-questions') AND NOT EXISTS (SELECT 1 FROM assistant_profile_skill ps WHERE ps.profile_id = p.id AND ps.skill_id = s.id); --- rollback DELETE FROM assistant_profile_skill WHERE profile_id IN (SELECT id FROM assistant_profile WHERE name = 'default') AND skill_id IN (SELECT id FROM assistant_skill WHERE name IN ('find-and-count', 'explain-a-record', 'safe-changes', 'access-questions')); +-- rollback empty -- changeset qbit-assistant:defaults-8 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -120,7 +123,7 @@ SELECT 'Monthly spend (whole instance)', TRUE, 'INSTANCE', 'MONTH', 'COST_USD', 'Starter budget. Warns everyone at 80% of $50 a month, switches to the cheapest allowed model at 100% (or stops when there is none), and stops at 120%. Set Downgrade Model to choose the fallback; add an override to raise it temporarily.', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP WHERE NOT EXISTS (SELECT 1 FROM assistant_budget); --- rollback DELETE FROM assistant_budget WHERE name = 'Monthly spend (whole instance)'; +-- rollback empty -- changeset qbit-assistant:defaults-9 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -143,7 +146,7 @@ FROM (VALUES ('anthropic', 'claude-haiku-4-5', 1.00, 5.00, 0.10, 1.25, 2.00, 'Anthropic list price, 2026-10') ) AS v (provider_type, model, input_price, output_price, cache_read, cache_write5m, cache_write1h, notes) WHERE NOT EXISTS (SELECT 1 FROM assistant_model_price p); --- rollback SELECT 1; -- prices may have been edited; keep them +-- rollback empty -- changeset qbit-assistant:defaults-10 -- preconditions onFail:MARK_RAN onError:MARK_RAN @@ -156,4 +159,4 @@ SELECT 'Example: weekly change summary', FALSE, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP FROM (SELECT MIN(id) AS id FROM assistant_profile WHERE name = 'default') p WHERE NOT EXISTS (SELECT 1 FROM assistant_scheduled_task WHERE name = 'Example: weekly change summary'); --- rollback DELETE FROM assistant_scheduled_task WHERE name = 'Example: weekly change summary' AND enabled = FALSE; +-- rollback empty diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantDefaultsChangelogTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantDefaultsChangelogTest.java index e66df40..081c77e 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantDefaultsChangelogTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantDefaultsChangelogTest.java @@ -144,6 +144,33 @@ class AssistantDefaultsChangelogTest extends BaseAssistantTest + /******************************************************************************* + ** Rolling the changelog back never deletes data: not the host's own rows + ** (changesets marked ran included), not the seeded ones. + *******************************************************************************/ + @Test + void rollbackKeepsEveryRow() throws Exception + { + try(Connection connection = freshSchema(); Statement statement = connection.createStatement()) + { + statement.execute("INSERT INTO assistant_policy (name, kill_switch) VALUES ('default', FALSE)"); + statement.execute("INSERT INTO assistant_profile (name, display_name, is_default, enabled) VALUES ('default', 'Hal', TRUE, TRUE)"); + statement.execute("INSERT INTO assistant_skill (name, instructions) VALUES ('find-and-count', 'The host wrote this.')"); + runChangelog(connection); + Map before = counts(connection); + + Database database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(new JdbcConnection(connection)); + new Liquibase(CHANGELOG.substring(1), new ClassLoaderResourceAccessor(), database).rollback(10, (String) null); + + assertThat(ranChangesets(connection)).isEmpty(); + assertThat(counts(connection)).isEqualTo(before); + assertThat(scalar(connection, "SELECT display_name FROM assistant_profile WHERE name = 'default'")).isEqualTo("Hal"); + assertThat(scalar(connection, "SELECT instructions FROM assistant_skill WHERE name = 'find-and-count'")).isEqualTo("The host wrote this."); + } + } + + + /******************************************************************************* ** A host without one of the tables: that changeset is marked ran, the ** rest apply, nothing fails. @@ -174,7 +201,7 @@ class AssistantDefaultsChangelogTest extends BaseAssistantTest for(String changeset : changesets) { assertThat(changeset).contains("-- preconditions onFail:MARK_RAN onError:MARK_RAN").contains("-- precondition-sql-check expectedResult:0 ") - .doesNotContain("precondition-table-exists").contains("-- rollback ").contains("NOT EXISTS"); + .doesNotContain("precondition-table-exists").contains("-- rollback empty").contains("NOT EXISTS"); } for(String table : SEEDED_TABLES) { -- 2.54.0