feat(agent): database-counted agent concurrency and model-call rate (#4, part 2) #18

Merged
KofTwentyTwo merged 1 commits from feature/distributed-limits into develop 2026-10-14 16:51:41 +00:00
10 changed files with 291 additions and 90 deletions

No files matched your search

+1
View File
@@ -5,6 +5,7 @@ All notable changes to this project are documented here. The format follows Keep
## [0.1.0] - Unreleased
### Changed
- Make the agent limits distributed (#4): `maxConcurrentAgentRuns` counts Running `marketing_agent_run` rows in the database (a new `heartbeatAt` column, updated on each model call, lets runs whose server stopped be marked Failed first), and `maxModelCallsPerMinute` counts model-call steps in the last 60 seconds and makes a run wait rather than fail. The in-memory semaphore and rate window are gone; the settings and Marketing Policy override keep their names and defaults (the policy `maxConcurrentAgentRuns` override now applies). Hosts add `heartbeat_at` to `marketing_agent_run`. The count-then-insert window is still racy; a conditional-update claim (qqq#1178) is a follow-up.
- Move to the assistant provider SPI `0.1.0-RC.12` (package `com.kingsrook.qbits.assistant.spi`) and `qbit-build-parent` `2.1.2`, and drop the per-repository guardrail allow-lists. `qbit-assistant-provider-openai` becomes a test dependency; hosts that register it add it themselves.
- Switch the optional qbit-email integration to the published `com.kingsrook.qbits.email.api` package in `0.1.0-RC.3` (fixes #13). Preserve MARKETING messages, idempotency keys, execution links and outcome mapping; use typed readiness failures so missing providers or brand postal addresses remain send-time rejections. Upgrade `qbit-build-parent` to `2.1.1` and remove all per-repository guardrail allow-lists.
- Resolve architecture finding F5 (#9): depend on `qbit-assistant-provider-api` and `qbit-assistant-provider-openai` `0.1.0-RC.11`, with no full assistant dependency, including in tests.
+1 -1
View File
@@ -121,7 +121,7 @@ When the optional [qbit-email](https://github.com/QRun-IO/qbit-email) qbit is pr
| CRM gateway and qbit-crm adapter, explicit lead rules | |
| Persistence across restart (H2 file, and PostgreSQL in the Voyage host) | |
Scale gaps (single worker per JVM, in-memory rate limits) are listed in [`docs/OPERATIONS.md`](docs/OPERATIONS.md#scale-and-limits).
Scale gaps (single execution worker per JVM; agent limits are database-counted but still racy across replicas) are listed in [`docs/OPERATIONS.md`](docs/OPERATIONS.md#scale-and-limits).
## Build
+2 -2
View File
@@ -12,7 +12,7 @@
| `providerName` / `model` / `effort` | none | Which provider and model the agent uses. The Marketing Policy can override them. |
| `maxOutputTokens` | 8000 | Per model call. |
| `maxToolCallsPerRun` / `maxModelCallsPerRun` / `maxRunSeconds` | 24 / 16 / 300 | Limits of one agent run. |
| `maxConcurrentAgentRuns` / `maxModelCallsPerMinute` | 2 / 30 | Agent concurrency and rate (per JVM). |
| `maxConcurrentAgentRuns` / `maxModelCallsPerMinute` | 2 / 30 | Agent concurrency and rate, counted in the database so they hold across replicas. |
| `maxRowsToModel` | 50 | Rows per read tool. |
| `pricing` | Anthropic list prices (qbit-assistant) | USD per million tokens, by model id. Models without a price show no cost. Prices are never invented. |
| `registerKernelTools` | `true` | Contribute the read tools to the QQQ AI tool kernel (when QQQ has `AiToolRegistry`). |
@@ -95,7 +95,7 @@ Fix the cause. If the content changed, it needs a new approval. Then schedule ag
These are real limitations of 0.1.0. Nothing here is claimed to be production-scale.
- **Single execution worker per JVM.** The worker takes an in-process lock and claims a row by re-reading its status before marking it Running. Two replicas could both claim the same due execution in the instant between read and update, and the idempotency key is then the only protection (it fully protects the test transport, but not email). **Run Marketing's worker on one replica**, or add a database-level claim (`UPDATE ... WHERE status = SCHEDULED`, or `SELECT ... FOR UPDATE SKIP LOCKED`) before scaling out.
- **Rate limits and agent concurrency are per JVM** (in memory). The channel hourly rate is computed from the database, so it holds across replicas, but only up to the race above.
- **Agent concurrency and the model-call rate are counted in the database**, so they hold across replicas. Concurrency counts `marketing_agent_run` rows in Running (all products); a run records a heartbeat on each model call, and a Running run silent for `maxRunSeconds` plus 60 seconds is marked Failed ("Stopped responding; the server may have restarted.") before the count. The rate counts model-call steps in the last 60 seconds; at the limit a run waits (bounded by `maxRunSeconds`, and Cancel Run still works) instead of failing. The channel hourly rate is computed the same way. Counting and then inserting is still racy: two replicas can both see room, so a limit can be overshot briefly. A conditional-update claim (`ConditionalUpdateAction`, qqq#1178, QQQ 4.1.0-RC.7) will close that window. Hosts add the new `heartbeat_at` column to `marketing_agent_run` (see `docs/ddl/postgres.sql`).
- **Agent runs are synchronous** inside the record action, bounded by `maxRunSeconds`. A restart interrupts them: they are marked Failed and can be run again. There is no queue or resume.
- **The worker queries due executions page by page** (`workerBatchSize` per run) and scans channel attempts for the hourly rate. That suits hundreds of executions per hour, not bulk email. Bulk outreach is deliberately not supported.
- **Measurement is computed in memory** from a campaign's events on every view. Add a rollup table before campaigns have hundreds of thousands of events.
+1
View File
@@ -321,6 +321,7 @@ CREATE TABLE IF NOT EXISTS marketing_agent_run (
model VARCHAR(100),
requested_by_user_id VARCHAR(100),
started_at TIMESTAMP WITHOUT TIME ZONE,
heartbeat_at TIMESTAMP WITHOUT TIME ZONE,
finished_at TIMESTAMP WITHOUT TIME ZONE,
model_calls INTEGER,
tool_calls INTEGER,
+1 -1
View File
@@ -171,7 +171,7 @@ TABLES = [
("MarketingAgentRun", "marketing_agent_run", "A run of the marketing agent: task, status, provider, model usage and outputs.", [
F("productId", "Fk:MarketingProduct"), F("campaignId", "Fk:MarketingCampaign"), F("contentId", "Fk:MarketingContent"),
F("task", "Enum:AgentTask", required=True), F("status", "Enum:AgentRunStatus", required=True), F("instructions", "Text"),
F("providerName", "String:64"), F("model", "String:100"), F("requestedByUserId", "String:100"), F("startedAt", "Instant"),
F("providerName", "String:64"), F("model", "String:100"), F("requestedByUserId", "String:100"), F("startedAt", "Instant"), F("heartbeatAt", "Instant"),
F("finishedAt", "Instant"), F("modelCalls", "Integer"), F("toolCalls", "Integer"), F("inputTokens", "Long"),
F("outputTokens", "Long"), F("cacheReadTokens", "Long"), F("costMicros", "Long"), F("draftsCreated", "Integer"),
F("findingsCreated", "Integer"), F("summary", "Text"), F("outputJson", "Text"), F("errorMessage", "Text")]),
@@ -18,14 +18,11 @@ package com.kingsrook.qbits.marketing;
import java.net.InetAddress;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Semaphore;
import java.util.concurrent.locks.ReentrantLock;
import com.kingsrook.qbits.assistant.spi.AssistantModelProvider;
import com.kingsrook.qbits.marketing.crm.MarketingCrmGateway;
@@ -50,14 +47,15 @@ import com.kingsrook.qqq.backend.core.utils.StringUtils;
/*******************************************************************************
** The live parts of the Marketing QBit for one instance: the transports, the
** model providers, the CRM gateway, and the in-process limits (agent
** concurrency, model-call rate, the single execution worker).
** model providers, the CRM gateway, and the single execution worker's
** in-process lock. (Agent concurrency and the model-call rate are counted in
** the database by MarketingAgentService.)
**
** Registered as supplemental instance meta-data (not published to
** frontends), so process steps and widgets find it from QContext.
**
** Scale note: the limits here are per JVM. One replica is the supported shape
** today; see docs/OPERATIONS.md ("Scale and limits").
** Scale note: the execution worker is per JVM. One worker replica is the
** supported shape today; see docs/OPERATIONS.md ("Scale and limits").
*******************************************************************************/
public class MarketingRuntime implements QSupplementalInstanceMetaData
{
@@ -69,9 +67,7 @@ public class MarketingRuntime implements QSupplementalInstanceMetaData
private final MarketingQBitConfig config;
private final Map<String, MarketingTransport> transports = new LinkedHashMap<>();
private final Semaphore agentPermits;
private final ReentrantLock workerLock = new ReentrantLock();
private final Deque<Long> modelCallTimes = new ArrayDeque<>();
private final String workerName;
private volatile MarketingTransport qbitEmailTransport;
@@ -91,7 +87,6 @@ public class MarketingRuntime implements QSupplementalInstanceMetaData
{
transports.put(transport.getName(), transport);
}
this.agentPermits = new Semaphore(config.getMaxConcurrentAgentRuns());
this.workerName = hostName() + ":" + UUID.randomUUID().toString().substring(0, 8);
}
@@ -328,47 +323,6 @@ public class MarketingRuntime implements QSupplementalInstanceMetaData
/*******************************************************************************
** Try to start an agent run (bounded concurrency). Release with
** releaseAgentPermit.
*******************************************************************************/
public boolean tryAcquireAgentPermit()
{
return (agentPermits.tryAcquire());
}
/*******************************************************************************
**
*******************************************************************************/
public void releaseAgentPermit()
{
agentPermits.release();
}
/*******************************************************************************
** Reserve a model call within the per-minute rate (false = over the limit).
*******************************************************************************/
public synchronized boolean tryReserveModelCall()
{
long now = System.currentTimeMillis();
while(!modelCallTimes.isEmpty() && modelCallTimes.peekFirst() < now - 60_000)
{
modelCallTimes.pollFirst();
}
if(modelCallTimes.size() >= config.getMaxModelCallsPerMinute())
{
return (false);
}
modelCallTimes.addLast(now);
return (true);
}
/*******************************************************************************
** The single execution worker's lock in this JVM.
*******************************************************************************/
@@ -78,11 +78,15 @@ import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
** A run:
** 1. checks the policy (kill switch), the provider and model, the daily
** token and monthly cost budgets, and the concurrency limit - failing the
** run with a plain reason when any is missing (missing-provider behavior);
** run with a plain reason when any is missing (missing-provider behavior).
** The concurrency limit counts RUNNING runs in the database (after failing
** those whose heartbeat stopped), so it holds across replicas;
** 2. runs the deterministic readiness checks (a CHECK step);
** 3. loops model call -> tool calls, bounded by model calls, tool calls,
** wall-clock time and the per-minute model-call rate; it re-reads its own
** status before every model call so Cancel Run stops it;
** wall-clock time and the per-minute model-call rate (counted from the
** database; at the limit the run waits, it does not fail); it re-reads its
** own status before every model call so Cancel Run stops it, and records a
** heartbeat on each model call;
** 4. records every model call (tokens, stop reason) and tool call (input,
** outcome, enveloped result), the usage and cost, the drafts and findings
** it created, and its final summary.
@@ -98,8 +102,29 @@ public class MarketingAgentService
static final int MAX_STEP_JSON = 20_000;
/////////////////////////////////////////////////////////////////////////////
// a RUNNING run silent for maxRunSeconds plus this long is not alive //
/////////////////////////////////////////////////////////////////////////////
public static final int STALE_GRACE_SECONDS = 60;
private final AiRedactor redactor = new AiRedactor(List.of());
private Duration modelCallWindow = Duration.ofSeconds(60);
private Duration modelCallPoll = Duration.ofSeconds(1);
/*******************************************************************************
** Change the per-minute window and how often a waiting run re-checks it
** (the defaults are 60 seconds and 1 second; tests use short ones).
*******************************************************************************/
public MarketingAgentService withModelCallWindow(Duration window, Duration poll)
{
this.modelCallWindow = window;
this.modelCallPoll = poll;
return (this);
}
/*******************************************************************************
@@ -158,19 +183,7 @@ public class MarketingAgentService
return (finish(run, AgentRunStatus.FAILED, null, refusal, null, new MarketingToolContext(run, 1), new ModelUsage(), 0, 0));
}
if(!runtime.tryAcquireAgentPermit())
{
return (finish(run, AgentRunStatus.FAILED, null, "Too many agent runs at once (limit " + runtime.getConfig().getMaxConcurrentAgentRuns() + "); start it again shortly.",
null, new MarketingToolContext(run, 1), new ModelUsage(), 0, 0));
}
try
{
return (execute(runtime, policy, provider, model, run, task));
}
finally
{
runtime.releaseAgentPermit();
}
return (execute(runtime, policy, provider, model, run, task));
}
@@ -213,6 +226,101 @@ public class MarketingAgentService
return ("This month's agent cost budget is used up.");
}
}
/////////////////////////////////////////////////////////////////////////////
// concurrency is counted in the database, across every replica and //
// product. NOTE: count-then-mark-RUNNING is still racy - two replicas can //
// both see room and both start. A conditional-update claim //
// (ConditionalUpdateAction, qqq#1178, QQQ 4.1.0-RC.7) will close that //
// window as a follow-up. //
/////////////////////////////////////////////////////////////////////////////
failStaleRuns(runtime.getConfig());
int limit = policy.getMaxConcurrentAgentRuns() != null ? policy.getMaxConcurrentAgentRuns() : runtime.getConfig().getMaxConcurrentAgentRuns();
if(runningRuns() >= limit)
{
return ("Too many agent runs at once (limit " + limit + "); start it again shortly.");
}
return (null);
}
/*******************************************************************************
** RUNNING runs (all products) whose heartbeat - or start, for a run that
** never beat - is older than maxRunSeconds plus a grace period are failed:
** their server stopped. Returns how many.
*******************************************************************************/
int failStaleRuns(MarketingQBitConfig config) throws QException
{
Instant cutoff = Instant.now().minusSeconds((long) config.getMaxRunSeconds() + STALE_GRACE_SECONDS);
List<QRecord> running = MarketingStore.asSystem(() -> MarketingStore.query(MarketingAgentRun.TABLE_NAME,
new QQueryFilter(new QFilterCriteria("status", QCriteriaOperator.EQUALS, AgentRunStatus.RUNNING.getId()))));
int count = 0;
for(QRecord record : running)
{
Instant last = record.getValueInstant("heartbeatAt");
last = last != null ? last : record.getValueInstant("startedAt");
last = last != null ? last : record.getValueInstant("createDate");
if(last != null && last.isBefore(cutoff))
{
MarketingStore.updateInternal(MarketingAgentRun.TABLE_NAME, new QRecord().withValue("id", record.getValue("id"))
.withValue("status", AgentRunStatus.FAILED.getId()).withValue("finishedAt", Instant.now())
.withValue("errorMessage", "Stopped responding; the server may have restarted."));
count++;
}
}
return (count);
}
/*******************************************************************************
** RUNNING runs across all products, counted as the system user.
*******************************************************************************/
static int runningRuns() throws QException
{
return (MarketingStore.asSystem(() -> MarketingStore.count(MarketingAgentRun.TABLE_NAME,
new QQueryFilter(new QFilterCriteria("status", QCriteriaOperator.EQUALS, AgentRunStatus.RUNNING.getId())))));
}
/*******************************************************************************
** Model calls (steps of the model-call kind, all runs, system user) recorded
** since the instant.
*******************************************************************************/
static int modelCallsSince(Instant since) throws QException
{
return (MarketingStore.asSystem(() -> MarketingStore.count(MarketingAgentStep.TABLE_NAME, new QQueryFilter(
new QFilterCriteria("kind", QCriteriaOperator.EQUALS, AgentStepKind.MODEL_CALL.getId()),
new QFilterCriteria("createDate", QCriteriaOperator.GREATER_THAN_OR_EQUALS, since)))));
}
/*******************************************************************************
** Wait until the per-minute model-call limit has room. Returns null to go
** ahead, CANCELLED if a person cancelled the run while waiting, or FAILED
** if the run's time ran out first.
**
** NOTE: counting then calling is racy across replicas (see preflight); the
** conditional-update claim follow-up (qqq#1178) will close it.
*******************************************************************************/
private AgentRunStatus awaitModelCallSlot(MarketingAgentRun run, int limit, Instant deadline) throws QException, InterruptedException
{
while(modelCallsSince(Instant.now().minus(modelCallWindow)) >= limit)
{
if(isCancelled(run.getId()))
{
return (AgentRunStatus.CANCELLED);
}
if(Instant.now().isAfter(deadline))
{
return (AgentRunStatus.FAILED);
}
Thread.sleep(modelCallPoll.toMillis());
}
return (null);
}
@@ -229,7 +337,7 @@ public class MarketingAgentService
int maxTools = policy != null && policy.getMaxToolCallsPerRun() != null ? policy.getMaxToolCallsPerRun() : config.getMaxToolCallsPerRun();
MarketingStore.updateInternal(MarketingAgentRun.TABLE_NAME, new QRecord().withValue("id", run.getId())
.withValue("status", AgentRunStatus.RUNNING.getId()).withValue("startedAt", Instant.now()));
.withValue("status", AgentRunStatus.RUNNING.getId()).withValue("startedAt", Instant.now()).withValue("heartbeatAt", Instant.now()));
MarketingToolContext context = new MarketingToolContext(run, config.getMaxRowsToModel());
ModelUsage usage = new ModelUsage();
@@ -278,11 +386,21 @@ public class MarketingAgentService
{
return (finish(run, AgentRunStatus.FAILED, null, "Stopped: the run reached its limit of " + config.getMaxModelCallsPerRun() + " model calls.", null, context, usage, modelCalls, toolCalls));
}
if(!runtime.tryReserveModelCall())
/////////////////////////////////////////////////////////////////
// at the per-minute limit, wait (bounded by the run's time) //
/////////////////////////////////////////////////////////////////
AgentRunStatus blocked = awaitModelCallSlot(run, config.getMaxModelCallsPerMinute(), deadline);
if(blocked == AgentRunStatus.CANCELLED)
{
return (finish(run, AgentRunStatus.FAILED, null, "Stopped: the model-call rate limit (" + config.getMaxModelCallsPerMinute() + " per minute) was reached.", null, context, usage, modelCalls,
toolCalls));
return (finish(run, AgentRunStatus.CANCELLED, null, "Cancelled by a person.", null, context, usage, modelCalls, toolCalls));
}
if(blocked == AgentRunStatus.FAILED)
{
return (finish(run, AgentRunStatus.FAILED, null, "Stopped: the run exceeded " + config.getMaxRunSeconds() + " seconds, waiting for the model-call rate limit ("
+ config.getMaxModelCallsPerMinute() + " per minute).", null, context, usage, modelCalls, toolCalls));
}
MarketingStore.updateInternal(MarketingAgentRun.TABLE_NAME, new QRecord().withValue("id", run.getId()).withValue("heartbeatAt", Instant.now()));
ModelRequest request = new ModelRequest()
.withModel(model)
@@ -68,6 +68,9 @@ public class MarketingAgentRun extends QRecordEntity
@QField()
private Instant startedAt;
@QField()
private Instant heartbeatAt;
@QField()
private Instant finishedAt;
@@ -460,6 +463,36 @@ public class MarketingAgentRun extends QRecordEntity
}
/*******************************************************************************
** 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 MarketingAgentRun withHeartbeatAt(Instant heartbeatAt)
{
this.heartbeatAt = heartbeatAt;
return (this);
}
/*******************************************************************************
** Getter for finishedAt
*******************************************************************************/
@@ -17,9 +17,14 @@
package com.kingsrook.qbits.marketing;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.spi.ModelPart;
import com.kingsrook.qbits.assistant.spi.ModelProviderException;
import com.kingsrook.qbits.assistant.spi.ModelResponse;
import com.kingsrook.qbits.assistant.spi.ModelUsage;
import com.kingsrook.qbits.marketing.agent.MarketingAgentService;
import com.kingsrook.qbits.marketing.agent.tools.MarketingTools;
@@ -44,6 +49,8 @@ import com.kingsrook.qqq.backend.core.actions.permissions.PermissionEnforcement;
import com.kingsrook.qqq.backend.core.actions.processes.QProcessCallbackFactory;
import com.kingsrook.qqq.backend.core.actions.processes.RunProcessAction;
import com.kingsrook.qqq.backend.core.actions.tables.QueryAction;
import com.kingsrook.qqq.backend.core.context.CapturedContext;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.exceptions.QPermissionDeniedException;
import com.kingsrook.qqq.backend.core.model.actions.processes.RunProcessInput;
@@ -280,11 +287,10 @@ class AgentAndSecurityTest extends BaseMarketingTest
/*******************************************************************************
** Runs stop at their model-call limit; a cancelled run stops; the
** concurrency limit refuses extra runs.
** Runs stop at their model-call limit; a cancelled run stops.
*******************************************************************************/
@Test
void testLimitsCancelAndConcurrency() throws QException
void testLimitsAndCancel() throws QException
{
config.withMaxModelCallsPerRun(2);
try
@@ -320,20 +326,107 @@ class AgentAndSecurityTest extends BaseMarketingTest
});
MarketingAgentRun cancelled = new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(AgentRunStatus.getById(cancelled.getStatus())).isEqualTo(AgentRunStatus.CANCELLED);
}
MarketingRuntime runtime = MarketingRuntime.current();
assertThat(runtime.tryAcquireAgentPermit()).isTrue();
assertThat(runtime.tryAcquireAgentPermit()).isTrue();
try
/*******************************************************************************
** The concurrency limit is counted in the database: a second service
** instance (another replica on the same backend) is refused while the
** first one's run is RUNNING.
*******************************************************************************/
@Test
void testConcurrencyHoldsAcrossReplicas() throws QException
{
config.withMaxConcurrentAgentRuns(1);
CapturedContext context = QContext.capture();
List<MarketingAgentRun> refused = new ArrayList<>();
model.thenRespond(request ->
{
MarketingAgentRun busy = new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(busy.getErrorMessage()).contains("Too many agent runs");
}
finally
{
runtime.releaseAgentPermit();
runtime.releaseAgentPermit();
}
try
{
QContext.init(context);
refused.add(new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null));
}
catch(QException e)
{
throw (new RuntimeException(e));
}
return (new ModelResponse().withParts(new ArrayList<>(List.of(ModelPart.text("Done.")))).withStopReason(ModelResponse.STOP_END_TURN));
});
MarketingAgentRun first = new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(AgentRunStatus.getById(first.getStatus())).isEqualTo(AgentRunStatus.SUCCEEDED);
assertThat(first.getHeartbeatAt()).isNotNull();
assertThat(refused).hasSize(1);
assertThat(AgentRunStatus.getById(refused.get(0).getStatus())).isEqualTo(AgentRunStatus.FAILED);
assertThat(refused.get(0).getErrorMessage()).contains("Too many agent runs").contains("limit 1");
/////////////////////////////////////////////////////////////
// the first run finished, so the next one is let through //
/////////////////////////////////////////////////////////////
model.thenText("Again.");
assertThat(AgentRunStatus.getById(new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null).getStatus())).isEqualTo(AgentRunStatus.SUCCEEDED);
}
/*******************************************************************************
** A RUNNING run whose heartbeat is older than maxRunSeconds plus the grace
** is failed with a plain reason (and no longer counts); a fresh one stays.
*******************************************************************************/
@Test
void testStaleRunningRunIsReaped() throws QException
{
config.withMaxConcurrentAgentRuns(1);
Instant stale = Instant.now().minusSeconds(config.getMaxRunSeconds() + MarketingAgentService.STALE_GRACE_SECONDS + 30);
Integer staleId = MarketingStore.insertInternal(new MarketingAgentRun().withTask(AgentTask.ANALYZE_RESULTS.getId()).withStatus(AgentRunStatus.RUNNING.getId())
.withStartedAt(stale).withHeartbeatAt(stale)).getValueInteger("id");
model.thenText("Ran.");
MarketingAgentRun run = new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(AgentRunStatus.getById(run.getStatus())).isEqualTo(AgentRunStatus.SUCCEEDED);
QRecord reaped = MarketingStore.get(MarketingAgentRun.TABLE_NAME, staleId);
assertThat(reaped.getValueInteger("status")).isEqualTo(AgentRunStatus.FAILED.getId());
assertThat(reaped.getValueString("errorMessage")).isEqualTo("Stopped responding; the server may have restarted.");
///////////////////////////////////////////////////////////////////
// a RUNNING run with a recent heartbeat is alive and counts //
///////////////////////////////////////////////////////////////////
MarketingStore.insertInternal(new MarketingAgentRun().withTask(AgentTask.ANALYZE_RESULTS.getId()).withStatus(AgentRunStatus.RUNNING.getId())
.withStartedAt(stale).withHeartbeatAt(Instant.now()));
MarketingAgentRun refused = new MarketingAgentService().run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(refused.getErrorMessage()).contains("Too many agent runs");
}
/*******************************************************************************
** At the per-minute model-call limit a run waits for the window to clear
** rather than failing; the wait is bounded by maxRunSeconds.
*******************************************************************************/
@Test
void testModelCallRateWaitsInsteadOfFailing() throws QException
{
config.withMaxModelCallsPerMinute(1);
MarketingAgentService service = new MarketingAgentService().withModelCallWindow(Duration.ofMillis(1500), Duration.ofMillis(50));
model.thenText("One.").thenText("Two.");
assertThat(AgentRunStatus.getById(service.run(AgentTask.ANALYZE_RESULTS, campaignId, null, null).getStatus())).isEqualTo(AgentRunStatus.SUCCEEDED);
long start = System.currentTimeMillis();
MarketingAgentRun second = service.run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(AgentRunStatus.getById(second.getStatus())).isEqualTo(AgentRunStatus.SUCCEEDED);
assertThat(System.currentTimeMillis() - start).isGreaterThanOrEqualTo(500);
///////////////////////////////////////////////////////////////////
// a wait longer than the run's time ends the run, not the JVM //
///////////////////////////////////////////////////////////////////
config.withMaxRunSeconds(1);
MarketingAgentService slow = new MarketingAgentService().withModelCallWindow(Duration.ofSeconds(60), Duration.ofMillis(50));
MarketingAgentRun timedOut = slow.run(AgentTask.ANALYZE_RESULTS, campaignId, null, null);
assertThat(AgentRunStatus.getById(timedOut.getStatus())).isEqualTo(AgentRunStatus.FAILED);
assertThat(timedOut.getErrorMessage()).contains("exceeded 1 seconds").contains("rate limit");
}
}
+1
View File
@@ -314,6 +314,7 @@ CREATE TABLE IF NOT EXISTS marketing_agent_run (
model VARCHAR(100),
requested_by_user_id VARCHAR(100),
started_at TIMESTAMP,
heartbeat_at TIMESTAMP,
finished_at TIMESTAMP,
model_calls INTEGER,
tool_calls INTEGER,