From f639c4cba6fd6fd821daea24472c4e351259929b Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Thu, 11 Jun 2026 12:40:07 +0700
Subject: [PATCH] fix(protocol): allow future agent capability flags

---
 packages/protocol/src/agent-types.ts          |  2 ++
 packages/protocol/src/messages.ts             | 31 ++++++++++---------
 .../src/server/agent/agent-sdk-types.ts       |  1 +
 .../src/server/agent/mcp-server.test.ts       | 11 +++++++
 4 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/packages/protocol/src/agent-types.ts b/packages/protocol/src/agent-types.ts
index e0728f592..2c5cf7224 100644
--- a/packages/protocol/src/agent-types.ts
+++ b/packages/protocol/src/agent-types.ts
@@ -131,8 +131,10 @@ export interface AgentFeatureSelect {
 export type AgentFeature = AgentFeatureToggle | AgentFeatureSelect;
 
 export interface AgentCapabilityFlags {
+  [capability: string]: boolean | undefined;
   supportsStreaming: boolean;
   supportsSessionPersistence: boolean;
+  supportsSessionListing?: boolean;
   supportsDynamicModes: boolean;
   supportsMcpServers: boolean;
   supportsReasoningStream: boolean;
diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts
index a241f71b2..fd161a840 100644
--- a/packages/protocol/src/messages.ts
+++ b/packages/protocol/src/messages.ts
@@ -237,20 +237,23 @@ export const ProviderSnapshotEntrySchema = z.object({
   defaultModeId: z.string().nullable().optional(),
 });
 
-const AgentCapabilityFlagsSchema: z.ZodType<AgentCapabilityFlags> = z.object({
-  supportsStreaming: z.boolean(),
-  supportsSessionPersistence: z.boolean(),
-  supportsDynamicModes: z.boolean(),
-  supportsMcpServers: z.boolean(),
-  supportsReasoningStream: z.boolean(),
-  supportsToolInvocations: z.boolean(),
-  // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
-  supportsRewindConversation: z.boolean().optional().default(false),
-  // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
-  supportsRewindFiles: z.boolean().optional().default(false),
-  // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
-  supportsRewindBoth: z.boolean().optional().default(false),
-});
+const AgentCapabilityFlagsSchema: z.ZodType<AgentCapabilityFlags> = z
+  .object({
+    supportsStreaming: z.boolean(),
+    supportsSessionPersistence: z.boolean(),
+    supportsSessionListing: z.boolean().optional(),
+    supportsDynamicModes: z.boolean(),
+    supportsMcpServers: z.boolean(),
+    supportsReasoningStream: z.boolean(),
+    supportsToolInvocations: z.boolean(),
+    // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
+    supportsRewindConversation: z.boolean().optional().default(false),
+    // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
+    supportsRewindFiles: z.boolean().optional().default(false),
+    // COMPAT(rewind): added in v0.1.X, drop when floor >= v0.1.X.
+    supportsRewindBoth: z.boolean().optional().default(false),
+  })
+  .catchall(z.boolean());
 
 const AgentUsageSchema: z.ZodType<AgentUsage> = z.object({
   inputTokens: z.number().optional(),
diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts
index 61869e43f..f6fd50d20 100644
--- a/packages/server/src/server/agent/agent-sdk-types.ts
+++ b/packages/server/src/server/agent/agent-sdk-types.ts
@@ -160,6 +160,7 @@ export interface AgentFeatureSelect {
 export type AgentFeature = AgentFeatureToggle | AgentFeatureSelect;
 
 export interface AgentCapabilityFlags {
+  [capability: string]: boolean | undefined;
   supportsStreaming: boolean;
   supportsSessionPersistence: boolean;
   supportsSessionListing?: boolean;
diff --git a/packages/server/src/server/agent/mcp-server.test.ts b/packages/server/src/server/agent/mcp-server.test.ts
index e9a0cde1f..f88c50e56 100644
--- a/packages/server/src/server/agent/mcp-server.test.ts
+++ b/packages/server/src/server/agent/mcp-server.test.ts
@@ -3306,6 +3306,16 @@ describe("agent snapshot MCP serialization", () => {
           provider: "codex",
           sessionId: "session-full",
         },
+        capabilities: {
+          supportsStreaming: false,
+          supportsSessionPersistence: false,
+          supportsSessionListing: true,
+          supportsDynamicModes: false,
+          supportsMcpServers: true,
+          supportsReasoningStream: false,
+          supportsToolInvocations: true,
+          supportsNewProviderCapability: true,
+        },
       }),
     );
 
@@ -3325,6 +3335,7 @@ describe("agent snapshot MCP serialization", () => {
         `get_agent_status response failed AgentSnapshotPayloadSchema: ${JSON.stringify(parsed.error.issues, null, 2)}`,
       );
     }
+    expectOutputSchemaAccepts(tool, response.structuredContent);
     expect(response.structuredContent.status).toBe("idle");
     expect(snapshot).toEqual(
       expect.objectContaining({