From 1264f4cfed75f1cf7a714fc368fce9e830dd7fa5 Mon Sep 17 00:00:00 2001 From: James Brunton Date: Wed, 3 Jun 2026 12:52:11 +0100 Subject: [PATCH] Set up document management for Stirling Engine (#6476) # Description of Changes Change Stirling Engine to support deleting documents automatically. This happens both on user logout and after an amount of time specified by the Java when ingesting a document (allowing for personal documents to have short lifetimes but org documents to be left in the db with no expiry date). Also sets up an [ACL policy](https://en.wikipedia.org/wiki/Access-control_list) for the documents so the database knows which users have access to which documents. This is not fully implemented in the Java, so currently all docs are treated as having a single owner, the uploader, but theoretically when we need to support org storage, we shouldn't need to change the db schema. --- .../controller/api/AiEngineController.java | 15 +- .../model/api/ai/AiDocumentIngestRequest.java | 13 + .../security/CustomLogoutSuccessHandler.java | 33 ++ .../configuration/SecurityConfiguration.java | 10 +- .../controller/api/AuthController.java | 8 +- .../security/service/JwtService.java | 15 + .../security/service/JwtServiceInterface.java | 13 + .../proprietary/service/AiEngineClient.java | 81 +++- .../service/AiUserDataService.java | 50 +++ .../service/AiWorkflowService.java | 74 +++- .../service/MathAuditorOrchestrator.java | 26 +- .../service/PdfCommentAgentOrchestrator.java | 26 +- .../api/AuthControllerLoginTest.java | 3 +- .../security/service/JwtServiceTest.java | 29 ++ .../service/AiEngineClientTest.java | 8 +- .../service/AiUserDataServiceTest.java | 76 ++++ .../service/AiWorkflowServiceTest.java | 19 +- .../PdfCommentAgentOrchestratorTest.java | 16 +- .../agents/contradiction/capability.py | 5 +- .../stirling/agents/contradiction/detector.py | 15 +- engine/src/stirling/agents/pdf_questions.py | 10 +- engine/src/stirling/agents/pdf_review.py | 13 +- .../agents/shared/whole_doc_reader.py | 5 +- engine/src/stirling/api/app.py | 51 +++ engine/src/stirling/api/dependencies.py | 22 +- engine/src/stirling/api/middleware.py | 3 +- engine/src/stirling/api/routes/documents.py | 60 ++- engine/src/stirling/config/settings.py | 4 + engine/src/stirling/contracts/__init__.py | 2 + engine/src/stirling/contracts/documents.py | 22 +- .../src/stirling/documents/pgvector_store.py | 288 ++++++++++--- .../src/stirling/documents/rag_capability.py | 22 +- engine/src/stirling/documents/service.py | 122 ++++-- .../stirling/documents/sqlite_vec_store.py | 400 ++++++++++++++---- engine/src/stirling/documents/store.py | 108 +++-- engine/src/stirling/models/__init__.py | 5 +- engine/src/stirling/models/base.py | 13 +- engine/src/stirling/services/__init__.py | 4 +- engine/src/stirling/services/runtime.py | 15 +- engine/src/stirling/services/tracking.py | 17 +- engine/tests/agents/test_whole_doc_reader.py | 24 +- engine/tests/contradiction/test_capability.py | 11 +- engine/tests/contradiction/test_detector.py | 41 +- .../test_question_integration.py | 20 +- .../contradiction/test_review_integration.py | 26 +- engine/tests/test_documents.py | 396 ++++++++++++++--- engine/tests/test_documents_routes.py | 183 +++++++- engine/tests/test_pdf_question_agent.py | 28 +- 48 files changed, 2039 insertions(+), 411 deletions(-) create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/AiUserDataService.java create mode 100644 app/proprietary/src/test/java/stirling/software/proprietary/service/AiUserDataServiceTest.java diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java index 28fc7168d..a96b3d124 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java @@ -5,6 +5,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Executor; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; import org.springframework.http.HttpStatus; @@ -30,6 +31,7 @@ import lombok.extern.slf4j.Slf4j; import stirling.software.common.model.job.ResultFile; import stirling.software.common.service.JobOwnershipService; import stirling.software.common.service.TaskManager; +import stirling.software.common.service.UserServiceInterface; import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent; import stirling.software.proprietary.model.api.ai.AiWorkflowRequest; import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; @@ -58,6 +60,7 @@ public class AiEngineController { private final TaskManager taskManager; private final JobOwnershipService jobOwnershipService; private final AiEngineEndpointResolver endpointResolver; + private final UserServiceInterface userService; /** * SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a @@ -74,7 +77,8 @@ public class AiEngineController { @Qualifier("aiStreamExecutor") Executor aiStreamExecutor, TaskManager taskManager, JobOwnershipService jobOwnershipService, - AiEngineEndpointResolver endpointResolver) { + AiEngineEndpointResolver endpointResolver, + @Autowired(required = false) UserServiceInterface userService) { this.aiEngineClient = aiEngineClient; this.aiWorkflowService = aiWorkflowService; this.objectMapper = objectMapper; @@ -82,6 +86,11 @@ public class AiEngineController { this.taskManager = taskManager; this.jobOwnershipService = jobOwnershipService; this.endpointResolver = endpointResolver; + this.userService = userService; + } + + private String currentUserId() { + return userService != null ? userService.getCurrentUsername() : null; } @GetMapping("/health") @@ -89,7 +98,7 @@ public class AiEngineController { summary = "AI engine health check", description = "Returns the health status of the AI engine including configured models") public ResponseEntity health() throws IOException { - String response = aiEngineClient.get("/health"); + String response = aiEngineClient.get("/health", currentUserId()); return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response); } @@ -243,7 +252,7 @@ public class AiEngineController { HttpStatus.BAD_REQUEST, "Request body must be a JSON object"); } String forwardedBody = withEnabledEndpoints((ObjectNode) parsed); - String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody); + String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody, currentUserId()); return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response); } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java index 47c431807..44091056c 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java @@ -1,5 +1,6 @@ package stirling.software.proprietary.model.api.ai; +import java.time.Instant; import java.util.List; import lombok.AllArgsConstructor; @@ -10,6 +11,12 @@ import lombok.NoArgsConstructor; * Body for {@code POST /api/v1/documents} on the AI engine. Sent by Java when the engine reports * {@code need_ingest} and the requested document's extracted content must be stored before the * workflow can continue. + * + *

{@code ownerId} is the tenant the doc belongs to (a user for personal uploads, an org for + * shared content). {@code readPrincipals} is the explicit list of principals granted read access. + * {@code expiresAt} is when the engine's reaper should delete this doc; {@code null} means + * "persistent until explicit delete" (used for org-shared content). Java picks the value per doc; + * the engine does not default it. */ @Data @NoArgsConstructor @@ -21,4 +28,10 @@ public class AiDocumentIngestRequest { private String source; private List pageText; + + private String ownerId; + + private List readPrincipals; + + private Instant expiresAt; } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/CustomLogoutSuccessHandler.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/CustomLogoutSuccessHandler.java index fa44b1337..4bfef06c9 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/CustomLogoutSuccessHandler.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/CustomLogoutSuccessHandler.java @@ -8,6 +8,8 @@ import java.util.List; import java.util.Locale; import org.springframework.core.io.Resource; +import org.springframework.security.authentication.AuthenticationTrustResolver; +import org.springframework.security.authentication.AuthenticationTrustResolverImpl; import org.springframework.security.authentication.UsernamePasswordAuthenticationToken; import org.springframework.security.core.Authentication; import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken; @@ -36,6 +38,7 @@ import stirling.software.proprietary.audit.Audited; import stirling.software.proprietary.security.saml2.CertificateUtils; import stirling.software.proprietary.security.saml2.CustomSaml2AuthenticatedPrincipal; import stirling.software.proprietary.security.service.JwtServiceInterface; +import stirling.software.proprietary.service.AiUserDataService; @Slf4j @RequiredArgsConstructor @@ -49,12 +52,22 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler { private final JwtServiceInterface jwtService; + private final AiUserDataService aiUserDataService; + + private static final AuthenticationTrustResolver TRUST_RESOLVER = + new AuthenticationTrustResolverImpl(); + @Override @Audited(type = AuditEventType.USER_LOGOUT, level = AuditLevel.BASIC) public void onLogoutSuccess( HttpServletRequest request, HttpServletResponse response, Authentication authentication) throws IOException { + String username = resolveUsername(request, authentication); + if (username != null) { + aiUserDataService.purgeUserDocuments(username); + } + if (!response.isCommitted()) { if (authentication != null) { if (authentication instanceof Saml2Authentication samlAuthentication) { @@ -88,6 +101,26 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler { } } + /** + * Pick the right name to purge under. JWT cookie wins if present and parseable; we fall through + * to whatever Spring handed us only when there's no cookie. Spring's anonymous principal is + * filtered out via {@link AuthenticationTrustResolver} so we don't purge under that + * pseudo-user. + */ + private String resolveUsername(HttpServletRequest request, Authentication authentication) { + if (jwtService != null) { + String fromCookie = jwtService.extractUsernameFromRequestAllowExpired(request); + if (fromCookie != null) { + return fromCookie; + } + } + if (authentication == null || TRUST_RESOLVER.isAnonymous(authentication)) { + return null; + } + String name = authentication.getName(); + return (name != null && !name.isBlank()) ? name : null; + } + // Redirect for SAML2 authentication logout private void getRedirect_saml2( HttpServletRequest request, diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/configuration/SecurityConfiguration.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/configuration/SecurityConfiguration.java index b53e3d5b7..3ca384126 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/configuration/SecurityConfiguration.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/configuration/SecurityConfiguration.java @@ -93,6 +93,7 @@ public class SecurityConfiguration { licenseSettingsService; private final ClientRegistrationRepository clientRegistrationRepository; private final PasswordEncoder passwordEncoder; + private final stirling.software.proprietary.service.AiUserDataService aiUserDataService; public SecurityConfiguration( PersistentLoginRepository persistentLoginRepository, @@ -115,7 +116,8 @@ public class SecurityConfiguration { OpenSaml5AuthenticationRequestResolver saml2AuthenticationRequestResolver, @Autowired(required = false) ClientRegistrationRepository clientRegistrationRepository, stirling.software.proprietary.service.UserLicenseSettingsService licenseSettingsService, - PasswordEncoder passwordEncoder) { + PasswordEncoder passwordEncoder, + stirling.software.proprietary.service.AiUserDataService aiUserDataService) { this.userDetailsService = userDetailsService; this.userService = userService; this.loginEnabledValue = loginEnabledValue; @@ -135,6 +137,7 @@ public class SecurityConfiguration { this.clientRegistrationRepository = clientRegistrationRepository; this.licenseSettingsService = licenseSettingsService; this.passwordEncoder = passwordEncoder; + this.aiUserDataService = aiUserDataService; } /** @@ -322,7 +325,10 @@ public class SecurityConfiguration { .matcher("/logout")) .logoutSuccessHandler( new CustomLogoutSuccessHandler( - securityProperties, appConfig, jwtService)) + securityProperties, + appConfig, + jwtService, + aiUserDataService)) .clearAuthentication(true) .invalidateHttpSession(true) .deleteCookies("JSESSIONID", "remember-me", "stirling_jwt")); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AuthController.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AuthController.java index df0817277..490dea255 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AuthController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AuthController.java @@ -44,6 +44,7 @@ import stirling.software.proprietary.security.service.RefreshRateLimitService; import stirling.software.proprietary.security.service.TotpService; import stirling.software.proprietary.security.service.UserService; import stirling.software.proprietary.security.util.DesktopClientUtils; +import stirling.software.proprietary.service.AiUserDataService; /** REST API Controller for authentication operations. */ @RestController @@ -62,6 +63,7 @@ public class AuthController { private final RefreshRateLimitService refreshRateLimitService; private final ApplicationProperties.Security securityProperties; private final ApplicationProperties applicationProperties; + private final AiUserDataService aiUserDataService; /** * Login endpoint - replaces Supabase signInWithPassword @@ -281,11 +283,13 @@ public class AuthController { */ @PreAuthorize("!hasAuthority('ROLE_DEMO_USER')") @PostMapping("/logout") - public ResponseEntity logout(HttpServletResponse response) { + public ResponseEntity logout(HttpServletRequest request, HttpServletResponse response) { try { + String username = jwtService.extractUsernameFromRequestAllowExpired(request); SecurityContextHolder.clearContext(); + aiUserDataService.purgeUserDocuments(username); - log.debug("User logged out successfully"); + log.debug("User logged out successfully (username={})", username); return ResponseEntity.ok(Map.of("message", "Logged out successfully")); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtService.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtService.java index 382d6273f..2c838ee12 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtService.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtService.java @@ -349,6 +349,21 @@ public class JwtService implements JwtServiceInterface { return null; } + @Override + public String extractUsernameFromRequestAllowExpired(HttpServletRequest request) { + try { + String token = extractToken(request); + if (token == null || token.isBlank()) { + return null; + } + String username = extractUsernameAllowExpired(token); + return (username != null && !username.isBlank()) ? username : null; + } catch (Exception e) { + log.debug("Could not extract username from request JWT: {}", e.getMessage()); + return null; + } + } + @Override public boolean isJwtEnabled() { return v2Enabled; diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtServiceInterface.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtServiceInterface.java index cded5b31f..6b6c62f8e 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtServiceInterface.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtServiceInterface.java @@ -93,6 +93,19 @@ public interface JwtServiceInterface { */ String extractToken(HttpServletRequest request); + /** + * Read the username off the request's JWT, allowing an expired token. Returns null when no + * token is present, the token can't be parsed, or the resulting username is blank. + * + *

Used by flows that need to identify the user without depending on {@code + * SecurityContextHolder} - for example logout, where the security filter chain may have left + * the anonymous principal in place by the time the handler runs. + * + * @param request HTTP servlet request + * @return username from the token, or null when one can't be safely derived + */ + String extractUsernameFromRequestAllowExpired(HttpServletRequest request); + /** * Check if JWT authentication is enabled * diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java index f5585ea5b..dedc558af 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java @@ -43,22 +43,23 @@ public class AiEngineClient { this.httpClient = httpClient; } - public String post(String path, String jsonBody) throws IOException { + public String post(String path, String jsonBody, String userId) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); - return postWithTimeout(path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds())); + return postWithTimeout( + path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds()), userId); } /** * POST with an explicit per-call timeout, for heavy operations (e.g. RAG ingestion of a large * document) that legitimately take longer than the default timeout. */ - public String postLongRunning(String path, String jsonBody) throws IOException { + public String postLongRunning(String path, String jsonBody, String userId) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); return postWithTimeout( - path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds())); + path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds()), userId); } - private String postWithTimeout(String path, String jsonBody, Duration timeout) + private String postWithTimeout(String path, String jsonBody, Duration timeout, String userId) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); if (!config.isEnabled()) { @@ -69,22 +70,32 @@ public class AiEngineClient { String url = config.getUrl().stripTrailing() + path; log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds()); - HttpRequest request = + HttpRequest.Builder builder = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Content-Type", "application/json") .header("Accept", "application/json") .timeout(timeout) - .POST(HttpRequest.BodyPublishers.ofString(jsonBody)) - .build(); - - HttpResponse response = sendRequest(request); + .POST(HttpRequest.BodyPublishers.ofString(jsonBody)); + addUserHeader(builder, userId); + HttpResponse response = sendRequest(builder.build()); log.debug("AI engine responded with status {}", response.statusCode()); checkResponseStatus(response); return response.body(); } + /** + * Attach the X-User-Id header so the engine can scope per-user storage (RAG documents, search + * results) to the caller. Skipped when {@code userId} is blank: the engine treats the request + * as anonymous and refuses any route that requires tenancy. + */ + private static void addUserHeader(HttpRequest.Builder builder, String userId) { + if (userId != null && !userId.isBlank()) { + builder.header("X-User-Id", userId); + } + } + /** * POST a JSON body and consume the response as a stream of NDJSON lines. Each line is passed to * {@code lineConsumer} in arrival order; the call returns when the engine closes the stream. @@ -94,7 +105,8 @@ public class AiEngineClient { * practice line arrival keeps the connection logically alive: as long as the engine emits * events, the work is progressing. Genuine engine hangs still hit the total timeout. */ - public void streamPost(String path, String jsonBody, Consumer lineConsumer) + public void streamPost( + String path, String jsonBody, String userId, Consumer lineConsumer) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); if (!config.isEnabled()) { @@ -109,14 +121,15 @@ public class AiEngineClient { url, timeout.toSeconds()); - HttpRequest request = + HttpRequest.Builder builder = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Content-Type", "application/json") .header("Accept", "application/x-ndjson") .timeout(timeout) - .POST(HttpRequest.BodyPublishers.ofString(jsonBody)) - .build(); + .POST(HttpRequest.BodyPublishers.ofString(jsonBody)); + addUserHeader(builder, userId); + HttpRequest request = builder.build(); HttpResponse> response; try { @@ -149,7 +162,36 @@ public class AiEngineClient { } } - public String get(String path) throws IOException { + /** + * DELETE with no body. Used for purging the caller's RAG content on logout. Wraps the same + * error envelope as {@link #post} / {@link #get} so callers see a consistent set of {@code + * ResponseStatusException}s. + */ + public String delete(String path, String userId) throws IOException { + ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); + if (!config.isEnabled()) { + throw new ResponseStatusException( + HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled"); + } + + String url = config.getUrl().stripTrailing() + path; + log.debug("Proxying AI engine DELETE request to {}", url); + + HttpRequest.Builder builder = + HttpRequest.newBuilder() + .uri(URI.create(url)) + .header("Accept", "application/json") + .timeout(Duration.ofSeconds(config.getTimeoutSeconds())) + .DELETE(); + addUserHeader(builder, userId); + HttpResponse response = sendRequest(builder.build()); + + log.debug("AI engine responded with status {}", response.statusCode()); + checkResponseStatus(response); + return response.body(); + } + + public String get(String path, String userId) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); if (!config.isEnabled()) { throw new ResponseStatusException( @@ -159,15 +201,14 @@ public class AiEngineClient { String url = config.getUrl().stripTrailing() + path; log.debug("Proxying AI engine GET request to {}", url); - HttpRequest request = + HttpRequest.Builder builder = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Accept", "application/json") .timeout(Duration.ofSeconds(config.getTimeoutSeconds())) - .GET() - .build(); - - HttpResponse response = sendRequest(request); + .GET(); + addUserHeader(builder, userId); + HttpResponse response = sendRequest(builder.build()); log.debug("AI engine responded with status {}", response.statusCode()); checkResponseStatus(response); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiUserDataService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiUserDataService.java new file mode 100644 index 000000000..0876c734f --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiUserDataService.java @@ -0,0 +1,50 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; + +import org.springframework.scheduling.annotation.Async; +import org.springframework.stereotype.Service; +import org.springframework.web.server.ResponseStatusException; + +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +/** + * Lifecycle hooks for a user's AI document data on the Python engine. + * + *

Today: cleanup on logout. The engine also runs a TTL reaper that catches sessions ended + * without a clean logout (tab close, JWT expiry, engine restart), so this service is the happy-path + * purge, not a hard guarantee. Calls are fire-and-forget on a background thread (Spring's default + * {@code @Async} executor) so an unavailable engine never delays the caller's response. + */ +@Slf4j +@Service +@RequiredArgsConstructor +public class AiUserDataService { + + private static final String PURGE_PATH = "/api/v1/documents/by-owner"; + + private final AiEngineClient aiEngineClient; + + /** + * Tell the engine to delete every collection owned by {@code userId}: vector chunks, page text, + * ACL rows, and the owner row itself. Runs asynchronously so the calling thread (typically a + * logout handler) returns immediately; engine errors are logged on the worker thread and never + * propagated. The engine's TTL reaper backstops any miss within ~24h. + */ + @Async + public void purgeUserDocuments(String userId) { + if (userId == null || userId.isBlank()) { + log.debug("Skipping user document purge: no user id"); + return; + } + try { + aiEngineClient.delete(PURGE_PATH, userId); + log.debug("Requested document purge for user {}", userId); + } catch (ResponseStatusException e) { + log.warn("AI engine refused document purge for {}: {}", userId, e.getReason()); + } catch (IOException e) { + log.warn("Failed to purge documents for {}: {}", userId, e.getMessage()); + } + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java index e219aedd4..8f49694d8 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java @@ -1,6 +1,8 @@ package stirling.software.proprietary.service; import java.io.IOException; +import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; @@ -9,6 +11,7 @@ import java.util.stream.Collectors; import org.apache.commons.io.FilenameUtils; import org.apache.pdfbox.pdmodel.PDDocument; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.core.io.FileSystemResource; import org.springframework.core.io.Resource; import org.springframework.http.HttpHeaders; @@ -24,14 +27,15 @@ import org.springframework.web.multipart.MultipartFile; import io.github.pixee.security.Filenames; import lombok.Data; -import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import stirling.software.common.model.ApplicationProperties; import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.service.FileStorage; import stirling.software.common.service.InternalApiClient; import stirling.software.common.service.InternalApiTimeoutException; import stirling.software.common.service.ToolMetadataService; +import stirling.software.common.service.UserServiceInterface; import stirling.software.common.util.ExceptionUtils; import stirling.software.common.util.TempFile; import stirling.software.common.util.TempFileManager; @@ -49,6 +53,7 @@ import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent; import stirling.software.proprietary.model.api.ai.AiWorkflowRequest; import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile; +import stirling.software.proprietary.security.util.DesktopClientUtils; import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile; import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult; import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact; @@ -59,7 +64,6 @@ import tools.jackson.databind.ObjectMapper; @Slf4j @Service -@RequiredArgsConstructor public class AiWorkflowService { private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents"; @@ -74,6 +78,56 @@ public class AiWorkflowService { private final TempFileManager tempFileManager; private final FileIdStrategy fileIdStrategy; private final AiEngineEndpointResolver endpointResolver; + private final UserServiceInterface userService; + private final ApplicationProperties applicationProperties; + + public AiWorkflowService( + CustomPDFDocumentFactory pdfDocumentFactory, + AiEngineClient aiEngineClient, + PdfContentExtractor pdfContentExtractor, + ObjectMapper objectMapper, + InternalApiClient internalApiClient, + FileStorage fileStorage, + ToolMetadataService toolMetadataService, + TempFileManager tempFileManager, + FileIdStrategy fileIdStrategy, + AiEngineEndpointResolver endpointResolver, + @Autowired(required = false) UserServiceInterface userService, + ApplicationProperties applicationProperties) { + this.pdfDocumentFactory = pdfDocumentFactory; + this.aiEngineClient = aiEngineClient; + this.pdfContentExtractor = pdfContentExtractor; + this.objectMapper = objectMapper; + this.internalApiClient = internalApiClient; + this.fileStorage = fileStorage; + this.toolMetadataService = toolMetadataService; + this.tempFileManager = tempFileManager; + this.fileIdStrategy = fileIdStrategy; + this.endpointResolver = endpointResolver; + this.userService = userService; + this.applicationProperties = applicationProperties; + } + + /** + * How long an AI-workflow-ingested personal doc lives on the engine before the reaper deletes + * it. Mirrors the configured web JWT lifetime, so a stale cookie can never see data the user + * has lost their session to. Org-shared content (when we add it) bypasses this and sends a null + * {@code expiresAt} so it's persistent. + */ + private Duration personalDocTtl() { + int minutes = DesktopClientUtils.getWebTokenExpiryMinutes(applicationProperties); + return Duration.ofMinutes(minutes); + } + + /** + * Resolve the currently-authenticated user's id for X-User-Id propagation to the AI engine. + * Returns null when security is disabled (no UserServiceInterface bean) or no one is logged in. + * The engine rejects per-user routes (ingest, search) when this is null; non-tenant routes + * (health, orchestrate without RAG) still work. + */ + private String currentUserId() { + return userService != null ? userService.getCurrentUsername() : null; + } @FunctionalInterface public interface ProgressListener { @@ -298,10 +352,21 @@ public class AiWorkflowService { } } } + // Personal-doc semantics for AI workflows today: caller owns the doc and is its only + // grantee, with a session-bounded expiry so the reaper cleans up if logout misses. + // When org / shared-doc ingestion lands, the caller chooses owner, grantees, and + // expiry (null = persistent) explicitly. + String callerId = currentUserId(); AiDocumentIngestRequest ingestRequest = - new AiDocumentIngestRequest(file.getId(), file.getName(), pages); + new AiDocumentIngestRequest( + file.getId(), + file.getName(), + pages, + callerId, + callerId == null ? List.of() : List.of(callerId), + Instant.now().plus(personalDocTtl())); String body = objectMapper.writeValueAsString(ingestRequest); - aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body); + aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body, callerId); log.debug( "Ingested document: id={}, name={}, pages={}", file.getId(), @@ -699,6 +764,7 @@ public class AiWorkflowService { aiEngineClient.streamPost( "/api/v1/orchestrator", requestBody, + currentUserId(), line -> handleStreamLine(line, listener, resultHolder, errorHolder)); if (errorHolder[0] != null) { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/MathAuditorOrchestrator.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/MathAuditorOrchestrator.java index 92c595b43..e8c6f7f6b 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/MathAuditorOrchestrator.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/MathAuditorOrchestrator.java @@ -8,13 +8,14 @@ import java.util.List; import java.util.UUID; import org.apache.pdfbox.pdmodel.PDDocument; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; import org.springframework.web.multipart.MultipartFile; -import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import stirling.software.common.service.CustomPDFDocumentFactory; +import stirling.software.common.service.UserServiceInterface; import stirling.software.proprietary.model.api.ai.Evidence; import stirling.software.proprietary.model.api.ai.Folio; import stirling.software.proprietary.model.api.ai.FolioManifest; @@ -40,7 +41,6 @@ import tools.jackson.databind.ObjectMapper; */ @Slf4j @Service -@RequiredArgsConstructor public class MathAuditorOrchestrator { private static final String EXAMINE_PATH = "/api/v1/ai/math-auditor-agent/examine"; @@ -50,6 +50,24 @@ public class MathAuditorOrchestrator { private final CustomPDFDocumentFactory pdfDocumentFactory; private final PdfContentExtractor pdfContentExtractor; private final ObjectMapper objectMapper; + private final UserServiceInterface userService; + + public MathAuditorOrchestrator( + AiEngineClient aiEngineClient, + CustomPDFDocumentFactory pdfDocumentFactory, + PdfContentExtractor pdfContentExtractor, + ObjectMapper objectMapper, + @Autowired(required = false) UserServiceInterface userService) { + this.aiEngineClient = aiEngineClient; + this.pdfDocumentFactory = pdfDocumentFactory; + this.pdfContentExtractor = pdfContentExtractor; + this.objectMapper = objectMapper; + this.userService = userService; + } + + private String currentUserId() { + return userService != null ? userService.getCurrentUsername() : null; + } /** * Run a full math audit against the supplied PDF file. @@ -110,7 +128,7 @@ public class MathAuditorOrchestrator { EXAMINE_PATH, manifest.sessionId(), manifest.round()); - String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody); + String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody, currentUserId()); return objectMapper.readValue(responseBody, Requisition.class); } @@ -123,7 +141,7 @@ public class MathAuditorOrchestrator { evidence.sessionId(), evidence.round(), evidence.finalRound()); - String responseBody = aiEngineClient.post(path, requestBody); + String responseBody = aiEngineClient.post(path, requestBody, currentUserId()); return objectMapper.readValue(responseBody, Verdict.class); } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfCommentAgentOrchestrator.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfCommentAgentOrchestrator.java index cbb0a1026..9fbcda00f 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfCommentAgentOrchestrator.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfCommentAgentOrchestrator.java @@ -10,18 +10,19 @@ import java.util.UUID; import org.apache.commons.io.FilenameUtils; import org.apache.pdfbox.pdmodel.PDDocument; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.http.HttpStatus; import org.springframework.stereotype.Service; import org.springframework.web.multipart.MultipartFile; import org.springframework.web.server.ResponseStatusException; -import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import stirling.software.common.model.api.comments.AnnotationLocation; import stirling.software.common.model.api.comments.StickyNoteSpec; import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.service.PdfAnnotationService; +import stirling.software.common.service.UserServiceInterface; import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineRequest; import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineResponse; import stirling.software.proprietary.model.api.ai.comments.PdfCommentInstruction; @@ -50,7 +51,6 @@ import tools.jackson.databind.ObjectMapper; */ @Slf4j @Service -@RequiredArgsConstructor public class PdfCommentAgentOrchestrator { private static final String GENERATE_PATH = "/api/v1/ai/pdf-comment-agent/generate"; @@ -80,6 +80,26 @@ public class PdfCommentAgentOrchestrator { private final CustomPDFDocumentFactory pdfDocumentFactory; private final ObjectMapper objectMapper; private final PdfAnnotationService pdfAnnotationService; + private final UserServiceInterface userService; + + public PdfCommentAgentOrchestrator( + AiEngineClient aiEngineClient, + PdfTextChunkExtractor pdfTextChunkExtractor, + CustomPDFDocumentFactory pdfDocumentFactory, + ObjectMapper objectMapper, + PdfAnnotationService pdfAnnotationService, + @Autowired(required = false) UserServiceInterface userService) { + this.aiEngineClient = aiEngineClient; + this.pdfTextChunkExtractor = pdfTextChunkExtractor; + this.pdfDocumentFactory = pdfDocumentFactory; + this.objectMapper = objectMapper; + this.pdfAnnotationService = pdfAnnotationService; + this.userService = userService; + } + + private String currentUserId() { + return userService != null ? userService.getCurrentUsername() : null; + } /** * Run the full PDF comment generation flow. @@ -164,7 +184,7 @@ public class PdfCommentAgentOrchestrator { PdfCommentEngineRequest engineRequest = new PdfCommentEngineRequest(sessionId, prompt, chunks); String requestBody = objectMapper.writeValueAsString(engineRequest); - String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody); + String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody, currentUserId()); PdfCommentEngineResponse engineResponse = objectMapper.readValue(responseBody, PdfCommentEngineResponse.class); diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AuthControllerLoginTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AuthControllerLoginTest.java index cf92ef5f3..73b52fcd4 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AuthControllerLoginTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AuthControllerLoginTest.java @@ -79,7 +79,8 @@ class AuthControllerLoginTest { totpService, refreshRateLimitService, securityProperties, - applicationProperties); + applicationProperties, + new stirling.software.proprietary.service.AiUserDataService(null)); mockMvc = MockMvcBuilders.standaloneSetup(controller).build(); } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/security/service/JwtServiceTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/security/service/JwtServiceTest.java index c24680789..8ae6b9d0b 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/security/service/JwtServiceTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/security/service/JwtServiceTest.java @@ -251,6 +251,35 @@ class JwtServiceTest { assertNull(jwtService.extractToken(request)); } + @Test + void testExtractUsernameFromRequestAllowExpiredReturnsUsername() throws Exception { + String username = "alice"; + when(keystoreService.getActiveKey()).thenReturn(testVerificationKey); + when(keystoreService.getKeyPair("test-key-id")).thenReturn(Optional.of(testKeyPair)); + + String token = jwtService.generateToken(username, Collections.emptyMap()); + when(request.getHeader("Authorization")).thenReturn("Bearer " + token); + + assertEquals(username, jwtService.extractUsernameFromRequestAllowExpired(request)); + } + + @Test + void testExtractUsernameFromRequestAllowExpiredReturnsNullWhenNoToken() { + when(request.getHeader("Authorization")).thenReturn(null); + + assertNull(jwtService.extractUsernameFromRequestAllowExpired(request)); + } + + @Test + void testExtractUsernameFromRequestAllowExpiredReturnsNullOnGarbageToken() { + // The logout path depends on this helper returning null - never throwing - so a + // malformed JWT never blocks a user from logging out. Any parse/validation problem + // collapses to "no resolvable user". + when(request.getHeader("Authorization")).thenReturn("Bearer not-a-real-jwt"); + + assertNull(jwtService.extractUsernameFromRequestAllowExpired(request)); + } + @Test void testGenerateTokenWithKeyId() throws Exception { String username = "testuser"; diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineClientTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineClientTest.java index 26b30fa2e..f4c9af592 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineClientTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineClientTest.java @@ -46,7 +46,7 @@ class AiEngineClientTest { when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause); ResponseStatusException ex = - assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); + assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null)); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); assertSame(cause, ex.getCause(), "Original cause should be preserved for diagnostics"); @@ -58,7 +58,7 @@ class AiEngineClientTest { when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause); ResponseStatusException ex = - assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); + assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null)); assertEquals(HttpStatus.GATEWAY_TIMEOUT, ex.getStatusCode()); assertSame(cause, ex.getCause()); @@ -70,7 +70,7 @@ class AiEngineClientTest { when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause); ResponseStatusException ex = - assertThrows(ResponseStatusException.class, () -> client.get("/x")); + assertThrows(ResponseStatusException.class, () -> client.get("/x", null)); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); } @@ -80,7 +80,7 @@ class AiEngineClientTest { applicationProperties.getAiEngine().setEnabled(false); ResponseStatusException ex = - assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); + assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null)); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiUserDataServiceTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiUserDataServiceTest.java new file mode 100644 index 000000000..0d948933d --- /dev/null +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiUserDataServiceTest.java @@ -0,0 +1,76 @@ +package stirling.software.proprietary.service; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import java.io.IOException; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.http.HttpStatus; +import org.springframework.web.server.ResponseStatusException; + +/** + * Tests the wiring of {@link AiUserDataService}: the path/userId combination forwarded to the + * engine, plus the swallowing-not-throwing behaviour for engine failures. Doesn't try to assert + * {@code @Async} dispatch - that's Spring infrastructure and only fires through the proxy, which a + * unit-level test bypasses by design. + */ +@ExtendWith(MockitoExtension.class) +class AiUserDataServiceTest { + + private static final String PURGE_PATH = "/api/v1/documents/by-owner"; + + @Mock private AiEngineClient aiEngineClient; + + private AiUserDataService service; + + @BeforeEach + void setUp() { + service = new AiUserDataService(aiEngineClient); + } + + @Test + void purgesUnderTheGivenUserIdOnTheByOwnerEndpoint() throws IOException { + service.purgeUserDocuments("alice"); + verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice")); + } + + @Test + void noOpWhenUserIdIsNull() throws IOException { + service.purgeUserDocuments(null); + verify(aiEngineClient, never()).delete(eq(PURGE_PATH), eq(null)); + } + + @Test + void noOpWhenUserIdIsBlank() throws IOException { + service.purgeUserDocuments(" "); + verify(aiEngineClient, never()).delete(eq(PURGE_PATH), eq(" ")); + } + + @Test + void swallowsIoExceptionFromEngine() throws IOException { + doThrow(new IOException("connection refused")) + .when(aiEngineClient) + .delete(eq(PURGE_PATH), eq("alice")); + // Must not throw - the logout path depends on this swallowing failures so the user + // can still log out when the engine is unreachable. + service.purgeUserDocuments("alice"); + verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice")); + } + + @Test + void swallowsResponseStatusExceptionFromEngine() throws IOException { + doThrow(new ResponseStatusException(HttpStatus.BAD_GATEWAY, "engine returned 502")) + .when(aiEngineClient) + .delete(eq(PURGE_PATH), eq("alice")); + service.purgeUserDocuments("alice"); + verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice")); + } +} diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java index 5c0939663..5b73ea5cc 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java @@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.nullable; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.never; @@ -117,7 +118,9 @@ class AiWorkflowServiceTest { toolMetadataService, tempFileManager, fileIdStrategy, - endpointResolver); + endpointResolver, + null, + new ApplicationProperties()); when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of()); } @@ -479,18 +482,20 @@ class AiWorkflowServiceTest { {"outcome":"answer","answer":"done","evidence":[]} """; } - Consumer consumer = inv.getArgument(2); + Consumer consumer = inv.getArgument(3); consumer.accept(wrapAsResultEvent(responseJson)); return null; }) .when(aiEngineClient) - .streamPost(eq("/api/v1/orchestrator"), anyString(), any()); + .streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any()); AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this")); assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome()); - verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString()); - verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any()); + verify(aiEngineClient, times(1)) + .postLongRunning(eq("/api/v1/documents"), anyString(), nullable(String.class)); + verify(aiEngineClient, times(2)) + .streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any()); } // --- helpers --- @@ -498,12 +503,12 @@ class AiWorkflowServiceTest { private void stubOrchestrator(String responseJson) throws IOException { doAnswer( inv -> { - Consumer consumer = inv.getArgument(2); + Consumer consumer = inv.getArgument(3); consumer.accept(wrapAsResultEvent(responseJson)); return null; }) .when(aiEngineClient) - .streamPost(eq("/api/v1/orchestrator"), anyString(), any()); + .streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any()); } /** diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/PdfCommentAgentOrchestratorTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/PdfCommentAgentOrchestratorTest.java index c7b97be4b..c30a201c2 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/service/PdfCommentAgentOrchestratorTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/PdfCommentAgentOrchestratorTest.java @@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.nullable; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -72,7 +73,8 @@ class PdfCommentAgentOrchestratorTest { pdfTextChunkExtractor, pdfDocumentFactory, objectMapper, - pdfAnnotationService); + pdfAnnotationService, + null); } @Test @@ -96,7 +98,7 @@ class PdfCommentAgentOrchestratorTest { new PdfCommentInstruction( "p1-c0", "Comment on page 1", null, null)), "reviewed"); - when(aiEngineClient.post(anyString(), anyString())) + when(aiEngineClient.post(anyString(), anyString(), nullable(String.class))) .thenReturn(objectMapper.writeValueAsString(engineResponse)); AnnotatedPdf result = orchestrator.applyComments(input, "please comment"); @@ -130,7 +132,7 @@ class PdfCommentAgentOrchestratorTest { new PdfCommentInstruction("p0-c0", "Valid", null, null), new PdfCommentInstruction("p999-c999", "Bogus", null, null)), "mixed"); - when(aiEngineClient.post(anyString(), anyString())) + when(aiEngineClient.post(anyString(), anyString(), nullable(String.class))) .thenReturn(objectMapper.writeValueAsString(engineResponse)); AnnotatedPdf result = orchestrator.applyComments(input, "test"); @@ -156,7 +158,7 @@ class PdfCommentAgentOrchestratorTest { PdfCommentEngineResponse engineResponse = new PdfCommentEngineResponse("s", List.of(), "no comments worth making"); - when(aiEngineClient.post(anyString(), anyString())) + when(aiEngineClient.post(anyString(), anyString(), nullable(String.class))) .thenReturn(objectMapper.writeValueAsString(engineResponse)); AnnotatedPdf result = orchestrator.applyComments(input, "test"); @@ -184,7 +186,7 @@ class PdfCommentAgentOrchestratorTest { ResponseStatusException.class, () -> orchestrator.applyComments(input, "whatever")); assertEquals(400, ex.getStatusCode().value()); - verify(aiEngineClient, never()).post(anyString(), anyString()); + verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class)); } @Test @@ -197,7 +199,7 @@ class PdfCommentAgentOrchestratorTest { ResponseStatusException.class, () -> orchestrator.applyComments(input, tooLong)); assertEquals(400, ex.getStatusCode().value()); - verify(aiEngineClient, never()).post(anyString(), anyString()); + verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class)); } @Test @@ -209,7 +211,7 @@ class PdfCommentAgentOrchestratorTest { ResponseStatusException.class, () -> orchestrator.applyComments(input, " ")); assertEquals(400, ex.getStatusCode().value()); - verify(aiEngineClient, never()).post(anyString(), anyString()); + verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class)); } // --------------------------------------------------------------------- diff --git a/engine/src/stirling/agents/contradiction/capability.py b/engine/src/stirling/agents/contradiction/capability.py index ad7b4ba21..1f116c006 100644 --- a/engine/src/stirling/agents/contradiction/capability.py +++ b/engine/src/stirling/agents/contradiction/capability.py @@ -21,6 +21,7 @@ from pydantic_ai.toolsets import AbstractToolset from stirling.agents.contradiction.detector import ContradictionDetector from stirling.contracts import AiFile from stirling.contracts.contradiction import Claim, ContradictionReport +from stirling.models import PrincipalId logger = logging.getLogger(__name__) @@ -52,6 +53,7 @@ class ContradictionCapability: self, detector: ContradictionDetector, files: list[AiFile], + principals: list[PrincipalId], *, max_audits: int = DEFAULT_MAX_AUDITS, ) -> None: @@ -59,6 +61,7 @@ class ContradictionCapability: raise ValueError("max_audits must be >= 1") self._detector = detector self._files = files + self._principals = principals self._max_audits = max_audits self._audit_count = 0 toolset: FunctionToolset[None] = FunctionToolset() @@ -121,7 +124,7 @@ class ContradictionCapability: if not self._files: return "No documents attached to audit." - report = await self._detector.detect(self._files, query=query) + report = await self._detector.detect(self._files, principals=self._principals, query=query) formatted = self.format_report(report) logger.info( "[contradiction-capability] audit query=%r files=%d -> %d findings, %d chars", diff --git a/engine/src/stirling/agents/contradiction/detector.py b/engine/src/stirling/agents/contradiction/detector.py index b77ed2575..c075ae342 100644 --- a/engine/src/stirling/agents/contradiction/detector.py +++ b/engine/src/stirling/agents/contradiction/detector.py @@ -41,6 +41,7 @@ from stirling.contracts.contradiction import ( ContradictionSeverity, ) from stirling.contracts.documents import Page +from stirling.models import PrincipalId from stirling.services import AppRuntime logger = logging.getLogger(__name__) @@ -205,8 +206,13 @@ class ContradictionDetector: # Public entry point # ------------------------------------------------------------------ - async def detect(self, files: list[AiFile], query: str | None = None) -> ContradictionReport: - """Run the full pipeline over the supplied files. + async def detect( + self, + files: list[AiFile], + principals: list[PrincipalId], + query: str | None = None, + ) -> ContradictionReport: + """Run the full pipeline over the supplied files for ``principals``. ``files`` must have already been ingested (the caller is responsible for the ``has_collection`` precheck — the question @@ -234,7 +240,7 @@ class ContradictionDetector: effective_query = query or "extract claims" per_file_results = await asyncio.gather( - *(self._extract_claims_for_file(file, effective_query) for file in files), + *(self._extract_claims_for_file(file, effective_query, principals) for file in files), return_exceptions=True, ) @@ -312,6 +318,7 @@ class ContradictionDetector: self, file: AiFile, query: str, + principals: list[PrincipalId], ) -> _FileExtractionResult: """Run the per-chunk extractor over one file's pages. @@ -324,7 +331,7 @@ class ContradictionDetector: ``asyncio.gather`` and the mapper's internal semaphore — this helper itself awaits each step sequentially within one file. """ - file_pages = await self._runtime.documents.read_pages(file.id) + file_pages = await self._runtime.documents.read_pages(file.id, principals=principals) if not file_pages: logger.info( "[contradiction] no stored pages for %s (id=%s); skipping", diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index c1203f6d2..5b956c468 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -27,8 +27,9 @@ from stirling.contracts import ( format_file_names, ) from stirling.documents import RagCapability +from stirling.models import PrincipalId from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams -from stirling.services import AppRuntime +from stirling.services import AppRuntime, require_current_user_id logger = logging.getLogger(__name__) @@ -161,9 +162,10 @@ class PdfQuestionAgent: ) async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]: + principals = [PrincipalId(require_current_user_id())] missing: list[AiFile] = [] for file in files: - if not await self.runtime.documents.has_collection(file.id): + if not await self.runtime.documents.has_collection(file.id, principals=principals): missing.append(file) return missing @@ -175,8 +177,10 @@ class PdfQuestionAgent: upstream classifier keeps that judgement in the same call that writes the answer, and lets the agent mix tools when the question warrants it. """ + principals = [PrincipalId(require_current_user_id())] rag = RagCapability( documents=self.runtime.documents, + principals=principals, collections=[file.id for file in request.files], top_k=self.runtime.settings.rag_default_top_k, max_searches=self.runtime.settings.rag_max_searches, @@ -184,11 +188,13 @@ class PdfQuestionAgent: whole_doc = WholeDocReaderCapability( runtime=self.runtime, files=request.files, + principals=principals, reasoner=self._chunked_reasoner, ) contradiction = ContradictionCapability( detector=self._contradiction_detector, files=request.files, + principals=principals, ) agent = Agent( model=self.runtime.smart_model, diff --git a/engine/src/stirling/agents/pdf_review.py b/engine/src/stirling/agents/pdf_review.py index e62fcb532..5aca029a8 100644 --- a/engine/src/stirling/agents/pdf_review.py +++ b/engine/src/stirling/agents/pdf_review.py @@ -49,14 +49,14 @@ from stirling.contracts import ( Verdict, ) from stirling.contracts.ledger import Discrepancy -from stirling.models import ApiModel, ToolEndpoint +from stirling.models import ApiModel, PrincipalId, ToolEndpoint from stirling.models.agent_tool_models import ( AgentToolId, MathAuditorAgentParams, PdfCommentAgentParams, ) from stirling.models.tool_models import AddCommentsParams -from stirling.services import AppRuntime +from stirling.services import AppRuntime, require_current_user_id # Fallback right-margin placement used when a finding has no usable # anchor text. A4/Letter portrait assumed. @@ -158,7 +158,11 @@ class PdfReviewAgent: files_to_ingest=missing, content_types=[PdfContentType.PAGE_TEXT], ) - report = await self._contradiction_detector.detect(request.files, query=request.user_message) + report = await self._contradiction_detector.detect( + request.files, + principals=[PrincipalId(require_current_user_id())], + query=request.user_message, + ) comments_json = await self._build_contradiction_comments_payload(request.user_message, report) return EditPlanResponse( summary="", @@ -193,9 +197,10 @@ class PdfReviewAgent: ) async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]: + principals = [PrincipalId(require_current_user_id())] missing: list[AiFile] = [] for file in files: - if not await self.runtime.documents.has_collection(file.id): + if not await self.runtime.documents.has_collection(file.id, principals=principals): missing.append(file) return missing diff --git a/engine/src/stirling/agents/shared/whole_doc_reader.py b/engine/src/stirling/agents/shared/whole_doc_reader.py index 324a5f535..ae641137d 100644 --- a/engine/src/stirling/agents/shared/whole_doc_reader.py +++ b/engine/src/stirling/agents/shared/whole_doc_reader.py @@ -19,6 +19,7 @@ from pydantic_ai.toolsets import AbstractToolset from stirling.agents.shared.chunked_reasoner import ChunkedReasoner from stirling.contracts import AiFile +from stirling.models import PrincipalId from stirling.services import AppRuntime logger = logging.getLogger(__name__) @@ -47,12 +48,14 @@ class WholeDocReaderCapability: self, runtime: AppRuntime, files: list[AiFile], + principals: list[PrincipalId], *, reasoner: ChunkedReasoner | None = None, max_reads: int = DEFAULT_MAX_READS, ) -> None: self._runtime = runtime self._files = files + self._principals = principals self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime) self._max_reads = max_reads self._read_count = 0 @@ -115,7 +118,7 @@ class WholeDocReaderCapability: sections: list[str] = [] for file in self._files: - pages = await self._runtime.documents.read_pages(file.id) + pages = await self._runtime.documents.read_pages(file.id, principals=self._principals) if not pages: logger.info( "[whole-doc-reader] no stored pages for %s (id=%s); skipping", diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index 3539379bb..1cf4b76c3 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -1,5 +1,7 @@ from __future__ import annotations +import asyncio +import logging from contextlib import asynccontextmanager from typing import Annotated @@ -29,8 +31,45 @@ from stirling.api.routes import ( ) from stirling.config import AppSettings, load_settings from stirling.contracts import HealthResponse +from stirling.documents import DocumentService from stirling.services import build_runtime, setup_posthog_tracking +logger = logging.getLogger(__name__) + + +async def _run_expired_doc_reaper( + documents: DocumentService, + interval_seconds: int, +) -> None: + """Periodically delete documents whose ``expires_at`` has passed. + + A reaped collection drops everything rooted at that document. Backstop + for the explicit logout purge: catches sessions that ended without a + clean logout (tab close, JWT expiry, engine restart). Persistent rows + (``expires_at`` null, the shape we use for org-shared docs) are never + touched. Runs until cancelled by the lifespan teardown. + """ + await _reap(documents) + while True: + await asyncio.sleep(interval_seconds) + await _reap(documents) + + +async def _reap(documents: DocumentService) -> None: + """One reaper iteration. Logs the deleted count on success and the full + exception with traceback on failure; never re-raises non-cancel errors so + a bad iteration doesn't kill the loop. ``asyncio.CancelledError`` is + re-raised so the lifespan teardown can cancel the task cleanly. + """ + try: + deleted = await documents.reap_expired() + if deleted: + logger.info("Reaped %d expired document collection(s)", deleted) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Document reaper iteration failed; will retry on next interval") + def _load_startup_settings(fast_api: FastAPI) -> AppSettings: override = fast_api.dependency_overrides.get(load_settings) @@ -56,7 +95,19 @@ async def lifespan(fast_api: FastAPI): tracer_provider = setup_posthog_tracking(settings) if tracer_provider: Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider)) + reaper_task = asyncio.create_task( + _run_expired_doc_reaper( + runtime.documents, + interval_seconds=settings.documents_reaper_interval_seconds, + ), + name="expired-document-reaper", + ) yield + reaper_task.cancel() + try: + await reaper_task + except asyncio.CancelledError: + pass await runtime.documents.close() if tracer_provider: tracer_provider.shutdown() diff --git a/engine/src/stirling/api/dependencies.py b/engine/src/stirling/api/dependencies.py index 70cbad877..cae20a8ae 100644 --- a/engine/src/stirling/api/dependencies.py +++ b/engine/src/stirling/api/dependencies.py @@ -1,6 +1,6 @@ from __future__ import annotations -from fastapi import Request +from fastapi import HTTPException, Request, status from stirling.agents import ( ExecutionPlanningAgent, @@ -12,7 +12,8 @@ from stirling.agents import ( from stirling.agents.ledger import MathAuditorAgent from stirling.agents.pdf_comment import PdfCommentAgent from stirling.documents import DocumentService -from stirling.services import AppRuntime +from stirling.models import UserId +from stirling.services import AppRuntime, current_user_id def get_runtime(request: Request) -> AppRuntime: @@ -49,3 +50,20 @@ def get_math_auditor_agent(request: Request) -> MathAuditorAgent: def get_pdf_comment_agent(request: Request) -> PdfCommentAgent: return request.app.state.pdf_comment_agent + + +def require_user_id() -> UserId: + """FastAPI dependency for routes that touch per-user storage. + + Reads ``X-User-Id`` (already extracted into a ContextVar by ``UserIdMiddleware``) + and returns it. Returns HTTP 401 if the caller didn't supply the header. Apply + to any route that ingests, searches, reads, or deletes document content so + the tenancy gate is enforced at the API boundary. + """ + user_id = current_user_id.get() + if user_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="X-User-Id header is required", + ) + return user_id diff --git a/engine/src/stirling/api/middleware.py b/engine/src/stirling/api/middleware.py index 8d5be8b03..30e329623 100644 --- a/engine/src/stirling/api/middleware.py +++ b/engine/src/stirling/api/middleware.py @@ -4,6 +4,7 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin from starlette.requests import Request from starlette.responses import Response +from stirling.models import UserId from stirling.services.tracking import current_user_id _USER_ID_HEADER = "X-User-Id" @@ -15,7 +16,7 @@ class UserIdMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: user_id = request.headers.get(_USER_ID_HEADER) if user_id: - token = current_user_id.set(user_id) + token = current_user_id.set(UserId(user_id)) try: return await call_next(request) finally: diff --git a/engine/src/stirling/api/routes/documents.py b/engine/src/stirling/api/routes/documents.py index 4887b7ac3..a9c863d12 100644 --- a/engine/src/stirling/api/routes/documents.py +++ b/engine/src/stirling/api/routes/documents.py @@ -1,49 +1,83 @@ from __future__ import annotations +import logging from typing import Annotated from fastapi import APIRouter, Depends -from stirling.api.dependencies import get_document_service +from stirling.api.dependencies import get_document_service, require_user_id from stirling.contracts import ( DeleteDocumentResponse, IngestDocumentRequest, IngestDocumentResponse, ) +from stirling.contracts.documents import PurgeOwnerResponse from stirling.documents import DocumentService -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId, UserId router = APIRouter(prefix="/api/v1/documents", tags=["documents"]) +logger = logging.getLogger(__name__) @router.post("", response_model=IngestDocumentResponse) async def ingest_document( request: IngestDocumentRequest, documents: Annotated[DocumentService, Depends(get_document_service)], + _user_id: Annotated[UserId, Depends(require_user_id)], ) -> IngestDocumentResponse: - """Replace-ingest a document's content under ``document_id``. + """Replace-ingest a document's content under ``(owner_id, read_principals)``. - Stores both representations in one shot: - * embedded chunks for RAG search, - * ordered page text for whole-document reading. - Any previously-stored content for this document is removed first. + ``owner_id`` and ``read_principals`` are required on the request body and + are stored verbatim. + + ``X-User-Id`` is still required so the caller is authenticated and + PostHog tracks the right user, but the value is not used to constrain + the body. """ - pages = request.page_text or [] chunks_indexed = await documents.ingest( collection=request.document_id, - pages=pages, + pages=request.page_text or [], source=request.source, + owner_id=request.owner_id, + read_principals=request.read_principals, + expires_at=request.expires_at, ) return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed) -@router.delete("/{document_id}", response_model=DeleteDocumentResponse) +@router.delete("/by-id/{document_id}", response_model=DeleteDocumentResponse) async def delete_document( document_id: FileId, documents: Annotated[DocumentService, Depends(get_document_service)], + user_id: Annotated[UserId, Depends(require_user_id)], ) -> DeleteDocumentResponse: - """Remove a document's content. Idempotent.""" - existed = await documents.has_collection(document_id) + """Remove the caller's copy of this document. + + Owner is inferred from the caller - only personal-doc deletes go through + here. Org-doc deletes will need an explicit owner_id once we add the + admin endpoints for them; for now this route can't reach docs owned by + a different principal. + """ + owner_id = OwnerId(user_id) + principals = [PrincipalId(user_id)] + existed = await documents.has_collection(document_id, principals=principals) if existed: - await documents.delete_collection(document_id) + await documents.delete_collection(document_id, owner_id=owner_id) return DeleteDocumentResponse(document_id=document_id, deleted=existed) + + +@router.delete("/by-owner", response_model=PurgeOwnerResponse) +async def purge_caller_documents( + documents: Annotated[DocumentService, Depends(get_document_service)], + user_id: Annotated[UserId, Depends(require_user_id)], +) -> PurgeOwnerResponse: + """Delete every personal-doc collection owned by the caller. + + Called by Java on logout so a user's document content disappears as soon as + the session ends. Org-owned docs (where the caller is a reader but not + the owner) are not touched - only collections whose ``owner_id`` matches + the calling user are removed. + """ + deleted = await documents.purge_owner(OwnerId(user_id)) + logger.info("Purged %d collection(s) for owner=%s", deleted, user_id) + return PurgeOwnerResponse(owner_id=OwnerId(user_id), deleted=deleted) diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index 37bce1637..bd0499e6b 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -37,6 +37,10 @@ class AppSettings(BaseSettings): rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES") + documents_reaper_interval_seconds: int = Field( + default=900, + validation_alias="STIRLING_DOCUMENTS_REAPER_INTERVAL_SECONDS", + ) # Chunked reasoner settings (whole-document map-reduce). chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE") diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index f0105ffe2..407f2ca51 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -42,6 +42,7 @@ from .documents import ( Page, PageRange, PageText, + PurgeOwnerResponse, ) from .execution import ( AgentExecutionRequest, @@ -137,6 +138,7 @@ __all__ = [ "ContradictionSeverity", "ConversationMessage", "DeleteDocumentResponse", + "PurgeOwnerResponse", "PdfToMarkdownCannotDoResponse", "PdfToMarkdownOrchestrateResponse", "PdfToMarkdownRequest", diff --git a/engine/src/stirling/contracts/documents.py b/engine/src/stirling/contracts/documents.py index 23e2ee8f4..99b5044e5 100644 --- a/engine/src/stirling/contracts/documents.py +++ b/engine/src/stirling/contracts/documents.py @@ -1,8 +1,10 @@ from __future__ import annotations +from datetime import datetime + from pydantic import Field -from stirling.models import ApiModel +from stirling.models import ApiModel, OwnerId, PrincipalId from .common import FileId @@ -35,20 +37,27 @@ class PageRange(ApiModel): class IngestDocumentRequest(ApiModel): - """Replace-ingest a document's content under the given ``document_id``. + """Replace-ingest a document's content under ``(document_id, owner_id)``. - Each call wipes any previously-stored content for the document and writes + Each call wipes any previously-stored content for the pair and writes both the vector-chunk and ordered-page representations from the supplied pages. ``source`` is a human-readable label (typically the original filename) that flows into chunk metadata so search results are readable when ``document_id`` is a hash. + + ``owner_id`` and ``read_principals`` are required: the engine never + defaults them. Callers must declare ownership and access explicitly. """ document_id: FileId = Field(min_length=1) source: str = Field(min_length=1) page_text: list[PageText] | None = None + owner_id: OwnerId = Field(min_length=1) + read_principals: list[PrincipalId] = Field(min_length=1) + # When to delete this doc. ``None`` means "persistent" (keep until an explicit delete) + expires_at: datetime | None class IngestDocumentResponse(ApiModel): @@ -59,3 +68,10 @@ class IngestDocumentResponse(ApiModel): class DeleteDocumentResponse(ApiModel): document_id: FileId deleted: bool + + +class PurgeOwnerResponse(ApiModel): + """Returned by ``DELETE /api/v1/documents/by-owner``.""" + + owner_id: OwnerId + deleted: int = Field(ge=0) diff --git a/engine/src/stirling/documents/pgvector_store.py b/engine/src/stirling/documents/pgvector_store.py index c79afafe4..8aff5ad3a 100644 --- a/engine/src/stirling/documents/pgvector_store.py +++ b/engine/src/stirling/documents/pgvector_store.py @@ -1,24 +1,34 @@ from __future__ import annotations import json +from datetime import datetime import psycopg from pgvector.psycopg import register_vector_async from stirling.contracts.documents import Page, PageRange from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage +from stirling.models import OwnerId, PrincipalId + +_READ_PERMISSION = "read" class PgVectorStore(DocumentStore): - """PostgreSQL + pgvector backed store. + """PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads. Connects to an external Postgres instance (DSN provided via config) and uses the `vector` extension for similarity search. The schema is created on first use. - Holds two tables under the same connection: + Owned tables: + * ``documents_meta`` - parent row, one per ``(collection, owner_id)``. * ``rag_documents`` - vector chunks for RAG search. * ``document_pages`` - ordered page text for whole-document reading. + * ``document_acl`` - principals (users, groups, orgs) granted permissions + on a ``(collection, owner_id)`` pair. Read methods take the caller's + principal set and join through this table. + + Writes are owner-scoped; reads are ACL-scoped. """ def __init__(self, dsn: str) -> None: @@ -41,60 +51,141 @@ class PgVectorStore(DocumentStore): await cur.execute( """ CREATE TABLE IF NOT EXISTS documents_meta ( - collection TEXT PRIMARY KEY, - source TEXT NOT NULL + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, + source TEXT NOT NULL, + expires_at TIMESTAMP WITH TIME ZONE, + PRIMARY KEY (collection, owner_id) ) """ ) + # Partial index over rows that can actually expire keeps the reaper + # scan tight even when most rows are persistent (org docs). + await cur.execute( + "CREATE INDEX IF NOT EXISTS idx_meta_expires_at " + "ON documents_meta(expires_at) WHERE expires_at IS NOT NULL" + ) await cur.execute( """ CREATE TABLE IF NOT EXISTS rag_documents ( id TEXT NOT NULL, - collection TEXT NOT NULL - REFERENCES documents_meta(collection) ON DELETE CASCADE, + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, text TEXT NOT NULL, metadata JSONB NOT NULL DEFAULT '{}'::jsonb, embedding vector NOT NULL, - PRIMARY KEY (id, collection) + PRIMARY KEY (id, collection, owner_id), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) - await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)") + await cur.execute( + "CREATE INDEX IF NOT EXISTS idx_rag_collection_owner ON rag_documents(collection, owner_id)" + ) await cur.execute( """ CREATE TABLE IF NOT EXISTS document_pages ( - collection TEXT NOT NULL - REFERENCES documents_meta(collection) ON DELETE CASCADE, + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, page_number INTEGER NOT NULL, text TEXT NOT NULL, char_count INTEGER NOT NULL, - PRIMARY KEY (collection, page_number) + PRIMARY KEY (collection, owner_id, page_number), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) - await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)") + await cur.execute( + "CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)" + ) + await cur.execute( + """ + CREATE TABLE IF NOT EXISTS document_acl ( + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, + principal_id TEXT NOT NULL, + permission TEXT NOT NULL, + PRIMARY KEY (collection, owner_id, principal_id, permission), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE + ) + """ + ) + # Hot path: every read joins ``WHERE principal_id = ANY(...) AND permission = ?``. + await cur.execute( + "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)" + ) await conn.commit() self._initialized = True - async def ensure_collection(self, collection: str, source: str) -> None: + # ── lifecycle of the (collection, owner_id) row ──────────────────────── + + async def ensure_collection( + self, + collection: str, + source: str, + owner_id: OwnerId, + expires_at: datetime | None, + ) -> None: await self._ensure_schema() async with await self._connect() as conn: async with conn.cursor() as cur: await cur.execute( """ - INSERT INTO documents_meta (collection, source) - VALUES (%s, %s) - ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source + INSERT INTO documents_meta (collection, owner_id, source, expires_at) + VALUES (%s, %s, %s, %s) + ON CONFLICT (collection, owner_id) DO UPDATE SET + source = EXCLUDED.source, + expires_at = EXCLUDED.expires_at """, - (collection, source), + (collection, owner_id, source, expires_at), ) await conn.commit() + async def purge_owner(self, owner_id: OwnerId) -> int: + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + await cur.execute( + "DELETE FROM documents_meta WHERE owner_id = %s", + (owner_id,), + ) + deleted = cur.rowcount + await conn.commit() + return deleted + + async def reap_expired(self) -> int: + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()") + deleted = cur.rowcount + await conn.commit() + return deleted + + async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + # Cascade FKs handle rag_documents, document_pages, document_acl. + await cur.execute( + "DELETE FROM documents_meta WHERE collection = %s AND owner_id = %s", + (collection, owner_id), + ) + deleted = cur.rowcount > 0 + await conn.commit() + return deleted + + # ── write paths ──────────────────────────────────────────────────────── + async def add_documents( self, collection: str, documents: list[Document], embeddings: list[list[float]], + owner_id: OwnerId, ) -> None: if len(documents) != len(embeddings): raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings") @@ -107,36 +198,100 @@ class PgVectorStore(DocumentStore): for doc, emb in zip(documents, embeddings): await cur.execute( """ - INSERT INTO rag_documents (id, collection, text, metadata, embedding) - VALUES (%s, %s, %s, %s::jsonb, %s) - ON CONFLICT (id, collection) + INSERT INTO rag_documents (id, collection, owner_id, text, metadata, embedding) + VALUES (%s, %s, %s, %s, %s::jsonb, %s) + ON CONFLICT (id, collection, owner_id) DO UPDATE SET text = EXCLUDED.text, metadata = EXCLUDED.metadata, embedding = EXCLUDED.embedding """, - (doc.id, collection, doc.text, json.dumps(doc.metadata), emb), + (doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), emb), ) await conn.commit() + async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + await cur.execute( + "DELETE FROM document_pages WHERE collection = %s AND owner_id = %s", + (collection, owner_id), + ) + if pages: + await cur.executemany( + """ + INSERT INTO document_pages (collection, owner_id, page_number, text, char_count) + VALUES (%s, %s, %s, %s, %s) + """, + [(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages], + ) + await conn.commit() + + # ── ACL management ───────────────────────────────────────────────────── + + async def grant_read( + self, + collection: str, + owner_id: OwnerId, + principals: list[PrincipalId], + ) -> None: + if not principals: + return + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + await cur.executemany( + """ + INSERT INTO document_acl (collection, owner_id, principal_id, permission) + VALUES (%s, %s, %s, %s) + ON CONFLICT (collection, owner_id, principal_id, permission) DO NOTHING + """, + [(collection, owner_id, p, _READ_PERMISSION) for p in principals], + ) + await conn.commit() + + async def revoke( + self, + collection: str, + owner_id: OwnerId, + principal: PrincipalId, + ) -> None: + await self._ensure_schema() + async with await self._connect() as conn: + async with conn.cursor() as cur: + await cur.execute( + "DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s", + (collection, owner_id, principal), + ) + await conn.commit() + + # ── read paths (ACL-gated) ───────────────────────────────────────────── + async def search( self, collection: str, query_embedding: list[float], - top_k: int = 5, + top_k: int, + principals: list[PrincipalId], ) -> list[SearchResult]: + if not principals: + return [] await self._ensure_schema() async with await self._connect() as conn: async with conn.cursor() as cur: + owner_id = await self._readable_owner_for(cur, collection, principals) + if owner_id is None: + return [] await cur.execute( """ SELECT id, text, metadata, 1 - (embedding <=> %s) AS score FROM rag_documents - WHERE collection = %s + WHERE collection = %s AND owner_id = %s ORDER BY embedding <=> %s LIMIT %s """, - (query_embedding, collection, query_embedding, top_k), + (query_embedding, collection, owner_id, query_embedding, top_k), ) rows = await cur.fetchall() @@ -148,71 +303,84 @@ class PgVectorStore(DocumentStore): for r in rows ] - async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: - await self._ensure_schema() - async with await self._connect() as conn: - async with conn.cursor() as cur: - await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,)) - if pages: - await cur.executemany( - """ - INSERT INTO document_pages (collection, page_number, text, char_count) - VALUES (%s, %s, %s, %s) - """, - [(collection, p.page_number, p.text, p.char_count) for p in pages], - ) - await conn.commit() + @staticmethod + async def _readable_owner_for( + cur: psycopg.AsyncCursor, + collection: str, + principals: list[PrincipalId], + ) -> str | None: + """Resolve which owner_id this caller is reading. ``None`` means no access.""" + await cur.execute( + """ + SELECT owner_id FROM document_acl + WHERE collection = %s + AND permission = %s + AND principal_id = ANY(%s) + ORDER BY owner_id + LIMIT 1 + """, + (collection, _READ_PERMISSION, list(principals)), + ) + row = await cur.fetchone() + return row[0] if row else None async def read_pages( self, collection: str, - page_range: PageRange | None = None, + page_range: PageRange | None, + principals: list[PrincipalId], ) -> list[Page]: + if not principals: + return [] await self._ensure_schema() async with await self._connect() as conn: async with conn.cursor() as cur: + owner_id = await self._readable_owner_for(cur, collection, principals) + if owner_id is None: + return [] if page_range is None: await cur.execute( "SELECT page_number, text, char_count FROM document_pages " - "WHERE collection = %s ORDER BY page_number", - (collection,), + "WHERE collection = %s AND owner_id = %s ORDER BY page_number", + (collection, owner_id), ) else: await cur.execute( "SELECT page_number, text, char_count FROM document_pages " - "WHERE collection = %s AND page_number BETWEEN %s AND %s " + "WHERE collection = %s AND owner_id = %s " + "AND page_number BETWEEN %s AND %s " "ORDER BY page_number", - (collection, page_range.start, page_range.end), + (collection, owner_id, page_range.start, page_range.end), ) rows = await cur.fetchall() return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows] - async def delete_collection(self, collection: str) -> None: + async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: + if not principals: + return False await self._ensure_schema() async with await self._connect() as conn: async with conn.cursor() as cur: - # Cascade FKs handle rag_documents and document_pages. - await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,)) - await conn.commit() + owner_id = await self._readable_owner_for(cur, collection, principals) + return owner_id is not None - async def list_collections(self) -> list[str]: - await self._ensure_schema() - async with await self._connect() as conn: - async with conn.cursor() as cur: - await cur.execute("SELECT collection FROM documents_meta ORDER BY collection") - rows = await cur.fetchall() - return [r[0] for r in rows] - - async def has_collection(self, collection: str) -> bool: + async def list_collections(self, principals: list[PrincipalId]) -> list[str]: + if not principals: + return [] await self._ensure_schema() async with await self._connect() as conn: async with conn.cursor() as cur: await cur.execute( - "SELECT 1 FROM documents_meta WHERE collection = %s", - (collection,), + """ + SELECT DISTINCT collection FROM document_acl + WHERE permission = %s + AND principal_id = ANY(%s) + ORDER BY collection + """, + (_READ_PERMISSION, list(principals)), ) - row = await cur.fetchone() - return row is not None + rows = await cur.fetchall() + return [r[0] for r in rows] async def close(self) -> None: # Connections are opened and closed per call, so nothing persistent to release. diff --git a/engine/src/stirling/documents/rag_capability.py b/engine/src/stirling/documents/rag_capability.py index bed22e94d..55efc99cd 100644 --- a/engine/src/stirling/documents/rag_capability.py +++ b/engine/src/stirling/documents/rag_capability.py @@ -8,7 +8,7 @@ from pydantic_ai.toolsets import AbstractToolset from stirling.documents.service import DocumentService from stirling.documents.store import SearchResult -from stirling.models import FileId +from stirling.models import FileId, PrincipalId logger = logging.getLogger(__name__) @@ -16,9 +16,13 @@ logger = logging.getLogger(__name__) class RagCapability: """Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection. - Agents consume this as:: + Agents construct one per run, scoped to the caller's principal set, as:: - rag = runtime.rag_capability + rag = RagCapability( + documents=self.runtime.documents, + principals=current_principals(), + collections=[file.id for file in request.files], + ) Agent( ..., instructions=rag.instructions, @@ -26,20 +30,22 @@ class RagCapability: ) When no collections are pinned, the instructions are generated dynamically at - run time so the agent sees the current list of collections in the store. + run time so the agent sees the collections the caller can read. Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a - single agent run. + single agent run and binds to one caller's principal set. """ def __init__( self, documents: DocumentService, + principals: list[PrincipalId], collections: list[FileId] | None = None, top_k: int = 5, max_searches: int = 5, ) -> None: self._documents = documents + self._principals = principals self._collections = collections self._top_k = top_k self._max_searches = max_searches @@ -74,7 +80,7 @@ class RagCapability: ) async def _dynamic_instructions(self) -> str: - collections = await self._documents.list_collections() + collections = await self._documents.list_collections(self._principals) if collections: names = ", ".join(collections) collection_desc = f"the following knowledge base collections: {names}" @@ -115,12 +121,12 @@ class RagCapability: if self._collections: all_results = [] for col in self._collections: - col_results = await self._documents.search(query, collection=col, top_k=k) + col_results = await self._documents.search(query, principals=self._principals, collection=col, top_k=k) all_results.extend(col_results) all_results.sort(key=lambda r: r.score, reverse=True) results = all_results[:k] else: - results = await self._documents.search(query, top_k=k) + results = await self._documents.search(query, principals=self._principals, top_k=k) if not results: logger.info("[rag] search_knowledge query=%r -> 0 results", query) diff --git a/engine/src/stirling/documents/service.py b/engine/src/stirling/documents/service.py index e542ffa58..262f5f7ff 100644 --- a/engine/src/stirling/documents/service.py +++ b/engine/src/stirling/documents/service.py @@ -1,11 +1,12 @@ from __future__ import annotations import logging +from datetime import datetime from stirling.contracts.documents import Page, PageRange, PageText from stirling.documents.embedder import EmbeddingService from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId logger = logging.getLogger(__name__) @@ -17,7 +18,8 @@ PAGE_TEXT_CONTENT_TYPE = "page_text" class DocumentService: """Top-level facade for stored document content. - Holds two representations of every document under a single ``collection``: + Holds two representations of every document under a single + ``(collection, owner_id)`` pair: * **Vector chunks** for RAG-style semantic retrieval (``search``). * **Ordered pages** for whole-document reading (``read_pages``). @@ -25,6 +27,15 @@ class DocumentService: Both are populated by :meth:`ingest` from a single ``pages`` payload. Agents pick the strategy that fits the question; they don't need to know which storage they're hitting. + + **Owner vs principal.** ``owner_id`` is the tenant (a user or an org). + ``principals`` is the caller's accessible principal set, matched against + the ACL on every read. Personal-doc behaviour: ingest with + ``owner_id=user:bob`` and the route grants ``user:bob`` read access; Bob's + later searches pass ``principals=[user:bob]`` and find his docs. + Org-doc behaviour: ingest with ``owner_id=org:acme`` and grant read to + whichever groups should see it (``group:engineering``, etc.); members' + principal sets pick the doc up automatically. """ def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None: @@ -37,20 +48,27 @@ class DocumentService: collection: FileId, pages: list[PageText], source: str, + owner_id: OwnerId, + read_principals: list[PrincipalId], + expires_at: datetime | None, ) -> int: """Replace-ingest a document. Returns the number of vector chunks indexed. - This wipes any previously-stored content for ``collection`` and writes - both the vector-chunk and page-text representations from the same - ``pages`` payload. Pages with empty/whitespace-only text are skipped - for chunking but still written to the page store so page numbering is - preserved end-to-end. + Wipes any previously-stored content for ``(collection, owner_id)`` and + writes both the vector-chunk and page-text representations from the + same ``pages`` payload. Pages with empty/whitespace-only text are + skipped for chunking but still written to the page store so page + numbering is preserved end-to-end. """ - await self._store.delete_collection(collection) - await self._store.ensure_collection(collection, source) + if not read_principals: + raise ValueError("read_principals must not be empty - every doc needs at least one reader") + + await self._store.delete_collection(collection, owner_id) + await self._store.ensure_collection(collection, source, owner_id, expires_at) stored_pages = [StoredPage(page_number=p.page_number, text=p.text, char_count=len(p.text)) for p in pages] - await self._store.add_pages(collection, stored_pages) + await self._store.add_pages(collection, stored_pages, owner_id) + await self._store.grant_read(collection, owner_id, read_principals) chunks: list[Document] = [] for page in pages: @@ -70,63 +88,103 @@ class DocumentService: if not chunks: return 0 embeddings = await self._embedder.embed_documents([doc.text for doc in chunks]) - await self._store.add_documents(collection, chunks, embeddings) + await self._store.add_documents(collection, chunks, embeddings, owner_id) return len(chunks) async def search( self, query: str, + principals: list[PrincipalId], collection: FileId | None = None, top_k: int | None = None, ) -> list[SearchResult]: - """Embed query and search across one or all collections. + """Embed query and search collections readable by ``principals``. - If collection is None, searches all available collections and merges results. + If ``collection`` is supplied, search only that collection (and only + if at least one principal can read it). If ``None``, search every + collection any principal in ``principals`` can read, merge, and + return the top-k. """ k = top_k if top_k is not None else self._default_top_k query_embedding = await self._embedder.embed_query(query) if collection is not None: - if not await self._store.has_collection(collection): + if not await self._store.has_collection(collection, principals): return [] - return await self._store.search(collection, query_embedding, k) + return await self._store.search(collection, query_embedding, k, principals) - # Search all collections, skipping any that error (e.g. dimension mismatch) - collections = await self._store.list_collections() + # Search every collection the caller can read, skipping any that error + # (e.g. dimension mismatch). + collections = await self._store.list_collections(principals) all_results: list[SearchResult] = [] for col_name in collections: try: - results = await self._store.search(col_name, query_embedding, k) + results = await self._store.search(col_name, query_embedding, k, principals) all_results.extend(results) except Exception: # noqa: BLE001 - any backend error on one collection should not stop the others - logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True) + logger.warning( + "Skipping collection %s during cross-collection search", + col_name, + exc_info=True, + ) - # Sort by score descending, return top_k across all collections + # Sort by score descending, return top_k across all the caller's collections all_results.sort(key=lambda r: r.score, reverse=True) return all_results[:k] async def read_pages( self, collection: FileId, + principals: list[PrincipalId], page_range: PageRange | None = None, ) -> list[Page]: - """Return ordered page text for ``collection``. + """Return ordered page text for ``collection`` if any principal can read it.""" + return await self._store.read_pages(collection, page_range, principals) - Empty list if the collection has no stored pages. + async def delete_collection(self, collection: FileId, owner_id: OwnerId) -> bool: + """Remove a collection (chunks, pages, ACL). + + Returns ``True`` if the collection was found and deleted, ``False`` if + no row matched ``(collection, owner_id)``. """ - return await self._store.read_pages(collection, page_range) + return await self._store.delete_collection(collection, owner_id) - async def delete_collection(self, collection: FileId) -> None: - """Remove a collection's chunks and pages.""" - await self._store.delete_collection(collection) + async def purge_owner(self, owner_id: OwnerId) -> int: + """Remove every collection ``owner_id`` owns, including vector chunks, + page text, and ACL rows. Returns the number of collections purged.""" + return await self._store.purge_owner(owner_id) - async def has_collection(self, collection: FileId) -> bool: - """Check whether a collection exists.""" - return await self._store.has_collection(collection) + async def reap_expired(self) -> int: + """Delete collections whose ``expires_at`` is set and in the past. + Persistent collections (``expires_at=null``) are never touched. + Returns the number of collections deleted.""" + return await self._store.reap_expired() - async def list_collections(self) -> list[FileId]: - """List all available collections.""" - return [FileId(name) for name in await self._store.list_collections()] + async def has_collection(self, collection: FileId, principals: list[PrincipalId]) -> bool: + """Check whether at least one principal can read this collection.""" + return await self._store.has_collection(collection, principals) + + async def list_collections(self, principals: list[PrincipalId]) -> list[FileId]: + """List collections readable by at least one of ``principals``.""" + return [FileId(name) for name in await self._store.list_collections(principals)] + + async def grant_read( + self, + collection: FileId, + owner_id: OwnerId, + principals: list[PrincipalId], + ) -> None: + """Grant read access to additional principals on an existing doc.""" + await self._store.grant_read(collection, owner_id, principals) + + async def revoke( + self, + collection: FileId, + owner_id: OwnerId, + principal: PrincipalId, + ) -> None: + """Revoke a principal's access on an existing doc.""" + await self._store.revoke(collection, owner_id, principal) async def close(self) -> None: """Release the underlying store's resources.""" diff --git a/engine/src/stirling/documents/sqlite_vec_store.py b/engine/src/stirling/documents/sqlite_vec_store.py index 49f5041a2..934662331 100644 --- a/engine/src/stirling/documents/sqlite_vec_store.py +++ b/engine/src/stirling/documents/sqlite_vec_store.py @@ -5,19 +5,35 @@ import json import math import re import sqlite3 +from datetime import UTC, datetime from pathlib import Path import sqlite_vec from stirling.contracts.documents import Page, PageRange from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage +from stirling.models import OwnerId, PrincipalId + +_READ_PERMISSION = "read" +# sqlite stores TIMESTAMP as TEXT. We normalise to UTC ISO 8601 ``YYYY-MM-DD HH:MM:SS`` +# so lexicographic comparison against ``datetime('now')`` matches chronological order. +_SQLITE_DATETIME_FMT = "%Y-%m-%d %H:%M:%S" + + +def _to_sqlite_utc(dt: datetime | None) -> str | None: + if dt is None: + return None + if dt.tzinfo is not None: + dt = dt.astimezone(UTC).replace(tzinfo=None) + return dt.strftime(_SQLITE_DATETIME_FMT) class SqliteVecStore(DocumentStore): """sqlite-vec backed vector store. Single-file SQLite database, embedded, no server. - Each collection gets its own `vec0` virtual table with a fixed embedding dimension - (detected on first insert). Document metadata lives in a regular table joined by rowid. + Each ``(collection, owner_id)`` pair gets its own `vec0` virtual table with a + fixed embedding dimension (detected on first insert). Document metadata lives + in a regular table joined by rowid. """ def __init__(self, db_path: str | Path) -> None: @@ -51,17 +67,30 @@ class SqliteVecStore(DocumentStore): self._conn.execute( """ CREATE TABLE IF NOT EXISTS documents_meta ( - collection TEXT PRIMARY KEY, - source TEXT NOT NULL + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, + source TEXT NOT NULL, + expires_at TIMESTAMP, + PRIMARY KEY (collection, owner_id) ) """ ) + # The reaper filters on ``expires_at IS NOT NULL AND expires_at < now`` so + # a partial index over non-null rows keeps the scan tight even when most + # rows are persistent (org docs). + self._conn.execute( + "CREATE INDEX IF NOT EXISTS idx_meta_expires_at ON documents_meta(expires_at) WHERE expires_at IS NOT NULL" + ) self._conn.execute( """ CREATE TABLE IF NOT EXISTS collections ( - name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE, + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, dim INTEGER NOT NULL, - table_name TEXT NOT NULL + table_name TEXT NOT NULL, + PRIMARY KEY (collection, owner_id), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) @@ -69,47 +98,154 @@ class SqliteVecStore(DocumentStore): """ CREATE TABLE IF NOT EXISTS documents ( id TEXT NOT NULL, - collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE, + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, text TEXT NOT NULL, metadata TEXT NOT NULL DEFAULT '{}', vec_rowid INTEGER NOT NULL, - PRIMARY KEY (id, collection) + PRIMARY KEY (id, collection, owner_id), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) - self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)") + self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection_owner ON documents(collection, owner_id)") self._conn.execute( """ CREATE TABLE IF NOT EXISTS document_pages ( - collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE, + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, page_number INTEGER NOT NULL, text TEXT NOT NULL, char_count INTEGER NOT NULL, - PRIMARY KEY (collection, page_number) + PRIMARY KEY (collection, owner_id, page_number), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) - self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)") - self._conn.commit() - - async def ensure_collection(self, collection: str, source: str) -> None: - async with self._lock: - await asyncio.to_thread(self._sync_ensure_collection, collection, source) - - def _sync_ensure_collection(self, collection: str, source: str) -> None: + self._conn.execute( + "CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)" + ) self._conn.execute( """ - INSERT INTO documents_meta(collection, source) VALUES (?, ?) - ON CONFLICT(collection) DO UPDATE SET source = excluded.source - """, - (collection, source), + CREATE TABLE IF NOT EXISTS document_acl ( + collection TEXT NOT NULL, + owner_id TEXT NOT NULL, + principal_id TEXT NOT NULL, + permission TEXT NOT NULL, + PRIMARY KEY (collection, owner_id, principal_id, permission), + FOREIGN KEY (collection, owner_id) + REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE + ) + """ + ) + # Lookup by principal is the hot path for search/list (every read + # joins through this index). Composite ordering matches the WHERE. + self._conn.execute( + "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)" ) self._conn.commit() + # ── lifecycle of the (collection, owner_id) row ──────────────────────── + + async def ensure_collection( + self, + collection: str, + source: str, + owner_id: OwnerId, + expires_at: datetime | None, + ) -> None: + async with self._lock: + await asyncio.to_thread(self._sync_ensure_collection, collection, source, owner_id, expires_at) + + def _sync_ensure_collection( + self, + collection: str, + source: str, + owner_id: OwnerId, + expires_at: datetime | None, + ) -> None: + self._conn.execute( + """ + INSERT INTO documents_meta(collection, owner_id, source, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(collection, owner_id) DO UPDATE SET + source = excluded.source, + expires_at = excluded.expires_at + """, + (collection, owner_id, source, _to_sqlite_utc(expires_at)), + ) + self._conn.commit() + + async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: + async with self._lock: + return await asyncio.to_thread(self._sync_delete_collection, collection, owner_id) + + def _sync_delete_collection(self, collection: str, owner_id: OwnerId) -> bool: + # Drop the sqlite-vec virtual table first; FK cascade handles the regular tables + # (collections, documents, document_pages, document_acl) when documents_meta is deleted. + row = self._conn.execute( + "SELECT table_name FROM collections WHERE collection = ? AND owner_id = ?", + (collection, owner_id), + ).fetchone() + if row is not None: + self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}") + cursor = self._conn.execute( + "DELETE FROM documents_meta WHERE collection = ? AND owner_id = ?", + (collection, owner_id), + ) + self._conn.commit() + return cursor.rowcount > 0 + + async def purge_owner(self, owner_id: OwnerId) -> int: + async with self._lock: + return await asyncio.to_thread(self._sync_purge_owner, owner_id) + + def _sync_purge_owner(self, owner_id: OwnerId) -> int: + # Drop all vec0 virtual tables for this owner first (FK cascade can't reach them). + vec_tables = [ + r[0] + for r in self._conn.execute("SELECT table_name FROM collections WHERE owner_id = ?", (owner_id,)).fetchall() + ] + for name in vec_tables: + self._conn.execute(f"DROP TABLE IF EXISTS {name}") + cursor = self._conn.execute("DELETE FROM documents_meta WHERE owner_id = ?", (owner_id,)) + self._conn.commit() + return cursor.rowcount + + async def reap_expired(self) -> int: + async with self._lock: + return await asyncio.to_thread(self._sync_reap_expired) + + def _sync_reap_expired(self) -> int: + # Drop vec0 virtual tables for expired collections first (FK cascade can't reach them). + vec_tables = [ + r[0] + for r in self._conn.execute( + """ + SELECT c.table_name FROM collections c + JOIN documents_meta m + ON m.collection = c.collection AND m.owner_id = c.owner_id + WHERE m.expires_at IS NOT NULL AND m.expires_at < datetime('now') + """ + ).fetchall() + ] + for name in vec_tables: + self._conn.execute(f"DROP TABLE IF EXISTS {name}") + cursor = self._conn.execute( + "DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < datetime('now')" + ) + self._conn.commit() + return cursor.rowcount + + # ── write paths ──────────────────────────────────────────────────────── + @staticmethod - def _sanitize_table_name(collection: str) -> str: - safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection) - return f"vec_{safe}" + def _sanitize_table_name(collection: str, owner_id: OwnerId) -> str: + safe_col = re.sub(r"[^a-zA-Z0-9_]", "_", collection) + safe_owner = re.sub(r"[^a-zA-Z0-9_]", "_", owner_id) + return f"vec_{safe_owner}_{safe_col}" @staticmethod def _normalize(vector: list[float]) -> list[float]: @@ -123,6 +259,7 @@ class SqliteVecStore(DocumentStore): collection: str, documents: list[Document], embeddings: list[list[float]], + owner_id: OwnerId, ) -> None: if len(documents) != len(embeddings): raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings") @@ -130,34 +267,40 @@ class SqliteVecStore(DocumentStore): return async with self._lock: - await asyncio.to_thread(self._sync_add, collection, documents, embeddings) + await asyncio.to_thread(self._sync_add, collection, documents, embeddings, owner_id) def _sync_add( self, collection: str, documents: list[Document], embeddings: list[list[float]], + owner_id: OwnerId, ) -> None: dim = len(embeddings[0]) - row = self._conn.execute("SELECT dim, table_name FROM collections WHERE name = ?", (collection,)).fetchone() + row = self._conn.execute( + "SELECT dim, table_name FROM collections WHERE collection = ? AND owner_id = ?", + (collection, owner_id), + ).fetchone() if row is None: - table_name = self._sanitize_table_name(collection) + table_name = self._sanitize_table_name(collection, owner_id) self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])") self._conn.execute( - "INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)", - (collection, dim, table_name), + "INSERT INTO collections(collection, owner_id, dim, table_name) VALUES (?, ?, ?, ?)", + (collection, owner_id, dim, table_name), ) else: existing_dim, table_name = row if existing_dim != dim: - raise ValueError(f"Collection {collection} has dim {existing_dim}, got embedding of dim {dim}") + raise ValueError( + f"Collection {collection} for owner {owner_id} has dim {existing_dim}, got embedding of dim {dim}" + ) # Upsert: delete existing docs with matching IDs first ids = [doc.id for doc in documents] placeholders = ",".join("?" * len(ids)) existing = self._conn.execute( - f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})", - (collection, *ids), + f"SELECT vec_rowid FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})", + (collection, owner_id, *ids), ).fetchall() if existing: vec_rowids = [r[0] for r in existing] @@ -167,8 +310,8 @@ class SqliteVecStore(DocumentStore): vec_rowids, ) self._conn.execute( - f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})", - (collection, *ids), + f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})", + (collection, owner_id, *ids), ) for doc, emb in zip(documents, embeddings): @@ -179,27 +322,101 @@ class SqliteVecStore(DocumentStore): ) vec_rowid = cursor.lastrowid self._conn.execute( - "INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)", - (doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid), + "INSERT INTO documents(id, collection, owner_id, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?, ?)", + (doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), vec_rowid), ) self._conn.commit() + async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: + async with self._lock: + await asyncio.to_thread(self._sync_add_pages, collection, pages, owner_id) + + def _sync_add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: + self._conn.execute( + "DELETE FROM document_pages WHERE collection = ? AND owner_id = ?", + (collection, owner_id), + ) + if pages: + self._conn.executemany( + "INSERT INTO document_pages(collection, owner_id, page_number, text, char_count) " + "VALUES (?, ?, ?, ?, ?)", + [(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages], + ) + self._conn.commit() + + # ── ACL management ───────────────────────────────────────────────────── + + async def grant_read( + self, + collection: str, + owner_id: OwnerId, + principals: list[PrincipalId], + ) -> None: + if not principals: + return + async with self._lock: + await asyncio.to_thread(self._sync_grant_read, collection, owner_id, principals) + + def _sync_grant_read( + self, + collection: str, + owner_id: OwnerId, + principals: list[PrincipalId], + ) -> None: + self._conn.executemany( + """ + INSERT INTO document_acl(collection, owner_id, principal_id, permission) + VALUES (?, ?, ?, ?) + ON CONFLICT(collection, owner_id, principal_id, permission) DO NOTHING + """, + [(collection, owner_id, p, _READ_PERMISSION) for p in principals], + ) + self._conn.commit() + + async def revoke( + self, + collection: str, + owner_id: OwnerId, + principal: PrincipalId, + ) -> None: + async with self._lock: + await asyncio.to_thread(self._sync_revoke, collection, owner_id, principal) + + def _sync_revoke(self, collection: str, owner_id: OwnerId, principal: PrincipalId) -> None: + self._conn.execute( + "DELETE FROM document_acl WHERE collection = ? AND owner_id = ? AND principal_id = ?", + (collection, owner_id, principal), + ) + self._conn.commit() + + # ── read paths (ACL-gated) ───────────────────────────────────────────── + async def search( self, collection: str, query_embedding: list[float], - top_k: int = 5, + top_k: int, + principals: list[PrincipalId], ) -> list[SearchResult]: async with self._lock: - return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k) + return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k, principals) def _sync_search( self, collection: str, query_embedding: list[float], top_k: int, + principals: list[PrincipalId], ) -> list[SearchResult]: - row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone() + if not principals: + return [] + owner_id = self._readable_owner_for(collection, principals) + if owner_id is None: + return [] + row = self._conn.execute( + "SELECT table_name, dim FROM collections WHERE collection = ? AND owner_id = ?", + (collection, owner_id), + ).fetchone() if row is None: return [] table_name, dim = row @@ -213,11 +430,14 @@ class SqliteVecStore(DocumentStore): f""" SELECT d.id, d.text, d.metadata, v.distance FROM {table_name} v - JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ? + JOIN documents d + ON d.vec_rowid = v.rowid + AND d.collection = ? + AND d.owner_id = ? WHERE v.embedding MATCH ? AND k = ? ORDER BY v.distance """, - (collection, query_blob, top_k), + (collection, owner_id, query_blob, top_k), ).fetchall() return [ @@ -233,77 +453,85 @@ class SqliteVecStore(DocumentStore): for r in results ] - async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: - async with self._lock: - await asyncio.to_thread(self._sync_add_pages, collection, pages) - - def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None: - self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,)) - if pages: - self._conn.executemany( - "INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)", - [(collection, p.page_number, p.text, p.char_count) for p in pages], - ) - self._conn.commit() + def _readable_owner_for(self, collection: str, principals: list[PrincipalId]) -> str | None: + """Resolve which owner_id this caller is reading. ``None`` means no access.""" + placeholders = ",".join("?" * len(principals)) + row = self._conn.execute( + f""" + SELECT owner_id FROM document_acl + WHERE collection = ? + AND permission = ? + AND principal_id IN ({placeholders}) + ORDER BY owner_id + LIMIT 1 + """, + (collection, _READ_PERMISSION, *principals), + ).fetchone() + return row[0] if row else None async def read_pages( self, collection: str, - page_range: PageRange | None = None, + page_range: PageRange | None, + principals: list[PrincipalId], ) -> list[Page]: async with self._lock: - return await asyncio.to_thread(self._sync_read_pages, collection, page_range) + return await asyncio.to_thread(self._sync_read_pages, collection, page_range, principals) def _sync_read_pages( self, collection: str, page_range: PageRange | None, + principals: list[PrincipalId], ) -> list[Page]: + if not principals: + return [] + owner_id = self._readable_owner_for(collection, principals) + if owner_id is None: + return [] if page_range is None: rows = self._conn.execute( - "SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number", - (collection,), + "SELECT page_number, text, char_count FROM document_pages " + "WHERE collection = ? AND owner_id = ? ORDER BY page_number", + (collection, owner_id), ).fetchall() else: rows = self._conn.execute( "SELECT page_number, text, char_count FROM document_pages " - "WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number", - (collection, page_range.start, page_range.end), + "WHERE collection = ? AND owner_id = ? AND page_number BETWEEN ? AND ? " + "ORDER BY page_number", + (collection, owner_id, page_range.start, page_range.end), ).fetchall() return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows] - async def delete_collection(self, collection: str) -> None: + async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: async with self._lock: - await asyncio.to_thread(self._sync_delete_collection, collection) + return await asyncio.to_thread(self._sync_has_collection, collection, principals) - def _sync_delete_collection(self, collection: str) -> None: - # Drop the sqlite-vec virtual table first; FK cascade handles the regular tables - # (collections, documents, document_pages) when documents_meta is deleted. - row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone() - if row is not None: - self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}") - self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,)) - self._conn.commit() + def _sync_has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: + if not principals: + return False + return self._readable_owner_for(collection, principals) is not None - async def list_collections(self) -> list[str]: + async def list_collections(self, principals: list[PrincipalId]) -> list[str]: async with self._lock: - return await asyncio.to_thread(self._sync_list_collections) + return await asyncio.to_thread(self._sync_list_collections, principals) - def _sync_list_collections(self) -> list[str]: - rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall() + def _sync_list_collections(self, principals: list[PrincipalId]) -> list[str]: + if not principals: + return [] + placeholders = ",".join("?" * len(principals)) + rows = self._conn.execute( + f""" + SELECT DISTINCT collection FROM document_acl + WHERE permission = ? + AND principal_id IN ({placeholders}) + ORDER BY collection + """, + (_READ_PERMISSION, *principals), + ).fetchall() return [r[0] for r in rows] - async def has_collection(self, collection: str) -> bool: - async with self._lock: - return await asyncio.to_thread(self._sync_has_collection, collection) - - def _sync_has_collection(self, collection: str) -> bool: - row = self._conn.execute( - "SELECT 1 FROM documents_meta WHERE collection = ?", - (collection,), - ).fetchone() - return row is not None - async def close(self) -> None: async with self._lock: await asyncio.to_thread(self._sync_close) diff --git a/engine/src/stirling/documents/store.py b/engine/src/stirling/documents/store.py index d59fd9d63..300d60bd8 100644 --- a/engine/src/stirling/documents/store.py +++ b/engine/src/stirling/documents/store.py @@ -2,8 +2,10 @@ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass, field +from datetime import datetime from stirling.contracts.documents import Page, PageRange +from stirling.models import OwnerId, PrincipalId @dataclass @@ -40,73 +42,113 @@ class DocumentStore(ABC): * **Vector chunks** - small, embedded chunks used for RAG search. * **Ordered pages** - the original page text retained in document order, used for whole-document reading. - - Both representations live under the same ``collection`` (file id) and are - rooted at a single parent row in ``documents_meta``. Removing that parent - row cascades to both child representations, so :meth:`delete_collection` - is one logical delete. """ + # ── lifecycle of the (collection, owner_id) row ──────────────────────── + @abstractmethod - async def ensure_collection(self, collection: str, source: str) -> None: - """Upsert the top-level ``documents_meta`` row for this collection. + async def ensure_collection( + self, + collection: str, + source: str, + owner_id: OwnerId, + expires_at: datetime | None, + ) -> None: + """Upsert the top-level ``documents_meta`` row for ``(collection, owner_id)``. Must be called before :meth:`add_pages` or :meth:`add_documents`. Both - of those write into child tables that hold a foreign key to the parent - row, so it must exist first. + write into child tables that hold a foreign key to the parent row, so + it must exist first. """ + @abstractmethod + async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: + """Remove a collection's chunks, pages, and ACL rows. + + Returns ``True`` when a row matched and was deleted, ``False`` when + ``(collection, owner_id)`` didn't exist. + """ + + @abstractmethod + async def purge_owner(self, owner_id: OwnerId) -> int: + """Remove every collection (and ACL row) belonging to ``owner_id``. + Returns the number of collections purged. + """ + + @abstractmethod + async def reap_expired(self) -> int: + """Remove collections whose ``expires_at`` is non-null and in the past.""" + + # ── write paths (scoped by owner) ────────────────────────────────────── + @abstractmethod async def add_documents( self, collection: str, documents: list[Document], embeddings: list[list[float]], + owner_id: OwnerId, ) -> None: - """Store vector chunks with their embeddings in the named collection.""" + """Store vector chunks with their embeddings under the owner's collection.""" + + @abstractmethod + async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: + """Replace the stored pages for ``(collection, owner_id)`` with the supplied pages.""" + + # ── ACL management ───────────────────────────────────────────────────── + + @abstractmethod + async def grant_read( + self, + collection: str, + owner_id: OwnerId, + principals: list[PrincipalId], + ) -> None: + """Grant read access on ``(collection, owner_id)`` to each principal.""" + + @abstractmethod + async def revoke( + self, + collection: str, + owner_id: OwnerId, + principal: PrincipalId, + ) -> None: + """Remove every permission this principal has on ``(collection, owner_id)``.""" + + # ── read paths (scoped by ACL principal set) ─────────────────────────── @abstractmethod async def search( self, collection: str, query_embedding: list[float], - top_k: int = 5, + top_k: int, + principals: list[PrincipalId], ) -> list[SearchResult]: - """Return the top_k most similar vector chunks from the collection.""" + """Top-k similar vector chunks from ``collection`` that any principal can read. - @abstractmethod - async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: - """Replace the stored pages for ``collection`` with the supplied pages. - - Implementations must remove any previously-stored pages for the - collection before writing, so callers can re-ingest by calling this - method again. + Returns an empty list when no principal in ``principals`` has a read + ACL row for this collection (regardless of which owner backs it). """ @abstractmethod async def read_pages( self, collection: str, - page_range: PageRange | None = None, + page_range: PageRange | None, + principals: list[PrincipalId], ) -> list[Page]: - """Return ordered pages for ``collection``. - - If ``page_range`` is ``None`` all pages are returned. Otherwise only - pages whose ``page_number`` falls within the inclusive range are - returned. Pages are always ordered by ``page_number`` ascending. - """ + """Return ordered pages for ``collection``. Empty list when no read ACL match.""" @abstractmethod - async def delete_collection(self, collection: str) -> None: - """Remove a collection's chunks and pages.""" + async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: + """Check whether any principal in ``principals`` can read this collection.""" @abstractmethod - async def list_collections(self) -> list[str]: - """Return names of all existing collections.""" + async def list_collections(self, principals: list[PrincipalId]) -> list[str]: + """Return collection names readable by at least one of ``principals``.""" - @abstractmethod - async def has_collection(self, collection: str) -> bool: - """Check whether a collection exists.""" + # ── lifecycle ────────────────────────────────────────────────────────── @abstractmethod async def close(self) -> None: diff --git a/engine/src/stirling/models/__init__.py b/engine/src/stirling/models/__init__.py index 9205f60a1..a1a5530af 100644 --- a/engine/src/stirling/models/__init__.py +++ b/engine/src/stirling/models/__init__.py @@ -1,12 +1,15 @@ from . import tool_models -from .base import ApiModel, FileId +from .base import ApiModel, FileId, OwnerId, PrincipalId, UserId from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint __all__ = [ "ApiModel", "FileId", "OPERATIONS", + "OwnerId", "ParamToolModel", + "PrincipalId", "ToolEndpoint", + "UserId", "tool_models", ] diff --git a/engine/src/stirling/models/base.py b/engine/src/stirling/models/base.py index 12c32b62c..a1d1e1081 100644 --- a/engine/src/stirling/models/base.py +++ b/engine/src/stirling/models/base.py @@ -5,11 +5,18 @@ from typing import NewType from pydantic import BaseModel, ConfigDict from pydantic.alias_generators import to_camel -# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's -# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key -# throughout the engine. +# Stable, opaque identifier for a file supplied by the caller FileId = NewType("FileId", str) +# Stable, opaque identifier for the calling user +UserId = NewType("UserId", str) + +# Tenant that owns a document +OwnerId = NewType("OwnerId", str) + +# An entity that can hold permissions on a document +PrincipalId = NewType("PrincipalId", str) + class ApiModel(BaseModel): """Base for every contract model crossing a service boundary.""" diff --git a/engine/src/stirling/services/__init__.py b/engine/src/stirling/services/__init__.py index 32189312c..1db54243f 100644 --- a/engine/src/stirling/services/__init__.py +++ b/engine/src/stirling/services/__init__.py @@ -7,14 +7,16 @@ from .progress import ( set_progress_emitter, ) from .runtime import AppRuntime, build_model_settings, build_runtime -from .tracking import setup_posthog_tracking +from .tracking import current_user_id, require_current_user_id, setup_posthog_tracking __all__ = [ "AppRuntime", "ProgressEmitter", "build_model_settings", "build_runtime", + "current_user_id", "emit_progress", + "require_current_user_id", "reset_progress_emitter", "set_progress_emitter", "setup_posthog_tracking", diff --git a/engine/src/stirling/services/runtime.py b/engine/src/stirling/services/runtime.py index 4d5841489..685509cfe 100644 --- a/engine/src/stirling/services/runtime.py +++ b/engine/src/stirling/services/runtime.py @@ -16,7 +16,6 @@ from stirling.documents import ( DocumentStore, EmbeddingService, PgVectorStore, - RagCapability, SqliteVecStore, ) @@ -47,7 +46,6 @@ class AppRuntime: fast_model: Model smart_model: Model documents: DocumentService - rag_capability: RagCapability @property def fast_model_settings(self) -> ModelSettings: @@ -89,8 +87,8 @@ def _build_document_store(settings: AppSettings) -> DocumentStore: assert_never(settings.rag_backend) -def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]: - """Build the document service and the RAG-search capability that wraps it.""" +def _build_documents(settings: AppSettings) -> DocumentService: + """Build the document service used by per-request RAG capabilities.""" logger.info("Documents: embedding_model=%s", settings.rag_embedding_model) embedder = EmbeddingService( model_name=settings.rag_embedding_model, @@ -98,9 +96,7 @@ def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapabil chunk_overlap=settings.rag_chunk_overlap, ) store = _build_document_store(settings) - service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k) - capability = RagCapability(documents=service, top_k=settings.rag_default_top_k) - return service, capability + return DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k) def build_runtime(settings: AppSettings) -> AppRuntime: @@ -109,14 +105,11 @@ def build_runtime(settings: AppSettings) -> AppRuntime: validate_structured_output_support(fast_model, settings.fast_model_name) validate_structured_output_support(smart_model, settings.smart_model_name) - documents, rag_capability = _build_documents(settings) - return AppRuntime( settings=settings, fast_model=fast_model, smart_model=smart_model, - documents=documents, - rag_capability=rag_capability, + documents=_build_documents(settings), ) diff --git a/engine/src/stirling/services/tracking.py b/engine/src/stirling/services/tracking.py index f8c93524a..0064611bb 100644 --- a/engine/src/stirling/services/tracking.py +++ b/engine/src/stirling/services/tracking.py @@ -27,10 +27,25 @@ from opentelemetry.trace import Span from posthog.client import Client as PostHogClient from stirling.config import AppSettings +from stirling.models import UserId # Per-request user ID, set by middleware from the X-User-Id header. # When not set, PostHog generates a random ID and marks the event as personless. -current_user_id: ContextVar[str | None] = ContextVar("current_user_id", default=None) +current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None) + + +def require_current_user_id() -> UserId: + """Return the request's user ID or raise if the X-User-Id header was missing. + + Use at the boundary of any code path that touches per-user document + storage (vector chunks, page text, ACL rows). Routes prefer the FastAPI dependency form + (``Depends(require_user_id)``); agent internals that don't have a request + object in scope call this helper directly to fail closed. + """ + user_id = current_user_id.get() + if user_id is None: + raise RuntimeError("X-User-Id is required for this operation but was not set on the request") + return user_id class LRUSet: diff --git a/engine/tests/agents/test_whole_doc_reader.py b/engine/tests/agents/test_whole_doc_reader.py index 588a43f37..749b847a1 100644 --- a/engine/tests/agents/test_whole_doc_reader.py +++ b/engine/tests/agents/test_whole_doc_reader.py @@ -15,9 +15,12 @@ import pytest from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability from stirling.contracts import AiFile, PageText from stirling.documents import Document, DocumentService, SqliteVecStore -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId from stirling.services.runtime import AppRuntime +OWNER = OwnerId("test-user") +PRINCIPALS = [PrincipalId("test-user")] + class StubEmbedder: """Deterministic embeddings so tests don't need a real provider.""" @@ -75,7 +78,9 @@ async def test_read_full_document_returns_formatted_notes_for_single_file( PageText(page_number=1, text="Chapter one prose."), PageText(page_number=2, text="Chapter two prose."), ] - await runtime_with_stub_docs.documents.ingest(FileId("doc-id"), pages, source="doc.pdf") + await runtime_with_stub_docs.documents.ingest( + FileId("doc-id"), pages, source="doc.pdf", owner_id=OWNER, read_principals=PRINCIPALS, expires_at=None + ) reasoner = ChunkedReasoner(runtime_with_stub_docs) canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])] @@ -83,6 +88,7 @@ async def test_read_full_document_returns_formatted_notes_for_single_file( capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("doc-id", "doc.pdf")], + principals=PRINCIPALS, reasoner=reasoner, ) result = await capability._read_full_document("what is in the document?") @@ -106,6 +112,9 @@ async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs FileId(cid), [PageText(page_number=1, text=f"contents of {cid}")], source=source, + owner_id=OWNER, + read_principals=PRINCIPALS, + expires_at=None, ) reasoner = ChunkedReasoner(runtime_with_stub_docs) @@ -121,6 +130,7 @@ async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")], + principals=PRINCIPALS, reasoner=reasoner, ) result = await capability._read_full_document("compare them") @@ -140,6 +150,9 @@ async def test_read_full_document_skips_files_without_pages(runtime_with_stub_do FileId("present"), [PageText(page_number=1, text="real content")], source="present.pdf", + owner_id=OWNER, + read_principals=PRINCIPALS, + expires_at=None, ) # 'missing' is never ingested -> read_pages returns []. @@ -149,6 +162,7 @@ async def test_read_full_document_skips_files_without_pages(runtime_with_stub_do capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")], + principals=PRINCIPALS, reasoner=reasoner, ) result = await capability._read_full_document("anything") @@ -167,6 +181,7 @@ async def test_read_full_document_returns_empty_message_when_no_pages_anywhere( capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("nope", "nope.pdf")], + principals=PRINCIPALS, reasoner=reasoner, ) result = await capability._read_full_document("anything") @@ -186,12 +201,16 @@ async def test_read_full_document_budget_hides_tool_when_exhausted( FileId("doc-id"), [PageText(page_number=1, text="content")], source="doc.pdf", + owner_id=OWNER, + read_principals=PRINCIPALS, + expires_at=None, ) reasoner = ChunkedReasoner(runtime_with_stub_docs) with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])): capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("doc-id", "doc.pdf")], + principals=PRINCIPALS, reasoner=reasoner, max_reads=1, ) @@ -212,6 +231,7 @@ async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRu capability = WholeDocReaderCapability( runtime=runtime_with_stub_docs, files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")], + principals=PRINCIPALS, ) text = capability.instructions diff --git a/engine/tests/contradiction/test_capability.py b/engine/tests/contradiction/test_capability.py index 0439411b1..3a27c6e6a 100644 --- a/engine/tests/contradiction/test_capability.py +++ b/engine/tests/contradiction/test_capability.py @@ -17,9 +17,11 @@ from stirling.contracts.contradiction import ( ContradictionReport, ContradictionSeverity, ) -from stirling.models import FileId +from stirling.models import FileId, PrincipalId from stirling.services.runtime import AppRuntime +PRINCIPALS = [PrincipalId("test-user")] + def _file(file_id: str, name: str) -> AiFile: return AiFile(id=FileId(file_id), name=name) @@ -58,7 +60,7 @@ async def test_find_contradictions_returns_formatted_text(runtime: AppRuntime) - canned = _canned_report() detector.detect = AsyncMock(return_value=canned) - capability = ContradictionCapability(detector=detector, files=[_file("doc-a", "a.pdf")]) + capability = ContradictionCapability(detector=detector, files=[_file("doc-a", "a.pdf")], principals=PRINCIPALS) result = await capability._find_contradictions("are there inconsistent deadlines?") detector.detect.assert_awaited_once() @@ -84,6 +86,7 @@ async def test_budget_gate_hides_tool_after_first_audit(runtime: AppRuntime) -> capability = ContradictionCapability( detector=detector, files=[_file("doc-a", "a.pdf")], + principals=PRINCIPALS, max_audits=1, ) # A real, minimal ToolDefinition — the prepare callback returns this @@ -107,7 +110,7 @@ async def test_budget_gate_hides_tool_after_first_audit(runtime: AppRuntime) -> async def test_find_contradictions_with_no_files_returns_message(runtime: AppRuntime) -> None: detector = ContradictionDetector(runtime) detector.detect = AsyncMock(return_value=_canned_report()) - capability = ContradictionCapability(detector=detector, files=[]) + capability = ContradictionCapability(detector=detector, files=[], principals=PRINCIPALS) result = await capability._find_contradictions("anything") @@ -120,6 +123,7 @@ def test_instructions_mention_attached_files(runtime: AppRuntime) -> None: capability = ContradictionCapability( detector=detector, files=[_file("doc-a", "alpha.pdf"), _file("doc-b", "beta.pdf")], + principals=PRINCIPALS, ) text = capability.instructions @@ -149,6 +153,7 @@ def test_instructions_escape_filename_injection_attempt(runtime: AppRuntime) -> capability = ContradictionCapability( detector=detector, files=[_file("doc-evil", evil_name)], + principals=PRINCIPALS, ) text = capability.instructions diff --git a/engine/tests/contradiction/test_detector.py b/engine/tests/contradiction/test_detector.py index d38d3dcad..b9532e1e3 100644 --- a/engine/tests/contradiction/test_detector.py +++ b/engine/tests/contradiction/test_detector.py @@ -27,7 +27,7 @@ from stirling.agents.shared.chunked_mapper import ChunkOutput from stirling.contracts import AiFile from stirling.contracts.contradiction import ContradictionSeverity from stirling.contracts.documents import Page, PageRange -from stirling.models import FileId +from stirling.models import FileId, PrincipalId from stirling.services.runtime import AppRuntime @@ -58,10 +58,17 @@ def pages_a() -> list[Page]: ] +PRINCIPALS = [PrincipalId("test-user")] + + def _install_documents_stub(runtime: AppRuntime, pages_by_id: dict[FileId, list[Page]]) -> None: """Patch ``runtime.documents.read_pages`` to return canned pages per file.""" - async def _read(collection: FileId, page_range: PageRange | None = None) -> list[Page]: + async def _read( + collection: FileId, + principals: list[PrincipalId], + page_range: PageRange | None = None, + ) -> list[Page]: return pages_by_id.get(collection, []) # AppRuntime is frozen; monkey-patch the documents service. @@ -76,7 +83,7 @@ async def test_no_pages_returns_clean_empty_report(runtime: AppRuntime, file_a: _install_documents_stub(runtime, {file_a.id: []}) detector = ContradictionDetector(runtime) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert report.contradictions == [] assert report.pages_examined == [] @@ -126,7 +133,7 @@ async def test_happy_path_finds_contradiction_across_two_pages( ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("Examined 2 pages; found 1 contradiction.")) - report = await detector.detect([file_a], query="check the deadline") + report = await detector.detect([file_a], principals=PRINCIPALS, query="check the deadline") assert len(report.contradictions) == 1 c = report.contradictions[0] @@ -150,12 +157,12 @@ async def test_zero_claims_returns_clean_report(runtime: AppRuntime, file_a: AiF return_value=[ChunkOutput(pages=[1, 2], output=_ExtractedClaims(claims=[]), label="pages=1-2")] ) # Stubbing the summary agent is unavoidable (the production code calls - # it on every detect()); we just don't assert on what it returns — - # asserting on the canned value here would only re-prove that AsyncMock + # it on every detect()); we just don't assert on what it returns. + # Asserting on the canned value here would only re-prove that AsyncMock # works. detector._summary_agent.run = AsyncMock(return_value=_stub_result("any text")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert report.contradictions == [] assert report.clean is True @@ -204,7 +211,7 @@ async def test_canonicaliser_accepts_empty_alias_list(runtime: AppRuntime, file_ ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert len(report.contradictions) == 1 @@ -332,7 +339,7 @@ async def test_canonicaliser_failure_falls_back_to_lexical_keys( ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) # Lexical key collapses both subjects so the bucket still forms. assert len(report.contradictions) == 1 @@ -389,7 +396,7 @@ async def test_same_page_contradiction_is_surfaced(runtime: AppRuntime, file_a: ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert len(report.contradictions) == 1 assert report.contradictions[0].severity == ContradictionSeverity.ERROR @@ -426,7 +433,7 @@ async def test_identical_quote_pair_is_still_dropped(runtime: AppRuntime, file_a ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert report.contradictions == [] @@ -455,7 +462,7 @@ async def test_summary_falls_back_to_deterministic_when_llm_unavailable( ) detector._summary_agent.run = AsyncMock(side_effect=failure) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) assert "No contradictions" in report.summary assert report.clean is True @@ -499,7 +506,7 @@ async def test_detector_chunk_timeout_falls_through(runtime: AppRuntime, file_a: detector._pair_detector.run = AsyncMock(side_effect=TimeoutError("simulated")) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) # Detector timed out so no pairs come back. Crucially: the pipeline # reached the summary stage rather than crashing earlier, so @@ -533,7 +540,7 @@ async def test_empty_chunk_with_substantial_content_logs_warning( detector._summary_agent.run = AsyncMock(return_value=_stub_result("ok")) with caplog.at_level(logging.WARNING, logger="stirling.agents.contradiction.detector"): - await detector.detect([file_a]) + await detector.detect([file_a], principals=PRINCIPALS) assert any( "produced 0 claims" in record.getMessage() and "pages=1" in record.getMessage() for record in caplog.records @@ -580,7 +587,7 @@ async def test_pages_examined_includes_every_attempted_page(runtime: AppRuntime, detector._pair_detector.run = AsyncMock(return_value=_stub_result(_BucketContradictions(pairs=[]))) detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) # Every page the extractor ran against is reported, even page 2 # (which produced no claim). @@ -655,7 +662,7 @@ async def test_oversized_bucket_windows_translate_indices_globally(runtime: AppR detector._pair_detector.run = _stub_detector # type: ignore[method-assign] detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) - report = await detector.detect([file_a]) + report = await detector.detect([file_a], principals=PRINCIPALS) # Both windows produced one valid pair each; dedup by global (i, j) # leaves exactly two contradictions. @@ -791,7 +798,7 @@ async def test_multi_file_pages_dont_collide_in_validation(runtime: AppRuntime) ) detector._summary_agent.run = AsyncMock(return_value=_stub_result("ok")) - report = await detector.detect([file_a, file_b]) + report = await detector.detect([file_a, file_b], principals=PRINCIPALS) # Both claims validated as verbatim — each against the right file's # page text. A collision bug would have produced "paraphrased" for at diff --git a/engine/tests/contradiction/test_question_integration.py b/engine/tests/contradiction/test_question_integration.py index a0018b63c..228f24525 100644 --- a/engine/tests/contradiction/test_question_integration.py +++ b/engine/tests/contradiction/test_question_integration.py @@ -9,6 +9,7 @@ detector when invoked. from __future__ import annotations +from collections.abc import Iterator from dataclasses import replace import pytest @@ -22,10 +23,24 @@ from stirling.contracts import ( ) from stirling.contracts.contradiction import Claim from stirling.documents import DocumentService, SqliteVecStore -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId, UserId +from stirling.services import current_user_id from stirling.services.runtime import AppRuntime from tests.test_pdf_question_agent import StubEmbedder +USER = UserId("test-user") +OWNER = OwnerId("test-user") +OWNER_PRINCIPALS = [PrincipalId("test-user")] + + +@pytest.fixture(autouse=True) +def _set_user_context() -> Iterator[None]: + token = current_user_id.set(USER) + try: + yield + finally: + current_user_id.reset(token) + def _file(file_id: str, name: str) -> AiFile: return AiFile(id=FileId(file_id), name=name) @@ -69,6 +84,9 @@ async def test_run_answer_agent_builds_agent_with_three_toolsets( file.id, [PageText(page_number=1, text="content")], source=file.name, + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = PdfQuestionAgent(runtime_with_stub_docs) diff --git a/engine/tests/contradiction/test_review_integration.py b/engine/tests/contradiction/test_review_integration.py index 9c843a53f..51038790a 100644 --- a/engine/tests/contradiction/test_review_integration.py +++ b/engine/tests/contradiction/test_review_integration.py @@ -8,6 +8,7 @@ contradiction and the right cross-references and anchor handling. from __future__ import annotations import json +from collections.abc import Iterator from dataclasses import replace from typing import Literal from unittest.mock import AsyncMock @@ -27,11 +28,25 @@ from stirling.contracts import ( ) from stirling.contracts.contradiction import Claim from stirling.documents import DocumentService, SqliteVecStore -from stirling.models import FileId, ToolEndpoint +from stirling.models import FileId, OwnerId, PrincipalId, ToolEndpoint, UserId from stirling.models.tool_models import AddCommentsParams +from stirling.services import current_user_id from stirling.services.runtime import AppRuntime from tests.test_pdf_question_agent import StubEmbedder +USER = UserId("test-user") +OWNER = OwnerId("test-user") +OWNER_PRINCIPALS = [PrincipalId("test-user")] + + +@pytest.fixture(autouse=True) +def _set_user_context() -> Iterator[None]: + token = current_user_id.set(USER) + try: + yield + finally: + current_user_id.reset(token) + def _file(file_id: str, name: str) -> AiFile: return AiFile(id=FileId(file_id), name=name) @@ -88,6 +103,9 @@ async def test_localiser_prompt_escapes_verdict_tag_injection( file.id, [PageText(page_number=1, text="x")], source=file.name, + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = PdfReviewAgent(runtime_with_stub_docs) @@ -159,6 +177,9 @@ async def test_contradiction_intent_emits_add_comments_plan( file.id, [PageText(page_number=1, text="ignored"), PageText(page_number=5, text="ignored")], source=file.name, + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = PdfReviewAgent(runtime_with_stub_docs) @@ -264,6 +285,9 @@ async def test_contradiction_takes_precedence_over_math( file.id, [PageText(page_number=1, text="x")], source=file.name, + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = PdfReviewAgent(runtime_with_stub_docs) diff --git a/engine/tests/test_documents.py b/engine/tests/test_documents.py index a5febf078..7da5c3480 100644 --- a/engine/tests/test_documents.py +++ b/engine/tests/test_documents.py @@ -8,7 +8,17 @@ from stirling.documents.rag_capability import RagCapability from stirling.documents.service import DocumentService from stirling.documents.sqlite_vec_store import SqliteVecStore from stirling.documents.store import Document, SearchResult -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId + +# Personal-doc tests reuse the same opaque string in all three roles — keeps the +# defaults of ingest() exercised and tests honest about what the same id means +# (caller, tenant, single ACL grantee). Org-doc tests construct ad-hoc owners +# and principal sets inline. +OWNER = OwnerId("test-user") +OWNER_PRINCIPALS = [PrincipalId("test-user")] +OTHER_OWNER = OwnerId("other-user") +OTHER_OWNER_PRINCIPALS = [PrincipalId("other-user")] + # chunk_text @@ -56,7 +66,8 @@ class TestSqliteVecStore: @pytest.mark.anyio async def test_add_and_search(self) -> None: store = SqliteVecStore.ephemeral() - await store.ensure_collection("test-col", "test.pdf") + await store.ensure_collection("test-col", "test.pdf", OWNER, None) + await store.grant_read("test-col", OWNER, OWNER_PRINCIPALS) docs = [ Document(id="1", text="Python is a programming language", metadata={"source": "test"}), Document(id="2", text="Java is another programming language", metadata={"source": "test"}), @@ -67,9 +78,9 @@ class TestSqliteVecStore: [0.9, 0.1, 0.0], [0.0, 0.0, 1.0], ] - await store.add_documents("test-col", docs, embeddings) + await store.add_documents("test-col", docs, embeddings, OWNER) - results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2) + results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2, principals=OWNER_PRINCIPALS) assert len(results) == 2 assert isinstance(results[0], SearchResult) assert results[0].document.id == "1" @@ -78,33 +89,36 @@ class TestSqliteVecStore: @pytest.mark.anyio async def test_list_and_has_collection(self) -> None: store = SqliteVecStore.ephemeral() - await store.ensure_collection("my-collection", "test.pdf") + await store.ensure_collection("my-collection", "test.pdf", OWNER, None) + await store.grant_read("my-collection", OWNER, OWNER_PRINCIPALS) docs = [Document(id="1", text="test", metadata={})] - await store.add_documents("my-collection", docs, [[1.0, 0.0]]) + await store.add_documents("my-collection", docs, [[1.0, 0.0]], OWNER) - collections = await store.list_collections() + collections = await store.list_collections(OWNER_PRINCIPALS) assert "my-collection" in collections - assert await store.has_collection("my-collection") is True - assert await store.has_collection("nonexistent") is False + assert await store.has_collection("my-collection", OWNER_PRINCIPALS) is True + assert await store.has_collection("nonexistent", OWNER_PRINCIPALS) is False @pytest.mark.anyio async def test_delete_collection(self) -> None: store = SqliteVecStore.ephemeral() - await store.ensure_collection("to-delete", "test.pdf") + await store.ensure_collection("to-delete", "test.pdf", OWNER, None) + await store.grant_read("to-delete", OWNER, OWNER_PRINCIPALS) docs = [Document(id="1", text="test", metadata={})] - await store.add_documents("to-delete", docs, [[1.0]]) + await store.add_documents("to-delete", docs, [[1.0]], OWNER) - assert await store.has_collection("to-delete") is True - await store.delete_collection("to-delete") - assert await store.has_collection("to-delete") is False + assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is True + await store.delete_collection("to-delete", OWNER) + assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is False @pytest.mark.anyio async def test_search_empty_collection(self) -> None: store = SqliteVecStore.ephemeral() - await store.ensure_collection("empty-test", "test.pdf") + await store.ensure_collection("empty-test", "test.pdf", OWNER, None) + await store.grant_read("empty-test", OWNER, OWNER_PRINCIPALS) docs = [Document(id="1", text="test", metadata={})] - await store.add_documents("empty-test", docs, [[1.0, 0.0]]) - results = await store.search("empty-test", [1.0, 0.0], top_k=5) + await store.add_documents("empty-test", docs, [[1.0, 0.0]], OWNER) + results = await store.search("empty-test", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS) assert len(results) == 1 @pytest.mark.anyio @@ -112,7 +126,118 @@ class TestSqliteVecStore: store = SqliteVecStore.ephemeral() docs = [Document(id="1", text="test", metadata={})] with pytest.raises(ValueError, match="documents.*embeddings"): - await store.add_documents("bad", docs, [[1.0], [2.0]]) + await store.add_documents("bad", docs, [[1.0], [2.0]], OWNER) + + @pytest.mark.anyio + async def test_collections_isolated_by_owner(self) -> None: + """Two owners can store the same collection id; reads stay scoped to ACL.""" + store = SqliteVecStore.ephemeral() + await store.ensure_collection("shared-id", "alice.pdf", OWNER, None) + await store.grant_read("shared-id", OWNER, OWNER_PRINCIPALS) + await store.ensure_collection("shared-id", "bob.pdf", OTHER_OWNER, None) + await store.grant_read("shared-id", OTHER_OWNER, OTHER_OWNER_PRINCIPALS) + await store.add_documents( + "shared-id", + [Document(id="1", text="alice content", metadata={})], + [[1.0, 0.0]], + OWNER, + ) + await store.add_documents( + "shared-id", + [Document(id="1", text="bob content", metadata={})], + [[1.0, 0.0]], + OTHER_OWNER, + ) + + alice_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS) + bob_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OTHER_OWNER_PRINCIPALS) + assert [r.document.text for r in alice_results] == ["alice content"] + assert [r.document.text for r in bob_results] == ["bob content"] + assert await store.list_collections(OWNER_PRINCIPALS) == ["shared-id"] + assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["shared-id"] + + await store.delete_collection("shared-id", OWNER) + assert await store.has_collection("shared-id", OWNER_PRINCIPALS) is False + assert await store.has_collection("shared-id", OTHER_OWNER_PRINCIPALS) is True + + @pytest.mark.anyio + async def test_purge_owner_removes_only_that_owners_collections(self) -> None: + """Logout path: one owner's purge must not touch another owner's docs.""" + store = SqliteVecStore.ephemeral() + for owner, principals, name in ( + (OWNER, OWNER_PRINCIPALS, "a.pdf"), + (OWNER, OWNER_PRINCIPALS, "b.pdf"), + (OTHER_OWNER, OTHER_OWNER_PRINCIPALS, "c.pdf"), + ): + await store.ensure_collection(name, name, owner, None) + await store.grant_read(name, owner, principals) + await store.add_documents(name, [Document(id="1", text="x", metadata={})], [[1.0, 0.0]], owner) + + deleted = await store.purge_owner(OWNER) + assert deleted == 2 + assert await store.list_collections(OWNER_PRINCIPALS) == [] + assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["c.pdf"] + + @pytest.mark.anyio + async def test_reap_expired_drops_collections_past_expires_at(self) -> None: + """TTL backstop: rows with ``expires_at`` in the past go away on reap.""" + from datetime import UTC, datetime, timedelta + + store = SqliteVecStore.ephemeral() + now = datetime.now(UTC) + await store.ensure_collection("fresh", "fresh.pdf", OWNER, now + timedelta(hours=1)) + await store.grant_read("fresh", OWNER, OWNER_PRINCIPALS) + await store.ensure_collection("stale", "stale.pdf", OWNER, now - timedelta(seconds=1)) + await store.grant_read("stale", OWNER, OWNER_PRINCIPALS) + + deleted = await store.reap_expired() + assert deleted == 1 + assert await store.list_collections(OWNER_PRINCIPALS) == ["fresh"] + + @pytest.mark.anyio + async def test_reap_expired_keeps_persistent_and_unexpired(self) -> None: + """Rows with ``expires_at`` in the future stay. Rows with null + ``expires_at`` (org docs) are persistent and never touched, regardless + of age.""" + from datetime import UTC, datetime, timedelta + + store = SqliteVecStore.ephemeral() + await store.ensure_collection("session", "s.pdf", OWNER, datetime.now(UTC) + timedelta(hours=1)) + await store.grant_read("session", OWNER, OWNER_PRINCIPALS) + await store.ensure_collection("persistent", "p.pdf", OTHER_OWNER, None) + await store.grant_read("persistent", OTHER_OWNER, OTHER_OWNER_PRINCIPALS) + + deleted = await store.reap_expired() + assert deleted == 0 + assert await store.list_collections(OWNER_PRINCIPALS) == ["session"] + assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["persistent"] + + @pytest.mark.anyio + async def test_acl_grants_read_to_extra_principal(self) -> None: + """A principal without an ACL row can't read; once granted, they can.""" + store = SqliteVecStore.ephemeral() + team_principal = PrincipalId("group:engineering") + await store.ensure_collection("doc", "engineering-runbook.pdf", OWNER, None) + await store.grant_read("doc", OWNER, OWNER_PRINCIPALS) + await store.add_documents( + "doc", + [Document(id="1", text="how to deploy", metadata={})], + [[1.0, 0.0]], + OWNER, + ) + + # Engineering can't see it yet. + assert await store.has_collection("doc", [team_principal]) is False + + # Owner grants engineering read access. + await store.grant_read("doc", OWNER, [team_principal]) + assert await store.has_collection("doc", [team_principal]) is True + + # Revoke kills it. + await store.revoke("doc", OWNER, team_principal) + assert await store.has_collection("doc", [team_principal]) is False + # Owner still can. + assert await store.has_collection("doc", OWNER_PRINCIPALS) is True # DocumentService (with stub embedder) @@ -163,39 +288,131 @@ class TestDocumentService: @pytest.mark.anyio async def test_ingest_and_search(self, documents: DocumentService) -> None: text = "Python is great for data science. It has many libraries like pandas and numpy." - count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf") + count = await documents.ingest( + FileId("docs"), + _pages(text), + source="guide.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) assert count > 0 - results = await documents.search("Python libraries", collection=FileId("docs")) + results = await documents.search("Python libraries", principals=OWNER_PRINCIPALS, collection=FileId("docs")) assert len(results) > 0 assert results[0].document.text @pytest.mark.anyio async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None: - count = await documents.ingest(FileId("docs"), _pages(""), source="empty.pdf") + count = await documents.ingest( + FileId("docs"), + _pages(""), + source="empty.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) assert count == 0 @pytest.mark.anyio async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None: - results = await documents.search("anything", collection=FileId("nonexistent")) + results = await documents.search("anything", principals=OWNER_PRINCIPALS, collection=FileId("nonexistent")) assert results == [] @pytest.mark.anyio - async def test_search_all_collections(self, documents: DocumentService) -> None: - await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf") - await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf") + async def test_search_all_collections_for_principal(self, documents: DocumentService) -> None: + await documents.ingest( + FileId("col-a"), + _pages("Machine learning overview."), + source="ml.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + await documents.ingest( + FileId("col-b"), + _pages("Deep learning with neural networks."), + source="dl.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) - results = await documents.search("neural networks") + results = await documents.search("neural networks", principals=OWNER_PRINCIPALS) assert len(results) > 0 + @pytest.mark.anyio + async def test_search_does_not_cross_owner_boundary(self, documents: DocumentService) -> None: + """A search by one principal set never returns docs owned by an unrelated owner.""" + await documents.ingest( + FileId("col-a"), + _pages("Alice's private notes."), + source="alice.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + await documents.ingest( + FileId("col-b"), + _pages("Bob's private notes."), + source="bob.pdf", + owner_id=OTHER_OWNER, + read_principals=OTHER_OWNER_PRINCIPALS, + expires_at=None, + ) + + alice_results = await documents.search("notes", principals=OWNER_PRINCIPALS) + bob_results = await documents.search("notes", principals=OTHER_OWNER_PRINCIPALS) + alice_texts = [r.document.text for r in alice_results] + bob_texts = [r.document.text for r in bob_results] + assert any("Alice" in t for t in alice_texts) + assert not any("Bob" in t for t in alice_texts) + assert any("Bob" in t for t in bob_texts) + assert not any("Alice" in t for t in bob_texts) + + @pytest.mark.anyio + async def test_org_doc_visible_only_to_granted_group(self, documents: DocumentService) -> None: + """Org-owned doc with explicit ACL: only granted principals see it.""" + org_owner = OwnerId("org:acme") + eng_group = PrincipalId("group:engineering") + hr_group = PrincipalId("group:hr") + + await documents.ingest( + FileId("runbook"), + _pages("Production deploy steps."), + source="runbook.pdf", + owner_id=org_owner, + read_principals=[eng_group], + expires_at=None, + ) + + # Engineering can read. + eng_results = await documents.search("deploy", principals=[eng_group]) + assert len(eng_results) > 0 + + # HR cannot. + hr_results = await documents.search("deploy", principals=[hr_group]) + assert hr_results == [] + + # Caller with both memberships (or org-wide principal) still sees it via eng. + multi_results = await documents.search("deploy", principals=[hr_group, eng_group]) + assert len(multi_results) > 0 + @pytest.mark.anyio async def test_delete_collection(self, documents: DocumentService) -> None: - await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf") - collections = await documents.list_collections() + await documents.ingest( + FileId("temp"), + _pages("Temporary data."), + source="tmp.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + collections = await documents.list_collections(OWNER_PRINCIPALS) assert "temp" in collections - await documents.delete_collection(FileId("temp")) - collections = await documents.list_collections() + await documents.delete_collection(FileId("temp"), owner_id=OWNER) + collections = await documents.list_collections(OWNER_PRINCIPALS) assert "temp" not in collections @pytest.mark.anyio @@ -205,9 +422,16 @@ class TestDocumentService: PageText(page_number=2, text="Second page text."), PageText(page_number=3, text="Third page text."), ] - await documents.ingest(FileId("ordered"), pages, source="ordered.pdf") + await documents.ingest( + FileId("ordered"), + pages, + source="ordered.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) - stored = await documents.read_pages(FileId("ordered")) + stored = await documents.read_pages(FileId("ordered"), principals=OWNER_PRINCIPALS) assert [p.page_number for p in stored] == [1, 2, 3] assert stored[0].text == "First page text." assert stored[0].char_count == len("First page text.") @@ -217,9 +441,13 @@ class TestDocumentService: from stirling.contracts import PageRange pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)] - await documents.ingest(FileId("ranged"), pages, source="r.pdf") + await documents.ingest( + FileId("ranged"), pages, source="r.pdf", owner_id=OWNER, read_principals=OWNER_PRINCIPALS, expires_at=None + ) - subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4)) + subset = await documents.read_pages( + FileId("ranged"), principals=OWNER_PRINCIPALS, page_range=PageRange(start=2, end=4) + ) assert [p.page_number for p in subset] == [2, 3, 4] @pytest.mark.anyio @@ -228,14 +456,20 @@ class TestDocumentService: FileId("doc"), [PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")], source="v1.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) await documents.ingest( FileId("doc"), [PageText(page_number=1, text="new")], source="v2.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) - stored = await documents.read_pages(FileId("doc")) + stored = await documents.read_pages(FileId("doc"), principals=OWNER_PRINCIPALS) assert [p.page_number for p in stored] == [1] assert stored[0].text == "new" @@ -248,9 +482,16 @@ class TestDocumentService: PageText(page_number=2, text=" "), PageText(page_number=3, text="Real text on page 3."), ] - await documents.ingest(FileId("with-blanks"), pages, source="blanks.pdf") + await documents.ingest( + FileId("with-blanks"), + pages, + source="blanks.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) - stored = await documents.read_pages(FileId("with-blanks")) + stored = await documents.read_pages(FileId("with-blanks"), principals=OWNER_PRINCIPALS) assert [p.page_number for p in stored] == [1, 2, 3] assert stored[1].text.strip() == "" @@ -270,22 +511,36 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re class TestRagCapability: def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None: - cap = RagCapability(documents, collections=[FileId("docs"), FileId("manuals")]) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("docs"), FileId("manuals")]) instructions = cap.instructions assert isinstance(instructions, str) assert "docs, manuals" in instructions assert "search_knowledge" in instructions def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None: - cap = RagCapability(documents) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) instructions = cap.instructions assert callable(instructions) @pytest.mark.anyio async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None: - await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf") - await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf") - cap = RagCapability(documents) + await documents.ingest( + FileId("col-a"), + _pages("Alpha content."), + source="a.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + await documents.ingest( + FileId("col-b"), + _pages("Beta content."), + source="b.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) instructions_fn = cap.instructions assert callable(instructions_fn) text = await instructions_fn() @@ -294,7 +549,7 @@ class TestRagCapability: @pytest.mark.anyio async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None: - cap = RagCapability(documents) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) instructions_fn = cap.instructions assert callable(instructions_fn) text = await instructions_fn() @@ -302,14 +557,21 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None: - cap = RagCapability(documents) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) output = await _invoke_search_knowledge(cap, "anything") assert output == "No relevant results found in the knowledge base." @pytest.mark.anyio async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None: - await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf") - cap = RagCapability(documents) + await documents.ingest( + FileId("docs"), + _pages("Python is a programming language."), + source="guide.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) output = await _invoke_search_knowledge(cap, "Python") assert "[Result 1" in output assert "source: guide.pdf" in output @@ -318,10 +580,24 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None: - await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf") - await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf") + await documents.ingest( + FileId("pinned"), + _pages("Pinned collection content."), + source="pinned.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + await documents.ingest( + FileId("other"), + _pages("Content in another collection."), + source="other.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) - cap = RagCapability(documents, collections=[FileId("pinned")]) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("pinned")]) output = await _invoke_search_knowledge(cap, "content") assert "pinned.pdf" in output assert "other.pdf" not in output @@ -329,9 +605,16 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None: paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10)) - await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf") + await documents.ingest( + FileId("bulk"), + _pages(paragraphs), + source="bulk.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) - cap = RagCapability(documents) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS) output = await _invoke_search_knowledge(cap, "topic", max_results=2) assert "[Result 1" in output assert "[Result 2" in output @@ -341,8 +624,15 @@ class TestRagCapability: async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None: """The prepare callback must return None once max_searches has been reached so the agent can no longer call the tool on subsequent turns.""" - await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf") - cap = RagCapability(documents, max_searches=2) + await documents.ingest( + FileId("docs"), + _pages("Some content."), + source="x.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, + ) + cap = RagCapability(documents, principals=OWNER_PRINCIPALS, max_searches=2) tool_def = _dummy_tool_def() assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type] diff --git a/engine/tests/test_documents_routes.py b/engine/tests/test_documents_routes.py index 844d78518..4f317bb32 100644 --- a/engine/tests/test_documents_routes.py +++ b/engine/tests/test_documents_routes.py @@ -8,7 +8,11 @@ from fastapi.testclient import TestClient from stirling.api import app from stirling.api.dependencies import get_document_service from stirling.documents import Document, DocumentService, SqliteVecStore -from stirling.models import FileId +from stirling.models import FileId, PrincipalId, UserId + +USER = UserId("test-user") +USER_PRINCIPALS = [PrincipalId("test-user")] +HEADERS = {"X-User-Id": USER} class StubEmbedder: @@ -78,7 +82,11 @@ def test_ingest_document_indexes_page_text(client: TestClient, service: Document {"pageNumber": 1, "text": "The introduction covers the main topic."}, {"pageNumber": 2, "text": "The conclusion summarises the findings."}, ], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) assert response.status_code == 200 body = response.json() @@ -94,7 +102,11 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser "documentId": "replace-me", "source": "replace-me.pdf", "pageText": [{"pageNumber": 1, "text": "Original content that existed before."}], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) # Second ingest with different content should replace the first entirely response = client.post( @@ -103,11 +115,15 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser "documentId": "replace-me", "source": "replace-me.pdf", "pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) assert response.status_code == 200 - results = await service.search("New content", collection=FileId("replace-me"), top_k=5) + results = await service.search("New content", principals=USER_PRINCIPALS, collection=FileId("replace-me"), top_k=5) texts = [r.document.text for r in results] assert any("New content" in t for t in texts) assert not any("Original content" in t for t in texts) @@ -123,14 +139,28 @@ def test_ingest_document_skips_empty_pages(client: TestClient) -> None: {"pageNumber": 1, "text": " "}, {"pageNumber": 2, "text": "Real content on page 2."}, ], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) assert response.status_code == 200 assert response.json()["chunksIndexed"] >= 1 def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None: - response = client.post("/api/v1/documents", json={"documentId": "empty", "source": "empty.pdf"}) + response = client.post( + "/api/v1/documents", + json={ + "documentId": "empty", + "source": "empty.pdf", + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, + }, + headers=HEADERS, + ) assert response.status_code == 200 assert response.json()["chunksIndexed"] == 0 @@ -139,6 +169,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None: response = client.post( "/api/v1/documents", json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]}, + headers=HEADERS, ) assert response.status_code == 422 @@ -147,6 +178,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None: response = client.post( "/api/v1/documents", json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]}, + headers=HEADERS, ) assert response.status_code == 422 @@ -155,6 +187,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None: response = client.post( "/api/v1/documents", json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]}, + headers=HEADERS, ) assert response.status_code == 422 @@ -166,11 +199,61 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> "documentId": "bad-page", "source": "bad-page.pdf", "pageText": [{"pageNumber": 0, "text": "something"}], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) assert response.status_code == 422 +def test_ingest_document_rejects_missing_owner_id(client: TestClient) -> None: + """ownerId is required — never derived from the caller. Forgetting it must 422, + not silently fall back to personal-doc semantics.""" + response = client.post( + "/api/v1/documents", + json={ + "documentId": "no-owner", + "source": "no-owner.pdf", + "pageText": [{"pageNumber": 1, "text": "something"}], + "readPrincipals": [USER], + "expiresAt": None, + }, + headers=HEADERS, + ) + assert response.status_code == 422 + + +def test_ingest_document_rejects_empty_read_principals(client: TestClient) -> None: + """readPrincipals is required and must not be empty — every doc needs at least one reader.""" + response = client.post( + "/api/v1/documents", + json={ + "documentId": "no-readers", + "source": "no-readers.pdf", + "pageText": [{"pageNumber": 1, "text": "something"}], + "ownerId": USER, + "readPrincipals": [], + }, + headers=HEADERS, + ) + assert response.status_code == 422 + + +def test_ingest_document_rejects_missing_user_header(client: TestClient) -> None: + """The route must refuse to write per-user data when the caller didn't identify themselves.""" + response = client.post( + "/api/v1/documents", + json={ + "documentId": "doc-1", + "source": "x.pdf", + "pageText": [{"pageNumber": 1, "text": "something"}], + }, + ) + assert response.status_code == 401 + + # ── DELETE /documents/{id} ────────────────────────────────────────────── @@ -181,15 +264,19 @@ def test_delete_document_reports_deleted_true_when_existed(client: TestClient) - "documentId": "to-delete", "source": "to-delete.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, }, + headers=HEADERS, ) - response = client.delete("/api/v1/documents/to-delete") + response = client.delete("/api/v1/documents/by-id/to-delete", headers=HEADERS) assert response.status_code == 200 assert response.json() == {"documentId": "to-delete", "deleted": True} def test_delete_document_is_idempotent(client: TestClient) -> None: - response = client.delete("/api/v1/documents/never-existed") + response = client.delete("/api/v1/documents/by-id/never-existed", headers=HEADERS) assert response.status_code == 200 assert response.json() == {"documentId": "never-existed", "deleted": False} @@ -198,8 +285,86 @@ def test_delete_document_is_idempotent(client: TestClient) -> None: async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None: client.post( "/api/v1/documents", - json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]}, + json={ + "documentId": "gone", + "source": "gone.pdf", + "pageText": [{"pageNumber": 1, "text": "Text."}], + "ownerId": USER, + "readPrincipals": [USER], + "expiresAt": None, + }, + headers=HEADERS, ) - assert await service.has_collection(FileId("gone")) - client.delete("/api/v1/documents/gone") - assert not await service.has_collection(FileId("gone")) + assert await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS) + client.delete("/api/v1/documents/by-id/gone", headers=HEADERS) + assert not await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS) + + +def test_delete_document_rejects_missing_user_header(client: TestClient) -> None: + response = client.delete("/api/v1/documents/by-id/anything") + assert response.status_code == 401 + + +@pytest.mark.anyio +async def test_purge_by_owner_removes_only_callers_collections(client: TestClient, service: DocumentService) -> None: + """Logout path: DELETE /api/v1/documents/by-owner purges only the caller's docs.""" + alice = { + "documentId": "alice-doc", + "source": "alice.pdf", + "pageText": [{"pageNumber": 1, "text": "alice content"}], + "ownerId": "alice", + "readPrincipals": ["alice"], + "expiresAt": None, + } + bob = { + "documentId": "bob-doc", + "source": "bob.pdf", + "pageText": [{"pageNumber": 1, "text": "bob content"}], + "ownerId": "bob", + "readPrincipals": ["bob"], + "expiresAt": None, + } + client.post("/api/v1/documents", json=alice, headers={"X-User-Id": "alice"}) + client.post("/api/v1/documents", json=bob, headers={"X-User-Id": "bob"}) + + response = client.delete("/api/v1/documents/by-owner", headers={"X-User-Id": "alice"}) + assert response.status_code == 200 + assert response.json() == {"ownerId": "alice", "deleted": 1} + + # Alice gone, Bob still there. + assert await service.has_collection(FileId("alice-doc"), principals=[PrincipalId("alice")]) is False + assert await service.has_collection(FileId("bob-doc"), principals=[PrincipalId("bob")]) is True + + +def test_purge_by_owner_is_idempotent(client: TestClient) -> None: + """Calling purge with no docs is fine — deleted=0 and no error.""" + response = client.delete("/api/v1/documents/by-owner", headers=HEADERS) + assert response.status_code == 200 + assert response.json() == {"ownerId": USER, "deleted": 0} + + +def test_purge_by_owner_rejects_missing_user_header(client: TestClient) -> None: + response = client.delete("/api/v1/documents/by-owner") + assert response.status_code == 401 + + +def test_delete_document_only_affects_calling_user(client: TestClient) -> None: + """Two users with the same document id: one user's delete must not remove the other's.""" + alice_body = { + "documentId": "shared", + "source": "shared.pdf", + "pageText": [{"pageNumber": 1, "text": "x"}], + "ownerId": "alice", + "readPrincipals": ["alice"], + "expiresAt": None, + } + bob_body = {**alice_body, "ownerId": "bob", "readPrincipals": ["bob"], "expiresAt": None} + client.post("/api/v1/documents", json=alice_body, headers={"X-User-Id": "alice"}) + client.post("/api/v1/documents", json=bob_body, headers={"X-User-Id": "bob"}) + + alice_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "alice"}) + assert alice_delete.json() == {"documentId": "shared", "deleted": True} + + # Bob's copy is still there + bob_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "bob"}) + assert bob_delete.json() == {"documentId": "shared", "deleted": True} diff --git a/engine/tests/test_pdf_question_agent.py b/engine/tests/test_pdf_question_agent.py index 26bd5c0ff..c26d4b504 100644 --- a/engine/tests/test_pdf_question_agent.py +++ b/engine/tests/test_pdf_question_agent.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Iterator from dataclasses import replace import pytest @@ -19,9 +20,25 @@ from stirling.contracts import ( SupportedCapability, ) from stirling.documents import Document, DocumentService, SqliteVecStore -from stirling.models import FileId +from stirling.models import FileId, OwnerId, PrincipalId, UserId +from stirling.services import current_user_id from stirling.services.runtime import AppRuntime +USER = UserId("test-user") +OWNER = OwnerId("test-user") +OWNER_PRINCIPALS = [PrincipalId("test-user")] + + +@pytest.fixture(autouse=True) +def _set_user_context() -> Iterator[None]: + """Set current_user_id to a fixed test user for every test in this module so + agent code calling require_current_user_id() doesn't fail closed.""" + token = current_user_id.set(USER) + try: + yield + finally: + current_user_id.reset(token) + class StubEmbedder: """Deterministic embeddings so RAG lookups work in tests without network.""" @@ -94,6 +111,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> FileId("present-id"), [PageText(page_number=1, text="Invoice total: 120.00.")], source="present.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = PdfQuestionAgent(runtime_with_stub_rag) @@ -111,6 +131,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub FileId("invoice-id"), [PageText(page_number=1, text="Invoice total: 120.00.")], source="invoice.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = StubPdfQuestionAgent( runtime_with_stub_rag, @@ -142,6 +165,9 @@ async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: A FileId("shipping-id"), [PageText(page_number=1, text="This page contains only a shipping address.")], source="shipping.pdf", + owner_id=OWNER, + read_principals=OWNER_PRINCIPALS, + expires_at=None, ) agent = StubPdfQuestionAgent( runtime_with_stub_rag,