Compare commits
57
Commits
develop
...
v0.1.0-RC.26
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
feab4caf17 | ||
|
|
ba21b72188
|
||
|
|
0a65702dab
|
||
|
|
739ca6981f | ||
|
|
b51dd4dad1
|
||
|
|
c3a03b8d84
|
||
|
|
2d25f0b769
|
||
|
|
c5125f9788
|
||
|
|
3de664e0fd | ||
|
|
2caa634044 | ||
|
|
c77e5b0123
|
||
|
|
f46735b1c0 | ||
|
|
86563c64d2 | ||
|
|
82febe0f19 | ||
|
|
1ea23dfc49
|
||
|
|
9d31de46a6 | ||
|
|
24b0b3b14a
|
||
|
|
1c33e97664 | ||
|
|
fcaf2ded7e
|
||
|
|
2cea8ffd44 | ||
|
|
6b98d5216b
|
||
|
|
6adf390bf8 | ||
|
|
bc92c89665
|
||
|
|
d43402e262 | ||
|
|
fe725ac8ff
|
||
|
|
dfb05d8935 | ||
|
|
523fd0128a
|
||
|
|
47115fe359 | ||
|
|
f2d2b50459
|
||
|
|
8c33db9fcf | ||
|
|
75a5dedf7f
|
||
|
|
c1cceaf4c3 | ||
|
|
f8a114bbb8
|
||
|
|
3a7ff1c659 | ||
|
|
ea81166454
|
||
|
|
cfcc044055 | ||
|
|
620dee488d
|
||
|
|
e23206f1e8 | ||
|
|
c6bae3bcf2
|
||
|
|
2d8459bb24 | ||
|
|
ac06a96787
|
||
|
|
d4a1f64480 | ||
|
|
c28e782075
|
||
|
|
5c76a13245 | ||
|
|
8d1348888d
|
||
|
|
47cc572c36 | ||
|
|
8099f64bf2
|
||
|
|
ac6ec75deb | ||
|
|
1f48a370e3
|
||
|
|
d0ad2aa0fa | ||
|
|
16c8500f59
|
||
|
|
a8c5f586ee | ||
|
|
3dfa775b21
|
||
|
|
67f78fd660 | ||
|
|
ff67c6f7cc | ||
|
|
19551ac2af
|
||
|
|
7263ac627d
|
No files matched your search
@@ -1,5 +1,11 @@
|
||||
# Changelog
|
||||
|
||||
## Unreleased - A turn always ends
|
||||
|
||||
- Fixed an answer that stayed on "Thinking..." forever (seen after the tool steps had finished) when the host database stopped answering, for example a connection pool with every connection checked out. The watchdog's heartbeat did its database read and write while holding the engine's monitor and before it checked the deadline, so one blocked heartbeat stopped the deadline from ever firing and stopped `finish` from running. The deadline is now checked first on every tick without touching the database; the heartbeat, cancel poll and kill-switch poll run on a separate bounded worker that is reported (WARN), interrupted and abandoned if it does not return within two heartbeat intervals; and the engine's monitor is no longer held across database I/O in the heartbeat or while waiting to finalize.
|
||||
- A finalize guard ends the live stream with an `error` event and `message.end` if the terminal save has not completed within `turnForceFinalizeGraceSeconds`, so the panel never waits on the database to be told the answer failed. The terminal row is left for lease recovery.
|
||||
- New INFO log lines for every phase: `Assistant turn started`, `Assistant model call started` / `finished`, `Assistant tool call started` / `finished`, `Assistant turn finalizing`, `Assistant turn finished` (no message content). WARN lines when the deadline passes, when the watchdog finalizes a turn, when a heartbeat does not return, and when the guard ends a stream.
|
||||
|
||||
## Unreleased - Answer visible at the end of the stream
|
||||
|
||||
- The turn's redacted transcript (the question and answer parts, pending actions, stored events) is now written before `message.end` is sent. The panel reloads the conversation the moment the stream ends, and used to read the empty placeholders (`[]`) the transcript redactor stores until its flush, so the newest question and answer showed as empty bubbles until the panel was closed and reopened.
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
<!-- Multi-module build. The qqq-orb sets the revision here; the modules -->
|
||||
<!-- inherit it, and the flatten plugin resolves it in the published POMs. -->
|
||||
<properties>
|
||||
<revision>0.1.0-SNAPSHOT</revision>
|
||||
<revision>0.1.0-RC.26</revision>
|
||||
<maven.compiler.source>21</maven.compiler.source>
|
||||
<maven.compiler.target>21</maven.compiler.target>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
|
||||
+228
-32
@@ -31,12 +31,15 @@ import java.util.Objects;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.CancellationException;
|
||||
import java.util.concurrent.ConcurrentLinkedQueue;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.ExecutionException;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.Semaphore;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.TimeoutException;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicLongArray;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
|
||||
import com.kingsrook.qbits.assistant.AssistantRuntime;
|
||||
import com.kingsrook.qbits.assistant.attachments.AttachmentService;
|
||||
@@ -156,6 +159,22 @@ public class TurnEngine implements Runnable
|
||||
private ProviderService.Active budgetOrigin;
|
||||
|
||||
private volatile boolean finished = false;
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Terminal-event bookkeeping, shared with the finalize guard. The guard //
|
||||
// ends the live stream when the database will not let finish() reach //
|
||||
// message.end, so these must be usable without the engine's monitor. //
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
private final CountDownLatch finishedLatch = new CountDownLatch(1);
|
||||
private final AtomicBoolean streamEnded = new AtomicBoolean();
|
||||
private final AtomicBoolean errorEmitted = new AtomicBoolean();
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// The watchdog's database duties (heartbeat, cancel poll, kill switch) //
|
||||
// run on one bounded worker at a time, never on the watchdog thread //
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
private final AtomicReference<Thread> housekeeper = new AtomicReference<>();
|
||||
private volatile long housekeeperStartedNanos;
|
||||
private UiApiCredentials credentials;
|
||||
private volatile TranscriptRedactor transcript;
|
||||
private boolean engineRunning;
|
||||
@@ -259,6 +278,8 @@ public class TurnEngine implements Runnable
|
||||
engineRunning = true;
|
||||
attachTranscript();
|
||||
}
|
||||
LOG.info("Assistant turn started", logPair("turnId", turn.getId()), logPair("conversationId", conversation.getId()), logPair("userId", userId),
|
||||
logPair("maxTurnSeconds", config.getMaxTurnSeconds()));
|
||||
startWatchdog();
|
||||
scheduledTask = runtime.getScheduledTurn().task(conversation.getId());
|
||||
if(conversation.getScheduledTaskId() != null && scheduledTask == null)
|
||||
@@ -520,8 +541,15 @@ public class TurnEngine implements Runnable
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** 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.
|
||||
** Enforce a hard deadline even if a provider or the database is wedged, and
|
||||
** renew the lease, poll cancellation and the kill switch with the user's own
|
||||
** session.
|
||||
**
|
||||
** The deadline is checked first on every tick and touches no database: a
|
||||
** connection pool with every connection checked out blocks callers forever,
|
||||
** so anything that reads or writes the database runs on a separate, bounded
|
||||
** worker (see tickHousekeeping) and can never stop this thread from
|
||||
** cancelling and finalizing the turn.
|
||||
*******************************************************************************/
|
||||
private void startWatchdog()
|
||||
{
|
||||
@@ -540,36 +568,15 @@ public class TurnEngine implements Runnable
|
||||
{
|
||||
break;
|
||||
}
|
||||
try
|
||||
enforceDeadline();
|
||||
if(running.isFinalized())
|
||||
{
|
||||
if(System.nanoTime() - lastHeartbeat >= heartbeatMillis * 1_000_000L)
|
||||
{
|
||||
lastHeartbeat = System.nanoTime();
|
||||
heartbeat();
|
||||
}
|
||||
if(runtime.getPolicyService().isKillSwitchOn())
|
||||
{
|
||||
runtime.cancelAll("KILL_SWITCH");
|
||||
}
|
||||
break;
|
||||
}
|
||||
catch(Exception e)
|
||||
boolean heartbeatDue = System.nanoTime() - lastHeartbeat >= heartbeatMillis * 1_000_000L;
|
||||
if(tickHousekeeping(heartbeatDue, heartbeatMillis) && heartbeatDue)
|
||||
{
|
||||
LOG.warn("Assistant watchdog tick failed", hipaaSafe(e), logPair("turnId", turn.getId()));
|
||||
}
|
||||
finally
|
||||
{
|
||||
// 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());
|
||||
}
|
||||
lastHeartbeat = System.nanoTime();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -577,6 +584,10 @@ public class TurnEngine implements Runnable
|
||||
{
|
||||
Thread.currentThread().interrupt();
|
||||
}
|
||||
catch(Exception | Error e)
|
||||
{
|
||||
LOG.warn("Assistant watchdog stopped unexpectedly", hipaaSafe(e), logPair("turnId", turn.getId()));
|
||||
}
|
||||
finally
|
||||
{
|
||||
QContext.clear();
|
||||
@@ -587,9 +598,103 @@ public class TurnEngine implements Runnable
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Renew only an active row, then observe persisted cancellation requests.
|
||||
** The deadline and the force-finalize grace period. No database access.
|
||||
*******************************************************************************/
|
||||
private synchronized void heartbeat() throws QException
|
||||
private void enforceDeadline()
|
||||
{
|
||||
Instant now = runtime.getClock().instant();
|
||||
if(!now.isBefore(deadline) && !running.isCancelled())
|
||||
{
|
||||
LOG.warn("Assistant turn passed its deadline; cancelling", logPair("turnId", turn.getId()), logPair("phase", running.getPhase()),
|
||||
logPair("tool", running.getToolName()), logPair("maxTurnSeconds", config.getMaxTurnSeconds()));
|
||||
cancelWithoutBlocking("TIMEOUT");
|
||||
}
|
||||
if(!now.isBefore(deadline.plusSeconds(config.getTurnForceFinalizeGraceSeconds()))
|
||||
|| running.cancelGraceElapsed(config.getTurnForceFinalizeGraceSeconds()))
|
||||
{
|
||||
LOG.warn("Assistant watchdog is finalizing the turn", logPair("turnId", turn.getId()), logPair("phase", running.getPhase()),
|
||||
logPair("tool", running.getToolName()), logPair("cancelReason", running.getCancelReason()), logPair("modelCalls", modelCallCount),
|
||||
logPair("toolCalls", toolCallCount), logPair("graceSeconds", config.getTurnForceFinalizeGraceSeconds()));
|
||||
Stop stop = running.isCancelled() ? cancelledStop() : new Stop("ERROR", "TIMEOUT", TIMEOUT_TEXT, "timeout");
|
||||
finish(stop.status(), stop.stopReason(), stop.errorCode(), stop.message());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Cancel on a throwaway thread: cancelling aborts the provider's stream,
|
||||
** which is a call into code that may itself be wedged.
|
||||
*******************************************************************************/
|
||||
private void cancelWithoutBlocking(String reason)
|
||||
{
|
||||
Thread.ofVirtual().name("assistant-cancel-" + turn.getId()).start(() -> safely(() -> running.cancel(reason)));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Start the database duties of one watchdog tick (the heartbeat when due,
|
||||
** then the kill-switch poll) on a worker thread, unless the previous one is
|
||||
** still running. A worker that has not returned after two heartbeat
|
||||
** intervals is reported, interrupted and abandoned, so a database that
|
||||
** stopped answering is visible and later ticks can try again.
|
||||
** Returns whether a worker was started.
|
||||
*******************************************************************************/
|
||||
private boolean tickHousekeeping(boolean heartbeatDue, long heartbeatMillis)
|
||||
{
|
||||
Thread current = housekeeper.get();
|
||||
if(current != null)
|
||||
{
|
||||
long runningMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - housekeeperStartedNanos);
|
||||
if(runningMillis < 2 * heartbeatMillis)
|
||||
{
|
||||
return (false);
|
||||
}
|
||||
LOG.warn("Assistant heartbeat did not finish in time; the database may not be answering", logPair("turnId", turn.getId()),
|
||||
logPair("phase", running.getPhase()), logPair("blockedMillis", runningMillis));
|
||||
current.interrupt();
|
||||
housekeeper.compareAndSet(current, null);
|
||||
}
|
||||
|
||||
housekeeperStartedNanos = System.nanoTime();
|
||||
Thread worker = Thread.ofVirtual().name("assistant-housekeeping-" + turn.getId()).unstarted(() ->
|
||||
{
|
||||
try
|
||||
{
|
||||
QContext.init(runtime.getQInstance(), session);
|
||||
if(heartbeatDue)
|
||||
{
|
||||
heartbeat();
|
||||
}
|
||||
if(runtime.getPolicyService().isKillSwitchOn())
|
||||
{
|
||||
runtime.cancelAll("KILL_SWITCH");
|
||||
}
|
||||
}
|
||||
catch(Exception | Error e)
|
||||
{
|
||||
LOG.warn("Assistant watchdog tick failed", hipaaSafe(e), logPair("turnId", turn.getId()));
|
||||
}
|
||||
finally
|
||||
{
|
||||
QContext.clear();
|
||||
housekeeper.compareAndSet(Thread.currentThread(), null);
|
||||
}
|
||||
});
|
||||
housekeeper.set(worker);
|
||||
worker.start();
|
||||
return (true);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Renew only an active row, then observe persisted cancellation requests.
|
||||
** Not synchronized on the engine: it does database I/O, and the engine's
|
||||
** monitor must never be held across it.
|
||||
*******************************************************************************/
|
||||
private void heartbeat() throws QException
|
||||
{
|
||||
if(running.isFinalized())
|
||||
{
|
||||
@@ -838,12 +943,17 @@ public class TurnEngine implements Runnable
|
||||
}
|
||||
modelCallCount++;
|
||||
phase("thinking", null);
|
||||
LOG.info("Assistant model call started", logPair("turnId", turn.getId()), logPair("call", modelCallCount), logPair("model", model),
|
||||
logPair("messages", request.getMessages().size()));
|
||||
StreamHandle handle = provider.stream(request.copy(), fallback.listener(listener));
|
||||
running.setCurrentHandle(handle);
|
||||
try
|
||||
{
|
||||
Duration remaining = Duration.between(runtime.getClock().instant(), deadline);
|
||||
response = handle.await(remaining.isNegative() ? Duration.ofMillis(1) : remaining);
|
||||
LOG.info("Assistant model call finished", logPair("turnId", turn.getId()), logPair("call", modelCallCount), logPair("model", model),
|
||||
logPair("stopReason", response.getStopReason()), logPair("toolUses", response.getToolUses().size()),
|
||||
logPair("durationMs", TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started)));
|
||||
}
|
||||
finally
|
||||
{
|
||||
@@ -936,6 +1046,7 @@ public class TurnEngine implements Runnable
|
||||
QContext.init(runtime.getQInstance(), session);
|
||||
permits.acquire();
|
||||
startedNanos.set(index, System.nanoTime());
|
||||
LOG.info("Assistant tool call started", logPair("turnId", turn.getId()), logPair("tool", call.getName()));
|
||||
try
|
||||
{
|
||||
if(!running.isCancelled())
|
||||
@@ -984,6 +1095,9 @@ public class TurnEngine implements Runnable
|
||||
finally
|
||||
{
|
||||
permits.release();
|
||||
LOG.info("Assistant tool call finished", logPair("turnId", turn.getId()), logPair("tool", call.getName()),
|
||||
logPair("outcome", results[index] == null || results[index].getOutcome() == null ? null : results[index].getOutcome().name()),
|
||||
logPair("durationMs", TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedNanos.get(index))));
|
||||
}
|
||||
}
|
||||
catch(InterruptedException e)
|
||||
@@ -1287,12 +1401,74 @@ public class TurnEngine implements Runnable
|
||||
/*******************************************************************************
|
||||
** Persist the message, turn, conversation and usage; emit the end events.
|
||||
*******************************************************************************/
|
||||
private synchronized void finish(String status, String stopReason, String errorCode, String errorText)
|
||||
private void finish(String status, String stopReason, String errorCode, String errorText)
|
||||
{
|
||||
if(!running.tryFinalize())
|
||||
{
|
||||
return;
|
||||
}
|
||||
LOG.info("Assistant turn finalizing", logPair("turnId", turn.getId()), logPair("status", status), logPair("errorCode", errorCode),
|
||||
logPair("phase", running.getPhase()), logPair("cancelReason", running.getCancelReason()));
|
||||
startFinalizeGuard(status, stopReason, errorCode, errorText);
|
||||
synchronized(this)
|
||||
{
|
||||
finishLocked(status, stopReason, errorCode, errorText);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The user must never wait on the database to be told the turn is over.
|
||||
** finish() persists the transcript, the message and the terminal row before
|
||||
** it sends message.end (the panel reloads the conversation at once), and a
|
||||
** database that stopped answering would hold it there forever with the
|
||||
** panel on "Thinking". If finish() has not completed within the
|
||||
** force-finalize grace period, end the live stream from here, without
|
||||
** touching the database or the engine's monitor: an error event with the
|
||||
** reason, then message.end. The terminal row is left for lease recovery.
|
||||
*******************************************************************************/
|
||||
private void startFinalizeGuard(String status, String stopReason, String errorCode, String errorText)
|
||||
{
|
||||
long graceMillis = config.getTurnForceFinalizeGraceSeconds() * 1000L;
|
||||
Thread.ofVirtual().name("assistant-finalize-guard-" + turn.getId()).start(() ->
|
||||
{
|
||||
try
|
||||
{
|
||||
if(finishedLatch.await(graceMillis, TimeUnit.MILLISECONDS) || !streamEnded.compareAndSet(false, true))
|
||||
{
|
||||
return;
|
||||
}
|
||||
LOG.warn("Assistant turn could not be saved in time; ending the stream without it (database not answering?)", logPair("turnId", turn.getId()),
|
||||
logPair("status", status), logPair("errorCode", errorCode), logPair("graceSeconds", config.getTurnForceFinalizeGraceSeconds()));
|
||||
TurnEventLog log = running.getLog();
|
||||
if(errorCode != null && !errorEmitted.get())
|
||||
{
|
||||
String message = StringUtils.hasContent(errorText) ? errorText : "Something went wrong while answering. Try again.";
|
||||
log.append(TurnEvent.ERROR, map("code", errorCode, "message", message));
|
||||
}
|
||||
log.append(TurnEvent.MESSAGE_END, map("messageId", assistantMessage.getId(), "status", status, "stopReason", stopReason));
|
||||
log.close();
|
||||
}
|
||||
catch(InterruptedException e)
|
||||
{
|
||||
Thread.currentThread().interrupt();
|
||||
}
|
||||
catch(Exception | Error e)
|
||||
{
|
||||
LOG.warn("Assistant finalize guard failed", hipaaSafe(e), logPair("turnId", turn.getId()));
|
||||
running.getLog().close();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The body of finish(), under the engine's monitor.
|
||||
*******************************************************************************/
|
||||
private void finishLocked(String status, String stopReason, String errorCode, String errorText)
|
||||
{
|
||||
TranscriptRedactor previousTranscript = TranscriptRedactor.current();
|
||||
try
|
||||
{
|
||||
@@ -1421,6 +1597,7 @@ public class TurnEngine implements Runnable
|
||||
finally
|
||||
{
|
||||
finished = true;
|
||||
finishedLatch.countDown();
|
||||
running.getLog().close();
|
||||
safely(() ->
|
||||
{
|
||||
@@ -1559,6 +1736,25 @@ public class TurnEngine implements Runnable
|
||||
{
|
||||
return;
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// the finalize guard may already have ended the stream; nothing may //
|
||||
// follow its message.end, and message.end is sent exactly once //
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
if(TurnEvent.MESSAGE_END.equals(type))
|
||||
{
|
||||
if(!streamEnded.compareAndSet(false, true))
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
else if(streamEnded.get())
|
||||
{
|
||||
return;
|
||||
}
|
||||
if(TurnEvent.ERROR.equals(type))
|
||||
{
|
||||
errorEmitted.set(true);
|
||||
}
|
||||
toPersist.add(running.getLog().append(type, data));
|
||||
}
|
||||
|
||||
|
||||
+395
@@ -0,0 +1,395 @@
|
||||
/*
|
||||
* 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.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
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.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.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.logging.QCollectingLogger;
|
||||
import com.kingsrook.qqq.backend.core.logging.QLogger;
|
||||
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.AfterEach;
|
||||
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;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** A turn must always end for the user, even when the database stops answering
|
||||
** (a connection pool with every connection checked out blocks forever). These
|
||||
** tests model that with a store whose calls park until the test releases
|
||||
** them, and prove the user still gets an error and message.end.
|
||||
*******************************************************************************/
|
||||
class StalledDatabaseTurnsTest extends BaseAssistantTest
|
||||
{
|
||||
private final CountDownLatch dbGate = new CountDownLatch(1);
|
||||
private final AtomicBoolean stallHeartbeat = new AtomicBoolean();
|
||||
private final AtomicBoolean stallAllCalls = new AtomicBoolean();
|
||||
|
||||
private QSession session;
|
||||
private AssistantApi api;
|
||||
private AssistantStore store;
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Short lease, deadline and grace settings keep the tests bounded.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
protected void customizeConfig(AssistantQBitConfig config)
|
||||
{
|
||||
config.withTurnHeartbeatSeconds(1).withTurnLeaseSeconds(3).withTurnForceFinalizeGraceSeconds(1)
|
||||
.withMaxTurnSeconds(2).withKillSwitchPollMillis(50);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Install the store whose calls can stall.
|
||||
*******************************************************************************/
|
||||
@BeforeEach
|
||||
void stallingBackend()
|
||||
{
|
||||
config.getRuntime().close();
|
||||
config.setRuntime(new AssistantRuntime(qInstance, config, Clock.systemUTC(), new StallingStore(config)));
|
||||
session = AssistantTestFixture.ownerSession(EDITOR, config);
|
||||
QContext.init(qInstance, session);
|
||||
api = new AssistantApi(config.getRuntime(), session);
|
||||
store = config.getRuntime().getStore();
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Never leave a parked thread behind.
|
||||
*******************************************************************************/
|
||||
@AfterEach
|
||||
void releaseDatabase()
|
||||
{
|
||||
dbGate.countDown();
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The incident's shape, part one: the heartbeat's database write never
|
||||
** returns while a model call is wedged. The deadline must still fire and
|
||||
** the turn must end as a TIMEOUT error, with the row written.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void blockedHeartbeatDoesNotStopTheDeadline() throws Exception
|
||||
{
|
||||
stallHeartbeat.set(true);
|
||||
config.withModelProvider(new WedgedProvider());
|
||||
var logger = QLogger.activateCollectingLoggerForClass(TurnEngine.class);
|
||||
try
|
||||
{
|
||||
RunningTurn run = api.startTurn(conversation(), Map.of("text", "how many scanners are connected?"));
|
||||
await(() -> run.getLog().isClosed(), 8);
|
||||
|
||||
List<TurnEvent> events = run.getLog().after(0, 1);
|
||||
assertThat(events).filteredOn(TurnEvent::isTerminal).hasSize(1);
|
||||
assertThat(events).filteredOn(event -> TurnEvent.ERROR.equals(event.getType())).singleElement()
|
||||
.satisfies(event -> assertThat(event.getData()).containsEntry("code", "TIMEOUT"));
|
||||
await(() -> config.getRuntime().getRunningTurn(run.getTurnId()) == null, 8);
|
||||
AssistantTurn row = store.getTurn(EDITOR, run.getTurnId());
|
||||
assertThat(row.getStatus()).isEqualTo("ERROR");
|
||||
assertThat(row.getErrorCode()).isEqualTo("TIMEOUT");
|
||||
assertThat(messages(logger)).anySatisfy(message -> assertThat(message).contains("Assistant watchdog is finalizing the turn"));
|
||||
}
|
||||
finally
|
||||
{
|
||||
QLogger.deactivateCollectingLoggerForClass(TurnEngine.class);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The incident's shape, part two: tools ran ("Described Mobile Devices"),
|
||||
** then the database stopped answering for everything, so the second model
|
||||
** call, the usage write and the terminal writes all hang. The stream must
|
||||
** still end with an error and message.end instead of "Thinking" forever.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void databaseStallAfterToolsStillEndsTheStream() throws Exception
|
||||
{
|
||||
stallHeartbeat.set(true);
|
||||
provider.thenToolUse("count_records", input("table", "blog_post", "filter", filter("AND")))
|
||||
.thenRespond(request ->
|
||||
{
|
||||
stallAllCalls.set(true);
|
||||
return (new ModelResponse().withParts(List.of(ModelPart.text("There are 3 scanners."))).withStopReason(ModelResponse.STOP_END_TURN));
|
||||
});
|
||||
var logger = QLogger.activateCollectingLoggerForClass(TurnEngine.class);
|
||||
try
|
||||
{
|
||||
RunningTurn run = api.startTurn(conversation(), Map.of("text", "how many scanners are connected?"));
|
||||
await(() -> run.getLog().isClosed(), 10);
|
||||
|
||||
List<TurnEvent> events = run.getLog().after(0, 1);
|
||||
assertThat(events).extracting(TurnEvent::getType).contains(TurnEvent.TOOL_START, TurnEvent.TOOL_RESULT);
|
||||
assertThat(events).filteredOn(TurnEvent::isTerminal).hasSize(1);
|
||||
assertThat(events.getLast().isTerminal()).isTrue();
|
||||
assertThat(events).filteredOn(event -> TurnEvent.ERROR.equals(event.getType())).isNotEmpty();
|
||||
assertThat(messages(logger)).anySatisfy(message -> assertThat(message).contains("Assistant tool call finished"));
|
||||
assertThat(messages(logger)).anySatisfy(message -> assertThat(message).contains("Assistant model call started"));
|
||||
}
|
||||
finally
|
||||
{
|
||||
QLogger.deactivateCollectingLoggerForClass(TurnEngine.class);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** A healthy turn logs every phase transition at INFO and is unchanged.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void healthyTurnLogsEveryPhase() throws Exception
|
||||
{
|
||||
provider.thenToolUse("count_records", input("table", "blog_post", "filter", filter("AND"))).thenText("Three.");
|
||||
var logger = QLogger.activateCollectingLoggerForClass(TurnEngine.class);
|
||||
try
|
||||
{
|
||||
RunningTurn run = api.startTurn(conversation(), Map.of("text", "how many posts?"));
|
||||
await(() -> run.getLog().isClosed(), 8);
|
||||
await(() -> config.getRuntime().getRunningTurn(run.getTurnId()) == null, 8);
|
||||
|
||||
assertThat(store.getTurn(EDITOR, run.getTurnId()).getStatus()).isEqualTo("COMPLETE");
|
||||
List<String> messages = messages(logger);
|
||||
assertThat(messages).anySatisfy(message -> assertThat(message).contains("Assistant turn started"));
|
||||
assertThat(messages.stream().filter(message -> message.contains("Assistant model call started")).count()).isEqualTo(2);
|
||||
assertThat(messages.stream().filter(message -> message.contains("Assistant model call finished")).count()).isEqualTo(2);
|
||||
assertThat(messages).anySatisfy(message -> assertThat(message).contains("Assistant tool call started"));
|
||||
assertThat(messages).anySatisfy(message -> assertThat(message).contains("Assistant tool call finished"));
|
||||
assertThat(messages).anySatisfy(message -> assertThat(message).contains("Assistant turn finished"));
|
||||
assertThat(messages).noneSatisfy(message -> assertThat(message).contains("watchdog is finalizing"));
|
||||
}
|
||||
finally
|
||||
{
|
||||
QLogger.deactivateCollectingLoggerForClass(TurnEngine.class);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The messages the collecting logger saw.
|
||||
*******************************************************************************/
|
||||
private static List<String> messages(QCollectingLogger logger)
|
||||
{
|
||||
return (logger.getCollectedMessages().stream().map(entry -> entry.getMessage()).toList());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Create a user-owned conversation.
|
||||
*******************************************************************************/
|
||||
private Integer conversation() throws Exception
|
||||
{
|
||||
return ((Integer) api.createConversation(Map.of()).get("id"));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Poll with a hard bound.
|
||||
*******************************************************************************/
|
||||
private static void await(BooleanSupplier condition, int seconds) throws InterruptedException
|
||||
{
|
||||
long until = System.nanoTime() + Duration.ofSeconds(seconds).toNanos();
|
||||
while(!condition.getAsBoolean() && System.nanoTime() < until)
|
||||
{
|
||||
Thread.sleep(10);
|
||||
}
|
||||
assertThat(condition.getAsBoolean()).as("condition reached within " + seconds + " seconds").isTrue();
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Park like a pool checkout that never returns: interruption does not help.
|
||||
*******************************************************************************/
|
||||
private static void awaitUninterruptibly(CountDownLatch latch)
|
||||
{
|
||||
boolean interrupted = false;
|
||||
while(latch.getCount() != 0)
|
||||
{
|
||||
try
|
||||
{
|
||||
latch.await(50, TimeUnit.MILLISECONDS);
|
||||
}
|
||||
catch(InterruptedException e)
|
||||
{
|
||||
interrupted = true;
|
||||
}
|
||||
}
|
||||
if(interrupted)
|
||||
{
|
||||
Thread.currentThread().interrupt();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** A provider that ignores both its timeout and cancellation until the test ends.
|
||||
*******************************************************************************/
|
||||
private class WedgedProvider extends ScriptedModelProvider
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Start a stream that never completes.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public StreamHandle stream(ModelRequest request, ModelEventListener listener)
|
||||
{
|
||||
return (new StreamHandle()
|
||||
{
|
||||
private volatile boolean cancelled;
|
||||
|
||||
|
||||
|
||||
@Override
|
||||
public void cancel()
|
||||
{
|
||||
cancelled = true;
|
||||
}
|
||||
|
||||
|
||||
|
||||
@Override
|
||||
public boolean isCancelled()
|
||||
{
|
||||
return (cancelled);
|
||||
}
|
||||
|
||||
|
||||
|
||||
@Override
|
||||
public ModelResponse await(Duration timeout)
|
||||
{
|
||||
awaitUninterruptibly(dbGate);
|
||||
return (new ModelResponse().withParts(List.of(ModelPart.text("late"))).withStopReason(ModelResponse.STOP_CANCELLED));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The real store, except that chosen calls park until the test ends.
|
||||
*******************************************************************************/
|
||||
private class StallingStore extends AssistantStore
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Keep the configured persistence behavior.
|
||||
*******************************************************************************/
|
||||
StallingStore(AssistantQBitConfig config)
|
||||
{
|
||||
super(config);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Heartbeat renewals (and everything, once the whole database is stalled).
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public void updateFields(String table, QRecord record) throws QException
|
||||
{
|
||||
stallIf(stallAllCalls.get() || (stallHeartbeat.get() && AssistantTurn.TABLE_NAME.equals(table) && record.getValue("heartbeatAt") != null));
|
||||
super.updateFields(table, record);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Reads stall too once the whole database is stalled.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AssistantTurn getTurn(String userId, Integer id) throws QException
|
||||
{
|
||||
stallIf(stallAllCalls.get());
|
||||
return (super.getTurn(userId, id));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The terminal row write.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public boolean updateRunningTurn(String userId, QRecord changes) throws QException
|
||||
{
|
||||
stallIf(stallAllCalls.get());
|
||||
return (super.updateRunningTurn(userId, changes));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Inserts (usage rows, tool-call audit rows, events) stall too.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public <T extends QRecordEntity> T insert(T entity) throws QException
|
||||
{
|
||||
stallIf(stallAllCalls.get());
|
||||
return (super.insert(entity));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Park when asked to.
|
||||
*******************************************************************************/
|
||||
private void stallIf(boolean stall)
|
||||
{
|
||||
if(stall)
|
||||
{
|
||||
awaitUninterruptibly(dbGate);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user