Merge pull request 'fix(store): the assistant works for people who do not administer it where core enforces table permissions' (#102) from feat/store-reads-as-trusted-code into develop
maven / build (push) Successful in 6m28s
build / maven (push) Successful in 6m28s

This commit was merged in pull request #102.
This commit is contained in:
james.maes committed 2026-11-11 17:10:42 +00:00
commit e02658de7e
9 files changed
+461 -15

No files matched your search

+7
View File
@@ -1,5 +1,12 @@
# Changelog
## Unreleased - The assistant works for people who do not administer it on a host that enforces table permissions in core actions; a person's own budget (UI API extension 1.8.0)
- Fixed: on a host with `QInstance.enforcePermissionsInCoreActions` on, a person with `assistant.hasAccess` and without `assistant.admin` was refused before any model call ("Not authorized to access the assistant admin data."). The assistant read its settings (policy, providers, models, prices, budgets, data access) and wrote that person's messages, tool calls and usage through tables whose permission rule is `assistant.admin`, with the person's session. `AssistantStore` now marks every one of its calls as the assistant's own service code (`withPermissionBypass`), the way qbit-collaboration and qbit-backups do.
- What that does not change: the session is still the person's, so writes are attributed to them and the owner record locks keep every other person's conversations, messages, pending actions and usage out (core does not relax record locks for a bypass). The tables' permission rules are unchanged, so the tables stay closed to that person in the UI, the UI API and the assistant's tools. Every operation of the assistant API still requires assistant access first.
- A person sees their own spending against their own limit and nobody else's. `GET /config` has `budgets` (UI API extension assistant 1.8.0): the budgets scoped to that user or to a permission that user holds, with used, limit, percent, state and when the count starts again. Budgets shared by several people (INSTANCE, PROVIDER, MODEL) are not listed. The same lines ride beside the person's question (`<assistant_budget>`), so "how much of my budget have I used?" has an answer. `BudgetService.own` and `ownLine` are public.
- New `CoreEnforcementTest` (fixture instance with core enforcement on): a non-administrator's question with a tool call completes and is accounted to them; the settings tables refuse that person and the assistant's tools do not reach them; one person cannot read another's conversation, messages or model calls, through the API or the store; own budget in `/config` and beside the question; a person at their own daily limit is refused with the sentence while another person is answered. Without the store change every one of them fails at the first store read.
## Unreleased - Save and schedule an analysis with one confirmation; confirmed actions reach the conversation; ask before guessing
- New tool `save_analysis`: keeps an `analyze_records` question as a saved analysis the user owns (`data_analysis_saved`) and, when the user asked for it to be sent regularly, a daily schedule emailed to them (`data_analysis_schedule`: DAILY, a local time, an IANA time zone, the confirming user as the recipient). It only ever proposes: ONE pending-action card lists the saved analysis, what it shows, the filter, the schedule, who is emailed and the numbers as they are now. Nothing is created until the user confirms, whatever the policy says about confirming record writes. On confirmation both records are created through the tables' own insert operations with the confirming user's session, so their permissions, owner locks and validation decide.
@@ -187,7 +187,7 @@ public class AssistantRuntime
this.knowledgeTool = new KnowledgeTool(this, ragService, vectors);
this.bulkLoadTools = new BulkLoadTools(qInstance, dataAccessService, config.getBulkLoadRunner());
this.attachmentService = new AttachmentService(config.getAttachments(), new UploadRefSource());
this.contextBuilder = new ContextBuilder(config, kernel, toolbox).withKnowledgeTool(knowledgeTool).withBulkLoadTools(bulkLoadTools);
this.contextBuilder = new ContextBuilder(config, kernel, toolbox).withKnowledgeTool(knowledgeTool).withBulkLoadTools(bulkLoadTools).withBudgets(budgetService);
if(Boolean.TRUE.equals(config.getRegisterShutdownHook()))
{
///////////////////////////////////////////////////////////////////////
@@ -36,6 +36,7 @@ import com.kingsrook.qbits.assistant.access.DataAccessMode;
import com.kingsrook.qbits.assistant.attachments.AttachmentService;
import com.kingsrook.qbits.assistant.attachments.UploadRefSource;
import com.kingsrook.qbits.assistant.engine.AbandonedTurns;
import com.kingsrook.qbits.assistant.engine.BudgetService;
import com.kingsrook.qbits.assistant.engine.ClientErrors;
import com.kingsrook.qbits.assistant.engine.ContextBuilder;
import com.kingsrook.qbits.assistant.engine.MessageParts;
@@ -207,11 +208,49 @@ public class AssistantApi
{
result.put("conversationsTable", AssistantConversation.TABLE_NAME);
}
if(hasAccess())
{
result.put("budgets", ownBudgets());
}
return (result);
}
/*******************************************************************************
** This person's own budgets with what they have used (BudgetService.own):
** their own spending against their own limits, never a shared budget's count
** and never another person's. Never throws.
*******************************************************************************/
private List<Map<String, Object>> ownBudgets()
{
List<Map<String, Object>> rs = new ArrayList<>();
try
{
for(BudgetService.Snapshot snapshot : runtime.getBudgetService().own(userId))
{
Map<String, Object> budget = new LinkedHashMap<>();
budget.put("name", snapshot.budget().getName());
budget.put("period", snapshot.budget().getPeriod());
budget.put("metric", snapshot.budget().getMetric());
budget.put("used", snapshot.used());
budget.put("limit", snapshot.limit());
budget.put("percent", snapshot.percent());
budget.put("state", snapshot.state());
budget.put("startsAgainAt", String.valueOf(BudgetService.periodEnd(snapshot.budget().getPeriod(), snapshot.periodStart())));
budget.put("text", BudgetService.ownLine(snapshot));
rs.add(budget);
}
}
catch(QException | RuntimeException e)
{
LOG.warn("Could not read the user's own assistant budgets", e);
}
return (rs);
}
/*******************************************************************************
** The example prompts for this user (see StarterPrompts): built from the
** tables the user may read, plus the configured ones they may see. Never
@@ -63,7 +63,7 @@ import io.javalin.http.HttpStatus;
public final class AssistantSpecs
{
public static final String EXTENSION_NAME = "assistant";
public static final String EXTENSION_VERSION = "1.7.0";
public static final String EXTENSION_VERSION = "1.8.0";
public static final String PREFIX = "/assistant";
@@ -208,7 +208,11 @@ public final class AssistantSpecs
+ "and sent as pageContext.attachments [{uploadRef, fileName}] are read for the question."),
"profile", ref("AssistantProfileV1"),
"conversationsTable", string("Since 1.6.0: the table name of the user's conversations (assistant_conversation) for the standard table UI "
+ "(manage, archive, restore and delete); present only when this session may query that table.").withNullable(true)));
+ "(manage, archive, restore and delete); present only when this session may query that table.").withNullable(true),
"budgets", array(freeObject("name, period (DAY, WEEK or MONTH, in UTC), metric (COST_USD or TOKENS), used, limit (null while an override makes it unlimited), percent, "
+ "state (OK, WARNED, DOWNGRADED, STOPPED or UNLIMITED), startsAgainAt (ISO-8601 instant) and text (the same as one sentence)."),
"Since 1.8.0: this user's own budgets with what this user has used in the current period: the ones scoped to the user or to a permission the user holds. "
+ "Budgets shared by several people are not listed. Present when the user may use the assistant; empty when the user has no budget of their own.")));
rs.put("AssistantConversationV1", object("One of the user's own conversations.",
"id", integer(null), "title", string(null).withNullable(true), "status", string("ACTIVE or ARCHIVED."), "pinned", bool(null),
"lastMessageAt", string("ISO-8601 instant.").withNullable(true), "unreadCount", integer(null), "createDate", string("ISO-8601 instant.").withNullable(true),
@@ -202,6 +202,54 @@ public class BudgetService
/*******************************************************************************
** The budgets that are this person's own: scoped to them (USER) or to a
** permission they hold (ROLE, where every holder is counted separately), with
** what that person has used in the current period. It is asked with that
** person's session, so the usage read is their own rows and nobody else's.
** Budgets shared by several people (INSTANCE, PROVIDER, MODEL) are left out:
** their count is other people's spending too.
*******************************************************************************/
public List<Snapshot> own(String userId) throws QException
{
List<Snapshot> rs = new ArrayList<>();
if(userId == null)
{
return (rs);
}
for(AssistantBudget budget : store.budgets())
{
boolean mine = switch(budget.getScopeType())
{
case "USER" -> Objects.equals(budget.getScopeValue(), userId);
case "ROLE" -> QContext.getQSession().hasPermission(budget.getScopeValue());
default -> false;
};
if(mine)
{
rs.add(snapshot(budget, userId));
}
}
return (rs);
}
/*******************************************************************************
** One of a person's own budgets as a sentence: what is used, of how much, and
** when the count starts again.
*******************************************************************************/
public static String ownLine(Snapshot snapshot)
{
AssistantBudget budget = snapshot.budget();
Instant end = periodEnd(budget.getPeriod(), snapshot.periodStart());
return (periodWord(budget.getPeriod()) + " budget '" + budget.getName() + "': " + amount(budget, snapshot.used()) + " used"
+ (snapshot.limit() == null ? ", no limit at the moment" : " of " + amount(budget, snapshot.limit()) + " (" + snapshot.percent().setScale(0, RoundingMode.DOWN) + "%)")
+ "; the count starts again on " + RESET_DAY.format(LocalDate.ofInstant(end, ZoneOffset.UTC)) + " at 00:00 UTC");
}
/*******************************************************************************
** Whether the caller and selected provider/model are within this scope.
*******************************************************************************/
@@ -27,6 +27,7 @@ import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantSecurity;
import com.kingsrook.qbits.assistant.model.AssistantMessage;
import com.kingsrook.qbits.assistant.spi.CacheHint;
import com.kingsrook.qbits.assistant.spi.ModelMessage;
@@ -41,6 +42,8 @@ import com.kingsrook.qbits.assistant.tools.bulkload.BulkLoadTools;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolDefinition;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolKernel;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.logging.QLogger;
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;
@@ -65,6 +68,8 @@ public class ContextBuilder
///////////////////////////////////////////////////////////////////////
public static final Set<String> KERNEL_TOOLS = Set.of("navigate", "search_help");
private static final QLogger LOG = QLogger.getLogger(ContextBuilder.class);
public static final int MAX_CONTEXT_CHARS = 512;
public static final int MAX_SELECTED_IDS = 100;
public static final int MAX_ATTACHMENTS = 5;
@@ -75,6 +80,7 @@ public class ContextBuilder
private volatile String appDigest;
private KnowledgeTool knowledgeTool;
private BulkLoadTools bulkLoadTools;
private BudgetService budgets;
@@ -296,6 +302,11 @@ public class ContextBuilder
newMessage.withPart(ModelPart.text("<page_context note=\"where the user is in the app; a hint from the browser, not verified\">"
+ JsonUtils.toJson(pageContext) + "</page_context>"));
}
String ownBudgets = ownBudgets(session);
if(ownBudgets != null)
{
newMessage.withPart(ModelPart.text(ownBudgets));
}
if(!messages.isEmpty() && messages.get(messages.size() - 1).getRole() == ModelMessage.Role.USER)
{
ModelMessage previous = messages.remove(messages.size() - 1);
@@ -345,6 +356,51 @@ public class ContextBuilder
/*******************************************************************************
** The asking person's own budgets (BudgetService.own), as a note beside their
** question, so "how much of my budget have I used?" has an answer: their own
** spending against their own limits and nobody else's. It rides with the new
** message, not the cached blocks, since it changes with every question. Null
** when the person has no budget of their own; a failure to read is logged
** and never fails the question.
*******************************************************************************/
String ownBudgets(QSession session)
{
if(budgets == null)
{
return (null);
}
try
{
List<BudgetService.Snapshot> own = budgets.own(AssistantSecurity.userId(session));
if(own.isEmpty())
{
return (null);
}
StringBuilder text = new StringBuilder("<assistant_budget note=\"the user's own assistant spending limits, from the server; use it only when they ask about their usage or limit\">");
own.forEach(snapshot -> text.append("\n").append(BudgetService.ownLine(snapshot)));
return (text.append("\n</assistant_budget>").toString());
}
catch(QException | RuntimeException e)
{
LOG.warn("Could not read the user's own assistant budgets", e);
return (null);
}
}
/*******************************************************************************
** The budget service, for the asking person's own budgets.
*******************************************************************************/
public ContextBuilder withBudgets(BudgetService budgets)
{
this.budgets = budgets;
return (this);
}
/*******************************************************************************
** Append text as a message, merging consecutive messages of one role.
*******************************************************************************/
@@ -62,6 +62,7 @@ import com.kingsrook.qqq.backend.core.actions.tables.QueryAction;
import com.kingsrook.qqq.backend.core.actions.tables.UpdateAction;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.AbstractActionInput;
import com.kingsrook.qqq.backend.core.model.actions.tables.ActionFlag;
import com.kingsrook.qqq.backend.core.model.actions.tables.count.CountInput;
import com.kingsrook.qqq.backend.core.model.actions.tables.delete.DeleteInput;
@@ -85,6 +86,22 @@ import com.kingsrook.qqq.backend.core.utils.CollectionUtils;
** own rows). Queries also filter on userId explicitly, so an admin's
** all-access key never mixes other users' conversations into the panel.
**
** Every call is the assistant's own service code and is marked as such
** (withPermissionBypass), so a host that enforces table permissions inside
** core actions (QInstance.enforcePermissionsInCoreActions) does not ask the
** person for assistant.admin when the assistant reads its policy, providers,
** models, prices, budgets and data access settings to decide, or writes that
** person's messages, tool calls, pending actions and usage. Only the table
** permission check is skipped, and only for these calls:
** - the session stays the person's, so a write is attributed to them and the
** owner record locks still keep every other person's rows out (core does
** not relax record locks for a bypass);
** - the tables' permission rules are unchanged, so the tables stay closed to
** that person in the UI, the UI API and the assistant's tools;
** - who may call is decided before the store is reached: the assistant API
** requires assistant access for every operation, and the save rules of the
** budgets and overrides require assistant.admin themselves.
**
** This is the only package allowed to call QQQ data actions (ArchUnit).
*******************************************************************************/
public class AssistantStore
@@ -111,12 +128,26 @@ public class AssistantStore
private static final Map<String, Object> USAGE_LOCKS = new ConcurrentHashMap<>();
static final String TRUSTED_REASON = "Assistant store: the assistant's own tables, behind the assistant API's access check; owner record locks still apply.";
/////////////////////////////////////////////////////////////////////////////
// generic helpers //
/////////////////////////////////////////////////////////////////////////////
/*******************************************************************************
** Mark a store call as the assistant's own service code (see the class
** comment). The reason is what core logs when it applies the bypass.
*******************************************************************************/
private static <T extends AbstractActionInput> T trusted(T input)
{
input.withPermissionBypass(TRUSTED_REASON);
return (input);
}
/*******************************************************************************
** The store only ever touches the QBit's own tables (assistant_*): the
** user's application data is reached only through the UI API.
@@ -137,7 +168,7 @@ public class AssistantStore
public <T extends QRecordEntity> List<T> query(String tableName, QQueryFilter filter, Class<T> entityClass) throws QException
{
requireOwnTable(tableName);
List<QRecord> records = new QueryAction().execute(new QueryInput(tableName).withFilter(filter)).getRecords();
List<QRecord> records = new QueryAction().execute(trusted(new QueryInput(tableName)).withFilter(filter)).getRecords();
List<T> result = new ArrayList<>();
for(QRecord record : CollectionUtils.nonNullList(records))
{
@@ -155,7 +186,7 @@ public class AssistantStore
public <T extends QRecordEntity> T insert(T entity) throws QException
{
requireOwnTable(entity.tableName());
InsertOutput output = new InsertAction().execute(new InsertInput(entity.tableName()).withRecord(cryptRecord(entity.tableName(), entity.toQRecord(), true)));
InsertOutput output = new InsertAction().execute(trusted(new InsertInput(entity.tableName())).withRecord(cryptRecord(entity.tableName(), entity.toQRecord(), true)));
QRecord record = output.getRecords().get(0);
if(CollectionUtils.nullSafeHasContents(record.getErrors()))
{
@@ -183,7 +214,7 @@ public class AssistantStore
{
requireOwnTable(entity.tableName());
rememberTranscript(entity.tableName(), entity.toQRecord(), entity.toQRecord().getValue("id"));
UpdateOutput output = new UpdateAction().execute(new UpdateInput(entity.tableName()).withRecord(cryptRecord(entity.tableName(), entity.toQRecord(), true)));
UpdateOutput output = new UpdateAction().execute(trusted(new UpdateInput(entity.tableName())).withRecord(cryptRecord(entity.tableName(), entity.toQRecord(), true)));
QRecord record = output.getRecords().get(0);
if(CollectionUtils.nullSafeHasContents(record.getErrors()))
{
@@ -200,7 +231,7 @@ public class AssistantStore
{
requireOwnTable(tableName);
rememberTranscript(tableName, record, record.getValue("id"));
UpdateOutput output = new UpdateAction().execute(new UpdateInput(tableName).withRecord(cryptRecord(tableName, record, true)));
UpdateOutput output = new UpdateAction().execute(trusted(new UpdateInput(tableName)).withRecord(cryptRecord(tableName, record, true)));
QRecord result = output.getRecords().get(0);
if(CollectionUtils.nullSafeHasContents(result.getErrors()))
{
@@ -221,7 +252,7 @@ public class AssistantStore
public void delete(String tableName, QQueryFilter filter) throws QException
{
requireOwnTable(tableName);
new DeleteAction().execute(new DeleteInput(tableName).withQueryFilter(filter));
new DeleteAction().execute(trusted(new DeleteInput(tableName)).withQueryFilter(filter));
}
@@ -743,7 +774,7 @@ public class AssistantStore
{
records.add(new QRecord().withValue("id", call.getId()).withValue("resultJson", null));
}
new UpdateAction().execute(new UpdateInput(AssistantToolCall.TABLE_NAME).withRecords(records));
new UpdateAction().execute(trusted(new UpdateInput(AssistantToolCall.TABLE_NAME)).withRecords(records));
return (calls.size());
}
@@ -764,7 +795,7 @@ public class AssistantStore
}
List<Serializable> ids = new ArrayList<>();
events.forEach(e -> ids.add(e.getId()));
new DeleteAction().execute(new DeleteInput(AssistantTurnEvent.TABLE_NAME).withPrimaryKeys(ids));
new DeleteAction().execute(trusted(new DeleteInput(AssistantTurnEvent.TABLE_NAME)).withPrimaryKeys(ids));
return (events.size());
}
@@ -776,7 +807,7 @@ public class AssistantStore
*******************************************************************************/
public int countClientErrorsSince(Instant since) throws QException
{
Integer count = new CountAction().execute(new CountInput(AssistantClientError.TABLE_NAME)
Integer count = new CountAction().execute(trusted(new CountInput(AssistantClientError.TABLE_NAME))
.withFilter(new QQueryFilter().withCriteria("createDate", QCriteriaOperator.GREATER_THAN_OR_EQUALS, since))).getCount();
return (count == null ? 0 : count);
}
@@ -797,7 +828,7 @@ public class AssistantStore
}
List<Serializable> ids = new ArrayList<>();
errors.forEach(e -> ids.add(e.getId()));
new DeleteAction().execute(new DeleteInput(AssistantClientError.TABLE_NAME).withPrimaryKeys(ids));
new DeleteAction().execute(trusted(new DeleteInput(AssistantClientError.TABLE_NAME)).withPrimaryKeys(ids));
return (errors.size());
}
@@ -923,7 +954,7 @@ public class AssistantStore
public List<QRecord> queryRecords(String tableName, QQueryFilter filter, int limit) throws QException
{
requireOwnTable(tableName);
List<QRecord> records = new QueryAction().execute(new QueryInput(tableName).withFilter(filter.withLimit(limit))).getRecords();
List<QRecord> records = new QueryAction().execute(trusted(new QueryInput(tableName)).withFilter(filter.withLimit(limit))).getRecords();
return (records == null ? List.of() : records.stream().map(record -> cryptRecord(tableName, record, false)).toList());
}
@@ -1160,7 +1191,7 @@ public class AssistantStore
*******************************************************************************/
public void scheduledUpdate(QRecord record) throws QException
{
UpdateOutput output = new UpdateAction().execute(new UpdateInput(AssistantScheduledTask.TABLE_NAME)
UpdateOutput output = new UpdateAction().execute(trusted(new UpdateInput(AssistantScheduledTask.TABLE_NAME))
.withRecord(record).withFlag(ScheduledWrite.INTERNAL));
if(CollectionUtils.nullSafeHasContents(output.getRecords().get(0).getErrors()))
{
@@ -212,6 +212,20 @@ components:
\ only when this session may query that table."
nullable: true
type: "string"
budgets:
description: "Since 1.8.0: this user's own budgets with what this user has\
\ used in the current period: the ones scoped to the user or to a permission\
\ the user holds. Budgets shared by several people are not listed. Present\
\ when the user may use the assistant; empty when the user has no budget\
\ of their own."
items:
additionalProperties: true
description: "name, period (DAY, WEEK or MONTH, in UTC), metric (COST_USD\
\ or TOKENS), used, limit (null while an override makes it unlimited),\
\ percent, state (OK, WARNED, DOWNGRADED, STOPPED or UNLIMITED), startsAgainAt\
\ (ISO-8601 instant) and text (the same as one sentence)."
type: "object"
type: "array"
type: "object"
AssistantConversationDetailResponseV1:
description: "A conversation and a page of its messages, oldest first."
@@ -476,7 +490,7 @@ info:
description: "QQQ Assistant: conversations, streamed turns, feedback, pending actions\
\ and data access explanations"
title: "QQQ Middleware API extension: assistant"
version: "1.7.0"
version: "1.8.0"
openapi: "3.0.3"
paths:
/qqq/v1/assistant/config:
@@ -0,0 +1,247 @@
/*
* 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.math.BigDecimal;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.api.AssistantApi;
import com.kingsrook.qbits.assistant.api.AssistantApiException;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetUsage;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantMessage;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantModelPrice;
import com.kingsrook.qbits.assistant.model.AssistantPolicy;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
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.TurnRunner;
import com.kingsrook.qqq.backend.core.actions.tables.QueryAction;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QPermissionDeniedException;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QueryInput;
import org.junit.jupiter.api.Test;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.ADMIN;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.EDITOR;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.OTHER_EDITOR;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/*******************************************************************************
** The assistant on a host that enforces table permissions inside core actions
** (QInstance.enforcePermissionsInCoreActions): a person who may use the
** assistant and does not administer it gets answers, within their own
** permissions and their own budget, and sees none of its configuration.
**
** The assistant's tables are of two kinds. A person's own data (conversations,
** messages, turns, tool calls, pending actions, usage) is owner-locked: the
** store reads and writes it with the person's session, so the record locks
** keep everybody else's rows out. The assistant's configuration and accounting
** (policy, providers, models, prices, budgets and their events, data access
** settings, the shared budget accounting) has no owner: the store reads it as
** trusted code, so the assistant can decide, and writes its accounting
** attributed to the person. In both cases only the table permission check is
** skipped, for the store's own calls; the tables' permission rules are
** unchanged for every other caller.
*******************************************************************************/
class CoreEnforcementTest extends BaseAssistantTest
{
/*******************************************************************************
** The fixture instance, enforcing table permissions inside core actions.
*******************************************************************************/
@Override
protected AssistantTestFixture.InstanceCustomizer instanceCustomizer()
{
return (qInstance -> qInstance.setEnforcePermissionsInCoreActions(true));
}
/*******************************************************************************
** A person with assistant access and without assistant.admin asks a question
** that needs a tool, and gets the answer. The model call is accounted to them.
*******************************************************************************/
@Test
void aPersonWhoDoesNotAdministerTheAssistantGetsAnAnswer() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
assertThat(QContext.getQSession().hasPermission("assistant.admin")).isFalse();
assertThat(QContext.getQSession().hasPermission("assistant.hasAccess")).isTrue();
provider.thenToolUse("count_records", input("table", "blog_post", "filter", filter("AND", "status", "EQUALS", List.of("PUBLISHED"))))
.thenRespond(request -> new ModelResponse().withParts(new ArrayList<>(List.of(ModelPart.text("There are 4 published posts.")))).withStopReason("end_turn"));
TurnRunner.Outcome outcome = runner.run(EDITOR, "how many published posts?");
assertThat(outcome.end()).containsEntry("status", "COMPLETE");
assertThat(outcome.text()).isEqualTo("There are 4 published posts.");
assertThat(outcome.toolCalls()).hasSize(1);
assertThat(outcome.toolCalls().get(0).getOutcome()).isEqualTo("OK");
QContext.init(qInstance, AssistantTestFixture.ownerSession(ADMIN, config));
List<AssistantUsage> calls = config.getRuntime().getStore().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class);
assertThat(calls).isNotEmpty().allMatch(call -> EDITOR.equals(call.getUserId()));
assertThat(config.getRuntime().getStore().query(AssistantBudgetUsage.TABLE_NAME, new QQueryFilter(), AssistantBudgetUsage.class)).hasSameSizeAs(calls);
}
/*******************************************************************************
** The configuration stays closed to that person: the tables refuse them, and
** the assistant's own tools do not know the tables.
*******************************************************************************/
@Test
void theConfigurationStaysClosedToThatPerson() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
for(String table : List.of(AssistantPolicy.TABLE_NAME, AssistantProvider.TABLE_NAME, AssistantModel.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantBudget.TABLE_NAME,
AssistantBudgetUsage.TABLE_NAME, AssistantUsage.TABLE_NAME, AssistantMessage.TABLE_NAME))
{
assertThatThrownBy(() -> new QueryAction().execute(new QueryInput(table))).as(table).isInstanceOf(QPermissionDeniedException.class);
}
provider.thenToolUse("query_records", input("table", AssistantProvider.TABLE_NAME))
.thenToolUse("count_records", input("table", AssistantBudget.TABLE_NAME))
.thenText("I cannot see that.");
TurnRunner.Outcome outcome = runner.run(EDITOR, "show me the assistant providers and count the budgets");
assertThat(outcome.end()).containsEntry("status", "COMPLETE");
assertThat(outcome.toolCalls()).hasSize(2).allMatch(call -> !"OK".equals(call.getOutcome()));
assertThat(outcome.toolCalls()).noneMatch(call -> String.valueOf(call.getResultJson()).contains("apiKey") || String.valueOf(call.getResultJson()).contains("limitValue"));
}
/*******************************************************************************
** One person's conversation is not another's: the API answers as if it did
** not exist, and the store, asked with the other person's session, finds no
** row (the owner record lock, which the store does not skip).
*******************************************************************************/
@Test
void aPersonsConversationIsNotAnothers() throws Exception
{
provider.thenText("Hello.");
TurnRunner.Outcome outcome = runner.run(EDITOR, "Hello");
Integer conversationId = outcome.conversationId();
QContext.init(qInstance, AssistantTestFixture.ownerSession(OTHER_EDITOR, config));
AssistantApi other = new AssistantApi(config.getRuntime(), QContext.getQSession());
assertThatThrownBy(() -> other.getConversation(conversationId, null, null)).isInstanceOf(AssistantApiException.class)
.satisfies(e -> assertThat(((AssistantApiException) e).getStatus()).isEqualTo(404));
AssistantStore store = config.getRuntime().getStore();
assertThat(store.query(AssistantConversation.TABLE_NAME, new QQueryFilter().withCriteria("id", QCriteriaOperator.EQUALS, conversationId), AssistantConversation.class)).isEmpty();
assertThat(store.query(AssistantMessage.TABLE_NAME, new QQueryFilter().withCriteria("conversationId", QCriteriaOperator.EQUALS, conversationId), AssistantMessage.class)).isEmpty();
assertThat(store.query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class)).as("nor the other person's model calls").isEmpty();
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
assertThat(store.query(AssistantMessage.TABLE_NAME, new QQueryFilter().withCriteria("conversationId", QCriteriaOperator.EQUALS, conversationId), AssistantMessage.class)).isNotEmpty();
}
/*******************************************************************************
** A person sees their own spending against their own limit, and nobody
** else's: in GET /config for the panel, and as a note beside their question
** so the assistant can answer "how much of my budget have I used?". The
** budget everyone shares is not listed, since its count is other people's
** spending too.
*******************************************************************************/
@Test
void aPersonSeesTheirOwnSpendingAgainstTheirOwnLimit() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.ownerSession(ADMIN, config));
AssistantStore store = config.getRuntime().getStore();
store.insert(new AssistantBudget().withName("Each person, a day").withScopeType("ROLE").withScopeValue("assistant.hasAccess")
.withPeriod("DAY").withMetric("COST_USD").withLimitValue(BigDecimal.valueOf(25)).withAtLimitAction("STOP").withHardCeilingPercent(100));
store.insert(new AssistantBudget().withName("Everyone, a week").withScopeType("INSTANCE").withPeriod("WEEK")
.withMetric("COST_USD").withLimitValue(BigDecimal.valueOf(500)).withAtLimitAction("STOP").withHardCeilingPercent(100));
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
store.insert(new AssistantUsage().withUserId(EDITOR).withPurpose("TURN").withStatus("OK").withCostMicros(3_200_000L).withInputTokens(1L).withOutputTokens(1L).withCreateDate(Instant.now()));
QContext.init(qInstance, AssistantTestFixture.ownerSession(OTHER_EDITOR, config));
store.insert(new AssistantUsage().withUserId(OTHER_EDITOR).withPurpose("TURN").withStatus("OK").withCostMicros(7_000_000L).withInputTokens(1L).withOutputTokens(1L).withCreateDate(Instant.now()));
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
List<Map<String, Object>> mine = budgetsOf(new AssistantApi(config.getRuntime(), QContext.getQSession()).config());
assertThat(mine).hasSize(1);
assertThat(mine.get(0)).containsEntry("name", "Each person, a day").containsEntry("period", "DAY").containsEntry("metric", "COST_USD").containsEntry("state", "OK");
assertThat((BigDecimal) mine.get(0).get("used")).isEqualByComparingTo("3.2");
assertThat((BigDecimal) mine.get(0).get("limit")).isEqualByComparingTo("25");
assertThat((String) mine.get(0).get("text")).startsWith("daily budget 'Each person, a day': $3.20 used of $25.00 (12%); the count starts again on ").endsWith(" at 00:00 UTC");
QContext.init(qInstance, AssistantTestFixture.ownerSession(OTHER_EDITOR, config));
List<Map<String, Object>> theirs = budgetsOf(new AssistantApi(config.getRuntime(), QContext.getQSession()).config());
assertThat(theirs).hasSize(1);
assertThat((BigDecimal) theirs.get(0).get("used")).as("the other person's own spending, not the sum").isEqualByComparingTo("7");
QContext.init(qInstance, AssistantTestFixture.sessionFor(AssistantTestFixture.NO_ACCESS));
assertThat(new AssistantApi(config.getRuntime(), QContext.getQSession()).config()).as("somebody who may not use the assistant is told nothing").doesNotContainKey("budgets");
provider.thenText("You have used $3.20 of your $25.00 for today.");
TurnRunner.Outcome outcome = runner.run(EDITOR, "how much of my assistant budget have I used?");
assertThat(outcome.end()).containsEntry("status", "COMPLETE");
String sent = provider.getRequestJson().get(provider.getRequestJson().size() - 1);
assertThat(sent).contains("<assistant_budget").contains("daily budget 'Each person, a day': $3.20 used of $25.00 (12%)")
.doesNotContain("Everyone, a week").doesNotContain("$7.00").doesNotContain("$10.20");
}
/*******************************************************************************
** The budgets block of a config answer.
*******************************************************************************/
@SuppressWarnings("unchecked")
private static List<Map<String, Object>> budgetsOf(Map<String, Object> config)
{
return ((List<Map<String, Object>>) config.get("budgets"));
}
/*******************************************************************************
** A person's own daily limit stops that person, with the sentence, and
** nobody else.
*******************************************************************************/
@Test
void aPersonIsStoppedByTheirOwnLimitAndNobodyElseIs() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.ownerSession(ADMIN, config));
AssistantStore store = config.getRuntime().getStore();
store.insert(new AssistantBudget().withName("Each person, a day").withScopeType("ROLE").withScopeValue("assistant.hasAccess")
.withPeriod("DAY").withMetric("COST_USD").withLimitValue(BigDecimal.valueOf(25)).withAtLimitAction("STOP").withHardCeilingPercent(100));
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
store.insert(new AssistantUsage().withUserId(EDITOR).withPurpose("TURN").withStatus("OK").withCostMicros(25_000_000L).withInputTokens(1L).withOutputTokens(1L).withCreateDate(Instant.now()));
assertThatThrownBy(() -> runner.run(EDITOR, "Hello")).isInstanceOf(AssistantApiException.class)
.hasMessageStartingWith("The assistant has reached the daily limit of $25.00 of the budget 'Each person, a day'. It starts again on ")
.satisfies(e -> assertThat(((AssistantApiException) e).getStatus()).isEqualTo(429));
provider.thenText("Hello.");
assertThat(runner.run(OTHER_EDITOR, "Hello").end()).containsEntry("status", "COMPLETE");
}
}