perf(execution): batch evidence attempt queries

This commit is contained in:
matevip 2026-09-11 04:39:40 -04:00
parent dfae7edbe0
commit 7e36ac7749
4 changed files with 131 additions and 7 deletions

View File

@ -15,6 +15,7 @@ import vip.mate.workspace.core.service.WorkspaceService;
import io.micrometer.core.instrument.MeterRegistry; import io.micrometer.core.instrument.MeterRegistry;
import vip.mate.exception.MateClawException; import vip.mate.exception.MateClawException;
import vip.mate.execution.evidence.model.ExecutionEvidence; import vip.mate.execution.evidence.model.ExecutionEvidence;
import vip.mate.execution.evidence.model.ExecutionAttempt;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.time.Instant; import java.time.Instant;
@ -64,7 +65,10 @@ public class ExecutionEvidenceQueryService {
before.id(), bounded + 1, goalId, teamTaskId); before.id(), bounded + 1, goalId, teamTaskId);
boolean hasMore = rows.size() > bounded; boolean hasMore = rows.size() > bounded;
List<ExecutionEvidence> page = rows.stream().limit(bounded).toList(); List<ExecutionEvidence> page = rows.stream().limit(bounded).toList();
return new Page(page.stream().map(row -> view(username, row)).toList(), if (page.isEmpty()) return new Page(List.of(), null);
var attempts = store.findAttempts(canonicalWorkspace, conversationId,
page.stream().map(ExecutionEvidence::attemptId).distinct().toList());
return new Page(page.stream().map(row -> view(username, row, attempts.get(row.attemptId()))).toList(),
hasMore ? encode(page.getLast()) : null); hasMore ? encode(page.getLast()) : null);
} finally { } finally {
metrics.timer("mateclaw.execution.evidence.query.latency").record( metrics.timer("mateclaw.execution.evidence.query.latency").record(
@ -77,7 +81,7 @@ public class ExecutionEvidenceQueryService {
ExecutionEvidence row = store.findById(id).orElseThrow(this::hidden); ExecutionEvidence row = store.findById(id).orElseThrow(this::hidden);
Long canonicalWorkspace = authorize(username, workspaceId, row.conversationId()); Long canonicalWorkspace = authorize(username, workspaceId, row.conversationId());
if (!canonicalWorkspace.equals(row.workspaceId())) throw hidden(); if (!canonicalWorkspace.equals(row.workspaceId())) throw hidden();
return view(username, row); return view(username, row, store.findAttempt(row.attemptId()).orElseThrow(this::hidden));
} }
private Long authorize(String username, Long workspaceId, String conversationId) { private Long authorize(String username, Long workspaceId, String conversationId) {
@ -90,8 +94,8 @@ public class ExecutionEvidenceQueryService {
return conversation.getWorkspaceId(); return conversation.getWorkspaceId();
} }
private View view(String username, ExecutionEvidence row) { private View view(String username, ExecutionEvidence row, ExecutionAttempt attempt) {
var attempt = store.findAttempt(row.attemptId()).orElseThrow(this::hidden); if (attempt == null || !Objects.equals(attempt.id(), row.attemptId())) throw hidden();
if (!Objects.equals(attempt.identity().workspaceId(), row.workspaceId()) if (!Objects.equals(attempt.identity().workspaceId(), row.workspaceId())
|| !Objects.equals(attempt.identity().conversationId(), row.conversationId())) throw hidden(); || !Objects.equals(attempt.identity().conversationId(), row.conversationId())) throw hidden();
var evidence = row.observation(); var evidence = row.observation();

View File

@ -30,6 +30,8 @@ import java.time.temporal.ChronoUnit;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.LinkedHashMap; import java.util.LinkedHashMap;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.Collections;
import java.util.Objects; import java.util.Objects;
import java.util.Optional; import java.util.Optional;
@ -199,6 +201,26 @@ public class ExecutionEvidenceStore {
return jdbc.query("SELECT * FROM mate_execution_attempt WHERE id=? AND deleted=0", this::attempt, id).stream().findFirst(); return jdbc.query("SELECT * FROM mate_execution_attempt WHERE id=? AND deleted=0", this::attempt, id).stream().findFirst();
} }
/** Load one page of attempts without allowing cross-conversation reads. */
public Map<Long, ExecutionAttempt> findAttempts(Long workspaceId, String conversationId, List<Long> ids) {
scope(workspaceId, conversationId);
Objects.requireNonNull(ids, "Attempt IDs required");
if (ids.size() > properties.getMaxListLimit())
throw new IllegalArgumentException("Attempt batch exceeds page limit");
if (ids.isEmpty()) return Map.of();
if (ids.stream().anyMatch(Objects::isNull))
throw new IllegalArgumentException("Attempt ID required");
var distinct = ids.stream().distinct().toList();
var args = new ArrayList<Object>(List.of(workspaceId, conversationId));
args.addAll(distinct);
String placeholders = String.join(",", Collections.nCopies(distinct.size(), "?"));
var result = new LinkedHashMap<Long, ExecutionAttempt>();
jdbc.query("SELECT * FROM mate_execution_attempt WHERE workspace_id=? AND conversation_id=?"
+ " AND deleted=0 AND id IN (" + placeholders + ")", this::attempt, args.toArray())
.forEach(attempt -> result.put(attempt.id(), attempt));
return result;
}
/** Internal lookup for source authorization; callers must authorize before exposing the result. */ /** Internal lookup for source authorization; callers must authorize before exposing the result. */
public Optional<ExecutionEvidence> findById(Long id) { public Optional<ExecutionEvidence> findById(Long id) {
return jdbc.query(EVIDENCE_QUERY + " AND e.id=?", this::evidence, id).stream().findFirst(); return jdbc.query(EVIDENCE_QUERY + " AND e.id=?", this::evidence, id).stream().findFirst();

View File

@ -17,6 +17,7 @@ import vip.mate.workspace.core.service.WorkspaceService;
import java.time.Instant; import java.time.Instant;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.Optional; import java.util.Optional;
import static org.junit.jupiter.api.Assertions.*; import static org.junit.jupiter.api.Assertions.*;
@ -39,10 +40,12 @@ class ExecutionEvidenceQueryTest {
conversation.setConversationId("conv"); conversation.setWorkspaceId(1L); conversation.setDeleted(0); conversation.setConversationId("conv"); conversation.setWorkspaceId(1L); conversation.setDeleted(0);
when(conversations.findByConversationId("conv")).thenReturn(conversation); when(conversations.findByConversationId("conv")).thenReturn(conversation);
when(conversations.isConversationOwner("conv", "owner")).thenReturn(true); when(conversations.isConversationOwner("conv", "owner")).thenReturn(true);
when(store.findAttempt(1L)).thenReturn(Optional.of(new ExecutionAttempt(1L, var attempt = new ExecutionAttempt(1L,
new ExecutionIdentity(1L, "conv", "native", null, "call", "call", 1, "provider", "tool", new ExecutionIdentity(1L, "conv", "native", null, "call", "call", 1, "provider", "tool",
null, null, null, null, null, null, "fence"), null, null, null, null, null, null, "fence"),
AttemptState.SUCCEEDED, EffectOutcome.UNCERTAIN, now, now))); AttemptState.SUCCEEDED, EffectOutcome.UNCERTAIN, now, now);
when(store.findAttempt(1L)).thenReturn(Optional.of(attempt));
when(store.findAttempts(1L, "conv", List.of(1L))).thenReturn(Map.of(1L, attempt));
} }
@Test void deniesUnscopedAnonymousAndWrongWorkspaceBeforeReadingEvidence() { @Test void deniesUnscopedAnonymousAndWrongWorkspaceBeforeReadingEvidence() {
@ -99,6 +102,39 @@ class ExecutionEvidenceQueryTest {
assertNull(unavailable.artifactDigest()); assertNull(unavailable.artifactDigest());
} }
@Test void emptyPageDoesNotLoadAttempts() {
when(store.list(eq(1L), eq("conv"), isNull(), isNull(), eq(21), isNull(), isNull()))
.thenReturn(List.of());
assertTrue(list("owner", 1L, "conv", null, null).items().isEmpty());
verify(store, never()).findAttempts(any(), any(), any());
verify(store, never()).findAttempt(any());
}
@Test void repeatedAttemptIsLoadedOnceAndLookaheadIsExcluded() {
when(store.list(eq(1L), eq("conv"), isNull(), isNull(), eq(3), isNull(), isNull()))
.thenReturn(List.of(evidence(id), evidence(id - 1),
new ExecutionEvidence(id - 2, 1L, 999L, "conv", evidence(id).observation())));
var page = list("owner", 1L, "conv", null, 2);
assertEquals(List.of(id, id - 1), page.items().stream().map(ExecutionEvidenceQueryService.View::id).toList());
assertNotNull(page.nextCursor());
verify(store).findAttempts(1L, "conv", List.of(1L));
verify(store, never()).findAttempt(any());
}
@Test void missingOrMismatchedBatchAttemptFailsClosed() {
when(store.list(eq(1L), eq("conv"), isNull(), isNull(), eq(21), isNull(), isNull()))
.thenReturn(List.of(evidence(id)));
when(store.findAttempts(1L, "conv", List.of(1L))).thenReturn(Map.of());
assertEquals(404, assertThrows(MateClawException.class,
() -> list("owner", 1L, "conv", null, null)).getCode());
var foreign = new ExecutionAttempt(1L, new ExecutionIdentity(2L, "other", "native", null,
"call", "call", 1, "provider", "tool", null, null, null, null, null, null, "fence"),
AttemptState.SUCCEEDED, EffectOutcome.UNCERTAIN, now, now);
when(store.findAttempts(1L, "conv", List.of(1L))).thenReturn(Map.of(1L, foreign));
assertEquals(404, assertThrows(MateClawException.class,
() -> list("owner", 1L, "conv", null, null)).getCode());
}
private ExecutionEvidenceQueryService.Page list(String user, Long workspace, String conversation, String cursor, Integer limit) { private ExecutionEvidenceQueryService.Page list(String user, Long workspace, String conversation, String cursor, Integer limit) {
return queries.list(user, workspace, conversation, cursor, limit, null, null); return queries.list(user, workspace, conversation, cursor, limit, null, null);
} }

View File

@ -1,6 +1,15 @@
package vip.mate.execution.evidence; package vip.mate.execution.evidence;
import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.BeforeEach;
import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
import vip.mate.execution.evidence.service.ExecutionEvidenceQueryService;
import vip.mate.workspace.conversation.ConversationService;
import vip.mate.workspace.conversation.model.ConversationEntity;
import vip.mate.team.service.TeamWorkerConversationGovernanceService;
import vip.mate.tool.document.GeneratedFileCache;
import vip.mate.auth.service.AuthService;
import vip.mate.workspace.core.service.WorkspaceService;
import static org.mockito.Mockito.*;
import org.h2.jdbcx.JdbcDataSource; import org.h2.jdbcx.JdbcDataSource;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.ParameterizedTest;
@ -26,6 +35,7 @@ import java.time.Instant;
import java.util.concurrent.Callable; import java.util.concurrent.Callable;
import org.springframework.dao.DuplicateKeyException; import org.springframework.dao.DuplicateKeyException;
import java.util.List; import java.util.List;
import java.util.Collections;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.Executors; import java.util.concurrent.Executors;
import static org.assertj.core.api.Assertions.*; import static org.assertj.core.api.Assertions.*;
@ -36,9 +46,61 @@ class ExecutionEvidenceStoreTest {
@BeforeEach void setup() { @BeforeEach void setup() {
var source = new DriverManagerDataSource("jdbc:h2:mem:" + UUID.randomUUID() + ";MODE=MySQL;DB_CLOSE_DELAY=-1", "sa", ""); var source = new DriverManagerDataSource("jdbc:h2:mem:" + UUID.randomUUID() + ";MODE=MySQL;DB_CLOSE_DELAY=-1", "sa", "");
new ResourceDatabasePopulator(new ClassPathResource("db/migration/h2/V191__execution_evidence_ledger.sql")).execute(source); new ResourceDatabasePopulator(new ClassPathResource("db/migration/h2/V191__execution_evidence_ledger.sql")).execute(source);
jdbc = new JdbcTemplate(source); jdbc = spy(new JdbcTemplate(source));
store = new ExecutionEvidenceStore(jdbc, new DataSourceTransactionManager(source), new ExecutionEvidenceProperties()); store = new ExecutionEvidenceStore(jdbc, new DataSourceTransactionManager(source), new ExecutionEvidenceProperties());
} }
@Test void fullPageUsesTwoLedgerQueriesAndPreservesOrder() {
for (int i = 0; i < 100; i++) {
var attempt = store.begin(identity("page-" + i));
store.finish(attempt.id(), "owner", AttemptState.SUCCEEDED, EffectOutcome.UNCERTAIN,
List.of(observation("result-" + i)));
}
var conversations = mock(ConversationService.class);
var conversation = new ConversationEntity();
conversation.setWorkspaceId(1L);
conversation.setConversationId("conversation");
when(conversations.findByConversationId("conversation")).thenReturn(conversation);
when(conversations.isConversationOwner("conversation", "owner")).thenReturn(true);
var queries = new ExecutionEvidenceQueryService(store, conversations,
mock(TeamWorkerConversationGovernanceService.class), mock(GeneratedFileCache.class),
mock(AuthService.class), mock(WorkspaceService.class), new ExecutionEvidenceProperties(),
new SimpleMeterRegistry());
clearInvocations(jdbc);
var page = queries.list("owner", 1L, "conversation", null, 100, null, null);
assertThat(page.items()).hasSize(100);
assertThat(page.nextCursor()).isNull();
assertThat(page.items()).allSatisfy(row -> {
assertThat(row.state()).isEqualTo(AttemptState.SUCCEEDED);
assertThat(row.toolName()).isEqualTo("shell");
});
assertThat(page.items().getFirst().summary()).isEqualTo("result-99");
assertThat(page.items().getLast().summary()).isEqualTo("result-0");
long selects = mockingDetails(jdbc).getInvocations().stream()
.filter(call -> call.getMethod().getName().equals("query") && call.getMethod().isVarArgs())
.filter(call -> call.getArgument(0) instanceof String sql && sql.startsWith("SELECT"))
.count();
assertThat(selects).isEqualTo(2);
}
@Test void batchAttemptsAreScopedBoundedAndExcludeDeletedRows() {
var first = store.begin(identity("first"));
var second = store.begin(identity("second"));
var foreign = store.begin(new ExecutionIdentity(2L, "other", "native", null, "foreign", "foreign",
1, null, "shell", null, null, null, null, null, null, "owner"));
jdbc.update("UPDATE mate_execution_attempt SET deleted=1 WHERE id=?", second.id());
assertThat(store.findAttempts(1L, "conversation", List.of(first.id(), first.id(), second.id(), foreign.id())))
.containsOnlyKeys(first.id());
assertThat(store.findAttempts(1L, "other", List.of(first.id()))).isEmpty();
assertThat(store.findAttempts(2L, "conversation", List.of(first.id()))).isEmpty();
clearInvocations(jdbc);
assertThat(store.findAttempts(1L, "conversation", List.of())).isEmpty();
assertThatThrownBy(() -> store.findAttempts(1L, "conversation", Collections.nCopies(101, first.id())))
.isInstanceOf(IllegalArgumentException.class);
verifyNoInteractions(jdbc);
}
private ExecutionIdentity identity(String invocation) { private ExecutionIdentity identity(String invocation) {
return new ExecutionIdentity(1L,"conversation","native","session",invocation,invocation,1,"provider-id","shell",null,null,null,null,null,null,"owner"); return new ExecutionIdentity(1L,"conversation","native","session",invocation,invocation,1,"provider-id","shell",null,null,null,null,null,null,"owner");
} }