improvement(search): rank an overflowing Confluence all-terms match with relevance search

apps/sim/lib/sim-search/live/account-session.ts

@@ -205,9 +205,11 @@ export async function openLiveAccountSession(
           client ? searchNativeProvider(provider, client, native) : searchMcp(native)
         )
       }
+      /** Providers see the session's cancellation, whatever its reason, on every request. */
+      const request = search.signal ? search : { ...search, signal }
       return client && requiresEveryTerm(provider)
-        ? searchWithAnyTermFallback(searchOnce, search, signal)
-        : searchOnce(search)
+        ? searchWithAnyTermFallback(searchOnce, request, signal)
+        : searchOnce(request)
     },
     async verify(document) {
       if (!(await boundary.verify(document))) return false

apps/sim/lib/sim-search/live/application.test.ts

@@ -697,6 +697,17 @@ describe('authorized live retrieval', () => {
     })
   })
 
+  it('gives provider searches the account session cancellation signal', async () => {
+    const signals: (AbortSignal | undefined)[] = []
+    mocks.search.mockImplementation(async (_provider, _client, search) => {
+      signals.push(search.signal)
+      return { documents: [document] }
+    })
+    await searchLiveKnowledge.execute({ principal, input: { ...input, query: 'launch checklist' } })
+    expect(signals.length).toBeGreaterThan(0)
+    for (const signal of signals) expect(signal).toBeInstanceOf(AbortSignal)
+  })
+
   it('retries an underfilled search with any term for providers that require every term', async () => {
     mocks.requiresEveryTerm.mockReturnValue(true)
     const loose = { ...document, id: 'loose', url: 'https://docs.google.com/document/d/loose/edit' }

apps/sim/lib/sim-search/live/atlassian.test.ts

@@ -245,3 +245,89 @@ describe('Jira any-term search', () => {
     expect(page.documents.map((document) => document.id)).toEqual(['5'])
   })
 })
+
+describe('Confluence plain search ranking', () => {
+  const search = '/ex/confluence/cloud/wiki/rest/api/search'
+  /** Answers the all-terms query with `exact` and the relevance query with `ranked`, or its failure. */
+  function confluence(exact: { ids: string[]; more: boolean }, ranked: string[] | (() => never)) {
+    const results = (ids: string[]) =>
+      ids.map((id) => ({ content: { id, type: 'page', title: id } }))
+    const json = vi.fn(async (path: string, options?: { query?: Record<string, string> }) => {
+      if (path === '/oauth/token/accessible-resources') return [SITE]
+      if (path !== search) throw new Error(`Unexpected request: ${path}`)
+      if (options?.query?.cql?.includes('siteSearch ~'))
+        return { results: results(typeof ranked === 'function' ? ranked() : ranked), _links: {} }
+      return {
+        results: results(exact.ids),
+        _links: exact.more ? { next: `${search}?cursor=page-2` } : {},
+      }
+    })
+    return { json, text: vi.fn(), bytes: vi.fn() } satisfies NativeClient
+  }
+  const query = { query: 'merge train service', limit: 2, scopes: [] }
+  const overflowing = { ids: ['old-1', 'old-2'], more: true }
+
+  it('ranks an all-terms match too large for one page with relevance search', async () => {
+    const api = confluence(overflowing, ['best', 'next'])
+    const page = await searchAtlassian(api, 'confluence', query)
+    expect(page.documents.map((document) => document.id)).toEqual(['best', 'next'])
+    expect(page.hasMore).toBe(true)
+    expect(page.partial).toBe(true)
+    expect(page.nextCursor).toBeUndefined()
+  })
+
+  it.each([
+    ['a provider refusal', new NativeSearchError('rate_limited', 'Rate limited.', 30)],
+    ['a request timeout', new Error('Request timed out after 10000ms')],
+  ])(
+    'keeps the all-terms page and its cursor when relevance ranking fails on %s',
+    async (_, failure) => {
+      const api = confluence(overflowing, () => {
+        throw failure
+      })
+      const page = await searchAtlassian(api, 'confluence', query)
+      expect(page.documents.map((document) => document.id)).toEqual(['old-1', 'old-2'])
+      expect(page.nextCursor).toBe('page-2')
+      expect(page.partial).toBe(true)
+    }
+  )
+
+  it.each([
+    ['an Error', new Error('Caller stopped')],
+    ['a string', 'AbortReason.UserStop'],
+    ['a DOMException', new DOMException('Search cancelled', 'AbortError')],
+  ])('propagates a cancellation whose reason is %s', async (_, reason) => {
+    const caller = new AbortController()
+    const api = confluence(overflowing, () => {
+      caller.abort(reason)
+      throw reason
+    })
+    await expect(
+      searchAtlassian(api, 'confluence', { ...query, signal: caller.signal })
+    ).rejects.toBe(reason)
+  })
+
+  it('keeps an all-terms match that fits on one page', async () => {
+    const api = confluence({ ids: ['exact'], more: false }, ['loose'])
+    const page = await searchAtlassian(api, 'confluence', query)
+    expect(page.documents.map((document) => document.id)).toEqual(['exact'])
+    expect(api.json).toHaveBeenCalledTimes(2)
+  })
+
+  it('keeps date order and native queries on the all-terms page', async () => {
+    const dated = confluence({ ids: ['newest', 'older'], more: true }, ['best'])
+    const byDate = await searchAtlassian(dated, 'confluence', {
+      ...query,
+      filters: { sortBy: 'newest' },
+    })
+    expect(byDate.documents.map((document) => document.id)).toEqual(['newest', 'older'])
+    expect(byDate.nextCursor).toBe('page-2')
+    const native = confluence({ ids: ['a', 'b'], more: true }, ['best'])
+    const own = await searchAtlassian(native, 'confluence', {
+      ...query,
+      native: { provider: 'confluence', query: 'text ~ "merge train"' },
+    })
+    expect(own.documents.map((document) => document.id)).toEqual(['a', 'b'])
+    expect(own.nextCursor).toBe('page-2')
+  })
+})

apps/sim/lib/sim-search/live/atlassian.ts

@@ -116,6 +116,17 @@ async function sites(client: NativeClient) {
   return array(await client.json('/oauth/token/accessible-resources', { memo: true }))
 }
 
+/**
+ * One Atlassian site's page; `ranking` marks a page whose order or completeness differs from the
+ * all-terms search.
+ */
+interface SitePage {
+  documents: NativeDocument[]
+  excluded: boolean
+  next?: string
+  ranking?: 'relevance' | 'unavailable'
+}
+
 export async function searchAtlassian(
   client: NativeClient,
   provider: 'jira' | 'confluence',
@@ -160,7 +171,7 @@ export async function searchAtlassian(
   // keep only issues holding at least half the terms.
   const relevant = input.textMatch === 'any' && text ? holdsHalfTheTerms(text) : undefined
   const pages = await Promise.all(
-    selected.map(async (site) => {
+    selected.map(async (site): Promise<SitePage> => {
       const cloudId = string(site.id)
       const origin = string(site.url).replace(/\/$/, '')
       if (provider === 'jira') {
@@ -193,47 +204,71 @@ export async function searchAtlassian(
           next: string(data.nextPageToken) || undefined,
         }
       }
-      const data = object(
-        await client.json(`/ex/confluence/${segment(cloudId)}/wiki/rest/api/search`, {
-          query: {
-            cql: scopeAtlassianQuery(
-              input.native?.query ||
-                `type IN (page, blogpost)${text ? ` AND ${confluenceTextClause(text, input.textMatch)}` : ''}`,
-              scope,
-              order
-            ),
-            limit: String(input.limit),
-            expand: 'content.version,content.space',
-            ...(input.native?.cursor && single ? { cursor: input.native.cursor } : {}),
-          },
-          signal: input.signal,
-        })
+      const searchPages = async (clause: string, cursor?: string) => {
+        const data = object(
+          await client.json(`/ex/confluence/${segment(cloudId)}/wiki/rest/api/search`, {
+            query: {
+              cql: scopeAtlassianQuery(
+                input.native?.query || `type IN (page, blogpost)${clause ? ` AND ${clause}` : ''}`,
+                scope,
+                order
+              ),
+              limit: String(input.limit),
+              expand: 'content.version,content.space',
+              ...(cursor ? { cursor } : {}),
+            },
+            signal: input.signal,
+          })
+        )
+        const next = string(object(data._links).next)
+        const rows = array(data.results)
+        const documents = rows
+          .map((row) => page(row, cloudId, origin))
+          .filter((document) => document !== undefined)
+        return {
+          documents,
+          excluded: documents.length < rows.length,
+          next: next
+            ? (new URL(next, 'https://api.atlassian.com').searchParams.get('cursor') ?? undefined)
+            : undefined,
+        }
+      }
+      const exact = await searchPages(
+        text ? confluenceTextClause(text, input.textMatch) : '',
+        input.native?.cursor && single ? input.native.cursor : undefined
       )
-      const next = string(object(data._links).next)
-      const rows = array(data.results)
-      const documents = rows
-        .map((row) => page(row, cloudId, origin))
-        .filter((document) => document !== undefined)
-      return {
-        documents,
-        excluded: documents.length < rows.length,
-        next: next
-          ? (new URL(next, 'https://api.atlassian.com').searchParams.get('cursor') ?? undefined)
-          : undefined,
+      // An all-terms match larger than a page comes back by raw term matches; rank it by relevance.
+      if (!text || input.native || order || input.textMatch === 'any' || !exact.next) return exact
+      try {
+        const ranked = await searchPages(confluenceTextClause(text, 'any'))
+        return { ...ranked, next: undefined, ranking: 'relevance' }
+      } catch {
+        // A cancelled search stops here; any other failure keeps the all-terms page.
+        input.signal?.throwIfAborted()
+        return { ...exact, ranking: 'unavailable' }
       }
     })
   )
   const documents = interleaveByRank(pages.map((result) => result.documents))
   const excluded = pages.some((result) => result.excluded)
   return {
     documents,
-    partial: excluded || (!input.native?.project && allSites.length > selected.length),
-    hasMore: pages.some((result) => Boolean(result.next)),
+    partial:
+      excluded ||
+      pages.some((result) => result.ranking) ||
+      (!input.native?.project && allSites.length > selected.length),
+    hasMore: pages.some((result) => Boolean(result.next) || result.ranking === 'relevance'),
     nextCursor: single ? pages[0]?.next : undefined,
     message:
       'Searches up to four accessible Atlassian sites. For a specific site and pagination, set project to its cloud ID.' +
       (excluded
         ? ' Matches other than pages, blog posts, and spaces (attachments, comments, whiteboards, folders, databases) were excluded because they cannot be read.'
+        : '') +
+      (pages.some((result) => result.ranking === 'relevance')
+        ? ' More pages contained every word than fit on one page, so this page is ordered by Confluence relevance search; refine the query, or page with a native CQL query.'
+        : '') +
+      (pages.some((result) => result.ranking === 'unavailable')
+        ? ' Relevance ranking was unavailable, so this page keeps the order of the all-terms search.'
         : ''),
   }
 }