1use serde::de::DeserializeOwned;
2use serde::{Deserialize, Serialize};
3use serde_json::{Value, json};
4
5use crate::store::WorkGraphEventFilter;
6use crate::types::{
7 AddEvidenceRequest, ClaimWorkItemRequest, CloseWorkItemRequest, LinkWorkItemsRequest,
8 ReadyWorkFilter, ReleaseWorkItemRequest, UpdateWorkItemRequest, WorkGraphSnapshotFilter,
9 WorkItemFilter, WorkItemId, WorkNamespace,
10};
11use crate::{CreateWorkItemRequest, WorkGraphError, WorkGraphService};
12
13pub const INVALID_ARGUMENTS: &str = "invalid_arguments";
14pub const NOT_FOUND: &str = "not_found";
15pub const CAPABILITY_UNAVAILABLE: &str = "capability_unavailable";
16
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
18pub struct WorkGraphToolError {
19 pub code: String,
20 pub message: String,
21}
22
23impl WorkGraphToolError {
24 fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
25 Self {
26 code: code.into(),
27 message: message.into(),
28 }
29 }
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33enum WorkGraphToolContract {
34 Create,
35 Get,
36 List,
37 Ready,
38 Snapshot,
39 Events,
40 Claim,
41 Release,
42 Update,
43 Block,
44 Close,
45 Link,
46 AddEvidence,
47}
48
49impl WorkGraphToolContract {
50 const ALL: &'static [Self] = &[
51 Self::Create,
52 Self::Get,
53 Self::List,
54 Self::Ready,
55 Self::Snapshot,
56 Self::Events,
57 Self::Claim,
58 Self::Release,
59 Self::Update,
60 Self::Block,
61 Self::Close,
62 Self::Link,
63 Self::AddEvidence,
64 ];
65
66 const fn name(self) -> &'static str {
67 match self {
68 Self::Create => "workgraph_create",
69 Self::Get => "workgraph_get",
70 Self::List => "workgraph_list",
71 Self::Ready => "workgraph_ready",
72 Self::Snapshot => "workgraph_snapshot",
73 Self::Events => "workgraph_events",
74 Self::Claim => "workgraph_claim",
75 Self::Release => "workgraph_release",
76 Self::Update => "workgraph_update",
77 Self::Block => "workgraph_block",
78 Self::Close => "workgraph_close",
79 Self::Link => "workgraph_link",
80 Self::AddEvidence => "workgraph_add_evidence",
81 }
82 }
83
84 const fn description(self) -> &'static str {
85 match self {
86 Self::Create => "Create a durable WorkGraph item.",
87 Self::Get => "Read one WorkGraph item.",
88 Self::List => "List WorkGraph items.",
89 Self::Ready => "List ready, claimable WorkGraph items.",
90 Self::Snapshot => "Read a WorkGraph observability snapshot.",
91 Self::Events => "Read WorkGraph event history.",
92 Self::Claim => "Claim a ready WorkGraph item with CAS revision checking.",
93 Self::Release => "Release a claimed WorkGraph item.",
94 Self::Update => "Update non-terminal WorkGraph item fields.",
95 Self::Block => "Mark a WorkGraph item blocked.",
96 Self::Close => "Close a WorkGraph item with a terminal status.",
97 Self::Link => "Create a dependency or relationship edge.",
98 Self::AddEvidence => "Attach a typed evidence reference to a WorkGraph item.",
99 }
100 }
101
102 fn schema(self) -> Value {
103 match self {
104 Self::Create => create_schema(),
105 Self::Get => id_schema(false),
106 Self::List => list_schema(),
107 Self::Ready => ready_schema(),
108 Self::Snapshot => snapshot_schema(),
109 Self::Events => events_schema(),
110 Self::Claim => claim_schema(),
111 Self::Release | Self::Block => revision_id_schema(),
112 Self::Update => update_schema(),
113 Self::Close => close_schema(),
114 Self::Link => link_schema(),
115 Self::AddEvidence => evidence_schema(),
116 }
117 }
118
119 fn parse(name: &str) -> Result<Self, WorkGraphToolError> {
120 Self::ALL
121 .iter()
122 .copied()
123 .find(|contract| contract.name() == name)
124 .ok_or_else(|| {
125 WorkGraphToolError::new(NOT_FOUND, format!("unknown WorkGraph tool '{name}'"))
126 })
127 }
128}
129
130pub fn workgraph_tools_list() -> Vec<Value> {
131 WorkGraphToolContract::ALL
132 .iter()
133 .map(|contract| tool(contract.name(), contract.description(), contract.schema()))
134 .collect()
135}
136
137pub async fn handle_workgraph_tools_call(
138 service: &WorkGraphService,
139 name: &str,
140 arguments: &Value,
141) -> Result<Value, WorkGraphToolError> {
142 match WorkGraphToolContract::parse(name)? {
143 WorkGraphToolContract::Create => {
144 let request: CreateWorkItemRequest = parse(arguments)?;
145 service
146 .create(request)
147 .await
148 .map(|item| json!({ "item": item }))
149 .map_err(map_error)
150 }
151 WorkGraphToolContract::Get => {
152 let request: IdParams = parse(arguments)?;
153 service
154 .get(request.realm_id, request.namespace, request.id)
155 .await
156 .map(|item| json!({ "item": item }))
157 .map_err(map_error)
158 }
159 WorkGraphToolContract::List => {
160 let filter: WorkItemFilter = parse(arguments)?;
161 service
162 .list(filter)
163 .await
164 .map(|items| json!({ "items": items }))
165 .map_err(map_error)
166 }
167 WorkGraphToolContract::Ready => {
168 let filter: ReadyWorkFilter = parse(arguments)?;
169 service
170 .ready(filter)
171 .await
172 .map(|items| json!({ "items": items }))
173 .map_err(map_error)
174 }
175 WorkGraphToolContract::Snapshot => {
176 let filter: WorkGraphSnapshotFilter = parse(arguments)?;
177 service
178 .snapshot(filter)
179 .await
180 .map(|snapshot| json!({ "snapshot": snapshot }))
181 .map_err(map_error)
182 }
183 WorkGraphToolContract::Claim => {
184 let request: ClaimWorkItemRequest = parse(arguments)?;
185 service
186 .claim(request)
187 .await
188 .map(|item| json!({ "item": item }))
189 .map_err(map_error)
190 }
191 WorkGraphToolContract::Release => {
192 let request: ReleaseWorkItemRequest = parse(arguments)?;
193 service
194 .release(request)
195 .await
196 .map(|item| json!({ "item": item }))
197 .map_err(map_error)
198 }
199 WorkGraphToolContract::Update => {
200 let request: UpdateWorkItemRequest = parse(arguments)?;
201 service
202 .update(request)
203 .await
204 .map(|item| json!({ "item": item }))
205 .map_err(map_error)
206 }
207 WorkGraphToolContract::Block => {
208 let request: RevisionIdParams = parse(arguments)?;
209 service
210 .block(
211 request.realm_id,
212 request.namespace,
213 request.id,
214 request.expected_revision,
215 )
216 .await
217 .map(|item| json!({ "item": item }))
218 .map_err(map_error)
219 }
220 WorkGraphToolContract::Close => {
221 let request: CloseWorkItemRequest = parse(arguments)?;
222 service
223 .close(request)
224 .await
225 .map(|item| json!({ "item": item }))
226 .map_err(map_error)
227 }
228 WorkGraphToolContract::Link => {
229 let request: LinkWorkItemsRequest = parse(arguments)?;
230 service
231 .link(request)
232 .await
233 .map(|edge| json!({ "edge": edge }))
234 .map_err(map_error)
235 }
236 WorkGraphToolContract::AddEvidence => {
237 let request: AddEvidenceRequest = parse(arguments)?;
238 service
239 .add_evidence(request)
240 .await
241 .map(|item| json!({ "item": item }))
242 .map_err(map_error)
243 }
244 WorkGraphToolContract::Events => {
245 let filter: WorkGraphEventFilterParams = parse(arguments)?;
246 service
247 .events(filter.into())
248 .await
249 .map(|events| json!({ "events": events }))
250 .map_err(map_error)
251 }
252 }
253}
254
255#[derive(Debug, Deserialize)]
256struct IdParams {
257 id: WorkItemId,
258 #[serde(default)]
259 realm_id: Option<String>,
260 #[serde(default)]
261 namespace: Option<WorkNamespace>,
262}
263
264#[derive(Debug, Deserialize)]
265struct RevisionIdParams {
266 id: WorkItemId,
267 expected_revision: u64,
268 #[serde(default)]
269 realm_id: Option<String>,
270 #[serde(default)]
271 namespace: Option<WorkNamespace>,
272}
273
274#[derive(Debug, Deserialize)]
275struct WorkGraphEventFilterParams {
276 #[serde(default)]
277 realm_id: Option<String>,
278 #[serde(default)]
279 namespace: Option<WorkNamespace>,
280 #[serde(default)]
281 all_namespaces: bool,
282 #[serde(default)]
283 after_seq: Option<i64>,
284 #[serde(default)]
285 limit: Option<usize>,
286}
287
288impl From<WorkGraphEventFilterParams> for WorkGraphEventFilter {
289 fn from(value: WorkGraphEventFilterParams) -> Self {
290 Self {
291 realm_id: value.realm_id,
292 namespace: value.namespace,
293 all_namespaces: value.all_namespaces,
294 after_seq: value.after_seq,
295 limit: value.limit,
296 }
297 }
298}
299
300fn parse<T: DeserializeOwned>(arguments: &Value) -> Result<T, WorkGraphToolError> {
301 serde_json::from_value(arguments.clone()).map_err(|err| {
302 WorkGraphToolError::new(
303 INVALID_ARGUMENTS,
304 format!("invalid WorkGraph arguments: {err}"),
305 )
306 })
307}
308
309fn map_error(error: WorkGraphError) -> WorkGraphToolError {
310 let code = match error {
311 WorkGraphError::NotFound { .. } => NOT_FOUND,
312 WorkGraphError::StaleRevision { .. } | WorkGraphError::Conflict(_) => "conflict",
313 WorkGraphError::InvalidTransition(_) => "invalid_transition",
314 WorkGraphError::InvalidInput(_) => INVALID_ARGUMENTS,
315 WorkGraphError::UnsupportedBackend(_) => CAPABILITY_UNAVAILABLE,
316 WorkGraphError::Store(_) => "store_error",
317 };
318 WorkGraphToolError::new(code, error.to_string())
319}
320
321fn tool(name: &str, description: &str, schema: Value) -> Value {
322 json!({
323 "name": name,
324 "description": description,
325 "inputSchema": schema,
326 })
327}
328
329fn base_properties() -> serde_json::Map<String, Value> {
330 serde_json::Map::from_iter([
331 ("realm_id".to_string(), json!({ "type": "string" })),
332 ("namespace".to_string(), json!({ "type": "string" })),
333 ])
334}
335
336fn external_ref_schema() -> Value {
337 json!({
338 "type": "object",
339 "properties": {
340 "kind": { "type": "string" },
341 "id": { "type": "string" },
342 "url": { "type": "string" }
343 },
344 "required": ["kind", "id"],
345 "additionalProperties": false
346 })
347}
348
349fn evidence_ref_schema() -> Value {
350 json!({
351 "type": "object",
352 "properties": {
353 "kind": { "type": "string" },
354 "id": { "type": "string" },
355 "label": { "type": "string" },
356 "summary": { "type": "string" }
357 },
358 "required": ["kind", "id"],
359 "additionalProperties": false
360 })
361}
362
363fn object(properties: serde_json::Map<String, Value>, required: &[&str]) -> Value {
364 json!({
365 "type": "object",
366 "properties": properties,
367 "required": required,
368 "additionalProperties": false,
369 })
370}
371
372fn id_schema(include_revision: bool) -> Value {
373 let mut properties = base_properties();
374 properties.insert("id".to_string(), json!({ "type": "string" }));
375 if include_revision {
376 properties.insert(
377 "expected_revision".to_string(),
378 json!({ "type": "integer", "minimum": 0 }),
379 );
380 object(properties, &["id", "expected_revision"])
381 } else {
382 object(properties, &["id"])
383 }
384}
385
386fn revision_id_schema() -> Value {
387 id_schema(true)
388}
389
390fn create_schema() -> Value {
391 let mut properties = base_properties();
392 properties.extend([
393 ("title".to_string(), json!({ "type": "string" })),
394 ("description".to_string(), json!({ "type": "string" })),
395 (
396 "priority".to_string(),
397 json!({ "type": "string", "enum": ["low", "medium", "high"] }),
398 ),
399 (
400 "labels".to_string(),
401 json!({ "type": "array", "items": { "type": "string" } }),
402 ),
403 (
404 "due_at".to_string(),
405 json!({ "type": "string", "format": "date-time" }),
406 ),
407 (
408 "not_before".to_string(),
409 json!({ "type": "string", "format": "date-time" }),
410 ),
411 (
412 "snoozed_until".to_string(),
413 json!({ "type": "string", "format": "date-time" }),
414 ),
415 (
416 "status".to_string(),
417 json!({ "type": "string", "enum": ["open", "blocked"] }),
418 ),
419 (
420 "external_refs".to_string(),
421 json!({ "type": "array", "items": external_ref_schema() }),
422 ),
423 (
424 "evidence_refs".to_string(),
425 json!({ "type": "array", "items": evidence_ref_schema() }),
426 ),
427 ]);
428 object(properties, &["title"])
429}
430
431fn list_schema() -> Value {
432 let mut properties = base_properties();
433 properties.extend([
434 ("all_namespaces".to_string(), json!({ "type": "boolean" })),
435 (
436 "statuses".to_string(),
437 json!({ "type": "array", "items": { "type": "string" } }),
438 ),
439 (
440 "labels".to_string(),
441 json!({ "type": "array", "items": { "type": "string" } }),
442 ),
443 ("include_terminal".to_string(), json!({ "type": "boolean" })),
444 (
445 "limit".to_string(),
446 json!({ "type": "integer", "minimum": 1 }),
447 ),
448 ]);
449 object(properties, &[])
450}
451
452fn ready_schema() -> Value {
453 let mut properties = base_properties();
454 properties.extend([
455 (
456 "labels".to_string(),
457 json!({ "type": "array", "items": { "type": "string" } }),
458 ),
459 (
460 "limit".to_string(),
461 json!({ "type": "integer", "minimum": 1 }),
462 ),
463 ]);
464 object(properties, &[])
465}
466
467fn snapshot_schema() -> Value {
468 list_schema()
469}
470
471fn events_schema() -> Value {
472 let mut properties = base_properties();
473 properties.extend([
474 ("all_namespaces".to_string(), json!({ "type": "boolean" })),
475 (
476 "after_seq".to_string(),
477 json!({ "type": "integer", "minimum": 0 }),
478 ),
479 (
480 "limit".to_string(),
481 json!({ "type": "integer", "minimum": 1 }),
482 ),
483 ]);
484 object(properties, &[])
485}
486
487fn claim_schema() -> Value {
488 let mut properties = base_properties();
489 properties.extend([
490 ("id".to_string(), json!({ "type": "string" })),
491 (
492 "expected_revision".to_string(),
493 json!({ "type": "integer", "minimum": 0 }),
494 ),
495 (
496 "owner".to_string(),
497 json!({
498 "type": "object",
499 "properties": {
500 "key": {
501 "type": "object",
502 "properties": {
503 "kind": {
504 "type": "string",
505 "enum": ["principal", "agent", "session", "mob", "label"]
506 },
507 "id": { "type": "string" }
508 },
509 "required": ["kind", "id"],
510 "additionalProperties": false
511 },
512 "display_name": { "type": "string" }
513 },
514 "required": ["key"],
515 "additionalProperties": false
516 }),
517 ),
518 (
519 "lease_seconds".to_string(),
520 json!({ "type": "integer", "minimum": 1 }),
521 ),
522 (
523 "lease_expires_at".to_string(),
524 json!({ "type": "string", "format": "date-time" }),
525 ),
526 ]);
527 object(properties, &["id", "expected_revision", "owner"])
528}
529
530fn update_schema() -> Value {
531 let mut properties = base_properties();
532 properties.extend([
533 ("id".to_string(), json!({ "type": "string" })),
534 (
535 "expected_revision".to_string(),
536 json!({ "type": "integer", "minimum": 0 }),
537 ),
538 ("title".to_string(), json!({ "type": "string" })),
539 ("description".to_string(), json!({ "type": "string" })),
540 (
541 "priority".to_string(),
542 json!({ "type": "string", "enum": ["low", "medium", "high"] }),
543 ),
544 (
545 "labels".to_string(),
546 json!({ "type": "array", "items": { "type": "string" } }),
547 ),
548 (
549 "due_at".to_string(),
550 json!({ "type": "string", "format": "date-time" }),
551 ),
552 (
553 "not_before".to_string(),
554 json!({ "type": "string", "format": "date-time" }),
555 ),
556 (
557 "snoozed_until".to_string(),
558 json!({ "type": "string", "format": "date-time" }),
559 ),
560 (
561 "external_refs".to_string(),
562 json!({ "type": "array", "items": external_ref_schema() }),
563 ),
564 ]);
565 object(properties, &["id", "expected_revision"])
566}
567
568fn close_schema() -> Value {
569 let mut properties = base_properties();
570 properties.extend([
571 ("id".to_string(), json!({ "type": "string" })),
572 (
573 "expected_revision".to_string(),
574 json!({ "type": "integer", "minimum": 0 }),
575 ),
576 (
577 "status".to_string(),
578 json!({ "type": "string", "enum": ["completed", "cancelled", "failed"] }),
579 ),
580 ]);
581 object(properties, &["id", "expected_revision"])
582}
583
584fn link_schema() -> Value {
585 let mut properties = base_properties();
586 properties.extend([
587 (
588 "kind".to_string(),
589 json!({
590 "type": "string",
591 "enum": ["blocks", "parent", "related", "supersedes", "derived_from"]
592 }),
593 ),
594 ("from_id".to_string(), json!({ "type": "string" })),
595 ("to_id".to_string(), json!({ "type": "string" })),
596 ]);
597 object(properties, &["kind", "from_id", "to_id"])
598}
599
600fn evidence_schema() -> Value {
601 let mut properties = base_properties();
602 properties.extend([
603 ("id".to_string(), json!({ "type": "string" })),
604 (
605 "expected_revision".to_string(),
606 json!({ "type": "integer", "minimum": 0 }),
607 ),
608 ("evidence".to_string(), evidence_ref_schema()),
609 ]);
610 object(properties, &["id", "expected_revision", "evidence"])
611}
612
613#[cfg(test)]
614#[allow(clippy::expect_used, clippy::unwrap_used)]
615mod tests {
616 use std::collections::BTreeSet;
617 use std::sync::Arc;
618
619 use serde_json::json;
620
621 use crate::{MemoryWorkGraphStore, WorkGraphService, WorkNamespace};
622
623 use super::*;
624
625 #[tokio::test]
626 async fn workgraph_tools_create_and_ready_round_trip() {
627 let service = WorkGraphService::with_scope(
628 Arc::new(MemoryWorkGraphStore::new()),
629 "realm",
630 WorkNamespace::default(),
631 );
632 let created = handle_workgraph_tools_call(
633 &service,
634 "workgraph_create",
635 &json!({ "title": "tool item", "labels": ["a"] }),
636 )
637 .await
638 .expect("create");
639 let id = created["item"]["id"].as_str().expect("id").to_string();
640 let ready =
641 handle_workgraph_tools_call(&service, "workgraph_ready", &json!({ "labels": ["a"] }))
642 .await
643 .expect("ready");
644 assert_eq!(ready["items"][0]["id"].as_str(), Some(id.as_str()));
645 }
646
647 #[test]
648 fn workgraph_tools_list_contains_requested_surface() {
649 let names = workgraph_tools_list()
650 .into_iter()
651 .filter_map(|tool| tool["name"].as_str().map(ToString::to_string))
652 .collect::<BTreeSet<_>>();
653 for expected in [
654 "workgraph_create",
655 "workgraph_get",
656 "workgraph_list",
657 "workgraph_ready",
658 "workgraph_snapshot",
659 "workgraph_events",
660 "workgraph_claim",
661 "workgraph_release",
662 "workgraph_update",
663 "workgraph_block",
664 "workgraph_close",
665 "workgraph_link",
666 "workgraph_add_evidence",
667 ] {
668 assert!(names.contains(expected), "missing {expected}");
669 }
670 }
671
672 #[test]
673 fn workgraph_tool_schemas_do_not_expose_bare_arrays_or_objects() {
674 fn assert_schema_is_provider_safe(path: &str, schema: &Value) {
675 match schema {
676 Value::Object(map) => {
677 let is_array = map.get("type").and_then(Value::as_str) == Some("array");
678 assert!(
679 !is_array || map.contains_key("items"),
680 "{path} is an array schema without items"
681 );
682
683 let is_object = map.get("type").and_then(Value::as_str) == Some("object");
684 assert!(
685 !is_object || map.contains_key("properties"),
686 "{path} is an object schema without properties"
687 );
688
689 for (key, value) in map {
690 assert_schema_is_provider_safe(&format!("{path}.{key}"), value);
691 }
692 }
693 Value::Array(items) => {
694 for (index, value) in items.iter().enumerate() {
695 assert_schema_is_provider_safe(&format!("{path}[{index}]"), value);
696 }
697 }
698 _ => {}
699 }
700 }
701
702 for tool in workgraph_tools_list() {
703 let name = tool["name"].as_str().expect("tool name");
704 assert_schema_is_provider_safe(name, &tool["inputSchema"]);
705 }
706 }
707}