diff --git a/src/mcp/server.ts b/src/mcp/server.ts index dbca07d9b..0563d0b82 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -253,7 +253,20 @@ export function registerMcpEndpoints( } case "memory_sessions": { - const sessions = await kv.list(KV.sessions); + const rawLimit = Number(args.limit); + const limit = + Number.isFinite(rawLimit) && rawLimit > 0 + ? Math.floor(rawLimit) + : 20; + const allSessions = await kv.list(KV.sessions); + const recencyKey = (s: Session): string => + [s.updatedAt, s.endedAt, s.lastCheckpointAt, s.startedAt].reduce( + (acc: string, t) => (typeof t === "string" && t > acc ? t : acc), + "", + ); + const sessions = allSessions + .sort((a, b) => (recencyKey(a) < recencyKey(b) ? 1 : -1)) + .slice(0, limit); return { status_code: 200, body: { diff --git a/src/mcp/tools-registry.ts b/src/mcp/tools-registry.ts index c4df3499c..6ab01bfd9 100644 --- a/src/mcp/tools-registry.ts +++ b/src/mcp/tools-registry.ts @@ -116,7 +116,15 @@ export const CORE_TOOLS: McpToolDef[] = [ name: "memory_sessions", description: "List recent sessions with their status and observation counts.", - inputSchema: { type: "object", properties: {} }, + inputSchema: { + type: "object", + properties: { + limit: { + type: "number", + description: "Max sessions to return (default 20)", + }, + }, + }, }, { name: "memory_smart_search", diff --git a/src/triggers/api.ts b/src/triggers/api.ts index 7b6c2bf2b..4b6791513 100644 --- a/src/triggers/api.ts +++ b/src/triggers/api.ts @@ -808,6 +808,7 @@ export function registerApiTriggers( async (req: ApiRequest): Promise => { const authErr = checkAuth(req, secret); if (authErr) return authErr; + const rawLimit = parseOptionalInt(req.query_params?.["limit"]); const sessions = await kv.list(KV.sessions); const normalizedAgentId = typeof req.query_params?.["agentId"] === "string" @@ -823,12 +824,24 @@ export function registerApiTriggers( const filtered = filterAgentId ? sessions.filter((s) => s.agentId === filterAgentId) : sessions; + const recencyKey = (s: Session): string => + [s.updatedAt, s.endedAt, s.lastCheckpointAt, s.startedAt].reduce( + (acc: string, t) => (typeof t === "string" && t > acc ? t : acc), + "", + ); + const sorted = filtered.sort((a, b) => + recencyKey(a) < recencyKey(b) ? 1 : -1, + ); + const recent = + rawLimit === undefined + ? sorted + : sorted.slice(0, Math.max(0, rawLimit)); const summaries = await Promise.all( - filtered.map((s) => + recent.map((s) => kv.get(KV.summaries, s.id).catch(() => null), ), ); - const withSummary = filtered.map((s, i) => + const withSummary = recent.map((s, i) => summaries[i] ? { ...s, summary: summaries[i] } : s, ); return { status_code: 200, body: { sessions: withSummary } }; diff --git a/src/types.ts b/src/types.ts index 6797dfaf9..8fe7b7471 100644 --- a/src/types.ts +++ b/src/types.ts @@ -12,6 +12,8 @@ export interface Session { summary?: string; commitShas?: string[]; agentId?: string; + updatedAt?: string; + lastCheckpointAt?: string; } export interface CommitLink { diff --git a/test/api-sessions-limit.test.ts b/test/api-sessions-limit.test.ts new file mode 100644 index 000000000..b910450f1 --- /dev/null +++ b/test/api-sessions-limit.test.ts @@ -0,0 +1,135 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../src/logger.js", () => ({ + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, +})); + +import { registerApiTriggers } from "../src/triggers/api.js"; +import { KV } from "../src/state/schema.js"; +import type { Session } from "../src/types.js"; + +function mockKV() { + const store = new Map>(); + return { + get: async (scope: string, key: string): Promise => + (store.get(scope)?.get(key) as T) ?? null, + set: async (scope: string, key: string, value: T): Promise => { + if (!store.has(scope)) store.set(scope, new Map()); + store.get(scope)!.set(key, value); + return value; + }, + update: async (scope: string, key: string): Promise => + store.get(scope)?.get(key) as T, + delete: async (scope: string, key: string): Promise => { + store.get(scope)?.delete(key); + }, + list: async (scope: string): Promise => { + const entries = store.get(scope); + return entries ? (Array.from(entries.values()) as T[]) : []; + }, + }; +} + +function mockSdk() { + const functions = new Map(); + return { + registerFunction: vi.fn( + (idOrOpts: string | { id: string }, handler: Function) => { + const id = typeof idOrOpts === "string" ? idOrOpts : idOrOpts.id; + functions.set(id, handler); + }, + ), + registerTrigger: vi.fn(), + trigger: vi.fn(), + getFunction: (id: string) => functions.get(id), + }; +} + +function session(id: string, updatedSecond: number): Session { + return { + id, + project: "/tmp/p", + cwd: "/tmp/p", + startedAt: "2026-06-01T00:00:00.000Z", + updatedAt: `2026-06-01T00:00:${String(updatedSecond).padStart(2, "0")}.000Z`, + status: "active", + observationCount: 1, + }; +} + +function reqWithLimit(limit?: string) { + return { + body: undefined, + headers: {}, + query_params: limit === undefined ? {} : { limit }, + }; +} + +type SessionsResponse = { + status_code: number; + body: { sessions: Array<{ id: string }> }; +}; + +describe("api::sessions limit + recency", () => { + let sdk: ReturnType; + let kv: ReturnType; + + async function seed(count: number) { + for (let i = 0; i < count; i++) { + await kv.set(KV.sessions, `s${i}`, session(`s${i}`, i % 60)); + } + } + + beforeEach(() => { + sdk = mockSdk(); + kv = mockKV(); + registerApiTriggers(sdk as never, kv as never); + }); + + it("defaults to all sessions, most-recent first", async () => { + await seed(25); + const fn = sdk.getFunction("api::sessions")!; + const res = (await fn(reqWithLimit())) as SessionsResponse; + + expect(res.status_code).toBe(200); + expect(res.body.sessions).toHaveLength(25); + expect(res.body.sessions[0].id).toBe("s24"); + const ids = new Set(res.body.sessions.map((s) => s.id)); + for (const old of ["s0", "s1", "s2", "s3", "s4"]) { + expect(ids.has(old)).toBe(true); + } + }); + + it("honors an explicit limit", async () => { + await seed(25); + const fn = sdk.getFunction("api::sessions")!; + const res = (await fn(reqWithLimit("5"))) as SessionsResponse; + + expect(res.body.sessions).toHaveLength(5); + expect(res.body.sessions.map((s) => s.id)).toEqual([ + "s24", + "s23", + "s22", + "s21", + "s20", + ]); + }); + + it("honors arbitrarily large limit", async () => { + await seed(205); + const fn = sdk.getFunction("api::sessions")!; + const res = (await fn(reqWithLimit("1000"))) as SessionsResponse; + + expect(res.body.sessions).toHaveLength(205); + }); + + it("treats non-positive limit as zero", async () => { + await seed(5); + const fn = sdk.getFunction("api::sessions")!; + const res = (await fn(reqWithLimit("0"))) as SessionsResponse; + + expect(res.body.sessions).toHaveLength(0); + const negative = (await fn(reqWithLimit("-1"))) as SessionsResponse; + expect(negative.body.sessions).toHaveLength(0); + }); +});