返回 CodeWhale
main.rs
根目录 / crates / memory / src / main.rs
1 use clap::{Parser, Subcommand};
2 use codewhale_memory::{import, protocol::ToolServer, workspace, *};
3 use serde::de::DeserializeOwned;
4 use std::{
5 fs::File,
6 io::{self, Read},
7 path::{Path, PathBuf},
8 };
9
10 #[derive(Parser)]
11 #[command(
12 name = "cw-memory",
13 version,
14 about = "CodeWhale local memory: evidence, review, recall and checkpoints"
15 )]
16 struct Args {
17 /// A NEW database, never the existing native index.sqlite3.
18 #[arg(long)]
19 db: PathBuf,
20 #[arg(long, default_value = "local")]
21 tenant: String,
22 #[arg(long, default_value = "local")]
23 user: String,
24 #[arg(long)]
25 workspace: Option<String>,
26 #[arg(long)]
27 branch: Option<String>,
28 #[arg(long)]
29 session: Option<String>,
30 #[arg(long)]
31 agent: Option<String>,
32 /// Trusted working-tree root for live dependency hashes.
33 #[arg(long)]
34 workspace_root: Option<PathBuf>,
35 /// Explicit local reviewer/admin authority. Never forward model arguments here.
36 #[arg(long, conflicts_with = "read_only")]
37 operator: bool,
38 #[arg(long)]
39 read_only: bool,
40 #[command(subcommand)]
41 command: Command,
42 }
43 #[derive(Subcommand)]
44 enum Command {
45 Status,
46 /// Read a scoped Context Lens snapshot without granting new permissions.
47 Lens {
48 #[arg(long)]
49 trace: Option<String>,
50 #[arg(long)]
51 after: Option<String>,
52 #[arg(long, default_value_t = 100)]
53 limit: usize,
54 },
55 /// Export metadata-only lifecycle observations, optionally as event-v1 JSONL.
56 Events {
57 #[arg(long, default_value_t = 0)]
58 after: i64,
59 #[arg(long, default_value_t = 200)]
60 limit: usize,
61 #[arg(long)]
62 jsonl: bool,
63 },
64 Preferences {
65 id: String,
66 #[arg(long)]
67 revision: i64,
68 #[arg(long)]
69 pinned: bool,
70 #[arg(long)]
71 suppressed: bool,
72 },
73 History {
74 id: String,
75 #[arg(long)]
76 known_at: i64,
77 #[arg(long)]
78 valid_at: i64,
79 },
80 List {
81 #[arg(long)]
82 after: Option<String>,
83 #[arg(long, default_value_t = 100)]
84 limit: usize,
85 },
86 Propose {
87 #[arg(long)]
88 request: String,
89 #[arg(long)]
90 file: PathBuf,
91 },
92 Search {
93 query: String,
94 #[arg(long, default_value_t = 12)]
95 limit: usize,
96 #[arg(long)]
97 include_stale: bool,
98 #[arg(long)]
99 embedding: Option<PathBuf>,
100 },
101 Get {
102 id: String,
103 },
104 Approve {
105 id: String,
106 #[arg(long)]
107 revision: i64,
108 #[arg(long)]
109 validation: Option<PathBuf>,
110 },
111 Reject {
112 id: String,
113 #[arg(long)]
114 revision: i64,
115 },
116 Supersede {
117 old: String,
118 new: String,
119 #[arg(long)]
120 old_revision: i64,
121 #[arg(long)]
122 new_revision: i64,
123 #[arg(long)]
124 validation: Option<PathBuf>,
125 },
126 Forget {
127 id: String,
128 #[arg(long)]
129 revision: i64,
130 },
131 Link {
132 from: String,
133 to: String,
134 #[arg(long, default_value = "related")]
135 relation: String,
136 },
137 Embed {
138 id: String,
139 #[arg(long)]
140 content_hash: String,
141 #[arg(long)]
142 file: PathBuf,
143 },
144 Context {
145 #[arg(default_value = "")]
146 query: String,
147 #[arg(long, default_value_t = 12000)]
148 max_bytes: usize,
149 },
150 CheckpointSave {
151 #[arg(long)]
152 file: PathBuf,
153 #[arg(long)]
154 expected_revision: Option<i64>,
155 },
156 Resume {
157 key: String,
158 },
159 ImportMarkdown {
160 file: PathBuf,
161 },
162 ImportJsonl {
163 file: PathBuf,
164 },
165 Export,
166 Reindex,
167 Expire,
168 /// MCP stdio compatibility server, revisions 2025-06-18 and 2025-11-25.
169 Serve,
170 }
171 fn read_bounded(path: &Path, max: usize) -> Result<String> {
172 let mut bytes = Vec::new();
173 File::open(path)?
174 .take((max + 1) as u64)
175 .read_to_end(&mut bytes)?;
176 if bytes.len() > max {
177 return Err(Error::Invalid("input file exceeds size limit".into()));
178 }
179 String::from_utf8(bytes).map_err(|_| Error::Invalid("input must be UTF-8".into()))
180 }
181 fn read_json<T: DeserializeOwned>(path: &Path) -> Result<T> {
182 Ok(serde_json::from_str(&read_bounded(path, 64 * 1024)?)?)
183 }
184 fn print<T: serde::Serialize>(value: &T) -> Result<()> {
185 let mut out = io::stdout().lock();
186 serde_json::to_writer_pretty(&mut out, value)?;
187 use std::io::Write;
188 out.write_all(b"\n")?;
189 Ok(())
190 }
191 fn run(args: Args) -> Result<()> {
192 let global = Scope::user(args.tenant, args.user);
193 let mut scopes = vec![global.clone()];
194 let mut current = global;
195 if let Some(w) = args.workspace {
196 current = current.workspace(w);
197 scopes.push(current.clone());
198 }
199 if let Some(b) = args.branch {
200 current = current.branch(b);
201 scopes.push(current.clone());
202 }
203 if let Some(s) = args.session {
204 current = current.session(s);
205 scopes.push(current.clone());
206 }
207 if let Some(a) = args.agent {
208 current = current.agent(a);
209 scopes.push(current.clone());
210 }
211 let checkpoint_scope = current.session.as_ref().map(|_| current.clone());
212 let access = if args.operator {
213 Access::operator(scopes)?
214 } else if args.read_only {
215 Access::readonly(scopes)?
216 } else {
217 Access::agent(scopes)?
218 };
219 let mut store = Store::open(&args.db)?;
220 let snapshot = match &args.workspace_root {
221 Some(root) => workspace::snapshot(root, store.dependency_paths(&access)?)?,
222 None => Snapshot::default(),
223 };
224 match args.command {
225 Command::Status => print(&store.status(&access)?),
226 Command::Lens {
227 trace,
228 after,
229 limit,
230 } => print(&store.lens_snapshot(
231 &access,
232 trace.as_deref(),
233 after.as_deref(),
234 limit,
235 &snapshot,
236 )?),
237 Command::Events {
238 after,
239 limit,
240 jsonl,
241 } => {
242 let page = store.event_page(&access, after, limit)?;
243 if jsonl {
244 use std::io::Write;
245 let mut out = io::stdout().lock();
246 for event in &page.events {
247 serde_json::to_writer(&mut out, event)?;
248 out.write_all(b"\n")?;
249 }
250 Ok(())
251 } else {
252 print(&page)
253 }
254 }
255 Command::Preferences {
256 id,
257 revision,
258 pinned,
259 suppressed,
260 } => print(&store.set_preferences(&access, &id, revision, pinned, suppressed)?),
261 Command::History {
262 id,
263 known_at,
264 valid_at,
265 } => print(&store.memory_as_of(&access, &id, known_at, valid_at)?),
266 Command::List { after, limit } => print(&store.list(&access, after.as_deref(), limit)?),
267 Command::Propose { request, file } => {
268 print(&store.capture(&access, &request, read_json(&file)?)?)
269 }
270 Command::Search {
271 query,
272 limit,
273 include_stale,
274 embedding,
275 } => {
276 let embedding = embedding.as_deref().map(read_json).transpose()?;
277 print(&store.recall(
278 &access,
279 &Recall {
280 query,
281 limit,
282 include_stale,
283 embedding,
284 snapshot,
285 ..Recall::default()
286 },
287 )?)
288 }
289 Command::Get { id } => {
290 let memory = store.get(&access, &id)?;
291 print(
292 &serde_json::json!({"freshness":store.freshness(&access,&memory,&snapshot)?,"memory":memory}),
293 )
294 }
295 Command::Approve {
296 id,
297 revision,
298 validation,
299 } => {
300 let receipt: Option<ValidationReceipt> =
301 validation.as_deref().map(read_json).transpose()?;
302 print(&store.approve(&access, &id, revision, receipt.as_ref(), &snapshot)?)
303 }
304 Command::Reject { id, revision } => print(&store.reject(&access, &id, revision)?),
305 Command::Supersede {
306 old,
307 new,
308 old_revision,
309 new_revision,
310 validation,
311 } => {
312 let receipt: Option<ValidationReceipt> =
313 validation.as_deref().map(read_json).transpose()?;
314 print(&store.supersede(
315 &access,
316 (&old, old_revision),
317 (&new, new_revision),
318 receipt.as_ref(),
319 &snapshot,
320 )?)
321 }
322 Command::Forget { id, revision } => print(&store.forget(&access, &id, revision)?),
323 Command::Link { from, to, relation } => {
324 store.link(&access, &from, &to, Relation::parse(&relation)?)?;
325 print(&serde_json::json!({"linked":true}))
326 }
327 Command::Embed {
328 id,
329 content_hash,
330 file,
331 } => {
332 store.set_embedding(&access, &id, &content_hash, &read_json(&file)?)?;
333 print(&serde_json::json!({"stored":true}))
334 }
335 Command::Context { query, max_bytes } => {
336 let report = store.recall(
337 &access,
338 &Recall {
339 query,
340 snapshot,
341 limit: 64,
342 ..Recall::default()
343 },
344 )?;
345 print(&compile_context(
346 &report.hits,
347 &ByteCounter,
348 &ContextBudget {
349 max_units: max_bytes,
350 max_bytes: max_bytes.min(64 * 1024),
351 max_entries: 32,
352 },
353 )?)
354 }
355 Command::CheckpointSave {
356 file,
357 expected_revision,
358 } => print(&store.save_checkpoint(
359 &access,
360 read_json(&file)?,
361 expected_revision,
362 &snapshot,
363 )?),
364 Command::Resume { key } => {
365 let scope =
366 checkpoint_scope.ok_or_else(|| Error::Invalid("--session is required".into()))?;
367 print(&store.resume(&access, &scope, &key, &snapshot)?)
368 }
369 Command::ImportMarkdown { file } => {
370 let text = read_bounded(&file, 1024 * 1024)?;
371 let uri = format!("file://{}", std::fs::canonicalize(&file)?.display());
372 print(&import::markdown(
373 &mut store, &access, &current, &uri, &text,
374 )?)
375 }
376 Command::ImportJsonl { file } => {
377 let text = read_bounded(&file, 16 * 1024 * 1024)?;
378 print(&import::jsonl(&mut store, &access, &current, &text)?)
379 }
380 Command::Export => {
381 store.export_jsonl(&access, io::stdout().lock())?;
382 Ok(())
383 }
384 Command::Reindex => {
385 store.reindex(&access)?;
386 print(&serde_json::json!({"reindexed":true}))
387 }
388 Command::Expire => print(&serde_json::json!({"marked_stale":store.expire(&access)?})),
389 Command::Serve => ToolServer::new(
390 store,
391 access,
392 current,
393 checkpoint_scope,
394 args.workspace_root,
395 )?
396 .serve(io::stdin().lock(), io::stdout().lock()),
397 }
398 }
399 fn main() {
400 if let Err(error) = run(Args::parse()) {
401 eprintln!("{}: {}", error.code(), error);
402 std::process::exit(2);
403 }
404 }
405
405 lines RUST