Skip to content

feat: treat query_knowledge_sources results as sensitive by default#4404

Open
abhinav-m22 wants to merge 25 commits into
archestra-ai:mainfrom
abhinav-m22:fix/untrust-knowledge-results
Open

feat: treat query_knowledge_sources results as sensitive by default#4404
abhinav-m22 wants to merge 25 commits into
archestra-ai:mainfrom
abhinav-m22:fix/untrust-knowledge-results

Conversation

@abhinav-m22
Copy link
Copy Markdown
Contributor

@abhinav-m22 abhinav-m22 commented May 6, 2026

This PR closes a critical prompt-injection and RAG-poisoning escalation path by ensuring that results from the built-in query_knowledge_sources tool are treated as sensitive by default.

Backend

  • Special-cased query_knowledge_sources in TrustedDataPolicyModel.evaluateBulk to skip the automatic trust bypass for built-in tools.
  • Implemented a fail-closed fallback: if the tool record is missing from the database, the system defaults to untrusted rather than trusted.

UI

  • Modified /api/tools/with-assignments to support an explicit inclusion filter for specific built-in tools.
  • Updated the Tool Guardrails page to expose query_knowledge_sources in the tools table while keeping other infrastructure tools hidden.

Docs and Tests

  • Updated codegen-archestra-mcp-server-docs.ts to reflect the trust exception for this tool in the auto-generated platform documentation.
  • Implemented comprehensive unit tests verifying that both standard and branded KB tool results are treated as untrusted and follow normal policy evaluation.
  • Added integration tests proving that restricted tools are successfully blocked when invoked after a KB query.

@abhinav-m22 abhinav-m22 marked this pull request as draft May 6, 2026 14:04
@abhinav-m22 abhinav-m22 marked this pull request as ready for review May 6, 2026 14:29
@abhinav-m22
Copy link
Copy Markdown
Contributor Author

Hi @joeyorlando!
PR is ready for review. The KB query results (query_knowledge_sources) are marked as sensitive by default and exposed the tool in the UI as suggested.

@Konstantinov-Innokentii Konstantinov-Innokentii requested review from joeyorlando and removed request for Konstantinov-Innokentii May 11, 2026 08:50
@Konstantinov-Innokentii
Copy link
Copy Markdown
Contributor

@joeyorlando - assigning you as a reviewer since you own the original issue.

Comment thread platform/backend/src/standalone-scripts/codegen-archestra-mcp-server-docs.ts Outdated
Comment thread platform/backend/src/standalone-scripts/codegen-archestra-mcp-server-docs.ts Outdated
Co-authored-by: Joey Orlando <joseph.t.orlando@gmail.com>
Comment thread docs/pages/platform-agents.md Outdated
Co-authored-by: Joey Orlando <joseph.t.orlando@gmail.com>
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I pushed a few minor changes to platform/backend/src/standalone-scripts/codegen-archestra-mcp-server-docs.ts - you'll just need to rerun pnpm codegen to regenerate the contents here

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sure, done

Comment thread docs/pages/platform-agents.md Outdated
Comment thread docs/pages/platform-knowledge-bases.md Outdated
Co-authored-by: Joey Orlando <joseph.t.orlando@gmail.com>
@joeyorlando joeyorlando changed the title feat: treat query_knowledge_sources results as sensitive by default feat: treat query_knowledge_sources results as sensitive by default May 11, 2026
Copy link
Copy Markdown
Contributor

@joeyorlando joeyorlando left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in terms of "migration", where are the default tool invocation + tool result policies being set/assigned for the query_knowledge_sources tool?

@abhinav-m22
Copy link
Copy Markdown
Contributor Author

abhinav-m22 commented May 12, 2026

in terms of "migration", where are the default tool invocation + tool result policies being set/assigned for the query_knowledge_sources tool?

@joeyorlando missed that point. I was leaning on the evaluator's "no policy = untrusted" fallback at evaluation time.

Just fixed it in seedArchestraTools which is the same startup-seed pattern as migratePlaywrightToolsToDynamicCredential.

It now ensures default rows in both tool_invocation_policies (allow_when_context_is_untrusted) and trusted_data_policies (mark_as_untrusted).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Prompt Injection via Knowledge Base Ingestion

3 participants