| 1 | package protocol |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "reflect" |
| 6 | "strings" |
| 7 | "testing" |
| 8 | ) |
| 9 | |
| 10 | // methodFixtures holds one fully populated representative value for every |
| 11 | // registered params and result DTO. Round-tripping through the strict |
| 12 | // direction decoders proves the JSON shape is lossless. |
| 13 | var methodFixtures = map[Method]struct { |
| 14 | params any |
| 15 | result any |
| 16 | }{ |
| 17 | MethodExtensionInitialize: { |
| 18 | params: InitializeParams{ |
| 19 | ProtocolVersion: "1", |
| 20 | ProtocolID: ProtocolID, |
| 21 | Manifest: ManifestExpectation{ |
| 22 | Intercepts: []string{"tool.before"}, |
| 23 | Replaces: []string{"tool:bash"}, |
| 24 | Providers: []string{"acme"}, |
| 25 | UIActions: []string{"acme.refresh"}, |
| 26 | Capabilities: []string{"content_refs"}, |
| 27 | }, |
| 28 | Session: SessionContext{SessionID: "s-1", WorkspaceRoot: "/repo", Generation: 3}, |
| 29 | Capabilities: HostCapabilities{ContentRefs: true, UIHost: UIHostDesktop, ProtocolVersion: "1"}, |
| 30 | }, |
| 31 | result: InitializeResult{ |
| 32 | ProtocolVersion: "1", Name: "acme", Version: "1.2.3", |
| 33 | Subscriptions: []string{"tool.before"}, |
| 34 | Replaces: []string{"tool:bash"}, |
| 35 | Providers: []ProviderDescriptor{{ |
| 36 | Ref: "acme", DisplayName: "Acme", Model: "acme-1", ContextWindow: 128000, |
| 37 | PricingCurrency: "$", CacheHitPerMillion: 0.1, InputPerMillion: 1, OutputPerMillion: 2, |
| 38 | Vision: true, Tools: true, Reasoning: true, |
| 39 | Efforts: []string{"low", "high"}, DefaultEffort: "low", |
| 40 | ToolCallReasoning: true, ReasoningRoundTrip: true, WarnOnMissingToolCallReasoning: true, |
| 41 | }}, |
| 42 | UIActions: []UIActionDecl{{ActionID: "acme.refresh", Label: "Refresh"}}, |
| 43 | StateSchemaVersion: 2, |
| 44 | }, |
| 45 | }, |
| 46 | MethodExtensionInitialized: {params: InitializedParams{}}, |
| 47 | MethodExtensionShutdown: { |
| 48 | params: ShutdownParams{TimeoutMillis: 5000}, |
| 49 | result: ShutdownResult{Accepted: true}, |
| 50 | }, |
| 51 | MethodExtensionIntercept: { |
| 52 | params: InterceptParams{ |
| 53 | Event: EventToolBefore, Seq: 7, |
| 54 | Payload: json.RawMessage(`{"tool":"bash"}`), |
| 55 | TimeoutMillis: 250, |
| 56 | }, |
| 57 | result: InterceptResult{ |
| 58 | Decision: DecisionReplace, |
| 59 | Reason: "rewritten", |
| 60 | Replacement: json.RawMessage(`{"tool":"read"}`), |
| 61 | }, |
| 62 | }, |
| 63 | MethodExtensionEvent: { |
| 64 | params: EventParams{Event: EventSessionStart, Payload: json.RawMessage(`{"sessionId":"s-1"}`)}, |
| 65 | }, |
| 66 | MethodExtensionResourcesChanged: { |
| 67 | params: ResourcesChangedParams{Paths: []string{"skills/a", "commands/b"}}, |
| 68 | }, |
| 69 | MethodExtensionProviderCatalog: { |
| 70 | params: ProviderCatalogParams{}, |
| 71 | result: ProviderCatalogResult{Providers: []ProviderDescriptor{{Ref: "acme"}}}, |
| 72 | }, |
| 73 | MethodExtensionProviderStreamOpen: { |
| 74 | params: StreamOpenParams{ |
| 75 | StreamID: "st-1", ProviderRef: "acme", Model: "acme-1", Effort: "high", SeqBase: 1, |
| 76 | Request: ProviderRequest{ |
| 77 | Messages: []ProviderMessage{{ |
| 78 | Role: ProviderRoleAssistant, Content: "hi", |
| 79 | Images: []string{"data:image/png;base64,AA=="}, |
| 80 | ReasoningContent: "thinking", |
| 81 | ReasoningSignature: "sig", |
| 82 | ToolCalls: []ProviderToolCall{{ID: "c1", Name: "bash", Arguments: "{}", ThoughtSignature: "ts"}}, |
| 83 | ToolCallID: "c1", |
| 84 | Name: "bash", |
| 85 | }}, |
| 86 | Tools: []ProviderToolSchema{{Name: "bash", Description: "run", Parameters: json.RawMessage(`{"type":"object"}`)}}, |
| 87 | Temperature: floatPtr(0.5), |
| 88 | MaxTokens: 1024, |
| 89 | }, |
| 90 | }, |
| 91 | result: StreamOpenResult{Accepted: true}, |
| 92 | }, |
| 93 | MethodExtensionProviderStreamCancel: { |
| 94 | params: StreamCancelParams{StreamID: "st-1"}, |
| 95 | result: StreamCancelResult{Cancelled: true}, |
| 96 | }, |
| 97 | MethodExtensionProviderStreamChunk: { |
| 98 | params: StreamChunkParams{ |
| 99 | StreamID: "st-1", Seq: 2, |
| 100 | Chunk: ProviderChunk{ |
| 101 | Type: ChunkUsage, ArgChars: 0, |
| 102 | Usage: &ProviderUsage{ |
| 103 | PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3, |
| 104 | CacheHitTokens: 4, CacheMissTokens: 5, ReasoningTokens: 6, |
| 105 | FinishReason: "stop", |
| 106 | }, |
| 107 | }, |
| 108 | }, |
| 109 | }, |
| 110 | MethodExtensionProviderStreamEnd: { |
| 111 | params: StreamEndParams{StreamID: "st-1", LastSeq: 9, Error: "", Interrupted: true}, |
| 112 | }, |
| 113 | MethodExtensionUIAction: { |
| 114 | params: UIActionParams{ |
| 115 | ActionID: "acme.refresh", SessionID: "s-1", Generation: 3, |
| 116 | Args: map[string]string{"k": "v"}, |
| 117 | }, |
| 118 | result: UIActionResult{Accepted: true, Message: "queued"}, |
| 119 | }, |
| 120 | MethodExtensionUISubmit: { |
| 121 | params: UISubmitParams{ |
| 122 | SurfaceID: "form-1", SessionID: "s-1", Generation: 3, |
| 123 | Values: map[string]any{"name": "x", "count": float64(2), "ok": true}, |
| 124 | }, |
| 125 | result: UISubmitResult{Accepted: true}, |
| 126 | }, |
| 127 | MethodHostUIPublish: { |
| 128 | params: UIPublishParams{ |
| 129 | SurfaceID: "card-1", SessionID: "s-1", Generation: 3, |
| 130 | Kind: UISurfaceCard, |
| 131 | Payload: json.RawMessage(`{"title":"t"}`), |
| 132 | }, |
| 133 | result: UIPublishResult{Accepted: true}, |
| 134 | }, |
| 135 | MethodHostUIRequest: { |
| 136 | params: UIRequestParams{ |
| 137 | SurfaceID: "ask-1", SessionID: "s-1", Generation: 3, |
| 138 | Kind: UIRequestSelect, |
| 139 | Payload: json.RawMessage(`{"fields":[]}`), |
| 140 | }, |
| 141 | result: UIRequestResult{Cancelled: false, Values: map[string]any{"choice": "a"}}, |
| 142 | }, |
| 143 | MethodHostContentRead: { |
| 144 | params: ContentReadParams{ContentRef: "cref-1", Offset: 0}, |
| 145 | result: ContentReadResult{ |
| 146 | ContentRef: "cref-1", Offset: 0, DataBase64: "aGk=", |
| 147 | NextOffset: int64Ptr(2), TotalBytes: 2, |
| 148 | SHA256: strings.Repeat("a", 64), |
| 149 | Encoding: ContentUTF8, |
| 150 | }, |
| 151 | }, |
| 152 | } |
| 153 | |
| 154 | func floatPtr(v float64) *float64 { return &v } |
| 155 | func int64Ptr(v int64) *int64 { return &v } |
| 156 | |
| 157 | func TestMethodDTORoundTripsAreLossless(t *testing.T) { |
| 158 | for _, spec := range Registry() { |
| 159 | fixture, ok := methodFixtures[spec.Name] |
| 160 | if !ok { |
| 161 | t.Fatalf("no fixture for %s", spec.Name) |
| 162 | } |
| 163 | if reflect.TypeOf(fixture.params) != spec.ParamsType { |
| 164 | t.Fatalf("%s fixture params type = %v, want %v", spec.Name, reflect.TypeOf(fixture.params), spec.ParamsType) |
| 165 | } |
| 166 | t.Run(string(spec.Name)+"/params", func(t *testing.T) { |
| 167 | roundTripThroughDecoder(t, spec, fixture.params, true) |
| 168 | }) |
| 169 | if spec.Notification() { |
| 170 | continue |
| 171 | } |
| 172 | if reflect.TypeOf(fixture.result) != spec.ResultType { |
| 173 | t.Fatalf("%s fixture result type = %v, want %v", spec.Name, reflect.TypeOf(fixture.result), spec.ResultType) |
| 174 | } |
| 175 | t.Run(string(spec.Name)+"/result", func(t *testing.T) { |
| 176 | roundTripThroughDecoder(t, spec, fixture.result, false) |
| 177 | }) |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | func roundTripThroughDecoder(t *testing.T, spec MethodSpec, value any, params bool) { |
| 182 | t.Helper() |
| 183 | raw, err := json.Marshal(value) |
| 184 | if err != nil { |
| 185 | t.Fatalf("marshal: %v", err) |
| 186 | } |
| 187 | var decoded any |
| 188 | switch spec.Direction { |
| 189 | case DirectionHostToExtensionRequest: |
| 190 | if params { |
| 191 | decoded, err = DecodeHostRequestParams(spec.Name, raw) |
| 192 | } else { |
| 193 | decoded, err = DecodeHostRequestResult(spec.Name, raw) |
| 194 | } |
| 195 | case DirectionExtensionToHostRequest: |
| 196 | if params { |
| 197 | decoded, err = DecodeExtensionRequestParams(spec.Name, raw) |
| 198 | } else { |
| 199 | decoded, err = DecodeExtensionRequestResult(spec.Name, raw) |
| 200 | } |
| 201 | case DirectionHostToExtensionNotification: |
| 202 | decoded, err = DecodeHostNotificationParams(spec.Name, raw) |
| 203 | case DirectionExtensionToHostNotification: |
| 204 | decoded, err = DecodeExtensionNotificationParams(spec.Name, raw) |
| 205 | } |
| 206 | if err != nil { |
| 207 | t.Fatalf("strict decode of own fixture failed: %v\njson: %s", err, raw) |
| 208 | } |
| 209 | if !reflect.DeepEqual(decoded, value) { |
| 210 | t.Fatalf("round trip not lossless:\n got: %#v\nwant: %#v\njson: %s", decoded, value, raw) |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | // TestPayloadDTORoundTrips covers the structured UI payload documents, which |
| 215 | // are not method DTOs but ride inside UIPublishParams/UIRequestParams. |
| 216 | func TestPayloadDTORoundTrips(t *testing.T) { |
| 217 | payloads := []any{ |
| 218 | UIStatusPayload{Label: "l", Detail: "d", Severity: UISeverityWarn, Progress: floatPtr(0.5)}, |
| 219 | UICardPayload{ |
| 220 | Title: "t", Markdown: "**m**", Text: "x", |
| 221 | Fields: []UIKeyValue{{Key: "k", Value: "v"}}, |
| 222 | Progress: floatPtr(1), |
| 223 | Actions: []UIActionRef{{ActionID: "a", Label: "go"}}, |
| 224 | }, |
| 225 | UIFormPayload{ |
| 226 | Title: "t", Message: "m", |
| 227 | Fields: []UIFormField{{ |
| 228 | Key: "f", Label: "l", Kind: UIFieldMultiselect, |
| 229 | Options: []string{"a", "b"}, Default: "a", Required: true, |
| 230 | }}, |
| 231 | }, |
| 232 | UINotificationPayload{Title: "t", Body: "b", Severity: UISeverityError}, |
| 233 | } |
| 234 | for _, payload := range payloads { |
| 235 | raw, err := json.Marshal(payload) |
| 236 | if err != nil { |
| 237 | t.Fatalf("marshal %T: %v", payload, err) |
| 238 | } |
| 239 | decoded, err := decodeAndValidate(raw, reflect.TypeOf(payload)) |
| 240 | if err != nil { |
| 241 | t.Fatalf("strict decode %T: %v\njson: %s", payload, err, raw) |
| 242 | } |
| 243 | if !reflect.DeepEqual(decoded, payload) { |
| 244 | t.Fatalf("round trip not lossless for %T:\n got: %#v\nwant: %#v", payload, decoded, payload) |
| 245 | } |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | func TestStrictDecodersRejectBadShapes(t *testing.T) { |
| 250 | tests := []struct { |
| 251 | name string |
| 252 | decode func() (any, error) |
| 253 | }{ |
| 254 | {"unknown field", func() (any, error) { |
| 255 | return DecodeHostRequestParams(MethodExtensionShutdown, []byte(`{"timeoutMillis":1,"bogus":1}`)) |
| 256 | }}, |
| 257 | {"missing required", func() (any, error) { |
| 258 | return DecodeHostRequestParams(MethodExtensionShutdown, []byte(`{}`)) |
| 259 | }}, |
| 260 | {"null for non-nullable", func() (any, error) { |
| 261 | return DecodeHostRequestParams(MethodExtensionShutdown, []byte(`{"timeoutMillis":null}`)) |
| 262 | }}, |
| 263 | {"bad enum", func() (any, error) { |
| 264 | return DecodeHostRequestResult(MethodExtensionIntercept, []byte(`{"decision":"bogus"}`)) |
| 265 | }}, |
| 266 | {"empty required enum", func() (any, error) { |
| 267 | return DecodeHostRequestResult(MethodExtensionIntercept, []byte(`{"decision":""}`)) |
| 268 | }}, |
| 269 | {"min violation seq", func() (any, error) { |
| 270 | return DecodeExtensionNotificationParams(MethodExtensionProviderStreamChunk, |
| 271 | []byte(`{"streamId":"s","seq":0,"chunk":{"type":"done"}}`)) |
| 272 | }}, |
| 273 | {"min violation offset", func() (any, error) { |
| 274 | return DecodeExtensionRequestParams(MethodHostContentRead, []byte(`{"contentRef":"c","offset":-1}`)) |
| 275 | }}, |
| 276 | {"nonempty violation", func() (any, error) { |
| 277 | return DecodeExtensionRequestParams(MethodHostContentRead, []byte(`{"contentRef":" ","offset":0}`)) |
| 278 | }}, |
| 279 | {"sha256 violation", func() (any, error) { |
| 280 | return DecodeExtensionRequestResult(MethodHostContentRead, []byte( |
| 281 | `{"contentRef":"c","offset":0,"dataBase64":"","totalBytes":0,"sha256":"zz","encoding":"utf8"}`)) |
| 282 | }}, |
| 283 | {"error chunk without error", func() (any, error) { |
| 284 | return DecodeExtensionNotificationParams(MethodExtensionProviderStreamChunk, |
| 285 | []byte(`{"streamId":"s","seq":1,"chunk":{"type":"error"}}`)) |
| 286 | }}, |
| 287 | {"usage chunk without usage", func() (any, error) { |
| 288 | return DecodeExtensionNotificationParams(MethodExtensionProviderStreamChunk, |
| 289 | []byte(`{"streamId":"s","seq":1,"chunk":{"type":"usage"}}`)) |
| 290 | }}, |
| 291 | {"nil request arrays", func() (any, error) { |
| 292 | return DecodeHostRequestParams(MethodExtensionProviderStreamOpen, |
| 293 | []byte(`{"streamId":"s","providerRef":"p","request":{"maxTokens":0},"seqBase":0}`)) |
| 294 | }}, |
| 295 | {"tool parameters not object", func() (any, error) { |
| 296 | return DecodeHostRequestParams(MethodExtensionProviderStreamOpen, |
| 297 | []byte(`{"streamId":"s","providerRef":"p","request":{"messages":[],"tools":[{"name":"t","parameters":[1]}],"maxTokens":0},"seqBase":0}`)) |
| 298 | }}, |
| 299 | {"trailing json", func() (any, error) { |
| 300 | return DecodeHostRequestParams(MethodExtensionShutdown, []byte(`{"timeoutMillis":1} {}`)) |
| 301 | }}, |
| 302 | } |
| 303 | for _, tt := range tests { |
| 304 | t.Run(tt.name, func(t *testing.T) { |
| 305 | if _, err := tt.decode(); err == nil { |
| 306 | t.Fatal("strict decoder accepted an invalid payload") |
| 307 | } |
| 308 | }) |
| 309 | } |
| 310 | } |
| 311 | |
| 312 | func TestExternalizableFieldsAcceptNullPlaceholder(t *testing.T) { |
| 313 | // A null payload is the content-ref placeholder shape; only |
| 314 | // externalizable-tagged fields may carry it. |
| 315 | if _, err := DecodeHostNotificationParams(MethodExtensionEvent, []byte(`{"event":"session.start","payload":null}`)); err != nil { |
| 316 | t.Fatalf("externalizable payload null rejected: %v", err) |
| 317 | } |
| 318 | if _, err := DecodeExtensionRequestParams(MethodHostUIPublish, |
| 319 | []byte(`{"surfaceId":"s","sessionId":"s","generation":0,"kind":"card","payload":null}`)); err == nil { |
| 320 | t.Fatal("non-externalizable payload accepted null") |
| 321 | } |
| 322 | pointers := ExternalizablePointers(reflect.TypeOf(InterceptParams{})) |
| 323 | if !reflect.DeepEqual(pointers, []string{"/payload"}) { |
| 324 | t.Fatalf("ExternalizablePointers(InterceptParams) = %v", pointers) |
| 325 | } |
| 326 | pointers = ExternalizablePointers(reflect.TypeOf(ProviderRequest{})) |
| 327 | if !reflect.DeepEqual(pointers, []string{"/messages/*/content"}) { |
| 328 | t.Fatalf("ExternalizablePointers(ProviderRequest) = %v", pointers) |
| 329 | } |
| 330 | } |
| 331 |