cookie_agent_protocol/
extension.rs

1use std::{borrow::Cow, collections::BTreeMap};
2
3use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7pub const EXTENSION_PROTOCOL_VERSION: &str = "0.0.6";
8pub const PLUGIN_PRODUCER_REGISTER_METHOD: &str = "plugin/producer/register";
9pub const PLUGIN_PRODUCER_SEND_METHOD: &str = "plugin/producer/send";
10pub const PLUGIN_PRODUCER_UNREGISTER_METHOD: &str = "plugin/producer/unregister";
11pub const PLUGIN_PRODUCER_DISCARD_METHOD: &str = "plugin/producer/discard";
12pub const PLUGIN_RECOVERY_START_METHOD: &str = "plugin/recovery/start";
13pub const PLUGIN_RECOVERY_COMPLETE_METHOD: &str = "plugin/recovery/complete";
14pub const PLUGIN_INITIALIZE_METHOD: &str = "plugin/initialize";
15pub const PLUGIN_PING_METHOD: &str = "plugin/ping";
16pub const PLUGIN_SHUTDOWN_METHOD: &str = "plugin/shutdown";
17pub const PLUGIN_TOOLS_CALL_METHOD: &str = "plugin/tools/call";
18pub const PLUGIN_EVENT_METHOD: &str = "plugin/event";
19pub const PLUGIN_BUS_EVENT_METHOD: &str = "plugin/bus_event";
20pub const PLUGIN_EMIT_METHOD: &str = "plugin/emit";
21pub const PLUGIN_EMIT_RESULT_METHOD: &str = "plugin/emit_result";
22pub const PLUGIN_INTERCEPT_TOOL_BEFORE_CALL_METHOD: &str = "plugin/intercept/tool_before_call";
23pub const PLUGIN_INTERCEPT_TOOL_AFTER_RESULT_METHOD: &str = "plugin/intercept/tool_after_result";
24pub const PLUGIN_INTERCEPT_AGENT_BEFORE_START_METHOD: &str = "plugin/intercept/agent_before_start";
25pub const PLUGIN_INTERCEPT_SESSION_BEFORE_COMPACT_METHOD: &str =
26    "plugin/intercept/session_before_compact";
27pub const PLUGIN_INTERCEPT_USER_BEFORE_INPUT_METHOD: &str = "plugin/intercept/user_before_input";
28pub const PLUGIN_INTERCEPT_MODEL_BEFORE_REQUEST_METHOD: &str =
29    "plugin/intercept/model_before_request";
30pub const PLUGIN_INTERCEPT_PROVIDER_BEFORE_HEADERS_METHOD: &str =
31    "plugin/intercept/provider_before_headers";
32pub const PLUGIN_INTERCEPT_PROVIDER_BEFORE_REQUEST_METHOD: &str =
33    "plugin/intercept/provider_before_request";
34pub const PLUGIN_INTERCEPT_PROVIDER_AFTER_RESPONSE_METHOD: &str =
35    "plugin/intercept/provider_after_response";
36pub const PLUGIN_INTERCEPT_MESSAGE_END_METHOD: &str = "plugin/intercept/message_end";
37pub const PLUGIN_INTERCEPT_MODEL_BEFORE_SELECT_METHOD: &str =
38    "plugin/intercept/model_before_select";
39pub const PLUGIN_INTERCEPT_SESSION_BEFORE_FORK_METHOD: &str =
40    "plugin/intercept/session_before_fork";
41pub const PLUGIN_INTERCEPT_SESSION_BEFORE_REVERT_METHOD: &str =
42    "plugin/intercept/session_before_revert";
43
44#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
45pub struct ExtensionProtocolVersion(());
46
47impl ExtensionProtocolVersion {
48    #[must_use]
49    pub const fn current() -> Self {
50        Self(())
51    }
52}
53
54impl Serialize for ExtensionProtocolVersion {
55    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
56    where
57        S: serde::Serializer,
58    {
59        serializer.serialize_str(EXTENSION_PROTOCOL_VERSION)
60    }
61}
62
63impl<'de> Deserialize<'de> for ExtensionProtocolVersion {
64    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
65    where
66        D: serde::Deserializer<'de>,
67    {
68        let found = String::deserialize(deserializer)?;
69        if found == EXTENSION_PROTOCOL_VERSION {
70            Ok(Self::current())
71        } else {
72            Err(serde::de::Error::custom(format!(
73                "extension protocol version must be exactly {EXTENSION_PROTOCOL_VERSION}; found {found}"
74            )))
75        }
76    }
77}
78
79impl JsonSchema for ExtensionProtocolVersion {
80    fn inline_schema() -> bool {
81        true
82    }
83
84    fn schema_name() -> Cow<'static, str> {
85        Cow::Borrowed("ExtensionProtocolVersion")
86    }
87
88    fn json_schema(_: &mut SchemaGenerator) -> Schema {
89        json_schema!({"type":"string","const":"0.0.6"})
90    }
91}
92
93#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
94#[serde(deny_unknown_fields)]
95pub struct ExtensionEngineCapabilities {
96    pub producer_messaging: bool,
97    pub ping: bool,
98    pub shutdown: bool,
99    pub tools: bool,
100    pub event_streaming: bool,
101    pub event_publishing: bool,
102    pub interception: bool,
103}
104
105#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, JsonSchema, PartialEq, Serialize)]
106#[serde(rename_all = "snake_case")]
107pub enum ExtensionInterceptionHook {
108    ToolBeforeCall,
109    ToolAfterResult,
110    AgentBeforeStart,
111    SessionBeforeCompact,
112    UserBeforeInput,
113    ModelBeforeRequest,
114    ProviderBeforeHeaders,
115    ProviderBeforeRequest,
116    ProviderAfterResponse,
117    MessageEnd,
118    ModelBeforeSelect,
119    SessionBeforeFork,
120    SessionBeforeRevert,
121}
122
123#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
124#[serde(deny_unknown_fields)]
125pub struct ExtensionPluginCapabilities {
126    // Explicit opt-in, off by default in plugin configuration and SDK builders.
127    pub producer_messaging: bool,
128    pub tools: bool,
129    pub resources: bool,
130    pub subscribe_events: bool,
131    pub subscribe_bus: bool,
132    pub publish_bus: bool,
133    pub publish_session_events: bool,
134    pub intercept: Vec<ExtensionInterceptionHook>,
135}
136
137#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
138#[serde(deny_unknown_fields)]
139pub struct ExtensionInitializeParams {
140    pub protocol_version: ExtensionProtocolVersion,
141    pub engine_version: String,
142    pub capabilities: ExtensionEngineCapabilities,
143}
144
145#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
146#[serde(deny_unknown_fields)]
147pub struct ExtensionToolDeclaration {
148    pub name: String,
149    pub description: String,
150    pub parameters: Value,
151    pub permission_name: String,
152    #[serde(default)]
153    pub output: crate::ToolOutputDeclaration,
154    #[serde(deserialize_with = "crate::deserialize_required_option")]
155    #[schemars(with = "crate::NullableSchema<String>", required)]
156    pub primary_resource_param: Option<String>,
157}
158
159#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
160#[serde(deny_unknown_fields)]
161pub struct ExtensionInitializeResult {
162    pub protocol_version: ExtensionProtocolVersion,
163    pub name: String,
164    pub version: String,
165    pub capabilities: ExtensionPluginCapabilities,
166    pub tools: Vec<ExtensionToolDeclaration>,
167}
168
169#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
170#[serde(deny_unknown_fields)]
171pub struct ExtensionPingParams {}
172
173#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
174#[serde(deny_unknown_fields)]
175pub struct ExtensionPingResult {}
176
177#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
178#[serde(deny_unknown_fields)]
179pub struct ExtensionShutdownParams {}
180
181/// Plugin-to-engine request; the connection supplies the stable plugin owner.
182/// Registration is allowed during recovery and does not wait for goal readiness.
183#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
184#[serde(deny_unknown_fields)]
185pub struct ExtensionProducerRegisterParams {
186    pub session_id: crate::SessionId,
187}
188
189#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
190#[serde(deny_unknown_fields)]
191pub struct ExtensionProducerRegisterResult {
192    pub producer_id: crate::ProducerId,
193}
194
195/// Send requires a live registration owned by this connection and session.
196/// Identical (session, stable owner, key) retries return the original receipt;
197/// different description, body, or mode is rejected. A missing ACK is commit-uncertain.
198#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
199#[serde(deny_unknown_fields)]
200pub struct ExtensionProducerSendParams {
201    pub session_id: crate::SessionId,
202    pub producer_id: crate::ProducerId,
203    pub mode: crate::ProducerDeliveryMode,
204    pub idempotency_key: crate::ProducerIdempotencyKey,
205    #[serde(deserialize_with = "deserialize_producer_description")]
206    pub description: crate::SafeDisplayText,
207    pub body: String,
208}
209
210fn deserialize_producer_description<'de, D>(
211    deserializer: D,
212) -> Result<crate::SafeDisplayText, D::Error>
213where
214    D: serde::Deserializer<'de>,
215{
216    let description = crate::SafeDisplayText::deserialize(deserializer)?;
217    if description.as_str().trim().is_empty() {
218        return Err(serde::de::Error::custom("description must not be blank"));
219    }
220    Ok(description)
221}
222
223/// Returned only after durable ProducerMessageAccepted append. Retries return
224/// this same receipt, including after consumption, with no second acceptance.
225#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
226#[serde(deny_unknown_fields)]
227pub struct ExtensionProducerSendResult {
228    pub message_id: crate::ProducerMessageId,
229}
230
231#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
232#[serde(deny_unknown_fields)]
233pub struct ExtensionProducerUnregisterParams {
234    pub session_id: crate::SessionId,
235    pub producer_id: crate::ProducerId,
236}
237
238/// Unregister does not remove accepted messages; zero-send unregister is valid.
239#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
240#[serde(deny_unknown_fields)]
241pub struct ExtensionProducerUnregisterResult {}
242
243/// Discard an owned message only while it remains waiting. Authority is the live
244/// connection's stable plugin owner and session, not a registration.
245///
246/// A durable actor claim reserves the message and removes it from waiting before
247/// request preparation and hooks. Discard rejects until that claim is released,
248/// even before network activity; a claim does not prove provider receipt or execution.
249/// Release after failed preparation or cancellation may return an unconsumed
250/// message to waiting. Consumed messages cannot return or be discarded.
251/// Discard rejection makes no exactly-once claim about execution or external effects.
252#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
253#[serde(deny_unknown_fields)]
254pub struct ExtensionProducerDiscardParams {
255    pub session_id: crate::SessionId,
256    pub message_id: crate::ProducerMessageId,
257}
258
259/// Already-discarded owned messages succeed; consumed or currently claimed messages reject.
260#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
261#[serde(deny_unknown_fields)]
262pub struct ExtensionProducerDiscardResult {}
263
264/// Engine-to-plugin notification, never a deadlineful pending request. Restore
265/// from plugin-owned storage/services and report completion on the same connection.
266/// No response, recovery deadline, replay, or eager session adoption is implied.
267#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
268#[serde(deny_unknown_fields)]
269pub struct ExtensionRecoveryStartParams {}
270
271#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
272#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
273pub enum ExtensionRecoveryOutcome {
274    Ready,
275    Failed { message: crate::SafeErrorMessage },
276}
277
278/// Plugin-to-engine completion request. The engine owns starting/disabled states.
279/// The transport supplies authenticated plugin name + connection epoch out of band.
280#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
281#[serde(deny_unknown_fields)]
282pub struct ExtensionRecoveryCompleteParams {
283    pub outcome: ExtensionRecoveryOutcome,
284}
285
286#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
287#[serde(deny_unknown_fields)]
288pub struct ExtensionRecoveryCompleteResult {}
289
290#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
291#[serde(deny_unknown_fields)]
292pub struct ExtensionToolCallParams {
293    pub tool: String,
294    pub session_id: crate::SessionId,
295    pub context_id: String,
296    pub invocation_id: crate::ToolCallId,
297    pub arguments: Value,
298    #[serde(deserialize_with = "crate::deserialize_required_option")]
299    #[schemars(with = "crate::NullableSchema<String>", required)]
300    pub resource: Option<String>,
301    #[serde(default, skip_serializing_if = "Option::is_none")]
302    pub cancellation_token: Option<String>,
303}
304
305#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
306#[serde(deny_unknown_fields)]
307pub struct ExtensionToolCallResult {
308    pub output: crate::ToolCompletionOutput,
309    pub display: String,
310    pub is_error: bool,
311}
312
313#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
314#[serde(deny_unknown_fields)]
315pub struct ExtensionEventParams {
316    pub session_id: crate::SessionId,
317    pub context_id: String,
318    #[schemars(range(min = 1))]
319    pub seq: u64,
320    pub event: crate::EventPayload,
321    pub timestamp: jiff::Timestamp,
322}
323
324#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
325#[serde(deny_unknown_fields)]
326pub struct ExtensionBusEventParams {
327    pub session_id: crate::SessionId,
328    #[serde(default, skip_serializing_if = "Option::is_none")]
329    pub context_id: Option<String>,
330    pub plugin: String,
331    pub name: String,
332    pub payload: Value,
333}
334
335#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
336#[serde(deny_unknown_fields)]
337/// Non-model publication only. Model-bound content must use producer.send with
338/// an explicitly registered producer; emit must never implicitly register one.
339pub struct ExtensionEmitParams {
340    pub session_id: crate::SessionId,
341    pub context_id: String,
342    #[schemars(length(min = 1, max = 128))]
343    pub name: String,
344    pub payload: Value,
345}
346
347#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
348#[serde(rename_all = "snake_case")]
349pub enum ExtensionEmitStatus {
350    Published,
351    Dropped,
352    Rejected,
353}
354
355#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
356#[serde(deny_unknown_fields)]
357pub struct ExtensionEmitResultParams {
358    pub name: String,
359    pub bus: ExtensionEmitStatus,
360    pub durable: ExtensionEmitStatus,
361    #[serde(default, skip_serializing_if = "Option::is_none")]
362    pub reason: Option<String>,
363}
364
365#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
366#[serde(deny_unknown_fields)]
367pub struct ExtensionToolBeforeCallParams {
368    pub session_id: crate::SessionId,
369    pub context_id: String,
370    pub tool: String,
371    pub arguments: Value,
372    pub permission_name: String,
373    #[serde(deserialize_with = "crate::deserialize_required_option")]
374    #[schemars(with = "crate::NullableSchema<String>", required)]
375    pub resource: Option<String>,
376}
377
378#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
379#[serde(rename_all = "snake_case")]
380pub enum ExtensionToolBeforeCallAction {
381    Allow,
382    Block,
383}
384
385#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
386#[serde(deny_unknown_fields)]
387pub struct ExtensionToolBeforeCallResult {
388    pub action: ExtensionToolBeforeCallAction,
389    #[serde(default, skip_serializing_if = "Option::is_none")]
390    pub reason: Option<String>,
391    #[serde(default, skip_serializing_if = "Option::is_none")]
392    pub modified_arguments: Option<Value>,
393    #[serde(default, skip_serializing_if = "Option::is_none")]
394    pub message_to_model: Option<String>,
395}
396
397#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
398#[serde(deny_unknown_fields)]
399pub struct ExtensionToolAfterResultParams {
400    pub session_id: crate::SessionId,
401    pub context_id: String,
402    pub tool: String,
403    pub arguments: Value,
404    pub result_content: String,
405    pub is_error: bool,
406}
407
408#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
409#[serde(rename_all = "snake_case")]
410pub enum ExtensionToolAfterResultAction {
411    Keep,
412    Replace,
413}
414
415#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
416#[serde(deny_unknown_fields)]
417pub struct ExtensionToolAfterResultResult {
418    pub action: ExtensionToolAfterResultAction,
419    #[serde(default, skip_serializing_if = "Option::is_none")]
420    pub replacement_content: Option<String>,
421    #[serde(default, skip_serializing_if = "Option::is_none")]
422    pub note: Option<String>,
423}
424
425#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
426#[serde(deny_unknown_fields)]
427pub struct ExtensionAgentBeforeStartParams {
428    pub session_id: crate::SessionId,
429    pub context_id: String,
430    pub agent_path: String,
431    pub prompt_context: Value,
432}
433
434#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
435#[serde(deny_unknown_fields)]
436pub struct ExtensionAgentBeforeStartResult {
437    #[serde(default, skip_serializing_if = "Option::is_none")]
438    pub addendum: Option<String>,
439    #[serde(default, skip_serializing_if = "Option::is_none")]
440    pub append_to_system_prompt: Option<String>,
441    #[serde(default, skip_serializing_if = "Option::is_none")]
442    pub replace_system_prompt: Option<String>,
443    #[serde(default, skip_serializing_if = "Option::is_none")]
444    pub inject_message: Option<ExtensionInjectedMessage>,
445}
446
447#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
448#[serde(rename_all = "snake_case")]
449pub enum ExtensionMessageRole {
450    System,
451    User,
452    Assistant,
453    Tool,
454}
455
456#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
457#[serde(deny_unknown_fields)]
458pub struct ExtensionInjectedMessage {
459    pub role: ExtensionMessageRole,
460    pub content: String,
461}
462
463#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
464#[serde(deny_unknown_fields)]
465pub struct ExtensionSessionBeforeCompactParams {
466    pub session_id: crate::SessionId,
467    pub context_id: String,
468    pub checkpoint_id: String,
469    pub additions: Vec<String>,
470    #[serde(default, skip_serializing_if = "Option::is_none")]
471    pub instructions: Option<String>,
472}
473
474#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
475#[serde(deny_unknown_fields)]
476pub struct ExtensionSessionBeforeCompactResult {
477    #[serde(default, skip_serializing_if = "Option::is_none")]
478    pub addendum: Option<String>,
479    #[serde(default)]
480    pub cancel: bool,
481    #[serde(default, skip_serializing_if = "Option::is_none")]
482    pub reason: Option<String>,
483    #[serde(default, skip_serializing_if = "Option::is_none")]
484    pub instructions_override: Option<String>,
485}
486
487#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
488#[serde(deny_unknown_fields)]
489pub struct ExtensionUserBeforeInputParams {
490    pub session_id: crate::SessionId,
491    pub context_id: String,
492    pub text: String,
493}
494
495#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
496#[serde(rename_all = "snake_case")]
497pub enum ExtensionUserBeforeInputAction {
498    Allow,
499    Transform,
500    Handled,
501}
502
503#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
504#[serde(deny_unknown_fields)]
505pub struct ExtensionUserBeforeInputResult {
506    pub action: ExtensionUserBeforeInputAction,
507    #[serde(default, skip_serializing_if = "Option::is_none")]
508    pub new_text: Option<String>,
509    #[serde(default, skip_serializing_if = "Option::is_none")]
510    pub reason: Option<String>,
511}
512
513#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
514#[serde(deny_unknown_fields)]
515pub struct ExtensionModelMessage {
516    pub role: ExtensionMessageRole,
517    pub content: Value,
518}
519
520#[derive(Clone, Debug, Default, Deserialize, JsonSchema, PartialEq, Serialize)]
521#[serde(deny_unknown_fields)]
522pub struct ExtensionModelParamsAdjustments {
523    #[serde(default, skip_serializing_if = "Option::is_none")]
524    pub max_tokens: Option<u64>,
525    #[serde(default, skip_serializing_if = "Option::is_none")]
526    pub temperature: Option<f64>,
527    #[serde(default, skip_serializing_if = "Option::is_none")]
528    pub top_p: Option<f64>,
529}
530
531#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
532#[serde(deny_unknown_fields)]
533pub struct ExtensionModelBeforeRequestParams {
534    pub session_id: crate::SessionId,
535    pub context_id: String,
536    pub attempt_id: crate::AttemptId,
537    pub messages: Vec<ExtensionModelMessage>,
538    pub model: crate::ResolvedModelRef,
539    pub params: ExtensionModelParamsAdjustments,
540}
541
542#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
543#[serde(rename_all = "snake_case")]
544pub enum ExtensionModelBeforeRequestAction {
545    Keep,
546    Replace,
547}
548
549#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
550#[serde(deny_unknown_fields)]
551pub struct ExtensionModelBeforeRequestResult {
552    pub action: ExtensionModelBeforeRequestAction,
553    #[serde(default, skip_serializing_if = "Option::is_none")]
554    pub messages: Option<Vec<ExtensionModelMessage>>,
555    #[serde(default, skip_serializing_if = "Option::is_none")]
556    pub params_adjustments: Option<ExtensionModelParamsAdjustments>,
557}
558
559#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
560#[serde(deny_unknown_fields)]
561pub struct ExtensionProviderBeforeHeadersParams {
562    pub session_id: crate::SessionId,
563    pub context_id: String,
564    pub attempt_id: crate::AttemptId,
565    pub headers: BTreeMap<String, String>,
566}
567
568#[derive(Clone, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
569#[serde(deny_unknown_fields)]
570pub struct ExtensionProviderBeforeHeadersResult {
571    #[serde(default)]
572    pub set: BTreeMap<String, String>,
573    #[serde(default)]
574    pub delete: Vec<String>,
575}
576
577#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
578#[serde(deny_unknown_fields)]
579pub struct ExtensionProviderBeforeRequestParams {
580    pub session_id: crate::SessionId,
581    pub context_id: String,
582    pub attempt_id: crate::AttemptId,
583    pub payload: Value,
584}
585
586#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
587#[serde(rename_all = "snake_case")]
588pub enum ExtensionProviderBeforeRequestAction {
589    Keep,
590    Replace,
591}
592
593#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
594#[serde(deny_unknown_fields)]
595pub struct ExtensionProviderBeforeRequestResult {
596    pub action: ExtensionProviderBeforeRequestAction,
597    #[serde(default, skip_serializing_if = "Option::is_none")]
598    pub payload: Option<Value>,
599}
600
601#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
602#[serde(deny_unknown_fields)]
603pub struct ExtensionProviderAfterResponseParams {
604    pub session_id: crate::SessionId,
605    pub context_id: String,
606    pub attempt_id: crate::AttemptId,
607    pub status: u16,
608    pub headers: BTreeMap<String, String>,
609}
610
611#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
612#[serde(deny_unknown_fields)]
613pub struct ExtensionProviderAfterResponseResult {}
614
615#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
616#[serde(deny_unknown_fields)]
617pub struct ExtensionMessageEndParams {
618    pub session_id: crate::SessionId,
619    pub context_id: String,
620    pub attempt_id: crate::AttemptId,
621    pub role: ExtensionMessageRole,
622    pub content: Vec<crate::PersistedAssistantPart>,
623}
624
625#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
626#[serde(rename_all = "snake_case")]
627pub enum ExtensionMessageEndAction {
628    Keep,
629    Replace,
630}
631
632#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
633#[serde(deny_unknown_fields)]
634pub struct ExtensionMessageEndResult {
635    pub action: ExtensionMessageEndAction,
636    #[serde(default, skip_serializing_if = "Option::is_none")]
637    pub content: Option<Vec<crate::PersistedAssistantPart>>,
638}
639
640#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
641#[serde(rename_all = "snake_case")]
642pub enum ExtensionModelSelectSource {
643    User,
644    Config,
645    FallbackRestore,
646}
647
648#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
649#[serde(deny_unknown_fields)]
650pub struct ExtensionModelBeforeSelectParams {
651    pub session_id: crate::SessionId,
652    pub context_id: String,
653    #[serde(default, skip_serializing_if = "Option::is_none")]
654    pub from: Option<cookie_agent_identity::ModelSelection>,
655    pub to: cookie_agent_identity::ModelSelection,
656    pub source: ExtensionModelSelectSource,
657}
658
659#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
660#[serde(rename_all = "snake_case")]
661pub enum ExtensionAllowBlockAction {
662    Allow,
663    Block,
664}
665
666#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
667#[serde(deny_unknown_fields)]
668pub struct ExtensionAllowBlockResult {
669    pub action: ExtensionAllowBlockAction,
670    #[serde(default, skip_serializing_if = "Option::is_none")]
671    pub reason: Option<String>,
672}
673
674#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
675#[serde(deny_unknown_fields)]
676pub struct ExtensionSessionBeforeForkParams {
677    pub session_id: crate::SessionId,
678    pub context_id: String,
679    pub through_seq: u64,
680}
681
682pub type ExtensionSessionBeforeForkResult = ExtensionAllowBlockResult;
683
684#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
685#[serde(deny_unknown_fields)]
686pub struct ExtensionSessionBeforeRevertParams {
687    pub session_id: crate::SessionId,
688    pub context_id: String,
689    pub through_seq: u64,
690    #[serde(default, skip_serializing_if = "Option::is_none")]
691    pub instructions: Option<String>,
692}
693
694#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
695#[serde(rename_all = "snake_case")]
696pub enum ExtensionSessionBeforeRevertAction {
697    Allow,
698    Block,
699    Override,
700}
701
702#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
703#[serde(deny_unknown_fields)]
704pub struct ExtensionSessionBeforeRevertResult {
705    pub action: ExtensionSessionBeforeRevertAction,
706    #[serde(default, skip_serializing_if = "Option::is_none")]
707    pub reason: Option<String>,
708    #[serde(default, skip_serializing_if = "Option::is_none")]
709    pub instructions_override: Option<String>,
710}
711
712#[must_use]
713pub fn extension_initialize_request(engine_version: impl Into<String>) -> crate::Request {
714    let params = ExtensionInitializeParams {
715        protocol_version: ExtensionProtocolVersion::current(),
716        engine_version: engine_version.into(),
717        capabilities: ExtensionEngineCapabilities {
718            producer_messaging: true,
719            ping: true,
720            shutdown: true,
721            tools: true,
722            event_streaming: true,
723            event_publishing: true,
724            interception: true,
725        },
726    };
727    crate::Request::new(
728        crate::JsonRpcId::Number(1),
729        PLUGIN_INITIALIZE_METHOD,
730        Some(serde_json::to_value(params).expect("extension initialize params serialize")),
731    )
732}
733
734#[must_use]
735pub fn extension_recovery_start_notification() -> crate::Notification {
736    crate::Notification::new(
737        PLUGIN_RECOVERY_START_METHOD,
738        Some(
739            serde_json::to_value(ExtensionRecoveryStartParams {})
740                .expect("recovery start params serialize"),
741        ),
742    )
743}
744
745#[must_use]
746pub fn extension_shutdown_notification() -> crate::Notification {
747    crate::Notification::new(PLUGIN_SHUTDOWN_METHOD, Some(serde_json::json!({})))
748}
749
750#[cfg(test)]
751mod tests {
752    use super::*;
753
754    #[test]
755    fn producer_discard_is_strict_and_registration_independent() {
756        let params = ExtensionProducerDiscardParams {
757            session_id: crate::SessionId::new_v7(),
758            message_id: crate::ProducerMessageId::new_v7(),
759        };
760        let wire = serde_json::to_value(params).unwrap();
761        assert_eq!(
762            serde_json::from_value::<ExtensionProducerDiscardParams>(wire.clone()).unwrap(),
763            params
764        );
765        assert_eq!(
766            wire,
767            serde_json::json!({"session_id": params.session_id, "message_id": params.message_id})
768        );
769        for field in ["producer_id", "plugin", "connection_epoch"] {
770            let mut forged = wire.clone();
771            forged[field] = serde_json::json!("foreign");
772            assert!(serde_json::from_value::<ExtensionProducerDiscardParams>(forged).is_err());
773        }
774        for field in ["session_id", "message_id"] {
775            let mut missing = wire.clone();
776            missing.as_object_mut().unwrap().remove(field);
777            assert!(serde_json::from_value::<ExtensionProducerDiscardParams>(missing).is_err());
778            let mut invalid = wire.clone();
779            invalid[field] = serde_json::json!("not-a-uuid");
780            assert!(serde_json::from_value::<ExtensionProducerDiscardParams>(invalid).is_err());
781        }
782        assert_eq!(PLUGIN_PRODUCER_DISCARD_METHOD, "plugin/producer/discard");
783        assert_eq!(
784            serde_json::to_value(ExtensionProducerDiscardResult {}).unwrap(),
785            serde_json::json!({})
786        );
787        assert!(
788            serde_json::from_value::<ExtensionProducerDiscardResult>(serde_json::json!({})).is_ok()
789        );
790        assert!(
791            serde_json::from_value::<ExtensionProducerDiscardResult>(
792                serde_json::json!({"message_id": params.message_id})
793            )
794            .is_err()
795        );
796    }
797
798    #[test]
799    fn protocol_version_rejection_reports_found_value() {
800        let error = serde_json::from_str::<ExtensionProtocolVersion>("\"0.0.1\"")
801            .expect_err("version mismatch");
802        assert!(error.to_string().contains("0.0.1"));
803    }
804}
cookie_agent_protocol/extension.rs

cookie_agent_protocol/
extension.rs
