返回 CodeWhale
shape.rs
根目录 / crates / tui / src / tools / github / shape.rs
1 //! Response shaping: what a `gh` payload looks like by the time the model
2 //! sees it.
3 //!
4 //! Long bodies and diffs spill to task artifacts and come back as summaries,
5 //! and every write action returns the task metadata that records what it did.
6
7 use std::path::{Path, PathBuf};
8
9 use chrono::Utc;
10 use serde_json::{Value, json};
11 use uuid::Uuid;
12
13 use crate::task_manager::{TaskArtifactRef, TaskGithubEvent};
14 use crate::tools::spec::{ToolContext, ToolError};
15
16 pub(super) const BODY_ARTIFACT_THRESHOLD: usize = 4_000;
17 pub(super) const DIFF_ARTIFACT_THRESHOLD: usize = 8_000;
18
19 pub(super) fn shape_large_text(
20 context: &ToolContext,
21 mut value: Value,
22 label: &str,
23 threshold: usize,
24 ) -> Result<Value, ToolError> {
25 let body = value
26 .get("body")
27 .and_then(Value::as_str)
28 .map(ToString::to_string);
29 if let Some(body) = body
30 && body.len() > threshold
31 {
32 let artifact = write_artifact_if_needed(context, label, &body, threshold)?;
33 value["body_summary"] = json!(summarize(&body, 900));
34 value["body_artifact"] = json!(artifact);
35 value["body"] = json!(summarize(&body, 1200));
36 }
37 Ok(value)
38 }
39
40 pub(super) fn write_artifact_if_needed(
41 context: &ToolContext,
42 label: &str,
43 content: &str,
44 threshold: usize,
45 ) -> Result<Option<PathBuf>, ToolError> {
46 if content.len() <= threshold {
47 return Ok(None);
48 }
49 let Some(task_id) = context.runtime.active_task_id.as_deref() else {
50 return Ok(None);
51 };
52 if let Some(manager) = context.runtime.task_manager.as_ref() {
53 return manager
54 .write_task_artifact(task_id, label, content)
55 .map(Some)
56 .map_err(|e| ToolError::execution_failed(e.to_string()));
57 }
58 let Some(data_dir) = context.runtime.task_data_dir.as_ref() else {
59 return Ok(None);
60 };
61 let dir = data_dir.join("artifacts").join(task_id);
62 std::fs::create_dir_all(&dir)
63 .map_err(|e| ToolError::execution_failed(format!("create artifact dir: {e}")))?;
64 let absolute = dir.join(format!(
65 "{}_{}.txt",
66 Utc::now().format("%Y%m%dT%H%M%S%.3fZ"),
67 sanitize_filename(label)
68 ));
69 std::fs::write(&absolute, content)
70 .map_err(|e| ToolError::execution_failed(format!("write artifact: {e}")))?;
71 Ok(Some(
72 absolute
73 .strip_prefix(data_dir)
74 .map(Path::to_path_buf)
75 .unwrap_or(absolute),
76 ))
77 }
78
79 pub(super) fn artifact_refs_from_context(content: &str, label: &str) -> Vec<TaskArtifactRef> {
80 let Ok(value) = serde_json::from_str::<Value>(content) else {
81 return Vec::new();
82 };
83 let (path_key, summary_key) = if label.ends_with("_diff") {
84 ("diff_artifact", "diff_summary")
85 } else {
86 ("body_artifact", "body_summary")
87 };
88 let mut refs = Vec::new();
89 collect_artifact_refs(&value, path_key, summary_key, label, &mut refs);
90 refs
91 }
92
93 fn collect_artifact_refs(
94 value: &Value,
95 path_key: &str,
96 summary_key: &str,
97 label: &str,
98 refs: &mut Vec<TaskArtifactRef>,
99 ) {
100 match value {
101 Value::Object(map) => {
102 if let Some(path) = map.get(path_key).and_then(Value::as_str) {
103 let summary = map
104 .get(summary_key)
105 .and_then(Value::as_str)
106 .map(ToString::to_string)
107 .unwrap_or_else(|| format!("GitHub {label} artifact"));
108 refs.push(TaskArtifactRef {
109 label: label.to_string(),
110 path: PathBuf::from(path),
111 summary,
112 created_at: Utc::now(),
113 });
114 }
115 for child in map.values() {
116 collect_artifact_refs(child, path_key, summary_key, label, refs);
117 }
118 }
119 Value::Array(items) => {
120 for child in items {
121 collect_artifact_refs(child, path_key, summary_key, label, refs);
122 }
123 }
124 _ => {}
125 }
126 }
127
128 pub(super) fn github_event_metadata(
129 action: &str,
130 target: &str,
131 number: u64,
132 summary: String,
133 url: Option<String>,
134 artifact: Option<PathBuf>,
135 ) -> Value {
136 let artifacts = artifact
137 .map(|path| {
138 json!([TaskArtifactRef {
139 label: format!("github_{action}"),
140 path,
141 summary: summary.clone(),
142 created_at: Utc::now(),
143 }])
144 })
145 .unwrap_or_else(|| json!([]));
146 json!({
147 "task_updates": {
148 "github_event": TaskGithubEvent {
149 id: format!("gh_{}", &Uuid::new_v4().to_string()[..8]),
150 action: action.to_string(),
151 target: target.to_string(),
152 number,
153 summary,
154 url,
155 recorded_at: Utc::now(),
156 },
157 "artifacts": artifacts
158 }
159 })
160 }
161
162 pub(super) fn summarize(text: &str, limit: usize) -> String {
163 let mut out = String::new();
164 for (idx, ch) in text.chars().enumerate() {
165 if idx >= limit.saturating_sub(3) {
166 out.push_str("...");
167 return out;
168 }
169 if ch.is_control() && ch != '\n' && ch != '\t' {
170 continue;
171 }
172 out.push(ch);
173 }
174 out
175 }
176
177 fn sanitize_filename(input: &str) -> String {
178 let mut out = String::new();
179 for ch in input.chars() {
180 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
181 out.push(ch);
182 } else {
183 out.push('_');
184 }
185 }
186 if out.is_empty() {
187 "artifact".to_string()
188 } else {
189 out
190 }
191 }
192
192 lines RUST